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

imap_processing/mag/l2/mag_l2.py

@@ -6,8 +6,7 @@ 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 DataMode
-from imap_processing.mag.l1b.mag_l1b import calibrate_vector
-from imap_processing.mag.l2.mag_l2_data import MagL2
+from imap_processing.mag.l2.mag_l2_data import MagL2, ValidFrames
 
 
 def mag_l2(
@@ -77,8 +76,8 @@ def mag_l2(
     always_output_mago = configuration.ALWAYS_OUTPUT_MAGO
 
     # TODO Check that the input file matches the offsets file
-    # if not np.array_equal(input_data["epoch"].data, offsets_dataset["epoch"].data):
-    #     raise ValueError("Input file and offsets file must have the same timestamps.")
+    if not np.array_equal(input_data["epoch"].data, offsets_dataset["epoch"].data):
+        raise ValueError("Input file and offsets file must have the same timestamps.")
 
     day: np.datetime64 = day_to_process.astype("datetime64[D]")
 
@@ -86,29 +85,35 @@ def mag_l2(
         calibration_dataset, day, always_output_mago
     )
 
-    vectors = np.apply_along_axis(
-        func1d=calibrate_vector,
-        axis=1,
-        arr=input_data["vectors"].data,
-        calibration_matrix=calibration_matrix,
+    cal_vectors = MagL2.apply_calibration(
+        vectors=input_data["vectors"].data, calibration_matrix=calibration_matrix
     )
-
-    input_data = MagL2(
-        vectors[:, :3],  # level 2 vectors don't include range
-        input_data["epoch"].data,
-        input_data["vectors"].data[:, 3],
-        {},
-        np.zeros(len(input_data["epoch"].data)),
-        np.zeros(len(input_data["epoch"].data)),
-        mode,
+    # level 2 vectors don't include range
+    vectors = cal_vectors[:, :3]
+
+    l2_data = MagL2(
+        vectors=vectors,
+        epoch=input_data["epoch"].data,
+        range=input_data["vectors"].data[:, 3],
+        global_attributes={},
+        quality_flags=offsets_dataset["quality_flag"].data,
+        quality_bitmask=offsets_dataset["quality_bitmask"].data,
+        data_mode=mode,
         offsets=offsets_dataset["offsets"].data,
         timedelta=offsets_dataset["timedeltas"].data,
     )
+
     attributes = ImapCdfAttributes()
     attributes.add_instrument_global_attrs("mag")
-    # temporarily point to l1c
     attributes.add_instrument_variable_attrs("mag", "l2")
-    return [input_data.generate_dataset(attributes, day)]
+
+    # Rotate from the MAG frame into the SRF frame
+    l2_data.rotate_frame(ValidFrames.SRF)
+    imap_srf = l2_data.generate_dataset(attributes, day)
+    l2_data.rotate_frame(ValidFrames.DSRF)
+    imap_dsrf = l2_data.generate_dataset(attributes, day)
+
+    return [imap_dsrf, imap_srf]
 
 
 def retrieve_matrix_from_l2_calibration(

imap_processing/mag/l2/mag_l2_data.py

@@ -7,36 +7,43 @@ import numpy as np
 import xarray as xr
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
-from imap_processing.mag.constants import DataMode
+from imap_processing.mag.constants import FILLVAL, DataMode
+from imap_processing.mag.l1b.mag_l1b import calibrate_vector
+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."""
 
-    dsrf = "dsrf"
-    srf = "srf"
-    rtn = "rtn"
-    gse = "gse"
+    MAG = SpiceFrame.IMAP_MAG
+    DSRF = SpiceFrame.IMAP_DPS
+    SRF = SpiceFrame.IMAP_SPACECRAFT
+    GSE = SpiceFrame.IMAP_GSE
+    RTN = SpiceFrame.IMAP_RTN
 
 
-@dataclass
-class MagL2:
+@dataclass(kw_only=True)
+class MagL2L1dBase:
     """
-    Dataclass for MAG L2 data.
+    Base class for MAG L2 and L1D data.
 
-    Since L2 and L1D should have the same structure, this can be used for either level.
+    Since these two data levels output identical files, and share some methods, this
+    superclass captures the tools in common, while allowing each subclass to define
+    individual attributes and algorithms.
 
-    Some of the methods are also static, so they can be used in i-ALiRT processing.
+    May also be extended for I-ALiRT.
 
     Attributes
     ----------
     vectors: np.ndarray
         Magnetic field vectors of size (n, 3) where n is the number of vectors.
-        Describes (x, y, z) components of the magnetic field.
+        Describes (x, y, z) components of the magnetic field. This field is the output
+        vectors, which are nominally from the MAGo sensor.
     epoch: np.ndarray
         Time of each vector in J2000 seconds. Should be of length n.
     range: np.ndarray
@@ -48,10 +55,11 @@ class MagL2:
     quality_bitmask: np.ndarray
         Quality bitmask for each vector. Should be of length n. Copied from offset
         file in L2, marked as good always in L1D.
-    magnitude: np.ndarray
-        Magnitude of each vector. Should be of length n. Calculated from L2 vectors.
-    is_l1d: bool
-        Flag to indicate if the data is L1D. Defaults to False.
+    frame:
+        The reference frame of the input vectors. Starts as the MAG 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
@@ -62,120 +70,13 @@ class MagL2:
     quality_bitmask: np.ndarray
     data_mode: DataMode
     magnitude: np.ndarray = field(init=False)
-    is_l1d: bool = False
-    offsets: InitVar[np.ndarray] = None
-    timedelta: InitVar[np.ndarray] = None
-
-    def __post_init__(self, offsets: np.ndarray, timedelta: np.ndarray) -> None:
-        """
-        Calculate the magnitude of the vectors after initialization.
-
-        Parameters
-        ----------
-        offsets : np.ndarray
-            Offsets to apply to the vectors. Should be of shape (n, 3) where n is the
-            number of vectors.
-        timedelta : np.ndarray
-            Time deltas to shift the timestamps by. Should be of length n.
-            Given in seconds.
-        """
-        if offsets is not None:
-            self.vectors = self.apply_offsets(self.vectors, offsets)
-        if timedelta is not None:
-            self.epoch = self.shift_timestamps(self.epoch, timedelta)
-
-        self.magnitude = self.calculate_magnitude(self.vectors)
-
-    @staticmethod
-    def calculate_magnitude(
-        vectors: np.ndarray,
-    ) -> np.ndarray:
-        """
-        Given a list of vectors (x, y, z), calculate the magnitude of each vector.
-
-        For an input list of vectors of size (n, 3) returns a list of magnitudes of
-        size (n,).
-
-        Parameters
-        ----------
-        vectors : np.ndarray
-            Array of vectors to calculate the magnitude of.
-
-        Returns
-        -------
-        np.ndarray
-            Array of magnitudes of the input vectors.
-        """
-        return np.linalg.norm(vectors, axis=1)  # type: ignore
-
-    @staticmethod
-    def apply_offsets(vectors: np.ndarray, offsets: np.ndarray) -> np.ndarray:
-        """
-        Apply the offsets to the vectors by adding them together.
-
-        These offsets are used to shift the vectors in the x, y, and z directions.
-        They can either be provided through a custom offsets datafile, or calculated
-        using a gradiometry algorithm.
-
-        Parameters
-        ----------
-        vectors : np.ndarray
-            Array of vectors to apply the offsets to. Should be of shape (n, 3) where n
-            is the number of vectors.
-        offsets : np.ndarray
-            Array of offsets to apply to the vectors. Should be of shape (n, 3) where n
-            is the number of vectors.
-
-        Returns
-        -------
-        np.ndarray
-            Array of vectors with offsets applied. Should be of shape (n, 3).
-        """
-        if vectors.shape[0] != offsets.shape[0]:
-            raise ValueError("Vectors and offsets must have the same length.")
-
-        offset_vectors: np.ndarray = vectors[:, :3] + offsets
-
-        # TODO: CDF files don't have NaNs. Emailed MAG to ask what this will look like.
-        # Any values where offsets is nan must also be nan
-        offset_vectors[np.isnan(offsets).any(axis=1)] = np.nan
-
-        return offset_vectors
-
-    @staticmethod
-    def shift_timestamps(epoch: np.ndarray, timedelta: np.ndarray) -> np.ndarray:
-        """
-        Shift the timestamps by the given timedelta.
-
-        If timedelta is positive, the epochs are shifted forward in time.
-
-        Parameters
-        ----------
-        epoch : np.ndarray
-            Array of timestamps to shift. Should be of length n.
-        timedelta : np.ndarray
-            Array of time deltas to shift the timestamps by. Should be the same length
-            as epoch. Given in seconds.
-
-        Returns
-        -------
-        np.ndarray
-            Shifted timestamps.
-        """
-        if epoch.shape[0] != timedelta.shape[0]:
-            raise ValueError(
-                "Input Epoch and offsets timedeltas must be the same length."
-            )
-
-        timedelta_ns = timedelta * 1e9
-        shifted_timestamps = epoch + timedelta_ns
-        return shifted_timestamps
+    frame: ValidFrames = ValidFrames.MAG
+    epoch_et: np.ndarray | None = field(init=False, default=None)
 
     def generate_dataset(
         self,
         attribute_manager: ImapCdfAttributes,
         day: np.datetime64,
-        frame: ValidFrames = ValidFrames.dsrf,
     ) -> xr.Dataset:
         """
         Generate an xarray dataset from the dataclass.
@@ -189,8 +90,6 @@ class MagL2:
             CDF attributes object for the correct level.
         day : np.datetime64
          The 24 hour day to process, as a numpy datetime format.
-        frame : ValidFrames
-            SPICE reference frame to rotate the data into.
 
         Returns
         -------
@@ -199,7 +98,9 @@ class MagL2:
         """
         self.truncate_to_24h(day)
 
-        logical_source_id = f"imap_mag_l2_{self.data_mode.value.lower()}-{frame.name}"
+        logical_source_id = (
+            f"imap_mag_l2_{self.data_mode.value.lower()}-{self.frame.name.lower()}"
+        )
         direction = xr.DataArray(
             np.arange(3),
             name="direction",
@@ -242,8 +143,8 @@ class MagL2:
         )
 
         quality_bitmask = xr.DataArray(
-            self.quality_flags,
-            name="quality_flags",
+            self.quality_bitmask,
+            name="quality_bitmask",
             dims=["epoch"],
             attrs=attribute_manager.get_variable_attributes("qf"),
         )
@@ -298,7 +199,6 @@ class MagL2:
         """
         if self.epoch.shape[0] != self.vectors.shape[0]:
             raise ValueError("Timestamps and vectors are not the same shape!")
-
         start_timestamp_j2000 = et_to_ttj2000ns(str_to_et(str(timestamp)))
         end_timestamp_j2000 = et_to_ttj2000ns(
             str_to_et(str(timestamp + np.timedelta64(1, "D")))
@@ -313,3 +213,172 @@ class MagL2:
         self.magnitude = self.magnitude[day_start_index:day_end_index]
         self.quality_flags = self.quality_flags[day_start_index:day_end_index]
         self.quality_bitmask = self.quality_bitmask[day_start_index:day_end_index]
+
+    @staticmethod
+    def calculate_magnitude(
+        vectors: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Given a list of vectors (x, y, z), calculate the magnitude of each vector.
+
+        For an input list of vectors of size (n, 3) returns a list of magnitudes of
+        size (n,).
+
+        Parameters
+        ----------
+        vectors : np.ndarray
+            Array of vectors to calculate the magnitude of.
+
+        Returns
+        -------
+        np.ndarray
+            Array of magnitudes of the input vectors.
+        """
+        return np.linalg.norm(vectors, axis=1)
+
+    @staticmethod
+    def apply_calibration(
+        vectors: np.ndarray, calibration_matrix: np.ndarray
+    ) -> np.ndarray:
+        """
+        Apply the calibration matrix to the vectors.
+
+        This works by repeatedly calling the function calibrate_vector on the vectors
+        input.
+
+        Parameters
+        ----------
+        vectors : np.ndarray
+            Array of vectors to apply the calibration to, including x,y,z and range.
+            Should be of shape (n, 4) where n is the number of vectors.
+        calibration_matrix : np.ndarray
+            Calibration matrix to apply to the vectors. Should be of shape (3, 3, 4).
+
+        Returns
+        -------
+        np.ndarray
+            Array of calibrated vectors. Should be of shape (n, 4).
+        """
+        calibrated_vectors = np.apply_along_axis(
+            func1d=calibrate_vector,
+            axis=1,
+            arr=vectors,
+            calibration_matrix=calibration_matrix,
+        )
+
+        return calibrated_vectors
+
+    @staticmethod
+    def shift_timestamps(epoch: np.ndarray, timedelta: np.ndarray) -> np.ndarray:
+        """
+        Shift the timestamps by the given timedelta.
+
+        If timedelta is positive, the epochs are shifted forward in time.
+
+        Parameters
+        ----------
+        epoch : np.ndarray
+            Array of timestamps to shift. Should be of length n.
+        timedelta : np.ndarray
+            Array of time deltas to shift the timestamps by. Should be the same length
+            as epoch. Given in seconds.
+
+        Returns
+        -------
+        np.ndarray
+            Shifted timestamps.
+        """
+        if epoch.shape[0] != timedelta.shape[0]:
+            raise ValueError(
+                "Input Epoch and offsets timedeltas must be the same length."
+            )
+
+        timedelta_ns = timedelta * 1e9
+        shifted_timestamps = epoch + timedelta_ns
+        return shifted_timestamps
+
+    def rotate_frame(self, end_frame: ValidFrames) -> None:
+        """
+        Rotate the vector data in the class to the output frame.
+
+        Parameters
+        ----------
+        end_frame : ValidFrames
+            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_et,
+            self.vectors,
+            from_frame=self.frame.value,
+            to_frame=end_frame.value,
+        )
+        self.frame = end_frame
+
+
+@dataclass(kw_only=True)
+class MagL2(MagL2L1dBase):
+    """
+    Dataclass for MAG L2 data.
+
+    Since L2 and L1D should have the same structure, this can be used for either level.
+
+    Some of the methods are also static, so they can be used in i-ALiRT processing.
+    """
+
+    offsets: InitVar[np.ndarray] = None
+    timedelta: InitVar[np.ndarray] = None
+
+    def __post_init__(self, offsets: np.ndarray, timedelta: np.ndarray) -> None:
+        """
+        Calculate the magnitude of the vectors after initialization.
+
+        Parameters
+        ----------
+        offsets : np.ndarray
+            Offsets to apply to the vectors. Should be of shape (n, 3) where n is the
+            number of vectors.
+        timedelta : np.ndarray
+            Time deltas to shift the timestamps by. Should be of length n.
+            Given in seconds.
+        """
+        if offsets is not None:
+            self.vectors = self.apply_offsets(self.vectors, offsets)
+        if timedelta is not None:
+            self.epoch = self.shift_timestamps(self.epoch, timedelta)
+
+        self.magnitude = self.calculate_magnitude(self.vectors)
+
+    @staticmethod
+    def apply_offsets(vectors: np.ndarray, offsets: np.ndarray) -> np.ndarray:
+        """
+        Apply the offsets to the vectors by adding them together.
+
+        These offsets are used to shift the vectors in the x, y, and z directions.
+        They can either be provided through a custom offsets datafile, or calculated
+        using a gradiometry algorithm.
+
+        Parameters
+        ----------
+        vectors : np.ndarray
+            Array of vectors to apply the offsets to. Should be of shape (n, 3) where n
+            is the number of vectors.
+        offsets : np.ndarray
+            Array of offsets to apply to the vectors. Should be of shape (n, 3) where n
+            is the number of vectors.
+
+        Returns
+        -------
+        np.ndarray
+            Array of vectors with offsets applied. Should be of shape (n, 3).
+        """
+        if vectors.shape[0] != offsets.shape[0]:
+            raise ValueError("Vectors and offsets must have the same length.")
+
+        offset_vectors: np.ndarray = vectors + offsets
+
+        # Any values where offsets is FILLVAL must also be FILLVAL
+        offset_vectors[(offsets == FILLVAL).any(axis=1), :] = FILLVAL
+        return offset_vectors

imap_processing/quality_flags.py

@@ -37,6 +37,14 @@ class ENAFlags(FlagNameMixin):
     BADSPIN = 2**2  # bit 2, Bad spin
 
 
+class ImapDEOutliersUltraFlags(FlagNameMixin):
+    """IMAP Ultra flags."""
+
+    NONE = CommonFlags.NONE
+    FOV = 2**0  # bit 0
+    PHCORR = 2**1  # bit 1
+
+
 class ImapHkUltraFlags(FlagNameMixin):
     """IMAP Ultra flags."""
 
@@ -53,14 +61,32 @@ class ImapAttitudeUltraFlags(FlagNameMixin):
     NONE = CommonFlags.NONE
     SPINRATE = 2**0  # bit 0
     AUXMISMATCH = 2**1  # bit 1 # aux packet does not match Universal Spin Table
+    SPINPHASE = 2**2  # bit 2 # spin phase flagged by Universal Spin Table
+    SPINPERIOD = 2**3  # bit 3 # spin period flagged by Universal Spin Table
 
 
 class ImapRatesUltraFlags(FlagNameMixin):
     """IMAP Ultra Rates flags."""
 
     NONE = CommonFlags.NONE
-    ZEROCOUNTS = 2**0  # bit 0
-    HIGHRATES = 2**1  # bit 1
+    HIGHRATES = 2**0  # bit 0
+    FIRSTSPIN = 2**1  # bit 1
+    LASTSPIN = 2**2  # bit 2
+    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 ImapInstrumentUltraFlags(FlagNameMixin):
+    """IMAP Ultra flags using other instruments."""
+
+    NONE = CommonFlags.NONE
 
 
 class ImapLoFlags(FlagNameMixin):