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

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,47 +1,33 @@
 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
 from imap_processing.swe.l1a.swe_l1a import swe_l1a
 from imap_processing.swe.l1b.swe_l1b import swe_l1b
 from imap_processing.swe.l2.swe_l2 import (
-    ENERGY_CONVERSION_FACTOR,
-    VELOCITY_CONVERSION_FACTOR,
     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.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 * ENERGY_CONVERSION_FACTOR
-    np.testing.assert_array_equal(all_energy["energy"], expected_energy)
+from imap_processing.swe.utils import swe_constants
 
 
-@patch("imap_processing.swe.l2.swe_l2.GEOMETRIC_FACTORS", new=np.full(7, 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),
-        }
-    ),
+    "imap_processing.swe.utils.swe_constants.GEOMETRIC_FACTORS",
+    new=np.full(swe_constants.N_CEMS, 1),
 )
-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
@@ -50,21 +36,35 @@ def test_calculate_phase_space_density(patch_get_particle_energy):
         {
             "science_data": (
                 ["epoch", "energy", "angle", "cem"],
-                np.full((total_sweeps, 24, 30, 7), 1),
+                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, 4), 80.0)),
-            "esa_table_num": (
-                ["epoch", "cycle"],
-                np.repeat([0, 1], 4).reshape(total_sweeps, 4),
+            "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,
-        24,
-        30,
-        7,
+        swe_constants.N_ESA_STEPS,
+        swe_constants.N_ANGLE_SECTORS,
+        swe_constants.N_CEMS,
     )
 
     # Test that first sweep has correct values. In patch,
@@ -72,43 +72,51 @@ def test_calculate_phase_space_density(patch_get_particle_energy):
     #   2. we have set energy to 1.
     #   3. we have set science_data to 1.
     # Using this in the formula, we calculate expected density value.
-    expected_calculated_density = (2 * 1) / (1 * VELOCITY_CONVERSION_FACTOR * 1**2)
-    expected_density = np.full((24, 30, 7), expected_calculated_density)
-    np.testing.assert_array_equal(
-        phase_space_density_ds["phase_space_density"][0].data, expected_density
+    expected_calculated_density = (2 * 1) / (
+        1 * swe_constants.VELOCITY_CONVERSION_FACTOR * 1**2
     )
-
-    # Test that second sweep has correct values, similar to first sweep,
-    # but with energy 2.
-    expected_calculated_density = (2 * 1) / (1 * VELOCITY_CONVERSION_FACTOR * 2**2)
-    expected_density = np.full((24, 30, 7), expected_calculated_density)
-    np.testing.assert_array_equal(
-        phase_space_density_ds["phase_space_density"][1].data, expected_density
+    expected_density = np.full(
+        (
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_SECTORS,
+            swe_constants.N_CEMS,
+        ),
+        expected_calculated_density,
     )
-    assert type(phase_space_density_ds) == xr.Dataset
+    np.testing.assert_array_equal(phase_space_density[0].data, expected_density)
 
 
 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, 24, 30, 7), 1),
-            ),
-            "acq_duration": (["epoch", "cycle"], np.full((total_sweeps, 4), 80.0)),
-            "esa_table_num": (
-                ["epoch", "cycle"],
-                np.repeat([0, 1], 4).reshape(total_sweeps, 4),
-            ),
-        }
+    phase_space_density = np.full(
+        (
+            total_sweeps,
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_SECTORS,
+            swe_constants.N_CEMS,
+        ),
+        1,
+    )
+
+    esa_energy = np.full(
+        (
+            total_sweeps,
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_SECTORS,
+        ),
+        1,
     )
 
-    flux = calculate_flux(l1b_dataset)
-    assert flux.shape == (total_sweeps, 24, 30, 7)
-    assert type(flux) == np.ndarray
+    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 isinstance(flux, np.ndarray)
 
 
 def test_find_angle_bin_indices():
@@ -142,18 +150,19 @@ def test_find_angle_bin_indices():
 def test_put_data_into_angle_bins():
     """Test put_data_into_angle_bins function."""
     num_cycles = 1
-    num_esa_step = 24
-    num_angle_bins = 30
-    num_cems = 7
     # Create test counts data to test
     # Find all even numbers in the range 0 to 30
     even_numbers = np.arange(0, 30, 2)
     # repeat it twice now to get:
     # [0, 0, 2, 2, ...., 28, 28]
     example_data = np.repeat(even_numbers, 2)
-    energy_angle_test_data = np.tile(example_data, (num_cycles, num_esa_step, 1))
+    energy_angle_test_data = np.tile(
+        example_data, (num_cycles, swe_constants.N_ESA_STEPS, 1)
+    )
     # Expand to include 7 CEMs by repeating across last dimension
-    test_data = np.repeat(energy_angle_test_data[..., np.newaxis], num_cems, axis=-1)
+    test_data = np.repeat(
+        energy_angle_test_data[..., np.newaxis], swe_constants.N_CEMS, axis=-1
+    )
 
     # Took this example from intermediate output from actual data
     angle_bins_example = [
@@ -190,16 +199,28 @@ def test_put_data_into_angle_bins():
     ]
     # Now data with every row to be same as angle_bins_example
     test_angle_bin_indices_data = np.full(
-        (num_cycles, num_esa_step, num_angle_bins), angle_bins_example
+        (num_cycles, swe_constants.N_ESA_STEPS, swe_constants.N_ANGLE_SECTORS),
+        angle_bins_example,
     )
 
     binned_data = put_data_into_angle_bins(test_data, test_angle_bin_indices_data)
-    assert binned_data.shape == (num_cycles, num_esa_step, num_angle_bins, num_cems)
+    assert binned_data.shape == (
+        num_cycles,
+        swe_constants.N_ESA_STEPS,
+        swe_constants.N_ANGLE_BINS,
+        swe_constants.N_CEMS,
+    )
 
     # Test that the binned data has correct values in correct bins by
     # checking that odd number columns are filled with nan
     expected_binned_data = np.full(
-        (num_cycles, num_esa_step, num_angle_bins, num_cems), np.nan
+        (
+            num_cycles,
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_BINS,
+            swe_constants.N_CEMS,
+        ),
+        np.nan,
     )
     np.testing.assert_array_equal(
         binned_data[0, 0, 1::2, 0], expected_binned_data[0, 0, 1::2, 0]
@@ -230,39 +251,80 @@ 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"
 
-    l1b_dataset = swe_l1b(l1a_datasets[0], "002")
-    l2_dataset = swe_l2(l1b_dataset[0], "002")
+    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_dataset = swe_l1b(dependencies)[0]
+    l1b_dataset.attrs["Data_version"] = "v000"
+    l2_dataset = swe_l2(l1b_dataset)
+
+    assert isinstance(l2_dataset, xr.Dataset)
+    assert l2_dataset["phase_space_density_spin_sector"].shape == (
+        6,
+        swe_constants.N_ESA_STEPS,
+        swe_constants.N_ANGLE_SECTORS,
+        swe_constants.N_CEMS,
+    )
+    assert l2_dataset["flux_spin_sector"].shape == (
+        6,
+        swe_constants.N_ESA_STEPS,
+        swe_constants.N_ANGLE_SECTORS,
+        swe_constants.N_CEMS,
+    )
+    assert l2_dataset["acquisition_time"].shape == (
+        6,
+        swe_constants.N_ESA_STEPS,
+        swe_constants.N_ANGLE_SECTORS,
+    )
 
-    assert type(l2_dataset) == xr.Dataset
-    assert l2_dataset["phase_space_density_spin_sector"].shape == (6, 24, 30, 7)
-    assert l2_dataset["flux_spin_sector"].shape == (6, 24, 30, 7)
-    assert l2_dataset["acquisition_time"].shape == (6, 24, 30)
+    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"