adjust combine_channels for EPT L3

solo_epd_loader/__init__.py

@@ -1837,35 +1837,48 @@ def combine_channels(df, energies, en_channel, sensor, viewing=None, species=Non
     > df_p, df_e, meta = epd_load('ept', 20200820, 20200821, 'l2', 'sun')
     > df_new, chan_new = combine_channels(df_p, meta, [9, 12], 'ept')
     """
-    # check for EPT level 3 data product:
-    if 'Pitch_Angle_A' in df.keys():
-        raise Exception('EPT level 3 data not supported yet!')  # TODO: Add support for EPT level 3 data
-        if not species or not viewing:
-            raise Exception('EPT level 3 data require defined "species" and "viewing"!')
-        # if sensor.lower() == 'ept':
-        #     if species.lower() in ['e', 'electrons']:
-        #         pass
-        #     if species.lower() in ['p', 'protons', 'i', 'ions', 'h']:
-        #         pass
-    else:
-        if sensor.lower() == 'step':
-            raise Exception('STEP data not supported yet!')
-            return pd.DataFrame(), ''
-        # if species.lower() in ['e', 'electrons']:
+    if sensor.lower() == 'step':
+        raise Exception('STEP data not supported yet!')
+        return pd.DataFrame(), ''
+
+    # for l2 EPT and HET data derive species from dataframe (they are separated)
+    if not species and 'Pitch_Angle_A' not in df.keys():
         if 'Electron_Flux' in df.keys():
-            en_str = energies['Electron_Bins_Text']
-            bins_width = 'Electron_Bins_Width'
-            flux_key = 'Electron_Flux'
-        # if species.lower() in ['p', 'protons', 'i', 'ions', 'h']:
+            species = 'e'
         else:
-            if sensor.lower() == 'het':
-                en_str = energies['H_Bins_Text']
-                bins_width = 'H_Bins_Width'
-                flux_key = 'H_Flux'
-            if sensor.lower() == 'ept':
-                en_str = energies['Ion_Bins_Text']
-                bins_width = 'Ion_Bins_Width'
-                flux_key = 'Ion_Flux'
+            species = 'p'
+
+    # for EPT l3 data
+    if 'Pitch_Angle_A' in df.keys():
+        if not species:
+            raise Exception("EPT level 3 data requires 'species' option!")
+        if not viewing:
+            raise Exception("EPT level 3 data requires 'viewing' option!")
+        else: 
+            viewing = viewing.lower().replace('-', '')
+            viewing = viewing.replace('asun', 'antisun')
+            viewing_short = {'sun': 'S',
+                             'antisun': 'A',
+                             'north': 'N',
+                             'south': 'D'}
+
+    if species.lower() in ['e', 'electrons']:
+        en_str = energies['Electron_Bins_Text']
+        bins_width = 'Electron_Bins_Width'
+        flux_key = 'Electron_Flux'
+        if 'Pitch_Angle_A' in df.keys():
+            flux_key = f'Electron_Corrected_Flux_{viewing_short[viewing]}'
+    elif species.lower() in ['p', 'protons', 'i', 'ions', 'h']:
+        if sensor.lower() == 'het':
+            en_str = energies['H_Bins_Text']
+            bins_width = 'H_Bins_Width'
+            flux_key = 'H_Flux'
+        if sensor.lower() == 'ept':
+            en_str = energies['Ion_Bins_Text']
+            bins_width = 'Ion_Bins_Width'
+            flux_key = 'Ion_Flux'
+            if 'Pitch_Angle_A' in df.keys():
+                flux_key = f'Ion_Flux_{viewing_short[viewing]}'
     if type(en_channel) == list:
         energy_low = en_str[en_channel[0]].flat[0].split('-')[0]
         energy_up = en_str[en_channel[-1]].flat[0].split('-')[-1]
@@ -1874,13 +1887,20 @@ def combine_channels(df, energies, en_channel, sensor, viewing=None, species=Non
         if len(en_channel) > 2:
             raise Exception('en_channel must have 2 elements: start channel and end channel, e.g. [1,3]!')
         if len(en_channel) == 2:
-            # catch overlapping EPT energy channels and cancel calculation:
-            if sensor.lower() == 'ept' and 'Electron_Flux' in df.keys() and en_channel[0] < 4:
+            # catch overlapping EPT L2 energy channels and cancel calculation:
+            if sensor.lower() == 'ept' and species == 'e' and en_channel[0] < 4 and 'Pitch_Angle_A' not in df.keys():
                 raise Exception('Lowest 4 EPT e channels not supported because of overlapping energies!')
                 return pd.DataFrame(), ''
-            if sensor.lower() == 'ept' and 'Electron_Flux' not in df.keys() and en_channel[0] < 9:
+            if sensor.lower() == 'ept' and species == 'p' and en_channel[0] < 9 and 'Pitch_Angle_A' not in df.keys():
                 raise Exception('Lowest 9 EPT ion channels not supported because of overlapping energies!')
                 return pd.DataFrame(), ''
+            # catch overlapping EPT L3 energy channels and cancel calculation:
+            if sensor.lower() == 'ept' and species == 'e' and en_channel[0] < 1 and 'Pitch_Angle_A' in df.keys():
+                raise Exception('Lowest EPT e channel not supported because of overlapping energies!')
+                return pd.DataFrame(), ''
+            if sensor.lower() == 'ept' and species == 'p' and en_channel[0] < 3 and 'Pitch_Angle_A' in df.keys():
+                raise Exception('Lowest 3 EPT ion channels not supported because of overlapping energies!')
+                return pd.DataFrame(), ''
             # try to convert multi-index dataframe to normal one. if this is already the case, just continue
             try:
                 df = df[flux_key]
@@ -1896,11 +1916,17 @@ def combine_channels(df, energies, en_channel, sensor, viewing=None, species=Non
             flux_out = pd.DataFrame({'flux': I_all/DE_total}, index=df.index)
         else:
             en_channel = en_channel[0]
-            flux_out = pd.DataFrame({'flux': df[flux_key][f'{flux_key}_{en_channel}']}, index=df.index)
-            en_channel_string = en_str[en_channel][0]
+            if 'Pitch_Angle_A' not in df.keys():
+                flux_out = pd.DataFrame({'flux': df[flux_key][f'{flux_key}_{en_channel}']}, index=df.index)
+            else:
+                flux_out = pd.DataFrame({'flux': df[f'{flux_key}_{en_channel}']}, index=df.index)
+            en_channel_string = en_str[en_channel]
     else:
-        flux_out = pd.DataFrame({'flux': df[flux_key][f'{flux_key}_{en_channel}']}, index=df.index)
-        en_channel_string = en_str[en_channel][0]
+        if 'Pitch_Angle_A' not in df.keys():
+            flux_out = pd.DataFrame({'flux': df[flux_key][f'{flux_key}_{en_channel}']}, index=df.index)
+        else:
+            flux_out = pd.DataFrame({'flux': df[f'{flux_key}_{en_channel}']}, index=df.index)
+        en_channel_string = en_str[en_channel]
     return flux_out, en_channel_string
 
 

solo_epd_loader/tests/tests.py

@@ -26,15 +26,22 @@ def test_ept_l3_load_online():
     assert np.sum(np.isnan(df['Ion_Flux_A_0'])) == np.int64(786)
     #
     # test combine_channels for ions
-    # df_p_new, chan_p_new = combine_channels(df=df_p, energies=meta, en_channel=[9, 12], sensor='ept')
-    # assert chan_p_new == '0.0809 - 0.1034 MeV'
-    # assert df_p_new.shape == (38809, 1)
-    # assert df_p_new['flux'].sum() == np.float32(28518708.0)
-    # # test combine_channels for electrons
-    # df_e_new, chan_e_new = combine_channels(df=df_e, energies=meta, en_channel=[4, 7], sensor='ept')
-    # assert chan_e_new == '0.0334 - 0.0420 MeV'
-    # assert df_e_new.shape == (38809, 1)
-    # assert df_e_new['flux'].sum() == np.float32(49434200.0)
+    df_p_new, chan_p_new = combine_channels(df=df, energies=energies_dict, en_channel=[9, 12], sensor='ept', viewing='asun', species='p')
+    assert chan_p_new == '0.1628 - 0.2980 MeV'
+    assert df_p_new.shape == (2880, 1)
+    assert df_p_new['flux'].sum() == np.float32(313939.88)
+    # test combine_channels for electrons
+    df_e_new, chan_e_new = combine_channels(df=df, energies=energies_dict, en_channel=[9, 12], sensor='ept', viewing='South', species='e')
+    assert chan_e_new == '0.1432 - 0.2826 MeV'
+    assert df_e_new.shape == (2880, 1)
+    assert df_e_new['flux'].sum() == np.float32(12301.297)
+    df_e_new, chan_e_new = combine_channels(df=df, energies=energies_dict, en_channel=[12], sensor='ept', viewing='South', species='e')
+    assert chan_e_new == '0.2379 - 0.2826 MeV'
+    assert df_e_new.shape == (2880, 1)
+    assert df_e_new['flux'].sum() == np.float32(7876.287)
+    df_e_new2, chan_e_new2 = combine_channels(df=df, energies=energies_dict, en_channel=12, sensor='ept', viewing='South', species='e')
+    assert df_e_new2.equals(df_e_new)
+    assert chan_e_new2 == chan_e_new
     #
     # test resampling
     df_res = resample_df(df=df, resample='1h')
@@ -61,6 +68,14 @@ def test_ept_l2_load_online():
     assert chan_p_new == '0.0809 - 0.1034 MeV'
     assert df_p_new.shape == (38809, 1)
     assert df_p_new['flux'].sum() == np.float32(28518708.0)
+    #
+    df_p_new, chan_p_new = combine_channels(df=df_p, energies=meta, en_channel=[1], sensor='ept')
+    assert chan_p_new == '0.0520 - 0.0602 MeV'
+    assert df_p_new.shape == (38809, 1)
+    assert df_p_new['flux'].sum() == np.float32(85663600.0)
+    df_p_new2, chan_p_new2 = combine_channels(df=df_p, energies=meta, en_channel=1, sensor='ept')
+    assert df_p_new2.equals(df_p_new)
+    assert chan_p_new2 == chan_p_new
     # test combine_channels for electrons
     df_e_new, chan_e_new = combine_channels(df=df_e, energies=meta, en_channel=[4, 7], sensor='ept')
     assert chan_e_new == '0.0408 - 0.0542 MeV'
@@ -147,6 +162,15 @@ def test_het_l2_load_online():
     assert chan_p_new == '11.8000 - 15.6500 MeV'
     assert df_p_new.shape == (7767, 1)
     assert df_p_new['flux'].sum() == np.float32(3.106687)
+    #
+    df_p_new, chan_p_new = combine_channels(df=df_p, energies=meta, en_channel=[1], sensor='het')
+    assert chan_p_new == '7.3540 - 7.8900 MeV'
+    assert df_p_new.shape == (7767, 1)
+    assert df_p_new['flux'].sum() == np.float32(1.4690328)
+    df_p_new2, chan_p_new2 = combine_channels(df=df_p, energies=meta, en_channel=1, sensor='het')
+    assert df_p_new2.equals(df_p_new)
+    assert chan_p_new2 == chan_p_new
+    #
     df_e_new, chan_e_new = combine_channels(df=df_e, energies=meta, en_channel=[1, 3], sensor='het')
     assert chan_e_new == '1.0530 - 18.8300 MeV'
     assert df_e_new.shape == (38839, 1)