Feat/update mag frame

imap_processing/ialirt/l0/parse_mag.py

@@ -24,7 +24,7 @@
     shift_time,
 )
 from imap_processing.mag.l1d.mag_l1d_data import MagL1d
-from imap_processing.mag.l2.mag_l2_data import MagL2L1dBase
+from imap_processing.mag.l2.mag_l2_data import MagL2L1dBase, ValidFrames
 from imap_processing.spice.geometry import (
     SpiceFrame,
     cartesian_to_spherical,
@@ -703,7 +703,7 @@ def process_packet(
             attitude_time,
             time_data["primary_epoch"],
             mago_out,
-            SpiceFrame.IMAP_MAG_O,
+            ValidFrames.MAGO.spice_frame,
         )
         magi_inertial_vector = transform_to_inertial(
             sc_spin_phase_rad.values,
@@ -712,7 +712,7 @@ def process_packet(
             attitude_time,
             time_data["secondary_epoch"],
             magi_out,
-            SpiceFrame.IMAP_MAG_I,
+            ValidFrames.MAGI.spice_frame,
         )
 
         met = grouped_data["met"][(grouped_data["group"] == group).values]

imap_processing/mag/l1d/mag_l1d_data.py

@@ -298,8 +298,8 @@ def rotate_frame(self, end_frame: ValidFrames) -> None:
         self.vectors = frame_transform(
             self.epoch_et,
             self.vectors,
-            from_frame=start_frame.value,
-            to_frame=end_frame.value,
+            from_frame=start_frame.spice_frame,
+            to_frame=end_frame.spice_frame,
             allow_spice_noframeconnect=True,
         )
 
@@ -311,8 +311,8 @@ def rotate_frame(self, end_frame: ValidFrames) -> None:
         self.magi_vectors = frame_transform(
             self.magi_epoch_et,
             self.magi_vectors,
-            from_frame=start_frame.value,
-            to_frame=end_frame.value,
+            from_frame=start_frame.spice_frame,
+            to_frame=end_frame.spice_frame,
             allow_spice_noframeconnect=True,
         )
 

imap_processing/mag/l2/mag_l2_data.py

@@ -20,13 +20,105 @@
 class ValidFrames(Enum):
     """SPICE reference frames for output."""
 
-    MAGO = SpiceFrame.IMAP_MAG_O
-    MAGI = SpiceFrame.IMAP_MAG_I
-    DSRF = SpiceFrame.IMAP_DPS
-    SRF = SpiceFrame.IMAP_SPACECRAFT
-    GSE = SpiceFrame.IMAP_GSE
-    GSM = SpiceFrame.IMAP_GSM
-    RTN = SpiceFrame.IMAP_RTN
+    """
+    Default MAGO and MAGI L1D and L2 frames both map to the same SPICE frame.
+    This is because the idealised IMAP_MAG_BASE frame is used for both sensors,
+    as the MAG team provides a calibration matrix to convert from the real mechanical
+    mount as assessed in flight into the idealised frame.
+
+    MAGO_GROUND_CAL and MAGI_GROUND_CAL additionally included for reference to the
+    ground assessed mount, and for future use if needed.
+    """
+    MAGO = ("MAGO", SpiceFrame.IMAP_MAG_BASE, "vector_attrs", "vectors")
+    MAGI = ("MAGI", SpiceFrame.IMAP_MAG_BASE, "vector_attrs", "vectors")
+
+    MAGO_GROUND_CAL = (
+        "MAGO_GROUND_CAL",
+        SpiceFrame.IMAP_MAG_O,
+        "vector_attrs",
+        "vectors",
+    )
+    MAGI_GROUND_CAL = (
+        "MAGI_GROUND_CAL",
+        SpiceFrame.IMAP_MAG_I,
+        "vector_attrs",
+        "vectors",
+    )
+
+    DSRF = ("DSRF", SpiceFrame.IMAP_DPS, "vector_attrs_dsrf", "b_dsrf")
+    SRF = ("SRF", SpiceFrame.IMAP_SPACECRAFT, "vector_attrs_srf", "b_srf")
+    GSE = ("GSE", SpiceFrame.IMAP_GSE, "vector_attrs_gse", "b_gse")
+    GSM = ("GSM", SpiceFrame.IMAP_GSM, "vector_attrs_gsm", "b_gsm")
+    RTN = ("RTN", SpiceFrame.IMAP_RTN, "vector_attrs_rtn", "b_rtn")
+
+    _spice_frame_: SpiceFrame
+    _vector_attrs_name_: str
+    _var_name_: str
+
+    def __new__(
+        cls, value: str, spice_frame: SpiceFrame, attrs_name: str, var_name: str
+    ) -> "ValidFrames":
+        """
+        Construct a new Valid Frame.
+
+        Parameters
+        ----------
+        value : str
+            Unique name of the frame.
+        spice_frame : str
+            The SPICE frame name corresponding to this frame.
+        attrs_name : str
+            The name of the variable attributes in the attribute manager for this frame.
+        var_name : str
+            The name of the variable in the output dataset for this frame.
+
+        Returns
+        -------
+        ValidFrame : ValidFrame
+            A ValidFrame enum member.
+        """
+        obj = object.__new__(cls)
+        obj._value_ = value
+        obj._spice_frame_ = spice_frame
+        obj._vector_attrs_name_ = attrs_name
+        obj._var_name_ = var_name
+        return obj
+
+    @property
+    def spice_frame(self) -> SpiceFrame:
+        """
+        Get the SPICE frame name for this ValidFrame.
+
+        Returns
+        -------
+        spice_frame : str
+            The frame's associated spice frame.
+        """
+        return self._spice_frame_
+
+    @property
+    def vector_attrs_name(self) -> str:
+        """
+        Get the vector attributes name for this valid frame.
+
+        Returns
+        -------
+        vector_attrs_name : str
+            The frame's associated vector attributes name.
+        """
+        return self._vector_attrs_name_
+
+    @property
+    def var_name(self) -> str:
+        """
+        Get the vector variable name for this valid frame.
+
+        Returns
+        -------
+        var_name : str
+            The frame's associated vectors variable name.
+        """
+        return self._var_name_
 
 
 @dataclass(kw_only=True)
@@ -109,16 +201,6 @@ def generate_dataset(
             f"{self.frame.name.lower()}"
         )
 
-        # Select the appropriate vector attributes based on the frame
-        frame_to_vector_attrs = {
-            ValidFrames.SRF: "vector_attrs_srf",
-            ValidFrames.DSRF: "vector_attrs_dsrf",
-            ValidFrames.GSE: "vector_attrs_gse",
-            ValidFrames.RTN: "vector_attrs_rtn",
-            ValidFrames.GSM: "vector_attrs_gsm",  # L2 Only
-        }
-        vector_attrs_name = frame_to_vector_attrs.get(self.frame, "vector_attrs")
-
         direction = xr.DataArray(
             np.arange(3),
             name="direction",
@@ -148,9 +230,11 @@ def generate_dataset(
 
         vectors = xr.DataArray(
             self.vectors,
-            name="vectors",
+            name=self.frame.var_name,
             dims=["epoch", "direction"],
-            attrs=attribute_manager.get_variable_attributes(vector_attrs_name),
+            attrs=attribute_manager.get_variable_attributes(
+                self.frame.vector_attrs_name
+            ),
         )
 
         quality_flags = xr.DataArray(
@@ -195,7 +279,7 @@ def generate_dataset(
             attrs=global_attributes,
         )
 
-        output["vectors"] = vectors
+        output[self.frame.var_name] = vectors
         output["quality_flags"] = quality_flags
         output["quality_bitmask"] = quality_bitmask
         output["range"] = rng
@@ -332,8 +416,8 @@ def rotate_frame(self, end_frame: ValidFrames) -> None:
         self.vectors = frame_transform(
             self.epoch_et,
             self.vectors,
-            from_frame=self.frame.value,
-            to_frame=end_frame.value,
+            from_frame=self.frame.spice_frame,
+            to_frame=end_frame.spice_frame,
             allow_spice_noframeconnect=True,
         )
         self.frame = end_frame

imap_processing/tests/mag/test_mag_l1d.py

@@ -94,8 +94,9 @@ def test_mag_l1d(mag_test_l1d_data, norm_dataset, furnish_kernels, fake_mag_spin
         )
     # Should have: 4 norm frames + 4 burst frames + spin offsets + 2 gradiometry offsets
 
+    frame = l1d[0].attrs["Logical_source"].split("-")[-1].lower()
     assert len(l1d) == 11
-    assert "vectors" in l1d[0].data_vars
+    assert f"b_{frame}" in l1d[0].data_vars
 
     # Check that expected logical sources are present
     logical_sources = [ds.attrs.get("Logical_source", "") for ds in l1d]
@@ -181,11 +182,11 @@ def test_mag_l1d_attributes(
             f"got '{logical_source_parts[2]}'"
         )
 
-        vectors_attrs = dataset["vectors"].attrs
-        assert "DICT_KEY" in vectors_attrs
-
         frame = dataset.attrs["Logical_source"].split("-")[-1].upper()
 
+        vectors_attrs = dataset[f"b_{frame.lower()}"].attrs
+        assert "DICT_KEY" in vectors_attrs
+
         assert f"CoordinateSystemName:{frame}" in vectors_attrs["DICT_KEY"]
 
         assert "magnitude" in dataset.data_vars
@@ -482,7 +483,7 @@ def test_mago_magi_swap_functionality(mag_l1d_test_class):
     assert np.array_equal(mag_l1d_test_class.vectors, mago_vectors)
     assert np.array_equal(mag_l1d_test_class.epoch, mago_epoch)
 
-    assert np.array_equal(result["vectors"].data, magi_vectors)
+    assert np.array_equal(result[mag_l1d_test_class.frame.var_name].data, magi_vectors)
     assert np.array_equal(result["epoch"].data, magi_epoch)
 
 
@@ -509,7 +510,7 @@ def test_mago_magi_no_swap_functionality(mag_l1d_test_class):
     assert np.array_equal(mag_l1d_test_class.vectors, mago_vectors)
     assert np.array_equal(mag_l1d_test_class.epoch, mago_epoch)
 
-    assert np.array_equal(result["vectors"].data, mago_vectors)
+    assert np.array_equal(result[mag_l1d_test_class.frame.var_name].data, mago_vectors)
     assert np.array_equal(result["epoch"].data, mago_epoch)
 
 
@@ -581,10 +582,8 @@ def test_rotate_frames(mag_l1d_test_class):
     def mock_frame_transform(
         epoch_et, vectors, from_frame, to_frame, allow_spice_noframeconnect
     ):
-        if from_frame == ValidFrames.MAGO.value:
+        if from_frame in [ValidFrames.MAGO.spice_frame, ValidFrames.MAGI.spice_frame]:
             return vectors + 100
-        elif from_frame == ValidFrames.MAGI.value:
-            return vectors + 200
         else:
             return vectors + 300
 
@@ -600,20 +599,22 @@ def mock_frame_transform(
 
         # First call should be for MAGO vectors
         first_call_args = mock_transform_l1d.call_args_list[0]
-        assert first_call_args[1]["from_frame"] == ValidFrames.MAGO.value
-        assert first_call_args[1]["to_frame"] == ValidFrames.SRF.value
+        assert first_call_args[1]["from_frame"] == ValidFrames.MAGO.spice_frame
+        assert first_call_args[1]["to_frame"] == ValidFrames.SRF.spice_frame
 
         # Second call should be for MAGI vectors
         second_call_args = mock_transform_l1d.call_args_list[1]
-        assert second_call_args[1]["from_frame"] == ValidFrames.MAGI.value
-        assert second_call_args[1]["to_frame"] == ValidFrames.SRF.value
+        assert second_call_args[1]["from_frame"] == ValidFrames.MAGI.spice_frame
+        assert second_call_args[1]["to_frame"] == ValidFrames.SRF.spice_frame
+
+        # MAGO frame and MAGi frame not necessarily different (and are now the same)
 
-        # Check that MAGO vectors were transformed from MAGO frame (+100)
+        # Check that MAGO vectors were transformed (+100)
         expected_mago_vectors = initial_vectors + 100
         np.testing.assert_array_equal(mag_l1d_test_class.vectors, expected_mago_vectors)
 
-        # Check that MAGI vectors were transformed from MAGI frame (+200)
-        expected_magi_vectors = initial_magi_vectors + 200
+        # Check that MAGI vectors were transformed (+100)
+        expected_magi_vectors = initial_magi_vectors + 100
         np.testing.assert_array_equal(
             mag_l1d_test_class.magi_vectors, expected_magi_vectors
         )

imap_processing/tests/mag/test_mag_l2.py

@@ -61,12 +61,12 @@ def test_mag_l2_attributes(norm_dataset, mag_test_l2_data, data_mode):
             f"got '{logical_source_parts[2]}'"
         )
 
-        vectors_attrs = dataset["vectors"].attrs
-        assert "DICT_KEY" in vectors_attrs
-
         # Extract frame from logical source
         frame = dataset.attrs["Logical_source"].split("-")[-1].upper()
 
+        vectors_attrs = dataset[f"b_{frame.lower()}"].attrs
+        assert "DICT_KEY" in vectors_attrs
+
         assert f"CoordinateSystemName:{frame}" in vectors_attrs["DICT_KEY"]
 
         assert "magnitude" in dataset.data_vars
@@ -105,7 +105,7 @@ def test_mag_l2(norm_dataset, mag_test_l2_data):
     )
 
     for i, dataset in enumerate(l2):
-        assert "vectors" in dataset.data_vars
+        assert expected_frames[i].var_name in dataset.data_vars
         assert expected_frames[i].name in dataset.attrs["Data_type"]
 
 
@@ -114,7 +114,7 @@ def return_some_nan_matrices_for_dsrf(
         et, from_frame, to_frame, allow_spice_noframeconnect
     ):
         matrices = np.tile(np.eye(3), (len(et), 1, 1))
-        if to_frame == ValidFrames.DSRF.value:
+        if to_frame == ValidFrames.DSRF.spice_frame:
             for i in range(10, matrices.shape[0], 10):  # every 10th matrix is NaN
                 matrices[i] = np.full((3, 3), np.nan)
         return matrices
@@ -135,14 +135,18 @@ def return_some_nan_matrices_for_dsrf(
 
     assert len(l2) == 5, "L2 should produce 5 frames"
 
+    all_vars = ["b_srf", "b_gse", "b_gsm", "b_rtn", "b_dsrf"]
+
     for dataset in l2:
-        assert "vectors" in dataset.data_vars
+        assert len(set(all_vars) & set(dataset.data_vars)) == 1, (
+            "Each dataset should have one of the expected vector variables"
+        )
 
     assert (
         l2[-1].attrs["Data_type"] == "L2_norm-dsrf>Level 2 normal rate data in DSRF"
     ), "Last frame should be DSRF"
 
-    dsrf_vectors = l2[-1]["vectors"].data
+    dsrf_vectors = l2[-1]["b_dsrf"].data
     for i in range(10, len(dsrf_vectors), 10):
         assert np.isnan(dsrf_vectors[i]).all(), f"Vectors at index {i} should be NaN"
 

imap_processing/tests/mag/test_mag_validation.py

@@ -398,19 +398,19 @@ def test_mag_l2_validation(test_number, mode):
 
         assert np.allclose(
             expected_output["x"].iloc[index],
-            l2["vectors"].data[index][0],
+            l2["b_srf"].data[index][0],
             atol=1e-5,
             rtol=0,
         )
         assert np.allclose(
             expected_output["y"].iloc[index],
-            l2["vectors"].data[index][1],
+            l2["b_srf"].data[index][1],
             atol=1e-5,
             rtol=0,
         )
         assert np.allclose(
             expected_output["z"].iloc[index],
-            l2["vectors"].data[index][2],
+            l2["b_srf"].data[index][2],
             atol=1e-5,
             rtol=0,
         )