FIX: Apply flux correction to background intensities for Lo (#2384)

Lo carries their background intensities separately, so we need to correct those as well.

Author: greglucas

imap_processing/lo/l2/lo_l2.py

@@ -1166,27 +1166,32 @@ def calculate_flux_corrections(dataset: xr.Dataset, flux_factors: Path) -> xr.Da
 
     # 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, ...
-    ]
+
+    # NOTE: We need to apply this to both total flux and background flux
+    for var in ["ena", "bg"]:
+        # 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[f"{var}_intensity"].shape[1:]  # Exclude epoch dimension
+        intensity = (
+            dataset[f"{var}_intensity"].values[0].reshape(len(dataset["energy"]), -1)
+        )
+        stat_uncert = (
+            dataset[f"{var}_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[f"{var}_intensity"].data = corrected_intensity.reshape(input_shape)[
+            np.newaxis, ...
+        ]
+        dataset[f"{var}_intensity_stat_uncert"].data = corrected_stat_unc.reshape(
+            input_shape
+        )[np.newaxis, ...]
 
     return dataset
 

imap_processing/tests/lo/test_lo_l2.py

@@ -566,6 +566,11 @@ def sample_dataset_with_intensities():
         intensity_values * 0.1,
     )
 
+    dataset["bg_intensity"] = dataset["ena_intensity"] * 0.5  # 50% background
+
+    # Add statistical uncertainties (10% of intensity)
+    dataset["bg_intensity_stat_uncert"] = dataset["bg_intensity"] * 0.1
+
     return dataset
 
 
@@ -1142,6 +1147,10 @@ def test_calculate_flux_corrections_energy_dimension_handling(
             ("epoch", "energy", "x", "y"),
             uncert_values.copy(),
         )
+        original_dataset["bg_intensity"] = original_dataset["ena_intensity"] * 0.5
+        original_dataset["bg_intensity_stat_uncert"] = (
+            original_dataset["bg_intensity"] * 0.1
+        )
 
         # Run flux correction on a copy
         dataset_copy = original_dataset.copy(deep=True)