Mag L2 rotation - propogate FILLVALS (#2903)

* Revert "A rotated NaN is still a NaN (#2901)"

This reverts commit cdfd0ae.

* Apply fix to propogate FILLVALs out from SPICE rotations

Author: maxinelasp

imap_processing/mag/l1d/mag_l1d_data.py

@@ -298,26 +298,40 @@ def rotate_frame(self, end_frame: ValidFrames) -> 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(
+        new_vectors = frame_transform(
             self.epoch_et,
             self.vectors,
             from_frame=start_frame.spice_frame,
             to_frame=end_frame.spice_frame,
             allow_spice_noframeconnect=True,
         )
+        if np.isnan(self.vectors).any() or (self.vectors == FILLVAL).any():
+            new_vectors = np.where(
+                np.isnan(self.vectors) | (self.vectors == FILLVAL),
+                FILLVAL,
+                new_vectors,
+            )
+        self.vectors = new_vectors
 
         # 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(
+        new_magi_vectors = frame_transform(
             self.magi_epoch_et,
             self.magi_vectors,
             from_frame=start_frame.spice_frame,
             to_frame=end_frame.spice_frame,
             allow_spice_noframeconnect=True,
         )
+        if np.isnan(self.magi_vectors).any() or (self.magi_vectors == FILLVAL).any():
+            new_magi_vectors = np.where(
+                np.isnan(self.magi_vectors) | (self.magi_vectors == FILLVAL),
+                FILLVAL,
+                new_magi_vectors,
+            )
+        self.magi_vectors = new_magi_vectors
 
         self.frame = end_frame
 

imap_processing/mag/l2/mag_l2_data.py

@@ -7,6 +7,7 @@
 import xarray as xr
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.mag import constants
 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
@@ -417,13 +418,21 @@ def rotate_frame(self, end_frame: ValidFrames) -> None:
         """
         if self.epoch_et is None:
             self.epoch_et = ttj2000ns_to_et(self.epoch)
-        self.vectors = frame_transform(
+        new_vectors = frame_transform(
             self.epoch_et,
             self.vectors,
             from_frame=self.frame.spice_frame,
             to_frame=end_frame.spice_frame,
             allow_spice_noframeconnect=True,
         )
+        if np.isnan(self.vectors).any() or (self.vectors == constants.FILLVAL).any():
+            new_vectors = np.where(
+                np.isnan(self.vectors) | (self.vectors == constants.FILLVAL),
+                constants.FILLVAL,
+                new_vectors,
+            )
+
+        self.vectors = new_vectors
         self.frame = end_frame
 
 

imap_processing/tests/mag/conftest.py

@@ -175,7 +175,11 @@ def norm_dataset(mag_test_l2_data):
     dataset.attrs["vectors_per_second"] = vectors_per_second_attr
     dataset["epoch"] = epoch_vals
     dataset.attrs["Logical_source"] = "imap_mag_l1c_norm-mago"
-    vectors = np.array([[i, i, i, 2] for i in range(1, 3505)])
+    # Actual dataset is a CDF_FLOAT which is a float32.
+    vectors = np.array(
+        [[i, i, i, 2] for i in range(1, 3505)],
+        dtype=np.float64,
+    )
     dataset["vectors"].data = vectors
 
     return dataset

imap_processing/tests/mag/test_mag_l1d.py

@@ -8,7 +8,7 @@
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.cdf.utils import write_cdf
 from imap_processing.cli import Mag
-from imap_processing.mag.constants import DataMode
+from imap_processing.mag.constants import FILLVAL, DataMode
 from imap_processing.mag.l1d.mag_l1d import mag_l1d
 from imap_processing.mag.l1d.mag_l1d_data import MagL1d, MagL1dConfiguration
 from imap_processing.mag.l2.mag_l2_data import ValidFrames
@@ -571,6 +571,57 @@ def test_enhanced_gradiometry_with_quality_flags_detailed():
     assert np.array_equal(grad_ds["quality_flags"].data, expected_flags)
 
 
+def test_rotate_frame_preserves_fillval_and_nan(mag_l1d_test_class):
+    """Test that L1D rotate_frame preserves FILLVAL and NaN vectors."""
+    mag_l1d_test_class.frame = ValidFrames.MAGO
+
+    vectors = mag_l1d_test_class.vectors.copy()
+    magi_vectors = mag_l1d_test_class.magi_vectors.copy()
+
+    # Set some MAGO vectors to FILLVAL and NaN
+    vectors[0] = [FILLVAL, FILLVAL, FILLVAL]
+    vectors[2] = [np.nan, np.nan, np.nan]
+    vectors[4] = [1.0, np.nan, 3.0]
+    mag_l1d_test_class.vectors = vectors
+
+    # Set some MAGI vectors to FILLVAL and NaN
+    magi_vectors[1] = [FILLVAL, FILLVAL, FILLVAL]
+    magi_vectors[3] = [np.nan, np.nan, np.nan]
+    mag_l1d_test_class.magi_vectors = magi_vectors
+
+    def mock_frame_transform(
+        epoch_et,
+        vecs,
+        from_frame,
+        to_frame,
+        allow_spice_noframeconnect,
+    ):
+        return np.full(vecs.shape, 99.0)
+
+    with patch(
+        "imap_processing.mag.l1d.mag_l1d_data.frame_transform",
+        side_effect=mock_frame_transform,
+    ):
+        mag_l1d_test_class.rotate_frame(ValidFrames.SRF)
+
+    assert mag_l1d_test_class.frame == ValidFrames.SRF
+
+    # MAGO: FILLVAL/NaN rows preserved as FILLVAL
+    assert np.all(mag_l1d_test_class.vectors[0] == FILLVAL)
+    assert np.all(mag_l1d_test_class.vectors[2] == FILLVAL)
+    assert mag_l1d_test_class.vectors[4, 1] == FILLVAL
+    # Normal MAGO vectors get rotated value
+    assert np.all(mag_l1d_test_class.vectors[1] == 99.0)
+    assert np.all(mag_l1d_test_class.vectors[3] == 99.0)
+
+    # MAGI: FILLVAL/NaN rows preserved as FILLVAL
+    assert np.all(mag_l1d_test_class.magi_vectors[1] == FILLVAL)
+    assert np.all(mag_l1d_test_class.magi_vectors[3] == FILLVAL)
+    # Normal MAGI vectors get rotated value
+    assert np.all(mag_l1d_test_class.magi_vectors[0] == 99.0)
+    assert np.all(mag_l1d_test_class.magi_vectors[2] == 99.0)
+
+
 def test_rotate_frames(mag_l1d_test_class):
     # Reset to initial MAGO frame for this test
     mag_l1d_test_class.frame = ValidFrames.MAGO

imap_processing/tests/mag/test_mag_l2.py

@@ -5,7 +5,7 @@
 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.l2.mag_l2 import mag_l2, retrieve_matrix_from_l2_calibration
 from imap_processing.mag.l2.mag_l2_data import MagL2, ValidFrames
 from imap_processing.spice.time import (
@@ -447,6 +447,55 @@ def test_spice_returns(norm_dataset):
         assert np.array_equal(l2.vectors[0], [-1, -1, -1])
 
 
+def test_rotate_frame_preserves_fillval_and_nan(norm_dataset):
+    """Test that rotate_frame preserves FILLVAL and NaN vectors."""
+
+    vectors = norm_dataset["vectors"].data[:, :3].copy()
+    n = len(vectors)
+
+    # Set some vectors to FILLVAL and NaN
+    vectors[0] = [FILLVAL, FILLVAL, FILLVAL]
+    vectors[2] = [np.nan, np.nan, np.nan]
+    # Partial NaN in a row
+    vectors[4] = [1.0, np.nan, 3.0]
+    # Partial FILLVAL in a row
+    vectors[5] = [FILLVAL, 2.0, 3.0]
+
+    l2 = MagL2(
+        vectors=vectors,
+        epoch=norm_dataset["epoch"].data,
+        range=norm_dataset["vectors"].data[:, 3],
+        global_attributes={},
+        quality_flags=np.zeros(n),
+        quality_bitmask=np.zeros(n),
+        data_mode=DataMode.NORM,
+        offsets=np.zeros((n, 3)),
+        timedelta=np.zeros(n),
+    )
+
+    rotated_values = np.full(vectors.shape, 99.0)
+    with patch(
+        "imap_processing.mag.l2.mag_l2_data.frame_transform",
+        return_value=rotated_values,
+    ):
+        l2.rotate_frame(ValidFrames.DSRF)
+
+    assert l2.frame == ValidFrames.DSRF
+
+    # Full FILLVAL row -> all components should be FILLVAL
+    assert np.all(l2.vectors[0] == FILLVAL)
+    # Full NaN row -> all components should be FILLVAL
+    assert np.all(l2.vectors[2] == FILLVAL)
+    # Partial NaN -> affected components should be FILLVAL
+    assert l2.vectors[4, 1] == FILLVAL
+    # Partial FILLVAL -> affected components should be FILLVAL
+    assert l2.vectors[5, 0] == FILLVAL
+
+    # Normal vectors should get the rotated value
+    assert np.all(l2.vectors[1] == 99.0)
+    assert np.all(l2.vectors[3] == 99.0)
+
+
 def test_qf(norm_dataset):
     qf = np.zeros(len(norm_dataset["epoch"].data), dtype=int)
     qf[1:4] = 1