diff --git a/AniMAIRE/AniMAIRE.py b/AniMAIRE/AniMAIRE.py index 0a848aa..f9a7da4 100644 --- a/AniMAIRE/AniMAIRE.py +++ b/AniMAIRE/AniMAIRE.py @@ -1,20 +1,13 @@ import numpy as np import datetime as dt -import plotly.express as px -import seaborn as sns -import geopandas -import matplotlib.pyplot as plt import spaceweather as sw import pandas as pd from .utils import get_correctly_formatted_particle_dist_list from .anisotropic_MAIRE_engine.spectralCalculations.rigiditySpectrum import DLRmodelSpectrum, MishevModifiedPowerLawSpectrum, MishevModifiedPowerLawSpectrumSplit from .anisotropic_MAIRE_engine.spectralCalculations.pitchAngleDistribution import gaussianBeeckPitchAngleDistribution, isotropicPitchAngleDistribution, gaussianPitchAngleDistribution - from .anisotropic_MAIRE_engine.generalEngineInstance import generalEngineInstance, default_array_of_lats_and_longs -from scipy.interpolate import NearestNDInterpolator - default_altitudes_in_kft = [0,10,20] + [i for i in range(25, 61 + 1, 3)] def run_from_spectra( @@ -161,151 +154,7 @@ def run_from_DLR_cosmic_ray_model(OULU_count_rate_in_seconds=None, return output_dose_rate_DF -def create_single_dose_map_plot(DF_to_use, - selected_altitude_in_km, - **kwargs): - - if selected_altitude_in_km is not None: - DF_to_use = DF_to_use[round(DF_to_use["altitude (km)"],4) == selected_altitude_in_km] - - if len(DF_to_use) == 0: - raise Exception("Error: specified altitude in kilometers did not match any of the altitudes in kilometers in the inputted DataFrame.") - - doseRateMap = px.scatter(DF_to_use, x="longitude",y="latitude",color="adose", - symbol_sequence=["square"], - range_y=[-90,90], - range_x=[0,360], - **kwargs) - - doseRateMap.update_traces(marker={'size': 10}) - doseRateMap.update_layout(yaxis=dict(scaleanchor="x", scaleratio=1, range=[-90,90])) - doseRateMap.update_xaxes(range=[0,360]) - doseRateMap.update_yaxes(range=[-90,90]) - doseRateMap.update_layout(xaxis_scaleanchor="y") - - doseRateMap.update_layout(autosize=False, - width=800, - height=600) - - doseRateMap.show() - - return doseRateMap - -#def create_single_dose_map_plot_plt(heatmap_DF_to_Plot, hue_range = (2,7), heatmap_s = 63,edgecolor='white', -# dose_type = "adose"): -def create_single_dose_map_plot_plt(heatmap_DF_to_Plot, - hue_range = None, #(0,100), - heatmap_s = 63, - edgecolor=None, - dose_type = "edose", - legend_label=r"Effective dose ($\mu Sv / hr$)", - palette="Spectral_r", - plot_longitude_east=False, - plot_colorbar=True): - - if not (heatmap_DF_to_Plot["altitude (km)"].nunique() == 1): - print() - print("\033[1mWARNING: multiple altitudes were supplied in the input dataframe, therefore only the map for the maximum altitude will be plotted!\033[0m") - print() - heatmap_DF_to_Plot = heatmap_DF_to_Plot.query(f"`altitude (km)` == {heatmap_DF_to_Plot['altitude (km)'].max()}") - - if hue_range is None: - hue_range = (0,heatmap_DF_to_Plot["edose"].max()) - - ############################ creating background world map and dose image - currentFigure = plt.gcf() - - currentFigure.set_figheight(10) - currentFigure.set_figwidth(10) - - #heatmap_DF_to_Plot = pd.read_csv(file_path_to_read, delimiter=',') - heatmap_DF_to_Plot["SEU (Upsets/hr/Gb)"] = heatmap_DF_to_Plot["SEU"] * (60.0 * 60.0) * 1e9 - heatmap_DF_to_Plot["SEL (Latch-ups/hr/device)"] = heatmap_DF_to_Plot["SEL"] * (60.0 * 60.0) - if plot_longitude_east is False: - heatmap_DF_to_Plot["longitudeTranslated"] = heatmap_DF_to_Plot["longitude"].apply(lambda x:x-360.0 if x > 180.0 else x) - else: - heatmap_DF_to_Plot["longitudeTranslated"] = heatmap_DF_to_Plot["longitude"] - - scatterPlotAxis = sns.scatterplot(data=heatmap_DF_to_Plot,x="longitudeTranslated",y="latitude", - hue=dose_type, hue_norm=hue_range, palette=palette, - zorder=10, - marker="s",s=heatmap_s,edgecolor=edgecolor, - legend=False, - )#ax=axToPlotOn) - - if plot_colorbar is True: - norm = plt.Normalize(hue_range[0], hue_range[1]) - sm = plt.cm.ScalarMappable(cmap=palette, norm=norm) - sm.set_array([]) - # Remove the legend and add a colorbar - #scatterPlotAxis.get_legend().remove() - #colorbar = scatterPlotAxis.figure.colorbar(sm,label=legend_label,shrink=0.4) - colorbar = scatterPlotAxis.figure.colorbar(sm,label=legend_label,orientation="horizontal") - else: - colorbar = None - - plt.ylim([-90,90]) - plt.xlim([-175,180]) - plt.grid(True) - plt.xlabel("Longitude (degrees)") - plt.ylabel("Latitude (degrees)") - - world = geopandas.read_file(geopandas.datasets.get_path('naturalearth_lowres')) - world.plot(color="None",edgecolor="black",lw=0.35,ax=scatterPlotAxis,zorder=20) - if plot_longitude_east is True: - world['geometry'] = world['geometry'].translate(xoff=360) - world.plot(color="None",edgecolor="black",lw=0.35,ax=scatterPlotAxis,zorder=20) - plt.xlim([0,355]) - - #################################################################### - - #plt.legend(title=dose_type,loc="center left",bbox_to_anchor=(1.1,0.5)) - - return scatterPlotAxis, colorbar - -def add_colorbar_to_plot(hue_range, palette, legend_label, scatterPlotAxis=None): - norm = plt.Normalize(hue_range[0], hue_range[1]) - sm = plt.cm.ScalarMappable(cmap=palette, norm=norm) - sm.set_array([]) - - # Remove the legend and add a colorbar - #scatterPlotAxis.get_legend().remove() - if scatterPlotAxis is None: - colorbar = plt.colorbar(sm,label=legend_label,shrink=0.4) - else: - colorbar = scatterPlotAxis.figure.colorbar(sm,label=legend_label,shrink=0.4) - return colorbar - -def plot_dose_map_contours(dose_map_to_plot,levels=3,**kwargs): - - dose_map_to_plot_sorted = dose_map_to_plot.sort_values(by=["longitudeTranslated","latitude"]) - - (contour_longs, contour_lats) = np.meshgrid(dose_map_to_plot_sorted["longitudeTranslated"].unique(), - dose_map_to_plot_sorted["latitude"].unique()) - - interp = NearestNDInterpolator(list(zip(dose_map_to_plot_sorted["longitudeTranslated"], dose_map_to_plot_sorted["latitude"])), - dose_map_to_plot_sorted["edose"]) - - contours = plt.contour(contour_longs,contour_lats,interp(contour_longs, contour_lats), - levels=levels,linestyles="dashed",colors="black",zorder=1000,**kwargs) - plt.clabel(contours, inline=True) #,fmt={"fontweight":"bold"}) - -def plot_dose_map(map_to_plot, - plot_title=None, - plot_contours=True, - levels=3, - **kwargs): - - #altitude_to_plot_in_km = altitude_to_plot_in_kft * 0.3048 - - axis_no_colorbar, colorbar = create_single_dose_map_plot_plt(map_to_plot, - **kwargs) - - plt.title(plot_title) - if plot_contours is True: - plot_dose_map_contours(map_to_plot,levels=levels,**kwargs) - return axis_no_colorbar, colorbar diff --git a/AniMAIRE/dose_plotting.py b/AniMAIRE/dose_plotting.py new file mode 100644 index 0000000..c1313f4 --- /dev/null +++ b/AniMAIRE/dose_plotting.py @@ -0,0 +1,152 @@ +import numpy as np +import plotly.express as px +from scipy.interpolate import NearestNDInterpolator +import matplotlib.pyplot as plt +import seaborn as sns +import geopandas + +def plot_dose_map_contours(dose_map_to_plot,levels=3,**kwargs): + + dose_map_to_plot_sorted = dose_map_to_plot.sort_values(by=["longitudeTranslated","latitude"]) + + (contour_longs, contour_lats) = np.meshgrid(dose_map_to_plot_sorted["longitudeTranslated"].unique(), + dose_map_to_plot_sorted["latitude"].unique()) + + interp = NearestNDInterpolator(list(zip(dose_map_to_plot_sorted["longitudeTranslated"], dose_map_to_plot_sorted["latitude"])), + dose_map_to_plot_sorted["edose"]) + + contours = plt.contour(contour_longs,contour_lats,interp(contour_longs, contour_lats), + levels=levels,linestyles="dashed",colors="black",zorder=1000,**kwargs) + plt.clabel(contours, inline=True) #,fmt={"fontweight":"bold"}) + +def create_single_dose_map_plot_plt(heatmap_DF_to_Plot, + hue_range = None, #(0,100), + heatmap_s = 63, + edgecolor=None, + dose_type = "edose", + legend_label=r"Effective dose ($\mu Sv / hr$)", + palette="Spectral_r", + plot_longitude_east=False, + plot_colorbar=True): + + if not (heatmap_DF_to_Plot["altitude (km)"].nunique() == 1): + print() + print("\033[1mWARNING: multiple altitudes were supplied in the input dataframe, therefore only the map for the maximum altitude will be plotted!\033[0m") + print() + heatmap_DF_to_Plot = heatmap_DF_to_Plot.query(f"`altitude (km)` == {heatmap_DF_to_Plot['altitude (km)'].max()}") + + if hue_range is None: + hue_range = (0,heatmap_DF_to_Plot["edose"].max()) + + ############################ creating background world map and dose image + currentFigure = plt.gcf() + + currentFigure.set_figheight(10) + currentFigure.set_figwidth(10) + + #heatmap_DF_to_Plot = pd.read_csv(file_path_to_read, delimiter=',') + heatmap_DF_to_Plot["SEU (Upsets/hr/Gb)"] = heatmap_DF_to_Plot["SEU"] * (60.0 * 60.0) * 1e9 + heatmap_DF_to_Plot["SEL (Latch-ups/hr/device)"] = heatmap_DF_to_Plot["SEL"] * (60.0 * 60.0) + if plot_longitude_east is False: + heatmap_DF_to_Plot["longitudeTranslated"] = heatmap_DF_to_Plot["longitude"].apply(lambda x:x-360.0 if x > 180.0 else x) + else: + heatmap_DF_to_Plot["longitudeTranslated"] = heatmap_DF_to_Plot["longitude"] + + scatterPlotAxis = sns.scatterplot(data=heatmap_DF_to_Plot,x="longitudeTranslated",y="latitude", + hue=dose_type, hue_norm=hue_range, palette=palette, + zorder=10, + marker="s",s=heatmap_s,edgecolor=edgecolor, + legend=False, + )#ax=axToPlotOn) + + if plot_colorbar is True: + norm = plt.Normalize(hue_range[0], hue_range[1]) + sm = plt.cm.ScalarMappable(cmap=palette, norm=norm) + sm.set_array([]) + + # Remove the legend and add a colorbar + #scatterPlotAxis.get_legend().remove() + #colorbar = scatterPlotAxis.figure.colorbar(sm,label=legend_label,shrink=0.4) + colorbar = scatterPlotAxis.figure.colorbar(sm,label=legend_label,orientation="horizontal") + else: + colorbar = None + + plt.ylim([-90,90]) + plt.xlim([-175,180]) + plt.grid(True) + plt.xlabel("Longitude (degrees)") + plt.ylabel("Latitude (degrees)") + + world = geopandas.read_file(geopandas.datasets.get_path('naturalearth_lowres')) + world.plot(color="None",edgecolor="black",lw=0.35,ax=scatterPlotAxis,zorder=20) + if plot_longitude_east is True: + world['geometry'] = world['geometry'].translate(xoff=360) + world.plot(color="None",edgecolor="black",lw=0.35,ax=scatterPlotAxis,zorder=20) + plt.xlim([0,355]) + + #################################################################### + + #plt.legend(title=dose_type,loc="center left",bbox_to_anchor=(1.1,0.5)) + + return scatterPlotAxis, colorbar + +def plot_dose_map(map_to_plot, + plot_title=None, + plot_contours=True, + levels=3, + **kwargs): + + #altitude_to_plot_in_km = altitude_to_plot_in_kft * 0.3048 + + axis_no_colorbar, colorbar = create_single_dose_map_plot_plt(map_to_plot, + **kwargs) + + plt.title(plot_title) + + if plot_contours is True: + plot_dose_map_contours(map_to_plot,levels=levels,**kwargs) + + return axis_no_colorbar, colorbar + +def add_colorbar_to_plot(hue_range, palette, legend_label, scatterPlotAxis=None): + norm = plt.Normalize(hue_range[0], hue_range[1]) + sm = plt.cm.ScalarMappable(cmap=palette, norm=norm) + sm.set_array([]) + + # Remove the legend and add a colorbar + #scatterPlotAxis.get_legend().remove() + if scatterPlotAxis is None: + colorbar = plt.colorbar(sm,label=legend_label,shrink=0.4) + else: + colorbar = scatterPlotAxis.figure.colorbar(sm,label=legend_label,shrink=0.4) + return colorbar + +def create_single_dose_map_plotly(DF_to_use, + selected_altitude_in_km, + **kwargs): + + if selected_altitude_in_km is not None: + DF_to_use = DF_to_use[round(DF_to_use["altitude (km)"],4) == selected_altitude_in_km] + + if len(DF_to_use) == 0: + raise Exception("Error: specified altitude in kilometers did not match any of the altitudes in kilometers in the inputted DataFrame.") + + doseRateMap = px.scatter(DF_to_use, x="longitude",y="latitude",color="adose", + symbol_sequence=["square"], + range_y=[-90,90], + range_x=[0,360], + **kwargs) + + doseRateMap.update_traces(marker={'size': 10}) + doseRateMap.update_layout(yaxis=dict(scaleanchor="x", scaleratio=1, range=[-90,90])) + doseRateMap.update_xaxes(range=[0,360]) + doseRateMap.update_yaxes(range=[-90,90]) + doseRateMap.update_layout(xaxis_scaleanchor="y") + + doseRateMap.update_layout(autosize=False, + width=800, + height=600) + + doseRateMap.show() + + return doseRateMap \ No newline at end of file diff --git a/README.md b/README.md index 733eb15..940ef65 100644 --- a/README.md +++ b/README.md @@ -211,17 +211,29 @@ Any particular altitudes the user wants to use can be supplied to `altitudes_in_ If you want to perform calculations only at a specific set of latitudes and longitudes you should use the `array_of_lats_and_longs` argument, supplying it as a 2 dimensional `list` or numpy array, where the first column refers to latitudes and the second column refers to longitudes. All longitudes in this case should be specified in terms of longitude east (i.e. 0.00 degrees - 359.99 degrees). **Using the `array_of_lats_and_longs` argument significantly speeds up the running of `AniMAIRE` if you're only interested in a small number of coordinates, so its use is highly recommended in those situations.** -There are many ways you could plot this data. An example function, `create_single_dose_map_plot`, has been supplied in `AniMAIRE` that uses plotly to plot the dose rates across Earth (i.e. as a function of latitude and longitude) at a given altitude. Its specification is the following: +There are many ways you could plot this data. Several example functions,`plot_dose_map` and `create_single_dose_map_plotly`, have been supplied in `AniMAIRE` that uses matplotlib or plotly to plot the dose rates across Earth (i.e. as a function of latitude and longitude) at a given altitude. Both of these functions are available in the `dose_plotting` submodule supplied with AniMAIRE. Their specifications are the following: ``` -def create_single_dose_map_plot(DF_to_use, +def plot_dose_map(map_to_plot, + plot_title=None, + plot_contours=True, + levels=3, + **kwargs) +``` +for matplotlib plots, where map_to_plot is the Pandas DataFrame outputted by a run of `AniMAIRE`, with only one altitude selected. `plot_contours` can be switched on or off to control whether contours are added to the plot, and `levels` can be used to specify to number of contours and/or dose rates for the contours to correspond to. `hue_range` can also be supplied with a 2-value tuple to specify the limits of the colorbar to be plotted with the plot. + +To generate a plotly plot, you can run +``` +def create_single_dose_map_plotly(DF_to_use, selected_altitude_in_km) ``` where `DF_to_use` is the Pandas DataFrame outputted by a run of `AniMAIRE` and altitude is one of the altitudes in kilometers supplied to/outputted by the run. To use this function to create a map of the isotropic situation as given as an example above, you could run ``` -AniMAIRE.plot_dose_map(test_isotropic_dose_rates.query("`altitude (km)` == 12.1920"), +from AniMAIRE import dose_plotting + +dose_plotting.plot_dose_map(test_isotropic_dose_rates.query("`altitude (km)` == 12.1920"), hue_range=(0,9)) ``` @@ -295,17 +307,21 @@ In this case printing `test_anisotropic_dose_rates` should output: which will produce the following plot when ``` -AniMAIRE.plot_dose_map(test_anisotropic_dose_rates.query("`altitude (km)` == 12.1920"), +from AniMAIRE import dose_plotting + +dose_plotting.plot_dose_map(test_anisotropic_dose_rates.query("`altitude (km)` == 12.1920"), hue_range=(0,9)) ``` -is run: +is run, as was described previously in this README for isotropic plotting: ![anisotropic_test_plot](https://raw.githubusercontent.com/ssc-maire/AniMAIRE-public/main/Anisotropic_40kft_example.svg) you can also produce similar plotly plots if you prefer plotly to matplotlib using: ``` -anisotropic_dose_rate_map = AniMAIRE.create_single_dose_map_plot(test_anisotropic_dose_rates, +from AniMAIRE import dose_plotting + +anisotropic_dose_rate_map = dose_plotting.create_single_dose_map_plotly(test_anisotropic_dose_rates, selected_altitude_in_km = 12.1920) ``` which generates the following plot: