imap-processing 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

imap_processing/codice/codice_l2.py

@@ -26,11 +26,16 @@ from imap_processing.codice.constants import (
     L2_GEOMETRIC_FACTOR,
     L2_HI_NUMBER_OF_SSD,
     L2_HI_SECTORED_ANGLE,
+    LO_NSW_ANGULAR_VARIABLE_NAMES,
     LO_NSW_SPECIES_VARIABLE_NAMES,
+    LO_POSITION_TO_ELEVATION_ANGLE,
+    LO_SW_ANGULAR_VARIABLE_NAMES,
     LO_SW_PICKUP_ION_SPECIES_VARIABLE_NAMES,
-    LO_SW_SPECIES_VARIABLE_NAMES,
+    LO_SW_SOLAR_WIND_SPECIES_VARIABLE_NAMES,
     NSW_POSITIONS,
+    PIXEL_ORIENTATIONS,
     PUI_POSITIONS,
+    SOLAR_WIND_POSITIONS,
     SW_POSITIONS,
 )
 
@@ -94,7 +99,7 @@ def get_efficiency_lut(dependencies: ProcessingInputCollection) -> pd.DataFrame:
     return pd.read_csv(dependencies.get_file_paths(descriptor="l2-lo-efficiency")[0])
 
 
-def get_species_efficiency(species: str, efficiency: pd.DataFrame) -> np.ndarray:
+def get_species_efficiency(species: str, efficiency: pd.DataFrame) -> xr.DataArray:
     """
     Get the efficiency values for a given species.
 
@@ -107,7 +112,7 @@ def get_species_efficiency(species: str, efficiency: pd.DataFrame) -> np.ndarray
 
     Returns
     -------
-    efficiency : np.ndarray
+    efficiency : xarray.DataArray
         A 2D array of efficiencies with shape (epoch, esa_steps).
     """
     species_efficiency = efficiency[efficiency["species"] == species].sort_values(
@@ -118,13 +123,16 @@ def get_species_efficiency(species: str, efficiency: pd.DataFrame) -> np.ndarray
         [col for col in species_efficiency if col.startswith("position")],
         key=lambda x: int(x.split("_")[-1]),
     )
-    # Shape: (esa_steps, positions)
-    return species_efficiency[position_names_sorted].to_numpy()
+    # Shape: (energy_table, inst_az)
+    return xr.DataArray(
+        species_efficiency[position_names_sorted].to_numpy(),
+        dims=("energy_table", "inst_az"),
+    )
 
 
 def compute_geometric_factors(
     dataset: xr.Dataset, geometric_factor_lookup: dict
-) -> np.ndarray:
+) -> xr.DataArray:
     """
     Calculate geometric factors needed for intensity calculations.
 
@@ -148,7 +156,7 @@ def compute_geometric_factors(
 
     Returns
     -------
-    geometric_factors : np.ndarray
+    geometric_factors : xarray.DataArray
         A 3D array of geometric factors with shape (epoch, esa_steps, positions).
     """
     # Convert the HALF_SPIN_LUT to a reverse mapping of esa_step to half_spin
@@ -170,22 +178,26 @@ def compute_geometric_factors(
 
     # Get the geometric factors based on the modes
     gf = np.where(
-        modes[:, :, np.newaxis],  # Shape (epoch, esa_step, 1)
-        geometric_factor_lookup["reduced"],  # Shape (1, esa_step, 24) - reduced mode
-        geometric_factor_lookup["full"],  # Shape (1, esa_step, 24) - full mode
-    )  # Shape: (epoch, esa_step, positions)
-    return gf
+        modes[:, :, np.newaxis],  # Shape (epoch, energy_table, 1)
+        geometric_factor_lookup[
+            "reduced"
+        ],  # Shape (1, energy_table, 24) - reduced mode
+        geometric_factor_lookup["full"],  # Shape (1, energy_table, 24) - full mode
+    )  # Shape: (epoch, energy_table, inst_az)
 
+    return xr.DataArray(gf, dims=("epoch", "energy_table", "inst_az"))
 
-def process_lo_species_intensity(
+
+def calculate_intensity(
     dataset: xr.Dataset,
     species_list: list,
-    geometric_factors: np.ndarray,
+    geometric_factors: xr.DataArray,
     efficiency: pd.DataFrame,
     positions: list,
+    average_across_positions: bool = False,
 ) -> xr.Dataset:
     """
-    Process the lo-species L2 dataset to calculate species intensities.
+    Calculate species or angular intensities.
 
     Parameters
     ----------
@@ -200,6 +212,9 @@ def process_lo_species_intensity(
     positions : list
         A list of position indices to select from the geometric factor and
         efficiency lookup tables.
+    average_across_positions : bool
+        Whether to average the efficiencies and geometric factors across the selected
+        positions. Default is False.
 
     Returns
     -------
@@ -207,33 +222,207 @@ def process_lo_species_intensity(
         The updated L2 dataset with species intensities calculated.
     """
     # Select the relevant positions from the geometric factors
-    geometric_factors = geometric_factors[:, :, positions]
-    # take the mean geometric factor across positions
-    geometric_factors = np.nanmean(geometric_factors, axis=-1)
-    scaler = len(positions)
-    # Calculate the species intensities using the provided geometric factors and
-    # efficiency. Species_intensity = species_rate / (gm * eff * esa_step)
+    geometric_factors = geometric_factors.isel(inst_az=positions)
+    if average_across_positions:
+        # take the mean geometric factor across positions
+        geometric_factors = geometric_factors.mean(dim="inst_az")
+        scalar = len(positions)
+    else:
+        scalar = 1
+    # Calculate the angular intensities using the provided geometric factors and
+    # efficiency.
+    # intensity = species_rate / (gm * eff * esa_step) for position and spin angle
     for species in species_list:
         # Select the relevant positions for the species from the efficiency LUT
-        # Shape: (epoch, esa_steps, positions)
-        species_eff = get_species_efficiency(species, efficiency)[
-            np.newaxis, :, positions
-        ]
+        # Shape: (epoch, energy_table, inst_az)
+        species_eff = get_species_efficiency(species, efficiency).isel(
+            inst_az=positions
+        )
         if species_eff.size == 0:
-            logger.warning("No efficiency data found for species {species}. Skipping.")
+            logger.warning(f"No efficiency data found for species {species}. Skipping.")
             continue
-        # Take the mean efficiency across positions
-        species_eff = np.nanmean(species_eff, axis=-1)
-        denominator = (
-            scaler * geometric_factors * species_eff * dataset["energy_table"].data
-        )
+
+        if average_across_positions:
+            # Take the mean efficiency across positions
+            species_eff = species_eff.mean(dim="inst_az")
+
+        # Shape: (epoch, energy_table, inst_az) or
+        # (epoch, energy_table) if averaged
+        denominator = scalar * geometric_factors * species_eff * dataset["energy_table"]
         if species not in dataset:
             logger.warning(
                 f"Species {species} not found in dataset. Filling with NaNS."
             )
             dataset[species] = np.full(dataset["energy_table"].data.shape, np.nan)
         else:
-            dataset[species] = dataset[species] / denominator[:, :, np.newaxis]
+            dataset[species] = dataset[species] / denominator
+
+        # Also calculate uncertainty if available
+        species_uncertainty = f"unc_{species}"
+        if species_uncertainty not in dataset:
+            logger.warning(
+                f"Uncertainty {species_uncertainty} not found in dataset."
+                f" Filling with NaNS."
+            )
+            dataset[species_uncertainty] = np.full(
+                dataset["energy_table"].data.shape, np.nan
+            )
+        else:
+            dataset[species_uncertainty] = dataset[species_uncertainty] / denominator
+
+    return dataset
+
+
+def process_lo_species_intensity(
+    dataset: xr.Dataset,
+    species_list: list,
+    geometric_factors: xr.DataArray,
+    efficiency: pd.DataFrame,
+    positions: list,
+) -> xr.Dataset:
+    """
+    Process the lo-species L2 dataset to calculate species intensities.
+
+    Parameters
+    ----------
+    dataset : xarray.Dataset
+        The L2 dataset to process.
+    species_list : list
+        List of species variable names to calculate intensity.
+    geometric_factors : xarray.DataArray
+        The geometric factors array with shape (epoch, esa_steps).
+    efficiency : pandas.DataFrame
+        The efficiency lookup table.
+    positions : list
+        A list of position indices to select from the geometric factor and
+        efficiency lookup tables.
+
+    Returns
+    -------
+    xarray.Dataset
+        The updated L2 dataset with species intensities calculated.
+    """
+    # Calculate the species intensities using the provided geometric factors and
+    # efficiency.
+    dataset = calculate_intensity(
+        dataset,
+        species_list,
+        geometric_factors,
+        efficiency,
+        positions,
+        average_across_positions=True,
+    )
+
+    return dataset
+
+
+def process_lo_angular_intensity(
+    dataset: xr.Dataset,
+    species_list: list,
+    geometric_factors: xr.DataArray,
+    efficiency: pd.DataFrame,
+    positions: list,
+) -> xr.Dataset:
+    """
+    Process the lo-species L2 dataset to calculate angular intensities.
+
+    Parameters
+    ----------
+    dataset : xarray.Dataset
+        The L2 dataset to process.
+    species_list : list
+        List of species variable names to calculate intensity.
+    geometric_factors : xarray.DataArray
+        The geometric factors array with shape (epoch, esa_steps).
+    efficiency : pandas.DataFrame
+        The efficiency lookup table.
+    positions : list
+        A list of position indices to select from the geometric factor and
+        efficiency lookup tables.
+
+    Returns
+    -------
+    xarray.Dataset
+        The updated L2 dataset with angular intensities calculated.
+    """
+    # Calculate the angular intensities using the provided geometric factors and
+    # efficiency.
+    dataset = calculate_intensity(
+        dataset,
+        species_list,
+        geometric_factors,
+        efficiency,
+        positions,
+        average_across_positions=False,
+    )
+
+    # transform positions to elevation angles
+    if positions == SW_POSITIONS:
+        pos_to_el = LO_POSITION_TO_ELEVATION_ANGLE["sw"]
+        position_index_to_adjust = 0
+    elif positions == NSW_POSITIONS:
+        pos_to_el = LO_POSITION_TO_ELEVATION_ANGLE["nsw"]
+        position_index_to_adjust = 9
+    else:
+        raise ValueError("Unknown positions for elevation angle mapping.")
+
+    # Create a new coordinate for elevation_angle based on inst_az
+    dataset = dataset.assign_coords(
+        elevation_angle=(
+            "inst_az",
+            [pos_to_el[pos] for pos in dataset["inst_az"].data],
+        )
+    )
+    # add uncertainties to species list
+    species_list = species_list + [f"unc_{var}" for var in species_list]
+    # Take the mean across elevation angles and restore the original dimension order
+    dataset_converted = (
+        dataset[species_list]
+        .groupby("elevation_angle")
+        .sum(keep_attrs=True)  # One position should always contain zeros so sum is safe
+        # Restore original dimension order because groupby moves the grouped
+        # dimension to the front
+        .transpose("epoch", "energy_table", "spin_sector", "elevation_angle", ...)
+    )
+    # Create a new coordinate for spin angle based on spin_sector
+    # Use equation from section 11.2.2 of algorithm document
+    dataset = dataset.assign_coords(
+        spin_angle=("spin_sector", dataset["spin_sector"].data * 15.0 + 7.5)
+    )
+    dataset = dataset.drop_vars(species_list).merge(dataset_converted)
+    # Positions 0 and 10 only observe half of the 24 spins for each esa step.
+    # To account for this, we replicate the counts observed in position 0 and 10 for
+    # each esa step to either spin angles 0-11 or 12-23, depending on the pixel
+    # orientation (A/B). See section 11.2.2 of the CoDICE algorithm document
+    a_inds = np.array(
+        [pos for pos, orientation in PIXEL_ORIENTATIONS.items() if orientation == "A"]
+    )
+    b_inds = np.array(
+        [pos for pos, orientation in PIXEL_ORIENTATIONS.items() if orientation == "B"]
+    )
+
+    position_index = position_index_to_adjust
+    for species in species_list:
+        # Determine the correct spin indices based on the position
+        spin_sectors = dataset["spin_sector"].data
+        spin_inds_1 = np.where(spin_sectors >= 12)[0]
+        spin_inds_2 = np.where(spin_sectors < 12)[0]
+        # if position_index is 9, swap the spin indices
+        if position_index == 9:
+            spin_inds_1, spin_inds_2 = spin_inds_2, spin_inds_1
+
+        # Assign the values to the correct positions and spin sectors
+        dataset[species].values[
+            :, a_inds[:, np.newaxis], spin_inds_1, position_index
+        ] = dataset[species].values[
+            :, a_inds[:, np.newaxis], spin_inds_2, position_index
+        ]
+
+        dataset[species].values[
+            :, b_inds[:, np.newaxis], spin_inds_2, position_index
+        ] = dataset[species].values[
+            :, b_inds[:, np.newaxis], spin_inds_1, position_index
+        ]
 
     return dataset
 
@@ -627,7 +816,6 @@ def process_codice_l2(
     # This should get science files since ancillary or spice doesn't have data_type
     # as data level.
     file_path = dependencies.get_file_paths(descriptor=descriptor)[0]
-
     # Now form product name from descriptor
     descriptor = ScienceFilePath(file_path).descriptor
     dataset_name = f"imap_codice_l2_{descriptor}"
@@ -637,7 +825,12 @@ def process_codice_l2(
     if dataset_name in [
         "imap_codice_l2_lo-sw-species",
         "imap_codice_l2_lo-nsw-species",
+        "imap_codice_l2_lo-nsw-angular",
+        "imap_codice_l2_lo-sw-angular",
     ]:
+        cdf_attrs = ImapCdfAttributes()
+        cdf_attrs.add_instrument_global_attrs("codice")
+
         l2_dataset = load_cdf(file_path).copy()
 
         geometric_factor_lookup = get_geometric_factor_lut(dependencies)
@@ -645,12 +838,13 @@ def process_codice_l2(
         geometric_factors = compute_geometric_factors(
             l2_dataset, geometric_factor_lookup
         )
+
         if dataset_name == "imap_codice_l2_lo-sw-species":
             # Filter the efficiency lookup table for solar wind efficiencies
             efficiencies = efficiency_lookup[efficiency_lookup["product"] == "sw"]
             # Calculate the pickup ion sunward solar wind intensities using equation
-            # described in section 11.2.4 of algorithm document.
-            process_lo_species_intensity(
+            # described in section 11.2.3 of algorithm document.
+            l2_dataset = process_lo_species_intensity(
                 l2_dataset,
                 LO_SW_PICKUP_ION_SPECIES_VARIABLE_NAMES,
                 geometric_factors,
@@ -658,27 +852,59 @@ def process_codice_l2(
                 PUI_POSITIONS,
             )
             # Calculate the sunward solar wind species intensities using equation
-            # described in section 11.2.4 of algorithm document.
-            process_lo_species_intensity(
+            # described in section 11.2.3 of algorithm document.
+            l2_dataset = process_lo_species_intensity(
                 l2_dataset,
-                LO_SW_SPECIES_VARIABLE_NAMES,
+                LO_SW_SOLAR_WIND_SPECIES_VARIABLE_NAMES,
                 geometric_factors,
                 efficiencies,
-                SW_POSITIONS,
+                SOLAR_WIND_POSITIONS,
             )
-        else:
-            # Filter the efficiency lookup table for non solar wind efficiencies
+            l2_dataset.attrs.update(
+                cdf_attrs.get_global_attributes("imap_codice_l2_lo-sw-species")
+            )
+        elif dataset_name == "imap_codice_l2_lo-nsw-species":
+            # Filter the efficiency lookup table for non-solar wind efficiencies
             efficiencies = efficiency_lookup[efficiency_lookup["product"] == "nsw"]
             # Calculate the non-sunward species intensities using equation
-            # described in section 11.2.4 of algorithm document.
-            process_lo_species_intensity(
+            # described in section 11.2.3 of algorithm document.
+            l2_dataset = process_lo_species_intensity(
                 l2_dataset,
                 LO_NSW_SPECIES_VARIABLE_NAMES,
                 geometric_factors,
                 efficiencies,
                 NSW_POSITIONS,
             )
-
+            l2_dataset.attrs.update(
+                cdf_attrs.get_global_attributes("imap_codice_l2_lo-nsw-species")
+            )
+        elif dataset_name == "imap_codice_l2_lo-sw-angular":
+            efficiencies = efficiency_lookup[efficiency_lookup["product"] == "sw"]
+            # Calculate the sunward solar wind angular intensities using equation
+            # described in section 11.2.2 of algorithm document.
+            l2_dataset = process_lo_angular_intensity(
+                l2_dataset,
+                LO_SW_ANGULAR_VARIABLE_NAMES,
+                geometric_factors,
+                efficiencies,
+                SW_POSITIONS,
+            )
+            l2_dataset.attrs.update(
+                cdf_attrs.get_global_attributes("imap_codice_l2_lo-sw-angular")
+            )
+        if dataset_name == "imap_codice_l2_lo-nsw-angular":
+            # Calculate the non sunward angular intensities
+            efficiencies = efficiency_lookup[efficiency_lookup["product"] == "nsw"]
+            l2_dataset = process_lo_angular_intensity(
+                l2_dataset,
+                LO_NSW_ANGULAR_VARIABLE_NAMES,
+                geometric_factors,
+                efficiencies,
+                NSW_POSITIONS,
+            )
+            l2_dataset.attrs.update(
+                cdf_attrs.get_global_attributes("imap_codice_l2_lo-nsw-angular")
+            )
     if dataset_name in [
         "imap_codice_l2_hi-counters-singles",
         "imap_codice_l2_hi-counters-aggregated",
@@ -727,65 +953,6 @@ def process_codice_l2(
         # See section 11.1.2 of algorithm document
         pass
 
-    elif dataset_name == "imap_codice_l2_lo-sw-angular":
-        # Calculate the sunward angular intensities using equation described in
-        # section 11.2.3 of algorithm document.
-        pass
-
-    elif dataset_name == "imap_codice_l2_lo-nsw-angular":
-        # Calculate the non-sunward angular intensities using equation described
-        # in section 11.2.3 of algorithm document.
-        pass
-
     # logger.info(f"\nFinal data product:\n{l2_dataset}\n")
 
     return l2_dataset
-
-
-def add_dataset_attributes(
-    dataset: xr.Dataset, dataset_name: str, cdf_attrs: ImapCdfAttributes
-) -> xr.Dataset:
-    """
-    Add the global and variable attributes to the dataset.
-
-    Parameters
-    ----------
-    dataset : xarray.Dataset
-        The dataset to update.
-    dataset_name : str
-        The name of the dataset.
-    cdf_attrs : ImapCdfAttributes
-        The attribute manager for CDF attributes.
-
-    Returns
-    -------
-    xarray.Dataset
-        The updated dataset.
-    """
-    cdf_attrs.add_instrument_global_attrs("codice")
-    cdf_attrs.add_instrument_variable_attrs("codice", "l2")
-
-    # Update the global attributes
-    dataset.attrs = cdf_attrs.get_global_attributes(dataset_name)
-
-    # Set the variable attributes
-    for variable_name in dataset.data_vars.keys():
-        try:
-            dataset[variable_name].attrs = cdf_attrs.get_variable_attributes(
-                variable_name, check_schema=False
-            )
-        except KeyError:
-            # Some variables may have a product descriptor prefix in the
-            # cdf attributes key if they are common to multiple products.
-            descriptor = dataset_name.split("imap_codice_l2_")[-1]
-            cdf_attrs_key = f"{descriptor}-{variable_name}"
-            try:
-                dataset[variable_name].attrs = cdf_attrs.get_variable_attributes(
-                    f"{cdf_attrs_key}", check_schema=False
-                )
-            except KeyError:
-                logger.error(
-                    f"Field '{variable_name}' and '{cdf_attrs_key}' not found in "
-                    f"attribute manager."
-                )
-    return dataset

imap_processing/codice/codice_new_l1a.py

@@ -0,0 +1,64 @@
+"""CoDICE L1A processing functions."""
+
+import logging
+
+import xarray as xr
+from imap_data_access import ProcessingInputCollection
+
+from imap_processing import imap_module_directory
+from imap_processing.codice.codice_l1a_lo_angular import l1a_lo_angular
+from imap_processing.codice.codice_l1a_lo_species import l1a_lo_species
+from imap_processing.codice.utils import (
+    CODICEAPID,
+)
+from imap_processing.utils import packet_file_to_datasets
+
+logger = logging.getLogger(__name__)
+
+
+def process_l1a(dependency: ProcessingInputCollection) -> list[xr.Dataset]:
+    """
+    Process L1A data based on descriptor and dependencies.
+
+    Parameters
+    ----------
+    dependency : ProcessingInputCollection
+        Collection of processing inputs required for L1A processing.
+
+    Returns
+    -------
+    list[xarray.Dataset]
+        List of processed L1A datasets generated from available APIDs.
+    """
+    # Get science data which is L0 packet file
+    science_file = dependency.get_file_paths(data_type="l0")[0]
+    # Get LUT file
+    lut_file = dependency.get_file_paths(descriptor="l1a-sci-lut")[0]
+
+    logger.info(f"Processing L1A for {science_file.name} with {lut_file.name}")
+
+    xtce_file = (
+        imap_module_directory / "codice/packet_definitions/codice_packet_definition.xml"
+    )
+    # Decom packet
+    datasets_by_apid = packet_file_to_datasets(
+        science_file,
+        xtce_file,
+    )
+
+    datasets = []
+    for apid in datasets_by_apid:
+        if apid == CODICEAPID.COD_LO_SW_SPECIES_COUNTS:
+            logger.info("Processing Lo SW Species Counts")
+            datasets.append(l1a_lo_species(datasets_by_apid[apid], lut_file))
+        elif apid == CODICEAPID.COD_LO_NSW_SPECIES_COUNTS:
+            logger.info("Processing Lo NSW Species Counts")
+            datasets.append(l1a_lo_species(datasets_by_apid[apid], lut_file))
+        elif apid == CODICEAPID.COD_LO_SW_ANGULAR_COUNTS:
+            logger.info("Processing Lo SW Angular Counts")
+            datasets.append(l1a_lo_angular(datasets_by_apid[apid], lut_file))
+        elif apid == CODICEAPID.COD_LO_NSW_ANGULAR_COUNTS:
+            logger.info("Processing Lo NSW Angular Counts")
+            datasets.append(l1a_lo_angular(datasets_by_apid[apid], lut_file))
+
+    return datasets

imap_processing/codice/constants.py

@@ -61,6 +61,8 @@ CODICEAPID_MAPPING = {
 SPIN_PERIOD_CONVERSION = 0.00032
 K_FACTOR = 5.76  # This is used to convert voltages to energies in L2
 HI_ACQUISITION_TIME = 0.59916
+NUM_ESA_STEPS = 128
+LO_DESPIN_SPIN_SECTORS = 24
 
 # CDF variable names used for lo data products
 LO_COUNTERS_SINGLES_VARIABLE_NAMES = ["apd_singles"]
@@ -2281,8 +2283,9 @@ HALF_SPIN_LUT = {
 }
 
 NSW_POSITIONS = [x for x in range(3, 22)]
-SW_POSITIONS = [0]
-PUI_POSITIONS = [0, 1, 2, 22, 23]
+SW_POSITIONS = [0, 1, 2, 22, 23]
+SOLAR_WIND_POSITIONS = [0]
+PUI_POSITIONS = SW_POSITIONS
 L2_GEOMETRIC_FACTOR = 0.013
 L2_HI_NUMBER_OF_SSD = 12.0
 
@@ -2320,3 +2323,35 @@ HI_L2_ELEVATION_ANGLE = np.array(
     ],
     dtype=float,
 )
+
+
+LO_POSITION_TO_ELEVATION_ANGLE = {
+    "sw": {
+        1: 0,
+        2: 15,
+        24: 15,
+        3: 30,
+        23: 30,
+    },
+    "nsw": {
+        4: 45,
+        22: 45,
+        5: 60,
+        21: 60,
+        6: 75,
+        20: 75,
+        7: 90,
+        19: 90,
+        8: 105,
+        18: 105,
+        9: 120,
+        17: 120,
+        10: 135,
+        16: 135,
+        11: 150,
+        15: 150,
+        12: 165,
+        14: 165,
+        13: 180,
+    },
+}