Lo flux correction (IMAP-Science-Operations-Center#2343)

* ENH: Add flux corrector to Lo processing

Author: greglucas

imap_processing/lo/l2/lo_l2.py

@@ -10,6 +10,7 @@
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.ena_maps import ena_maps
 from imap_processing.ena_maps.ena_maps import AbstractSkyMap, RectangularSkyMap
+from imap_processing.ena_maps.utils.corrections import PowerLawFluxCorrector
 from imap_processing.ena_maps.utils.naming import MapDescriptor
 from imap_processing.lo import lo_ancillary
 from imap_processing.spice.time import et_to_datetime64, ttj2000ns_to_et
@@ -77,15 +78,23 @@ def lo_l2(
     logger.info("Step 4: Calculating rates and intensities")
 
     # Determine if corrections are needed and prepare oxygen data if required
-    sputtering_correction, bootstrap_correction, o_map_dataset = _prepare_corrections(
+    (
+        sputtering_correction,
+        bootstrap_correction,
+        flux_correction,
+        o_map_dataset,
+        flux_factors,
+    ) = _prepare_corrections(
         map_descriptor, descriptor, sci_dependencies, anc_dependencies
     )
 
     dataset = calculate_all_rates_and_intensities(
         dataset,
         sputtering_correction=sputtering_correction,
         bootstrap_correction=bootstrap_correction,
+        flux_correction=flux_correction,
         o_map_dataset=o_map_dataset,
+        flux_factors=flux_factors,
     )
 
     logger.info("Step 5: Finalizing dataset with attributes")
@@ -100,7 +109,7 @@ def _prepare_corrections(
     descriptor: str,
     sci_dependencies: dict,
     anc_dependencies: list,
-) -> tuple[bool, bool, xr.Dataset | None]:
+) -> tuple[bool, bool, bool, xr.Dataset | None, Path | None]:
     """
     Determine what corrections are needed and prepare oxygen dataset if required.
 
@@ -130,7 +139,9 @@ def _prepare_corrections(
     # Default values - no corrections needed
     sputtering_correction = False
     bootstrap_correction = False
+    flux_correction = False
     o_map_dataset = None
+    flux_factors: None | Path = None
 
     # Sputtering and bootstrap corrections are only applied to hydrogen ENA data
     # Guard against recursion: don't process oxygen for oxygen maps
@@ -145,7 +156,24 @@ def _prepare_corrections(
         sputtering_correction = True
         bootstrap_correction = True
 
-    return sputtering_correction, bootstrap_correction, o_map_dataset
+    if "raw" not in map_descriptor.principal_data:
+        flux_correction = True
+        try:
+            flux_factors = next(
+                x for x in anc_dependencies if "esa-eta-fit-factors" in str(x)
+            )
+        except StopIteration:
+            raise ValueError(
+                "No flux correction factor file found in ancillary dependencies"
+            ) from None
+
+    return (
+        sputtering_correction,
+        bootstrap_correction,
+        flux_correction,
+        o_map_dataset,
+        flux_factors,
+    )
 
 
 # =============================================================================
@@ -664,7 +692,9 @@ def calculate_all_rates_and_intensities(
     dataset: xr.Dataset,
     sputtering_correction: bool = False,
     bootstrap_correction: bool = False,
+    flux_correction: bool = False,
     o_map_dataset: xr.Dataset | None = None,
+    flux_factors: Path | None = None,
 ) -> xr.Dataset:
     """
     Calculate rates and intensities with proper error propagation.
@@ -679,8 +709,13 @@ def calculate_all_rates_and_intensities(
     bootstrap_correction : bool, optional
         Whether to apply bootstrap corrections to intensities.
         Default is False.
+    flux_correction : bool, optional
+        Whether to apply flux corrections to intensities.
+        Default is False.
     o_map_dataset : xr.Dataset, optional
         Dataset specifically for oxygen, needed for sputtering corrections.
+    flux_factors : Path, optional
+        Path to flux factor file for flux corrections.
 
     Returns
     -------
@@ -705,7 +740,13 @@ def calculate_all_rates_and_intensities(
     if bootstrap_correction:
         dataset = calculate_bootstrap_corrections(dataset)
 
-    # Step 6: Clean up intermediate variables
+    # Optional Step 6: Calculate flux corrections
+    if flux_correction:
+        if flux_factors is None:
+            raise ValueError("Flux factors file must be provided for flux corrections")
+        dataset = calculate_flux_corrections(dataset, flux_factors)
+
+    # Step 7: Clean up intermediate variables
     dataset = cleanup_intermediate_variables(dataset)
 
     return dataset
@@ -1084,6 +1125,56 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
     return dataset
 
 
+def calculate_flux_corrections(dataset: xr.Dataset, flux_factors: Path) -> xr.Dataset:
+    """
+    Calculate flux corrections for intensities.
+
+    Uses the shared ena maps ``PowerLawFluxCorrector`` class to do the
+    correction calculations.
+
+    Parameters
+    ----------
+    dataset : xr.Dataset
+        Dataset with count rates, geometric factors, and center energies.
+    flux_factors : Path
+        Path to the eta flux factor file to use for corrections. Read in as
+        an ancillary file in the preprocessing step.
+
+    Returns
+    -------
+    xr.Dataset
+        Dataset with calculated flux-corrected intensities and their
+        uncertainties for the specified species.
+    """
+    logger.info("Applying flux corrections")
+
+    # Flux correction
+    corrector = PowerLawFluxCorrector(flux_factors)
+    # FluxCorrector works on (energy, :) arrays, so we need to flatten the map
+    # spatial dimensions for the correction and then reshape back after.
+    input_shape = dataset["ena_intensity"].shape[1:]  # Exclude epoch dimension
+    intensity = dataset["ena_intensity"].values[0].reshape(len(dataset["energy"]), -1)
+    stat_uncert = (
+        dataset["ena_intensity_stat_uncert"]
+        .values[0]
+        .reshape(len(dataset["energy"]), -1)
+    )
+    corrected_intensity, corrected_stat_unc = corrector.apply_flux_correction(
+        intensity,
+        stat_uncert,
+        dataset["energy"].data,
+    )
+    # Add the size 1 epoch dimension back in to the corrected fluxes.
+    dataset["ena_intensity"].data = corrected_intensity.reshape(input_shape)[
+        np.newaxis, ...
+    ]
+    dataset["ena_intensity_stat_uncert"].data = corrected_stat_unc.reshape(input_shape)[
+        np.newaxis, ...
+    ]
+
+    return dataset
+
+
 def cleanup_intermediate_variables(dataset: xr.Dataset) -> xr.Dataset:
     """
     Remove intermediate variables that were only needed for calculations.

imap_processing/tests/lo/test_lo_l2.py

@@ -1,5 +1,6 @@
 """Comprehensive test suite for IMAP-Lo L2 data processing."""
 
+from pathlib import Path
 from unittest.mock import Mock, patch
 
 import numpy as np
@@ -23,6 +24,7 @@
     calculate_backgrounds,
     calculate_bootstrap_corrections,
     calculate_efficiency_corrected_quantities,
+    calculate_flux_corrections,
     calculate_intensities,
     calculate_rates,
     calculate_sputtering_corrections,
@@ -532,6 +534,56 @@ def sample_dataset_with_bootstrap_data():
     return dataset
 
 
+@pytest.fixture
+def lo_flux_factors_file():
+    """Path to the LO flux factors test file."""
+    # Use the actual test data file from the ena_maps test data
+    test_data_path = Path(__file__).parent.parent / "ena_maps" / "data"
+    return test_data_path / "imap_lo_esa-eta-fit-factors_20240101_v001.csv"
+
+
+@pytest.fixture
+def sample_dataset_with_intensities():
+    """Create a dataset with intensities for flux correction testing."""
+    n_energy = 7
+    n_lon, n_lat = 6, 4  # Small for testing
+
+    # Create realistic energy values matching the flux factors file
+    energy_values = np.array([16.35, 30.56, 56.4, 105, 199.8, 407.5, 795.3])
+
+    coords = {
+        "epoch": [8.1794907049e17],
+        "energy": energy_values,
+        "longitude": np.linspace(0, 360, n_lon, endpoint=False),
+        "latitude": np.linspace(-90, 90, n_lat),
+    }
+
+    # Create intensity values with some spatial and energy structure
+    intensity_values = np.ones((1, n_energy, n_lon, n_lat))
+    for i in range(n_energy):
+        # Power law: I = I0 * (E/E0)^(-2.0)
+        intensity_values[0, i, :, :] = 1e6 * (energy_values[i] / 100.0) ** (-2.0)
+
+    # Add some spatial structure
+    for j in range(n_lon):
+        for k in range(n_lat):
+            intensity_values[0, :, j, k] *= 1.0 + 0.1 * np.sin(j) * np.cos(k)
+
+    dataset = xr.Dataset(coords=coords)
+    dataset["ena_intensity"] = (
+        ("epoch", "energy", "longitude", "latitude"),
+        intensity_values,
+    )
+
+    # Add statistical uncertainties (10% of intensity)
+    dataset["ena_intensity_stat_uncert"] = (
+        ("epoch", "energy", "longitude", "latitude"),
+        intensity_values * 0.1,
+    )
+
+    return dataset
+
+
 # =============================================================================
 # UNIT TESTS FOR INDIVIDUAL FUNCTIONS
 # =============================================================================
@@ -1002,6 +1054,129 @@ def test_calculate_backgrounds_zero_exposure(self):
         assert np.all(np.isinf(result["bg_rates_stat_uncert"].values))
 
 
+class TestCalculateFluxCorrections:
+    """Tests for the calculate_flux_corrections function."""
+
+    def test_calculate_flux_corrections_basic(
+        self, sample_dataset_with_intensities, lo_flux_factors_file
+    ):
+        """Test basic flux correction calculation."""
+        # Make a copy to avoid modifying the original fixture
+        original_dataset = sample_dataset_with_intensities.copy(deep=True)
+
+        # Run flux correction
+        result = calculate_flux_corrections(original_dataset, lo_flux_factors_file)
+
+        # Verify that the function returns a dataset
+        assert isinstance(result, xr.Dataset)
+
+        # Verify that intensity variables are present
+        assert "ena_intensity" in result.data_vars
+        assert "ena_intensity_stat_uncert" in result.data_vars
+
+        # Verify that data shape is preserved
+        original_shape = sample_dataset_with_intensities["ena_intensity"].shape
+        assert result["ena_intensity"].shape == original_shape
+
+        # Check that corrections were applied by comparing to the original fixture
+        # (not the potentially modified copy)
+        original_intensity = sample_dataset_with_intensities["ena_intensity"].values
+        corrected_intensity = result["ena_intensity"].values
+
+        # Check for meaningful differences
+        relative_diff = np.abs(
+            (corrected_intensity - original_intensity) / original_intensity
+        )
+        max_relative_diff = np.max(relative_diff)
+        # Should have at least 10% change somewhere
+        assert max_relative_diff > 0.1, (
+            f"Max relative difference was only {max_relative_diff}"
+        )
+
+        # Verify that uncertainties were also corrected
+        original_uncert = sample_dataset_with_intensities[
+            "ena_intensity_stat_uncert"
+        ].values
+        corrected_uncert = result["ena_intensity_stat_uncert"].values
+        uncert_relative_diff = np.abs(
+            (corrected_uncert - original_uncert) / original_uncert
+        )
+        max_uncert_diff = np.max(uncert_relative_diff)
+        # Should have at least 10% change in uncertainties too
+        assert max_uncert_diff > 0.1, (
+            f"Max uncertainty relative difference was only {max_uncert_diff}"
+        )
+
+    def test_calculate_flux_corrections_preserves_other_vars(
+        self, sample_dataset_with_intensities, lo_flux_factors_file
+    ):
+        """Test that flux correction preserves other variables in the dataset."""
+        # Add an extra variable to the dataset
+        sample_dataset_with_intensities["extra_var"] = (("energy",), np.ones(7))
+
+        result = calculate_flux_corrections(
+            sample_dataset_with_intensities, lo_flux_factors_file
+        )
+
+        # Verify that other variables are preserved
+        assert "extra_var" in result.data_vars
+        np.testing.assert_array_equal(
+            result["extra_var"].values,
+            sample_dataset_with_intensities["extra_var"].values,
+        )
+
+    def test_calculate_flux_corrections_energy_dimension_handling(
+        self, lo_flux_factors_file
+    ):
+        """Test that flux correction properly handles energy dimension reshaping."""
+        # Create a dataset with different spatial dimensions
+        n_energy = 7
+        n_x, n_y = 12, 8  # Different spatial dimensions
+
+        energy_values = np.array([16.35, 30.56, 56.4, 105, 199.8, 407.5, 795.3])
+
+        coords = {
+            "epoch": [8.1794907049e17],
+            "energy": energy_values,
+            "x": np.arange(n_x),
+            "y": np.arange(n_y),
+        }
+
+        # Create intensity values with energy-dependent structure (power law)
+        intensity_values = np.ones((1, n_energy, n_x, n_y))
+        for i in range(n_energy):
+            intensity_values[0, i, :, :] = 1e6 * (energy_values[i] / 100.0) ** (-2.0)
+        uncert_values = intensity_values * 0.1
+
+        original_dataset = xr.Dataset(coords=coords)
+        original_dataset["ena_intensity"] = (
+            ("epoch", "energy", "x", "y"),
+            intensity_values.copy(),
+        )
+        original_dataset["ena_intensity_stat_uncert"] = (
+            ("epoch", "energy", "x", "y"),
+            uncert_values.copy(),
+        )
+
+        # Run flux correction on a copy
+        dataset_copy = original_dataset.copy(deep=True)
+        result = calculate_flux_corrections(dataset_copy, lo_flux_factors_file)
+
+        # Verify shape is preserved
+        assert result["ena_intensity"].shape == (1, n_energy, n_x, n_y)
+        assert result["ena_intensity_stat_uncert"].shape == (1, n_energy, n_x, n_y)
+
+        # Verify corrections were applied by checking for meaningful differences
+        original_values = original_dataset["ena_intensity"].values
+        corrected_values = result["ena_intensity"].values
+        relative_diff = np.abs((corrected_values - original_values) / original_values)
+        max_relative_diff = np.max(relative_diff)
+        # Should have at least 10% change somewhere (flux corrections are significant)
+        assert max_relative_diff > 0.1, (
+            f"Max relative difference was only {max_relative_diff}"
+        )
+
+
 class TestCalculateSputteringCorrections:
     """Tests for the calculate_sputtering_corrections function."""
 
@@ -1970,11 +2145,13 @@ def test_calculate_all_rates_and_intensities_complete(self):
 class TestIntegrationWithMocks:
     """Integration tests using mocked external dependencies."""
 
-    def test_lo_l2_integration_minimal(self, minimal_pset_for_species):
+    def test_lo_l2_integration_minimal(
+        self, minimal_pset_for_species, lo_flux_factors_file
+    ):
         """Test the main lo_l2 function with minimal mocking."""
         # Test with hydrogen data
         sci_dependencies = {"imap_lo_l1c_pset": [minimal_pset_for_species]}
-        anc_dependencies = []
+        anc_dependencies = [lo_flux_factors_file]  # Include flux factors file
         descriptor = "l090-ena-h-sf-nsp-ram-hae-6deg-3mo"
 
         # Mock the complex external dependencies to return simple results