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

imap_processing/tests/ialirt/unit/test_parse_mag.py

@@ -0,0 +1,168 @@
+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.parse_mag import (
+    extract_magnetic_vectors,
+    get_pkt_counter,
+    get_status_data,
+    get_time,
+    parse_packet,
+)
+from imap_processing.utils import packet_file_to_datasets
+
+
+@pytest.fixture(scope="session")
+def xtce_mag_path():
+    """Returns the xtce directory."""
+    return imap_module_directory / "ialirt" / "packet_definitions" / "ialirt_mag.xml"
+
+
+@pytest.fixture(scope="session")
+def binary_packet_path():
+    """Returns the paths to the binary packets."""
+    directory = imap_module_directory / "tests" / "ialirt" / "test_data" / "l0"
+    filenames = [
+        "461971383-404.bin",
+        "461971384-405.bin",
+        "461971385-406.bin",
+        "461971386-407.bin",
+        "461971387-408.bin",
+        "461971388-409.bin",
+        "461971389-410.bin",
+        "461971390-411.bin",
+        "461971391-412.bin",
+    ]
+    return tuple(directory / fname for fname in filenames)
+
+
+@pytest.fixture(scope="session")
+def mag_test_data():
+    """Returns the test data directory."""
+    data_path = (
+        imap_module_directory
+        / "tests"
+        / "ialirt"
+        / "test_data"
+        / "l0"
+        / "sample_decoded_i-alirt_data.csv"
+    )
+    data = pd.read_csv(data_path)
+
+    return data
+
+
+@pytest.fixture()
+def xarray_data(binary_packet_path, xtce_mag_path):
+    """Create xarray data for multiple packets."""
+    apid = 1001
+
+    xarray_data = tuple(
+        packet_file_to_datasets(packet, xtce_mag_path, use_derived_value=False)[apid]
+        for packet in binary_packet_path
+    )
+
+    merged_xarray_data = xr.concat(xarray_data, dim="epoch")
+    return merged_xarray_data
+
+
+@pytest.fixture()
+def grouped_data():
+    """Creates grouped data for tests."""
+    epoch = np.arange(12)
+
+    # Example `src_seq_ctr` values for 3 groups:
+    # Group 0 - valid, all diffs = 1
+    # Group 1 - invalid, has a jump of 5
+    # Group 2 - valid, wraps at -16383
+    src_seq_ctr = np.concatenate(
+        [
+            np.arange(100, 104),
+            np.array([200, 205, 206, 207]),
+            np.array([16382, 16383, 0, 1]),
+        ],
+        dtype=np.int32,
+    )
+    mag_acq_tm_coarse = np.repeat(
+        np.array([461971382, 461971386, 461971390], dtype=np.uint32), repeats=4
+    )
+
+    mag_acq_tm_fine = np.array(
+        [1502, 1502, 1505, 1505, 1500, 1500, 1503, 1503, 1497, 1497, 1491, 1491]
+    )
+
+    group = np.tile(np.arange(3), 4).reshape(4, 3).T.ravel()
+
+    grouped_data = xr.Dataset(
+        data_vars={
+            "src_seq_ctr": ("epoch", src_seq_ctr),
+            "mag_acq_tm_coarse": ("epoch", mag_acq_tm_coarse),
+            "mag_acq_tm_fine": ("epoch", mag_acq_tm_fine),
+        },
+        coords={"epoch": epoch, "group": ("epoch", group)},
+    )
+
+    return grouped_data
+
+
+def test_get_pkt_counter(xarray_data):
+    """Tests the get_pkt_counter function."""
+    status_values = xarray_data["mag_status"].values
+    pkt_counter = get_pkt_counter(status_values)
+    assert np.array_equal(pkt_counter, np.array([0, 1, 2, 3, 0, 1, 2, 3, 0]))
+
+
+def test_get_status_data(xarray_data, mag_test_data):
+    """Tests the get_status_data function."""
+
+    status_data = get_status_data(
+        xarray_data["mag_status"].values[0:4], np.array([0, 1, 2, 3])
+    )
+    index = mag_test_data["PRI_COARSETM"] == 461971382
+    matching_row = mag_test_data[index]
+
+    for key in status_data.keys():
+        assert status_data[key] == matching_row[key.upper()].values[0]
+
+
+def test_get_time(grouped_data):
+    """Tests the get_time function."""
+    time_data = get_time(grouped_data, 1, np.array([0, 1, 2, 3]))
+    assert time_data == {
+        "pri_coarsetm": 461971386,
+        "pri_fintm": 1500,
+        "sec_coarsetm": 461971386,
+        "sec_fintm": 1503,
+    }
+
+
+def test_extract_magnetic_vectors():
+    """Tests the extract_magnetic_vectors function."""
+    science_values = xr.DataArray(
+        data=np.array([15797207, 5750698, 15921110, 2342918], dtype=np.uint32)
+    )
+
+    vectors = extract_magnetic_vectors(science_values)
+
+    assert vectors == {
+        "pri_x": 61707,
+        "pri_y": 55127,
+        "pri_z": 49066,
+        "sec_x": 62191,
+        "sec_y": 54819,
+        "sec_z": 49158,
+    }
+
+
+def test_parse_packet(xarray_data, mag_test_data):
+    """Tests the parse_packet function."""
+    parsed_packets = parse_packet(xarray_data)
+
+    for packet in parsed_packets:
+        index = packet["pri_coarsetm"] == mag_test_data["PRI_COARSETM"]
+        matching_rows = mag_test_data[index]
+
+        for key in packet.keys():
+            assert packet[key] == matching_rows[key.upper()].values[0]

imap_processing/tests/ialirt/unit/test_process_swe.py

@@ -1,8 +1,23 @@
+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 (
+    decompress_counts,
+    get_ialirt_energies,
+    normalize_counts,
+    phi_to_bin,
+    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
 
 
@@ -52,9 +67,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 +106,19 @@ 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
+
+
+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 +132,180 @@ 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_norm_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)
+
+
+@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"])))
+    swe_data = process_swe(ds)
+
+    assert swe_data == []

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,5 +1,3 @@
-from pathlib import Path
-
 import numpy as np
 import pytest
 import xarray as xr
@@ -7,6 +5,13 @@ import xarray as xr
 from imap_processing import imap_module_directory
 from imap_processing.idex.idex_l1a import PacketParser
 
+TEST_DATA_PATH = imap_module_directory / "tests" / "idex" / "test_data"
+
+TEST_L0_FILE = TEST_DATA_PATH / "imap_idex_l0_raw_20231218_v001.pkts"
+L1A_EXAMPLE_FILE = TEST_DATA_PATH / "idex_l1a_validation_file.h5"
+
+pytestmark = pytest.mark.external_test_data
+
 SPICE_ARRAYS = [
     "ephemeris_position_x",
     "ephemeris_position_y",
@@ -30,10 +35,20 @@ def decom_test_data() -> xr.Dataset:
     dataset : xarray.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"
-    )
-    return PacketParser(test_file, "001").data
+    return PacketParser(TEST_L0_FILE, "001").data
+
+
+@pytest.fixture(scope="session")
+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)
 
 
 def get_spice_data_side_effect_func(l1a_ds, idex_attrs):
@@ -48,3 +63,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.dims.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