Harrison L3 2881 - GSTU/EYEA - Lo: Add spxnbs support in processor

Author: Menlo Innovations - CAVA Project

imap_l3_processing/lo/lo_processor.py

@@ -1,6 +1,5 @@
 import dataclasses
 from pathlib import Path
-from typing import Optional
 
 import numpy as np
 from imap_processing.spice.geometry import SpiceFrame
@@ -29,7 +28,11 @@ def process(self, spice_frame_name: SpiceFrame = SpiceFrame.ECLIPJ2000) -> list[
         match descriptor:
             case MapDescriptorParts(quantity=MapQuantity.SpectralIndex):
                 deps = LoL3SpectralFitDependencies.fetch_dependencies(self.dependencies)
-                spectral_fit_data = perform_spectral_fit(deps.map_data, descriptor.spectral_index_energy_range)
+                spectral_fit_data = perform_spectral_fit(deps.map_data, descriptor.spectral_index_energy_range or (4, 8))
+                data_product = RectangularSpectralIndexDataProduct(self.input_metadata, spectral_fit_data)
+            case MapDescriptorParts(quantity=MapQuantity.SpectralIndexNBS):
+                deps = LoL3SpectralFitDependencies.fetch_dependencies(self.dependencies)
+                spectral_fit_data = perform_spectral_fit(deps.map_data, descriptor.spectral_index_energy_range or (0, 8))
                 data_product = RectangularSpectralIndexDataProduct(self.input_metadata, spectral_fit_data)
             case MapDescriptorParts(quantity=MapQuantity.ISNBackgroundSubtracted):
                 deps = LoL3ISNBackgroundSubtractedDependencies.fetch_dependencies(self.dependencies)
@@ -58,12 +61,8 @@ def _collapse_pset_dimension(pset: InputRectangularPointingSet) -> InputRectangu
                                    hae_latitude=np.mean(pset.hae_latitude, axis=-1))
 
 
-def perform_spectral_fit(data: RectangularIntensityMapData, spectral_index_range: Optional[tuple[int,int]]) -> RectangularSpectralIndexMapData:
-    if spectral_index_range is not None:
-        start, end = spectral_index_range
-        intensity_data = slice_energy_range_by_bin(data.intensity_map_data, start, end)
-    else:
-        intensity_data = data.intensity_map_data
+def perform_spectral_fit(data: RectangularIntensityMapData, spectral_index_range: tuple[int,int]) -> RectangularSpectralIndexMapData:
+    intensity_data = slice_energy_range_by_bin(data.intensity_map_data, spectral_index_range[0], spectral_index_range[1])
     return RectangularSpectralIndexMapData(
         spectral_index_map_data=fit_spectral_index_map(intensity_data),
         coords=data.coords

tests/lo/test_lo_processor.py

@@ -11,24 +11,26 @@
 from imap_l3_processing.maps.map_models import RectangularSpectralIndexDataProduct, RectangularIntensityDataProduct, \
     InputRectangularPointingSet, ISNRateData, ISNBackgroundSubtractedData, ISNBackgroundSubtractedMapData
 from imap_l3_processing.models import InputMetadata, Instrument
+from tests.maps import test_builders
 
 
 class TestLoProcessor(unittest.TestCase):
 
     @patch('imap_l3_processing.hi.hi_processor.MapProcessor.get_parent_file_names')
+    @patch('imap_l3_processing.lo.lo_processor.slice_energy_range_by_bin')
     @patch('imap_l3_processing.lo.lo_processor.LoL3SpectralFitDependencies.fetch_dependencies')
     @patch('imap_l3_processing.lo.lo_processor.fit_spectral_index_map')
     @patch('imap_l3_processing.lo.lo_processor.save_data')
-    def test_process_spectral_index(self, mock_save_data,
-                                    mock_fit_spectral_index_map, mock_fetch_dependencies,
-                                    mock_get_parent_file_names):
+    def test_process_spectral_index(self, mock_save_data, mock_fit_spectral_index_map, mock_fetch_dependencies,
+                                    mock_slice_energy_range_by_bin, mock_get_parent_file_names):
         mock_get_parent_file_names.return_value = ["some_input_file_name"]
 
         input_collection = Mock()
-        lo_l3_spectral_fit_dependency = mock_fetch_dependencies.return_value
+        lo_l3_spectral_fit_dependency = Mock()
+        lo_l3_spectral_fit_dependency.map_data.intensity_map_data = test_builders.create_intensity_map_data()
         lo_l3_spectral_fit_dependency.map_data.intensity_map_data.energy = np.array(
-            [1, 10, 1000, 10000, 100000, 1000000, 10000000])
-        mock_fit_spectral_index_map.return_value = Mock()
+            [1, 10, 100, 1000, 10000, 100000, 1000000, 10000000])
+        mock_fetch_dependencies.return_value = lo_l3_spectral_fit_dependency
 
         metadata = InputMetadata(instrument="lo",
                                  data_level="l3",
@@ -40,8 +42,10 @@ def test_process_spectral_index(self, mock_save_data,
         processor = LoProcessor(input_collection, input_metadata=metadata)
         product = processor.process()
 
-        mock_fetch_dependencies.assert_called_with(input_collection)
-        mock_fit_spectral_index_map.assert_called_once_with(lo_l3_spectral_fit_dependency.map_data.intensity_map_data)
+        mock_fetch_dependencies.assert_called_once_with(input_collection)
+        mock_slice_energy_range_by_bin.assert_called_once_with(
+            lo_l3_spectral_fit_dependency.map_data.intensity_map_data, 4, 8)
+        mock_fit_spectral_index_map.assert_called_once_with(mock_slice_energy_range_by_bin.return_value)
 
         data_product = mock_save_data.call_args_list[0].args[0]
 
@@ -59,42 +63,57 @@ def test_process_spectral_index(self, mock_save_data,
     @patch('imap_l3_processing.lo.lo_processor.save_data')
     @patch('imap_l3_processing.lo.lo_processor.slice_energy_range_by_bin')
     def test_process_spectral_index_with_range_specified(self, mock_slice_energy_range_by_bin, mock_save_data,
-                                    mock_fit_spectral_index_map, mock_fetch_dependencies,
-                                    mock_get_parent_file_names):
-        mock_get_parent_file_names.return_value = ["some_input_file_name"]
+                                                         mock_fit_spectral_index_map, mock_fetch_dependencies,
+                                                         mock_get_parent_file_names):
+        test_cases = [
+            ("l090-spx0104-h-hf-sp-ram-hae-6deg-1yr", (1, 4)),
+            ("l090-spxnbs0105-h-hf-sp-ram-hae-6deg-1yr", (1, 5)),
+            ("l090-spx-h-hf-sp-ram-hae-6deg-1yr", (4, 8)),
+            ("l090-spxnbs-h-hf-sp-ram-hae-6deg-1yr", (0, 8)),
+        ]
 
-        input_collection = Mock()
-        lo_l3_spectral_fit_dependency = mock_fetch_dependencies.return_value
-        lo_l3_spectral_fit_dependency.map_data.intensity_map_data.energy = np.array(
-            [1, 10, 1000, 10000, 100000, 1000000, 10000000])
+        for descriptor, (expected_energy_start, expected_energy_end) in test_cases:
+            mock_fit_spectral_index_map.reset_mock()
+            mock_fetch_dependencies.reset_mock()
+            mock_slice_energy_range_by_bin.reset_mock()
+            mock_save_data.reset_mock()
 
-        mock_fit_spectral_index_map.return_value = Mock()
+            with self.subTest(descriptor):
+                mock_get_parent_file_names.return_value = ["some_input_file_name"]
 
+                input_collection = Mock()
+                lo_l3_spectral_fit_dependency = mock_fetch_dependencies.return_value
+                lo_l3_spectral_fit_dependency.map_data.intensity_map_data.energy = np.array(
+                    [1, 10, 1000, 10000, 100000, 1000000, 10000000])
 
-        metadata = InputMetadata(instrument="lo",
-                                 data_level="l3",
-                                 version="v000",
-                                 start_date=datetime(2020, 1, 1, 1),
-                                 end_date=datetime(2020, 1, 1, 1),
-                                 descriptor="l090-spx0104-h-hf-sp-ram-hae-6deg-1yr")
+                metadata = InputMetadata(instrument="lo",
+                                         data_level="l3",
+                                         version="v000",
+                                         start_date=datetime(2020, 1, 1, 1),
+                                         end_date=datetime(2020, 1, 1, 1),
+                                         descriptor=descriptor)
 
-        processor = LoProcessor(input_collection, input_metadata=metadata)
-        product = processor.process()
+                processor = LoProcessor(input_collection, input_metadata=metadata)
+                product = processor.process()
 
-        mock_fetch_dependencies.assert_called_with(input_collection)
-        mock_slice_energy_range_by_bin.assert_called_with(lo_l3_spectral_fit_dependency.map_data.intensity_map_data, 1, 4)
+                mock_fetch_dependencies.assert_called_once_with(input_collection)
+                mock_slice_energy_range_by_bin.assert_called_once_with(
+                    lo_l3_spectral_fit_dependency.map_data.intensity_map_data,
+                    expected_energy_start,
+                    expected_energy_end
+                )
 
-        mock_fit_spectral_index_map.assert_called_once_with(mock_slice_energy_range_by_bin.return_value)
+                mock_fit_spectral_index_map.assert_called_once_with(mock_slice_energy_range_by_bin.return_value)
 
-        data_product = mock_save_data.call_args_list[0].args[0]
+                data_product = mock_save_data.call_args_list[0].args[0]
 
-        self.assertIsInstance(data_product, RectangularSpectralIndexDataProduct)
-        self.assertEqual(data_product.data.spectral_index_map_data,
-                         mock_fit_spectral_index_map.return_value)
-        self.assertEqual(data_product.data.coords, lo_l3_spectral_fit_dependency.map_data.coords)
-        self.assertEqual(data_product.input_metadata, processor.input_metadata)
-        self.assertEqual(data_product.parent_file_names, ["some_input_file_name"])
-        self.assertEqual([mock_save_data.return_value], product)
+                self.assertIsInstance(data_product, RectangularSpectralIndexDataProduct)
+                self.assertEqual(data_product.data.spectral_index_map_data,
+                                 mock_fit_spectral_index_map.return_value)
+                self.assertEqual(data_product.data.coords, lo_l3_spectral_fit_dependency.map_data.coords)
+                self.assertEqual(data_product.input_metadata, processor.input_metadata)
+                self.assertEqual(data_product.parent_file_names, ["some_input_file_name"])
+                self.assertEqual([mock_save_data.return_value], product)
 
     @patch('imap_l3_processing.lo.lo_processor.MapProcessor.get_parent_file_names')
     @patch("imap_l3_processing.lo.lo_processor.HiLoL3SurvivalDependencies.fetch_dependencies")
@@ -180,7 +199,6 @@ def test_rejects_unimplemented_descriptors(self):
                 self.assertEqual(exception_args, cm.exception.args)
 
     def test_isn_background_subtraction(self):
-
         input_data: ISNRateData = ISNRateData(
             epoch=sentinel.epoch,
             solid_angle=sentinel.solid_angle,