imap-processing 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

imap_processing/codice/codice_l1a_lo_species.py

@@ -0,0 +1,282 @@
+"""CoDICE Lo Species L1A processing functions."""
+
+import logging
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+
+from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.codice import constants
+from imap_processing.codice.decompress import decompress
+from imap_processing.codice.utils import (
+    CODICEAPID,
+    ViewTabInfo,
+    calculate_acq_time_per_step,
+    get_codice_epoch_time,
+    get_collapse_pattern_shape,
+    get_view_tab_info,
+    read_sci_lut,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def l1a_lo_species(unpacked_dataset: xr.Dataset, lut_file: Path) -> xr.Dataset:
+    """
+    L1A processing code.
+
+    Parameters
+    ----------
+    unpacked_dataset : xarray.Dataset
+        The decompressed and unpacked data from the packet file.
+    lut_file : pathlib.Path
+        Path to the LUT (Lookup Table) file used for processing.
+
+    Returns
+    -------
+    xarray.Dataset
+        The processed L1A dataset for the given species product.
+    """
+    # Get these values from unpacked data. These are used to
+    # lookup in LUT table.
+    table_id = unpacked_dataset["table_id"].values[0]
+    view_id = unpacked_dataset["view_id"].values[0]
+    apid = unpacked_dataset["pkt_apid"].values[0]
+    plan_id = unpacked_dataset["plan_id"].values[0]
+    plan_step = unpacked_dataset["plan_step"].values[0]
+
+    logger.info(
+        f"Processing species with - APID: {apid}, View ID: {view_id}, "
+        f"Table ID: {table_id}, Plan ID: {plan_id}, Plan Step: {plan_step}"
+    )
+
+    # ========== Get LUT Data ===========
+    # Read information from LUT
+    sci_lut_data = read_sci_lut(lut_file, table_id)
+
+    view_tab_info = get_view_tab_info(sci_lut_data, view_id, apid)
+    view_tab_obj = ViewTabInfo(
+        apid=apid,
+        view_id=view_id,
+        sensor=view_tab_info["sensor"],
+        three_d_collapsed=view_tab_info["3d_collapse"],
+        collapse_table=view_tab_info["collapse_table"],
+    )
+
+    if view_tab_obj.sensor != 0:
+        raise ValueError("Unsupported sensor ID for Lo species processing.")
+
+    # ========= Decompress and Reshape Data ===========
+    # Lookup SW or NSW species based on APID
+    if view_tab_obj.apid == CODICEAPID.COD_LO_SW_SPECIES_COUNTS:
+        species_names = sci_lut_data["data_product_lo_tab"]["0"]["species"]["sw"][
+            "species_names"
+        ]
+        logical_source_id = "imap_codice_l1a_lo-sw-species"
+    elif view_tab_obj.apid == CODICEAPID.COD_LO_NSW_SPECIES_COUNTS:
+        species_names = sci_lut_data["data_product_lo_tab"]["0"]["species"]["nsw"][
+            "species_names"
+        ]
+        logical_source_id = "imap_codice_l1a_lo-nsw-species"
+    else:
+        raise ValueError(f"Unknown apid {view_tab_obj.apid} in Lo species processing.")
+
+    compression_algorithm = constants.LO_COMPRESSION_ID_LOOKUP[view_tab_obj.view_id]
+    # Decompress data using byte count information from decommed data
+    binary_data_list = unpacked_dataset["data"].values
+    byte_count_list = unpacked_dataset["byte_count"].values
+
+    # The decompressed data in the shape of (epoch, n). Then reshape later.
+    decompressed_data = [
+        decompress(
+            packet_data[:byte_count],
+            compression_algorithm,
+        )
+        for (packet_data, byte_count) in zip(
+            binary_data_list, byte_count_list, strict=False
+        )
+    ]
+
+    # Look up collapse pattern using LUT table. This should return collapsed shape.
+    # For Lo species, it will be (1,)
+    collapsed_shape = get_collapse_pattern_shape(
+        sci_lut_data, view_tab_obj.sensor, view_tab_obj.collapse_table
+    )
+
+    # Reshape decompressed data to:
+    #   (num_packets, num_species, esa_steps, *collapsed_shape)
+    # where collapsed_shape is usually (1,) for Lo species.
+    num_packets = len(binary_data_list)
+    num_species = len(species_names)
+    esa_steps = constants.NUM_ESA_STEPS
+    species_data = np.array(decompressed_data).reshape(
+        num_packets, num_species, esa_steps, *collapsed_shape
+    )
+
+    # ========== Get Voltage Data from LUT ===========
+    # Use plan id and plan step to get voltage data's table_number in ESA sweep table.
+    # Voltage data is (128,)
+    esa_table_number = sci_lut_data["plan_tab"][f"({plan_id}, {plan_step})"][
+        "lo_stepping"
+    ]
+    voltage_data = sci_lut_data["esa_sweep_tab"][f"{esa_table_number}"]
+
+    # ========= Get Epoch Time Data ===========
+    # Epoch center time and delta
+    epoch_center, deltas = get_codice_epoch_time(
+        unpacked_dataset["acq_start_seconds"].values,
+        unpacked_dataset["acq_start_subseconds"].values,
+        unpacked_dataset["spin_period"].values,
+        view_tab_obj,
+    )
+
+    # ========== Create CDF Dataset with Metadata ===========
+    cdf_attrs = ImapCdfAttributes()
+    cdf_attrs.add_instrument_global_attrs("codice")
+    cdf_attrs.add_instrument_variable_attrs("codice", "l1a")
+
+    l1a_dataset = xr.Dataset(
+        coords={
+            "epoch": xr.DataArray(
+                epoch_center,
+                dims=("epoch",),
+                attrs=cdf_attrs.get_variable_attributes("epoch", check_schema=False),
+            ),
+            "epoch_delta_minus": xr.DataArray(
+                deltas,
+                dims=("epoch",),
+                attrs=cdf_attrs.get_variable_attributes(
+                    "epoch_delta_minus", check_schema=False
+                ),
+            ),
+            "epoch_delta_plus": xr.DataArray(
+                deltas,
+                dims=("epoch",),
+                attrs=cdf_attrs.get_variable_attributes(
+                    "epoch_delta_plus", check_schema=False
+                ),
+            ),
+            "esa_step": xr.DataArray(
+                np.arange(128),
+                dims=("esa_step",),
+                attrs=cdf_attrs.get_variable_attributes("esa_step", check_schema=False),
+            ),
+            "esa_step_label": xr.DataArray(
+                np.arange(128).astype(str),
+                dims=("esa_step",),
+                attrs=cdf_attrs.get_variable_attributes(
+                    "esa_step_label", check_schema=False
+                ),
+            ),
+            "k_factor": xr.DataArray(
+                np.array([constants.K_FACTOR]),
+                dims=("k_factor",),
+                attrs=cdf_attrs.get_variable_attributes(
+                    "k_factor_attrs", check_schema=False
+                ),
+            ),
+            "spin_sector": xr.DataArray(
+                np.array([0], dtype=np.uint8),
+                dims=("spin_sector",),
+                attrs=cdf_attrs.get_variable_attributes(
+                    "spin_sector", check_schema=False
+                ),
+            ),
+            "spin_sector_label": xr.DataArray(
+                np.array(["0"]).astype(str),
+                dims=("spin_sector",),
+                attrs=cdf_attrs.get_variable_attributes(
+                    "spin_sector_label", check_schema=False
+                ),
+            ),
+        },
+        attrs=cdf_attrs.get_global_attributes(logical_source_id),
+    )
+    # Add first few unique variables
+    l1a_dataset["spin_period"] = xr.DataArray(
+        unpacked_dataset["spin_period"].values * constants.SPIN_PERIOD_CONVERSION,
+        dims=("epoch",),
+        attrs=cdf_attrs.get_variable_attributes("spin_period"),
+    )
+    l1a_dataset["k_factor"] = xr.DataArray(
+        np.array([constants.K_FACTOR]),
+        dims=("k_factor",),
+        attrs=cdf_attrs.get_variable_attributes("k_factor_attrs", check_schema=False),
+    )
+    l1a_dataset["voltage_table"] = xr.DataArray(
+        np.array(voltage_data),
+        dims=("esa_step",),
+        attrs=cdf_attrs.get_variable_attributes("voltage_table", check_schema=False),
+    )
+    l1a_dataset["data_quality"] = xr.DataArray(
+        unpacked_dataset["suspect"].values,
+        dims=("epoch",),
+        attrs=cdf_attrs.get_variable_attributes("data_quality"),
+    )
+    l1a_dataset["acquisition_time_per_step"] = xr.DataArray(
+        calculate_acq_time_per_step(sci_lut_data["lo_stepping_tab"]),
+        dims=("esa_step",),
+        attrs=cdf_attrs.get_variable_attributes(
+            "acquisition_time_per_step", check_schema=False
+        ),
+    )
+
+    # Carry over these variables from unpacked data to l1a_dataset
+    l1a_carryover_vars = [
+        "sw_bias_gain_mode",
+        "st_bias_gain_mode",
+        "rgfo_half_spin",
+        "nso_half_spin",
+    ]
+    # Loop through them since we need to set their attrs too
+    for var in l1a_carryover_vars:
+        l1a_dataset[var] = xr.DataArray(
+            unpacked_dataset[var].values,
+            dims=("epoch",),
+            attrs=cdf_attrs.get_variable_attributes(var),
+        )
+
+    # Finally, add species data variables and their uncertainties
+    for idx, species in enumerate(species_names):
+        if view_tab_obj.apid == CODICEAPID.COD_LO_SW_SPECIES_COUNTS and species in [
+            "heplus",
+            "cnoplus",
+        ]:
+            species_attrs = cdf_attrs.get_variable_attributes("lo-pui-species-attrs")
+            unc_attrs = cdf_attrs.get_variable_attributes("lo-pui-species-unc-attrs")
+        else:
+            species_attrs = cdf_attrs.get_variable_attributes("lo-species-attrs")
+            unc_attrs = cdf_attrs.get_variable_attributes("lo-species-unc-attrs")
+
+        direction = (
+            "Sunward"
+            if view_tab_obj.apid == CODICEAPID.COD_LO_SW_SPECIES_COUNTS
+            else "Non-Sunward"
+        )
+        # Replace {species} and {direction} in attrs
+        species_attrs["CATDESC"] = species_attrs["CATDESC"].format(
+            species=species, direction=direction
+        )
+        species_attrs["FIELDNAM"] = species_attrs["FIELDNAM"].format(
+            species=species, direction=direction
+        )
+        l1a_dataset[species] = xr.DataArray(
+            species_data[:, idx, :, :],
+            dims=("epoch", "esa_step", "spin_sector"),
+            attrs=species_attrs,
+        )
+        # Uncertainty data
+        unc_attrs["CATDESC"] = unc_attrs["CATDESC"].format(
+            species=species, direction=direction
+        )
+        unc_attrs["FIELDNAM"] = unc_attrs["FIELDNAM"].format(
+            species=species, direction=direction
+        )
+        l1a_dataset[f"unc_{species}"] = xr.DataArray(
+            np.sqrt(species_data[:, idx, :, :]),
+            dims=("epoch", "esa_step", "spin_sector"),
+            attrs=unc_attrs,
+        )
+
+    return l1a_dataset

imap_processing/codice/codice_l1b.py

@@ -15,20 +15,15 @@ from pathlib import Path
 import numpy as np
 import xarray as xr
 
-from imap_processing import imap_module_directory
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.cdf.utils import load_cdf
 from imap_processing.codice import constants
-from imap_processing.codice.utils import CODICEAPID
-from imap_processing.utils import packet_file_to_datasets
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 
-def convert_to_rates(
-    dataset: xr.Dataset, descriptor: str, variable_name: str
-) -> np.ndarray:
+def convert_to_rates(dataset: xr.Dataset, descriptor: str) -> np.ndarray:
     """
     Apply a conversion from counts to rates.
 
@@ -41,14 +36,17 @@ def convert_to_rates(
         The L1b dataset containing the data to convert.
     descriptor : str
         The descriptor of the data product of interest.
-    variable_name : str
-        The variable name to apply the conversion to.
 
     Returns
     -------
     rates_data : np.ndarray
         The converted data array.
     """
+    # Variables to convert based on descriptor
+    variables_to_convert = getattr(
+        constants, f"{descriptor.upper().replace('-', '_')}_VARIABLE_NAMES"
+    )
+
     if descriptor in [
         "lo-counters-aggregated",
         "lo-counters-singles",
@@ -58,41 +56,49 @@ def convert_to_rates(
         "lo-sw-priority",
         "lo-ialirt",
     ]:
-        # Applying rate calculation described in section 10.2 of the algorithm
-        # document
-        # In order to divide by acquisition times, we must reshape the acq
-        # time data array to match the data variable shape
-        dims = [1] * dataset[variable_name].data.ndim
-        dims[1] = 128
-        acq_times = dataset.acquisition_time_per_step.data.reshape(dims)  # (128)
-        # Now perform the calculation
-        rates_data = dataset[variable_name].data / (
-            acq_times
-            * 1e-3  # Converting from milliseconds to seconds
+        # Denominator to convert counts to rates
+        denominator = (
+            dataset.acquisition_time_per_step
             * constants.L1B_DATA_PRODUCT_CONFIGURATIONS[descriptor]["num_spin_sectors"]
         )
+
+        # Do not carry these variable attributes from L1a to L1b for above products
+        drop_variables = [
+            "k_factor",
+            "nso_half_spin",
+            "sw_bias_gain_mode",
+            "st_bias_gain_mode",
+            "spin_period",
+        ]
+        dataset = dataset.drop_vars(drop_variables)
     elif descriptor in [
         "lo-nsw-species",
         "lo-sw-species",
     ]:
-        # Applying rate calculation described in section 10.2 of the algorithm
-        # document
-        # In order to divide by acquisition times, we must reshape the acq
-        # time data array to match the data variable shape (epoch, esa_step, sector)
-        dims = [1] * dataset[variable_name].data.ndim
-        dims[1] = 128
-        acq_times = dataset.acquisition_time_per_step.data.reshape(dims)  # (128)
-        # acquisition time have an array of shape (128,). We match n_sector to that.
+        # Create n_sector with 'esa_step' dimension. This is done by xr.full_like
+        # with input dataset.acquisition_time_per_step. This ensures that the resulting
+        # n_sector has the same dimensions as acquisition_time_per_step.
         # Per CoDICE, fill first 127 with default value of 12. Then fill last with 11.
-        n_sector = np.full(128, 12, dtype=int)
-        n_sector[-1] = 11
-
-        # Now perform the calculation
-        rates_data = dataset[variable_name].data / (
-            acq_times
-            * 1e-3  # Converting from milliseconds to seconds
-            * n_sector[:, np.newaxis]  # Spin sectors
+        n_sector = xr.full_like(
+            dataset.acquisition_time_per_step, 12.0, dtype=np.float64
         )
+        n_sector[-1] = 11.0
+
+        # Denominator to convert counts to rates
+        denominator = dataset.acquisition_time_per_step * n_sector
+
+        # Do not carry these variable attributes from L1a to L1b for above products
+        drop_variables = [
+            "k_factor",
+            "nso_half_spin",
+            "sw_bias_gain_mode",
+            "st_bias_gain_mode",
+            "spin_period",
+            "voltage_table",
+            "acquisition_time_per_step",
+        ]
+        dataset = dataset.drop_vars(drop_variables)
+
     elif descriptor in [
         "hi-counters-aggregated",
         "hi-counters-singles",
@@ -101,15 +107,27 @@ def convert_to_rates(
         "hi-sectored",
         "hi-ialirt",
     ]:
-        # Applying rate calculation described in section 10.1 of the algorithm
-        # document
-        rates_data = dataset[variable_name].data / (
+        # Denominator to convert counts to rates
+        denominator = (
             constants.L1B_DATA_PRODUCT_CONFIGURATIONS[descriptor]["num_spin_sectors"]
             * constants.L1B_DATA_PRODUCT_CONFIGURATIONS[descriptor]["num_spins"]
             * constants.HI_ACQUISITION_TIME
         )
 
-    return rates_data
+    # For each variable, convert counts and uncertainty to rates
+    for variable in variables_to_convert:
+        dataset[variable].data = dataset[variable].astype(np.float64) / denominator
+        # Carry over attrs and update as needed
+        dataset[variable].attrs["UNITS"] = "counts/s"
+
+        # Uncertainty calculation
+        unc_variable = f"unc_{variable}"
+        dataset[unc_variable].data = (
+            dataset[unc_variable].astype(np.float64) / denominator
+        )
+        dataset[unc_variable].attrs["UNITS"] = "1/s"
+
+    return dataset
 
 
 def process_codice_l1b(file_path: Path) -> xr.Dataset:
@@ -136,70 +154,17 @@ def process_codice_l1b(file_path: Path) -> xr.Dataset:
     dataset_name = l1a_dataset.attrs["Logical_source"].replace("_l1a_", "_l1b_")
     descriptor = dataset_name.removeprefix("imap_codice_l1b_")
 
-    # Direct event data products do not have a level L1B
-    if descriptor in ["lo-direct-events", "hi-direct-events"]:
-        logger.warning("Encountered direct event data product. Skipping L1b processing")
-        return None
-
     # Get the L1b CDF attributes
     cdf_attrs = ImapCdfAttributes()
     cdf_attrs.add_instrument_global_attrs("codice")
     cdf_attrs.add_instrument_variable_attrs("codice", "l1b")
 
     # Use the L1a data product as a starting point for L1b
-    l1b_dataset = l1a_dataset.copy()
+    l1b_dataset = l1a_dataset.copy(deep=True)
 
     # Update the global attributes
     l1b_dataset.attrs = cdf_attrs.get_global_attributes(dataset_name)
-
-    # TODO: This was thrown together quickly and should be double-checked
-    if descriptor == "hskp":
-        xtce_filename = "codice_packet_definition.xml"
-        xtce_packet_definition = Path(
-            f"{imap_module_directory}/codice/packet_definitions/{xtce_filename}"
-        )
-        packet_file = (
-            imap_module_directory
-            / "tests"
-            / "codice"
-            / "data"
-            / "imap_codice_l0_raw_20241110_v001.pkts"
-        )
-        datasets: dict[int, xr.Dataset] = packet_file_to_datasets(
-            packet_file, xtce_packet_definition, use_derived_value=True
-        )
-        l1b_dataset = datasets[CODICEAPID.COD_NHK]
-
-        # TODO: Drop the same variables as we do in L1a? (see line 1103 in
-        #       codice_l1a.py
-
-    else:
-        variables_to_convert = getattr(
-            constants, f"{descriptor.upper().replace('-', '_')}_VARIABLE_NAMES"
-        )
-
-        # Apply the conversion to rates
-        for variable_name in variables_to_convert:
-            l1b_dataset[variable_name].data = convert_to_rates(
-                l1b_dataset, descriptor, variable_name
-            )
-            # Set the variable attributes
-            cdf_attrs_key = f"{descriptor}-{variable_name}"
-            l1b_dataset[variable_name].attrs = cdf_attrs.get_variable_attributes(
-                cdf_attrs_key, check_schema=False
-            )
-
-        if descriptor in ["lo-sw-species", "lo-nsw-species"]:
-            # Do not carry these variable attributes from L1a to L1b
-            drop_variables = [
-                "k_factor",
-                "nso_half_spin",
-                "sw_bias_gain_mode",
-                "st_bias_gain_mode",
-                "spin_period",
-            ]
-            l1b_dataset = l1b_dataset.drop_vars(drop_variables)
-
-    logger.info(f"\nFinal data product:\n{l1b_dataset}\n")
-
-    return l1b_dataset
+    return convert_to_rates(
+        l1b_dataset,
+        descriptor,
+    )