imap-processing 0.19.3__py3-none-any.whl → 0.19.4__py3-none-any.whl

imap_processing/hi/hi_l1c.py

@@ -327,14 +327,18 @@ def pset_counts(
         fill_value=0,
     )
 
-    # Convert list of DEs to pandas dataframe for ease indexing/filtering
-    de_df = l1b_de_dataset.drop_dims("epoch").to_pandas()
+    # Drop events with FILLVAL for trigger_id. This should only occur for a
+    # pointing with no events that gets a single fill event
+    de_ds = l1b_de_dataset.drop_dims("epoch")
 
+    # Remove DEs with invalid trigger_id. This should only occur for a
+    # pointing with no events that gets a single fill event
+    good_mask = de_ds["trigger_id"].data != de_ds["trigger_id"].attrs["FILLVAL"]
     # Remove DEs not in Goodtimes/angles
-    good_mask = good_time_and_phase_mask(
+    good_mask &= good_time_and_phase_mask(
         l1b_de_dataset.event_met.values, l1b_de_dataset.spin_phase.values
     )
-    de_df = de_df[good_mask]
+    de_ds = de_ds.isel(event_met=good_mask)
 
     # The calibration product configuration potentially has different coincidence
     # types for each ESA and different TOF windows for each calibration product,
@@ -346,17 +350,17 @@ def pset_counts(
         # esa_energy_step is recorded for each packet rather than for each DE,
         # so we use ccsds_index to get the esa_energy_step for each DE
         esa_mask = (
-            l1b_de_dataset["esa_energy_step"].data[de_df["ccsds_index"].to_numpy()]
+            l1b_de_dataset["esa_energy_step"].data[de_ds["ccsds_index"].data]
             == esa_energy
         )
         # Now loop over the calibration products for the current ESA energy
         for config_row in esa_df.itertuples():
             # Remove DEs that are not at the current ESA energy and in the list
             # of coincidence types for the current calibration product
-            type_mask = de_df["coincidence_type"].isin(
+            type_mask = de_ds["coincidence_type"].isin(
                 config_row.coincidence_type_values
             )
-            filtered_de_df = de_df[(esa_mask & type_mask)]
+            filtered_de_ds = de_ds.isel(event_met=(esa_mask & type_mask))
 
             # Use the TOF window mask to remove DEs with TOFs outside the allowed range
             tof_fill_vals = {
@@ -366,17 +370,17 @@ def pset_counts(
                 for detector_pair in CalibrationProductConfig.tof_detector_pairs
             }
             tof_in_window_mask = get_tof_window_mask(
-                filtered_de_df, config_row, tof_fill_vals
+                filtered_de_ds, config_row, tof_fill_vals
             )
-            filtered_de_df = filtered_de_df[tof_in_window_mask]
+            filtered_de_ds = filtered_de_ds.isel(event_met=tof_in_window_mask)
 
             # Bin remaining DEs into spin-bins
             i_esa = np.flatnonzero(pset_coords["esa_energy_step"].data == esa_energy)[0]
             # spin_phase is in the range [0, 1). Multiplying by N_SPIN_BINS and
             # truncating to an integer gives the correct bin index
-            spin_bin_indices = (
-                filtered_de_df["spin_phase"].to_numpy() * N_SPIN_BINS
-            ).astype(int)
+            spin_bin_indices = (filtered_de_ds["spin_phase"].data * N_SPIN_BINS).astype(
+                int
+            )
             # When iterating over rows of a dataframe, the names of the multi-index
             # are not preserved. Below, `config_row.Index[0]` gets the
             # calibration_prod value from the namedtuple representing the
@@ -390,15 +394,15 @@ def pset_counts(
 
 
 def get_tof_window_mask(
-    de_df: pd.DataFrame, prod_config_row: NamedTuple, fill_vals: dict
+    de_ds: xr.Dataset, prod_config_row: NamedTuple, fill_vals: dict
 ) -> NDArray[bool]:
     """
     Generate a mask indicating which DEs to keep based on TOF windows.
 
     Parameters
     ----------
-    de_df : pandas.DataFrame
-        The Direct Event dataframe for the DEs to filter based on the TOF
+    de_ds : xarray.Dataset
+        The Direct Event Dataset for the DEs to filter based on the TOF
         windows.
     prod_config_row : NamedTuple
         A single row of the prod config dataframe represented as a named tuple.
@@ -415,11 +419,13 @@ def get_tof_window_mask(
         The mask is intended to directly filter the DE dataframe.
     """
     detector_pairs = CalibrationProductConfig.tof_detector_pairs
-    tof_in_window_mask = np.empty((len(detector_pairs), len(de_df)), dtype=bool)
+    tof_in_window_mask = np.empty(
+        (len(detector_pairs), len(de_ds["event_met"])), dtype=bool
+    )
     for i_pair, detector_pair in enumerate(detector_pairs):
         low_limit = getattr(prod_config_row, f"tof_{detector_pair}_low")
         high_limit = getattr(prod_config_row, f"tof_{detector_pair}_high")
-        tof_array = de_df[f"tof_{detector_pair}"].to_numpy()
+        tof_array = de_ds[f"tof_{detector_pair}"].data
         # The TOF in window mask contains True wherever the TOF is within
         # the configuration low/high bounds OR the FILLVAL is present. The
         # FILLVAL indicates that the detector pair was not hit. DEs with

imap_processing/hi/hi_l2.py

@@ -8,10 +8,10 @@ import pandas as pd
 import xarray as xr
 
 from imap_processing.ena_maps.ena_maps import (
-    AbstractSkyMap,
     HiPointingSet,
     RectangularSkyMap,
 )
+from imap_processing.ena_maps.utils.corrections import PowerLawFluxCorrector
 from imap_processing.ena_maps.utils.naming import MapDescriptor
 from imap_processing.hi.utils import CalibrationProductConfig
 
@@ -23,8 +23,7 @@ VARS_TO_EXPOSURE_TIME_AVERAGE = ["bg_rates", "bg_rates_unc", "obs_date"]
 
 def hi_l2(
     psets: list[str | Path],
-    geometric_factors_path: str | Path,
-    esa_energies_path: str | Path,
+    l2_ancillary_path_dict: dict[str, Path],
     descriptor: str,
 ) -> list[xr.Dataset]:
     """
@@ -34,10 +33,9 @@ def hi_l2(
     ----------
     psets : list of str or pathlib.Path
         List of input PSETs to make a map from.
-    geometric_factors_path : str or pathlib.Path
-        Where to get the geometric factors from.
-    esa_energies_path : str or pathlib.Path
-        Where to get the energies from.
+    l2_ancillary_path_dict : dict[str, pathlib.Path]
+        Mapping containing ancillary file descriptors as keys and file paths as
+        values. Require keys are: ["cal-prod", "esa-energies", "esa-eta-fit-factors"].
     descriptor : str
         Output filename descriptor. Contains full configuration for the options
         of how to generate the map.
@@ -47,33 +45,27 @@ def hi_l2(
     l2_dataset : list[xarray.Dataset]
         Level 2 IMAP-Hi dataset ready to be written to a CDF file.
     """
-    cg_corrected = False
-    map_descriptor = MapDescriptor.from_string(descriptor)
-
-    sky_map = generate_hi_map(
-        psets,
-        geometric_factors_path,
-        esa_energies_path,
-        spin_phase=map_descriptor.spin_phase,
-        output_map=map_descriptor.to_empty_map(),
-        cg_corrected=cg_corrected,
+    logger.info(
+        f"Hi L2 processing running for descriptor: {descriptor} with"
+        f"{len(psets)} PSETs input."
     )
 
-    # Get the map dataset with variables/coordinates in the correct shape
-    # TODO get the correct descriptor and frame
-
-    if not isinstance(sky_map, RectangularSkyMap):
-        raise NotImplementedError("HEALPix map output not supported for Hi")
+    map_descriptor = MapDescriptor.from_string(descriptor)
     if not isinstance(map_descriptor.sensor, str):
         raise ValueError(
             "Invalid map_descriptor. Sensor attribute must be of type str "
             "and be either '45' or '90'"
         )
 
+    sky_map = generate_hi_map(
+        psets,
+        l2_ancillary_path_dict,
+        map_descriptor,
+    )
+
     l2_ds = sky_map.build_cdf_dataset(
         "hi",
         "l2",
-        map_descriptor.frame_descriptor,
         descriptor,
         sensor=map_descriptor.sensor,
     )
@@ -83,12 +75,9 @@ def hi_l2(
 
 def generate_hi_map(
     psets: list[str | Path],
-    geometric_factors_path: str | Path,
-    esa_energies_path: str | Path,
-    output_map: AbstractSkyMap,
-    cg_corrected: bool = False,
-    spin_phase: str = "full",
-) -> AbstractSkyMap:
+    l2_ancillary_path_dict: dict[str, Path],
+    descriptor: MapDescriptor,
+) -> RectangularSkyMap:
     """
     Project Hi PSET data into a sky map.
 
@@ -96,32 +85,30 @@ def generate_hi_map(
     ----------
     psets : list of str or pathlib.Path
         List of input PSETs to make a map from.
-    geometric_factors_path : str or pathlib.Path
-        Where to get the geometric factors from.
-    esa_energies_path : str or pathlib.Path
-        Where to get the energies from.
-    output_map : AbstractSkyMap
-        The map object to collect data into. Determines pixel spacing,
-        coordinate system, etc.
-    cg_corrected : bool, Optional
-        Whether to apply Compton-Getting correction to the energies. Defaults to
-        False.
-    spin_phase : str, Optional
-        Apply filtering to PSET data include ram or anti-ram or full spin data.
-        Defaults to "full".
+    l2_ancillary_path_dict : dict[str, pathlib.Path]
+        Mapping containing ancillary file descriptors as keys and file paths as
+        values. Require keys are: ["cal-prod", "esa-energies", "esa-eta-fit-factors"].
+    descriptor : imap_processing.ena_maps.utils.naming.MapDescriptor
+        Output filename descriptor. Contains full configuration for the options
+        of how to generate the map.
 
     Returns
     -------
-    sky_map : AbstractSkyMap
+    sky_map : RectangularSkyMap
         The sky map with all the PSET data projected into the map.
     """
+    output_map = descriptor.to_empty_map()
+
+    if not isinstance(output_map, RectangularSkyMap):
+        raise NotImplementedError("Healpix map output not supported for Hi")
+
     # TODO: Implement Compton-Getting correction
-    if cg_corrected:
-        raise NotImplementedError
+    if descriptor.frame_descriptor != "sf":
+        raise NotImplementedError("CG correction not implemented for Hi")
 
     for pset_path in psets:
         logger.info(f"Processing {pset_path}")
-        pset = HiPointingSet(pset_path, spin_phase=spin_phase)
+        pset = HiPointingSet(pset_path, spin_phase=descriptor.spin_phase)
 
         # Background rate and uncertainty are exposure time weighted means in
         # the map.
@@ -141,10 +128,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, l2_ancillary_path_dict, descriptor
     )
 
     output_map.data_1d["obs_date"].data = output_map.data_1d["obs_date"].data.astype(
@@ -155,7 +140,8 @@ def generate_hi_map(
 
     # Rename and convert coordinate from esa_energy_step energy
     esa_df = esa_energy_df(
-        esa_energies_path, output_map.data_1d["esa_energy_step"].data
+        l2_ancillary_path_dict["esa-energies"],
+        output_map.data_1d["esa_energy_step"].data,
     )
     output_map.data_1d = output_map.data_1d.rename({"esa_energy_step": "energy"})
     output_map.data_1d = output_map.data_1d.assign_coords(
@@ -222,9 +208,9 @@ def calculate_ena_signal_rates(map_ds: xr.Dataset) -> dict[str, xr.DataArray]:
 
 def calculate_ena_intensity(
     map_ds: xr.Dataset,
-    geometric_factors_path: str | Path,
-    esa_energies_path: str | Path,
-) -> dict[str, xr.DataArray]:
+    l2_ancillary_path_dict: dict[str, Path],
+    descriptor: MapDescriptor,
+) -> xr.Dataset:
     """
     Calculate the ena intensities.
 
@@ -232,52 +218,210 @@ def calculate_ena_intensity(
     ----------
     map_ds : xarray.Dataset
         Map dataset that has ena_signal_rate fields calculated.
-    geometric_factors_path : str or pathlib.Path
-        Where to get the geometric factors from.
-    esa_energies_path : str or pathlib.Path
-        Where to get the esa energies, energy deltas, and geometric factors.
+    l2_ancillary_path_dict : dict[str, pathlib.Path]
+        Mapping containing ancillary file descriptors as keys and file paths as
+        values. Require keys are: ["cal-prod", "esa-energies", "esa-eta-fit-factors"].
+    descriptor : imap_processing.ena_maps.utils.naming.MapDescriptor
+        Output filename descriptor. Contains full configuration for the options
+        of how to generate the map. For this function, the principal data string
+        is used to determine if a flux correction should be applied.
 
     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)
+    cal_prod_df = CalibrationProductConfig.from_csv(l2_ancillary_path_dict["cal-prod"])
     # reindex_like removes esa_energy_steps and calibration products not in the
     # map_ds esa_energy_step and calibration_product coordinates
     geometric_factor = cal_prod_df.to_xarray().reindex_like(map_ds)["geometric_factor"]
     geometric_factor = geometric_factor.transpose(
         *[coord for coord in map_ds.coords if coord in geometric_factor.coords]
     )
-    energy_df = esa_energy_df(esa_energies_path, map_ds["esa_energy_step"].data)
+    energy_df = esa_energy_df(
+        l2_ancillary_path_dict["esa-energies"], map_ds["esa_energy_step"].data
+    )
     esa_energy = energy_df.to_xarray()["nominal_central_energy"]
 
     # 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
     )
-    intensity_vars["ena_intensity_stat_unc"] = np.sqrt(
-        (intensity_vars["ena_intensity_stat_unc"] ** 2).sum(dim="calibration_prod")
+    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,
     )
-    intensity_vars["ena_intensity_sys_err"] = np.sqrt(
-        (intensity_vars["ena_intensity_sys_err"] ** 2).sum(dim="calibration_prod")
+
+    if "raw" not in descriptor.principal_data:
+        # Flux correction
+        corrector = PowerLawFluxCorrector(l2_ancillary_path_dict["esa-eta-fit-factors"])
+        # FluxCorrector does not accept the size 1 epoch dimension. Remove that
+        # dimension by passing the zeroth element.
+        corrected_intensity, corrected_stat_unc = corrector.apply_flux_correction(
+            map_ds["ena_intensity"].values[0],
+            map_ds["ena_intensity_stat_unc"].values[0],
+            esa_energy.data,
+        )
+        # Add the size 1 epoch dimension back in to the corrected fluxes.
+        map_ds["ena_intensity"].data = corrected_intensity[np.newaxis, ...]
+        map_ds["ena_intensity_stat_unc"].data = corrected_stat_unc[np.newaxis, ...]
+
+    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
+    )
+
+    # 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
+    )
+
+    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(

imap_processing/ialirt/utils/create_xarray.py

@@ -144,7 +144,17 @@ def create_xarray_from_records(records: list[dict]) -> xr.Dataset:  # noqa: PLR0
     # Populate the dataset variables
     for i, record in enumerate(records):
         for key, val in record.items():
-            if key in ["apid", "met", "met_in_utc", "ttj2000ns", "last_modified"]:
+            if key in [
+                "apid",
+                "met",
+                "met_in_utc",
+                "ttj2000ns",
+                "last_modified",
+                "sc_position_GSM",
+                "sc_position_GSE",
+                "sc_velocity_GSM",
+                "sc_velocity_GSE",
+            ]:
                 continue
             elif key in ["mag_B_GSE", "mag_B_GSM", "mag_B_RTN"]:
                 dataset[key].data[i, :] = val