Allow customization of MAG L2 frames

imap_processing/mag/l2/mag_l2.py

@@ -12,13 +12,22 @@
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_L2_FRAMES = [
+    ValidFrames.SRF,
+    ValidFrames.GSE,
+    ValidFrames.GSM,
+    ValidFrames.RTN,
+    ValidFrames.DSRF,  # should be last as some vectors may become NaN
+]
+
 
 def mag_l2(
     calibration_dataset: xr.Dataset,
     offsets_dataset: xr.Dataset,
     input_data: xr.Dataset,
     day_to_process: np.datetime64,
     mode: DataMode = DataMode.NORM,
+    frames: list[ValidFrames] = DEFAULT_L2_FRAMES,
 ) -> list[xr.Dataset]:
     """
     Complete MAG L2 processing.
@@ -70,6 +79,9 @@ def mag_l2(
     mode : DataMode
         The data mode to process. Default is DataMode.NORM (normal mode).
         Can also be DataMode.BURST for burst mode processing.
+    frames : list[ValidFrames]
+        List of frames to output. DEFAULT_L2_FRAMES is [SRF, GSE, GSM, RTN, DSRF]
+        Note that DSRF should be last as some vectors may become NaN after rotation.
 
     Returns
     -------
@@ -79,6 +91,9 @@ def mag_l2(
     """
     always_output_mago = configuration.ALWAYS_OUTPUT_MAGO
 
+    if not frames:
+        frames = DEFAULT_L2_FRAMES
+
     # 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.")
@@ -118,19 +133,13 @@ def mag_l2(
     attributes.add_instrument_variable_attrs("mag", "l2")
 
     # Rotate from the MAG frame into the SRF frame
-    frames: list[xr.Dataset] = []
-
-    for frame in [
-        ValidFrames.SRF,
-        ValidFrames.GSE,
-        ValidFrames.GSM,
-        ValidFrames.RTN,
-        ValidFrames.DSRF,  # should be last as some vectors may become NaN
-    ]:
+    datasets: list[xr.Dataset] = []
+
+    for frame in frames:
         l2_data.rotate_frame(frame)
-        frames.append(l2_data.generate_dataset(attributes, day))
+        datasets.append(l2_data.generate_dataset(attributes, day))
 
-    return frames
+    return datasets
 
 
 def retrieve_matrix_from_l2_calibration(

imap_processing/tests/mag/test_mag_l2.py

@@ -18,9 +18,29 @@
 from imap_processing.tests.mag.conftest import mag_l1a_dataset_generator
 
 
-@pytest.mark.parametrize("data_mode", ["norm", "burst"])
-def test_mag_l2_attributes(norm_dataset, mag_test_l2_data, data_mode):
-    """Test that L2 datasets have correct attributes based on frame and mode."""
+@pytest.mark.parametrize(
+    "data_mode,frames,expected_frames",
+    [
+        (
+            "norm",
+            [
+                ValidFrames.SRF,
+                ValidFrames.GSE,
+                ValidFrames.GSM,
+                ValidFrames.RTN,
+                ValidFrames.DSRF,
+            ],
+            5,
+        ),
+        ("norm", [], 5),
+        ("burst", [ValidFrames.SRF], 1),
+        ("burst", [], 5),
+    ],
+)
+def test_mag_l2_attributes(
+    norm_dataset, mag_test_l2_data, data_mode, frames, expected_frames
+):
+    """Test that correct L2 datasets have correct attributes based on frame and mode."""
     calibration_dataset = mag_test_l2_data[0]
     offset_dataset = mag_test_l2_data[1]
 
@@ -35,18 +55,30 @@ def test_mag_l2_attributes(norm_dataset, mag_test_l2_data, data_mode):
         "imap_processing.mag.l2.mag_l2_data.frame_transform",
         side_effect=lambda *args, **kwargs: args[1],
     ):
-        l2_datasets = mag_l2(
-            calibration_dataset,
-            offset_dataset,
-            test_dataset,
-            np.datetime64("2025-10-17"),
-            mode=mode,
-        )
+        if frames:
+            # ensure when a subset of frames is needed only those are generated
+            l2_datasets = mag_l2(
+                calibration_dataset,
+                offset_dataset,
+                test_dataset,
+                np.datetime64("2025-10-17"),
+                mode=mode,
+                frames=frames,
+            )
+        else:
+            # be default all frames are generated
+            l2_datasets = mag_l2(
+                calibration_dataset,
+                offset_dataset,
+                test_dataset,
+                np.datetime64("2025-10-17"),
+                mode=mode,
+            )
 
     # Verify we have the expected number of datasets
     # L2 produces 5 frames: SRF, GSE, GSM, RTN, DSRF
-    assert len(l2_datasets) == 5, (
-        f"Expected 5 {data_mode} datasets, got {len(l2_datasets)}"
+    assert len(l2_datasets) == expected_frames, (
+        f"Expected {expected_frames} {data_mode} datasets, got {len(l2_datasets)}"
     )
 
     for dataset in l2_datasets: