Source code for pydivide.create_model_maps

# Copyright 2017 Regents of the University of Colorado. All Rights Reserved.
# Released under the MIT license.
# This software was developed at the University of Colorado's Laboratory for Atmospheric and Space Physics.
# Verify current version before use at: https://github.com/MAVENSDC/Pydivide

import numpy as np
import scipy
import os
from .utilities import mvn_kp_sc_traj_xyz
from scipy import interpolate, spatial
from .read_model_results import read_model_results


[docs]def create_model_maps(altitude, variable=None, model=None, file=None, numcontours=25, fill=False, ct='viridis', # https://matplotlib.org/examples/color/colormaps_reference.html transparency=1, nearest=False, linear=True, savefig=True): ''' Generates a .png contour map of a model at a specific altitude. These can be used as a background in map2d. The models must be downloaded manually from the SDC website: https://lasp.colorado.edu/maven/sdc/public/pages/models.html. Parameters: altitude: int Specified altitude of output map. variable: str Plots specified chemical species (Appendix A). model: dict Model variable produced from prior call to read_model_results. file: str If model not provided (produced from read_model_results), full path to model can be set and read. numContours: int Specifies number of contour lines. Default is 25. fill: bool If True, fills in contour levels instead of generating lines. ct: str Sets color table. Valid color tables can be found here: https://matplotlib.org/examples/color/colormaps_reference.html transparency: int, float Sets transparency between [0,1] inclusive. 0 is completely transparent, and 1 is completely opaque. nearest: bool If True, instead of interpolating nearby values, this returns the value of the nearest neighbor altitude. linear: bool If True, performs linear interpolation between 2 altitude layers. saveFig: bool If True, saves figure as .png file. Returns: None Examples: >>> # Interpolate all model tracers to spacecraft trajectory using nearest neighbor interpolation. >>> pydivide.create_model_maps(altitude=170, file = '<dir_path>/MAMPS_LS180_F130_081216.nc', variable='geo_x', saveFig=True) ''' import matplotlib.pyplot as plt if model is None and file is None: print("Please input either a model or the file path/name to a model.") return if file is not None: model = read_model_results(file) if variable is None or variable not in model.keys(): print("Variable not entered or not found, Please select one from: ") index = 0 name_index_dict = {} for name in model: if name.lower() == 'dim': continue if name.lower() == 'meta': continue print(index, ": ", name) name_index_dict[index] = name index += 1 chosen_var = False while not chosen_var: i_choice = input("Enter Selection (integer): ") try: i_choice = int(i_choice) dataname = name_index_dict[i_choice].lower() print('You chose the variable {}'.format(dataname)) chosen_var = True except ValueError: print('You must print the integer associated with your selection, try again.') continue else: dataname = variable.lower() mars_radius = model['meta']['mars_radius'] lats = np.arange(181) - 90 lons = np.arange(361) - 180 sc_lat_mso = np.repeat(lats, len(lons)) sc_lon_mso = np.tile(lons, len(lats)) r = np.full(len(sc_lon_mso), altitude + mars_radius) sc_alt_array = np.full(len(sc_lon_mso), altitude) sc_mso_x = r * np.sin((90 - sc_lat_mso) * (np.pi / 180)) * np.cos(sc_lon_mso * (np.pi / 180)) sc_mso_y = r * np.sin((90 - sc_lat_mso) * (np.pi / 180)) * np.sin(sc_lon_mso * (np.pi / 180)) sc_mso_z = r * np.cos((90 - sc_lat_mso) * (np.pi / 180)) sc_path = np.array([sc_mso_x, sc_mso_y, sc_mso_z]).T if nearest: interp_method = 'nearest' else: interp_method = 'linear' if model is None and file is None: print("Please input either a model dictionary from read_model_results, or a model file.") return if 'lon' in model['dim']: if 'mso' == model['meta']['coord_sys'].lower(): # Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows lat_mso_model = model['dim']['lat'] lon_mso_model = model['dim']['lon'] alt_mso_model = model['dim']['alt'] lat_array = np.repeat(lat_mso_model, len(lon_mso_model)) lon_array = np.tile(lon_mso_model, len(lat_mso_model)) data_points = np.transpose(np.array([lon_array, lat_array])) index = 0 for point in sc_path: r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2) alt_mso = altitude lat_mso = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180)) lon_mso = np.arctan2(point[1], point[0]) / (np.pi / 180) sc_path[index] = np.array([lon_mso, lat_mso, alt_mso]) index += 1 latlon_triangulation = spatial.Delaunay(data_points) for var in model: if var.lower() != dataname: continue print("Interpolating variable " + var) if var == 'dim' or var == 'meta': continue # Rearrange the data to lon/lat/alt data = model[var]['data'] dim_order_array = [0, 1, 2] for j in [0, 1, 2]: if model[var]['dim_order'][j] == 'longitude': dim_order_array[0] = j elif model[var]['dim_order'][j] == 'latitude': dim_order_array[1] = j elif model[var]['dim_order'][j] == 'altitude': dim_order_array[2] = j data_new = np.transpose(data, dim_order_array) # Build an array of values that correspond to the points in data point values = np.empty([len(lat_mso_model) * len(lon_mso_model), len(alt_mso_model)]) index = 0 for alt in range(0, len(alt_mso_model)): for lat in range(0, len(lat_mso_model)): for lon in range(0, len(lon_mso_model)): values[index, alt] = data_new[lon][lat][alt] index += 1 index = 0 x = np.empty(len(sc_path)) index = 0 for sc_pos in sc_path: if sc_pos[2] > np.max(alt_mso_model): x[index] = np.NaN index += 1 continue if sc_pos[2] < np.min(alt_mso_model): x[index] = np.NaN index += 1 continue sorted_x_distance = np.argsort(np.abs(alt_mso_model - sc_pos[2])) alti1 = sorted_x_distance[0] if nearest: x[index] = interpolate.griddata(data_points, values[:, alti1], [sc_pos[0], sc_pos[1]], method='nearest') else: if alt_mso_model[alti1] < sc_pos[2]: alti2 = alti1 + 1 else: temp = alti1 - 1 alti2 = alti1 alti1 = temp # Interpolate through space first_val_calc = interpolate.LinearNDInterpolator(latlon_triangulation, values[:, alti1]) second_val_calc = interpolate.LinearNDInterpolator(latlon_triangulation, values[:, alti2]) first_val = first_val_calc([sc_pos[0], sc_pos[1]]) second_val = second_val_calc([sc_pos[0], sc_pos[1]]) delta_1 = sc_pos[2] - alt_mso_model[alti1] delta_2 = alt_mso_model[alti2] - sc_pos[2] delta_tot = alt_mso_model[alti2] - alt_mso_model[alti1] x[index] = ((first_val * delta_2) + (second_val * delta_1)) / delta_tot index += 1 tracer = np.array(x) if 'geo' == model['meta']['coord_sys'].lower(): # Build the Matrix that transforms GEO to MSO coordinates ls_rad = model['meta']['ls'] * np.pi / 180 rads_tilted_y = 25.19 * np.sin(ls_rad) * np.pi / 180 rads_tilted_x = -25.19 * np.cos(ls_rad) * np.pi / 180 lonsubsol_rad = -model['meta']['longsubsol'] * np.pi / 180 z_rotation = np.matrix([[np.cos(lonsubsol_rad), -np.sin(lonsubsol_rad), 0], [np.sin(lonsubsol_rad), np.cos(lonsubsol_rad), 0], [0, 0, 1]]) y_rotation = np.matrix([[np.cos(rads_tilted_y), 0, np.sin(rads_tilted_y)], [0, 1, 0], [-np.sin(rads_tilted_y), 0, np.cos(rads_tilted_y)]]) x_rotation = np.matrix([[1, 0, 0], [0, np.cos(rads_tilted_x), -np.sin(rads_tilted_x)], [0, np.sin(rads_tilted_x), np.cos(rads_tilted_x)]]) geo_to_mso_matrix = np.dot(x_rotation, np.dot(y_rotation, z_rotation)) # Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows lat_geo_model = model['dim']['lat'] lon_geo_model = model['dim']['lon'] alt_geo_model = model['dim']['alt'] alt_array = np.repeat(alt_geo_model, len(lon_geo_model) * len(lat_geo_model)) lat_array = np.tile(np.repeat(lat_geo_model, len(lon_geo_model)), len(alt_geo_model)) lon_array = np.tile(lon_geo_model, len(lat_geo_model) * len(alt_geo_model)) data_points = np.transpose(np.array([lon_array, lat_array, alt_array])) # Convert to GEO coordinates, then to MSO index = 0 for point in data_points: r = point[2] + mars_radius x = r * np.sin((90 - point[1]) * np.pi / 180) * np.cos(point[0] * np.pi / 180) y = r * np.sin((90 - point[1]) * np.pi / 180) * np.sin(point[0] * np.pi / 180) z = r * np.cos((90 - point[1]) * np.pi / 180) data_points[index] = np.dot(geo_to_mso_matrix, np.array([x, y, z])) index += 1 # Convert to MSO Lon/Lat/Alt in order to weight the interpolation better lat_mso = np.empty(len(lon_geo_model) * len(lat_geo_model)) lon_mso = np.empty(len(lon_geo_model) * len(lat_geo_model)) index = 0 for point in data_points: r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2) lat_mso[index] = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180)) lon_mso[index] = np.arctan2(point[1], point[0]) / (np.pi / 180) index += 1 if index >= len(lon_geo_model) * len(lat_geo_model): break latlon_points = np.transpose(np.array([lon_mso, lat_mso])) index = 0 for point in sc_path: r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2) alt_mso = altitude lat_mso = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180)) lon_mso = np.arctan2(point[1], point[0]) / (np.pi / 180) sc_path[index] = np.array([lon_mso, lat_mso, alt_mso]) index += 1 latlon_triangulation = spatial.Delaunay(latlon_points) # Loop through the variables in the model for var in model: if var.lower() != dataname: continue print("Interpolating variable " + var) # Rearrange the data to lon/lat/alt data = model[var]['data'] dim_order_array = [0, 1, 2] for j in [0, 1, 2]: if model[var]['dim_order'][j] == 'longitude': dim_order_array[0] = j elif model[var]['dim_order'][j] == 'latitude': dim_order_array[1] = j elif model[var]['dim_order'][j] == 'altitude': dim_order_array[2] = j data_new = np.transpose(data, dim_order_array) # Build an array of values that correspond to the points in data point values = np.empty([len(lat_geo_model) * len(lon_geo_model), len(alt_geo_model)]) index = 0 for alt in range(0, len(alt_geo_model)): for lat in range(0, len(lat_geo_model)): for lon in range(0, len(lon_geo_model)): values[index, alt] = data_new[lon][lat][alt] index += 1 index = 0 x = np.empty(len(sc_path)) index = 0 for sc_pos in sc_path: if sc_pos[2] > np.max(alt_geo_model): x[index] = np.NaN index += 1 continue if sc_pos[2] < np.min(alt_geo_model): x[index] = np.NaN index += 1 continue sorted_x_distance = np.argsort(np.abs(alt_geo_model - sc_pos[2])) alti1 = sorted_x_distance[0] if nearest: x[index] = interpolate.griddata(latlon_points, values[:, alti1], [sc_pos[0], sc_pos[1]], method='nearest') else: if alt_geo_model[alti1] < sc_pos[2]: alti2 = alti1 + 1 else: temp = alti1 - 1 alti2 = alti1 alti1 = temp # Interpolate through space first_val_calc = interpolate.LinearNDInterpolator(latlon_triangulation, values[:, alti1]) second_val_calc = interpolate.LinearNDInterpolator(latlon_triangulation, values[:, alti2]) first_val = first_val_calc([sc_pos[0], sc_pos[1]]) second_val = second_val_calc([sc_pos[0], sc_pos[1]]) delta_1 = sc_pos[2] - alt_geo_model[alti1] delta_2 = alt_geo_model[alti2] - sc_pos[2] delta_tot = alt_geo_model[alti2] - alt_geo_model[alti1] x[index] = ((first_val * delta_2) + (second_val * delta_1)) / delta_tot index += 1 tracer = np.array(x) else: if 'mso' == model['meta']['coord_sys'].lower(): # Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows x_mso_model = model['dim']['x'] y_mso_model = model['dim']['y'] z_mso_model = model['dim']['z'] # Loop through the variables in the model for var in model: if var.lower() != dataname: continue # Rearrange the data to lon/lat/alt data = model[var]['data'] dim_order_array = [0, 1, 2] for j in [0, 1, 2]: if model[var]['dim_order'][j] == 'x': dim_order_array[0] = j elif model[var]['dim_order'][j] == 'y': dim_order_array[1] = j elif model[var]['dim_order'][j] == 'z': dim_order_array[2] = j data_new = np.transpose(data, dim_order_array) # Interpolate through space tracer = mvn_kp_sc_traj_xyz(x_mso_model, y_mso_model, z_mso_model, data_new, sc_mso_x, sc_mso_y, sc_mso_z, nn=interp_method) xi, yi = np.linspace(sc_lon_mso.min(), sc_lon_mso.max(), 300), np.linspace(sc_lat_mso.min(), sc_lat_mso.max(), 300) xi, yi = np.meshgrid(xi, yi) zi = scipy.interpolate.griddata((sc_lon_mso, sc_lat_mso), tracer, (xi, yi), method=interp_method) if savefig: fig = plt.figure() ax = fig.add_subplot(1, 1, 1) if fill: plt.contourf(xi, yi, zi, numcontours, alpha=transparency, cmap=ct, extent=(-180, -90, 180, 90)) else: cs = plt.contour(xi, yi, zi, numcontours, alpha=transparency, cmap=ct) plt.clabel(cs, inline=1, fontsize=7, fmt='%1.0f') extent = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) plt.axis('off') save_name = "ModelData_" + dataname + "_" + str(altitude) + "km" if fill: save_name = save_name + "_filled" print("Saving file as:" + os.path.join(os.path.dirname(file), save_name + ".png")) plt.savefig(os.path.join(os.path.dirname(file), save_name+".png"), transparent=False, bbox_inches=extent, pad_inches=0, dpi=150) plt.show(block=True) else: # convert back to east longitude (as in model file) return dict(lon=xi[0, :], elon=((xi[0, :] + 360) % 360), longsubsol=model['meta']['longsubsol'], dec=model['meta']['dec'], Ls=model['meta']['ls'], lat=yi[:, 0], param=zi)