fixed some minor internal bugs for alpha particle inputs

ssc-maire · Mar 13, 2024 · 8c02929 · 8c02929
1 parent f48e0e9
commit 8c02929
Show file tree

Hide file tree

Showing 5 changed files with 3,678 additions and 60 deletions.
diff --git a/AniMAIRE/AniMAIRE.py b/AniMAIRE/AniMAIRE.py
@@ -202,6 +202,15 @@ def create_single_dose_map_plot_plt(heatmap_DF_to_Plot,
                                     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()

diff --git a/AniMAIRE/anisotropic_MAIRE_engine/singleParticleEngineInstance.py b/AniMAIRE/anisotropic_MAIRE_engine/singleParticleEngineInstance.py
@@ -247,7 +247,7 @@ def runOverSpecifiedAltitudes(self):
         # assign each asymptotic direction a weighting factor in accordance with its pitch angle
         print("assigning pitch angle weighting factors...")
         df_with_weighting_factors_full_angles = acquireWeightingFactors(self.dfOfAllAsymptoticDirections,
-                                                            self.particle_distribution.momentum_distribution)
+                                                            self.particle_distribution)
         #######################################################
 
         # print("df_with_weighting_factors:")

diff --git a/AniMAIRE/anisotropic_MAIRE_engine/spectralCalculations/particleSpecies.py b/AniMAIRE/anisotropic_MAIRE_engine/spectralCalculations/particleSpecies.py
@@ -3,10 +3,12 @@
 class particleSpecies():
 
     particle_atomic_number_dict = {"proton":1,"alpha":2}
+    particle_atomic_mass_dict = {"proton":1,"alpha":4}
 
     def __init__(self, particleName, atomicNumber=None):
         self.particleName = particleName
         if atomicNumber == None:
             atomicNumber = self.particle_atomic_number_dict[particleName]
 
-        self.atomicNumber = atomicNumber
+        self.atomicNumber = atomicNumber
+        self.atomicMass = self.particle_atomic_mass_dict[particleName]
diff --git a/AniMAIRE/anisotropic_MAIRE_engine/utils/AsymptoticDirectionDataframe.py b/AniMAIRE/anisotropic_MAIRE_engine/utils/AsymptoticDirectionDataframe.py
@@ -40,7 +40,7 @@ def generate_asymp_dir_DF(dataframeToFillFrom:pd.DataFrame, IMFlatitude:float, I
 
     new_asymp_dir_DF = dataframeToFillFrom.copy()
 
-    new_asymp_dir_DF["Energy"] = PRCT.convertParticleRigidityToEnergy(dataframeToFillFrom["Rigidity"], particleMassAU = 1, particleChargeAU = 1)
+    #new_asymp_dir_DF["Energy"] = PRCT.convertParticleRigidityToEnergy(dataframeToFillFrom["Rigidity"], particleMassAU = 1, particleChargeAU = 1)
 
     print("assigning asymptotic coordinates")
     if cache == False:
@@ -132,7 +132,9 @@ def get_pitch_angle_for_DF_row(IMFlatitude, IMFlongitude, dataframeRow, datetime
     pitch_angle_for_row = calculatePitchAngle_from_IMF_dir(interplanetary_mag_field, rowInSpacepyCoords, datetime_to_run_across_UTC)
     return pitch_angle_for_row
 
-def acquireWeightingFactors(asymptotic_direction_DF:pd.DataFrame, momentaDist):
+def acquireWeightingFactors(asymptotic_direction_DF:pd.DataFrame, particle_distribution):
+
+    momentaDist = particle_distribution.momentum_distribution
 
     new_asymptotic_direction_DF = asymptotic_direction_DF.copy()
 
@@ -170,9 +172,13 @@ def acquireWeightingFactors(asymptotic_direction_DF:pd.DataFrame, momentaDist):
     #print(new_asymptotic_direction_DF)
 
     print("calculating energy + pitch combined weighting factors...")
+    new_asymptotic_direction_DF["Energy"] = PRCT.convertParticleRigidityToEnergy(new_asymptotic_direction_DF["Rigidity"], 
+                                                                      particleMassAU = particle_distribution.particle_species.atomicMass, 
+                                                                      particleChargeAU = particle_distribution.particle_species.atomicNumber)
     energySpectrum = PRCT.convertParticleRigiditySpecToEnergySpec(new_asymptotic_direction_DF["Rigidity"],
                                                                   new_asymptotic_direction_DF["fullRigidityPitchWeightingFactor"],
-                                                                  particleMassAU=1, particleChargeAU=1)
+                                                                  particleMassAU=particle_distribution.particle_species.atomicMass, 
+                                                                  particleChargeAU=particle_distribution.particle_species.atomicNumber)
 
     new_asymptotic_direction_DF["fullEnergyPitchWeightingFactor"] = energySpectrum["Energy distribution values"]
     #print(new_asymptotic_direction_DF)

diff --git a/AniMAIRE_examples.ipynb b/AniMAIRE_examples.ipynb