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

imap_processing/mag/l1d/mag_l1d_data.py

@@ -1,17 +1,23 @@
 # mypy: disable-error-code="unused-ignore"
 """Data classes for MAG L1D processing."""
 
+import logging
 from dataclasses import InitVar, dataclass
 
 import numpy as np
 import xarray as xr
 
+from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.mag import imap_mag_sdc_configuration_v001 as configuration
 from imap_processing.mag.constants import FILLVAL, DataMode
 from imap_processing.mag.l1c.interpolation_methods import linear
 from imap_processing.mag.l2.mag_l2 import retrieve_matrix_from_l2_calibration
 from imap_processing.mag.l2.mag_l2_data import MagL2L1dBase, ValidFrames
 from imap_processing.spice import spin
 from imap_processing.spice.geometry import frame_transform
+from imap_processing.spice.time import ttj2000ns_to_et, ttj2000ns_to_met
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -57,7 +63,7 @@ class MagL1dConfiguration:
     mago_calibration: np.ndarray
     magi_calibration: np.ndarray
     spin_count_calibration: int
-    quality_flag_threshold: np.float64
+    quality_flag_threshold: float
     spin_average_application_factor: np.float64
     gradiometer_factor: np.ndarray
     apply_gradiometry: bool = True
@@ -144,8 +150,6 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
         truncate the data to exactly 24 hours.
     """
 
-    # TODO Quality flags
-    # TODO generate and output ancillary files
     magi_vectors: np.ndarray
     magi_range: np.ndarray
     magi_epoch: np.ndarray
@@ -165,6 +169,9 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
             The day we are processing, in np.datetime64[D] format. This is used to
             truncate the data to exactly 24 hours.
         """
+        # The main data frame is MAGO, even though we have MAGI data included.
+        self.frame = ValidFrames.MAGO
+
         # set the magnitude before truncating
         self.magnitude = np.zeros(self.vectors.shape[0], dtype=np.float64)  # type: ignore[has-type]
         self.truncate_to_24h(day)
@@ -182,7 +189,7 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
 
         self.vectors = self.apply_spin_offsets(
             self.spin_offsets,
-            self.epoch,
+            self.epoch,  # type: ignore[has-type]
             self.vectors,
             self.config.spin_average_application_factor,
         )
@@ -198,7 +205,11 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
 
         if self.config.apply_gradiometry:
             self.gradiometry_offsets = self.calculate_gradiometry_offsets(
-                self.vectors, self.epoch, self.magi_vectors, self.magi_epoch
+                self.vectors,
+                self.epoch,  # type: ignore[has-type]
+                self.magi_vectors,
+                self.magi_epoch,
+                self.config.quality_flag_threshold,
             )
             self.vectors = self.apply_gradiometry_offsets(
                 self.gradiometry_offsets, self.vectors, self.config.gradiometer_factor
@@ -207,6 +218,54 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
         self.magnitude = MagL2L1dBase.calculate_magnitude(vectors=self.vectors)
         self.is_l1d = True
 
+    def generate_dataset(
+        self,
+        attribute_manager: ImapCdfAttributes,
+        day: np.datetime64,
+    ) -> xr.Dataset:
+        """
+        Generate an xarray dataset from the dataclass.
+
+        This overrides the parent method to conditionally swap MAGO/MAGI data
+        based on the always_output_mago configuration setting.
+
+        Parameters
+        ----------
+        attribute_manager : ImapCdfAttributes
+            CDF attributes object for the correct level.
+        day : np.datetime64
+            The 24 hour day to process, as a numpy datetime format.
+
+        Returns
+        -------
+        xr.Dataset
+            Complete dataset ready to write to CDF file.
+        """
+        always_output_mago = configuration.ALWAYS_OUTPUT_MAGO
+
+        if not always_output_mago:
+            # Swap vectors and epochs to use MAGI data instead of MAGO
+            original_vectors: np.ndarray = self.vectors.copy()
+            original_epoch: np.ndarray = self.epoch.copy()  # type: ignore[has-type]
+            original_range: np.ndarray = self.range.copy()  # type: ignore[has-type]
+
+            self.vectors = self.magi_vectors  # type: ignore[no-redef]
+            self.epoch = self.magi_epoch  # type: ignore[no-redef]
+            self.range = self.magi_range  # type: ignore[no-redef]
+
+            # Call parent generate_dataset method
+            dataset = super().generate_dataset(attribute_manager, day)
+
+            # Restore original vectors for any further processing
+            self.vectors = original_vectors
+            self.epoch = original_epoch
+            self.range = original_range
+        else:
+            # Use MAGO data (default behavior)
+            dataset = super().generate_dataset(attribute_manager, day)
+
+        return dataset
+
     def rotate_frame(self, end_frame: ValidFrames) -> None:
         """
         Rotate the vectors to the desired frame.
@@ -219,15 +278,42 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
         end_frame : ValidFrames
             The frame to rotate to. Should be one of the ValidFrames enum.
         """
+        # Self.frame should refer to the main data in self.vectors, which is MAGO
+        # data. For most frames, MAGO and MAGI are in the same frame, except the
+        # instrument reference frame.
+        if ValidFrames.MAGI in (self.frame, end_frame):
+            raise ValueError(
+                "MAGL1d.frame should never be equal to MAGI frame. If the "
+                "data is in the instrument frame, use MAGO."
+            )
+
         start_frame = self.frame
-        super().rotate_frame(end_frame)
+
+        if self.epoch_et is None:
+            self.epoch_et: np.ndarray = ttj2000ns_to_et(self.epoch)
+            self.magi_epoch_et: np.ndarray = ttj2000ns_to_et(self.magi_epoch)
+
+        self.vectors = frame_transform(
+            self.epoch_et,
+            self.vectors,
+            from_frame=start_frame.value,
+            to_frame=end_frame.value,
+        )
+
+        # If we were in MAGO frame, we need to rotate MAGI vectors from MAGI to
+        # end_frame
+        if start_frame == ValidFrames.MAGO:
+            start_frame = ValidFrames.MAGI
+
         self.magi_vectors = frame_transform(
-            self.magi_epoch,
+            self.magi_epoch_et,
             self.magi_vectors,
             from_frame=start_frame.value,
             to_frame=end_frame.value,
         )
 
+        self.frame = end_frame
+
     def _calibrate_and_offset_vectors(
         self,
         mago_calibration: np.ndarray,
@@ -312,7 +398,7 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
             The offset vector, shape (4,) where the last element is unchanged.
         """
         # Offsets are in shape (sensor, range, axis)
-        updated_vector = input_vector.copy().astype(np.int64)
+        updated_vector = input_vector.copy()
         rng = int(input_vector[3])
         x_y_z = input_vector[:3]
         updated_vector[:3] = x_y_z - offsets[int(is_magi), rng, :]
@@ -348,64 +434,132 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
                 "Spin offsets can only be calculated in NORM mode and SRF frame."
             )
 
-        # TODO: get the spin numbers which correspond to the epoch values for output
-        sc_spin_phase: np.ndarray = spin.get_spacecraft_spin_phase(self.epoch)  # type: ignore
+        epoch_met = ttj2000ns_to_met(self.epoch)
+        sc_spin_phase = spin.get_spacecraft_spin_phase(epoch_met)
         # mark vectors as nan where they are nan in sc_spin_phase
         vectors = self.vectors.copy().astype(np.float64)
 
         vectors[np.isnan(sc_spin_phase), :] = np.nan
 
-        # TODO: currently fully skipping spins with no valid data (not including
-        #  them in the averaging OR IN SPIN COUNTING!) is this correct?
-
         # first timestamp where spin phase is less than the previous value
         # this is when the spin crosses zero
         spin_starts = np.where(np.diff(sc_spin_phase) < 0)[0] + 1
 
-        # if the value switches from nan to a number, that is also a spin start (for an
-        # invalid spin)
-        nan_to_number = (
-            np.where(np.isnan(sc_spin_phase[:-1]) & ~np.isnan(sc_spin_phase[1:]))[0] + 1
-        )
+        # if the value switches from nan to a number, or from a number to nan, that
+        # is also a spin start
+        nan_to_number = np.where(np.diff(np.isnan(sc_spin_phase)) != 0)[0] + 1
 
         # find the places spins start while skipping over invalid or missing data
         # (marked as nan by get_spacecraft_spin_phase)
         spin_starts = np.sort(np.concatenate((spin_starts, nan_to_number)))
 
+        # Get the expected spin period from the spin table
+        # Convert to nanoseconds to match epoch
+        spin_data = spin.get_spin_data()
+        # Use the median spin period as the expected value
+        expected_spin = np.median(spin_data["spin_period_sec"]) * 1e9
+
+        paired_nans = nan_to_number.reshape(-1, 2)
+
+        for start_of_gap, end_of_gap in paired_nans:
+            # in nan_to_number, we have the start and end for every nan gap
+            # if this gap spans more than 1 spin period, we need to insert
+            # additional spin_starts into spin_starts.
+
+            gap_start_time = self.epoch[start_of_gap]
+            gap_end_time = self.epoch[end_of_gap]
+
+            # Calculate the number of spins in this gap
+            number_of_spins = int((gap_end_time - gap_start_time) // expected_spin)
+            if number_of_spins > 1:
+                # Insert new spin starts into spin_starts
+                for i in range(1, number_of_spins):
+                    estimated_start = gap_start_time + i * expected_spin
+                    new_spin_index = (np.abs(self.epoch - estimated_start)).argmin()
+
+                    spin_starts = np.append(spin_starts, new_spin_index)
+
+        # Now spin_starts contains all the indices where spins begin, including
+        # estimating skipped or missing spins.
+        spin_starts = np.sort(spin_starts)
+
         chunk_start = 0
         offset_epochs = []
-        x_avg = []
-        y_avg = []
+        x_avg_calcs: list[np.float64] = []
+        y_avg_calcs: list[np.float64] = []
+        validity_start_times = []
+        validity_end_times = []
+        start_spin_counters = []
+        end_spin_counters = []
+
         while chunk_start < len(spin_starts):
             # Take self.spin_count_calibration number of spins and put them into a chunk
             chunk_indices = spin_starts[
                 chunk_start : chunk_start + self.config.spin_count_calibration + 1
             ]
-            chunk_start = chunk_start + self.config.spin_count_calibration
-
-            # If we are in the end of the chunk, just grab all remaining data
-            if chunk_start >= len(spin_starts):
-                chunk_indices = np.append(chunk_indices, len(self.epoch))
+            chunk_start_idx = chunk_start
 
             chunk_vectors = self.vectors[chunk_indices[0] : chunk_indices[-1]]
             chunk_epoch = self.epoch[chunk_indices[0] : chunk_indices[-1]]
 
+            # Check if more than half of the chunk data is NaN before processing
+            x_valid_count: int = int(np.sum(~np.isnan(chunk_vectors[:, 0])))
+            y_valid_count: int = int(np.sum(~np.isnan(chunk_vectors[:, 1])))
+            total_points = len(chunk_vectors)
+
             # average the x and y axes (z is fixed, as the spin axis)
-            # TODO: is z the correct axis here?
             avg_x = np.nanmean(chunk_vectors[:, 0])
             avg_y = np.nanmean(chunk_vectors[:, 1])
 
+            # Skip chunk if more than half of x or y data is NaN, or if we have less
+            # than half a spin.
+            # in this case, we should reuse the previous averages.
+            if (
+                x_valid_count <= total_points / 2
+                or y_valid_count <= total_points / 2
+                or total_points <= self.config.spin_count_calibration / 2
+            ):
+                avg_x = x_avg_calcs[-1] if x_avg_calcs else np.float64(FILLVAL)
+                avg_y = y_avg_calcs[-1] if y_avg_calcs else np.float64(FILLVAL)
+
             if not np.isnan(avg_x) and not np.isnan(avg_y):
                 offset_epochs.append(chunk_epoch[0])
-                x_avg.append(avg_x)
-                y_avg.append(avg_y)
+                x_avg_calcs.append(avg_x)
+                y_avg_calcs.append(avg_y)
+
+                # Add validity time range for this chunk
+                validity_start_times.append(chunk_epoch[0])
+                validity_end_times.append(chunk_epoch[-1])
+
+                # Add spin counter information
+                start_spin_counters.append(chunk_start_idx)
+                end_spin_counters.append(
+                    min(
+                        chunk_start_idx + self.config.spin_count_calibration - 1,
+                        len(spin_starts) - 1,
+                    )
+                )
+
+            chunk_start = chunk_start + self.config.spin_count_calibration
 
         spin_epoch_dataarray = xr.DataArray(np.array(offset_epochs))
 
         spin_offsets = xr.Dataset(coords={"epoch": spin_epoch_dataarray})
 
-        spin_offsets["x_offset"] = xr.DataArray(np.array(x_avg), dims=["epoch"])
-        spin_offsets["y_offset"] = xr.DataArray(np.array(y_avg), dims=["epoch"])
+        spin_offsets["x_offset"] = xr.DataArray(np.array(x_avg_calcs), dims=["epoch"])
+        spin_offsets["y_offset"] = xr.DataArray(np.array(y_avg_calcs), dims=["epoch"])
+        spin_offsets["validity_start_time"] = xr.DataArray(
+            np.array(validity_start_times), dims=["epoch"]
+        )
+        spin_offsets["validity_end_time"] = xr.DataArray(
+            np.array(validity_end_times), dims=["epoch"]
+        )
+        spin_offsets["start_spin_counter"] = xr.DataArray(
+            np.array(start_spin_counters), dims=["epoch"]
+        )
+        spin_offsets["end_spin_counter"] = xr.DataArray(
+            np.array(end_spin_counters), dims=["epoch"]
+        )
 
         return spin_offsets
 
@@ -462,7 +616,7 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
         if spin_offsets is None:
             raise ValueError("No spin offsets calculated to apply.")
 
-        output_vectors = np.full(vectors.shape, FILLVAL, dtype=np.int64)
+        output_vectors = np.full(vectors.shape, FILLVAL, dtype=np.float64)
 
         for index in range(spin_offsets["epoch"].data.shape[0] - 1):
             timestamp = spin_offsets["epoch"].data[index]
@@ -483,7 +637,6 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
             if not np.any(mask):
                 continue
 
-            # TODO: should vectors be a float?
             x_offset = (
                 spin_offsets["x_offset"].data[index] * spin_average_application_factor
             )
@@ -504,6 +657,7 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
         mago_epoch: np.ndarray,
         magi_vectors: np.ndarray,
         magi_epoch: np.ndarray,
+        quality_flag_threshold: float = np.inf,
     ) -> xr.Dataset:
         """
         Calculate the gradiometry offsets between MAGo and MAGi.
@@ -525,6 +679,10 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
             The MAGi vectors, shape (N, 3).
         magi_epoch : np.ndarray
             The MAGi epoch values, shape (N,).
+        quality_flag_threshold : np.float64, optional
+            Threshold for quality flags. If the magnitude of gradiometer offset
+            exceeds this threshold, quality flag will be set. Default is np.inf
+            (no quality flags set).
 
         Returns
         -------
@@ -532,6 +690,8 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
             The gradiometer offsets dataset, with variables:
             - epoch: the timestamp of the MAGo data
             - gradiometer_offsets: the offset values (MAGi - MAGo) for each axis
+            - gradiometer_offset_magnitude: magnitude of the offset vector
+            - quality_flags: quality flags (1 if magnitude > threshold, 0 otherwise)
         """
         aligned_magi = linear(
             magi_vectors,
@@ -541,10 +701,20 @@ class MagL1d(MagL2L1dBase):  # type: ignore[misc]
 
         diff = aligned_magi - mago_vectors
 
+        # Calculate magnitude of gradiometer offset for each vector
+        magnitude = np.linalg.norm(diff, axis=1)
+
+        # Set quality flags: 0 = good data (below threshold), 1 = bad data
+        quality_flags = (magnitude > quality_flag_threshold).astype(int)
+
         grad_epoch = xr.DataArray(mago_epoch, dims=["epoch"])
         direction = xr.DataArray(["x", "y", "z"], dims=["axis"])
         grad_ds = xr.Dataset(coords={"epoch": grad_epoch, "direction": direction})
         grad_ds["gradiometer_offsets"] = xr.DataArray(diff, dims=["epoch", "direction"])
+        grad_ds["gradiometer_offset_magnitude"] = xr.DataArray(
+            magnitude, dims=["epoch"]
+        )
+        grad_ds["quality_flags"] = xr.DataArray(quality_flags, dims=["epoch"])
 
         return grad_ds
 

imap_processing/mag/l2/mag_l2.py

@@ -90,6 +90,7 @@ def mag_l2(
     )
     # level 2 vectors don't include range
     vectors = cal_vectors[:, :3]
+    instrument_frame = ValidFrames.MAGO if always_output_mago else ValidFrames.MAGI
 
     l2_data = MagL2(
         vectors=vectors,
@@ -101,6 +102,7 @@ def mag_l2(
         data_mode=mode,
         offsets=offsets_dataset["offsets"].data,
         timedelta=offsets_dataset["timedeltas"].data,
+        frame=instrument_frame,
     )
 
     attributes = ImapCdfAttributes()

imap_processing/mag/l2/mag_l2_data.py

@@ -13,16 +13,19 @@ from imap_processing.spice.geometry import SpiceFrame, frame_transform
 from imap_processing.spice.time import (
     et_to_ttj2000ns,
     str_to_et,
+    ttj2000ns_to_et,
 )
 
 
 class ValidFrames(Enum):
     """SPICE reference frames for output."""
 
-    MAG = SpiceFrame.IMAP_MAG
+    MAGO = SpiceFrame.IMAP_MAG_O
+    MAGI = SpiceFrame.IMAP_MAG_I
     DSRF = SpiceFrame.IMAP_DPS
     SRF = SpiceFrame.IMAP_SPACECRAFT
-    # TODO: include RTN and GSE as valid frames
+    GSE = SpiceFrame.IMAP_GSE
+    RTN = SpiceFrame.IMAP_RTN
 
 
 @dataclass(kw_only=True)
@@ -54,7 +57,10 @@ class MagL2L1dBase:
         Quality bitmask for each vector. Should be of length n. Copied from offset
         file in L2, marked as good always in L1D.
     frame:
-        The reference frame of the input vectors. Starts as the MAG instrument frame.
+        The reference frame of the input vectors. Defaults to the MAGO instrument frame.
+    epoch_et: np.ndarray
+        The epoch timestamps converted to ET format. Used for frame transformations.
+        Calculated on first use and then saved. Should not be passed in.
     """
 
     vectors: np.ndarray
@@ -65,7 +71,8 @@ class MagL2L1dBase:
     quality_bitmask: np.ndarray
     data_mode: DataMode
     magnitude: np.ndarray = field(init=False)
-    frame: ValidFrames = ValidFrames.MAG
+    frame: ValidFrames = ValidFrames.MAGO
+    epoch_et: np.ndarray | None = field(init=False, default=None)
 
     def generate_dataset(
         self,
@@ -301,8 +308,10 @@ class MagL2L1dBase:
             The frame to rotate the data to. Must be one of the ValidFrames enum
             values.
         """
+        if self.epoch_et is None:
+            self.epoch_et = ttj2000ns_to_et(self.epoch)
         self.vectors = frame_transform(
-            self.epoch,
+            self.epoch_et,
             self.vectors,
             from_frame=self.frame.value,
             to_frame=end_frame.value,

imap_processing/quality_flags.py

@@ -37,12 +37,13 @@ class ENAFlags(FlagNameMixin):
     BADSPIN = 2**2  # bit 2, Bad spin
 
 
-class ImapDEUltraFlags(FlagNameMixin):
+class ImapDEOutliersUltraFlags(FlagNameMixin):
     """IMAP Ultra flags."""
 
     NONE = CommonFlags.NONE
     FOV = 2**0  # bit 0
     PHCORR = 2**1  # bit 1
+    COINPH = 2**2  # bit 4 # Event validity
 
 
 class ImapHkUltraFlags(FlagNameMixin):
@@ -75,6 +76,21 @@ class ImapRatesUltraFlags(FlagNameMixin):
     PARTIALSPIN = 2**2  # bit 2
 
 
+class ImapDEScatteringUltraFlags(FlagNameMixin):
+    """IMAP Ultra Scattering flags."""
+
+    NONE = CommonFlags.NONE
+    ABOVE_THRESHOLD = 2**0  # bit 0
+    NAN_PHI_OR_THETA = 2**1  # bit 1
+
+
+class ImapPSETUltraFlags(FlagNameMixin):
+    """IMAP Ultra Rates flags."""
+
+    NONE = CommonFlags.NONE
+    EARTH_FOV = 2**0  # bit 0
+
+
 class ImapInstrumentUltraFlags(FlagNameMixin):
     """IMAP Ultra flags using other instruments."""
 
@@ -123,3 +139,9 @@ class SWAPIFlags(
     SCEM_V_ST = 2**12  # bit 12
     SCEM_I_ST = 2**13  # bit 13
     SCEM_INT_ST = 2**14  # bit 14
+
+
+class GLOWSL1bFlags(FlagNameMixin):
+    """Glows L1b flags."""
+
+    NONE = CommonFlags.NONE