Merge pull request #9 from ChrisSWDavis/main

fixed another tzinfo issue
ssc-maire · Feb 21, 2025 · 2b126b0 · 2b126b0
2 parents 8ff00b0 + 0681f6b
commit 2b126b0
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 3 deletions.
diff --git a/AniMAIRE/AniMAIRE.py b/AniMAIRE/AniMAIRE.py
@@ -9,6 +9,8 @@
 from .anisotropic_MAIRE_engine.spectralCalculations.pitchAngleDistribution import gaussianBeeckPitchAngleDistribution, isotropicPitchAngleDistribution, gaussianPitchAngleDistribution
 from .anisotropic_MAIRE_engine.generalEngineInstance import generalEngineInstance, default_array_of_lats_and_longs
 
+import logging
+
 def run_from_spectra(
         proton_rigidity_spectrum: Optional[Callable[[float], float]] = None,
         alpha_rigidity_spectrum: Optional[Callable[[float], float]] = None,
@@ -84,6 +86,10 @@ def run_from_spectra(
     if date_and_time is None:
         date_and_time = dt.datetime.utcnow()
 
+    if date_and_time.tzinfo is None:
+        logging.warning("The inputted date and time does not have timezone info. Assuming UTC.")
+        date_and_time = date_and_time.replace(tzinfo=dt.timezone.utc)
+
     if Kp_index is None:
         Kp_index = get_kp_index(date_and_time)
 

diff --git a/AniMAIRE/utils.py b/AniMAIRE/utils.py
@@ -33,7 +33,10 @@ def get_correctly_formatted_particle_dist_list(proton_rigidity_spectrum,
     return list_of_particle_distributions
 
 def get_kp_index(date_and_time: dt.datetime) -> int:
-    return sw.ap_kp_3h()[pd.to_datetime(sw.ap_kp_3h().index) < date_and_time].iloc[-1]["Kp"]
+    kp_data = sw.ap_kp_3h()
+    kp_data.index = kp_data.index.tz_localize('UTC')  # Localize the index to UTC
+
+    return round(kp_data[pd.to_datetime(kp_data.index) < date_and_time].iloc[-1]["Kp"])
 
 default_altitudes_in_kft = [0,10,20] + [i for i in range(25, 61 + 1, 3)]
 

diff --git a/pytests/test_anisotropic_MAIRE.py b/pytests/test_anisotropic_MAIRE.py
@@ -201,7 +201,7 @@ def test_Common_spec_max_asymp_dir():
     IN_GITHUB_ACTIONS,
     reason="Skipping optional test. Set the environment variable RUN_OPTIONAL_TESTS to run this test."
 )
-def test_DLR_spec():
+def test_DLR_spec_no_tzinfo():
     expected_dlr_spec = [
         [-28.3, -92.7, 0.0, 0.058988868823598634, 0.06080286996669734, 0.09344092993619182, 0.004531869853011256, 0.002768310853422817, 0.0019206239073194524, 3.26731705969241e-16, 3.2673170596924094e-11],
         [-28.3, -92.7, 0.3048, 0.06348032253799271, 0.0651336490962403, 0.09987392145246152, 0.005276616600637716, 0.0033975138199494118, 0.002356729649253225, 4.0141608835156814e-16, 4.014160883515682e-11],
@@ -243,6 +243,52 @@ def test_DLR_spec():
 
     assert np.allclose(np.array(result), np.array(expected_dlr_spec))
 
+@pytest.mark.skipif(
+    IN_GITHUB_ACTIONS,
+    reason="Skipping optional test. Set the environment variable RUN_OPTIONAL_TESTS to run this test."
+)
+def test_DLR_spec():
+    expected_dlr_spec = [
+        [-28.3, -92.7, 0.0, 0.058988868823598634, 0.06080286996669734, 0.09344092993619182, 0.004531869853011256, 0.002768310853422817, 0.0019206239073194524, 3.26731705969241e-16, 3.2673170596924094e-11],
+        [-28.3, -92.7, 0.3048, 0.06348032253799271, 0.0651336490962403, 0.09987392145246152, 0.005276616600637716, 0.0033975138199494118, 0.002356729649253225, 4.0141608835156814e-16, 4.014160883515682e-11],
+        [-28.3, -92.7, 0.6096, 0.06962530767408609, 0.07110977194012219, 0.10641008345663887, 0.006480742649458201, 0.004224252220127024, 0.0029715737299096236, 4.988695917906835e-16, 4.988695917906835e-11],
+        [-28.3, -92.7, 0.9144000000000001, 0.07685126882212685, 0.07833490279826147, 0.11671506485042812, 0.008140541391866922, 0.005290665368927664, 0.0037000334511939946, 6.237867068469103e-16, 6.237867068469101e-11],
+        [-28.3, -92.7, 1.2192, 0.08544527921710764, 0.08723433764061217, 0.12857229159878877, 0.010087151785442248, 0.0065957849826050255, 0.004579691089970698, 7.755801396150202e-16, 7.755801396150202e-11],
+        [-28.3, -92.7, 1.524, 0.09556610301769881, 0.09727093263130687, 0.1371439952258924, 0.012441988077917883, 0.008098934990302172, 0.005669393635363651, 9.520667277693957e-16, 9.520667277693957e-11],
+        [-28.3, -92.7, 1.8288000000000002, 0.10723704504976893, 0.10896753077268706, 0.1530910695744342, 0.015146686567144341, 0.009980012035880111, 0.006921676990472123, 1.1681976333089315e-15, 1.1681976333089315e-10],
+        [-28.3, -92.7, 2.1336, 0.1208472571456993, 0.12250055788157016, 0.16554757537900006, 0.018366236517442656, 0.011928245674245034, 0.008360253800142272, 1.400469240794348e-15, 1.400469240794348e-10],
+        [-28.3, -92.7, 2.4384, 0.13666788662858922, 0.13806942904728953, 0.19086865407406953, 0.021903370028946094, 0.014286237406805243, 0.009969195071770154, 1.671233824831793e-15, 1.671233824831793e-10],
+        [-28.3, -92.7, 2.7432000000000003, 0.1546351494379481, 0.15550977805502827, 0.20444645515930493, 0.025897022938217863, 0.01692827031465076, 0.011896459855279692, 1.981126523617674e-15, 1.9811265236176742e-10],
+        [-28.3, -92.7, 3.048, 0.17593838626887068, 0.1763018192292031, 0.22348587949339338, 0.031077591993027873, 0.020170197324535517, 0.014101618102528489, 2.354375200281265e-15, 2.354375200281265e-10],
+        [-28.3, -92.7, 3.3528000000000002, 0.1995474845150312, 0.19871933694107488, 0.25774908008277886, 0.03604072732456024, 0.023769715197039278, 0.016575324756554193, 2.7691209899848138e-15, 2.769120989984814e-10],
+        [-28.3, -92.7, 3.6576000000000004, 0.22689307794316027, 0.22593853475009562, 0.2894654790464485, 0.04223176077123161, 0.027560159456400522, 0.01937364401629483, 3.2117282348855534e-15, 3.2117282348855533e-10],
+        [-28.3, -92.7, 11.28, 2.3571842072916946, 2.0140180898295945, 2.630644853848645, 0.36285040904171806, 0.23613001211804185, 0.16794048103369616, 2.691748763913162e-14, 2.6917487639131617e-09],
+        [46.2, 187.4, 0.0, 0.06671019978636156, 0.06992548238549574, 0.10354471905996884, 0.007043967133951855, 0.004316608043173917, 0.002976619700967251, 5.363467916315131e-16, 5.363467916315131e-11],
+        [46.2, 187.4, 0.3048, 0.07284088288991165, 0.07567825016081399, 0.10827909438174309, 0.008283351439271345, 0.005360552724598643, 0.0037244297381090643, 6.681170995380951e-16, 6.681170995380952e-11],
+        [46.2, 187.4, 0.6096, 0.08167146527309645, 0.08449937050397927, 0.11618341302969032, 0.010377069611073373, 0.006762637687298858, 0.004745008051649888, 8.44049399093005e-16, 8.44049399093005e-11],
+        [46.2, 187.4, 0.9144000000000001, 0.09231227508105715, 0.09565693138863546, 0.12872709695413875, 0.013248537037775693, 0.008624792993379371, 0.006003462270635623, 1.0720559489489072e-15, 1.0720559489489071e-10],
+        [46.2, 187.4, 1.2192, 0.10494485762208239, 0.10856606232462561, 0.14404071215196426, 0.01648954709104448, 0.010820965260003299, 0.007513607594123712, 1.3444729581971438e-15, 1.3444729581971438e-10],
+        [46.2, 187.4, 1.524, 0.12044588046589635, 0.1243458363119665, 0.15865701209984695, 0.020698038390611302, 0.013512472680105388, 0.00944292084576848, 1.6796173456396709e-15, 1.6796173456396708e-10],
+        [46.2, 187.4, 1.8288000000000002, 0.13856071200530126, 0.14273631106492435, 0.1747352650257913, 0.02559046167451603, 0.016820750907536913, 0.011675958764439961, 2.0841758637764628e-15, 2.084175863776463e-10],
+        [46.2, 187.4, 2.1336, 0.1598753231052355, 0.16454497043137206, 0.19217773250008807, 0.03134496875496027, 0.020461140624593383, 0.014269361754313835, 2.5383793129609276e-15, 2.5383793129609276e-10],
+        [46.2, 187.4, 2.4384, 0.184949398554177, 0.18976144752716587, 0.22501797443352736, 0.038085732128113045, 0.024810226620964042, 0.017235204105509286, 3.07376150529956e-15, 3.0737615052995595e-10],
+        [46.2, 187.4, 2.7432000000000003, 0.21409042068518414, 0.21893183436514435, 0.24077449738687462, 0.04568440962866431, 0.029793311003672995, 0.020849667729028968, 3.694527081841085e-15, 3.694527081841084e-10],
+        [46.2, 187.4, 3.048, 0.24878028566598537, 0.25330667482573466, 0.2762339078511764, 0.055307994505423404, 0.035958467188018095, 0.02508062572260881, 4.448732406000993e-15, 4.4487324060009926e-10],
+        [46.2, 187.4, 3.3528000000000002, 0.28870416297205626, 0.2928445945722284, 0.3424323227214122, 0.0652533727565486, 0.04275837758249945, 0.02973062735855638, 5.2902148242008795e-15, 5.29021482420088e-10],
+        [46.2, 187.4, 3.6576000000000004, 0.33438236973665236, 0.33720379491784713, 0.3639080352296126, 0.07693858325485527, 0.05030576500792614, 0.03518647376068323, 6.223177997820377e-15, 6.223177997820378e-10],
+        [46.2, 187.4, 11.28, 4.596310901937148, 3.8613534445369284, 4.152791457357709, 0.85531816872186, 0.5565098591856836, 0.3941742711337946, 6.809159306290927e-14, 6.809159306290927e-09]
+    ]
+
+    result = AniMAIRE.run_from_DLR_cosmic_ray_model(
+        OULU_count_rate_in_seconds=106.0,
+        Kp_index=3,
+        date_and_time=dt.datetime(2006, 12, 13, 3, 0, tzinfo=dt.timezone.utc),
+        altitudes_in_km=np.append(np.array(range(0, 13)) * 0.3048, 11.28),
+        array_of_lats_and_longs=[[46.2, 187.4], [-28.3, -92.7]]
+    ).values.tolist()
+
+    assert np.allclose(np.array(result), np.array(expected_dlr_spec))
+
 def test_isotropic_dose_rates():
     test_isotropic_dose_rates = AniMAIRE.run_from_spectra(
         proton_rigidity_spectrum=lambda x:2.56*(x**-3.41),

diff --git a/setup.py b/setup.py
@@ -27,7 +27,7 @@
     package_data={"AniMAIRE":[
                                 "anisotropic_MAIRE_engine/data/*.csv"
                                          ]},
-    version='1.2.3',
+    version='1.2.4',
     description='Python library for running the anisotropic version of MAIRE+',
     long_description=long_description,
     long_description_content_type='text/markdown',