Mag L2 first pass (IMAP-Science-Operations-Center#1551)

maxinelasp · web-flow · commit 1494922f7c91 · 2025-04-04T02:49:09.000-06:00
* L2 first pass

* Setting up empty mag L2
diff --git a/imap_processing/cdf/config/imap_mag_global_cdf_attrs.yaml b/imap_processing/cdf/config/imap_mag_global_cdf_attrs.yaml
@@ -83,3 +83,18 @@ imap_mag_l1c_norm-magi:
     Data_type: L1C_norm-magi>Level 1C MAGi normal rate
     Logical_source: imap_mag_l1c_norm-magi
     Logical_source_description: IMAP Mission MAGi Normal Rate Instrument Level-1C Data.
+
+imap_mag_l2_norm-dsrf:
+  <<: *default
+  Data_type: L2_norm-dsrf>Level 2 normal rate data in DSRF
+  Logical_source: imap_mag_l2_norm-dsrf
+  Logical_source_description: IMAP Mission Normal Rate Instrument Level-2 Data in
+    Despun Spacecraft Reference Frame.
+
+
+imap_mag_l2_burst-dsrf:
+  <<: *default
+  Data_type: L2_burst-dsrf>Level 2 burst rate data in DSRF
+  Logical_source: imap_mag_l2_burst-dsrf
+  Logical_source_description: IMAP Mission Burst Rate Instrument Level-2 Data in
+    Despun Spacecraft Reference Frame.
diff --git a/imap_processing/cli.py b/imap_processing/cli.py
@@ -55,6 +55,7 @@
 from imap_processing.mag.l1a.mag_l1a import mag_l1a
 from imap_processing.mag.l1b.mag_l1b import mag_l1b
 from imap_processing.mag.l1c.mag_l1c import mag_l1c
+from imap_processing.mag.l2.mag_l2 import mag_l2
 from imap_processing.swapi.l1.swapi_l1 import swapi_l1
 from imap_processing.swapi.l2.swapi_l2 import swapi_l2
 from imap_processing.swe.l1a.swe_l1a import swe_l1a
@@ -830,6 +831,24 @@ def do_processing(
             # Input datasets can be in any order
             datasets = [mag_l1c(input_data[0], input_data[1], self.version)]
 
+        if self.data_level == "l2":
+            # TODO: Overwrite dependencies with versions from offsets file
+            input_data = load_cdf(dependencies[0])
+            # TODO: use ancillary from input
+            calibration_dataset = load_cdf(
+                Path(__file__).parent
+                / "tests"
+                / "mag"
+                / "validation"
+                / "calibration"
+                / "imap_mag_l1b-calibration_20240229_v001.cdf"
+            )
+            # TODO: Test data missing
+            offset_dataset = xr.Dataset()
+            datasets = [
+                mag_l2(calibration_dataset, offset_dataset, input_data, self.version)
+            ]
+
         return datasets
 
 
diff --git a/imap_processing/mag/l1a/mag_l1a.py b/imap_processing/mag/l1a/mag_l1a.py
@@ -297,8 +297,6 @@ def generate_dataset(
         ),
     )
 
-    # TODO: Epoch here refers to the start of the sample. Confirm that this is
-    # what mag is expecting, and if it is, CATDESC needs to be updated.
     epoch_time = xr.DataArray(
         time_data,
         name="epoch",
@@ -338,7 +336,6 @@ def generate_dataset(
         ),
     )
     global_attributes = attribute_manager.get_global_attributes(logical_file_id)
-    # TODO: this method won't work because these values are not in the schema.
     global_attributes["is_mago"] = str(bool(single_file_l1a.is_mago))
     global_attributes["is_active"] = str(bool(single_file_l1a.is_active))
     global_attributes["vectors_per_second"] = (
diff --git a/imap_processing/mag/l1b/__init__.py b/imap_processing/mag/l1b/__init__.py
diff --git a/imap_processing/mag/l2/mag_l2.py b/imap_processing/mag/l2/mag_l2.py
@@ -0,0 +1,77 @@
+"""Module to run MAG L2 processing."""
+
+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.l2.mag_l2_data import MagL2
+
+
+def mag_l2(
+    calibration_dataset: xr.Dataset,
+    offset_dataset: xr.Dataset,
+    input_data: xr.Dataset,
+    data_version: str,
+) -> list[xr.Dataset]:
+    """
+    Complete MAG L2 processing.
+
+    Input data can be burst or normal mode, but MUST match the file in offset_dataset.
+    TODO: retrieve the file from offset_dataset in cli.py.
+
+    Parameters
+    ----------
+    calibration_dataset : xr.Dataset
+        Calibration ancillary file input.
+    offset_dataset : xr.Dataset
+        Offset ancillary file input.
+    input_data : xr.Dataset
+        Input data from MAG L1C or L1B.
+    data_version : str
+        Version of output file.
+
+    Returns
+    -------
+    list[xr.Dataset]
+        List of xarray datasets ready to write to CDF file. Expected to be four outputs
+        for different frames.
+    """
+    # TODO we may need to combine multiple calibration datasets into one timeline.
+
+    basic_test_data = MagL2(
+        input_data["vectors"].data[:, :3],  # level 2 vectors don't include range
+        input_data["epoch"].data,
+        input_data["vectors"].data[:, 3],
+        {"Data_version": data_version},
+        np.zeros(len(input_data["epoch"].data)),
+        np.zeros(len(input_data["epoch"].data)),
+        DataMode.NORM,
+    )
+    attributes = ImapCdfAttributes()
+    attributes.add_instrument_global_attrs("mag")
+    # temporarily point to l1c
+    attributes.add_instrument_variable_attrs("mag", "l1c")
+
+    return [basic_test_data.generate_dataset(attributes)]
+
+
+def apply_calibration_matrix(
+    calibration_dataset: xr.Dataset, vectors: np.ndarray
+) -> np.ndarray:
+    """
+    Apply the calibration file to the vectors to rotate them in space.
+
+    Parameters
+    ----------
+    calibration_dataset : xr.Dataset
+        Ancillary file input for calibration.
+    vectors : np.ndarray
+        (n, 4) array of vectors to rotate and timeshift.
+
+    Returns
+    -------
+    np.ndarray
+        Rotated and timeshifted vectors.
+    """
+    raise NotImplementedError
diff --git a/imap_processing/mag/l2/mag_l2_data.py b/imap_processing/mag/l2/mag_l2_data.py
@@ -0,0 +1,205 @@
+"""Data structures for MAG L2 and L1D processing."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+import numpy as np
+import xarray as xr
+
+from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.mag.constants import DataMode
+
+
+class ValidFrames(Enum):
+    """SPICE reference frames for output."""
+
+    dsrf = "dsrf"
+    srf = "srf"
+    rtn = "rtn"
+    gse = "gse"
+
+
+@dataclass
+class MagL2:
+    """
+    Dataclass for MAG L2 data.
+
+    Since L2 and L1D should have the same structure, this can be used for either level.
+
+    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.
+    epoch: np.ndarray
+        Time of each vector in J2000 seconds. Should be of length n.
+    range: np.ndarray
+        Range of each vector. Should be of length n.
+    global_attributes: dict
+        Any global attributes we want to carry forward into the output CDF file.
+    quality_flags: np.ndarray
+        Quality flags for each vector. Should be of length n.
+    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.
+    """
+
+    vectors: np.ndarray
+    epoch: np.ndarray
+    range: np.ndarray
+    global_attributes: dict
+    quality_flags: np.ndarray
+    quality_bitmask: np.ndarray
+    data_mode: DataMode
+    magnitude: np.ndarray = field(init=False)
+    is_l1d: bool = False
+
+    def __post_init__(self) -> None:
+        """Calculate the magnitude of the vectors after initialization."""
+        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.zeros(vectors.shape[0])  # type: ignore
+
+    def truncate_to_24h(self, timestamp: str) -> None:
+        """
+        Truncate all data to a 24 hour period.
+
+        24 hours is given by timestamp in the format YYYYmmdd.
+
+        Parameters
+        ----------
+        timestamp : str
+            Timestamp in the format YYYYMMDD.
+        """
+        pass
+
+    def generate_dataset(
+        self,
+        attribute_manager: ImapCdfAttributes,
+        frame: ValidFrames = ValidFrames.dsrf,
+    ) -> xr.Dataset:
+        """
+        Generate an xarray dataset from the dataclass.
+
+        This method can be used for L2 and L1D, since they have extremely similar
+        output.
+
+        Parameters
+        ----------
+        attribute_manager : ImapCdfAttributes
+            CDF attributes object for the correct level.
+        frame : ValidFrames
+            SPICE reference frame to rotate the data into.
+
+        Returns
+        -------
+        xr.Dataset
+            Complete dataset ready to write to CDF file.
+        """
+        logical_source_id = f"imap_mag_l2_{self.data_mode.value.lower()}-{frame.name}"
+        direction = xr.DataArray(
+            np.arange(3),
+            name="direction",
+            dims=["direction"],
+            attrs=attribute_manager.get_variable_attributes(
+                "direction_attrs", check_schema=False
+            ),
+        )
+
+        direction_label = xr.DataArray(
+            direction.values.astype(str),
+            name="direction_label",
+            dims=["direction_label"],
+            attrs=attribute_manager.get_variable_attributes(
+                "direction_label", check_schema=False
+            ),
+        )
+
+        epoch_time = xr.DataArray(
+            self.epoch,
+            name="epoch",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("epoch"),
+        )
+
+        vectors = xr.DataArray(
+            self.vectors,
+            name="vectors",
+            dims=["epoch", "direction"],
+            attrs=attribute_manager.get_variable_attributes("vector_attrs"),
+        )
+
+        quality_flags = xr.DataArray(
+            self.quality_flags,
+            name="quality_flags",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("compression"),
+        )
+
+        quality_bitmask = xr.DataArray(
+            self.quality_flags,
+            name="quality_flags",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("compression"),
+        )
+
+        rng = xr.DataArray(
+            self.range,
+            name="range",
+            dims=["epoch"],
+            # TODO temp attrs
+            attrs=attribute_manager.get_variable_attributes("compression_width"),
+        )
+
+        magnitude = xr.DataArray(
+            self.magnitude,
+            name="magnitude",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("compression_width"),
+        )
+
+        global_attributes = (
+            attribute_manager.get_global_attributes(logical_source_id)
+            | self.global_attributes
+        )
+
+        output = xr.Dataset(
+            coords={
+                "epoch": epoch_time,
+                "direction": direction,
+                "direction_label": direction_label,
+            },
+            attrs=global_attributes,
+        )
+
+        output["vectors"] = vectors
+        output["quality_flags"] = quality_flags
+        output["quality_bitmask"] = quality_bitmask
+        output["range"] = rng
+        output["magnitude"] = magnitude
+
+        return output
diff --git a/imap_processing/tests/mag/conftest.py b/imap_processing/tests/mag/conftest.py
@@ -58,9 +58,14 @@ def mag_l1a_dataset_generator(length):
 
 
 @pytest.fixture()
-def mag_test_calibration_data():
+def mag_test_l1b_calibration_data():
     imap_dir = Path(__file__).parent
-    cal_file = imap_dir / "validation" / "imap_calibration_mag_20240229_v01.cdf"
+    cal_file = (
+        imap_dir
+        / "validation"
+        / "calibration"
+        / "imap_mag_l1b-calibration_20240229_v001.cdf"
+    )
     calibration_data = load_cdf(cal_file)
     return calibration_data
 
diff --git a/imap_processing/tests/mag/test_mag_l1b.py b/imap_processing/tests/mag/test_mag_l1b.py
@@ -17,7 +17,7 @@
 )
 
 
-def test_mag_processing(mag_test_calibration_data):
+def test_mag_processing(mag_test_l1b_calibration_data):
     # All specific test values come from MAG team to accommodate various cases.
     # Each vector is multiplied by the matrix in the calibration data for the given
     # range to get the calibrated vector.
@@ -31,7 +31,7 @@ def test_mag_processing(mag_test_calibration_data):
     mag_attributes.add_instrument_variable_attrs("mag", "l1b")
     mag_l1b = mag_l1b_processing(
         mag_l1a_dataset,
-        mag_test_calibration_data,
+        mag_test_l1b_calibration_data,
         mag_attributes,
         "imap_mag_l1b_norm-mago",
     )
@@ -50,7 +50,7 @@ def test_mag_processing(mag_test_calibration_data):
 
     mag_l1b = mag_l1b_processing(
         mag_l1a_dataset,
-        mag_test_calibration_data,
+        mag_test_l1b_calibration_data,
         mag_attributes,
         "imap_mag_l1b_norm-magi",
     )
diff --git a/imap_processing/tests/mag/test_mag_l2.py b/imap_processing/tests/mag/test_mag_l2.py
diff --git a/imap_processing/tests/mag/validation/calibration/imap_mag_l1b-calibration_20240229_v001.cdf b/imap_processing/tests/mag/validation/calibration/imap_mag_l1b-calibration_20240229_v001.cdf
