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

imap_processing/lo/l2/lo_l2.py

@@ -75,7 +75,18 @@ def lo_l2(
     dataset = add_geometric_factors(dataset, map_descriptor.species)
 
     logger.info("Step 4: Calculating rates and intensities")
-    dataset = calculate_all_rates_and_intensities(dataset)
+
+    # Determine if corrections are needed and prepare oxygen data if required
+    sputtering_correction, bootstrap_correction, o_map_dataset = _prepare_corrections(
+        map_descriptor, descriptor, sci_dependencies, anc_dependencies
+    )
+
+    dataset = calculate_all_rates_and_intensities(
+        dataset,
+        sputtering_correction=sputtering_correction,
+        bootstrap_correction=bootstrap_correction,
+        o_map_dataset=o_map_dataset,
+    )
 
     logger.info("Step 5: Finalizing dataset with attributes")
     dataset = finalize_dataset(dataset, descriptor)
@@ -84,6 +95,59 @@ def lo_l2(
     return [dataset]
 
 
+def _prepare_corrections(
+    map_descriptor: MapDescriptor,
+    descriptor: str,
+    sci_dependencies: dict,
+    anc_dependencies: list,
+) -> tuple[bool, bool, xr.Dataset | None]:
+    """
+    Determine what corrections are needed and prepare oxygen dataset if required.
+
+    This helper function encapsulates the logic for determining when sputtering
+    and bootstrap corrections should be applied, and handles the creation of
+    the oxygen dataset needed for sputtering corrections.
+
+    Parameters
+    ----------
+    map_descriptor : MapDescriptor
+        The parsed map descriptor containing species and data type information.
+    descriptor : str
+        The original descriptor string for creating the oxygen variant.
+    sci_dependencies : dict
+        Dictionary of datasets needed for L2 data product creation.
+    anc_dependencies : list
+        List of ancillary file paths.
+
+    Returns
+    -------
+    tuple[bool, bool, xr.Dataset | None]
+        A tuple containing:
+        - sputtering_correction: Whether to apply sputtering corrections
+        - bootstrap_correction: Whether to apply bootstrap corrections
+        - o_map_dataset: Oxygen dataset if needed, None otherwise
+    """
+    # Default values - no corrections needed
+    sputtering_correction = False
+    bootstrap_correction = False
+    o_map_dataset = None
+
+    # Sputtering and bootstrap corrections are only applied to hydrogen ENA data
+    # Guard against recursion: don't process oxygen for oxygen maps
+    if (
+        map_descriptor.species == "h"
+        and map_descriptor.principal_data == "ena"
+        and "-o-" not in descriptor
+    ):  # Safety check to prevent infinite recursion
+        logger.info("Creating map for oxygen for sputtering corrections")
+        o_descriptor = descriptor.replace("-h-", "-o-")
+        o_map_dataset = lo_l2(sci_dependencies, anc_dependencies, o_descriptor)[0]
+        sputtering_correction = True
+        bootstrap_correction = True
+
+    return sputtering_correction, bootstrap_correction, o_map_dataset
+
+
 # =============================================================================
 # SETUP AND INITIALIZATION HELPERS
 # =============================================================================
@@ -600,6 +664,7 @@ def calculate_all_rates_and_intensities(
     dataset: xr.Dataset,
     sputtering_correction: bool = False,
     bootstrap_correction: bool = False,
+    o_map_dataset: xr.Dataset | None = None,
 ) -> xr.Dataset:
     """
     Calculate rates and intensities with proper error propagation.
@@ -614,6 +679,8 @@ def calculate_all_rates_and_intensities(
     bootstrap_correction : bool, optional
         Whether to apply bootstrap corrections to intensities.
         Default is False.
+    o_map_dataset : xr.Dataset, optional
+        Dataset specifically for oxygen, needed for sputtering corrections.
 
     Returns
     -------
@@ -632,12 +699,9 @@ def calculate_all_rates_and_intensities(
 
     # Optional Step 4: Calculate sputtering corrections
     if sputtering_correction:
-        # TODO: The second dataset is for Oxygen specifically,
-        #       if we get an H dataset in, we may need to calculate
-        #       the O dataset separately before calling here.
-        dataset = calculate_sputtering_corrections(dataset, dataset)
+        dataset = calculate_sputtering_corrections(dataset, o_map_dataset)
 
-    # Optional Step 5: Clean up intermediate variables
+    # Optional Step 5: Calculate bootstrap corrections
     if bootstrap_correction:
         dataset = calculate_bootstrap_corrections(dataset)
 
@@ -764,7 +828,7 @@ def calculate_sputtering_corrections(
     ----------
     dataset : xr.Dataset
         Dataset with count rates, geometric factors, and center energies.
-        This could be either an H or O dataset.
+        This is an H dataset that we are applying the corrections to.
     o_dataset : xr.Dataset
         Dataset specifically for oxygen, needed to access oxygen intensities
         and uncertainties.
@@ -773,9 +837,9 @@ def calculate_sputtering_corrections(
     -------
     xr.Dataset
         Dataset with calculated sputtering-corrected intensities and their
-        uncertainties for hydrogen and oxygen.
+        uncertainties.
     """
-    logger.info("Applying sputtering corrections to oxygen intensities")
+    logger.info("Applying sputtering corrections to hydrogen intensities")
     # Only apply sputtering correction to esa levels 5 and 6 (indices 4 and 5)
     energy_indices = [4, 5]
     small_dataset = dataset.isel(epoch=0, energy=energy_indices)
@@ -789,6 +853,10 @@ def calculate_sputtering_corrections(
         "bg_rates_stat_uncert"
     ] / (o_small_dataset["geometric_factor"] * o_small_dataset["energy"])
 
+    # We need to align the energy dimensions from the oxygen dataset to the
+    # Hydrogen dataset so the calculations below get aligned by xarray correctly.
+    o_small_dataset["energy"] = small_dataset["energy"]
+
     # Equation 9
     j_o_prime = o_small_dataset["ena_intensity"] - o_small_dataset["bg_intensity"]
     j_o_prime.values[j_o_prime.values < 0] = 0  # No negative intensities
@@ -801,10 +869,10 @@ def calculate_sputtering_corrections(
 
     # NOTE: From table 2 of the mapping document, for energy level 5 and 6
     sputter_correction_factor = xr.DataArray(
-        [0.15, 0.01], dims=["energy"], coords={"energy": energy_indices}
+        [0.15, 0.01], dims=["energy"], coords={"energy": small_dataset["energy"]}
     )
     # Equation 11
-    # Remove the sputtered oxygen intensity to correct the original O intensity
+    # Remove the sputtered oxygen intensity to correct the original H intensity
     sputter_corrected_intensity = (
         small_dataset["ena_intensity"] - sputter_correction_factor * j_o_prime
     )
@@ -870,8 +938,10 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
         bootstrap_factor_array,
         dims=["energy_i", "energy_k"],
         coords={
-            "energy_i": list(range(7)),
-            "energy_k": list(range(8)),  # Include virtual channel 7 (index 7)
+            "energy_i": dataset["energy"].values,
+            # Add an extra coordinate for the virtual E8 channel, unused
+            # in the broadcasting calculations
+            "energy_k": np.concatenate([dataset["energy"].values, [np.nan]]),
         },
     )
 
@@ -933,7 +1003,7 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
         # NOTE: The paper uses 1-based indexing and we use 0-based indexing
         #       so there is an off-by-one difference in the indices.
         bootstrap_intensity_i[:] = (
-            j_c_prime_i - bootstrap_factor.sel(energy_i=i, energy_k=7) * j_8_b[0, ...]
+            j_c_prime_i - bootstrap_factor.isel(energy_i=i, energy_k=7) * j_8_b[0, ...]
         )
         # NOTE: We will square root at the end to get the uncertainty, but
         #       all equations are with variances
@@ -945,7 +1015,7 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
 
         # Get bootstrap factors for this i and the relevant k values
         # Rename energy_k dimension to energy for alignment with intensity
-        bootstrap_factors_k = bootstrap_factor.sel(
+        bootstrap_factors_k = bootstrap_factor.isel(
             energy_i=i, energy_k=k_indices
         ).rename({"energy_k": "energy"})
 

imap_processing/mag/l1a/mag_l1a.py

@@ -43,7 +43,7 @@ def mag_l1a(packet_filepath: Path) -> list[xr.Dataset]:
         A list of generated filenames.
     """
     packets = decom_mag.decom_packets(packet_filepath)
-
+    logger.info("Packet decoding complete, beginning L1A processing.")
     norm_data = packets["norm"]
     burst_data = packets["burst"]
 
@@ -188,7 +188,7 @@ def process_packets(
             secondary_packet_data.start_time,
         )
 
-        # Sort primary and secondary into MAGo and MAGi by 24 hour chunks
+        # Sort primary and secondary into MAGo and MAGi
 
         if mago is None:
             mago = MagL1a(

imap_processing/mag/l1a/mag_l1a_data.py

@@ -205,13 +205,25 @@ class MagL1a:
         1 if the sensor is active, 0 if not
     shcoarse : int
         Mission elapsed time for the first packet, the start time for the whole day
-    vectors : numpy.ndarray
-        List of magnetic vector samples, starting at start_time. [x, y, z, range, time],
-        where time is numpy.datetime64[ns]
+    starting_vectors : InitVar[numpy.ndarray]
+        Initvar to create the first entry in the vector list. This is to preserve the
+        external API of creating the object with the first set of vectors.
+        This cannot be accessed from an instance of the class. Instead, vectors
+        should be used.
     starting_packet : InitVar[MagL1aPacketProperties]
         The packet properties for the first packet in the day. As an InitVar, this
         cannot be accessed from an instance of the class. Instead, packet_definitions
         should be used.
+    vectors : numpy.ndarray
+        List of magnetic vector samples, starting at start_time. [x, y, z, range, time],
+        where time is numpy.datetime64[ns]. This is a property that concatenates the
+        internal vector list on demand.
+    compression_flags : numpy.ndarray
+        Array of flags to indicate compression and width for all timestamps in the
+        L1A file. Shaped like (n, 2) where n is the number of vectors. First value
+        is a boolean for compressed/uncompressed, second vector is a number between 0-20
+        if the data is compressed, which is the width in bits of the compressed data.
+        This is a property that concatenates the internal compression flags list.
     packet_definitions : dict[numpy.datetime64, MagL1aPacketProperties]
         Dictionary of packet properties for each packet in the day. The key is the start
         time of the packet, and the value is a dataclass of packet properties.
@@ -221,11 +233,20 @@ class MagL1a:
         List of missing sequence numbers in the day
     start_time : numpy.int64
         Start time of the day, in ns since J2000 epoch
-    compression_flags : np.ndarray
+    _compression_flags_list : np.ndarray
         Array of flags to indication compression and width for all timestamps in the
         L1A file. Shaped like (n, 2) where n is the number of vectors. First value
         is a boolean for compressed/uncompressed, second vector is a number between 0-20
         if the data is compressed, which is the width in bits of the compressed data.
+        Transformed into a numpy array upon retrieval.
+    _vector_list : list
+        Internal list of vectors, used to build the final vectors attribute.
+        This is a list of numpy arrays, each with shape (n, 5) where n is the
+        number of vectors in that packet, and each vector is (x, y, z, range, time).
+    _vector_cache : numpy.ndarray | None
+        A cache of the concatenated vector list. This is None until the vectors
+        property is accessed, at which point it is created and stored here for future
+        access.
 
     Methods
     -------
@@ -248,23 +269,30 @@ class MagL1a:
     is_mago: bool
     is_active: int
     shcoarse: int
-    vectors: np.ndarray
+    starting_vectors: InitVar[np.ndarray]
     starting_packet: InitVar[MagL1aPacketProperties]
     packet_definitions: dict[np.int64, MagL1aPacketProperties] = field(init=False)
     most_recent_sequence: int = field(init=False)
     missing_sequences: list[int] = field(default_factory=list)
     start_time: np.int64 = field(init=False)
-    compression_flags: np.ndarray | None = field(init=False, default=None)
+    _compression_flags_list: list = field(default_factory=list)
+    _vector_list: list = field(init=False)
+    _vector_cache: np.ndarray | None = field(init=False, default=None)
 
-    def __post_init__(self, starting_packet: MagL1aPacketProperties) -> None:
+    def __post_init__(
+        self, starting_vectors: np.ndarray, starting_packet: MagL1aPacketProperties
+    ) -> None:
         """
-        Initialize the packet_definition dictionary and most_recent_sequence.
+        Initialize the vector list, packet_definition dictionary & most_recent_sequence.
 
         Parameters
         ----------
+        starting_vectors : numpy.ndarray
+            The vectors for the first packet in the day.
         starting_packet : MagL1aPacketProperties
             The packet properties for the first packet in the day, including start time.
         """
+        self._vector_list = [starting_vectors]
         self.start_time = np.int64(met_to_ttj2000ns(starting_packet.shcoarse))
         self.packet_definitions = {self.start_time: starting_packet}
         # most_recent_sequence is the sequence number of the packet used to initialize
@@ -272,6 +300,36 @@ class MagL1a:
         self.most_recent_sequence = starting_packet.src_seq_ctr
         self.update_compression_array(starting_packet, self.vectors.shape[0])
 
+    @property
+    def vectors(self) -> np.ndarray:
+        """
+        Concatenate the internal vector list into a numpy array.
+
+        If the array has already been created, return the cached version.
+
+        Returns
+        -------
+        np.ndarray
+            Array of vectors with shape (n, 5) where n is the number of vectors,
+            and each vector is (x, y, z, range, time).
+        """
+        if self._vector_cache is None:
+            self._vector_cache = np.concatenate(self._vector_list, axis=0)
+        return self._vector_cache
+
+    @property
+    def compression_flags(self) -> np.ndarray:
+        """
+        Return the compression flags array.
+
+        Returns
+        -------
+        np.ndarray
+            Array of compression flags with shape (n, 2) where n is the number of
+            vectors, and each entry is (is_compressed, compression_width).
+        """
+        return np.concatenate(self._compression_flags_list, axis=0)
+
     def append_vectors(
         self, additional_vectors: np.ndarray, packet_properties: MagL1aPacketProperties
     ) -> None:
@@ -285,9 +343,12 @@ class MagL1a:
         packet_properties : MagL1aPacketProperties
             Additional vector definition to add to the l0_packets dictionary.
         """
+        self._vector_list.append(additional_vectors)
+        # Invalidate the cache
+        self._vector_cache = None
+
         vector_sequence = packet_properties.src_seq_ctr
 
-        self.vectors = np.concatenate([self.vectors, additional_vectors])
         start_time = np.int64(met_to_ttj2000ns(packet_properties.shcoarse))
         self.packet_definitions[start_time] = packet_properties
 
@@ -322,10 +383,7 @@ class MagL1a:
             [packet_properties.compression, packet_properties.compression_width],
             dtype=np.int8,
         )
-        if self.compression_flags is None:
-            self.compression_flags = new_flags
-        else:
-            self.compression_flags = np.concatenate([self.compression_flags, new_flags])
+        self._compression_flags_list.append(new_flags)
 
     @staticmethod
     def calculate_vector_time(

imap_processing/mag/l1c/interpolation_methods.py

@@ -33,9 +33,11 @@ def remove_invalid_output_timestamps(
     numpy.ndarray
         All valid output timestamps where there exists input data.
     """
-    if input_timestamps[0] > output_timestamps[0]:
-        # Chop data where we don't have input timestamps to interpolate
-        output_timestamps = output_timestamps[output_timestamps >= input_timestamps[0]]
+    # Chop data where we don't have input timestamps to interpolate
+    output_timestamps = output_timestamps[
+        (output_timestamps >= input_timestamps[0])
+        & (output_timestamps <= input_timestamps[-1])
+    ]
     return output_timestamps
 
 
@@ -45,7 +47,8 @@ def linear(
     output_timestamps: np.ndarray,
     input_rate: VecSec | None = None,
     output_rate: VecSec | None = None,
-) -> np.ndarray:
+    extrapolate: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Linear interpolation of input vectors to output timestamps.
 
@@ -63,6 +66,9 @@ def linear(
         Not required for this interpolation method.
     output_rate : VecSec, optional
         Not required for this interpolation method.
+    extrapolate : bool, optional
+        Whether to allow extrapolation of output timestamps outside the range of input
+        timestamps. Default is False.
 
     Returns
     -------
@@ -70,9 +76,12 @@ def linear(
         Interpolated vectors of shape (m, 3) where m is equal to the number of output
         timestamps. Contains x, y, z components of the vector.
     """
-    # TODO: Remove invalid timestamps using remove_invalid_output_timestamps
+    if not extrapolate:
+        output_timestamps = remove_invalid_output_timestamps(
+            input_timestamps, output_timestamps
+        )
     spline = make_interp_spline(input_timestamps, input_vectors, k=1)
-    return spline(output_timestamps)
+    return output_timestamps, spline(output_timestamps)
 
 
 def quadratic(
@@ -81,7 +90,7 @@ def quadratic(
     output_timestamps: np.ndarray,
     input_rate: VecSec | None = None,
     output_rate: VecSec | None = None,
-) -> np.ndarray:
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Quadratic interpolation of input vectors to output timestamps.
 
@@ -106,8 +115,11 @@ def quadratic(
         Interpolated vectors of shape (m, 3) where m is equal to the number of output
         timestamps. Contains x, y, z components of the vector.
     """
+    output_timestamps = remove_invalid_output_timestamps(
+        input_timestamps, output_timestamps
+    )
     spline = make_interp_spline(input_timestamps, input_vectors, k=2)
-    return spline(output_timestamps)
+    return output_timestamps, spline(output_timestamps)
 
 
 def cubic(
@@ -116,7 +128,7 @@ def cubic(
     output_timestamps: np.ndarray,
     input_rate: VecSec | None = None,
     output_rate: VecSec | None = None,
-) -> np.ndarray:
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Cubic interpolation of input vectors to output timestamps.
 
@@ -141,8 +153,11 @@ def cubic(
         Interpolated vectors of shape (m, 3) where m is equal to the number of output
         timestamps. Contains x, y, z components of the vector.
     """
+    output_timestamps = remove_invalid_output_timestamps(
+        input_timestamps, output_timestamps
+    )
     spline = make_interp_spline(input_timestamps, input_vectors, k=3)
-    return spline(output_timestamps)
+    return output_timestamps, spline(output_timestamps)
 
 
 def estimate_rate(timestamps: np.ndarray) -> VecSec:
@@ -245,7 +260,7 @@ def linear_filtered(
     output_timestamps: np.ndarray,
     input_rate: VecSec | None = None,
     output_rate: VecSec | None = None,
-) -> np.ndarray:
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Linear filtered interpolation of input vectors to output timestamps.
 
@@ -290,7 +305,7 @@ def quadratic_filtered(
     output_timestamps: np.ndarray,
     input_rate: VecSec | None = None,
     output_rate: VecSec | None = None,
-) -> np.ndarray:
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Quadratic filtered interpolation of input vectors to output timestamps.
 
@@ -317,6 +332,9 @@ def quadratic_filtered(
         Interpolated vectors of shape (m, 3) where m is equal to the number of output
         timestamps. Contains x, y, z components of the vector.
     """
+    output_timestamps = remove_invalid_output_timestamps(
+        input_timestamps, output_timestamps
+    )
     input_filtered, vectors_filtered = cic_filter(
         input_vectors, input_timestamps, output_timestamps, input_rate, output_rate
     )
@@ -329,7 +347,7 @@ def cubic_filtered(
     output_timestamps: np.ndarray,
     input_rate: VecSec | None = None,
     output_rate: VecSec | None = None,
-) -> np.ndarray:
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Cubic filtered interpolation of input vectors to output timestamps.
 
@@ -356,6 +374,9 @@ def cubic_filtered(
         Interpolated vectors of shape (m, 3) where m is equal to the number of output
         timestamps. Contains x, y, z components of the vector.
     """
+    output_timestamps = remove_invalid_output_timestamps(
+        input_timestamps, output_timestamps
+    )
     input_filtered, vectors_filtered = cic_filter(
         input_vectors, input_timestamps, output_timestamps, input_rate, output_rate
     )