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

imap_processing/tests/ialirt/unit/test_process_swe.py

@@ -1,8 +1,32 @@
+"""Tests to support I-ALiRT SWE packet parsing."""
+
+from unittest.mock import patch
+
 import numpy as np
 import pandas as pd
 import pytest
+import xarray as xr
 
 from imap_processing import imap_module_directory
+from imap_processing.ialirt.l0.process_swe import (
+    average_counts,
+    azimuthal_check_counterstreaming,
+    compute_bidirectional,
+    decompress_counts,
+    determine_streaming,
+    find_bin_offsets,
+    find_min_counts,
+    get_ialirt_energies,
+    normalize_counts,
+    phi_to_bin,
+    polar_check_counterstreaming,
+    prepare_raw_counts,
+    process_swe,
+)
+from imap_processing.swe.utils.swe_constants import (
+    ESA_VOLTAGE_ROW_INDEX_DICT,
+    GEOMETRIC_FACTORS,
+)
 from imap_processing.utils import packet_file_to_datasets
 
 
@@ -19,7 +43,7 @@ def binary_packet_path():
         imap_module_directory
         / "tests"
         / "ialirt"
-        / "test_data"
+        / "data"
         / "l0"
         / "20240827095047_SWE_IALIRT_packet.bin"
     )
@@ -32,7 +56,7 @@ def swe_test_data():
         imap_module_directory
         / "tests"
         / "ialirt"
-        / "test_data"
+        / "data"
         / "l0"
         / "idle_export_eu.SWE_IALIRT_20240827_093852.csv"
     )
@@ -41,7 +65,7 @@ def swe_test_data():
     return data
 
 
-@pytest.fixture()
+@pytest.fixture
 def xarray_data(binary_packet_path, xtce_swe_path):
     """Create xarray data"""
     apid = 1360
@@ -52,9 +76,9 @@ def xarray_data(binary_packet_path, xtce_swe_path):
     return xarray_data
 
 
-def test_decom_packets(xarray_data, swe_test_data):
-    """This function checks that all instrument parameters are accounted for."""
-
+@pytest.fixture
+def fields_to_test():
+    """Create a dictionary to convert names"""
     fields_to_test = {
         "swe_shcoarse": "SHCOARSE",
         "swe_acq_sec": "ACQUISITION_TIME",
@@ -91,6 +115,38 @@ def test_decom_packets(xarray_data, swe_test_data):
         "swe_cem7_e3": "ELEC_COUNTS_SPIN_I_POL_6_E_2J",
         "swe_cem7_e4": "ELEC_COUNTS_SPIN_I_POL_6_E_3J",
     }
+    return fields_to_test
+
+
+@pytest.fixture
+def summed_half_cycle():
+    """Create test set with known peaks"""
+
+    summed_half_cycle = np.zeros((8, 30))
+
+    for i in range(8):
+        peak = i + 5
+        summed_half_cycle[i, peak] = 100
+        summed_half_cycle[i, (peak + 6) % 30] = 60  # +90 offset
+        summed_half_cycle[i, (peak + 8) % 30] = 80  # +90 offset
+        summed_half_cycle[i, (peak + 14) % 30] = 20  # +180 offset
+        summed_half_cycle[i, (peak + 16) % 30] = 40  # +180 offset
+        summed_half_cycle[i, (peak - 6) % 30] = 5  # -90 offset
+        summed_half_cycle[i, (peak - 8) % 30] = 15  # -90 offset
+
+    return summed_half_cycle
+
+
+def test_get_energy():
+    """Tests get_alirt_energies function."""
+    energies = get_ialirt_energies()
+
+    for i in range(len(energies)):
+        assert i + 11 == ESA_VOLTAGE_ROW_INDEX_DICT[energies[i]]
+
+
+def test_decom_packets(xarray_data, swe_test_data, fields_to_test):
+    """This function checks that all instrument parameters are accounted for."""
     _, index, test_index = np.intersect1d(
         xarray_data["swe_shcoarse"], swe_test_data["SHCOARSE"], return_indices=True
     )
@@ -104,3 +160,260 @@ def test_decom_packets(xarray_data, swe_test_data):
             f"Mismatch found in {xarray_field}: "
             f"actual {actual_values}, expected {expected_values}"
         )
+
+
+def test_decompress_counts():
+    """Test that we get correct decompressed counts from the algorithm."""
+    expected_value = 24063
+    input_count = 230
+    returned_value = decompress_counts(np.array([input_count]))
+    assert np.all(expected_value == returned_value)
+
+
+def test_phi_to_bin():
+    """Test phi_to_bin function."""
+
+    # Define expected phi-to-bin mapping for one full spin
+    phis = np.arange(12, 361, 12).tolist()
+
+    expected_bins = np.arange(30)
+
+    for phi, expected_bin in zip(phis, expected_bins):
+        assert phi_to_bin(phi) == expected_bin
+
+
+def test_prepare_raw_counts():
+    """Test that prepare_raw_counts correctly bins counts into (30, 7, 4) array."""
+
+    # 2 rows = 4 phis (12, 24, 36, 48)
+    epochs = [0, 1]
+
+    data = {
+        "group": ("epoch", [1, 1]),  # Both rows belong to group 1
+        # CEM 1 (Phi 12, 24, 36, 48)
+        "swe_cem1_e1": ("epoch", [1, 9]),
+        "swe_cem1_e2": ("epoch", [2, 10]),
+        "swe_cem1_e3": ("epoch", [3, 11]),
+        "swe_cem1_e4": ("epoch", [4, 12]),
+        # CEM 2
+        "swe_cem2_e1": ("epoch", [5, 13]),
+        "swe_cem2_e2": ("epoch", [6, 14]),
+        "swe_cem2_e3": ("epoch", [7, 15]),
+        "swe_cem2_e4": ("epoch", [8, 16]),
+        "swe_seq": ("epoch", [0, 1]),
+    }
+
+    grouped_data = xr.Dataset(data, coords={"epoch": epochs})
+    group_mask = grouped_data["group"] == 1
+    grouped = grouped_data.sel(epoch=group_mask)
+
+    raw_counts = prepare_raw_counts(grouped, cem_number=2)
+
+    # Expected shape (8, 7, 30) but only some CEMs are used.
+    expected = np.zeros((8, 2, 30), dtype=np.uint8)
+
+    # Phi bins for 12, 24, 36, 48)
+    phi_bin_12 = 0  # Phi 12
+    phi_bin_24 = 1  # Phi 24
+    phi_bin_36 = 2  # Phi 36
+    phi_bin_48 = 3  # Phi 48
+
+    # CEM 1, Phi 12 (E1, E2)
+    expected[1, 0, phi_bin_12] = 1
+    expected[5, 0, phi_bin_12] = 2
+
+    # CEM 1, Phi 24 (E3, E4)
+    expected[7, 0, phi_bin_24] = 3
+    expected[3, 0, phi_bin_24] = 4
+
+    # CEM 1, Phi 36 (E1, E2)
+    expected[1, 0, phi_bin_36] = 9
+    expected[5, 0, phi_bin_36] = 10
+
+    # CEM 1, Phi 48 (E3, E4)
+    expected[7, 0, phi_bin_48] = 11
+    expected[3, 0, phi_bin_48] = 12
+
+    # CEM 2, Phi 12 (E1, E2)
+    expected[1, 1, phi_bin_12] = 5
+    expected[5, 1, phi_bin_12] = 6
+
+    # CEM 2, Phi 24 (E3, E4)
+    expected[7, 1, phi_bin_24] = 7
+    expected[3, 1, phi_bin_24] = 8
+
+    # CEM 2, Phi 36 (E1, E2)
+    expected[1, 1, phi_bin_36] = 13
+    expected[5, 1, phi_bin_36] = 14
+
+    # CEM 2, Phi 48 (E3, E4)
+    expected[7, 1, phi_bin_48] = 15
+    expected[3, 1, phi_bin_48] = 16
+
+    assert np.array_equal(raw_counts, expected)
+
+
+def test_normalize_counts():
+    """Tests normalize_counts function"""
+
+    # Shape (2, 7, 3) for a small test case
+    corrected_counts = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+                [10, 11, 12],
+                [13, 14, 15],
+                [16, 17, 18],
+                [19, 20, 21],
+            ],
+            [
+                [2, 4, 6],
+                [8, 10, 12],
+                [14, 16, 18],
+                [20, 22, 24],
+                [26, 28, 30],
+                [32, 34, 36],
+                [38, 40, 42],
+            ],
+        ],
+        dtype=np.uint8,
+    )
+
+    latest_cal = pd.Series(
+        [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0],
+        index=[
+            "cem1",
+            "cem2",
+            "cem3",
+            "cem4",
+            "cem5",
+            "cem6",
+            "cem7",
+        ],  # Simulating real data structure
+        dtype=np.float64,
+    )
+    expected = np.zeros((2, 7, 3), dtype=np.float64)
+
+    for i in range(2):
+        for j in range(7):
+            for k in range(3):
+                if corrected_counts[i][j][k] < 0:
+                    expected[i][j][k] = 0.0
+                else:
+                    expected[i][j][k] = (
+                        corrected_counts[i][j][k] * latest_cal[j] / GEOMETRIC_FACTORS[j]
+                    )
+
+    norm_counts = normalize_counts(corrected_counts, latest_cal)
+
+    assert np.allclose(norm_counts, expected, atol=1e-9)
+
+
+def test_find_bin_offsets():
+    """Tests find_bin_offsets function"""
+
+    peak_bins = np.array([5, 6, 29])
+    bins = find_bin_offsets(peak_bins, (2, 3))
+    np.testing.assert_array_equal(bins, np.array([[7, 8, 1], [8, 9, 2]]))
+
+
+def test_average_counts(summed_half_cycle):
+    """Tests average_values_and_azimuth function"""
+
+    # Find the azimuth angle that corresponds to the maximum counts at each energy
+    peak_az_bin = np.argmax(summed_half_cycle, axis=1)
+
+    # Bins +6 and +8 correspond to +90 degrees.
+    counts_90 = average_counts(peak_az_bin, summed_half_cycle, (6, 8))
+
+    assert np.allclose(counts_90, np.full(8, 70), atol=1e-9)
+
+
+def test_find_min_counts(summed_half_cycle):
+    """Tests find_min_counts function"""
+
+    cpeak, cmin, counts = find_min_counts(summed_half_cycle)
+    np.testing.assert_array_equal(cpeak, np.full(8, 100))
+    np.testing.assert_array_equal(cmin, counts[0])
+
+
+def test_determine_streaming_summed_cems():
+    """Tests determine_streaming_summed_cems function."""
+
+    # Case where streaming should be True
+    cpeak = np.array([100, 80, 50, 60])
+    cmin = np.array([20, 70, 40, 60])
+    counts_180 = np.array([40, 60, 90, 110])
+    assert np.array_equal(
+        determine_streaming(cpeak, counts_180, cmin), np.array([1, 0, 0, 0])
+    )
+
+
+def test_compute_bidirectional():
+    """Tests compute_bidirectional function."""
+
+    first_half = np.array([1, 0, 0, 0, 1, 0, 0, 0])
+    second_half = np.array([1, 0, 0, 0, 1, 0, 0, 0])
+    assert compute_bidirectional(first_half, second_half) == (0, 0)
+
+    first_half = np.array([1, 1, 1, 0, 0, 0, 0, 0])
+    second_half = np.array([0, 0, 0, 0, 0, 0, 0, 0])
+    assert compute_bidirectional(first_half, second_half) == (1, 0)
+
+
+def test_azimuthal_check_counterstreaming(summed_half_cycle):
+    """Tests azimuthal_check_counterstreaming function."""
+
+    bde = azimuthal_check_counterstreaming(summed_half_cycle, summed_half_cycle)
+
+    assert bde == (1, 1)
+
+
+def test_polar_check_counterstreaming():
+    """Tests polar_check_counterstreaming function."""
+
+    # cem0 (cem1) and cem6 (cem7) have high values
+    # cem2, cem3, cem4 (cem3 to cem5) are low and used for cmin
+    row = np.array([100, 20, 5, 5, 5, 20, 100])
+    summed_half = np.tile(row, (8, 1))
+
+    bde = polar_check_counterstreaming(summed_half, summed_half)
+
+    assert bde == (1, 1)
+
+
+@patch(
+    "imap_processing.ialirt.l0.process_swe.read_in_flight_cal_data",
+    return_value=pd.DataFrame(
+        {
+            "met_time": [453051300, 453051900],
+            "cem1": [1, 2],
+            "cem2": [1, 2],
+            "cem3": [1, 2],
+            "cem4": [1, 2],
+            "cem5": [1, 2],
+            "cem6": [1, 2],
+            "cem7": [1, 2],
+        }
+    ),
+)
+def test_process_swe(mock_read_cal, swe_test_data, fields_to_test):
+    """Test processing for swe."""
+    swe_test_data = swe_test_data.rename(
+        columns={v: k for k, v in fields_to_test.items()}
+    )
+    swe_test_data.index.name = "epoch"
+    ds = swe_test_data.to_xarray()
+    ds["src_seq_ctr"] = ("epoch", np.arange(len(ds["swe_shcoarse"])))
+    in_flight_cal_file = (
+        imap_module_directory
+        / "tests/swe/lut/imap_swe_l1b-in-flight-cal_20240510_20260716_v000.csv"
+    )
+    swe_data = process_swe(ds, [in_flight_cal_file])
+
+    # TODO: add tests with test data here.
+
+    # Check that all groups in the data are accounted for.
+    assert len(swe_data) == 912 // 60

imap_processing/tests/ialirt/unit/test_time.py

@@ -0,0 +1,16 @@
+"""Tests grouping functions for I-ALiRT instruments."""
+
+import numpy as np
+
+from imap_processing.ialirt.utils.time import calculate_time
+
+
+def test_calculate_time():
+    """Tests calculate_time function."""
+
+    coarse = np.array([1, 2, 3, 4, 5])
+    fine = np.array([0, 1, 2, 3, 4])
+
+    time = calculate_time(coarse, fine, 6553)
+
+    np.testing.assert_allclose(time, coarse + fine / 65535.0, atol=1e-03, rtol=0)

imap_processing/tests/idex/conftest.py

@@ -1,4 +1,4 @@
-from pathlib import Path
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -6,6 +6,19 @@ import xarray as xr
 
 from imap_processing import imap_module_directory
 from imap_processing.idex.idex_l1a import PacketParser
+from imap_processing.idex.idex_l1b import idex_l1b
+from imap_processing.idex.idex_l2a import idex_l2a
+
+TEST_DATA_PATH = imap_module_directory / "tests" / "idex" / "test_data"
+
+TEST_L0_FILE_SCI = TEST_DATA_PATH / "imap_idex_l0_raw_20231218_v001.pkts"
+TEST_L0_FILE_EVT = TEST_DATA_PATH / "imap_idex_l0_raw_20250108_v001.pkts"  # 1418
+TEST_L0_FILE_CATLST = TEST_DATA_PATH / "imap_idex_l0_raw_20241206_v001.pkts"  # 1419
+
+L1A_EXAMPLE_FILE = TEST_DATA_PATH / "idex_l1a_validation_file.h5"
+L1B_EXAMPLE_FILE = TEST_DATA_PATH / "idex_l1b_validation_file.h5"
+
+pytestmark = pytest.mark.external_test_data
 
 SPICE_ARRAYS = [
     "ephemeris_position_x",
@@ -21,19 +34,86 @@ SPICE_ARRAYS = [
 ]
 
 
-@pytest.fixture(scope="module")
-def decom_test_data() -> xr.Dataset:
+@pytest.fixture
+def decom_test_data_sci() -> xr.Dataset:
     """Return a ``xarray`` dataset containing test data.
 
     Returns
     -------
     dataset : xarray.Dataset
+        A ``xarray`` dataset containing the science test data
+    """
+    return PacketParser(TEST_L0_FILE_SCI).data[0]
+
+
+@pytest.fixture
+def decom_test_data_catlst() -> xr.Dataset:
+    """Return a ``xarray`` dataset containing the catalog list summary data.
+
+    Returns
+    -------
+    dataset : xarray.Dataset
+        A ``xarray`` dataset containing the catalog list summary data.
+    """
+    return PacketParser(TEST_L0_FILE_CATLST).data[0]
+
+
+@pytest.fixture
+def decom_test_data_evt() -> xr.Dataset:
+    """Return a ``xarray`` dataset containing the event log data.
+
+    Returns
+    -------
+    dataset : xarray.Dataset
+        A ``xarray`` dataset containing the event log data.
+    """
+    return PacketParser(TEST_L0_FILE_EVT).data[0]
+
+
+@pytest.fixture
+def l1a_example_data(_download_test_data):
+    """
+    Pytest fixture to load example L1A data (produced by the IDEX team) for testing.
+
+    Returns
+    -------
+    dict
+      A dictionary containing the 6 waveform and telemetry arrays
+    """
+    return load_hdf_file(L1A_EXAMPLE_FILE)
+
+
+@pytest.fixture
+def l2a_dataset(decom_test_data_sci: xr.Dataset) -> xr.Dataset:
+    """Return a ``xarray`` dataset containing test data.
+
+    Returns
+    -------
+    dataset : xr.Dataset
         A ``xarray`` dataset containing the test data
     """
-    test_file = Path(
-        f"{imap_module_directory}/tests/idex/test_data/imap_idex_l0_raw_20231214_v001.pkts"
+    spin_phase_angles = xr.DataArray(
+        np.random.randint(0, 360, len(decom_test_data_sci.epoch))
     )
-    return PacketParser(test_file, "001").data
+    with mock.patch(
+        "imap_processing.idex.idex_l1b.get_spice_data",
+        return_value={"spin_phase": spin_phase_angles},
+    ):
+        dataset = idex_l2a(idex_l1b(decom_test_data_sci))
+    return dataset
+
+
+@pytest.fixture
+def l1b_example_data(_download_test_data):
+    """
+    Pytest fixture to load example L1B data (produced by the IDEX team) for testing.
+
+    Returns
+    -------
+    dict
+      A dictionary containing the 6 waveform and telemetry arrays
+    """
+    return load_hdf_file(L1B_EXAMPLE_FILE)
 
 
 def get_spice_data_side_effect_func(l1a_ds, idex_attrs):
@@ -48,3 +128,44 @@ def get_spice_data_side_effect_func(l1a_ds, idex_attrs):
         )
         for name in SPICE_ARRAYS
     }
+
+
+def load_hdf_file(path: str) -> xr.Dataset:
+    """
+    Loads an HDF5 file produced by the IDEX team into a dataset.
+
+    Parameters
+    ----------
+    path : str
+        The file path to the HDF5 file.
+
+    Returns
+    -------
+    dataset
+        A dataset containing the extracted data.
+    """
+    # Load hdf5 data into a datatree
+    datatree = xr.open_datatree(path, engine="netcdf4")
+    datasets = []
+    # Sort datatree by the event number
+    datatree = sorted(datatree.items(), key=lambda x: int(x[0]))
+    # Iterate through every nested tree in the datatree (Each nested tree represents
+    # data from one event).
+    # Rename the dimensions across every tree to be the same
+    # Add an "event" dimension which will allow them all to be concatenated together.
+    for event, tree in datatree:
+        event_num = int(event)
+        # Extract the metadata
+        metadata = tree.Metadata.to_dataset().expand_dims({"event": [event_num]})
+        ds = tree.to_dataset()
+        # Sort dimensions by shape. The high sampling time dimension is always less
+        # than the low sampling time.
+        dims = [k for k, v in sorted(ds.sizes.items(), key=lambda item: item[1])]
+        ds = ds.rename({dims[0]: "time_low", dims[1]: "time_high"}).expand_dims(
+            {"event": [event_num]}
+        )
+        datasets.append(xr.merge([ds, metadata]))
+
+    example_dataset = xr.concat(datasets, dim="event")
+
+    return example_dataset