diff --git a/aetherpy/io/read_routines.py b/aetherpy/io/read_routines.py index e4c955e..7196709 100644 --- a/aetherpy/io/read_routines.py +++ b/aetherpy/io/read_routines.py @@ -157,12 +157,12 @@ def read_aether_netcdf_header(filename, epoch_name='time'): nlats - number of latitude grids nalts - number of altitude grids vars - list of data variable names - time - list of datetimes with data + time - datetime for the data filename - filename of file containing header data See Also -------- - read_aether_header + read_aether_headers """ @@ -234,7 +234,7 @@ def read_aether_ascii_header(filename): See Also -------- - read_aether_header + read_aether_headers """ @@ -297,7 +297,7 @@ def read_aether_one_binary_file(header, ifile, vars_to_read): ifile : int Integer corresponding to filename in the header dict vars_to_read : list - List of desired variable names to read + List of desired variable name indices to read Returns ------- @@ -523,7 +523,7 @@ def read_gitm_file(filename, file_vars=None): filename : str GITM file to read file_vars : list or NoneType - List of desired variable names to read or None to read all + List of desired variable name indices to read or None to read all (default=None) Returns @@ -604,3 +604,38 @@ def read_gitm_file(filename, file_vars=None): (data["nlons"], data["nlats"], data["nalts"]), order="F") return data + + +def read_headers(filelist, has_header=False, is_gitm=False): + """Load header data from a list of model files. + + Parameters + ---------- + filelist : list-like + List of model filenames to load + has_header : bool + Flag indicating that a separate header file contains the header data, + as is the case for binary files (default=False) + is_gitm : bool + Flag indicating whether this is a GITM file, if True, or and Aether + file, if False (default=False) + + Returns + ------- + header : dict + Dictionary of header data + + See Also + -------- + read_gitm_headers, read_aether_ascii_header, read_aether_headers + + """ + + # Load the header data using the appropriate routine + if is_gitm and not has_header: + header = read_gitm_headers(filelist, finds=0) + else: + filetype = "ascii" if has_header else "netcdf" + header = read_aether_headers(filelist, filetype=filetype) + + return header diff --git a/aetherpy/plot/__init__.py b/aetherpy/plot/__init__.py index 6fb1f1d..f4cc367 100644 --- a/aetherpy/plot/__init__.py +++ b/aetherpy/plot/__init__.py @@ -5,3 +5,4 @@ from aetherpy.plot import data_prep from aetherpy.plot import movie_routines +from aetherpy.plot import standard_plots diff --git a/aetherpy/plot/data_prep.py b/aetherpy/plot/data_prep.py index 731367d..0224ed2 100644 --- a/aetherpy/plot/data_prep.py +++ b/aetherpy/plot/data_prep.py @@ -1,12 +1,12 @@ #!/usr/bin/env python # Copyright 2020, the Aether Development Team (see doc/dev_team.md for members) # Full license can be found in License.md -""" Utilities for slicing and preparing data for plotting -""" +"""Utilities for slicing and preparing data for plotting.""" import numpy as np from aetherpy import logger +from aetherpy.io import read_routines def get_cut_index(lons, lats, alts, cut_val, isgrid=False, cut_coord='alt'): @@ -80,7 +80,10 @@ def get_cut_index(lons, lats, alts, cut_val, isgrid=False, cut_coord='alt'): icut = cut_val else: if cut_val < z_coord.min() or cut_val > z_coord.max(): - raise ValueError('Requested cut is outside the coordinate range') + raise ValueError(''.join(['Requested cut is outside the ', + 'coordinate range [{:} '.format(cut_val), + 'not in {:} to {:}]'.format( + z_coord.min(), z_coord.max())])) icut = abs(z_coord - cut_val).argmin() @@ -91,8 +94,9 @@ def get_cut_index(lons, lats, alts, cut_val, isgrid=False, cut_coord='alt'): # Warn the user if they selected a suspect index if cut_coord == "alt": if icut > len(z_coord) - 3: - logger.warning(''.join(['Requested altitude slice is above ', - 'the recommended upper limit'])) + logger.warning(''.join(['Requested altitude slice is above the ', + 'recommended upper limit of {:}'.format( + len(z_coord) - 3)])) else: if icut == 0 or icut == len(z_coord) - 1: logger.warning(''.join(['Requested ', cut_coord, ' slice is ', @@ -163,3 +167,175 @@ def calc_tec(alt, ne, ialt_min=2, ialt_max=-4): tec /= 1.0e16 return tec + + +def load_data_for_plotting(filelist, plot_var, cut_var='alt', cut_val=400, + has_header=False, is_gitm=False, winds=False, + tec=False): + """Load model data for plotting. + + Parameters + ---------- + filelist : list-like + List of model filenames to load + plot_var : str or int + Variable name or index of data to plot + cut_var : str + Variable name along which data should be sliced (default='alt') + cut_val : int or float + Data value along that will be held constant + has_header : bool + Flag indicating that a separate header file contains the header data, + as is the case for binary files (default=False) + is_gitm : bool + Flag indicating whether this is a GITM file, if True, or and Aether + file, if False (default=False) + winds : bool + If True prepare winds or drifts for quiver-style plotting + (default=False) + tec : bool + If True calculate the TEC for plotting (default=False) + + Returns + ------- + all_times : array-like + 1D array of datetimes to plot + all_2dim_data : array-like + 3D array of data variables (time, x, and y) + all_winds_x : array-like or NoneType + 1D array of winds along x-axis or None if `winds` is False + all_winds_y : array-like or NoneType + 1D array of winds along y-axis or None if `winds` is False + icut : int + Cut index + x_coord : array-like + Array of data to include along the x-axis + y_coord : array-like + Array of data to include along the y-axis + z_val : float + Data value for cut index + var_name : str + Long name of the data variable + + Raises + ------ + ValueError + If a bad `plot_var` value is provided + + """ + + # Load the header data + header = read_routines.read_headers(filelist, has_header=has_header, + is_gitm=is_gitm) + + if is_gitm and has_header: + is_gitm = False + + if isinstance(plot_var, int): + if plot_var >= len(header["vars"]): + raise ValueError("requested bad variable index: {:d}>={:d}".format( + plot_var, len(header["vars"]))) + elif plot_var not in header["vars"]: + raise ValueError("unknown variable requested: {:s} not in {:}".format( + plot_var, header["vars"])) + + # Define the plotting inputs + var_list = ['lon', 'lat', 'z', plot_var] + plot_vars = [0, 1, 2, plot_var] + + # Update plotting variables to include the wind, if desired + if winds: + if cut_var in ['alt', 'lat']: + plot_vars.append(16) + var_list.append('Zonal Wind') + else: + plot_vars.append(17) + var_list.append('Meridional Wind') + + if cut_var in ['lat', 'lon']: + plot_vars.append(18) + var_list.append('Vertical Wind') + else: + plot_vars.append(17) + var_list.append('Meridional Wind') + + all_winds_x = [] + all_winds_y = [] + else: + all_winds_x = None + all_winds_y = None + + # Prepare to load the desired file data + all_2dim_data = [] + all_times = [] + var_name = None + convert_lat = True + convert_lon = True + + for j, filename in enumerate(filelist): + # Read in the data file + if is_gitm: + data = read_routines.read_gitm_file(filename, plot_vars) + var_list[3] = data['vars'][3] + else: + if has_header: + data = read_routines.read_aether_one_binary_file(header, j, + plot_vars) + var_list[3] = data['vars'][3] + else: + data = read_routines.read_aether_file(filename, var_list) + plot_vars[3] = 3 + + if data['units'][0].find('degree') >= 0: + convert_lon = False + + if data['units'][1].find('degree') >= 0: + convert_lat = False + + # Get the z-variable name, if needed + if var_name is None: + vkey = 'long_name' if 'long_name' in data.keys() else 'vars' + var_name = data[vkey][plot_vars[3]] + + # For the first file, initialize the necessary plotting data + if j == 0: + # Get 1D arrays for the coordinates + alts = data[2][0, 0] / 1000.0 # Convert from m to km + + # Convert from rad to deg, if necessary, and reshape lat and lon + lons = data[0][:, 0, 0] + lats = data[1][0, :, 0] + + if convert_lat: + lats = np.degrees(lats) + + if convert_lon: + lons = np.degrees(lons) + + # Find the desired index to cut along to get a 2D slice + icut, cut_data, x_coord, y_coord, z_val = get_cut_index( + lons, lats, alts, cut_val, cut_coord=cut_var, isgrid=False) + + # Save the time data + all_times.append(data["time"]) + + # Save the z-axis data + if tec: + all_2dim_data.append(calc_tec(alts, data[plot_vars[3]], 2, -4)) + else: + all_2dim_data.append(data[plot_vars[3]][cut_data]) + + if winds: + all_winds_x.append(data[plot_vars[-1]][cut_data]) + all_winds_y.append(data[plot_vars[-1]][cut_data]) + + # Convert data list to a numpy array + all_2dim_data = np.array(all_2dim_data) + + if winds: + all_winds_x = np.array(all_winds_x) + all_winds_y = np.array(all_winds_y) + + # Return data for plotting + return(all_times, all_2dim_data, all_winds_x, all_winds_y, icut, x_coord, + y_coord, z_val, var_name) diff --git a/aetherpy/plot/standard_plots.py b/aetherpy/plot/standard_plots.py new file mode 100644 index 0000000..8e0a1ad --- /dev/null +++ b/aetherpy/plot/standard_plots.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python +# Copyright 2020, the Aether Development Team (see doc/dev_team.md for members) +# Full license can be found in License.md +"""Standard plot functions for model results.""" + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + +from aetherpy import logger +from aetherpy.plot import movie_routines +from aetherpy.utils import time_conversion + + +def plot_model_results(all_times, all_2dim_data, all_winds_x, all_winds_y, + plot_var, cut_var, icut, x_coord, y_coord, z_val, + log_scale=False, cmap=None, prefix='', ext='png', + movie_rate=0, save_plots=True, close_plots=True): + """Default plot for visualizing model data. + + Parameters + ---------- + all_times : array-like + 1D array of datetimes to plot + all_2dim_data : array-like + 3D array of data variables (time, x, and y) + all_winds_x : array-like or NoneType + 1D array of winds along x-axis or None if `winds` is False + all_winds_y : array-like or NoneType + 1D array of winds along y-axis or None if `winds` is False + plot_var : str + Variable name of data to plot + cut_var : str + Variable name along which data should be sliced + icut : int + Cut index + x_coord : array-like + Array of data to include along the x-axis + y_coord : array-like + Array of data to include along the y-axis + z_val : float + Data value for cut index + log_scale : bool + Use a base-10 logarithmic scale for the data (default=False) + cmap : matplotlib.colors.ListedColormap or NoneType + Colormap for plotting, if None defaults to matplotlib's 'plasma' for + asymmetric data and 'bwr' for symmetric data (default=mpl.cm.plasma) + prefix : str + Prefix for output figure filenames (default='') + ext : str + Extention for output figure files, do not include period (default='png') + movie_rate : int + Movie frame rate or zero if no movie is desired (default=0) + save_plots : bool + Save plots to file (default=True) + close_plots : bool + Close figure handles for plots (default=True) + + Returns + ------- + figs : list + List of open figure handles, empty list if `close_plots` is True + + See Also + -------- + aetherpy.plots.movie_routines.save_movie, matplotlib.cm + + """ + + # If desired, take the log of the data + if log_scale: + all_2dim_data = np.log10(all_2dim_data) + + # Define plotting limits + mini = all_2dim_data.min() * 0.99 + + if mini < 0.0: + symmetric = True + maxi = abs(all_2dim_data).max() * 1.05 + mini = -maxi + else: + symmetric = False + maxi = all_2dim_data.max() * 1.01 + + if cut_var == 'alt': + mask_north = ((y_coord >= 40) & (y_coord <= 90.0)) + mask_south = ((y_coord <= -40) & (y_coord >= -90.0)) + plot_north = mask_north.max() + plot_south = mask_south.max() + + if plot_north: + maxi_north = abs(all_2dim_data[:, :, mask_north]).max() * 1.05 + + if symmetric: + mini_north = -maxi_north + else: + mini_north = all_2dim_data[:, :, mask_north].min() * 0.95 + + if plot_south: + maxi_south = abs(all_2dim_data[:, :, mask_south]).max() * 1.05 + + if symmetric: + mini_south = -maxi_south + else: + mini_south = all_2dim_data[:, :, mask_south].min() * 0.95 + + # Define the color scale, if desired + if cmap is None: + cmap = mpl.cm.bwr if symmetric else mpl.cm.plasma + + # Define plot range + minx = (x_coord[1] + x_coord[2]) / 2.0 + maxx = (x_coord[-2] + x_coord[-3]) / 2.0 + miny = (y_coord[1] + y_coord[2]) / 2.0 + maxy = (y_coord[-2] + y_coord[-3]) / 2.0 + + # Prepare the output filename + if save_plots: + filename = "{:s}{:s}_{:s}{:03d}".format( + prefix, plot_var.replace(" ", "_"), cut_var, icut) + + if movie_rate > 0: + img_file_fmt = movie_routines.setup_movie_dir(filename) + else: + img_file_fmt = ''.join([filename, '_{:}.', ext]) + + # Create a plot for each time + figs = [] + for itime, utime in enumerate(all_times): + # Initialize the figure + figs.append(plt.figure(constrained_layout=False, + tight_layout=True, figsize=(10, 8.5))) + + gs1 = mpl.gridspec.GridSpec(nrows=2, ncols=2, wspace=0.0, hspace=0) + gs = mpl.gridspec.GridSpec(nrows=2, ncols=2, wspace=-0.05, left=0.02, + right=0.95, top=0.99, bottom=0.05) + ax = figs[-1].add_subplot(gs1[1, 0:2]) + + # Plot the global data set (square plot at bottom if three plots): + dx = (x_coord[1] - x_coord[0]) / 2.0 + xp = np.append(x_coord - dx, x_coord[-1:] + dx) + dy = (y_coord[1] - y_coord[0]) / 2.0 + yp = np.append(y_coord - dy, y_coord[-1] + dy) + con = ax.pcolormesh(xp, yp, all_2dim_data[itime].transpose(), + vmin=mini, vmax=maxi, cmap=cmap, shading='auto') + + # Add the winds, if desired + if all_winds_x is not None and all_winds_y is not None: + ax.quiver(x_coord, y_coord, all_winds_x[itime].transpose(), + all_winds_y[itime].transpose()) + ax.set_ylim([miny, maxy]) + ax.set_xlim([minx, maxx]) + + # Set the labels and aspect ratio + ax.set_title("{:s}; {:s}: {:.2f} {:s}".format( + utime.strftime("%d %b %Y %H:%M:%S UT"), cut_var, z_val, + 'km' if cut_var == 'alt' else r'$^\circ$'), fontsize='medium') + ax.set_xlabel(r'Latitude ($^\circ$)' if cut_var == 'lon' + else r'Longitude ($^\circ$)') + ax.set_ylabel(r'Latitude ($^\circ$)' if cut_var == 'alt' + else r'Altitude (km)') + if cut_var == 'alt': + ax.set_aspect(1.0) + + # Set the colorbar + cbar = figs[-1].colorbar(con, ax=ax, shrink=0.75, pad=0.02) + cbar.set_label(plot_var, rotation=90) + + # If this is an altitude slice, add polar dials + if cut_var == 'alt' and (plot_north or plot_south): + # Set the common inputs + shift = time_conversion.calc_time_shift(utime) + + xlabels = [] + xlabelpos = [] + ylabels = [r'80$^\circ$', r'70$^\circ$', r'60$^\circ$', + r'50$^\circ$'] + + ylabelpos = [10.0, 20.0, 30.0, 40.0] + xticks = np.arange(0, 2 * np.pi, np.pi / 2.0) + yticks = np.arange(10, 50, 10) + + if plot_north: + # Top Left Graph Northern Hemisphere + ax2 = figs[-1].add_subplot(gs[0, 0], projection='polar') + yp = 90.0 - y_coord[mask_north] + dy = (np.floor(100.0 * (yp[1] - yp[0])) / 100.0) / 2.0 + yp = np.append(yp - dy, yp[-1] + dy) + xp = np.radians(x_coord + shift - 90.0) + dx = (xp[1] - xp[0]) / 2 + xp = np.append(xp - dx, xp[-1] + dx) + z = all_2dim_data[itime][:, mask_north].transpose() + conn = ax2.pcolormesh(xp, yp, z, shading='auto', cmap=cmap, + vmin=mini_north, vmax=maxi_north) + ax2.set_xticks(xlabelpos) + ax2.set_xticklabels(xlabels) + ax2.text(-np.pi / 2, 46.0, '00 LT', verticalalignment='top', + horizontalalignment='center') + ax2.text(np.pi / 2, 46.0, '12 LT', verticalalignment='bottom', + horizontalalignment='center') + ax2.set_yticks(ylabelpos) + ax2.set_yticklabels(ylabels) + ax2.grid(linestyle=':', color='black') + ax2.set_xticks(xticks) + ax2.set_yticks(yticks) + ax2.set_ylim([0, 45]) + figs[-1].colorbar(conn, ax=ax2, shrink=0.5, pad=0.01) + + if plot_south: + # Top Right Graph Southern Hemisphere + rad, theta = np.meshgrid(90.0 + y_coord[mask_south], + np.radians(x_coord + shift - 90.0)) + ax3 = figs[-1].add_subplot(gs[0, 1], projection='polar') + + yp = 90.0 + y_coord[mask_south] + dy = (int(100.0 * (yp[1] - yp[0])) / 100.0) / 2.0 + yp = np.append(yp - dy, yp[-1] + dy) + xp = np.radians(x_coord + shift - 90.0) + dx = (xp[1] - xp[0]) / 2.0 + xp = np.append(xp - dx, xp[-1] + dx) + z = all_2dim_data[itime][:, mask_south].transpose() + cons = ax3.pcolormesh(xp, yp, z, shading='auto', cmap=cmap, + vmin=mini_south, vmax=maxi_south) + ax3.set_xticks(xlabelpos) + ax3.set_xticklabels(xlabels) + ax3.text(-np.pi / 2, 46.0, '00 LT', verticalalignment='top', + horizontalalignment='center') + ax3.text(np.pi / 2, 46.0, '12 LT', verticalalignment='bottom', + horizontalalignment='center') + ax3.set_yticks(ylabelpos) + ax3.set_yticklabels(ylabels) + ax3.grid(linestyle=':', color='black') + ax3.set_xticks(xticks) + ax3.set_yticks(yticks) + ax3.set_ylim([0, 45]) + figs[-1].colorbar(cons, ax=ax3, shrink=0.5, pad=0.01) + + # Format the output filename + fmt_input = itime if movie_rate > 0 else utime.strftime( + '%y%m%d_%H%M%S') + outfile = img_file_fmt.format(fmt_input) + + # Save the output file + if save_plots: + logger.info("Writing file : ", outfile) + figs[-1].savefig(outfile) + + if close_plots: + plt.close(figs[-1]) + figs.pop() + + # Create a movie, if desired + if save_plots and movie_rate > 0: + movie_routines.save_movie(filename, ext=ext, rate=movie_rate) + + return figs diff --git a/aetherpy/run_plot_model_results.py b/aetherpy/run_plot_model_results.py index 3e507da..8b9ef9b 100755 --- a/aetherpy/run_plot_model_results.py +++ b/aetherpy/run_plot_model_results.py @@ -10,10 +10,8 @@ import os import re -from aetherpy import logger from aetherpy.io import read_routines -from aetherpy.plot import data_prep -from aetherpy.plot import movie_routines +from aetherpy import plot from aetherpy.utils import inputs from aetherpy.utils import time_conversion @@ -37,20 +35,20 @@ def get_help(file_vars=None): """ mname = os.path.join( - os.path.commonpath([inputs.__file__, data_prep.__file__]), + os.path.commonpath([inputs.__file__, plot.data_prep.__file__]), 'run_plot_model_results.py') if __name__ == '__main__' else __name__ help_str = ''.join(['Usage:\n{:s} -[flags] [filenames]\n'.format(mname), 'Flags:\n', ' -help : print this message, include filename ', 'for variable names and indices\n', - ' -var=number : index of variable to plot\n', + ' -var=string : name of variable to plot\n', + ' -ivar=number : index of variable to plot\n', ' -cut=alt, lat, or lon : which cut you would ', 'like\n', - ' -alt=number : alt in km or grid number ', - '(closest)\n', + ' -alt=number : alt in km (closest)\n', ' -lat=number : latitude in degrees (closest)\n', - ' -lon=number: longitude in degrees (closest)\n', + ' -lon=number : longitude in degrees (closest)\n', ' -log : plot the log of the variable\n', ' -winds : overplot winds\n', ' -tec : plot the TEC variable\n', @@ -59,7 +57,8 @@ def get_help(file_vars=None): ' -ext=str : figure or movie extension\n', 'At end, list the files you want to plot. This code ', 'should work with either GITM files (*.bin) or Aether', - ' netCDF files (*.nc)']) + ' netCDF files (*.nc)\n', + 'If both -var and -ivar are provided, -var is used']) if file_vars is not None: help_str += "File Variables (index, name):\n" @@ -93,14 +92,14 @@ def get_command_line_args(argv): """ # Initialize the arguments to their default values - args = {'filelist': [], 'log': False, 'var': 15, 'alt': 400, 'tec': False, - 'lon': np.nan, 'lat': np.nan, 'cut': 'alt', 'winds': False, - 'diff': False, 'is_gitm': False, 'has_header': False, 'movie': 0, + args = {'filelist': [], 'log': False, 'var': '', 'ivar': -1, 'alt': 400, + 'tec': False, 'lon': np.nan, 'lat': np.nan, 'cut': 'alt', + 'winds': False, 'is_gitm': False, 'has_header': False, 'movie': 0, 'ext': 'png'} - arg_type = {'filelist': list, 'log': bool, 'var': int, 'alt': int, - 'tec': bool, 'lon': float, 'lat': float, 'cut': str, - 'help': bool, 'winds': bool, 'diff': bool, 'is_gitm': bool, + arg_type = {'filelist': list, 'log': bool, 'var': str, 'ivar': int, + 'alt': int, 'tec': bool, 'lon': float, 'lat': float, + 'cut': str, 'help': bool, 'winds': bool, 'is_gitm': bool, 'has_header': bool, 'tec': bool, 'movie': int, 'ext': str} # If there is input, set default help to False @@ -171,7 +170,15 @@ def get_command_line_args(argv): # ---------------------------------------------------------------------------- # Define the main plotting routine -def plot_model_results(): +def main(): + """Main routine for creating standard plots. + + See Also + -------- + aetherpy.plots.standard_plots.plot_model_results + + """ + # Get the input arguments args = get_command_line_args(inputs.process_command_line_input()) @@ -183,270 +190,36 @@ def plot_model_results(): is_gitm = args['is_gitm'] has_header = args['has_header'] - if is_gitm and not has_header: - header = read_routines.read_gitm_headers(args["filelist"], finds=0) - else: - if has_header: - header = read_routines.read_aether_ascii_header(args["filelist"]) - is_gitm = False - else: - header = read_routines.read_aether_header(args["filelist"]) - # If help is requested for a specific file, return it here - if args['help']: + if args['help'] or (len(args['var']) == 0 and args['ivar'] < 0): + header = read_routines.read_headers(args['filelist'], + has_header=has_header, + is_gitm=is_gitm) help_str = get_help(header['vars']) print(help_str) return - if args["var"] >= len(header["vars"]): - raise ValueError("requested variable doesn't exist: {:d}>{:d}".format( - args["var"], len(header["vars"]))) - - # Define the plotting inputs - plot_vars = [0, 1, 2, args["var"]] - - # Update plotting variables to include the wind, if desired - if args["winds"]: - plot_vars.append(16 if args['cut'] in ['alt', 'lat'] else 17) - plot_vars.append(18 if args['cut'] in ['lat', 'lon'] else 17) - all_winds_x = [] - all_winds_y = [] - - # Prepare to load the desired file data - all_2dim_data = [] - all_times = [] - - for j, filename in enumerate(args['filelist']): - # Read in the data file - if is_gitm: - data = read_routines.read_gitm_file(filename, plot_vars) - ivar = args["var"] - else: - if j == 0: - var_list = [] - for pvar in plot_vars: - var_list.append(header["vars"][pvar]) - if has_header: - data = read_routines.read_aether_one_binary_file(header, j, - plot_vars) - ivar = args["var"] - else: - data = read_routines.read_aether_file(filename, var_list) - ivar = 3 - - # For the first file, initialize the necessary plotting data - if j == 0: - # Get 1D arrays for the coordinates - alts = data[2][0][0] / 1000.0 # Convert from m to km - lons = np.degrees(data[0][:, 0, 0]) # Convert from rad to deg - lats = np.degrees(data[1][0, :, 0]) # Convert from rad to deg - - # Find the desired index to cut along to get a 2D slice - icut, cut_data, x_pos, y_pos, z_val = data_prep.get_cut_index( - lons, lats, alts, args[args['cut']], args['cut']) - - # Save the time data - all_times.append(data["time"]) - - # Save the z-axis data - if args["tec"]: - all_2dim_data.append(data_prep.calc_tec(alts, data[ivar], 2, -4)) - else: - all_2dim_data.append(data[ivar][cut_data]) - - if args["winds"]: - all_winds_x.append(data[plot_vars[-1]][cut_data]) - all_winds_y.append(data[plot_vars[-1]][cut_data]) - - # Convert data list to a numpy array - all_2dim_data = np.array(all_2dim_data) - - if args["winds"]: - all_winds_x = np.array(all_winds_x) - all_winds_y = np.array(all_winds_y) - - # If desired, take the log of the data - if args['log']: - all_2dim_data = np.log10(all_2dim_data) - - # Define plotting limits - symmetric = False - cmap = mpl.cm.plasma - - maxi = all_2dim_data.max() * 1.01 - mini = all_2dim_data.min() * 0.99 - - if mini < 0.0: - symmetric = True - cmap = mpl.cm.bwr - maxi = abs(all_2dim_data).max() * 1.05 - mini = -maxi - - if args['cut'] == 'alt': - mask_north = ((y_pos >= 40) & (y_pos <= 90.0)) - mask_south = ((y_pos <= -40) & (y_pos >= -90.0)) - plot_north = mask_north.max() - plot_south = mask_south.max() - - if plot_north: - maxi_north = abs(all_2dim_data[:, :, mask_north]).max() * 1.05 - - if symmetric: - mini_north = -maxi_north - else: - mini_north = all_2dim_data[:, :, mask_north].min() * 0.95 - - if plot_south: - maxi_south = abs(all_2dim_data[:, :, mask_south]).max() * 1.05 - - if symmetric: - mini_south = -maxi_south - else: - mini_south = all_2dim_data[:, :, mask_south].min() * 0.95 - - # Define plot range - minx = (x_pos[1] + x_pos[2]) / 2.0 - maxx = (x_pos[-2] + x_pos[-3]) / 2.0 - miny = (y_pos[1] + y_pos[2]) / 2.0 - maxy = (y_pos[-2] + y_pos[-3]) / 2.0 + if len(args['var']) == 0: + plot_var = args['ivar'] + else: + plot_var = args['var'] - # Prepare the output filename - filename = "var{:02d}_{:s}{:03d}".format(args["var"], args['cut'], icut) + # Load the data needed for plotting + (all_times, all_2dim_data, all_winds_x, all_winds_y, icut, x_coord, + y_coord, z_val, var_name) = plot.data_prep.load_data_for_plotting( + args['filelist'], plot_var, args['cut'], cut_val=args[args['cut']], + has_header=has_header, is_gitm=is_gitm, winds=args['winds'], + tec=args['tec']) - if args['movie'] > 0: - img_file_fmt = movie_routines.setup_movie_dir(filename) - else: - img_file_fmt = ''.join([filename, '_{:}.', args['ext']]) - - # Create a plot for each time - for itime, utime in enumerate(all_times): - # Initialize the figure - fig = plt.figure(constrained_layout=False, - tight_layout=True, figsize=(10, 8.5)) - - gs1 = mpl.gridspec.GridSpec(nrows=2, ncols=2, wspace=0.0, hspace=0) - gs = mpl.gridspec.GridSpec(nrows=2, ncols=2, wspace=-0.05, left=0.02, - right=0.95, top=0.99, bottom=0.05) - ax = fig.add_subplot(gs1[1, 0:2]) - - # Plot the global data set (square plot at bottom if three plots): - - dx = (x_pos[1] - x_pos[0]) / 2.0 - xp = np.append(x_pos - dx, x_pos[-1:] + dx) - dy = (y_pos[1] - y_pos[0]) / 2.0 - yp = np.append(y_pos - dy, y_pos[-1] + dy) - con = ax.pcolormesh(xp, yp, all_2dim_data[itime].transpose(), - vmin=mini, vmax=maxi, cmap=cmap, shading='auto') - - # Add the winds, if desired - if args["winds"]: - ax.quiver(x_pos, y_pos, all_winds_x[itime].transpose(), - all_winds_y[itime].transpose()) - ax.set_ylim([miny, maxy]) - ax.set_xlim([minx, maxx]) - - # Set the labels and aspect ratio - ax.set_title("{:s}; {:s}: {:.2f} {:s}".format( - utime.strftime("%d %b %Y %H:%M:%S UT"), args['cut'], z_val, - 'km' if args['cut'] == 'alt' else r'$^\circ$')) - ax.set_xlabel(r'Latitude ($^\circ$)' if args['cut'] == 'lon' - else r'Longitude ($^\circ$)') - ax.set_ylabel(r'Latitude ($^\circ$)' if args['cut'] == 'alt' - else r'Altitude (km)') - if args['cut'] == 'alt': - ax.set_aspect(1.0) - - # Set the colorbar - cbar = fig.colorbar(con, ax=ax, shrink=0.75, pad=0.02) - cbar.set_label(header["vars"][args["var"]], rotation=90) - - # If this is an altitude slice, add polar dials - if args['cut'] == 'alt' and (plot_north or plot_south): - # Set the common inputs - shift = time_conversion.calc_time_shift(utime) - - xlabels = [] - xlabelpos = [] - ylabels = [r'80$^\circ$', r'70$^\circ$', r'60$^\circ$', - r'50$^\circ$'] - - ylabelpos = [10.0, 20.0, 30.0, 40.0] - xticks = np.arange(0, 2 * np.pi, np.pi / 2.0) - yticks = np.arange(10, 50, 10) - - if plot_north: - # Top Left Graph Northern Hemisphere - ax2 = fig.add_subplot(gs[0, 0], projection='polar') - yp = 90.0 - y_pos[mask_north] - dy = (np.floor(100.0 * (yp[1] - yp[0])) / 100.0) / 2.0 - yp = np.append(yp - dy, yp[-1] + dy) - xp = np.radians(x_pos + shift - 90.0) - dx = (xp[1] - xp[0]) / 2 - xp = np.append(xp - dx, xp[-1] + dx) - z = all_2dim_data[itime][:, mask_north].transpose() - conn = ax2.pcolormesh(xp, yp, z, shading='auto', cmap=cmap, - vmin=mini_north, vmax=maxi_north) - ax2.set_xticks(xlabelpos) - ax2.set_xticklabels(xlabels) - ax2.text(-np.pi / 2, 46.0, '00 LT', verticalalignment='top', - horizontalalignment='center') - ax2.text(np.pi / 2, 46.0, '12 LT', verticalalignment='bottom', - horizontalalignment='center') - ax2.set_yticks(ylabelpos) - ax2.set_yticklabels(ylabels) - ax2.grid(linestyle=':', color='black') - ax2.set_xticks(xticks) - ax2.set_yticks(yticks) - ax2.set_ylim([0, 45]) - fig.colorbar(conn, ax=ax2, shrink=0.5, pad=0.01) - - if plot_south: - # Top Right Graph Southern Hemisphere - rad, theta = np.meshgrid(90.0 + y_pos[mask_south], - np.radians(x_pos + shift - 90.0)) - ax3 = fig.add_subplot(gs[0, 1], projection='polar') - - yp = 90.0 + y_pos[mask_south] - dy = (int(100.0 * (yp[1] - yp[0])) / 100.0) / 2.0 - yp = np.append(yp - dy, yp[-1] + dy) - xp = np.radians(x_pos + shift - 90.0) - dx = (xp[1] - xp[0]) / 2.0 - xp = np.append(xp - dx, xp[-1] + dx) - z = all_2dim_data[itime][:, mask_south].transpose() - cons = ax3.pcolormesh(xp, yp, z, shading='auto', cmap=cmap, - vmin=mini_south, vmax=maxi_south) - ax3.set_xticks(xlabelpos) - ax3.set_xticklabels(xlabels) - ax3.text(-np.pi / 2, 46.0, '00 LT', verticalalignment='top', - horizontalalignment='center') - ax3.text(np.pi / 2, 46.0, '12 LT', verticalalignment='bottom', - horizontalalignment='center') - ax3.set_yticks(ylabelpos) - ax3.set_yticklabels(ylabels) - ax3.grid(linestyle=':', color='black') - ax3.set_xticks(xticks) - ax3.set_yticks(yticks) - ax3.set_ylim([0, 45]) - fig.colorbar(cons, ax=ax3, shrink=0.5, pad=0.01) - - # Format the output filename - fmt_input = itime if args['movie'] > 0 else utime.strftime( - '%y%m%d_%H%M%S') - outfile = img_file_fmt.format(fmt_input) - - # Save the output file - logger.info("Writing file : ", outfile) - fig.savefig(outfile) - plt.close(fig) - - # Create a movie, if desired - if args['movie'] > 0: - movie_routines.save_movie(filename, ext=args['ext'], - rate=args['movie']) + # Plot the data using the specified keyword arguments from the command line + plot.standard_plots.plot_model_results( + all_times, all_2dim_data, all_winds_x, all_winds_y, var_name, + args['cut'], icut, x_coord, y_coord, z_val, log_scale=args['log'], + ext=args['ext'], movie_rate=args['movie']) return # Needed to run main script as the default executable from the command line if __name__ == '__main__': - plot_model_results() + main() diff --git a/aetherpy/tests/test_io_read.py b/aetherpy/tests/test_io_read.py index 4c2960c..f4d0be8 100644 --- a/aetherpy/tests/test_io_read.py +++ b/aetherpy/tests/test_io_read.py @@ -68,7 +68,8 @@ def eval_header(self, file_list=True): for htime in self.header[hkey]: assert isinstance(htime, dt.datetime), \ - "unexpected time format for: {:}".format(htime) + "unexpected time format for: {:} from {:}".format( + htime, self.header[hkey]) elif hkey == 'filename': if file_list: assert len(self.header[hkey]) > 0 @@ -177,8 +178,12 @@ def test_read_aether_netcdf_header_bad_file(self): @pytest.mark.parametrize('fbase', ['3DALL_*.nc', '3DBFI_*.nc']) @pytest.mark.parametrize('ftype', ['netcdf']) - @pytest.mark.parametrize('finds', [(-1), (None), ([0, 1]), (slice(1))]) - def test_read_aether_headers(self, fbase, ftype, finds): + @pytest.mark.parametrize('finds, is_single', [(-1, True), (None, True), + (-1, False), (None, False), + ([0, 1], False), + (slice(1), True), + (slice(1), False)]) + def test_read_aether_headers(self, fbase, ftype, finds, is_single): """Test successful Aether header reading. Parameters @@ -189,12 +194,18 @@ def test_read_aether_headers(self, fbase, ftype, finds): File type finds : int, NoneType, list, slice File indexers + is_single : bool + Uses a single file as a string input if True, or a list if False """ filenames = glob(os.path.join(self.test_dir, fbase)) + if is_single: + filenames = filenames[0] + self.header = read_routines.read_aether_headers(filenames, finds, ftype) + self.eval_header(file_list=True) return