2138 hi l2 combine calibration products (IMAP-Science-Operations-Center#2227)

* Add code to correctly combine fluxes and uncertainties across calibration products

* Add logging messages
Temporary solution for uncertainty
Fix hi_l2 tests

* Add comments
fix uncertainty
fix l2 test

* Fix error in improved_stat_unc. Square was already being calculated.

* rename unc_sq to variance

* Add test coverage for combining hi calibration products

* Fix doc build: xarray.DataSet -> xarray.Dataset

* Make sure comments consistently distinguish between uncertainty and variance
Fix bug in improved variance with single calibration product

Author: tmplummer

imap_processing/hi/hi_l2.py

@@ -141,10 +141,8 @@ def generate_hi_map(
             output_map.data_1d[var] /= output_map.data_1d["exposure_factor"]
 
     output_map.data_1d.update(calculate_ena_signal_rates(output_map.data_1d))
-    output_map.data_1d.update(
-        calculate_ena_intensity(
-            output_map.data_1d, geometric_factors_path, esa_energies_path
-        )
+    output_map.data_1d = calculate_ena_intensity(
+        output_map.data_1d, geometric_factors_path, esa_energies_path
     )
 
     output_map.data_1d["obs_date"].data = output_map.data_1d["obs_date"].data.astype(
@@ -224,7 +222,7 @@ def calculate_ena_intensity(
     map_ds: xr.Dataset,
     geometric_factors_path: str | Path,
     esa_energies_path: str | Path,
-) -> dict[str, xr.DataArray]:
+) -> xr.Dataset:
     """
     Calculate the ena intensities.
 
@@ -239,8 +237,9 @@ def calculate_ena_intensity(
 
     Returns
     -------
-    intensity_vars : dict[str, xarray.DataArray]
-        ENA Intensity with statistical and systematic uncertainties.
+    map_ds : xarray.Dataset
+        Map dataset with new variables: ena_intensity, ena_intensity_stat_unc,
+        ena_intensity_sys_err.
     """
     # read calibration product configuration file
     cal_prod_df = CalibrationProductConfig.from_csv(geometric_factors_path)
@@ -255,29 +254,167 @@ def calculate_ena_intensity(
 
     # Convert ENA Signal Rate to Flux
     flux_conversion_divisor = geometric_factor * esa_energy
-    intensity_vars = {
-        "ena_intensity": map_ds["ena_signal_rates"] / flux_conversion_divisor,
-        "ena_intensity_stat_unc": map_ds["ena_signal_rate_stat_unc"]
-        / flux_conversion_divisor,
-        "ena_intensity_sys_err": map_ds["bg_rates_unc"] / flux_conversion_divisor,
-    }
-
-    # TODO: Correctly implement combining of calibration products. For now, just sum
-    # Hi groups direct events into distinct calibration products based on coincidence
-    # type. (See L1B processing and Hi Algorithm Document section 6.1.2) When adding
-    # together different calibration products, a different weighting must be used
-    # than exposure time. (See Hi Algorithm Document Section 3.1.2)
-    intensity_vars["ena_intensity"] = intensity_vars["ena_intensity"].sum(
-        dim="calibration_prod"
+    map_ds["ena_intensity"] = map_ds["ena_signal_rates"] / flux_conversion_divisor
+    map_ds["ena_intensity_stat_unc"] = (
+        map_ds["ena_signal_rate_stat_unc"] / flux_conversion_divisor
+    )
+    map_ds["ena_intensity_sys_err"] = map_ds["bg_rates_unc"] / flux_conversion_divisor
+
+    # Combine calibration products using proper weighted averaging
+    # as described in Hi Algorithm Document Section 3.1.2
+    map_ds = combine_calibration_products(
+        map_ds,
+        geometric_factor,
+        esa_energy,
+    )
+
+    return map_ds
+
+
+def combine_calibration_products(
+    map_ds: xr.Dataset,
+    geometric_factors: xr.DataArray,
+    esa_energies: xr.DataArray,
+) -> xr.Dataset:
+    """
+    Combine calibration products using weighted averaging.
+
+    Implements the algorithm described in Hi Algorithm Document Section 3.1.2
+    for properly combining data from multiple calibration products.
+
+    Parameters
+    ----------
+    map_ds : xarray.Dataset
+        Map dataset that has preliminary intensity variables computed for each
+        calibration product.
+    geometric_factors : xarray.DataArray
+        Geometric factors for each calibration product and energy step.
+    esa_energies : xarray.DataArray
+        Central energies for each energy step.
+
+    Returns
+    -------
+    map_ds : xarray.Dataset
+        Map dataset with updated variables: ena_intensity, ena_intensity_stat_unc,
+        ena_intensity_sys_err now combined across calibration products at each
+        energy level.
+    """
+    ena_flux = map_ds["ena_intensity"]
+    sys_err = map_ds["ena_intensity_sys_err"]
+
+    # Calculate improved statistical variance estimates using geometric factor
+    # ratios to reduce bias from Poisson uncertainty estimation
+    improved_stat_variance = _calculate_improved_stat_variance(
+        map_ds, geometric_factors, esa_energies
     )
-    intensity_vars["ena_intensity_stat_unc"] = np.sqrt(
-        (intensity_vars["ena_intensity_stat_unc"] ** 2).sum(dim="calibration_prod")
+
+    # Calculate total variance
+    # Note that sys_err contains uncertainty, so it must be squared to get
+    # the systematic variance needed in this equation.
+    total_variance = improved_stat_variance + sys_err**2
+
+    # Perform inverse-variance weighted averaging
+    # Handle divide by zero and invalid values
+    with np.errstate(divide="ignore", invalid="ignore"):
+        # Calculate weights for statistical variance combination using only
+        # statistical variance
+        stat_weights = 1.0 / improved_stat_variance
+
+        # Combined statistical uncertainty from inverse-variance formula
+        combined_stat_unc = np.sqrt(1.0 / stat_weights.sum(dim="calibration_prod"))
+
+        # Use total variance weights for flux combination
+        flux_weights = 1.0 / total_variance
+        weighted_flux_sum = (ena_flux * flux_weights).sum(dim="calibration_prod")
+        combined_flux = weighted_flux_sum / flux_weights.sum(dim="calibration_prod")
+
+    map_ds["ena_intensity"] = combined_flux
+    map_ds["ena_intensity_stat_unc"] = combined_stat_unc
+    # For systematic error, just do quadrature sum over the systematic error for
+    # each calibration product.
+    map_ds["ena_intensity_sys_err"] = np.sqrt((sys_err**2).sum(dim="calibration_prod"))
+
+    return map_ds
+
+
+def _calculate_improved_stat_variance(
+    map_ds: xr.Dataset,
+    geometric_factors: xr.DataArray,
+    esa_energies: xr.DataArray,
+) -> xr.DataArray:
+    """
+    Calculate improved statistical variances using geometric factor ratios.
+
+    This implements the algorithm from Hi Algorithm Document Section 3.1.2:
+    For calibration product X, replace N_X in the uncertainty calculation with
+    an improved estimate using geometric factor ratios from all calibration products.
+
+    The key insight is that we can vectorize this by first computing a geometric
+    factor normalized signal rate, then scaling it back for each calibration product.
+
+    Parameters
+    ----------
+    map_ds : xarray.Dataset
+        Map dataset.
+    geometric_factors : xr.DataArray
+        Geometric factors for each calibration product.
+    esa_energies : xarray.DataArray
+        Central energies for each energy step.
+
+    Returns
+    -------
+    improved_variance : xr.DataArray
+        Improved statistical variance estimates.
+    """
+    n_calib_prods = map_ds["ena_intensity"].sizes.get("calibration_prod", 1)
+
+    if n_calib_prods <= 1:
+        # No improvement possible with single calibration product
+        return map_ds["ena_intensity_stat_unc"] ** 2
+
+    logger.debug("Computing geometric factor normalized signal rates")
+
+    # signal_rates = counts / exposure_factor - bg_rates
+    # signal_rates shape is: (n_epoch, n_energy, n_cal_prod, n_spatial_pixels)
+    signal_rates = map_ds["ena_signal_rates"]
+
+    # Compute geometric factor normalized signal rate (vectorized approach)
+    # This represents the weighted average signal rate per unit geometric factor
+    # geometric_factor_norm_signal_rates shape is: (n_epoch, n_energy, n_spatial_pixels)
+    geometric_factor_norm_signal_rates = signal_rates.sum(
+        dim="calibration_prod"
+    ) / geometric_factors.sum(dim="calibration_prod")
+
+    # For each calibration product, the averaged signal rate estimate is:
+    # averaged_signal_rate_i = geometric_factor_norm_signal_rates * geometric_factor_i
+    # averaged_signal_rates shape is: (n_epoch, n_energy, n_cal_prod, n_spatial_pixels)
+    averaged_signal_rates = geometric_factor_norm_signal_rates * geometric_factors
+
+    logger.debug("Including background rates in uncertainty calculation")
+    # Convert averaged signal rates back to flux uncertainties
+    # Total count rates for Poisson uncertainty calculation
+    total_count_rates_for_uncertainty = map_ds["bg_rates"] + averaged_signal_rates
+
+    # Ensure non-negative values for sqrt and minimum of 1 for uncertainty calculation
+    total_count_rates_for_uncertainty = xr.where(
+        total_count_rates_for_uncertainty < 1, 1, total_count_rates_for_uncertainty
     )
-    intensity_vars["ena_intensity_sys_err"] = np.sqrt(
-        (intensity_vars["ena_intensity_sys_err"] ** 2).sum(dim="calibration_prod")
+
+    logger.debug("Computing improved flux uncertainties")
+    # Statistical variance:
+    with np.errstate(divide="ignore", invalid="ignore"):
+        improved_variance = total_count_rates_for_uncertainty / (
+            map_ds["exposure_factor"] * (geometric_factors * esa_energies)
+        )
+
+    # Handle invalid cases by falling back to original uncertainties
+    improved_variance = xr.where(
+        ~np.isfinite(improved_variance) | (geometric_factors == 0),
+        map_ds["ena_intensity_stat_unc"],
+        improved_variance,
     )
 
-    return intensity_vars
+    return improved_variance
 
 
 def esa_energy_df(