diff --git a/imap_processing/mag/l2/mag_l2.py b/imap_processing/mag/l2/mag_l2.py index 2c47dd5026..65632ae1a5 100644 --- a/imap_processing/mag/l2/mag_l2.py +++ b/imap_processing/mag/l2/mag_l2.py @@ -12,6 +12,14 @@ 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 +27,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. @@ -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( diff --git a/imap_processing/tests/mag/test_mag_l2.py b/imap_processing/tests/mag/test_mag_l2.py index 6be5a7150a..473df52718 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 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: