Harrison 3029 - PMCL/SWYN - CoDICE Lo direct events: Updated math for computing normalization factor and refactor the computation to its own function

Author: Menlo Innovations - CAVA Project

imap_l3_processing/codice/l3/lo/codice_lo_processor.py

@@ -18,7 +18,7 @@
     calculate_mass, calculate_mass_per_charge, \
     rebin_to_counts_by_species_elevation_and_spin_sector, normalize_counts, \
     combine_priorities_and_convert_to_rate, rebin_3d_distribution_azimuth_to_elevation, convert_count_rate_to_intensity, \
-    rebin_direct_events_for_normalization
+    calculate_normalization_factor
 from imap_l3_processing.data_utils import safe_divide
 from imap_l3_processing.models import InputMetadata
 from imap_l3_processing.processor import Processor
@@ -186,30 +186,18 @@ def process_l3a_direct_event_data_product(self, dependencies: CodiceLoL3aDirectE
         ]
         num_priorities = len(priority_counts)
 
-        mass_per_charge = np.full((len(codice_direct_events.epoch), num_priorities, event_buffer),
-                                  np.nan)
-        mass = np.full((len(codice_direct_events.epoch), num_priorities, event_buffer), np.nan)
-
         spin_angle_lut = SpinAngleLookup()
-        normalization = np.full((len(codice_direct_events.epoch), num_priorities,
-                                 num_energies, num_spin_sectors), np.nan)
-
-        try:
-            mass_per_charge = calculate_mass_per_charge(codice_direct_events.energy_per_charge,
-                                                        codice_direct_events.tof)
-            mass = calculate_mass(codice_direct_events.apd_energy, codice_direct_events.tof, mass_coefficient_lookup)
-
-            direct_events_binned_by_energy_and_spin = rebin_direct_events_for_normalization(
-                codice_direct_events.num_events,
-                codice_direct_events.spin_sector,
-                codice_direct_events.energy_step,
-                num_spin_sectors,
-                num_energies)
-            stacked_priorities = np.stack(priority_counts, axis=1)
-            numerator = np.concatenate((stacked_priorities, stacked_priorities), axis=3)
-            normalization = numerator / direct_events_binned_by_energy_and_spin
-        except Exception as e:
-            print(e)
+
+        mass_per_charge = calculate_mass_per_charge(codice_direct_events.energy_per_charge,
+                                                    codice_direct_events.tof)
+        mass = calculate_mass(codice_direct_events.apd_energy, codice_direct_events.tof, mass_coefficient_lookup)
+        stacked_priorities = np.stack(priority_counts, axis=1)
+        normalization = calculate_normalization_factor(
+            stacked_priorities,
+            codice_direct_events.num_events,
+            codice_direct_events.spin_sector,
+            codice_direct_events.energy_step,
+        )
 
         return CodiceLoL3aDirectEventDataProduct(
             input_metadata=self.input_metadata,

imap_l3_processing/codice/l3/lo/science/codice_lo_calculations.py

@@ -53,10 +53,10 @@ def calculate_mass_per_charge(energy_per_charge: np.ndarray, tof: np.ndarray) ->
             tof ** 2) * CONVERSION_CONSTANT_K
 
 
-def rebin_direct_events_for_normalization(num_events: np.ndarray, spin_angle: np.ndarray, energy_step: np.ndarray,
+def rebin_direct_events_for_normalization(num_events: np.ndarray, spin_sector: np.ndarray, energy_step: np.ndarray,
                                           num_spin_sectors: int,
                                           num_energies: int) -> np.ndarray:
-    base_counts = rebin_direct_events_by_energy_and_spin_sector(num_events, spin_angle, energy_step, num_spin_sectors, num_energies)
+    base_counts = rebin_direct_events_by_energy_and_spin_sector(num_events, spin_sector, energy_step, num_spin_sectors, num_energies)
     half_spin = num_spin_sectors//2
     result = np.zeros_like(base_counts)
     result[:,:,:,0:half_spin] = base_counts[:,:,:,0:half_spin] + base_counts[:,:,:,half_spin:num_spin_sectors]
@@ -83,6 +83,20 @@ def rebin_direct_events_by_energy_and_spin_sector(num_events: np.ndarray, spin_s
             np.add.at(rebinned_output[time_index, priority_index], (energy_indices, spin_sectors), 1)
     return rebinned_output
 
+def calculate_normalization_factor(priority_counts: np.ndarray, num_events: np.ndarray, spin_sectors: np.ndarray, energy_steps: np.ndarray) -> np.ndarray:
+    numerator = np.concatenate((priority_counts, priority_counts), axis=3)
+    num_energies = priority_counts.shape[2]
+    num_spin_sectors = 2 * priority_counts.shape[3]
+    denominator = rebin_direct_events_for_normalization(num_events, spin_sectors, energy_steps, num_spin_sectors, num_energies)
+
+    division_result = np.zeros(numerator.shape, dtype=float)
+    np.divide(numerator, denominator, out=division_result, where=denominator!=0)
+
+    output = np.zeros_like(division_result)
+    output[denominator!=0] = np.maximum(1.0, division_result[denominator!=0])
+    output[(denominator==0)&(numerator==0)] = 0.0
+    output[(denominator==0)&(numerator!=0)] = np.nan
+    return output
 
 def rebin_to_counts_by_species_elevation_and_spin_sector(num_events: np.ndarray, mass: np.ndarray,
                                                          mass_per_charge: np.ndarray,

tests/codice/l3/lo/science/test_codice_lo_calculations.py

@@ -14,7 +14,8 @@
     calculate_total_number_of_events, calculate_normalization_ratio, calculate_mass, calculate_mass_per_charge, \
     rebin_to_counts_by_species_elevation_and_spin_sector, rebin_direct_events_by_energy_and_spin_sector, \
     CODICE_LO_NUM_AZIMUTH_BINS, normalize_counts, combine_priorities_and_convert_to_rate, \
-    rebin_3d_distribution_azimuth_to_elevation, convert_count_rate_to_intensity, rebin_direct_events_for_normalization
+    rebin_3d_distribution_azimuth_to_elevation, convert_count_rate_to_intensity, rebin_direct_events_for_normalization, \
+    calculate_normalization_factor
 from imap_l3_processing.constants import ONE_SECOND_IN_MICROSECONDS
 
 
@@ -316,6 +317,96 @@ def test_rebin_direct_events_for_normalization(self):
 
         np.testing.assert_array_equal(result, expected_rebinned_counts)
 
+    def test_calculate_normalization_factor(self):
+        num_epochs = 2
+        num_priorities = 7
+        num_esa_steps = 128
+        num_l1_spin_sectors = 12
+        num_l2_spin_sectors = 2*num_l1_spin_sectors
+        priority_counts = np.zeros((num_epochs, num_priorities, num_esa_steps, num_l1_spin_sectors))
+        priority_counts[0,0,0,0] = 4
+        priority_counts[0,0,0,2] = 20
+
+        event_buffer_len = 15
+        num_events = np.zeros((num_epochs, num_priorities), dtype=int)
+        spin_sectors = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=int)
+        energy_steps = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=int)
+
+        num_events[0, 0] = 3
+        spin_sectors[0, 0, :3] = [0, 2, 12]
+        energy_steps[0, 0, :3] = [0, 0, 0]
+
+        result = calculate_normalization_factor(priority_counts, num_events, spin_sectors, energy_steps)
+
+        expected_shape = (num_epochs, num_priorities, num_esa_steps, num_l2_spin_sectors)
+        expected = np.full(expected_shape, 0.0)
+        expected[0,0,0,0] = 4 / 2
+        expected[0,0,0,0+12] = 4 / 2
+        expected[0,0,0,2] = 20 / 1
+        expected[0,0,0,2+12] = 20 / 1
+
+        np.testing.assert_equal(result, expected)
+
+    def test_calculate_normalization_factor_floor_is_one(self):
+        num_epochs = 2
+        num_priorities = 7
+        num_esa_steps = 128
+        num_l1_spin_sectors = 12
+        num_l2_spin_sectors = 2*num_l1_spin_sectors
+        priority_counts = np.zeros((num_epochs, num_priorities, num_esa_steps, num_l1_spin_sectors))
+        priority_counts[0,0,0,0] = 1.
+        priority_counts[0,0,0,2] = 0.
+
+
+        event_buffer_len = 15
+        num_events = np.zeros((num_epochs, num_priorities), dtype=int)
+        spin_sectors = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=int)
+        energy_steps = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=int)
+
+        num_events[0, 0] = 3
+        spin_sectors[0, 0, :3] = [0, 12, 2]
+        energy_steps[0, 0, :3] = [0, 0, 0]
+
+        result = calculate_normalization_factor(priority_counts, num_events, spin_sectors, energy_steps)
+
+        expected_shape = (num_epochs, num_priorities, num_esa_steps, num_l2_spin_sectors)
+        expected = np.full(expected_shape, 0.0)
+        self.assertEqual(1.0, result[0,0,0,0])
+        self.assertEqual(1.0, result[0,0,0,0+12])
+        self.assertEqual(1.0, result[0,0,0,2])
+        self.assertEqual(1.0, result[0,0,0,2+12])
+
+        expected[0, 0, 0, (0, 2, 12, 14)] = 1.0
+        np.testing.assert_equal(result, expected)
+
+
+    def test_normalization_factor_is_fill_when_priority_nonzero_and_direct_zero(self):
+        num_epochs = 2
+        num_priorities = 7
+        num_esa_steps = 128
+        num_l1_spin_sectors = 12
+        num_l2_spin_sectors = 2*num_l1_spin_sectors
+        priority_counts = np.zeros((num_epochs, num_priorities, num_esa_steps, num_l1_spin_sectors))
+        priority_counts[0,0,0,0] = 1.
+        priority_counts[0,0,0,2] = 50.
+
+        event_buffer_len = 15
+        num_events = np.zeros((num_epochs, num_priorities), dtype=int)
+        spin_sectors = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=int)
+        energy_steps = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=int)
+
+        result = calculate_normalization_factor(priority_counts, num_events, spin_sectors, energy_steps)
+
+        expected_shape = (num_epochs, num_priorities, num_esa_steps, num_l2_spin_sectors)
+        expected = np.full(expected_shape, 0.0)
+        np.testing.assert_equal(result[0,0,0,0], np.nan)
+        np.testing.assert_equal(result[0,0,0,0+12], np.nan)
+        np.testing.assert_equal(result[0,0,0,2], np.nan)
+        np.testing.assert_equal(result[0,0,0,2+12], np.nan)
+
+        expected[0, 0, 0, (0, 2, 12, 14)] = np.nan
+        np.testing.assert_equal(result, expected)
+
     def test_normalize_counts(self):
         num_epochs = 2
         num_priorities = 3

tests/codice/l3/lo/test_codice_lo_processor.py

@@ -499,11 +499,11 @@ def test_process_l3a_partial_densities(self, mock_calculate_partial_densities):
         self.assertEqual(fe_hiq_partial_density, result_data.fe_hiq_partial_density),
 
     @patch('imap_l3_processing.codice.l3.lo.codice_lo_processor.SpinAngleLookup')
-    @patch('imap_l3_processing.codice.l3.lo.codice_lo_processor.rebin_direct_events_for_normalization')
+    @patch('imap_l3_processing.codice.l3.lo.codice_lo_processor.calculate_normalization_factor')
     @patch('imap_l3_processing.codice.l3.lo.codice_lo_processor.calculate_mass_per_charge')
     @patch('imap_l3_processing.codice.l3.lo.codice_lo_processor.calculate_mass')
     def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mass_per_charge,
-                                       mock_rebin_direct_events_for_normalization,
+                                       mock_calculate_normalization_factor,
                                        mock_spin_angle_lookup_class):
         rng = np.random.default_rng()
 
@@ -517,42 +517,39 @@ def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mas
         mock_spin_angle_lookup_class.return_value = mock_spin_angle_lookup
 
         epochs = np.array([datetime.now(), datetime.now() + timedelta(hours=1)])
+        num_epochs = len(epochs)
 
-        priority_counts_variable_shape = (len(epochs), num_energy_bins, num_spin_angle_bins//2)
+        priority_counts_variable_shape = (num_epochs, num_energy_bins, num_spin_angle_bins // 2)
         sw_priority_rates = create_dataclass_mock(CodiceLoL1aSWPriorityRates)
         sw_priority_rates.epoch = epochs
         sw_priority_rates.p0_tcrs = rng.random(priority_counts_variable_shape)
         sw_priority_rates.p1_hplus = rng.random(priority_counts_variable_shape)
         sw_priority_rates.p2_heplusplus = rng.random(priority_counts_variable_shape)
         sw_priority_rates.p3_heavies = rng.random(priority_counts_variable_shape)
         sw_priority_rates.p4_dcrs = rng.random(priority_counts_variable_shape)
-        sw_priority_rates.half_spin_per_esa_step = rng.random((len(epochs), num_energy_bins))
-        sw_priority_rates.rgfo_spin_sector = rng.random(len(epochs))
-        sw_priority_rates.rgfo_esa_step = rng.random(len(epochs))
-        sw_priority_rates.nso_spin_sector = rng.random(len(epochs))
-        sw_priority_rates.nso_esa_step = rng.random(len(epochs))
+        sw_priority_rates.half_spin_per_esa_step = rng.random((num_epochs, num_energy_bins))
+        sw_priority_rates.rgfo_spin_sector = rng.random(num_epochs)
+        sw_priority_rates.rgfo_esa_step = rng.random(num_epochs)
+        sw_priority_rates.nso_spin_sector = rng.random(num_epochs)
+        sw_priority_rates.nso_esa_step = rng.random(num_epochs)
 
         nsw_priority_rates = create_dataclass_mock(CodiceLoL1aNSWPriorityRates)
         nsw_priority_rates.epoch = epochs
         nsw_priority_rates.p5_heavies = rng.random(priority_counts_variable_shape)
         nsw_priority_rates.p6_hplus_heplusplus = rng.random(priority_counts_variable_shape)
 
-        counts_rebinned_by_energy_and_spin = rng.random(
-            (len(epochs), CODICE_LO_L2_NUM_PRIORITIES, num_energy_bins, num_spin_angle_bins))
-        mock_rebin_direct_events_for_normalization.return_value = counts_rebinned_by_energy_and_spin
-
-        mass_per_charge = rng.random((len(epochs), CODICE_LO_L2_NUM_PRIORITIES, event_buffer_size))
+        mass_per_charge = rng.random((num_epochs, CODICE_LO_L2_NUM_PRIORITIES, event_buffer_size))
         mock_calculate_mass_per_charge.return_value = mass_per_charge
-        mass = rng.random((len(epochs), CODICE_LO_L2_NUM_PRIORITIES, event_buffer_size))
+        mass = rng.random((num_epochs, CODICE_LO_L2_NUM_PRIORITIES, event_buffer_size))
         mock_calculate_mass.return_value = mass
 
-        codice_l2_variables = {f.name: rng.random((len(epochs), CODICE_LO_L2_NUM_PRIORITIES, event_buffer_size)) for f
+        codice_l2_variables = {f.name: rng.random((num_epochs, CODICE_LO_L2_NUM_PRIORITIES, event_buffer_size)) for f
                                in fields(CodiceLoL2DirectEventData)}
 
         codice_l2_variables["epoch"] = epochs
         codice_l2_variables["spin_angle"] *= 360
-        codice_l2_variables["data_quality"] = rng.random((len(epochs)))
-        codice_l2_variables["num_events"] = rng.random((len(epochs)))
+        codice_l2_variables["data_quality"] = rng.random((num_epochs))
+        codice_l2_variables["num_events"] = rng.random((num_epochs))
 
         expected_energy_per_charge = codice_l2_variables["energy_per_charge"]
         expected_apd_energy = codice_l2_variables["apd_energy"]
@@ -578,6 +575,9 @@ def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mas
         mock_energy_lookup.bin_centers = np.geomspace(14100, 88.082825, 128)
         mock_energy_lookup.num_bins = num_energy_bins
 
+        expected_normalization = rng.random((num_epochs, num_priorities, num_energy_bins, num_spin_angle_bins))
+        mock_calculate_normalization_factor.return_value = expected_normalization
+
         dependencies = CodiceLoL3aDirectEventsDependencies(sw_priority_rates, nsw_priority_rates, direct_events, Mock(),
                                                            mock_energy_lookup)
 
@@ -595,12 +595,21 @@ def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mas
         np.testing.assert_equal(mock_calculate_mass_per_charge.call_args.args[0], expected_energy_per_charge)
         np.testing.assert_equal(mock_calculate_mass_per_charge.call_args.args[1], expected_tof)
 
-        self.assertEqual(1, mock_rebin_direct_events_for_normalization.call_count)
-        np.testing.assert_equal(mock_rebin_direct_events_for_normalization.call_args.args[0], expected_num_events)
-        np.testing.assert_equal(mock_rebin_direct_events_for_normalization.call_args.args[1], expected_spin_sector)
-        np.testing.assert_equal(mock_rebin_direct_events_for_normalization.call_args.args[2], expected_energy_step)
-        np.testing.assert_equal(mock_rebin_direct_events_for_normalization.call_args.args[3], num_spin_angle_bins)
-        np.testing.assert_equal(mock_rebin_direct_events_for_normalization.call_args.args[4], num_energy_bins)
+        self.assertEqual(1, mock_calculate_normalization_factor.call_count)
+        actual_stacked_priorities = mock_calculate_normalization_factor.call_args.args[0]
+
+        np.testing.assert_equal(actual_stacked_priorities[:, 0], sw_priority_rates.p0_tcrs)
+        np.testing.assert_equal(actual_stacked_priorities[:, 1], sw_priority_rates.p1_hplus)
+        np.testing.assert_equal(actual_stacked_priorities[:, 2], sw_priority_rates.p2_heplusplus)
+        np.testing.assert_equal(actual_stacked_priorities[:, 3], sw_priority_rates.p3_heavies)
+        np.testing.assert_equal(actual_stacked_priorities[:, 4], sw_priority_rates.p4_dcrs)
+        np.testing.assert_equal(actual_stacked_priorities[:, 5], nsw_priority_rates.p5_heavies)
+        np.testing.assert_equal(actual_stacked_priorities[:, 6], nsw_priority_rates.p6_hplus_heplusplus)
+
+        np.testing.assert_equal(mock_calculate_normalization_factor.call_args.args[1], expected_num_events)
+        np.testing.assert_equal(mock_calculate_normalization_factor.call_args.args[2], expected_spin_sector)
+        np.testing.assert_equal(mock_calculate_normalization_factor.call_args.args[3], expected_energy_step)
+
 
         self.assertIsInstance(l3a_direct_event_data_product, CodiceLoL3aDirectEventDataProduct)
         self.assertEqual(input_metadata, l3a_direct_event_data_product.input_metadata)
@@ -613,20 +622,6 @@ def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mas
         np.testing.assert_array_equal(expected_mass_per_charge, l3a_direct_event_data_product.mass_per_charge)
         np.testing.assert_array_equal(expected_mass, l3a_direct_event_data_product.mass)
 
-        priority_rates = np.stack([
-            sw_priority_rates.p0_tcrs,
-            sw_priority_rates.p1_hplus,
-            sw_priority_rates.p2_heplusplus,
-            sw_priority_rates.p3_heavies,
-            sw_priority_rates.p4_dcrs,
-            nsw_priority_rates.p5_heavies,
-            nsw_priority_rates.p6_hplus_heplusplus,
-        ], axis=1)
-        expected_numerator = np.zeros((len(epochs), num_priorities, num_energy_bins, num_spin_angle_bins))
-        expected_numerator[:,:,:,:12] = priority_rates
-        expected_numerator[:,:,:,12:] = priority_rates
-        expected_normalization = expected_numerator / counts_rebinned_by_energy_and_spin
-
         np.testing.assert_array_equal(l3a_direct_event_data_product.normalization,
                                       np.flip(expected_normalization, axis=2))