imap-processing 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

imap_processing/tests/mag/test_mag_l1b.py

@@ -5,7 +5,11 @@ import pytest
 import xarray as xr
 
 from imap_processing.cdf.utils import load_cdf, write_cdf
-from imap_processing.mag.l1b.mag_l1b import mag_l1b, mag_l1b_processing
+from imap_processing.mag.l1b.mag_l1b import (
+    calibrate_vector,
+    mag_l1b,
+    mag_l1b_processing,
+)
 
 
 @pytest.fixture(scope="module")
@@ -44,6 +48,7 @@ def mag_l1a_dataset():
     output_dataset["compression_flags"] = compression_flags
     output_dataset["direction_label"] = direction_label
     output_dataset["compression_label"] = compression_label
+    output_dataset.attrs["Logical_source"] = ["imap_mag_l1a_norm-mago"]
 
     return output_dataset
 
@@ -52,9 +57,8 @@ def test_mag_processing(mag_l1a_dataset):
     mag_l1a_dataset.attrs["Logical_source"] = ["imap_mag_l1a_norm-mago"]
 
     mag_l1b = mag_l1b_processing(mag_l1a_dataset)
-
     np.testing.assert_allclose(
-        mag_l1b["vectors"][0].values, [2.29819857, 2.22914442, 2.24950008, 0]
+        mag_l1b["vectors"][0].values, [2.2972, 2.2415, 2.2381, 0], atol=1e-4
     )
     np.testing.assert_allclose(mag_l1b["vectors"][1].values, [0, 0, 0, 0])
 
@@ -65,7 +69,7 @@ def test_mag_processing(mag_l1a_dataset):
     mag_l1b = mag_l1b_processing(mag_l1a_dataset)
 
     np.testing.assert_allclose(
-        mag_l1b["vectors"][0].values, [2.27615106, 2.22638234, 2.24382211, 0]
+        mag_l1b["vectors"][0].values, [2.27538, 2.23416, 2.23682, 0], atol=1e-5
     )
     np.testing.assert_allclose(mag_l1b["vectors"][1].values, [0, 0, 0, 0])
 
@@ -118,7 +122,7 @@ def test_mag_compression_scale(mag_l1a_dataset):
     mag_l1a_dataset.attrs["Logical_source"] = ["imap_mag_l1a_norm-mago"]
     output = mag_l1b(mag_l1a_dataset, "v001")
 
-    calibrated_vectors = np.matmul(np.array([1, 1, 1]), test_calibration)
+    calibrated_vectors = np.matmul(test_calibration, np.array([1, 1, 1]))
     # 16 bit width is the standard
     assert np.allclose(output["vectors"].data[0][:3], calibrated_vectors)
     # uncompressed data is uncorrected
@@ -132,3 +136,33 @@ def test_mag_compression_scale(mag_l1a_dataset):
     # width of 14 should be multiplied by 4
     scaled_vectors = calibrated_vectors * 4
     assert np.allclose(output["vectors"].data[3][:3], scaled_vectors)
+
+
+def test_calibrate_vector():
+    # from MFOTOURFO
+    cal_array = np.array(
+        [
+            [
+                [2.29722020e00, 7.38200160e-02, 1.88479865e-02, 4.59777333e-03],
+                [0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
+                [0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
+            ],
+            [
+                [3.48624576e-03, 1.09224000e-04, 3.26118600e-05, 5.02830000e-06],
+                [2.23802879e00, 7.23781440e-02, 1.84842873e-02, 4.50744060e-03],
+                [0.00000000e00, 0.00000000e00, 0.00000000e00, 0.00000000e00],
+            ],
+            [
+                [-2.50787532e-03, -8.33760000e-05, -2.71240200e-05, 2.50509000e-06],
+                [-8.88437262e-03, -2.84256000e-04, -7.41600000e-05, -2.29399200e-05],
+                [2.24950008e00, 7.23836160e-02, 1.84847323e-02, 4.50945192e-03],
+            ],
+        ]
+    )
+
+    calibration_matrix = xr.DataArray(cal_array)
+
+    cal_vector = calibrate_vector(np.array([1.0, 1.0, 1.0, 0]), calibration_matrix)
+    expected_vector = np.array([2.2972, 2.2415, 2.2381, 0])
+
+    assert np.allclose(cal_vector, expected_vector, atol=1e-4)

imap_processing/tests/spice/test_geometry.py

@@ -204,6 +204,19 @@ def test_get_spacecraft_to_instrument_spin_phase_offset(instrument, expected_off
             SpiceFrame.IMAP_SPACECRAFT,
             SpiceFrame.IMAP_DPS,
         ),
+        # single et, multiple position vectors
+        (
+            ["2025-04-30T12:00:00.000"],
+            np.array(
+                [
+                    [1, 0, 0],
+                    [0, 1, 0],
+                    [0, 0, 1],
+                ]
+            ),
+            SpiceFrame.IMAP_SPACECRAFT,
+            SpiceFrame.IMAP_DPS,
+        ),
     ],
 )
 def test_frame_transform(et_strings, position, from_frame, to_frame, furnish_kernels):
@@ -223,11 +236,24 @@ def test_frame_transform(et_strings, position, from_frame, to_frame, furnish_ker
         et = np.array([spice.utc2et(et_str) for et_str in et_strings])
         et_arg = et[0] if len(et) == 1 else et
         result = frame_transform(et_arg, position, from_frame, to_frame)
-        # check the result shape before modifying for value checking
-        assert result.shape == (3,) if len(et) == 1 else (len(et), 3)
-        # compare against pure SPICE calculation
-        position = np.broadcast_to(position, (len(et), 3))
-        result = np.broadcast_to(result, (len(et), 3))
+        # check the result shape before modifying for value checking.
+        # There are 3 cases to consider:
+
+        # 1 event time, multiple position vectors:
+        if len(et) == 1 and position.ndim > 1:
+            assert result.shape == position.shape
+        # multiple event times, single position vector:
+        elif len(et) > 1 and position.ndim == 1:
+            assert result.shape == (len(et), 3)
+        # multiple event times, multiple position vectors (same number of each)
+        elif len(et) > 1 and position.ndim > 1:
+            assert result.shape == (len(et), 3)
+
+        # compare against pure SPICE calculation.
+        # If the result is a single position vector, broadcast it to first.
+        if position.ndim == 1:
+            position = np.broadcast_to(position, (len(et), 3))
+            result = np.broadcast_to(result, (len(et), 3))
         for spice_et, spice_position, test_result in zip(et, position, result):
             rotation_matrix = spice.pxform(from_frame.name, to_frame.name, spice_et)
             spice_result = spice.mxv(rotation_matrix, spice_position)
@@ -254,7 +280,7 @@ def test_frame_transform_exceptions():
         match="Mismatch in number of position vectors and Ephemeris times provided.",
     ):
         frame_transform(
-            1,
+            [1, 2],
             np.arange(9).reshape((3, 3)),
             SpiceFrame.ECLIPJ2000,
             SpiceFrame.IMAP_HIT,

imap_processing/tests/spice/test_time.py

@@ -2,22 +2,41 @@
 
 import numpy as np
 import pytest
-import spiceypy as spice
+import spiceypy
 
 from imap_processing.spice import IMAP_SC_ID
-from imap_processing.spice.time import _sct2e_wrapper, j2000ns_to_j2000s, met_to_j2000ns
+from imap_processing.spice.time import (
+    TICK_DURATION,
+    _sct2e_wrapper,
+    et_to_utc,
+    j2000ns_to_j2000s,
+    met_to_datetime64,
+    met_to_j2000ns,
+    met_to_sclkticks,
+    met_to_utc,
+    str_to_et,
+)
+
+
+@pytest.mark.parametrize("met", [1, np.arange(10)])
+def test_met_to_sclkticks(met):
+    """Test coverage for met_to_sclkticks."""
+    # Tick duration is 20us as specified in imap_sclk_0000.tsc
+    expected = met * 1 / 20e-6
+    ticks = met_to_sclkticks(met)
+    np.testing.assert_array_equal(ticks, expected)
 
 
 def test_met_to_j2000ns(furnish_time_kernels):
     """Test coverage for met_to_j2000ns function."""
     utc = "2026-01-01T00:00:00.125"
-    et = spice.str2et(utc)
-    sclk_str = spice.sce2s(IMAP_SC_ID, et)
+    et = spiceypy.str2et(utc)
+    sclk_str = spiceypy.sce2s(IMAP_SC_ID, et)
     seconds, ticks = sclk_str.split("/")[1].split(":")
     # There is some floating point error calculating tick duration from 1 clock
     # tick so average over many clock ticks for better accuracy
     spice_tick_duration = (
-        spice.sct2e(IMAP_SC_ID, 1e12) - spice.sct2e(IMAP_SC_ID, 0)
+        spiceypy.sct2e(IMAP_SC_ID, 1e12) - spiceypy.sct2e(IMAP_SC_ID, 0)
     ) / 1e12
     met = float(seconds) + float(ticks) * spice_tick_duration
     j2000ns = met_to_j2000ns(met)
@@ -30,7 +49,7 @@ def test_j2000ns_to_j2000s(furnish_time_kernels):
     # Use spice to come up with reasonable J2000 values
     utc = "2025-09-23T00:00:00.000"
     # Test single value input
-    et = spice.str2et(utc)
+    et = spiceypy.str2et(utc)
     epoch = int(et * 1e9)
     j2000s = j2000ns_to_j2000s(epoch)
     assert j2000s == et
@@ -42,6 +61,61 @@ def test_j2000ns_to_j2000s(furnish_time_kernels):
     )
 
 
+@pytest.mark.parametrize(
+    "expected_utc, precision",
+    [
+        ("2024-01-01T00:00:00.000", 3),
+        (
+            [
+                "2024-01-01T00:00:00.000555",
+                "2025-09-23T00:00:00.000111",
+                "2040-11-14T10:23:48.156980",
+            ],
+            6,
+        ),
+    ],
+)
+def test_met_to_utc(furnish_time_kernels, expected_utc, precision):
+    """Test coverage for met_to_utc function."""
+    if isinstance(expected_utc, list):
+        et_arr = spiceypy.str2et(expected_utc)
+        sclk_ticks = np.array([spiceypy.sce2c(IMAP_SC_ID, et) for et in et_arr])
+    else:
+        et = spiceypy.str2et(expected_utc)
+        sclk_ticks = spiceypy.sce2c(IMAP_SC_ID, et)
+    met = sclk_ticks * TICK_DURATION
+    utc = met_to_utc(met, precision=precision)
+    np.testing.assert_array_equal(utc, expected_utc)
+
+
+@pytest.mark.parametrize(
+    "utc",
+    [
+        "2024-01-01T00:00:00.000",
+        [
+            "2024-01-01T00:00:00.000",
+            "2025-09-23T00:00:00.000",
+            "2040-11-14T10:23:48.15698",
+        ],
+    ],
+)
+def test_met_to_datetime64(furnish_time_kernels, utc):
+    """Test coverage for met_to_datetime64 function."""
+    if isinstance(utc, list):
+        expected_dt64 = np.array([np.datetime64(utc_str) for utc_str in utc])
+        et_arr = spiceypy.str2et(utc)
+        sclk_ticks = np.array([spiceypy.sce2c(IMAP_SC_ID, et) for et in et_arr])
+    else:
+        expected_dt64 = np.asarray(np.datetime64(utc))
+        et = spiceypy.str2et(utc)
+        sclk_ticks = spiceypy.sce2c(IMAP_SC_ID, et)
+    met = sclk_ticks * TICK_DURATION
+    dt64 = met_to_datetime64(met)
+    np.testing.assert_array_equal(
+        dt64.astype("datetime64[us]"), expected_dt64.astype("datetime64[us]")
+    )
+
+
 @pytest.mark.parametrize("sclk_ticks", [0.0, np.arange(10)])
 def test_sct2e_wrapper(sclk_ticks):
     """Test for `_sct2e_wrapper` function."""
@@ -50,3 +124,58 @@ def test_sct2e_wrapper(sclk_ticks):
         assert isinstance(et, float)
     else:
         assert len(et) == len(sclk_ticks)
+
+
+def test_str_to_et(furnish_time_kernels):
+    """Test coverage for string to et conversion function."""
+    utc = "2017-07-14T19:46:00"
+    # Test single value input
+    expected_et = 553333629.1837274
+    actual_et = str_to_et(utc)
+    assert expected_et == actual_et
+
+    # Test list input
+    list_of_utc = [
+        "2017-08-14T19:46:00.000",
+        "2017-09-14T19:46:00.000",
+        "2017-10-14T19:46:00.000",
+    ]
+
+    expected_et_array = np.array(
+        (556012029.1829445, 558690429.1824446, 561282429.1823651)
+    )
+    actual_et_array = str_to_et(list_of_utc)
+    assert np.array_equal(expected_et_array, actual_et_array)
+
+    # Test array input
+    array_of_utc = np.array(
+        [
+            "2017-08-14T19:46:00.000",
+            "2017-09-14T19:46:00.000",
+            "2017-10-14T19:46:00.000",
+        ]
+    )
+
+    actual_et_array = str_to_et(array_of_utc)
+    assert np.array_equal(expected_et_array, actual_et_array)
+
+
+def test_et_to_utc(furnish_time_kernels):
+    """Test coverage for et to utc conversion function."""
+    et = 553333629.1837274
+    # Test single value input
+    expected_utc = "2017-07-14T19:46:00.000"
+    actual_utc = et_to_utc(et)
+    assert expected_utc == actual_utc
+
+    # Test array input
+    array_of_et = np.array((556012029.1829445, 558690429.1824446, 561282429.1823651))
+    expected_utc_array = np.array(
+        (
+            "2017-08-14T19:46:00.000",
+            "2017-09-14T19:46:00.000",
+            "2017-10-14T19:46:00.000",
+        )
+    )
+    actual_utc_array = et_to_utc(array_of_et)
+    assert np.array_equal(expected_utc_array, actual_utc_array)

imap_processing/tests/swapi/test_swapi_decom.py

@@ -2,36 +2,34 @@ import pandas as pd
 import pytest
 
 from imap_processing import imap_module_directory
-from imap_processing.decom import decom_packets
 from imap_processing.swapi.l1.swapi_l1 import (
     SWAPIAPID,
 )
-from imap_processing.utils import group_by_apid
+from imap_processing.utils import packet_file_to_datasets
 
 
 @pytest.fixture(scope="session")
 def decom_test_data(swapi_l0_test_data_path):
-    """Read test data from file"""
+    """Read test data from file with derived values"""
     test_file = "imap_swapi_l0_raw_20240924_v001.pkts"
     packet_file = imap_module_directory / swapi_l0_test_data_path / test_file
     packet_definition = (
         f"{imap_module_directory}/swapi/packet_definitions/swapi_packet_definition.xml"
     )
-    data_list = []
-    data_list.extend(decom_packets(packet_file, packet_definition))
-    return data_list
+    return packet_file_to_datasets(
+        packet_file, packet_definition, use_derived_value=False
+    )
 
 
 def test_number_of_packets(decom_test_data):
     """This test and validate number of packets."""
-    grouped_data = group_by_apid(decom_test_data)
-    sci_packets = grouped_data[SWAPIAPID.SWP_SCI]
+    sci_packets = decom_test_data[SWAPIAPID.SWP_SCI]
     expected_sci_packets = 153
-    assert len(sci_packets) == expected_sci_packets
+    assert len(sci_packets["epoch"]) == expected_sci_packets
 
-    hk_packets = grouped_data[SWAPIAPID.SWP_HK]
+    hk_packets = decom_test_data[SWAPIAPID.SWP_HK]
     expected_hk_packets = 17
-    assert len(hk_packets) == expected_hk_packets
+    assert len(hk_packets["epoch"]) == expected_hk_packets
 
 
 def test_swapi_sci_data(decom_test_data, swapi_l0_validation_data_path):
@@ -42,32 +40,30 @@ def test_swapi_sci_data(decom_test_data, swapi_l0_validation_data_path):
         index_col="SHCOARSE",
     )
 
-    grouped_data = group_by_apid(decom_test_data)
-    sci_packets = grouped_data[SWAPIAPID.SWP_SCI]
-    first_data = sci_packets[0]
-    validation_data = raw_validation_data.loc[first_data["SHCOARSE"]]
+    sci_packets = decom_test_data[SWAPIAPID.SWP_SCI]
+    first_data = sci_packets.isel(epoch=0)
+    validation_data = raw_validation_data.loc[first_data["shcoarse"].values]
 
     # compare raw values of validation data
-    for key, value in first_data.items():
-        # check if the data is the same
-        if key == "PLAN_ID_SCIENCE":
-            # We had to work around this because HK and SCI packet uses
-            # PLAN_ID but they uses different length of bits.
-            assert value == validation_data["PLAN_ID"]
-        elif key == "SPARE_2_SCIENCE":
-            # Same for this SPARE_2 as above case
-            assert value == validation_data["SPARE_2"]
-        elif key == "MODE":
-            assert value.raw_value == validation_data[key]
-        elif "RNG" in key:
-            assert value.raw_value == validation_data[key]
-        else:
-            # for SHCOARSE we need the name of the column.
-            # This is done because pandas removed it from the
-            # main columns to make it the index.
-            assert value.raw_value == (
-                validation_data[key] if key != "SHCOARSE" else validation_data.name
-            )
+    for key in raw_validation_data.columns:
+        if key in [
+            "PHAPID",
+            "timestamp",
+            "PHGROUPF",
+            "PHSHF",
+            "PHVERNO",
+            "PHSEQCNT",
+            "PHDLEN",
+            "PHTYPE",
+        ]:
+            continue
+
+        # for SHCOARSE we need the name of the column.
+        # This is done because pandas removed it from the
+        # main columns to make it the index.
+        assert first_data[key.lower()].values == (
+            validation_data[key] if key != "SHCOARSE" else validation_data.name
+        )
 
 
 def test_swapi_hk_data(decom_test_data, swapi_l0_validation_data_path):
@@ -78,39 +74,49 @@ def test_swapi_hk_data(decom_test_data, swapi_l0_validation_data_path):
         index_col="SHCOARSE",
     )
 
-    grouped_data = group_by_apid(decom_test_data)
-    hk_packets = grouped_data[SWAPIAPID.SWP_HK]
-    first_data = hk_packets[0]
-    validation_data = raw_validation_data.loc[first_data["SHCOARSE"]]
-    bad_keys = [
-        "N5_V",
-        "SCEM_I",
-        "P5_I",
-        "PHD_LLD1_V",
-        "SPARE_4",
-        "P_CEM_CMD_LVL_MON",
-        "S_CEM_CMD_LVL_MON",
-        "ESA_CMD_LVL_MON",
-        "PHD_LLD2_V",
-        "CHKSUM",
-    ]
+    hk_packets = decom_test_data[SWAPIAPID.SWP_HK]
+    first_data = hk_packets.isel(epoch=0)
+    validation_data = raw_validation_data.loc[first_data["shcoarse"].values]
+
     # compare raw values of validation data
-    for key, value in first_data.items():
-        if key == "PLAN_ID_HK":
-            # We had to work around this because HK and SCI packet uses
-            # PLAN_ID but they uses different length of bits.
-            assert value == validation_data["PLAN_ID"]
-        elif key == "SPARE_2_HK":
-            # Same for this SPARE_2 as PLAN_ID
-            assert value == validation_data["SPARE_2"]
-        elif key == "SHCOARSE":
-            # for SHCOARSE we need the name of the column.
-            # This is done because pandas removed it from the main columns
-            # to make it the index.
-            assert value == validation_data.name
-        elif key in bad_keys:
-            # TODO: remove this elif after getting good validation data
-            # Validation data has wrong value for N5_V
+    for key in raw_validation_data.columns:
+        if key in [
+            "PHAPID",
+            "timestamp",
+            "PHGROUPF",
+            "PHSHF",
+            "PHVERNO",
+            "PHSEQCNT",
+            "PHDLEN",
+            "PHTYPE",
+        ]:
+            continue
+
+        value_mismatching_keys = [
+            "SCEM_I",
+            "N5_V",
+            "P5_I",
+            "PHD_LLD1_V",
+            "P_CEM_CMD_LVL_MON",
+            "S_CEM_CMD_LVL_MON",
+            "ESA_CMD_LVL_MON",
+            "PHD_LLD2_V",
+            "CHKSUM",
+        ]
+
+        extra_keys_val_data = [
+            "ESA_GATE_SET",
+            "P5V_ESA_V_MON",
+            "M5V_ESA_V_MON",
+            "P5V_ESA_I_MON",
+            "M5V_ESA_I_MON",
+        ]
+
+        if key in extra_keys_val_data or key in value_mismatching_keys:
             continue
-        else:
-            assert value.raw_value == validation_data[key]
+        # for SHCOARSE we need the name of the column.
+        # This is done because pandas removed it from the
+        # main columns to make it the index.
+        assert first_data[key.lower()].values == (
+            validation_data[key] if key != "SHCOARSE" else validation_data.name
+        )

imap_processing/tests/swapi/test_swapi_l1.py

@@ -39,7 +39,7 @@ def test_filter_good_data():
     total_sweeps = 3
     ds = xr.Dataset(
         {
-            "plan_id_science": xr.DataArray(np.full((total_sweeps * 12), 1)),
+            "plan_id": xr.DataArray(np.full((total_sweeps * 12), 1)),
             "sweep_table": xr.DataArray(np.repeat(np.arange(total_sweeps), 12)),
             "mode": xr.DataArray(np.full((total_sweeps * 12), SWAPIMODE.HVSCI.value)),
         },
@@ -69,9 +69,9 @@ def test_filter_good_data():
     expected = [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
     np.testing.assert_array_equal(filter_good_data(ds), expected)
 
-    # Check for bad plan_id_science data.
+    # Check for bad plan_id data.
     ds["sweep_table"] = xr.DataArray(np.repeat(np.arange(total_sweeps), 12))
-    ds["plan_id_science"][24 : total_sweeps * 12] = np.arange(0, 12)
+    ds["plan_id"][24 : total_sweeps * 12] = np.arange(0, 12)
     np.testing.assert_array_equal(filter_good_data(ds), np.arange(0, 24))
 
 
@@ -157,7 +157,7 @@ def test_process_swapi_science(decom_test_data):
     )
 
     # make PLAN_ID data incorrect. Now processed data should have less sweeps
-    ds_data["plan_id_science"].data[:24] = np.arange(24)
+    ds_data["plan_id"].data[:24] = np.arange(24)
     processed_data = process_swapi_science(
         ds_data, decom_test_data[SWAPIAPID.SWP_HK], data_version="001"
     )

imap_processing/tests/swe/test_swe_l2.py

@@ -2,13 +2,19 @@ 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.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 (
-    ELECTRON_MASS,
     ENERGY_CONVERSION_FACTOR,
+    VELOCITY_CONVERSION_FACTOR,
+    calculate_flux,
     calculate_phase_space_density,
     get_particle_energy,
+    swe_l2,
 )
 from imap_processing.swe.utils.swe_utils import read_lookup_table
 
@@ -17,7 +23,7 @@ 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
-    assert np.all(all_energy["energy"] == expected_energy)
+    np.testing.assert_array_equal(all_energy["energy"], expected_energy)
 
 
 @patch("imap_processing.swe.l2.swe_l2.GEOMETRIC_FACTORS", new=np.full(7, 1))
@@ -50,20 +56,70 @@ def test_calculate_phase_space_density(patch_get_particle_energy):
             ),
         }
     )
-    density = calculate_phase_space_density(l1b_dataset)
-    assert density.shape == (total_sweeps, 24, 30, 7)
+    phase_space_density_ds = calculate_phase_space_density(l1b_dataset)
+    assert phase_space_density_ds["phase_space_density"].shape == (
+        total_sweeps,
+        24,
+        30,
+        7,
+    )
 
     # Test that first sweep has correct values. In patch,
     #   1. we have set GEOMETRIC_FACTORS to 1.
     #   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 * (np.sqrt(2 * 1 / ELECTRON_MASS)) ** 4)
+    expected_calculated_density = (2 * 1) / (1 * VELOCITY_CONVERSION_FACTOR * 1**2)
     expected_density = np.full((24, 30, 7), expected_calculated_density)
-    assert np.all(density[0].data == expected_density)
+    np.testing.assert_array_equal(
+        phase_space_density_ds["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 * (np.sqrt(2 * 2 / ELECTRON_MASS)) ** 4)
+    expected_calculated_density = (2 * 1) / (1 * VELOCITY_CONVERSION_FACTOR * 2**2)
     expected_density = np.full((24, 30, 7), expected_calculated_density)
-    assert np.all(density[1].data == expected_density)
+    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, 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),
+            ),
+        }
+    )
+
+    flux = calculate_flux(l1b_dataset)
+    assert flux.shape == (total_sweeps, 24, 30, 7)
+    assert type(flux) == np.ndarray
+
+
+@pytest.mark.usefixtures("use_fake_spin_data_for_time")
+def test_swe_l2(use_fake_spin_data_for_time):
+    """Test L2 processing."""
+    data_start_time = 453051293.099714
+    data_end_time = 453066734
+    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")
+
+    l1b_dataset = swe_l1b(l1a_datasets, "002")
+    l2_dataset = swe_l2(l1b_dataset, "002")
+
+    assert type(l2_dataset) == xr.Dataset
+    assert l2_dataset["spin_phase"].shape == (6, 24, 30, 7)

imap_processing/tests/test_utils.py

@@ -91,7 +91,7 @@ def test_packet_file_to_datasets(use_derived_value, expected_mode):
         packet_files, packet_definition, use_derived_value=use_derived_value
     )
     # 3 apids in the test data
-    assert len(datasets_by_apid) == 4
+    assert len(datasets_by_apid) == 2
     data = datasets_by_apid[1188]
     assert data["sec_hdr_flg"].dtype == np.uint8
     assert data["pkt_apid"].dtype == np.uint16