imap-processing 0.6.0__py3-none-any.whl

imap_processing/tests/ultra/unit/test_ultra_l1b_extended.py

@@ -0,0 +1,342 @@
+"""Tests Extended Raw Events for ULTRA L1b."""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from imap_processing.ultra.l1b.ultra_l1b_extended import (
+    CoinType,
+    StartType,
+    StopType,
+    calculate_etof_xc,
+    determine_species_pulse_height,
+    determine_species_ssd,
+    get_coincidence_positions,
+    get_ctof,
+    get_energy_pulse_height,
+    get_energy_ssd,
+    get_front_x_position,
+    get_front_y_position,
+    get_particle_velocity,
+    get_path_length,
+    get_ph_tof_and_back_positions,
+    get_ssd_back_position_and_tof_offset,
+    get_ssd_tof,
+)
+
+
+@pytest.fixture()
+def yf_fixture(de_dataset, events_fsw_comparison_theta_0):
+    """Fixture to compute and return yf and related data."""
+    df = pd.read_csv(events_fsw_comparison_theta_0)
+    df_filt = df[df["StartType"] != -1]
+
+    d, yf = get_front_y_position(
+        de_dataset["START_TYPE"].data, df_filt.Yb.values.astype("float")
+    )
+
+    return df_filt, d, yf
+
+
+def test_get_front_x_position(
+    de_dataset,
+    yf_fixture,
+):
+    """Tests get_front_x_position function."""
+
+    df_filt, _, _ = yf_fixture
+
+    xf = get_front_x_position(
+        de_dataset["START_TYPE"].data,
+        de_dataset["START_POS_TDC"].data,
+    )
+
+    assert xf == pytest.approx(df_filt["Xf"].astype("float"), 1e-5)
+
+
+def test_get_front_y_position(yf_fixture):
+    """Tests get_front_y_position function."""
+    df_filt, d, yf = yf_fixture
+
+    assert yf == pytest.approx(df_filt["Yf"].astype("float"), abs=1e-5)
+    assert d == pytest.approx(df_filt["d"].astype("float"), abs=1e-5)
+
+
+def test_get_path_length(de_dataset, yf_fixture):
+    """Tests get_path_length function."""
+
+    df_filt, d, yf = yf_fixture
+
+    test_xf = df_filt["Xf"].astype("float").values
+    test_yf = df_filt["Yf"].astype("float").values
+
+    test_xb = df_filt["Xb"].astype("float").values
+    test_yb = df_filt["Yb"].astype("float").values
+    r = get_path_length((test_xf, test_yf), (test_xb, test_yb), d)
+    assert r == pytest.approx(df_filt["r"].astype("float"), abs=1e-5)
+
+
+def test_get_ph_tof_and_back_positions(
+    de_dataset,
+    yf_fixture,
+):
+    """Tests get_ph_tof_and_back_positions function."""
+
+    df_filt, _, _ = yf_fixture
+
+    ph_tof, _, ph_xb, ph_yb = get_ph_tof_and_back_positions(
+        de_dataset, df_filt.Xf.astype("float").values, "ultra45"
+    )
+
+    ph_indices = np.nonzero(
+        np.isin(de_dataset["STOP_TYPE"], [StopType.Top.value, StopType.Bottom.value])
+    )[0]
+
+    selected_rows = df_filt.iloc[ph_indices]
+
+    np.testing.assert_array_equal(ph_xb, selected_rows["Xb"].astype("float"))
+    np.testing.assert_array_equal(ph_yb, selected_rows["Yb"].astype("float"))
+    np.testing.assert_allclose(
+        ph_tof, selected_rows["TOF"].astype("float"), atol=1e-5, rtol=0
+    )
+
+
+def test_get_ssd_back_position_and_tof_offset(
+    de_dataset,
+    events_fsw_comparison_theta_0,
+):
+    """Tests get_ssd_back_position function."""
+    yb, tof_offset, ssd_number = get_ssd_back_position_and_tof_offset(de_dataset)
+
+    df = pd.read_csv(events_fsw_comparison_theta_0)
+    df_filt = df[(df["StartType"] != -1) & (df["StopType"] >= 8)]
+
+    np.testing.assert_array_equal(yb, df_filt["Yb"].astype("float"))
+
+    tof_offset_lt = tof_offset[df_filt["StartType"] == StartType.Left.value]
+    tof_offset_rt = tof_offset[df_filt["StartType"] == StartType.Right.value]
+
+    ssd_number_lt = ssd_number[df_filt["StartType"] == StartType.Left.value]
+    ssd_number_rt = ssd_number[df_filt["StartType"] == StartType.Right.value]
+
+    np.testing.assert_array_equal(
+        tof_offset_lt[ssd_number_lt == 3],
+        np.full(len(tof_offset_lt[ssd_number_lt == 3]), -4.2),
+    )
+    np.testing.assert_array_equal(
+        tof_offset_rt[ssd_number_rt == 7],
+        np.full(len(tof_offset_rt[ssd_number_rt == 7]), -6),
+    )
+    np.testing.assert_array_equal(
+        tof_offset_rt[ssd_number_rt == 4],
+        np.full(len(tof_offset_rt[ssd_number_rt == 4]), -4),
+    )
+
+    assert np.all(ssd_number_lt >= 0), "Values in ssd_number_lt out of range."
+
+    assert np.all(ssd_number_lt <= 7), "Values in ssd_number_lt out of range."
+
+    assert np.all(ssd_number_rt >= 0), "Values in ssd_number_rt out of range."
+
+    assert np.all(ssd_number_rt <= 7), "Values in ssd_number_rt out of range."
+
+
+def test_get_coincidence_positions(de_dataset, yf_fixture):
+    """Tests get_coincidence_positions function."""
+    df_filt, _, _ = yf_fixture
+    # Get particle tof (t2).
+    _, t2, _, _ = get_ph_tof_and_back_positions(
+        de_dataset, df_filt.Xf.astype("float").values, "ultra45"
+    )
+
+    # Filter for stop type.
+    indices = np.nonzero(
+        np.isin(de_dataset["STOP_TYPE"], [StopType.Top.value, StopType.Bottom.value])
+    )[0]
+    de_filtered = de_dataset.isel(epoch=indices)
+    rows = df_filt.iloc[indices]
+
+    # Get coincidence position and eTOF.
+    etof, xc = get_coincidence_positions(de_filtered, t2, "ultra45")
+
+    np.testing.assert_allclose(xc, rows["Xc"].astype("float"), atol=1e-4, rtol=0)
+    np.testing.assert_allclose(
+        etof, rows["eTOF"].astype("float").values, rtol=0, atol=1e-06
+    )
+
+
+def test_calculate_etof_xc(de_dataset, yf_fixture):
+    """Tests calculate_etof_xc function."""
+    df_filt, _, _ = yf_fixture
+    # Get particle tof (t2).
+    _, t2, _, _ = get_ph_tof_and_back_positions(
+        de_dataset, df_filt.Xf.astype("float").values, "ultra45"
+    )
+    # Filter based on STOP_TYPE.
+    indices = np.nonzero(
+        np.isin(de_dataset["STOP_TYPE"], [StopType.Top.value, StopType.Bottom.value])
+    )[0]
+    de_filtered = de_dataset.isel(epoch=indices)
+    df_filtered = df_filt.iloc[indices]
+
+    # Filter for COIN_TYPE Top and Bottom.
+    index_top = np.nonzero(np.isin(de_filtered["COIN_TYPE"], CoinType.Top.value))[0]
+    de_top = de_filtered.isel(epoch=index_top)
+    df_top = df_filtered.iloc[index_top]
+
+    index_bottom = np.nonzero(np.isin(de_filtered["COIN_TYPE"], CoinType.Bottom.value))[
+        0
+    ]
+    de_bottom = de_filtered.isel(epoch=index_bottom)
+    df_bottom = df_filtered.iloc[index_bottom]
+
+    # Calculate for Top and Bottom
+    etof_top, xc_top = calculate_etof_xc(de_top, t2[index_top], "ultra45", "TP")
+    etof_bottom, xc_bottom = calculate_etof_xc(
+        de_bottom, t2[index_bottom], "ultra45", "BT"
+    )
+
+    # Assertions for Top
+    np.testing.assert_allclose(
+        xc_top * 100, df_top["Xc"].astype("float"), atol=1e-4, rtol=0
+    )
+    np.testing.assert_allclose(
+        etof_top, df_top["eTOF"].astype("float").values, atol=1e-06, rtol=0
+    )
+
+    # Assertions for Bottom
+    np.testing.assert_allclose(
+        xc_bottom * 100, df_bottom["Xc"].astype("float"), atol=1e-4, rtol=0
+    )
+    np.testing.assert_allclose(
+        etof_bottom, df_bottom["eTOF"].astype("float").values, atol=1e-06, rtol=0
+    )
+
+
+def test_get_particle_velocity(de_dataset, yf_fixture):
+    """Tests get_particle_velocity function."""
+    df_filt, _, _ = yf_fixture
+
+    ph_indices = np.nonzero(
+        np.isin(de_dataset["STOP_TYPE"], [StopType.Top.value, StopType.Bottom.value])
+    )[0]
+
+    ph_rows = df_filt.iloc[ph_indices]
+    test_xf = ph_rows["Xf"].astype("float").values
+    test_yf = ph_rows["Yf"].astype("float").values
+    test_xb = ph_rows["Xb"].astype("float").values
+    test_yb = ph_rows["Yb"].astype("float").values
+    test_d = ph_rows["d"].astype("float").values
+    test_tof = ph_rows["TOF"].astype("float").values
+
+    vhat_x, vhat_y, vhat_z = get_particle_velocity(
+        (test_xf, test_yf),
+        (test_xb, test_yb),
+        test_d,
+        test_tof,
+    )
+    # FSW test data should be negative and not have an analysis
+    # for negative tof values.
+    assert vhat_x[test_tof > 0] == pytest.approx(
+        -df_filt["vhatX"].iloc[ph_indices].astype("float").values[test_tof > 0],
+        rel=1e-2,
+    )
+    assert vhat_y[test_tof > 0] == pytest.approx(
+        -df_filt["vhatY"].iloc[ph_indices].astype("float").values[test_tof > 0],
+        rel=1e-2,
+    )
+    assert vhat_z[test_tof > 0] == pytest.approx(
+        -df_filt["vhatZ"].iloc[ph_indices].astype("float").values[test_tof > 0],
+        rel=1e-2,
+    )
+
+
+def test_get_ssd_tof(de_dataset, yf_fixture):
+    """Tests get_ssd_tof function."""
+    df_filt, _, _ = yf_fixture
+    df_ssd = df_filt[df_filt["StopType"].isin(StopType.SSD.value)]
+    test_xf = df_filt["Xf"].astype("float").values
+
+    ssd_tof = get_ssd_tof(de_dataset, test_xf)
+
+    np.testing.assert_allclose(
+        ssd_tof, df_ssd["TOF"].astype("float"), atol=1e-05, rtol=0
+    )
+
+
+def test_get_energy_ssd(de_dataset, yf_fixture):
+    """Tests get_energy_ssd function."""
+    df_filt, _, _ = yf_fixture
+    df_ssd = df_filt[df_filt["StopType"].isin(StopType.SSD.value)]
+    _, _, ssd_number = get_ssd_back_position_and_tof_offset(de_dataset)
+    energy = get_energy_ssd(de_dataset, ssd_number)
+    test_energy = df_ssd["Energy"].astype("float")
+
+    assert np.array_equal(test_energy, energy)
+
+
+def test_get_energy_pulse_height(de_dataset, yf_fixture):
+    """Tests get_energy_ssd function."""
+    df_filt, _, _ = yf_fixture
+    df_ph = df_filt[df_filt["StopType"].isin(StopType.PH.value)]
+    ph_indices = np.nonzero(
+        np.isin(de_dataset["STOP_TYPE"], [StopType.Top.value, StopType.Bottom.value])
+    )[0]
+
+    test_xb = df_filt["Xb"].astype("float").values
+    test_yb = df_filt["Yb"].astype("float").values
+
+    energy = get_energy_pulse_height(
+        de_dataset["STOP_TYPE"].data, de_dataset["ENERGY_PH"].data, test_xb, test_yb
+    )
+    test_energy = df_ph["Energy"].astype("float")
+
+    assert np.array_equal(test_energy, energy[ph_indices])
+
+
+def test_get_ctof(yf_fixture):
+    """Tests get_ctof function."""
+    df_filt, _, _ = yf_fixture
+
+    df_ph_ssd = df_filt[
+        df_filt["StopType"].isin([StopType.SSD.value, StopType.PH.value])
+    ]
+
+    ctof = get_ctof(
+        df_ph_ssd["TOF"].astype("float").to_numpy(),
+        df_ph_ssd["r"].astype("float").to_numpy(),
+    )
+
+    np.testing.assert_allclose(
+        ctof, df_ph_ssd["cTOF"].astype("float"), atol=1e-05, rtol=0
+    )
+
+
+def test_determine_species_ph(yf_fixture):
+    """Tests determine_species_ph function."""
+    df_filt, _, _ = yf_fixture
+    df_ph = df_filt[df_filt["StopType"].isin(StopType.PH.value)]
+
+    bin = determine_species_pulse_height(
+        df_ph["Energy"].astype("float").to_numpy(),
+        df_ph["TOF"].astype("float").to_numpy(),
+        df_ph["r"].astype("float").to_numpy(),
+    )
+
+    # TODO: add in bin values.
+    np.testing.assert_allclose(bin, np.zeros(len(bin)), atol=1e-05, rtol=0)
+
+
+def test_determine_species_ssd(yf_fixture):
+    """Tests determine_species_ssd function."""
+    df_filt, _, _ = yf_fixture
+    df_ssd = df_filt[df_filt["StopType"].isin(StopType.SSD.value)]
+
+    bin = determine_species_ssd(
+        df_ssd["Energy"].astype("float").to_numpy(),
+        df_ssd["TOF"].astype("float").to_numpy(),
+        df_ssd["r"].astype("float").to_numpy(),
+    )
+
+    # TODO: add in bin values.
+    np.testing.assert_allclose(bin, np.zeros(len(bin)), atol=1e-05, rtol=0)

imap_processing/tests/ultra/unit/test_ultra_l1c.py

@@ -0,0 +1,104 @@
+import numpy as np
+import pytest
+import xarray as xr
+
+from imap_processing.ultra.l1c.ultra_l1c import ultra_l1c
+from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
+
+
+@pytest.fixture()
+def mock_data_l1b_dict():
+    # Create sample data for the xarray Dataset
+    epoch = np.arange(
+        "2024-02-07T15:28:37", "2024-02-07T15:28:42", dtype="datetime64[s]"
+    ).astype("datetime64[ns]")
+
+    data_vars_histogram = {
+        "sid": ("epoch", np.zeros(5)),
+        "row": ("epoch", np.zeros(5)),
+        "column": ("epoch", np.zeros(5)),
+        "shcoarse": ("epoch", np.zeros(5)),
+        "spin": ("epoch", np.zeros(5)),
+        "packetdata": ("epoch", np.zeros(5)),
+    }
+
+    coords = {"epoch": epoch}
+
+    attrs_histogram = {
+        "Logical_source": "imap_ultra_l1a_45sensor-histogram",
+        "Logical_source_description": "IMAP Mission ULTRA Instrument "
+        "Level-1A Single-Sensor Data",
+    }
+
+    dataset_histogram = xr.Dataset(
+        data_vars=data_vars_histogram, coords=coords, attrs=attrs_histogram
+    )
+
+    data_vars_cullingmask = {
+        "spin_number": ("epoch", np.zeros(5)),
+    }
+
+    attrs_cullingmask = {
+        "Logical_source": "imap_ultra_l1b_45sensor-cullingmask",
+        "Logical_source_description": "IMAP Mission ULTRA Instrument "
+        "Level-1B Culling Mask Data",
+    }
+
+    dataset_cullingmask = xr.Dataset(
+        data_vars_cullingmask, coords={"epoch": epoch}, attrs=attrs_cullingmask
+    )
+
+    data_dict = {
+        "imap_ultra_l1b_45sensor-cullingmask": dataset_cullingmask,
+        "imap_ultra_l1a_45sensor-histogram": dataset_histogram,
+    }
+    return data_dict
+
+
+@pytest.fixture()
+def mock_data_l1c_dict():
+    epoch = np.array(
+        [760591786368000000, 760591787368000000, 760591788368000000],
+        dtype="datetime64[ns]",
+    )
+    data_dict = {"epoch": epoch, "sid": np.zeros(3)}
+    return data_dict
+
+
+def test_create_dataset(mock_data_l1c_dict):
+    """Tests that dataset is created as expected."""
+    dataset = create_dataset(
+        mock_data_l1c_dict, "imap_ultra_l1c_45sensor-histogram", "l1c"
+    )
+
+    assert "epoch" in dataset.coords
+    assert dataset.coords["epoch"].dtype == "datetime64[ns]"
+    assert dataset.attrs["Logical_source"] == "imap_ultra_l1c_45sensor-histogram"
+    assert dataset["sid"].attrs["UNITS"] == " "
+    np.testing.assert_array_equal(dataset["sid"], np.zeros(3))
+
+
+def test_ultra_l1c(mock_data_l1b_dict):
+    """Tests that L1c data is created."""
+    output_datasets = ultra_l1c(mock_data_l1b_dict, data_version="001")
+
+    assert len(output_datasets) == 1
+    assert (
+        output_datasets[0].attrs["Logical_source"]
+        == "imap_ultra_l1c_45sensor-histogram"
+    )
+    assert (
+        output_datasets[0].attrs["Logical_source_description"]
+        == "IMAP-Ultra Instrument Level-1C Pointing Set Grid Histogram Data."
+    )
+
+
+def test_ultra_l1c_error(mock_data_l1b_dict):
+    """Tests that L1b data throws an error."""
+    mock_data_l1b_dict["bad_key"] = mock_data_l1b_dict.pop(
+        "imap_ultra_l1a_45sensor-histogram"
+    )
+    with pytest.raises(
+        ValueError, match="Data dictionary does not contain the expected keys."
+    ):
+        ultra_l1c(mock_data_l1b_dict, data_version="001")

imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py

@@ -0,0 +1,35 @@
+"Tests bins for pointing sets"
+
+import numpy as np
+
+from imap_processing.ultra.l1c.ultra_l1c_pset_bins import (
+    build_energy_bins,
+    build_spatial_bins,
+)
+
+
+def test_build_energy_bins():
+    """Tests build_energy_bins function."""
+    energy_bin_start, energy_bin_end = build_energy_bins()
+
+    assert energy_bin_start[0] == 3.5
+    assert len(energy_bin_start) == 90
+    assert len(energy_bin_end) == 90
+
+    # Comparison to expected values.
+    np.testing.assert_allclose(energy_bin_end[0], 3.6795, atol=1e-4)
+    np.testing.assert_allclose(energy_bin_start[-1], 299.9724, atol=1e-4)
+    np.testing.assert_allclose(energy_bin_end[-1], 315.3556, atol=1e-4)
+
+
+def test_build_spatial_bins():
+    """Tests build_spatial_bins function."""
+    az_bin_edges, el_bin_edges = build_spatial_bins()
+
+    assert az_bin_edges[0] == 0
+    assert az_bin_edges[-1] == 360
+    assert len(az_bin_edges) == 721
+
+    assert el_bin_edges[0] == -90
+    assert el_bin_edges[-1] == 90
+    assert len(el_bin_edges) == 361

imap_processing/ultra/__init__.py

@@ -0,0 +1 @@
+__version__ = "001"

imap_processing/ultra/constants.py

@@ -0,0 +1,60 @@
+"""Module for constants and useful shared classes used in Ultra."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class UltraConstants:
+    """
+    Constants for Ultra which can be used across different levels or classes.
+
+    Attributes
+    ----------
+    D_SLIT_FOIL: float
+        Shortest distance from slit to foil [mm]
+    SLIT_Z: float
+        Position of slit on Z axis [mm]
+    YF_ESTIMATE_LEFT: float
+        Front position of particle for left shutter [mm]
+    YF_ESTIMATE_RIGHT: float
+        Front position of particle for right shutter [mm]
+    N_ELEMENTS: int
+        Number of elements in lookup table
+    TRIG_CONSTANT: float
+        Trigonometric constant [mm]
+    COMPOSITE_ENERGY_THRESHOLD: int
+        DN threshold for composite energy
+    Z_DSTOP: float
+        Position of stop foil on Z axis [mm]
+    Z_DS: float
+        Position of slit on Z axis [mm]
+    DF: float
+        Distance from slit to foil [mm]
+    DMIN: float
+        Minimum distance between front and back detectors [mm]
+    DMIN_SSD_CTOF: float
+        SSD-specific correction to DMIN for time-of-flight normalization
+    """
+
+    D_SLIT_FOIL: float = 3.39
+    SLIT_Z: float = 44.89
+    YF_ESTIMATE_LEFT: float = 40.0
+    YF_ESTIMATE_RIGHT: float = -40.0
+    N_ELEMENTS: int = 256
+    TRIG_CONSTANT: float = 81.92
+
+    # Composite energy threshold for SSD events
+    COMPOSITE_ENERGY_THRESHOLD: int = 1707
+
+    # Geometry-related constants
+    Z_DSTOP: float = 2.6 / 2  # Position of stop foil on Z axis [mm]
+    Z_DS: float = 46.19 - (2.6 / 2)  # Position of slit on Z axis [mm]
+    DF: float = 3.39  # Distance from slit to foil [mm]
+
+    # Derived constants
+    DMIN: float = (
+        Z_DS - (2**0.5) * DF
+    )  # Minimum distance between front and back detectors [mm]
+    DMIN_SSD_CTOF: float = (DMIN**2) / (
+        DMIN - Z_DSTOP
+    )  # SSD-specific correction to DMIN [mm]