imap-processing 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

imap_processing/codice/codice_l1a.py

@@ -54,6 +54,8 @@ class CoDICEL1aPipeline:
 
     Methods
     -------
+    apply_despinning()
+        Apply the despinning algorithm to lo- angular and priority products.
     decompress_data(science_values)
         Perform decompression on the data.
     define_coordinates()
@@ -87,6 +89,73 @@ class CoDICEL1aPipeline:
         self.plan_step = plan_step
         self.view_id = view_id
 
+    def apply_despinning(self) -> None:
+        """
+        Apply the despinning algorithm to lo- angular and priority products.
+
+        This only applies to CoDICE-Lo angular and priority data products. See
+        sections 9.3.4 and 9.3.5 of the algorithm document for more details.
+        """
+        # Determine the appropriate dimensions for the despun data
+        num_energies = self.config["dims"]["esa_step"]
+        num_spin_sectors = self.config["dims"]["spin_sector"]
+        num_spins = num_spin_sectors * 2
+        num_counters = self.config["num_counters"]
+        num_positions = self.config["dims"].get(
+            "inst_az"
+        )  # Defaults to None if not present
+
+        # The dimensions are dependent on the specific data product
+        if "angular" in self.config["dataset_name"]:
+            despun_dims: tuple[int, ...] = (
+                num_energies,
+                num_positions,
+                num_spins,
+                num_counters,
+            )
+        elif "priority" in self.config["dataset_name"]:
+            despun_dims = (num_energies, num_spins, num_counters)
+
+        # Placeholder for finalized despun data
+        self.data: list[np.ndarray]  # Needed to appease mypy
+        despun_data = [np.zeros(despun_dims) for _ in range(len(self.data))]
+
+        # Iterate over the energy and spin sector indices, and determine the
+        # appropriate pixel orientation. The combination of the pixel
+        # orientation and the azimuth determine which spin sector the data
+        # gets stored in.
+        # TODO: All these nested for-loops are bad. Try to find a better
+        #       solution.
+        for i, epoch_data in enumerate(self.data):
+            for energy_index in range(num_energies):
+                pixel_orientation = constants.PIXEL_ORIENTATIONS[energy_index]
+                for spin_sector_index in range(num_spin_sectors):
+                    for azimuth_index in range(num_spins):
+                        if pixel_orientation == "A" and azimuth_index < 12:
+                            despun_spin_sector = spin_sector_index
+                        elif pixel_orientation == "A" and azimuth_index >= 12:
+                            despun_spin_sector = spin_sector_index + 12
+                        elif pixel_orientation == "B" and azimuth_index < 12:
+                            despun_spin_sector = spin_sector_index + 12
+                        elif pixel_orientation == "B" and azimuth_index >= 12:
+                            despun_spin_sector = spin_sector_index
+
+                        if "angular" in self.config["dataset_name"]:
+                            spin_data = epoch_data[
+                                energy_index, :, spin_sector_index, :
+                            ]  # (5, 4)
+                            despun_data[i][energy_index, :, despun_spin_sector, :] = (
+                                spin_data
+                            )
+                        elif "priority" in self.config["dataset_name"]:
+                            spin_data = epoch_data[energy_index, spin_sector_index, :]
+                            despun_data[i][energy_index, despun_spin_sector, :] = (
+                                spin_data
+                            )
+
+        # Replace original data
+        self.data = despun_data
+
     def decompress_data(self, science_values: list[NDArray[str]] | list[str]) -> None:
         """
         Perform decompression on the data.
@@ -122,7 +191,7 @@ class CoDICEL1aPipeline:
 
         else:
             for packet_data, byte_count in zip(
-                science_values, self.dataset.byte_count.data
+                science_values, self.dataset.byte_count.data, strict=False
             ):
                 # Convert from numpy array to byte object
                 values = packet_data[()]
@@ -134,17 +203,20 @@ class CoDICEL1aPipeline:
                 decompressed_values = decompress(values, compression_algorithm)
                 self.raw_data.append(decompressed_values)
 
-    def define_coordinates(self) -> None:
+    def define_coordinates(self) -> None:  # noqa: PLR0912 (too many branches)
         """
         Create ``xr.DataArrays`` for the coords needed in the final dataset.
 
         The coordinates for the dataset depend on the data product being made.
+
+        # TODO: Split this function up or simplify it to avoid too many branches
+        #       error.
         """
         self.coords = {}
 
         coord_names = [
-            *self.config["output_dims"].keys(),
-            *[key + "_label" for key in self.config["output_dims"].keys()],
+            *self.config["dims"].keys(),
+            *[key + "_label" for key in self.config["dims"].keys()],
         ]
 
         # Define epoch coordinates
@@ -169,12 +241,17 @@ class CoDICEL1aPipeline:
             if name in [
                 "esa_step",
                 "inst_az",
-                "spin_sector",
                 "spin_sector_pairs",
                 "spin_sector_index",
                 "ssd_index",
             ]:
-                values = np.arange(self.config["output_dims"][name])
+                values = np.arange(self.config["dims"][name])
+                dims = [name]
+            elif name == "spin_sector":
+                if self.config["dataset_name"] in constants.REQUIRES_DESPINNING:
+                    values = np.arange(24)
+                else:
+                    values = np.arange(self.config["dims"][name])
                 dims = [name]
             elif name == "spin_sector_pairs_label":
                 values = np.array(
@@ -188,16 +265,29 @@ class CoDICEL1aPipeline:
                     ]
                 )
                 dims = [name]
+            elif name == "inst_az_label":
+                if self.config["dataset_name"] == "imap_codice_l1a_lo-nsw-angular":
+                    values = [str(x) for x in range(4, 23)]
+                elif self.config["dataset_name"] == "imap_codice_l1a_lo-sw-angular":
+                    values = ["1", "2", "3", "23", "24"]
+                else:
+                    values = np.arange(self.config["dims"]["inst_az"]).astype(str)
+                dims = ["inst_az"]
             elif name in [
-                "spin_sector_label",
                 "esa_step_label",
-                "inst_az_label",
                 "spin_sector_index_label",
                 "ssd_index_label",
             ]:
                 key = name.removesuffix("_label")
-                values = np.arange(self.config["output_dims"][key]).astype(str)
+                values = np.arange(self.config["dims"][key]).astype(str)
                 dims = [key]
+            elif name == "spin_sector_label":
+                key = name.removesuffix("_label")
+                dims = [key]
+                if self.config["dataset_name"] in constants.REQUIRES_DESPINNING:
+                    values = np.arange(24).astype(str)
+                else:
+                    values = np.arange(self.config["dims"][key]).astype(str)
 
             coord = xr.DataArray(
                 values,
@@ -230,16 +320,16 @@ class CoDICEL1aPipeline:
         # Stack the data so that it is easier to reshape and iterate over
         all_data = np.stack(self.data)
 
-        # The dimension of all_data is something like (epoch, num_counters,
-        # num_energy_steps, num_positions, num_spin_sectors) (or may be slightly
+        # The dimension of all_data is something like (epoch, num_energy_steps,
+        # num_positions, num_spin_sectors, num_counters) (or may be slightly
         # different depending on the data product). In any case, iterate over
         # the num_counters dimension to isolate the data for each counter so
         # each counter's data can be placed in a separate CDF data variable.
         for counter, variable_name in zip(
-            range(all_data.shape[1]), self.config["variable_names"]
+            range(all_data.shape[-1]), self.config["variable_names"], strict=False
         ):
             # Extract the counter data
-            counter_data = all_data[:, counter, ...]
+            counter_data = all_data[..., counter]
 
             # Get the CDF attributes
             descriptor = self.config["dataset_name"].split("imap_codice_l1a_")[-1]
@@ -249,7 +339,7 @@ class CoDICEL1aPipeline:
             # For most products, the final CDF dimensions always has "epoch" as
             # the first dimension followed by the dimensions for the specific
             # data product
-            dims = ["epoch", *list(self.config["output_dims"].keys())]
+            dims = ["epoch", *list(self.config["dims"].keys())]
 
             # However, CoDICE-Hi products use specific energy bins for the
             # energy dimension
@@ -306,7 +396,7 @@ class CoDICEL1aPipeline:
             ``xarray`` dataset for the data product, with added energy variables.
         """
         energy_bin_name = f"energy_{species}"
-        centers, deltas = self.get_hi_energy_table_data(
+        centers, deltas_minus, deltas_plus = self.get_hi_energy_table_data(
             energy_bin_name.split("energy_")[-1]
         )
 
@@ -319,11 +409,19 @@ class CoDICEL1aPipeline:
                 check_schema=False,
             ),
         )
-        dataset[f"{energy_bin_name}_delta"] = xr.DataArray(
-            deltas,
-            dims=[f"{energy_bin_name}_delta"],
+        dataset[f"{energy_bin_name}_minus"] = xr.DataArray(
+            deltas_minus,
+            dims=[f"{energy_bin_name}_minus"],
+            attrs=self.cdf_attrs.get_variable_attributes(
+                f"{self.config['dataset_name'].split('_')[-1]}-{energy_bin_name}_minus",
+                check_schema=False,
+            ),
+        )
+        dataset[f"{energy_bin_name}_plus"] = xr.DataArray(
+            deltas_plus,
+            dims=[f"{energy_bin_name}_plus"],
             attrs=self.cdf_attrs.get_variable_attributes(
-                f"{self.config['dataset_name'].split('_')[-1]}-{energy_bin_name}_delta",
+                f"{self.config['dataset_name'].split('_')[-1]}-{energy_bin_name}_plus",
                 check_schema=False,
             ),
         )
@@ -398,6 +496,12 @@ class CoDICEL1aPipeline:
                 dims = ["epoch"]
                 attrs = self.cdf_attrs.get_variable_attributes("spin_period")
 
+            # The k-factor is a constant that maps voltages to energies
+            elif variable_name == "k_factor":
+                variable_data = np.array([constants.K_FACTOR], dtype=np.float32)
+                dims = [""]
+                attrs = self.cdf_attrs.get_variable_attributes("k_factor")
+
             # Add variable to the dataset
             dataset[variable_name] = xr.DataArray(
                 variable_data,
@@ -475,7 +579,7 @@ class CoDICEL1aPipeline:
 
     def get_hi_energy_table_data(
         self, species: str
-    ) -> tuple[NDArray[float], NDArray[float]]:
+    ) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
         """
         Retrieve energy table data for CoDICE-Hi products.
 
@@ -493,22 +597,25 @@ class CoDICEL1aPipeline:
         -------
         centers : NDArray[float]
             An array whose values represent the centers of the energy bins.
-        deltas : NDArray[float]
-            An array whose values represent the deltas of the energy bins.
+        deltas_minus : NDArray[float]
+            An array whose values represent the minus deltas of the energy bins.
+        deltas_plus : NDArray[float]
+            An array whose values represent the plus deltas of the energy bins.
         """
         data_product = self.config["dataset_name"].split("-")[-1].upper()
-        energy_table = getattr(constants, f"{data_product}_ENERGY_TABLE")[species]
-
-        # Find the centers and deltas of the energy bins
-        centers = np.array(
-            [
-                (energy_table[i] + energy_table[i + 1]) / 2
-                for i in range(len(energy_table) - 1)
-            ]
+        energy_table = np.array(
+            getattr(constants, f"{data_product}_ENERGY_TABLE")[species]
         )
-        deltas = energy_table[1:] - centers
 
-        return centers, deltas
+        # Find the geometric centers and deltas of the energy bins
+        # The delta minus is the difference between the center of the bin
+        # and the 'left edge' of the bin. The delta plus is the difference
+        # between the 'right edge' of the bin and the center of the bin
+        centers = np.sqrt(energy_table[:-1] * energy_table[1:])
+        deltas_minus = centers - energy_table[:-1]
+        deltas_plus = energy_table[1:] - centers
+
+        return centers, deltas_minus, deltas_plus
 
     def reshape_binned_data(self, dataset: xr.Dataset) -> dict[str, list]:
         """
@@ -589,52 +696,31 @@ class CoDICEL1aPipeline:
 
         These data need to be divided up by species or priorities (or
         what I am calling "counters" as a general term), and re-arranged into
-        4D arrays representing dimensions such as time, spin sectors, positions,
-        and energies (depending on the data product).
+        multidimensional arrays representing dimensions such as time,
+        spin sectors, positions, and energies (depending on the data product).
 
         However, the existence and order of these dimensions can vary depending
-        on the specific data product, so we define this in the "input_dims"
-        and "output_dims" values configuration dictionary; the "input_dims"
-        defines how the dimensions are written into the packet data, while
-        "output_dims" defines how the dimensions should be written to the final
-        CDF product.
+        on the specific data product, so we define this in the "dims" key of the
+        configuration dictionary.
         """
         # This will contain the reshaped data for all counters
         self.data = []
 
-        # First reshape the data based on how it is written to the data array of
-        # the packet data. The number of counters is the first dimension / axis,
-        # with the exception of lo-counters-aggregated which is treated slightly
-        # differently
-        if self.config["dataset_name"] != "imap_codice_l1a_lo-counters-aggregated":
-            reshape_dims = (
-                self.config["num_counters"],
-                *self.config["input_dims"].values(),
-            )
-        else:
-            reshape_dims = (
-                *self.config["input_dims"].values(),
-                self.config["num_counters"],
-            )
-
-        # Then, transpose the data based on how the dimensions should be written
-        # to the CDF file. Since this is specific to each data product, we need
-        # to determine this dynamically based on the "output_dims" config.
-        # Again, lo-counters-aggregated is treated slightly differently
-        input_keys = ["num_counters", *self.config["input_dims"].keys()]
-        output_keys = ["num_counters", *self.config["output_dims"].keys()]
-        if self.config["dataset_name"] != "imap_codice_l1a_lo-counters-aggregated":
-            transpose_axes = [input_keys.index(dim) for dim in output_keys]
-        else:
-            transpose_axes = [1, 2, 0]  # [esa_step, spin_sector_pairs, num_counters]
-
+        # Reshape the data based on how it is written to the data array of
+        # the packet data. The number of counters is the last dimension / axis.
+        reshape_dims = (
+            *self.config["dims"].values(),
+            self.config["num_counters"],
+        )
         for packet_data in self.raw_data:
             reshaped_packet_data = np.array(packet_data, dtype=np.uint32).reshape(
                 reshape_dims
             )
-            reshaped_cdf_data = np.transpose(reshaped_packet_data, axes=transpose_axes)
+            self.data.append(reshaped_packet_data)
 
-            self.data.append(reshaped_cdf_data)
+        # Apply despinning if necessary
+        if self.config["dataset_name"] in constants.REQUIRES_DESPINNING:
+            self.apply_despinning()
 
         # No longer need to keep the raw data around
         del self.raw_data
@@ -950,9 +1036,9 @@ def create_direct_event_dataset(apid: int, packets: xr.Dataset) -> xr.Dataset:
 
     # Create the dataset to hold the data variables
     if apid == CODICEAPID.COD_LO_PHA:
-        attrs = cdf_attrs.get_global_attributes("imap_codice_l1a_lo-pha")
+        attrs = cdf_attrs.get_global_attributes("imap_codice_l1a_lo-direct-events")
     elif apid == CODICEAPID.COD_HI_PHA:
-        attrs = cdf_attrs.get_global_attributes("imap_codice_l1a_hi-pha")
+        attrs = cdf_attrs.get_global_attributes("imap_codice_l1a_hi-direct-events")
     dataset = xr.Dataset(
         coords={
             "epoch": epoch,
@@ -967,9 +1053,9 @@ def create_direct_event_dataset(apid: int, packets: xr.Dataset) -> xr.Dataset:
     # Create the CDF data variables for each Priority and Field
     for i in range(constants.DE_DATA_PRODUCT_CONFIGURATIONS[apid]["num_priorities"]):
         for field in constants.DE_DATA_PRODUCT_CONFIGURATIONS[apid]["cdf_fields"]:
-            variable_name = f"P{i}_{field}"
+            variable_name = f"p{i}_{field}"
             attrs = cdf_attrs.get_variable_attributes(variable_name)
-            if field in ["NumEvents", "DataQuality"]:
+            if field in ["num_events", "data_quality"]:
                 dims = ["epoch"]
             else:
                 dims = ["epoch", "event_num"]
@@ -1427,13 +1513,15 @@ def reshape_de_data(
     for priority_num in range(num_priorities):
         for field in bit_structure:
             if field not in ["Priority", "Spare"]:
-                data[f"P{priority_num}_{field}"] = np.full(
+                data[f"p{priority_num}_{field}"] = np.full(
                     (num_epochs, 10000),
                     bit_structure[field]["fillval"],
                     dtype=bit_structure[field]["dtype"],
                 )
-        data[f"P{priority_num}_NumEvents"] = np.full(num_epochs, 65535, dtype=np.uint16)
-        data[f"P{priority_num}_DataQuality"] = np.full(num_epochs, 255, dtype=np.uint8)
+        data[f"p{priority_num}_num_events"] = np.full(
+            num_epochs, 65535, dtype=np.uint16
+        )
+        data[f"p{priority_num}_data_quality"] = np.full(num_epochs, 255, dtype=np.uint8)
 
     # decompressed_data is one large list of values of length
     # (<number of epochs> * <number of priorities>)
@@ -1457,8 +1545,8 @@ def reshape_de_data(
 
             # Number of events and data quality can be determined at this stage
             num_events = num_events_arr[epoch_start:epoch_end][i]
-            data[f"P{priority_num}_NumEvents"][epoch_index] = num_events
-            data[f"P{priority_num}_DataQuality"][epoch_index] = data_quality[i]
+            data[f"p{priority_num}_num_events"][epoch_index] = num_events
+            data[f"p{priority_num}_data_quality"][epoch_index] = data_quality[i]
 
             # Iterate over each event
             for event_index in range(num_events):
@@ -1489,7 +1577,7 @@ def reshape_de_data(
                     )
 
                     # Set the value into the data array
-                    data[f"P{priority_num}_{field_name}"][epoch_index, event_index] = (
+                    data[f"p{priority_num}_{field_name}"][epoch_index, event_index] = (
                         value
                     )
                     bit_position += field_components["bit_length"]

imap_processing/codice/codice_l1b.py

@@ -116,7 +116,7 @@ def process_codice_l1b(file_path: Path) -> xr.Dataset:
     descriptor = dataset_name.removeprefix("imap_codice_l1b_")
 
     # Direct event data products do not have a level L1B
-    if descriptor in ["lo-pha", "hi-pha"]:
+    if descriptor in ["lo-direct-events", "hi-direct-events"]:
         logger.warning("Encountered direct event data product. Skipping L1b processing")
         return None
 

imap_processing/codice/codice_l2.py

@@ -45,22 +45,112 @@ def process_codice_l2(file_path: Path) -> xr.Dataset:
     # TODO: Could clean this up by using imap-data-access methods?
     dataset_name = l1_dataset.attrs["Logical_source"]
     data_level = dataset_name.removeprefix("imap_codice_").split("_")[0]
-    descriptor = dataset_name.removeprefix(f"imap_codice_{data_level}_")
     dataset_name = dataset_name.replace(data_level, "l2")
 
-    # TODO: Temporary work-around to replace "PHA" naming convention with
-    #       "direct events" This will eventually be changed at the L1 level and
-    #       thus this will eventually be removed.
-    if descriptor == "lo-pha":
-        dataset_name = dataset_name.replace("lo-pha", "lo-direct-events")
-    elif descriptor == "hi-pha":
-        dataset_name = dataset_name.replace("hi-pha", "hi-direct-events")
-
     # Use the L1 data product as a starting point for L2
     l2_dataset = l1_dataset.copy()
 
     # Get the L2 CDF attributes
     cdf_attrs = ImapCdfAttributes()
+    l2_dataset = add_dataset_attributes(l2_dataset, dataset_name, cdf_attrs)
+
+    if dataset_name in [
+        "imap_codice_l2_hi-counters-singles",
+        "imap_codice_l2_hi-counters-aggregated",
+        "imap_codice_l2_lo-counters-singles",
+        "imap_codice_l2_lo-counters-aggregated",
+        "imap_codice_l2_lo-sw-priority",
+        "imap_codice_l2_lo-nsw-priority",
+    ]:
+        # No changes needed. Just save to an L2 CDF file.
+        pass
+
+    elif dataset_name == "imap_codice_l2_hi-direct-events":
+        # Convert the following data variables to physical units using
+        # calibration data:
+        #    - ssd_energy
+        #    - tof
+        #    - elevation_angle
+        #    - spin_angle
+        # These converted variables are *in addition* to the existing L1 variables
+        # The other data variables require no changes
+        # See section 11.1.2 of algorithm document
+        pass
+
+    elif dataset_name == "imap_codice_l2_hi-sectored":
+        # Convert the sectored count rates using equation described in section
+        # 11.1.3 of algorithm document.
+        pass
+
+    elif dataset_name == "imap_codice_l2_hi-omni":
+        # Calculate the omni-directional intensity for each species using
+        # equation described in section 11.1.4 of algorithm document
+        # hopefully this can also apply to hi-ialirt
+        pass
+
+    elif dataset_name == "imap_codice_l2_lo-direct-events":
+        # Convert the following data variables to physical units using
+        # calibration data:
+        #    - apd_energy
+        #    - elevation_angle
+        #    - tof
+        #    - spin_sector
+        #    - esa_step
+        # These converted variables are *in addition* to the existing L1 variables
+        # The other data variables require no changes
+        # 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
+
+    elif dataset_name == "imap_codice_l2_lo-sw-species":
+        # Calculate the sunward solar wind species intensities using equation
+        # described in section 11.2.4 of algorithm document.
+        # Calculate the pickup ion sunward solar wind intensities using equation
+        # described in section 11.2.4 of algorithm document.
+        # Hopefully this can also apply to lo-ialirt
+        pass
+
+    elif dataset_name == "imap_codice_l2_lo-nsw-species":
+        # Calculate the non-sunward solar wind species intensities using
+        # equation described in section 11.2.4 of algorithm document.
+        # Calculate the pickup ion non-sunward solar wind intensities using
+        # equation described in section 11.2.4 of algorithm document.
+        pass
+
+    logger.info(f"\nFinal data product:\n{l2_dataset}\n")
+
+    return l2_dataset
+
+
+def add_dataset_attributes(
+    l2_dataset: xr.Dataset, dataset_name: str, cdf_attrs: ImapCdfAttributes
+) -> xr.Dataset:
+    """
+    Add the global and variable attributes to the dataset.
+
+    Parameters
+    ----------
+    l2_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")
 
@@ -68,14 +158,23 @@ def process_codice_l2(file_path: Path) -> xr.Dataset:
     l2_dataset.attrs = cdf_attrs.get_global_attributes(dataset_name)
 
     # Set the variable attributes
-    for variable_name in l2_dataset:
-        l2_dataset[variable_name].attrs = cdf_attrs.get_variable_attributes(
-            variable_name, check_schema=False
-        )
-
-    # TODO: Add L2-specific algorithms/functionality here. For SIT-4, we can
-    #       just keep the data as-is.
-
-    logger.info(f"\nFinal data product:\n{l2_dataset}\n")
-
+    for variable_name in l2_dataset.data_vars.keys():
+        try:
+            l2_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:
+                l2_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 l2_dataset