From f58ece467917854925f9b329da39a36a66aa1185 Mon Sep 17 00:00:00 2001 From: Alastair Crabtree Date: Mon, 30 Mar 2026 12:26:32 +0100 Subject: [PATCH 1/3] feat: Allow customization of MAG L2 frames - in mag_l2 function and update tests accordingly - used (externally) by MAG team to generate l2 files without needing all files --- imap_processing/mag/l2/mag_l2.py | 25 ++++++----- imap_processing/tests/mag/test_mag_l2.py | 56 +++++++++++++++++++----- 2 files changed, 59 insertions(+), 22 deletions(-) diff --git a/imap_processing/mag/l2/mag_l2.py b/imap_processing/mag/l2/mag_l2.py index 2c47dd5026..be8d054b9e 100644 --- a/imap_processing/mag/l2/mag_l2.py +++ b/imap_processing/mag/l2/mag_l2.py @@ -12,6 +12,13 @@ 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, @@ -19,6 +26,7 @@ def mag_l2( 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. @@ -79,6 +87,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 +129,13 @@ def mag_l2( attributes.add_instrument_variable_attrs("mag", "l2") # Rotate from the MAG frame into the SRF frame - frames: list[xr.Dataset] = [] + datasets: 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 - ]: + 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( diff --git a/imap_processing/tests/mag/test_mag_l2.py b/imap_processing/tests/mag/test_mag_l2.py index 6be5a7150a..bb2827920b 100644 --- a/imap_processing/tests/mag/test_mag_l2.py +++ b/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 are generated and 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: From f3a354b6d7269e2cb205c11e30374a43ebae3e11 Mon Sep 17 00:00:00 2001 From: Alastair Crabtree Date: Mon, 30 Mar 2026 12:43:24 +0100 Subject: [PATCH 2/3] chore: QA fixes --- imap_processing/mag/l2/mag_l2.py | 15 ++++++++------- imap_processing/tests/mag/test_mag_l2.py | 2 +- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/imap_processing/mag/l2/mag_l2.py b/imap_processing/mag/l2/mag_l2.py index be8d054b9e..0db05eaa15 100644 --- a/imap_processing/mag/l2/mag_l2.py +++ b/imap_processing/mag/l2/mag_l2.py @@ -13,12 +13,13 @@ 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 - ] + 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, @@ -26,7 +27,7 @@ def mag_l2( input_data: xr.Dataset, day_to_process: np.datetime64, mode: DataMode = DataMode.NORM, - frames: list[ValidFrames] = DEFAULT_L2_FRAMES + frames: list[ValidFrames] = DEFAULT_L2_FRAMES, ) -> list[xr.Dataset]: """ Complete MAG L2 processing. diff --git a/imap_processing/tests/mag/test_mag_l2.py b/imap_processing/tests/mag/test_mag_l2.py index bb2827920b..473df52718 100644 --- a/imap_processing/tests/mag/test_mag_l2.py +++ b/imap_processing/tests/mag/test_mag_l2.py @@ -40,7 +40,7 @@ def test_mag_l2_attributes( norm_dataset, mag_test_l2_data, data_mode, frames, expected_frames ): - """Test that correct L2 datasets are generated and have correct attributes based on frame and mode.""" + """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] From 590e53b012f345e655356b43e0f3b83126cf4bf3 Mon Sep 17 00:00:00 2001 From: Alastair Crabtree Date: Mon, 30 Mar 2026 12:52:39 +0100 Subject: [PATCH 3/3] docs: Add docs for new parameter in MAG L2 --- imap_processing/mag/l2/mag_l2.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/imap_processing/mag/l2/mag_l2.py b/imap_processing/mag/l2/mag_l2.py index 0db05eaa15..65632ae1a5 100644 --- a/imap_processing/mag/l2/mag_l2.py +++ b/imap_processing/mag/l2/mag_l2.py @@ -79,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 -------