Harrison 3029 - PMCL - CoDICE Lo direct events: Finish new normalization factor formula

Author: Menlo Innovations - CAVA Project

imap_l3_processing/codice/l3/lo/codice_lo_processor.py

@@ -16,8 +16,9 @@
     CodiceLoL3ChargeStateDistributionsDataProduct, CodiceLoL3a3dDistributionDataProduct
 from imap_l3_processing.codice.l3.lo.science.codice_lo_calculations import calculate_partial_densities, \
     calculate_mass, calculate_mass_per_charge, \
-    rebin_counts_by_energy_and_spin_angle, 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_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
 from imap_l3_processing.data_utils import safe_divide
 from imap_l3_processing.models import InputMetadata
 from imap_l3_processing.processor import Processor
@@ -171,8 +172,10 @@ def process_l3a_direct_event_data_product(self, dependencies: CodiceLoL3aDirectE
         esa_energy_per_charge_lookup = dependencies.energy_lookup
 
         event_buffer = codice_direct_events.tof[0].shape[-1]
+        num_energies = codice_sw_priority_counts_l1a_data.p0_tcrs.shape[1]
+        num_spin_sectors = 2*codice_sw_priority_counts_l1a_data.p0_tcrs.shape[2]
         mass_coefficient_lookup = dependencies.mass_coefficient_lookup
-        priority_counts_for_events = [
+        priority_counts = [
             codice_sw_priority_counts_l1a_data.p0_tcrs,
             codice_sw_priority_counts_l1a_data.p1_hplus,
             codice_sw_priority_counts_l1a_data.p2_heplusplus,
@@ -181,28 +184,30 @@ def process_l3a_direct_event_data_product(self, dependencies: CodiceLoL3aDirectE
             codice_nsw_priority_counts_l1a_data.p5_heavies,
             codice_nsw_priority_counts_l1a_data.p6_hplus_heplusplus
         ]
+        num_priorities = len(priority_counts)
 
-        mass_per_charge = np.full((len(codice_direct_events.epoch), len(priority_counts_for_events), event_buffer),
+        mass_per_charge = np.full((len(codice_direct_events.epoch), num_priorities, event_buffer),
                                   np.nan)
-        mass = np.full((len(codice_direct_events.epoch), len(priority_counts_for_events), 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), len(priority_counts_for_events),
-                                 esa_energy_per_charge_lookup.num_bins, spin_angle_lut.num_bins), np.nan)
+        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_counts_by_energy_and_spin_angle(
+            direct_events_binned_by_energy_and_spin = rebin_direct_events_for_normalization(
                 codice_direct_events.num_events,
-                codice_direct_events.spin_angle,
+                codice_direct_events.spin_sector,
                 codice_direct_events.energy_step,
-                spin_angle_lut,
-                esa_energy_per_charge_lookup)
-            total_events_per_priority = np.sum(np.stack(priority_counts_for_events, axis=1), axis=(2, 3), keepdims=True)
-            normalization = total_events_per_priority / direct_events_binned_by_energy_and_spin
+                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)
 

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

@@ -53,31 +53,34 @@ def calculate_mass_per_charge(energy_per_charge: np.ndarray, tof: np.ndarray) ->
             tof ** 2) * CONVERSION_CONSTANT_K
 
 
-def rebin_counts_by_energy_and_spin_angle(num_events: np.ndarray, spin_angle: np.ndarray, energy_step: np.ndarray,
-                                          spin_angle_lookup: SpinAngleLookup,
-                                          energy_lookup: EnergyLookup) -> np.ndarray:
+def rebin_direct_events_for_normalization(num_events: np.ndarray, spin_angle: 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)
+    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]
+    result[:,:,:,half_spin:num_spin_sectors] = result[:,:,:,0:half_spin]
+    return result
+
+def rebin_direct_events_by_energy_and_spin_sector(num_events: np.ndarray, spin_sector: np.ndarray, energy_step: np.ndarray,
+                                                  num_spin_sectors: int,
+                                                  num_energies: int) -> np.ndarray:
     num_epochs = num_events.shape[0]
     num_priorities = num_events.shape[1]
-    num_energies = energy_lookup.num_bins
-    num_spin_bins = spin_angle_lookup.num_bins
 
-    rebinned_output = np.zeros((num_epochs, num_priorities, num_energies, num_spin_bins))
+    rebinned_output = np.zeros((num_epochs, num_priorities, num_energies, num_spin_sectors))
 
     for time_index in range(num_epochs):
         for priority_index in range(num_priorities):
             events_at_index = num_events[time_index, priority_index]
             if events_at_index is np.ma.masked:
                 continue
 
-            spin_angle_in_degrees = spin_angle[time_index, priority_index, :events_at_index]
-            energy_in_keV = energy_step[time_index, priority_index, :events_at_index]
+            spin_sectors = spin_sector[time_index, priority_index, :events_at_index]
+            energy_indices = energy_step[time_index, priority_index, :events_at_index]
 
-            spin_angle_indices = spin_angle_lookup.get_spin_angle_index(spin_angle_in_degrees)
-            energy_indices = energy_lookup.get_energy_index(energy_in_keV)
-            spin_angle_indices_first_half = spin_angle_indices % (num_spin_bins//2)
-
-            np.add.at(rebinned_output[time_index, priority_index], (energy_indices, spin_angle_indices_first_half), 1)
-            np.add.at(rebinned_output[time_index, priority_index], (energy_indices, spin_angle_indices_first_half+(num_spin_bins//2)), 1)
+            np.add.at(rebinned_output[time_index, priority_index], (energy_indices, spin_sectors), 1)
     return rebinned_output
 
 

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

@@ -12,9 +12,9 @@
 from imap_l3_processing.codice.l3.lo.models import EnergyAndSpinAngle, CodiceLo3dData
 from imap_l3_processing.codice.l3.lo.science.codice_lo_calculations import calculate_partial_densities, \
     calculate_total_number_of_events, calculate_normalization_ratio, calculate_mass, calculate_mass_per_charge, \
-    rebin_to_counts_by_species_elevation_and_spin_sector, rebin_counts_by_energy_and_spin_angle, \
+    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_3d_distribution_azimuth_to_elevation, convert_count_rate_to_intensity, rebin_direct_events_for_normalization
 from imap_l3_processing.constants import ONE_SECOND_IN_MICROSECONDS
 
 
@@ -262,33 +262,48 @@ def mock_get_species_index(species):
         expected_o_counts[0, 1, 0, 0, 0] = 1
         np.testing.assert_array_equal(actual_counts_3d_data.get_3d_distribution("O+5"), expected_o_counts)
 
-    def test_rebin_counts_to_by_energy_and_spin_angle(self):
-        mock_energy_lookup = Mock(spec=EnergyLookup)
+    def test_rebin_direct_events_by_energy_and_spin_sector(self):
         num_energy_bins = 30
-        mock_energy_lookup.num_bins = num_energy_bins
-        mock_energy_lookup.get_energy_index.side_effect = [
-            np.array([7, 7]),
-            np.array([1, 3, 1, 4])
-        ]
+        num_spin_angle_bins = 20
+        num_priorities = 3
+        num_epochs = 1
+        event_buffer_len = 15
+        num_events = np.ma.masked_array(np.array([[1, 2, 4]]), mask=[[False, True, False]])
+        spin_sector = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=np.uint8)
+        spin_sector[0,0,:1] = [3]
+        spin_sector[0,2,:4] = [5, 6, 5, 7]
 
-        mock_spin_angle_lookup = Mock(spec=SpinAngleLookup)
-        num_spin_angle_bins = 24
-        mock_spin_angle_lookup.num_bins = num_spin_angle_bins
-        mock_spin_angle_lookup.get_spin_angle_index.side_effect = [
-            np.array([3, 3+12]),
-            np.array([5, 6, 5, 7+12])
-        ]
+        energy_step = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=np.uint8)
+        energy_step[0,0,:1] = [7]
+        energy_step[0,2,:4] = [1, 3, 1, 4]
 
-        rng = np.random.default_rng()
+        result = rebin_direct_events_by_energy_and_spin_sector(num_events, spin_sector, energy_step, num_spin_angle_bins, num_energy_bins)
+
+        expected_rebinned_counts = np.zeros((num_epochs, num_priorities, num_energy_bins, num_spin_angle_bins))
+        expected_rebinned_counts[0, 0, 7, 3] = 1
+        expected_rebinned_counts[0, 2, 1, 5] = 2
+        expected_rebinned_counts[0, 2, 3, 6] = 1
+        expected_rebinned_counts[0, 2, 4, 7] = 1
+
+        np.testing.assert_array_equal(result, expected_rebinned_counts)
+
+    def test_rebin_direct_events_for_normalization(self):
+        num_energy_bins = 30
+        num_spin_sectors = 24
         num_priorities = 3
-        num_events = np.ma.masked_array(np.array([[2, 2, 4]]), mask=[[False, True, False]])
         num_epochs = 1
-        spin_angle = rng.random((num_epochs, num_priorities, 15))
-        energy_step = rng.random((num_epochs, num_priorities, 15))
-        result = rebin_counts_by_energy_and_spin_angle(num_events, spin_angle, energy_step, mock_spin_angle_lookup,
-                                                       mock_energy_lookup)
+        event_buffer_len = 15
+        num_events = np.ma.masked_array(np.array([[2, 2, 4]]), mask=[[False, True, False]])
+        spin_sector = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=np.uint8)
+        spin_sector[0,0,:2] = [3, 3+12]
+        spin_sector[0,2,:4] = [5, 6, 5, 7+12]
 
-        expected_rebinned_counts = np.zeros((num_epochs, num_priorities, num_energy_bins, num_spin_angle_bins))
+        energy_step = np.zeros((num_epochs, num_priorities, event_buffer_len), dtype=np.uint8)
+        energy_step[0,0,:2] = [7, 7]
+        energy_step[0,2,:4] = [1, 3, 1, 4]
+        result = rebin_direct_events_for_normalization(num_events, spin_sector, energy_step, num_spin_sectors, num_energy_bins)
+
+        expected_rebinned_counts = np.zeros((num_epochs, num_priorities, num_energy_bins, num_spin_sectors))
         expected_rebinned_counts[0, 0, 7, 3] = 2
         expected_rebinned_counts[0, 2, 1, 5] = 2
         expected_rebinned_counts[0, 2, 3, 6] = 1

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_counts_by_energy_and_spin_angle')
+    @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_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_counts_by_energy_and_spin,
+                                       mock_rebin_direct_events_for_normalization,
                                        mock_spin_angle_lookup_class):
         rng = np.random.default_rng()
 
@@ -539,7 +539,7 @@ def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mas
 
         counts_rebinned_by_energy_and_spin = rng.random(
             (len(epochs), CODICE_LO_L2_NUM_PRIORITIES, num_energy_bins, num_spin_angle_bins))
-        mock_rebin_counts_by_energy_and_spin.return_value = np.array(counts_rebinned_by_energy_and_spin)
+        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))
         mock_calculate_mass_per_charge.return_value = mass_per_charge
@@ -586,8 +586,6 @@ def test_process_l3a_direct_events(self, mock_calculate_mass, mock_calculate_mas
         processor = CodiceLoProcessor(dependencies=input_collection, input_metadata=input_metadata)
         l3a_direct_event_data_product = processor.process_l3a_direct_event_data_product(dependencies)
 
-        mock_spin_angle_lookup_class.assert_called_once()
-
         self.assertEqual(1, mock_calculate_mass.call_count)
         np.testing.assert_equal(mock_calculate_mass.call_args.args[0], expected_apd_energy)
         np.testing.assert_equal(mock_calculate_mass.call_args.args[1], expected_tof)
@@ -597,12 +595,12 @@ 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_counts_by_energy_and_spin.call_count)
-        np.testing.assert_equal(mock_rebin_counts_by_energy_and_spin.call_args.args[0], expected_num_events)
-        np.testing.assert_equal(mock_rebin_counts_by_energy_and_spin.call_args.args[1], expected_spin_angle)
-        np.testing.assert_equal(mock_rebin_counts_by_energy_and_spin.call_args.args[2], expected_energy_step)
-        np.testing.assert_equal(mock_rebin_counts_by_energy_and_spin.call_args.args[3], mock_spin_angle_lookup)
-        np.testing.assert_equal(mock_rebin_counts_by_energy_and_spin.call_args.args[4], dependencies.energy_lookup)
+        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.assertIsInstance(l3a_direct_event_data_product, CodiceLoL3aDirectEventDataProduct)
         self.assertEqual(input_metadata, l3a_direct_event_data_product.input_metadata)