imap-processing 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

imap_processing/tests/swe/test_swe_l1a_science.py

@@ -75,7 +75,7 @@ def test_data_order(decom_test_data):
     )
 
     # Get unpacked science data
-    processed_data = swe_science(decom_test_data, "001")
+    processed_data = swe_science(decom_test_data)
 
     quarter_cycle = processed_data["quarter_cycle"].isel(epoch=slice(0, 4))
     np.testing.assert_array_equal(quarter_cycle, [0, 1, 2, 3])
@@ -84,7 +84,7 @@ def test_data_order(decom_test_data):
 def test_swe_science_algorithm(decom_test_data):
     """Test general shape of return dataset from swe_science."""
     # Get unpacked science data
-    processed_data = swe_science(decom_test_data, "001")
+    processed_data = swe_science(decom_test_data)
 
     # science data should have this shape, 15x12x7.
     science_data = processed_data["science_data"].data[0]
@@ -101,7 +101,7 @@ def test_decompress_counts(decom_test_data, l1a_validation_df):
     raw_counts = l1a_validation_df.iloc[:, 1:8]
     decompressed_counts = l1a_validation_df.iloc[:, 8:15]
 
-    l1a_dataset = swe_science(decom_test_data, "001")
+    l1a_dataset = swe_science(decom_test_data)
 
     # compare raw counts
     assert np.all(

imap_processing/tests/swe/test_swe_l1b.py

@@ -2,19 +2,112 @@ from unittest.mock import patch
 
 import numpy as np
 import pandas as pd
+import pytest
+from imap_data_access.processing_input import (
+    AncillaryInput,
+    ProcessingInputCollection,
+    ScienceInput,
+)
 
 from imap_processing import imap_module_directory
 from imap_processing.cdf.utils import load_cdf, write_cdf
 from imap_processing.swe.l1a.swe_l1a import swe_l1a
 from imap_processing.swe.l1a.swe_science import swe_science
-from imap_processing.swe.l1b.swe_l1b import swe_l1b
-from imap_processing.swe.l1b.swe_l1b_science import (
+from imap_processing.swe.l1b.swe_l1b import (
+    apply_in_flight_calibration,
     convert_counts_to_rate,
     deadtime_correction,
+    get_checker_board_pattern,
+    get_indices_of_full_cycles,
+    swe_l1b,
 )
+from imap_processing.swe.utils import swe_constants
+
+
+@pytest.fixture(scope="session")
+def l1a_test_data(decom_test_data):
+    """Read test data from file and process to l1a"""
+    processed_data = swe_science(decom_test_data)
+    return processed_data
+
+
+def test_get_full_cycle_data_indices():
+    q = np.array([0, 1, 2, 0, 1, 2, 3, 2, 3, 0, 2, 3, 0, 1, 2, 3, 2, 3, 1, 0])
+    filtered_q = get_indices_of_full_cycles(q)
+    np.testing.assert_array_equal(filtered_q, np.array([3, 4, 5, 6, 12, 13, 14, 15]))
+
+    q = np.array([0, 1, 0, 1, 2, 3, 0, 2])
+    filtered_q = get_indices_of_full_cycles(q)
+    np.testing.assert_array_equal(filtered_q, np.array([2, 3, 4, 5]))
+
+    q = np.array([0, 1, 2, 3])
+    filtered_q = get_indices_of_full_cycles(q)
+    np.testing.assert_array_equal(filtered_q, np.array([0, 1, 2, 3]))
+
+    q = np.array([1, 2])
+    filtered_q = get_indices_of_full_cycles(q)
+    np.testing.assert_array_equal(filtered_q, np.array([]))
 
 
-def test_swe_l1b(decom_test_data_derived):
+def test_in_flight_calibration_factor(l1a_test_data):
+    """Test that the L1B processing is working as expected."""
+    # create sample data
+
+    input_time = 453051355.0
+    input_count = 19967
+    one_full_cycle_data = np.full(
+        (
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_SECTORS,
+            swe_constants.N_CEMS,
+        ),
+        input_count,
+    )
+    acquisition_time = np.full(
+        (swe_constants.N_ESA_STEPS, swe_constants.N_ANGLE_SECTORS), input_time
+    )
+
+    # Test that calibration factor is within correct range given test data
+    expected_cal_factor = 1 + ((1 - 1) / (453051900 - 453051300)) * (
+        input_time - 453051300
+    )
+
+    in_flight_cal_files = [
+        imap_module_directory
+        / "tests/swe/lut/imap_swe_l1b-in-flight-cal_20240510_20260716_v000.csv"
+    ]
+    calibrated_count = apply_in_flight_calibration(
+        one_full_cycle_data,
+        acquisition_time,
+        in_flight_cal_files,
+    )
+
+    np.testing.assert_allclose(
+        calibrated_count,
+        np.full(
+            (
+                swe_constants.N_ESA_STEPS,
+                swe_constants.N_ANGLE_SECTORS,
+                swe_constants.N_CEMS,
+            ),
+            input_count * expected_cal_factor,
+        ),
+        rtol=1e-9,
+    )
+
+    # Check for value outside of calibration time range
+    input_time = 1.0
+    acquisition_time = np.full(
+        (swe_constants.N_ESA_STEPS, swe_constants.N_ANGLE_SECTORS), input_time
+    )
+
+    with pytest.raises(ValueError, match="Acquisition min/max times: "):
+        apply_in_flight_calibration(
+            one_full_cycle_data, acquisition_time, in_flight_cal_files
+        )
+
+
+def test_swe_l1b_conversion(decom_test_data_derived):
     """Test that calculate engineering unit(EU) matches validation data.
 
     Parameters
@@ -22,7 +115,7 @@ def test_swe_l1b(decom_test_data_derived):
     decom_test_data_derived : xarray.dataset
         Dataset with derived values
     """
-    science_l1a_ds = swe_science(decom_test_data_derived, "001")
+    science_l1a_ds = swe_science(decom_test_data_derived)
 
     # read science validation data
     test_data_path = imap_module_directory / "tests/swe/l0_validation_data"
@@ -47,29 +140,73 @@ def test_swe_l1b(decom_test_data_derived):
         )
 
 
-@patch(
-    "imap_processing.swe.l1b.swe_l1b_science.read_in_flight_cal_data",
-    return_value=pd.DataFrame(
-        {
-            "met_time": [452051300, 454051900],
-            "cem1": [1, 1],
-            "cem2": [1, 1],
-            "cem3": [1, 1],
-            "cem4": [1, 1],
-            "cem5": [1, 1],
-            "cem6": [1, 1],
-            "cem7": [1, 1],
-        }
-    ),
-)
-def test_cdf_creation(mock_read_in_flight_cal_data, l1b_validation_df):
+def test_get_checker_board_pattern():
+    """Test that the checkerboard pattern is generated correctly."""
+    # First read the checkerboard pattern from the file
+    expected_checkerboard = pd.read_csv(
+        imap_module_directory / "tests/swe/lut/checker-board-indices.csv", header=None
+    ).values
+    esa_lut_file = (
+        imap_module_directory / "tests/swe/lut/imap_swe_esa-lut_20250301_v000.csv"
+    )
+    checkerboard_pattern = get_checker_board_pattern(esa_lut_file)
+
+    # Use assert_array_equal to compare the arrays
+    np.testing.assert_array_equal(
+        checkerboard_pattern,
+        expected_checkerboard,
+        err_msg="Checkerboard pattern is not the same as the one in the LUT file.",
+    )
+
+
+@patch("imap_data_access.processing_input.ProcessingInputCollection.get_file_paths")
+def test_swe_l1b(mock_get_file_paths, l1b_validation_df):
     """Test that CDF file is created and has the correct name."""
     test_data_path = "tests/swe/l0_data/2024051010_SWE_SCIENCE_packet.bin"
-    l1a_datasets = swe_l1a(imap_module_directory / test_data_path, "002")
+    l1a_datasets = swe_l1a(imap_module_directory / test_data_path)
+
+    l1b_input = l1a_datasets[0]
+
+    # Set these two as cli.py -> post_processing would have.
+    l1b_input.attrs["Data_version"] = "v002"
+    # write data to CDF
+    l1a_cdf_filepath = write_cdf(l1b_input)
+    assert l1a_cdf_filepath.name == "imap_swe_l1a_sci_20240510_v002.cdf"
+
+    def get_file_paths_side_effect(descriptor):
+        if descriptor == "sci":
+            return [l1a_cdf_filepath]
+        elif descriptor == "l1b-in-flight-cal":
+            return [
+                imap_module_directory
+                / "tests/swe/lut/imap_swe_l1b-in-flight-cal_20240510_20260716_v000.csv"
+            ]
+        elif descriptor == "eu-conversion":
+            return [
+                imap_module_directory
+                / "tests/swe/lut/imap_swe_eu-conversion_20240510_v000.csv"
+            ]
+        elif descriptor == "esa-lut":
+            return [
+                imap_module_directory
+                / "tests/swe/lut/imap_swe_esa-lut_20250301_v000.csv"
+            ]
+        else:
+            raise ValueError(f"Unknown descriptor: {descriptor}")
+
+    mock_get_file_paths.side_effect = get_file_paths_side_effect
+    science_input = ScienceInput(l1a_cdf_filepath.name)
+    inflight_anc = AncillaryInput(
+        "imap_swe_l1b-in-flight-cal_20240510_20260716_v000.csv"
+    )
+    eu_anc = AncillaryInput("imap_swe_eu-conversion_20240510_v000.csv")
+    dependencies = ProcessingInputCollection(science_input, inflight_anc, eu_anc)
+    l1b_datasets = swe_l1b(dependencies)
 
-    l1b_dataset = swe_l1b(l1a_datasets[0], "002")
+    l1b_write_ds = l1b_datasets[0]
+    l1b_write_ds.attrs["Data_version"] = "v002"
 
-    sci_l1b_filepath = write_cdf(l1b_dataset[0])
+    sci_l1b_filepath = write_cdf(l1b_write_ds)
 
     assert sci_l1b_filepath.name == "imap_swe_l1b_sci_20240510_v002.cdf"
     # load the CDF file and compare the values
@@ -81,9 +218,10 @@ def test_cdf_creation(mock_read_in_flight_cal_data, l1b_validation_df):
 
 def test_count_rate():
     x = np.array([1, 10, 100, 1000, 10000, 38911, 65535])
-    acq_duration = 80000
+    acq_duration = np.array([80000])
     deatime_corrected = deadtime_correction(x, acq_duration)
     count_rate = convert_counts_to_rate(deatime_corrected, acq_duration)
+    count_rate = count_rate.flatten()  # Ensure the shape matches expected_output
     # Ruth provided the expected output for this test
     expected_output = [
         12.50005653,

imap_processing/tests/swe/test_swe_l2.py

@@ -1,9 +1,13 @@
 from unittest.mock import patch
 
 import numpy as np
-import pandas as pd
 import pytest
 import xarray as xr
+from imap_data_access.processing_input import (
+    AncillaryInput,
+    ProcessingInputCollection,
+    ScienceInput,
+)
 
 from imap_processing import imap_module_directory
 from imap_processing.cdf.utils import write_cdf
@@ -13,41 +17,17 @@ from imap_processing.swe.l2.swe_l2 import (
     calculate_flux,
     calculate_phase_space_density,
     find_angle_bin_indices,
-    get_particle_energy,
     put_data_into_angle_bins,
     swe_l2,
 )
 from imap_processing.swe.utils import swe_constants
-from imap_processing.swe.utils.swe_utils import (
-    read_lookup_table,
-)
-
-
-def test_get_particle_energy():
-    """Test get_particle_energy function."""
-    all_energy = get_particle_energy()
-    expected_energy = (
-        read_lookup_table()["esa_v"].values * swe_constants.ENERGY_CONVERSION_FACTOR
-    )
-    np.testing.assert_array_equal(all_energy["energy"], expected_energy)
 
 
 @patch(
     "imap_processing.swe.utils.swe_constants.GEOMETRIC_FACTORS",
     new=np.full(swe_constants.N_CEMS, 1),
 )
-@patch(
-    "imap_processing.swe.l2.swe_l2.get_particle_energy",
-    return_value=pd.DataFrame(
-        {
-            "table_index": np.repeat([0, 1], 720),
-            "e_step": np.tile(np.arange(720), 2),
-            "esa_v": np.repeat([1, 2], 720),
-            "energy": np.repeat([1, 2], 720),
-        }
-    ),
-)
-def test_calculate_phase_space_density(patch_get_particle_energy):
+def test_calculate_phase_space_density():
     """Test calculate_phase_space_density function."""
     # Create a dummy l1b dataset
     total_sweeps = 2
@@ -66,20 +46,21 @@ def test_calculate_phase_space_density(patch_get_particle_energy):
                     1,
                 ),
             ),
-            "acq_duration": (
-                ["epoch", "cycle"],
-                np.full((total_sweeps, swe_constants.N_QUARTER_CYCLES), 80.0),
-            ),
-            "esa_table_num": (
-                ["epoch", "cycle"],
-                np.repeat([0, 1], swe_constants.N_QUARTER_CYCLES).reshape(
-                    total_sweeps, swe_constants.N_QUARTER_CYCLES
+            "esa_energy": (
+                ["epoch", "energy", "angle"],
+                np.full(
+                    (
+                        total_sweeps,
+                        swe_constants.N_ESA_STEPS,
+                        swe_constants.N_ANGLE_SECTORS,
+                    ),
+                    1,
                 ),
             ),
         }
     )
-    phase_space_density_ds = calculate_phase_space_density(l1b_dataset)
-    assert phase_space_density_ds["phase_space_density"].shape == (
+    phase_space_density = calculate_phase_space_density(l1b_dataset)
+    assert phase_space_density.shape == (
         total_sweeps,
         swe_constants.N_ESA_STEPS,
         swe_constants.N_ANGLE_SECTORS,
@@ -102,68 +83,40 @@ def test_calculate_phase_space_density(patch_get_particle_energy):
         ),
         expected_calculated_density,
     )
-    np.testing.assert_array_equal(
-        phase_space_density_ds["phase_space_density"][0].data, expected_density
-    )
+    np.testing.assert_array_equal(phase_space_density[0].data, expected_density)
 
-    # Test that second sweep has correct values, similar to first sweep,
-    # but with energy 2.
-    expected_calculated_density = (2 * 1) / (
-        1 * swe_constants.VELOCITY_CONVERSION_FACTOR * 2**2
-    )
-    expected_density = np.full(
+
+def test_calculate_flux():
+    """Test calculate_flux function."""
+    # Create a dummy l1b dataset
+    total_sweeps = 2
+    phase_space_density = np.full(
         (
+            total_sweeps,
             swe_constants.N_ESA_STEPS,
             swe_constants.N_ANGLE_SECTORS,
             swe_constants.N_CEMS,
         ),
-        expected_calculated_density,
+        1,
     )
-    np.testing.assert_array_equal(
-        phase_space_density_ds["phase_space_density"][1].data, expected_density
-    )
-    assert type(phase_space_density_ds) == xr.Dataset
 
-
-def test_calculate_flux():
-    """Test calculate_flux function."""
-    # Create a dummy l1b dataset
-    total_sweeps = 2
-    l1b_dataset = xr.Dataset(
-        {
-            "science_data": (
-                ["epoch", "energy", "angle", "cem"],
-                np.full(
-                    (
-                        total_sweeps,
-                        swe_constants.N_ESA_STEPS,
-                        swe_constants.N_ANGLE_SECTORS,
-                        swe_constants.N_CEMS,
-                    ),
-                    1,
-                ),
-            ),
-            "acq_duration": (
-                ["epoch", "cycle"],
-                np.full((total_sweeps, swe_constants.N_QUARTER_CYCLES), 80.0),
-            ),
-            "esa_table_num": (
-                ["epoch", "cycle"],
-                np.repeat([0, 1], swe_constants.N_QUARTER_CYCLES).reshape(
-                    total_sweeps, swe_constants.N_QUARTER_CYCLES
-                ),
-            ),
-        }
+    esa_energy = np.full(
+        (
+            total_sweeps,
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_SECTORS,
+        ),
+        1,
     )
 
-    flux = calculate_flux(l1b_dataset)
+    flux = calculate_flux(phase_space_density, esa_energy)
     assert flux.shape == (
         total_sweeps,
         swe_constants.N_ESA_STEPS,
         swe_constants.N_ANGLE_SECTORS,
         swe_constants.N_CEMS,
     )
-    assert type(flux) == np.ndarray
+    assert isinstance(flux, np.ndarray)
 
 
 def test_find_angle_bin_indices():
@@ -298,35 +251,54 @@ def test_put_data_into_angle_bins():
     np.testing.assert_array_equal(even_col_mean_data, expected_mean_data)
 
 
-@patch(
-    "imap_processing.swe.l1b.swe_l1b_science.read_in_flight_cal_data",
-    return_value=pd.DataFrame(
-        {
-            "met_time": [453050300, 453077900],
-            "cem1": [1, 1],
-            "cem2": [1, 1],
-            "cem3": [1, 1],
-            "cem4": [1, 1],
-            "cem5": [1, 1],
-            "cem6": [1, 1],
-            "cem7": [1, 1],
-        }
-    ),
-)
+@patch("imap_data_access.processing_input.ProcessingInputCollection.get_file_paths")
 @pytest.mark.usefixtures("use_fake_spin_data_for_time")
-def test_swe_l2(mock_read_in_flight_cal_data, use_fake_spin_data_for_time):
+def test_swe_l2(mock_get_file_paths, use_fake_spin_data_for_time):
     """Test L2 processing."""
     data_start_time = 453051293.099714
     data_end_time = 453070000.0
     use_fake_spin_data_for_time(data_start_time, data_end_time)
 
     test_data_path = "tests/swe/l0_data/2024051010_SWE_SCIENCE_packet.bin"
-    l1a_datasets = swe_l1a(imap_module_directory / test_data_path, "002")
+    l1a_datasets = swe_l1a(imap_module_directory / test_data_path)
+    l1a_ds = l1a_datasets[0]
+    l1a_ds.attrs["Data_version"] = "v000"
+    l1a_cdf_filepath = write_cdf(l1a_ds)
+    assert l1a_cdf_filepath.name == "imap_swe_l1a_sci_20240510_v000.cdf"
+
+    def get_file_paths_side_effect(descriptor):
+        if descriptor == "sci":
+            return [l1a_cdf_filepath]
+        elif descriptor == "l1b-in-flight-cal":
+            return [
+                imap_module_directory
+                / "tests/swe/lut/imap_swe_l1b-in-flight-cal_20240510_20260716_v000.csv"
+            ]
+        elif descriptor == "eu-conversion":
+            return [
+                imap_module_directory
+                / "tests/swe/lut/imap_swe_eu-conversion_20240510_v000.csv"
+            ]
+        elif descriptor == "esa-lut":
+            return [
+                imap_module_directory
+                / "tests/swe/lut/imap_swe_esa-lut_20250301_v000.csv"
+            ]
+        else:
+            raise ValueError(f"Unknown descriptor: {descriptor}")
 
-    l1b_dataset = swe_l1b(l1a_datasets[0], "002")
-    l2_dataset = swe_l2(l1b_dataset[0], "002")
+    mock_get_file_paths.side_effect = get_file_paths_side_effect
+    science_input = ScienceInput(l1a_cdf_filepath.name)
+    inflight_anc = AncillaryInput(
+        "imap_swe_l1b-in-flight-cal_20240510_20260716_v000.csv"
+    )
+    eu_anc = AncillaryInput("imap_swe_eu-conversion_20240510_v000.csv")
+    dependencies = ProcessingInputCollection(science_input, inflight_anc, eu_anc)
+    l1b_dataset = swe_l1b(dependencies)[0]
+    l1b_dataset.attrs["Data_version"] = "v000"
+    l2_dataset = swe_l2(l1b_dataset)
 
-    assert type(l2_dataset) == xr.Dataset
+    assert isinstance(l2_dataset, xr.Dataset)
     assert l2_dataset["phase_space_density_spin_sector"].shape == (
         6,
         swe_constants.N_ESA_STEPS,
@@ -345,6 +317,14 @@ def test_swe_l2(mock_read_in_flight_cal_data, use_fake_spin_data_for_time):
         swe_constants.N_ANGLE_SECTORS,
     )
 
+    rate = l1b_dataset.science_data.to_numpy()
+    psd = l2_dataset.phase_space_density_spin_sector.to_numpy()
+    rate = rate[2, :, :, 3]  # nonzero counts at all energy & spin
+    psd = psd[2, :, :, 3]
+    cal_factor = psd / rate  # same CEM, should be constant at a given energy
+    assert np.allclose(cal_factor, cal_factor[:, 0:1], rtol=1e-9, atol=0)
+
     # Write L2 to CDF
+    l2_dataset.attrs["Data_version"] = "v002"
     l2_cdf_filepath = write_cdf(l2_dataset)
     assert l2_cdf_filepath.name == "imap_swe_l2_sci_20240510_v002.cdf"