imap-processing 0.18.0__py3-none-any.whl → 0.19.2__py3-none-any.whl

imap_processing/lo/l1c/lo_l1c.py

@@ -1,15 +1,30 @@
 """IMAP-Lo L1C Data Processing."""
 
+import logging
 from dataclasses import Field
 from enum import Enum
 
 import numpy as np
-import pandas as pd
 import xarray as xr
 from scipy.stats import binned_statistic_dd
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
-from imap_processing.spice.time import met_to_ttj2000ns
+from imap_processing.lo import lo_ancillary
+from imap_processing.spice.repoint import get_pointing_times
+from imap_processing.spice.spin import get_spin_number
+from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_met
+
+N_ESA_ENERGY_STEPS = 7
+N_SPIN_ANGLE_BINS = 3600
+N_OFF_ANGLE_BINS = 40
+# 1 time, 7 energy steps, 3600 spin angle bins, and 40 off angle bins
+PSET_SHAPE = (1, N_ESA_ENERGY_STEPS, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)
+PSET_DIMS = ["epoch", "esa_energy_step", "spin_angle", "off_angle"]
+ESA_ENERGY_STEPS = np.arange(N_ESA_ENERGY_STEPS) + 1  # 1 to 7 inclusive
+SPIN_ANGLE_BIN_EDGES = np.linspace(0, 360, N_SPIN_ANGLE_BINS + 1)
+SPIN_ANGLE_BIN_CENTERS = (SPIN_ANGLE_BIN_EDGES[:-1] + SPIN_ANGLE_BIN_EDGES[1:]) / 2
+OFF_ANGLE_BIN_EDGES = np.linspace(-2, 2, N_OFF_ANGLE_BINS + 1)
+OFF_ANGLE_BIN_CENTERS = (OFF_ANGLE_BIN_EDGES[:-1] + OFF_ANGLE_BIN_EDGES[1:]) / 2
 
 
 class FilterType(str, Enum):
@@ -52,10 +67,63 @@ def lo_l1c(sci_dependencies: dict, anc_dependencies: list) -> list[xr.Dataset]:
     if "imap_lo_l1b_de" in sci_dependencies:
         logical_source = "imap_lo_l1c_pset"
         l1b_de = sci_dependencies["imap_lo_l1b_de"]
-
         l1b_goodtimes_only = filter_goodtimes(l1b_de, anc_dependencies)
-        pset = initialize_pset(l1b_goodtimes_only, attr_mgr, logical_source)
-        full_counts = create_pset_counts(l1b_goodtimes_only)
+
+        # Set the pointing start and end times based on the first epoch
+        pointing_start_met, pointing_end_met = get_pointing_times(
+            ttj2000ns_to_met(l1b_goodtimes_only["epoch"][0].item())
+        )
+
+        pset = xr.Dataset(
+            coords={"epoch": np.array([met_to_ttj2000ns(pointing_start_met)])},
+            attrs=attr_mgr.get_global_attributes(logical_source),
+        )
+
+        # ESA mode needs to be added to L1B DE. Adding try statement
+        # to avoid error until it's available in the dataset
+        if "esa_mode" not in l1b_de:
+            logging.debug(
+                "ESA mode not found in L1B DE dataset. \
+                Setting to default value of 0 for Hi-Res."
+            )
+            pset["esa_mode"] = xr.DataArray(
+                np.array([0]),
+                dims=["epoch"],
+                attrs=attr_mgr.get_variable_attributes("esa_mode"),
+            )
+        else:
+            pset["esa_mode"] = xr.DataArray(
+                np.array([l1b_de["esa_mode"].values[0]]),
+                dims=["epoch"],
+                attrs=attr_mgr.get_variable_attributes("esa_mode"),
+            )
+
+        pset["pointing_start_met"] = xr.DataArray(
+            np.array([pointing_start_met]),
+            dims="epoch",
+            attrs=attr_mgr.get_variable_attributes("pointing_start_met"),
+        )
+        pset["pointing_end_met"] = xr.DataArray(
+            np.array([pointing_end_met]),
+            dims="epoch",
+            attrs=attr_mgr.get_variable_attributes("pointing_end_met"),
+        )
+
+        # Get the start and end spin numbers based on the pointing start and end MET
+        pset["start_spin_number"] = xr.DataArray(
+            [get_spin_number(pset["pointing_start_met"].item())],
+            dims="epoch",
+            attrs=attr_mgr.get_variable_attributes("start_spin_number"),
+        )
+        pset["end_spin_number"] = xr.DataArray(
+            [get_spin_number(pset["pointing_end_met"].item())],
+            dims="epoch",
+            attrs=attr_mgr.get_variable_attributes("end_spin_number"),
+        )
+
+        full_counts = create_pset_counts(l1b_de, FilterType.NONE)
+
+        # Set the counts
         pset["triples_counts"] = create_pset_counts(
             l1b_goodtimes_only, FilterType.TRIPLES
         )
@@ -64,63 +132,50 @@ def lo_l1c(sci_dependencies: dict, anc_dependencies: list) -> list[xr.Dataset]:
         )
         pset["h_counts"] = create_pset_counts(l1b_goodtimes_only, FilterType.HYDROGEN)
         pset["o_counts"] = create_pset_counts(l1b_goodtimes_only, FilterType.OXYGEN)
+
+        # Set the exposure time
         pset["exposure_time"] = calculate_exposure_times(
             full_counts, l1b_goodtimes_only
         )
+
+        # Set backgrounds
+        (
+            pset["h_background_rates"],
+            pset["h_background_rates_stat_uncert"],
+            pset["h_background_rates_sys_err"],
+        ) = set_background_rates(
+            pset["pointing_start_met"].item(),
+            pset["pointing_end_met"].item(),
+            FilterType.HYDROGEN,
+            anc_dependencies,
+            attr_mgr,
+        )
+
+        (
+            pset["o_background_rates"],
+            pset["o_background_rates_stat_uncert"],
+            pset["o_background_rates_sys_err"],
+        ) = set_background_rates(
+            pset["pointing_start_met"].item(),
+            pset["pointing_end_met"].item(),
+            FilterType.OXYGEN,
+            anc_dependencies,
+            attr_mgr,
+        )
+
     pset.attrs = attr_mgr.get_global_attributes(logical_source)
-    # TODO: Temp fix before adding attribute variables.
-    #  CDF won't open if DEPEND_0 is not deleted currently.
-    del pset["epoch"].attrs["DEPEND_0"]
 
     pset = pset.assign_coords(
         {
-            "energy": np.arange(1, 8),
-            "longitude": np.arange(3600),
-            "latitude": np.arange(40),
+            "esa_energy_step": ESA_ENERGY_STEPS,
+            "spin_angle": SPIN_ANGLE_BIN_CENTERS,
+            "off_angle": OFF_ANGLE_BIN_CENTERS,
         }
     )
 
     return [pset]
 
 
-def initialize_pset(
-    l1b_de: xr.Dataset, attr_mgr: ImapCdfAttributes, logical_source: str
-) -> xr.Dataset:
-    """
-    Initialize the PSET dataset and set the Epoch.
-
-    The Epoch time is set to the first of the L1B
-    Direct Event times. There is one Epoch per PSET file.
-
-    Parameters
-    ----------
-    l1b_de : xarray.Dataset
-        L1B Direct Event dataset.
-    attr_mgr : ImapCdfAttributes
-        Attribute manager used to get the L1C attributes.
-    logical_source : str
-        The logical source of the pset.
-
-    Returns
-    -------
-    pset : xarray.Dataset
-        Initialized PSET dataset.
-    """
-    pset = xr.Dataset(
-        attrs=attr_mgr.get_global_attributes(logical_source),
-    )
-    # TODO: Need to create utility to get start of repointing to use
-    #  for the pset epoch time. Setting to first DE for now
-    pset_epoch = l1b_de["epoch"][0].item()
-    pset["epoch"] = xr.DataArray(
-        np.array([pset_epoch]),
-        dims=["epoch"],
-        attrs=attr_mgr.get_variable_attributes("epoch"),
-    )
-
-    return pset
-
-
 def filter_goodtimes(l1b_de: xr.Dataset, anc_dependencies: list) -> xr.Dataset:
     """
     Filter the L1B Direct Event dataset to only include good times.
@@ -141,17 +196,19 @@ def filter_goodtimes(l1b_de: xr.Dataset, anc_dependencies: list) -> xr.Dataset:
         Filtered L1B Direct Event dataset.
     """
     # the goodtimes are currently the only ancillary file needed for L1C processing
-    goodtimes_table_df = pd.read_csv(anc_dependencies[0])
+    goodtimes_table_df = lo_ancillary.read_ancillary_file(
+        next(str(s) for s in anc_dependencies if "good-times" in str(s))
+    )
 
     # convert goodtimes from MET to TTJ2000
-    goodtimes_start = met_to_ttj2000ns(goodtimes_table_df["GoodTime_strt"])
+    goodtimes_start = met_to_ttj2000ns(goodtimes_table_df["GoodTime_start"])
     goodtimes_end = met_to_ttj2000ns(goodtimes_table_df["GoodTime_end"])
 
     # Create a mask for epochs within any of the start/end time ranges
     goodtimes_mask = np.zeros_like(l1b_de["epoch"], dtype=bool)
 
     # Iterate over the good times and create a mask
-    for start, end in zip(goodtimes_start, goodtimes_end):
+    for start, end in zip(goodtimes_start, goodtimes_end, strict=False):
         goodtimes_mask |= (l1b_de["epoch"] >= start) & (l1b_de["epoch"] < end)
 
     # Filter the dataset using the mask
@@ -206,9 +263,9 @@ def create_pset_counts(
             "001000",
         ],
         # hydrogen species identifier
-        FilterType.HYDROGEN: "h",
+        FilterType.HYDROGEN: "H",
         # oxygen species identifier
-        FilterType.OXYGEN: "o",
+        FilterType.OXYGEN: "O",
     }
 
     # if the filter string is triples or doubles, filter using the coincidence type
@@ -249,7 +306,7 @@ def create_pset_counts(
 
     counts = xr.DataArray(
         data=hist.astype(np.int16),
-        dims=["epoch", "energy", "longitude", "latitude"],
+        dims=PSET_DIMS,
     )
 
     return counts
@@ -275,11 +332,6 @@ def calculate_exposure_times(counts: xr.DataArray, l1b_de: xr.Dataset) -> xr.Dat
     exposure_time : xarray.DataArray
         The exposure times for the L1B Direct Event dataset.
     """
-    # Create bin edges
-    lon_edges = np.arange(3601)
-    lat_edges = np.arange(41)
-    energy_edges = np.arange(8)
-
     data = np.column_stack(
         (l1b_de["esa_step"], l1b_de["pointing_bin_lon"], l1b_de["pointing_bin_lat"])
     )
@@ -289,14 +341,19 @@ def calculate_exposure_times(counts: xr.DataArray, l1b_de: xr.Dataset) -> xr.Dat
         # exposure time equation from Lo Alg Document 10.1.1.4
         4 * l1b_de["avg_spin_durations"].to_numpy() / 3600,
         statistic="mean",
-        bins=[energy_edges, lon_edges, lat_edges],
+        # NOTE: The l1b pointing_bin_lon is bin number, not actual angle
+        bins=[
+            np.arange(N_ESA_ENERGY_STEPS + 1),
+            np.arange(N_SPIN_ANGLE_BINS + 1),
+            np.arange(N_OFF_ANGLE_BINS + 1),
+        ],
     )
 
     stat = result.statistic[np.newaxis, :, :, :]
 
     exposure_time = xr.DataArray(
         data=stat.astype(np.float16),
-        dims=["epoch", "energy", "longitude", "latitude"],
+        dims=PSET_DIMS,
     )
 
     return exposure_time
@@ -328,8 +385,6 @@ def create_datasets(
     #  can be used direction
     epoch_converted_time = [1]
 
-    # Create a data array for the epoch time
-    # TODO: might need to update the attrs to use new YAML file
     epoch_time = xr.DataArray(
         data=epoch_converted_time,
         name="epoch",
@@ -338,38 +393,54 @@ def create_datasets(
     )
 
     if logical_source == "imap_lo_l1c_pset":
-        esa_step = xr.DataArray(
-            data=[1, 2, 3, 4, 5, 6, 7],
-            name="esa_step",
-            dims=["esa_step"],
-            attrs=attr_mgr.get_variable_attributes("esa_step"),
-        )
-        pointing_bins = xr.DataArray(
-            data=np.arange(3600),
-            name="pointing_bins",
-            dims=["pointing_bins"],
-            attrs=attr_mgr.get_variable_attributes("pointing_bins"),
+        esa_energy_step = xr.DataArray(
+            data=ESA_ENERGY_STEPS,
+            name="esa_energy_step",
+            dims=["esa_energy_step"],
+            attrs=attr_mgr.get_variable_attributes("esa_energy_step"),
         )
-
-        esa_step_label = xr.DataArray(
-            esa_step.values.astype(str),
+        esa_energy_step_label = xr.DataArray(
+            esa_energy_step.values.astype(str),
             name="esa_step_label",
             dims=["esa_step_label"],
             attrs=attr_mgr.get_variable_attributes("esa_step_label"),
         )
-        pointing_bins_label = xr.DataArray(
-            pointing_bins.values.astype(str),
-            name="pointing_bins_label",
-            dims=["pointing_bins_label"],
-            attrs=attr_mgr.get_variable_attributes("pointing_bins_label"),
+
+        spin_angle = xr.DataArray(
+            data=SPIN_ANGLE_BIN_CENTERS,
+            name="spin_angle",
+            dims=["spin_angle"],
+            attrs=attr_mgr.get_variable_attributes("spin_angle"),
+        )
+        spin_angle_label = xr.DataArray(
+            spin_angle.values.astype(str),
+            name="spin_angle_label",
+            dims=["spin_angle_label"],
+            attrs=attr_mgr.get_variable_attributes("spin_angle_label"),
+        )
+
+        off_angle = xr.DataArray(
+            data=OFF_ANGLE_BIN_CENTERS,
+            name="off_angle",
+            dims=["off_angle"],
+            attrs=attr_mgr.get_variable_attributes("off_angle"),
+        )
+        off_angle_label = xr.DataArray(
+            off_angle.values.astype(str),
+            name="off_angle_label",
+            dims=["off_angle_label"],
+            attrs=attr_mgr.get_variable_attributes("off_angle_label"),
         )
+
         dataset = xr.Dataset(
             coords={
                 "epoch": epoch_time,
-                "pointing_bins": pointing_bins,
-                "pointing_bins_label": pointing_bins_label,
-                "esa_step": esa_step,
-                "esa_step_label": esa_step_label,
+                "esa_energy_step": esa_energy_step,
+                "esa_energy_step_label": esa_energy_step_label,
+                "spin_angle": spin_angle,
+                "spin_angle_label": spin_angle_label,
+                "off_angle": off_angle,
+                "off_angle_label": off_angle_label,
             },
             attrs=attr_mgr.get_global_attributes(logical_source),
         )
@@ -389,32 +460,141 @@ def create_datasets(
 
         # Create a data array for the current field and add it to the dataset
         # TODO: TEMPORARY. need to update to use l1b data once that's available.
-        if field in ["pointing_start", "pointing_end", "mode", "pivot_angle"]:
+        if field in [
+            "pointing_start_met",
+            "pointing_end_met",
+            "esa_mode",
+            "pivot_angle",
+        ]:
             dataset[field] = xr.DataArray(
                 data=[1],
                 dims=dims,
                 attrs=attr_mgr.get_variable_attributes(field),
             )
         # TODO: This is temporary.
-        #  The data type will be set in the data class when that's created
         elif field == "exposure_time":
             dataset[field] = xr.DataArray(
-                data=np.ones((1, 7), dtype=np.float16),
+                data=np.ones((1, 7, 3600, 40), dtype=np.float16),
                 dims=dims,
                 attrs=attr_mgr.get_variable_attributes(field),
             )
 
-        elif "rate" in field:
+        elif "rates" in field:
             dataset[field] = xr.DataArray(
-                data=np.ones((1, 3600, 7), dtype=np.float16),
+                data=np.ones(PSET_SHAPE, dtype=np.float16),
                 dims=dims,
                 attrs=attr_mgr.get_variable_attributes(field),
             )
         else:
             dataset[field] = xr.DataArray(
-                data=np.ones((1, 3600, 7), dtype=np.int16),
+                data=np.ones(PSET_SHAPE, dtype=np.int16),
                 dims=dims,
                 attrs=attr_mgr.get_variable_attributes(field),
             )
 
     return dataset
+
+
+def set_background_rates(
+    pointing_start_met: float,
+    pointing_end_met: float,
+    species: FilterType,
+    anc_dependencies: list,
+    attr_mgr: ImapCdfAttributes,
+) -> tuple[xr.DataArray, xr.DataArray, xr.DataArray]:
+    """
+    Set the background rates for the specified species.
+
+    The background rates are set to a constant value of 0.01 counts/s for all bins.
+
+    Parameters
+    ----------
+    pointing_start_met : float
+        The start MET time of the pointing.
+    pointing_end_met : float
+        The end MET time of the pointing.
+    species : FilterType
+        The species to set the background rates for. Can be "h" or "o".
+    anc_dependencies : list
+        Ancillary files needed for L1C data product creation.
+    attr_mgr : ImapCdfAttributes
+        Attribute manager used to get the L1C attributes.
+
+    Returns
+    -------
+    background_rates : tuple[xr.DataArray, xr.DataArray, xr.DataArray]
+        Tuple containing:
+        - The background rates for the specified species.
+        - The statistical uncertainties for the background rates.
+        - The systematic errors for the background rates.
+    """
+    if species not in {FilterType.HYDROGEN, FilterType.OXYGEN}:
+        raise ValueError(f"Species must be 'h' or 'o', but got {species.value}.")
+
+    bg_rates = np.zeros(
+        (N_ESA_ENERGY_STEPS, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS), dtype=np.float16
+    )
+    bg_stat_uncert = np.zeros(
+        (N_ESA_ENERGY_STEPS, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS), dtype=np.float16
+    )
+    bg_sys_err = np.zeros(
+        (N_ESA_ENERGY_STEPS, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS), dtype=np.float16
+    )
+
+    # read in the background rates from ancillary file
+    if species == FilterType.HYDROGEN:
+        background_df = lo_ancillary.read_ancillary_file(
+            next(str(s) for s in anc_dependencies if "hydrogen-background" in str(s))
+        )
+    else:
+        background_df = lo_ancillary.read_ancillary_file(
+            next(str(s) for s in anc_dependencies if "oxygen-background" in str(s))
+        )
+
+    # find to the rows for the current pointing
+    pointing_bg_df = background_df[
+        (background_df["GoodTime_strt"] >= pointing_start_met)
+        & (background_df["GoodTime_end"] <= pointing_end_met)
+    ]
+
+    # convert the bin start and end resolution from 6 degrees to .1 degrees
+    pointing_bg_df["bin_strt"] = pointing_bg_df["bin_strt"] * 60
+    # The last bin end in the file is 0, which means 60 degrees. This is
+    # converted to 0.1 degree resolution of 3600
+    pointing_bg_df["bin_end"] = pointing_bg_df["bin_end"] * 60
+    pointing_bg_df.loc[pointing_bg_df["bin_end"] == 0, "bin_end"] = 3600
+    # for each row in the bg ancillary file for this pointing
+    for _, row in pointing_bg_df.iterrows():
+        bin_start = int(row["bin_strt"])
+        bin_end = int(row["bin_end"])
+        # for each energy step, set the background rate and uncertainty
+        for esa_step in range(0, 7):
+            value = row[f"E-Step{esa_step + 1}"]
+            if row["type"] == "rate":
+                bg_rates[esa_step, bin_start:bin_end, :] = value
+            elif row["type"] == "sigma":
+                bg_stat_uncert[esa_step, bin_start:bin_end, :] = value
+            else:
+                raise ValueError("Unknown background type in ancillary file.")
+    # set the background rates, uncertainties, and systematic errors
+    bg_rates_data = xr.DataArray(
+        data=bg_rates,
+        dims=["esa_energy_step", "spin_angle", "off_angle"],
+        attrs=attr_mgr.get_variable_attributes(f"{species.value}_background_rates"),
+    )
+    bg_stat_uncert_data = xr.DataArray(
+        data=bg_stat_uncert,
+        dims=["esa_energy_step", "spin_angle", "off_angle"],
+        attrs=attr_mgr.get_variable_attributes(
+            f"{species.value}_background_rates_stat_uncert"
+        ),
+    )
+    bg_sys_err_data = xr.DataArray(
+        data=bg_sys_err,
+        dims=["esa_energy_step", "spin_angle", "off_angle"],
+        attrs=attr_mgr.get_variable_attributes(
+            f"{species.value}_background_rates_sys_err"
+        ),
+    )
+
+    return bg_rates_data, bg_stat_uncert_data, bg_sys_err_data