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

imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py

@@ -1,7 +1,9 @@
 "Tests pointing sets"
 
+import astropy_healpix.healpy as hp
 import cdflib
 import numpy as np
+import pandas as pd
 import pytest
 from cdflib import CDF
 
@@ -9,13 +11,15 @@ from imap_processing import imap_module_directory
 from imap_processing.ena_maps.utils.spatial_utils import build_spatial_bins
 from imap_processing.ultra.l1c.ultra_l1c_pset_bins import (
     build_energy_bins,
+    get_background_rates,
     get_helio_exposure_times,
-    get_histogram,
-    get_pointing_frame_exposure_times,
-    get_pointing_frame_sensitivity,
+    get_spacecraft_exposure_times,
+    get_spacecraft_histogram,
+    get_spacecraft_sensitivity,
 )
 
 BASE_PATH = imap_module_directory / "ultra" / "lookup_tables"
+TEST_PATH = imap_module_directory / "tests" / "ultra" / "data" / "l1"
 
 
 @pytest.fixture()
@@ -24,7 +28,7 @@ def test_data():
     vx_sc = np.array([-186.5575, 508.5697, 508.5697, 508.5697])
     vy_sc = np.array([-707.5707, -516.0282, -516.0282, -516.0282])
     vz_sc = np.array([618.0569, 892.6931, 892.6931, 892.6931])
-    energy = np.array([3.384, 3.385, 200, 200])
+    energy = np.array([3.384, 3.385, 4.138, 4.138])
     v = np.column_stack((vx_sc, vy_sc, vz_sc))
 
     return v, energy
@@ -32,13 +36,13 @@ def test_data():
 
 def test_build_energy_bins():
     """Tests build_energy_bins function."""
-    energy_bin_edges, energy_midpoints = build_energy_bins()
-    energy_bin_start = [interval[0] for interval in energy_bin_edges]
-    energy_bin_end = [interval[1] for interval in energy_bin_edges]
+    intervals, energy_midpoints, energy_bin_geometric_means = build_energy_bins()
+    energy_bin_start = [interval[0] for interval in intervals]
+    energy_bin_end = [interval[1] for interval in intervals]
 
     assert energy_bin_start[0] == 0
     assert energy_bin_start[1] == 3.385
-    assert len(energy_bin_edges) == 24
+    assert len(intervals) == 24
     assert energy_midpoints[0] == (energy_bin_start[0] + energy_bin_end[0]) / 2
 
     # Comparison to expected values.
@@ -46,42 +50,67 @@ def test_build_energy_bins():
     np.testing.assert_allclose(energy_bin_start[-1], 279.810, atol=1e-4)
     np.testing.assert_allclose(energy_bin_end[-1], 341.989, atol=1e-4)
 
+    expected_geometric_means = np.sqrt(
+        np.array(energy_bin_start) * np.array(energy_bin_end)
+    )
+    np.testing.assert_allclose(
+        energy_bin_geometric_means, expected_geometric_means, atol=1e-4
+    )
+
 
-def test_get_histogram(test_data):
+def test_get_spacecraft_histogram(test_data):
     """Tests get_histogram function."""
     v, energy = test_data
 
-    az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
-        np.rad2deg(angle_radians) for angle_radians in (build_spatial_bins())
-    )
-    energy_bin_edges, _ = build_energy_bins()
-
-    hist = get_histogram(v, energy, az_bin_edges, el_bin_edges, energy_bin_edges)
+    energy_bin_edges, _, _ = build_energy_bins()
+    subset_energy_bin_edges = energy_bin_edges[:3]
 
-    assert hist.shape == (
-        len(az_bin_edges) - 1,
-        len(el_bin_edges) - 1,
-        len(energy_bin_edges),
+    hist, latitude, longitude, n_pix = get_spacecraft_histogram(
+        v, energy, subset_energy_bin_edges, nside=1
     )
-
-
-def test_get_pointing_frame_exposure_times():
-    """Tests get_pointing_frame_exposure_times function."""
-
-    constant_exposure = BASE_PATH / "dps_grid45_compressed.cdf"
-    spins_per_pointing = 5760
-    exposure = get_pointing_frame_exposure_times(
-        constant_exposure, spins_per_pointing, "45"
+    assert hist.shape == (len(subset_energy_bin_edges), hp.nside2npix(1))
+    assert n_pix == hp.nside2npix(1)
+    assert latitude.shape == (n_pix,)
+    assert longitude.shape == (n_pix,)
+
+    # Spot check that 2 counts are in the third energy bin
+    assert np.sum(hist[2, :]) == 2
+
+    # Test overlapping energy bins
+    overlapping_bins = [
+        (0.0, 3.385),
+        (2.5, 4.137),
+        (3.385, 5.057),
+    ]
+    hist, latitude, longitude, n_pix = get_spacecraft_histogram(
+        v, energy, overlapping_bins, nside=1
     )
-
-    assert exposure.shape == (720, 360)
-    # Assert that the exposure time at the highest azimuth is
-    # 15s x spins per pointing.
-    assert np.array_equal(
-        exposure[:, 359], np.full_like(exposure[:, 359], spins_per_pointing * 15)
+    # Spot check that 3 counts are in the third energy bin
+    assert np.sum(hist[2, :]) == 3
+    assert n_pix == hp.nside2npix(1)
+    assert latitude.shape == (n_pix,)
+    assert longitude.shape == (n_pix,)
+
+
+def test_get_background_rates():
+    """Tests get_background_rates function."""
+    background_rates = get_background_rates(nside=128)
+    assert background_rates.shape == hp.nside2npix(128)
+
+
+@pytest.mark.external_test_data()
+def test_get_spacecraft_exposure_times():
+    """Test get_spacecraft_exposure_times function."""
+    constant_exposure = TEST_PATH / "ultra_90_dps_exposure.csv"
+    df_exposure = pd.read_csv(constant_exposure)
+    exposure_pointing = get_spacecraft_exposure_times(df_exposure)
+    assert exposure_pointing.shape == (196608,)
+
+    np.testing.assert_allclose(
+        exposure_pointing.values[22684:22686],
+        np.array([1.035, 1.035]) * 5760,
+        atol=1e-6,
     )
-    # Assert that the exposure time at the lowest azimuth is 0 (no exposure).
-    assert np.array_equal(exposure[:, 0], np.full_like(exposure[:, 359], 0.0))
 
 
 @pytest.mark.external_kernel()
@@ -99,9 +128,9 @@ def test_get_helio_exposure_times():
 
     exposure_3d = get_helio_exposure_times(mid_time, sc_exposure)
 
-    energy_bin_edges, energy_midpoints = build_energy_bins()
+    energy_bin_edges, energy_midpoints, _ = build_energy_bins()
     az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
-        np.rad2deg(angle_radians) for angle_radians in (build_spatial_bins())
+        build_spatial_bins()
     )
 
     assert exposure_3d.shape == (
@@ -116,7 +145,7 @@ def test_get_helio_exposure_times():
         ("dps_exposure_helio_45_E24.cdf", "dps_exposure_helio_45_E24"),
     ]
 
-    cdf_directory = imap_module_directory / "tests" / "ultra" / "test_data" / "l1"
+    cdf_directory = imap_module_directory / "tests" / "ultra" / "data" / "l1"
 
     exposures = []
 
@@ -132,21 +161,34 @@ def test_get_helio_exposure_times():
     assert np.array_equal(np.squeeze(exposures[2]), exposure_3d[:, :, 23])
 
 
-def test_get_pointing_frame_sensitivity():
-    """Tests get_pointing_frame_sensitivity function."""
+@pytest.mark.external_test_data()
+def test_get_spacecraft_sensitivity():
+    """Tests get_spacecraft_sensitivity function."""
+    # TODO: remove below here with lookup table aux api
+    efficiences = TEST_PATH / "Ultra_90_DPS_efficiencies_all.csv"
+    geometric_function = TEST_PATH / "ultra_90_dps_gf.csv"
+
+    df_efficiencies = pd.read_csv(efficiences)
+    df_geometric_function = pd.read_csv(geometric_function)
+
+    sensitivity = get_spacecraft_sensitivity(df_efficiencies, df_geometric_function)
+
+    assert sensitivity.shape == df_efficiencies.shape
 
-    # TODO: energy bins need to be modified from N=90 to N=24.
-    constant_sensitivity = BASE_PATH / "dps_sensitivity45.cdf"
-    spins_per_pointing = 5760
-    sensitivity = get_pointing_frame_sensitivity(
-        constant_sensitivity,
-        spins_per_pointing,
-        "45",
+    df_efficiencies_test = pd.DataFrame(
+        {"3.0keV": [1.0, 2.0], "3.5keV": [3.0, 4.0], "4.0keV": [5.0, 6.0]}
     )
 
-    assert sensitivity.shape == (90, 720, 360)
+    df_geometric_function_test = pd.DataFrame({"Response": [0.1, 0.2]})
 
-    with cdflib.CDF(constant_sensitivity) as cdf_file:
-        expected_sensitivity = cdf_file.varget("dps_sensitivity45") * spins_per_pointing
+    df_sensitivity_test = df_efficiencies_test.mul(
+        df_geometric_function_test["Response"], axis=0
+    )
+
+    expected_sensitivity = pd.DataFrame(
+        {"3.0keV": [0.1, 0.4], "3.5keV": [0.3, 0.8], "4.0keV": [0.5, 1.2]}
+    )
 
-    assert np.array_equal(sensitivity, expected_sensitivity)
+    assert np.allclose(
+        df_sensitivity_test.to_numpy(), expected_sensitivity.to_numpy(), atol=1e-6
+    )

imap_processing/ultra/l0/decom_tools.py

@@ -130,14 +130,15 @@ def decompress_binary(
     current_position = 0
     decompressed_values: list = []
 
-    while current_position < len(binary):
+    while current_position < len(binary) and len(decompressed_values) < array_length:
         # Read the width of the block
         width, current_position = read_and_advance(binary, width_bit, current_position)
-        # If width is 0 or None, we don't have enough bits left
-        if width is None or len(decompressed_values) >= array_length:
-            break
+        # If width is 0, add 'block' number of zeroes and continue
+        if width == 0:
+            decompressed_values.extend([0] * block)
+            continue
 
-        # For each block, read 16 values of the given width
+        # For each block, read 'block' values of the given width
         for _ in range(block):
             # Ensure there are enough bits left to read the width
             if len(binary) - current_position < width:

imap_processing/ultra/l1a/ultra_l1a.py

@@ -44,7 +44,7 @@ def initiate_data_arrays(decom_ultra: dict, apid: int) -> xr.Dataset:
         index = ULTRA_EVENTS.apid.index(apid)
         logical_source = ULTRA_EVENTS.logical_source[index]
         addition_to_logical_desc = ULTRA_EVENTS.addition_to_logical_desc
-        raw_time = decom_ultra["EVENTTIMES"]
+        raw_time = decom_ultra["SHCOARSE"]
     elif apid in ULTRA_TOF.apid:
         index = ULTRA_TOF.apid.index(apid)
         logical_source = ULTRA_TOF.logical_source[index]
@@ -120,59 +120,43 @@ def initiate_data_arrays(decom_ultra: dict, apid: int) -> xr.Dataset:
     return dataset
 
 
-def get_event_time(decom_ultra_dict: dict) -> dict:
+def get_event_id(decom_ultra_dict: dict) -> dict:
     """
-    Get event times using data from events and aux packets.
+    Get unique event IDs using data from events packets.
 
     Parameters
     ----------
     decom_ultra_dict : dict
-        Events and aux data.
+        Events data.
 
     Returns
     -------
     decom_events : dict
-        Ultra events data with calculated events timestamps.
-
-    Notes
-    -----
-    Equation for event time:
-    t = t_(spin start) + t_(spin start sub)/1000 +
-    t_(spin duration)/1000 * phase_angle/720
+        Ultra events data with calculated unique event IDs as 64-bit integers.
     """
-    event_times, durations, spin_starts = ([] for _ in range(3))
-    decom_aux = decom_ultra_dict[ULTRA_AUX.apid[0]]
     decom_events: dict = decom_ultra_dict[ULTRA_EVENTS.apid[0]]
 
-    timespinstart_array = np.array(decom_aux["TIMESPINSTART"])
-    timespinstartsub_array = np.array(decom_aux["TIMESPINSTARTSUB"]) / 1000
-
-    # spin start according to aux data
-    aux_spin_starts = timespinstart_array + timespinstartsub_array
-
-    for time in np.unique(decom_events["SHCOARSE"]):
-        # Get the nearest spin start and duration prior to the event
-        spin_start = aux_spin_starts[aux_spin_starts <= time][-1]
-        duration = np.array(decom_aux["DURATION"])[aux_spin_starts <= time][-1]
+    event_ids = []
+    packet_counters = {}
 
-        # Find the events
-        event_indices = np.where(np.array(decom_events["SHCOARSE"]) == time)
-
-        for event_index in event_indices[0]:
-            phase_angle = decom_events["PHASE_ANGLE"][event_index]
-
-            durations.append(duration)
-            spin_starts.append(spin_start)
+    for met in decom_events["SHCOARSE"]:
+        # Initialize the counter for a new packet (MET value)
+        if met not in packet_counters:
+            packet_counters[met] = 0
+        else:
+            packet_counters[met] += 1
 
-            # If there were no events, the time is set to 'SHCOARSE'
-            if decom_events["COUNT"][event_index] == 0:
-                event_times.append(decom_events["SHCOARSE"][event_index])
-            else:
-                event_times.append(spin_start + (duration / 1000) * (phase_angle / 720))
+        # Left shift SHCOARSE (u32) by 31 bits, to make room for our event counters
+        # (31 rather than 32 to keep it positive in the int64 representation)
+        # Append the current number of events in this packet to the right-most bits
+        # This makes each event a unique value including the MET and event number
+        # in the packet
+        # NOTE: CDF does not allow for uint64 values,
+        # so we use int64 representation here
+        event_id = (np.int64(met) << np.int64(31)) | np.int64(packet_counters[met])
+        event_ids.append(event_id)
 
-    decom_events["DURATION"] = durations
-    decom_events["TIMESPINSTART"] = spin_starts
-    decom_events["EVENTTIMES"] = event_times
+    decom_events["EVENTID"] = event_ids
 
     return decom_events
 
@@ -194,7 +178,7 @@ def create_dataset(decom_ultra_dict: dict) -> xr.Dataset:
     # Combine events and aux datasets so we can have proper event timestamps
     if ULTRA_EVENTS.apid[0] in decom_ultra_dict.keys():
         apid = ULTRA_EVENTS.apid[0]
-        decom_ultra = get_event_time(decom_ultra_dict)
+        decom_ultra = get_event_id(decom_ultra_dict)
     else:
         apid = next(iter(decom_ultra_dict.keys()))
         decom_ultra = decom_ultra_dict[apid]
@@ -286,31 +270,19 @@ def ultra_l1a(
     output_datasets = []
 
     # This is used for two purposes currently:
-    # 1. For testing purposes to only generate a dataset for a single apid.
+    #    For testing purposes to only generate a dataset for a single apid.
     #    Each test dataset is only for a single apid while the rest of the apids
     #    contain zeros. Ideally we would have
     #    test data for all apids and remove this parameter.
-    # 2. When we are generating the l1a dataset for the events packet since
-    #    right now we need to combine the events and aux packets to get the
-    #    correct event timestamps (get_event_time). This part will change
-    #    when we begin using the spin table in the database instead of the aux packet.
     if apid is not None:
         apids = [apid]
     else:
         apids = list(grouped_data.keys())
 
     for apid in apids:
-        if apid == ULTRA_EVENTS.apid[0]:
-            decom_ultra_dict = {
-                apid: process_ultra_apids(grouped_data[apid], apid),
-                ULTRA_AUX.apid[0]: process_ultra_apids(
-                    grouped_data[ULTRA_AUX.apid[0]], ULTRA_AUX.apid[0]
-                ),
-            }
-        else:
-            decom_ultra_dict = {
-                apid: process_ultra_apids(grouped_data[apid], apid),
-            }
+        decom_ultra_dict = {
+            apid: process_ultra_apids(grouped_data[apid], apid),
+        }
         dataset = create_dataset(decom_ultra_dict)
         # TODO: move this to use ImapCdfAttributes().add_global_attribute()
         dataset.attrs["Data_version"] = data_version

imap_processing/ultra/l1b/de.py

@@ -13,15 +13,16 @@ from imap_processing.ultra.l1b.ultra_l1b_extended import (
     determine_species,
     get_coincidence_positions,
     get_ctof,
-    get_de_az_el,
     get_de_energy_kev,
     get_de_velocity,
     get_energy_pulse_height,
     get_energy_ssd,
+    get_eventtimes,
     get_front_x_position,
     get_front_y_position,
     get_path_length,
     get_ph_tof_and_back_positions,
+    get_phi_theta,
     get_ssd_back_position_and_tof_offset,
     get_ssd_tof,
 )
@@ -49,18 +50,52 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
     de_dict = {}
     sensor = parse_filename_like(name)["sensor"][0:2]
 
-    # Drop events with invalid start type.
-    de_dataset = de_dataset.where(
-        de_dataset["START_TYPE"] != np.iinfo(np.int64).min, drop=True
+    # Define epoch and spin.
+    de_dict["epoch"] = de_dataset["epoch"].data
+    de_dict["spin"] = de_dataset["SPIN"].data
+
+    # Add already populated fields.
+    keys = [
+        "coincidence_type",
+        "start_type",
+        "event_type",
+        "de_event_met",
+        "phase_angle",
+        "spin",
+    ]
+    dataset_keys = [
+        "COIN_TYPE",
+        "START_TYPE",
+        "STOP_TYPE",
+        "SHCOARSE",
+        "PHASE_ANGLE",
+        "SPIN",
+    ]
+
+    de_dict.update(
+        {key: de_dataset[dataset_key] for key, dataset_key in zip(keys, dataset_keys)}
     )
 
+    valid_mask = de_dataset["START_TYPE"].data != np.iinfo(np.int64).min
+    ph_mask = np.isin(
+        de_dataset["STOP_TYPE"].data, [StopType.Top.value, StopType.Bottom.value]
+    )
+    ssd_mask = np.isin(de_dataset["STOP_TYPE"].data, [StopType.SSD.value])
+
+    valid_indices = np.nonzero(valid_mask)[0]
+    ph_indices = np.nonzero(valid_mask & ph_mask)[0]
+    ssd_indices = np.nonzero(valid_mask & ssd_mask)[0]
+
     # Instantiate arrays
+    xf = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     yf = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     xb = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     yb = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     xc = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     d = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float64)
     r = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
+    phi = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
+    theta = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     tof = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     etof = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     ctof = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
@@ -68,19 +103,25 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
     energy = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
     species_bin = np.full(len(de_dataset["epoch"]), "UNKNOWN", dtype="U10")
     t2 = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
+    event_times = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float64)
+    spin_starts = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float64)
+    spin_period_sec = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float64)
 
-    # Define epoch.
-    de_dict["epoch"] = de_dataset["epoch"].data
+    xf[valid_indices] = get_front_x_position(
+        de_dataset["START_TYPE"].data[valid_indices],
+        de_dataset["START_POS_TDC"].data[valid_indices],
+    )
 
-    xf = get_front_x_position(
-        de_dataset["START_TYPE"].data,
-        de_dataset["START_POS_TDC"].data,
+    (
+        event_times[valid_indices],
+        spin_starts[valid_indices],
+        spin_period_sec[valid_indices],
+    ) = get_eventtimes(
+        de_dataset["SPIN"].data[valid_indices],
+        de_dataset["PHASE_ANGLE"].data[valid_indices],
     )
 
     # Pulse height
-    ph_indices = np.nonzero(
-        np.isin(de_dataset["STOP_TYPE"], [StopType.Top.value, StopType.Bottom.value])
-    )[0]
     tof[ph_indices], t2[ph_indices], xb[ph_indices], yb[ph_indices] = (
         get_ph_tof_and_back_positions(de_dataset, xf, f"ultra{sensor}")
     )
@@ -98,6 +139,11 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
         (xb[ph_indices], yb[ph_indices]),
         d[ph_indices],
     )
+    phi[ph_indices], theta[ph_indices] = get_phi_theta(
+        (xf[ph_indices], yf[ph_indices]),
+        (xb[ph_indices], yb[ph_indices]),
+        d[ph_indices],
+    )
     species_bin[ph_indices] = determine_species(tof[ph_indices], r[ph_indices], "PH")
     etof[ph_indices], xc[ph_indices] = get_coincidence_positions(
         de_dataset.isel(epoch=ph_indices), t2[ph_indices], f"ultra{sensor}"
@@ -107,7 +153,6 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
     )
 
     # SSD
-    ssd_indices = np.nonzero(np.isin(de_dataset["STOP_TYPE"], StopType.SSD.value))[0]
     tof[ssd_indices] = get_ssd_tof(de_dataset, xf)
     yb[ssd_indices], _, ssd_number = get_ssd_back_position_and_tof_offset(de_dataset)
     xc[ssd_indices] = np.zeros(len(ssd_indices))
@@ -122,6 +167,11 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
         (xb[ssd_indices], yb[ssd_indices]),
         d[ssd_indices],
     )
+    phi[ssd_indices], theta[ssd_indices] = get_phi_theta(
+        (xf[ssd_indices], yf[ssd_indices]),
+        (xb[ssd_indices], yb[ssd_indices]),
+        d[ssd_indices],
+    )
     species_bin[ssd_indices] = determine_species(
         tof[ssd_indices], r[ssd_indices], "SSD"
     )
@@ -131,6 +181,9 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
 
     # Combine ph_yb and ssd_yb along with their indices
     de_dict["x_front"] = xf.astype(np.float32)
+    de_dict["event_times"] = event_times
+    de_dict["spin_starts"] = spin_starts
+    de_dict["spin_period"] = spin_period_sec
     de_dict["y_front"] = yf
     de_dict["x_back"] = xb
     de_dict["y_back"] = yb
@@ -141,19 +194,8 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
     de_dict["velocity_magnitude"] = magnitude_v
     de_dict["front_back_distance"] = d
     de_dict["path_length"] = r
-
-    keys = [
-        "coincidence_type",
-        "start_type",
-        "event_type",
-        "de_event_met",
-        "event_times",
-    ]
-    dataset_keys = ["COIN_TYPE", "START_TYPE", "STOP_TYPE", "SHCOARSE", "EVENTTIMES"]
-
-    de_dict.update(
-        {key: de_dataset[dataset_key] for key, dataset_key in zip(keys, dataset_keys)}
-    )
+    de_dict["phi"] = phi
+    de_dict["theta"] = theta
 
     v = get_de_velocity(
         (de_dict["x_front"], de_dict["y_front"]),
@@ -164,14 +206,13 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
     de_dict["direct_event_velocity"] = v.astype(np.float32)
 
     de_dict["tof_energy"] = get_de_energy_kev(v, species_bin)
-    de_dict["azimuth"], de_dict["elevation"] = get_de_az_el(v)
     de_dict["energy"] = energy
     de_dict["species"] = species_bin
 
     # Annotated Events.
     ultra_frame = getattr(SpiceFrame, f"IMAP_ULTRA_{sensor}")
     sc_velocity, sc_dps_velocity, helio_velocity = get_annotated_particle_velocity(
-        de_dataset.data_vars["EVENTTIMES"].values,
+        event_times,
         de_dict["direct_event_velocity"],
         ultra_frame,
         SpiceFrame.IMAP_DPS,
@@ -182,6 +223,9 @@ def calculate_de(de_dataset: xr.Dataset, name: str, data_version: str) -> xr.Dat
     de_dict["velocity_dps_sc"] = sc_dps_velocity
     de_dict["velocity_dps_helio"] = helio_velocity
 
+    de_dict["energy_spacecraft"] = get_de_energy_kev(sc_dps_velocity, species_bin)
+    de_dict["energy_heliosphere"] = get_de_energy_kev(helio_velocity, species_bin)
+
     # TODO: TBD.
     de_dict["event_efficiency"] = np.full(
         len(de_dataset["epoch"]), np.nan, dtype=np.float32

imap_processing/ultra/l1b/extendedspin.py

@@ -6,50 +6,48 @@ from imap_processing.ultra.l1b.ultra_l1b_culling import (
     flag_attitude,
     flag_spin,
     get_energy_histogram,
-    get_spin,
 )
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
 
 def calculate_extendedspin(
-    hk_dataset: xr.Dataset,
-    rates_dataset: xr.Dataset,
-    de_dataset: xr.Dataset,
+    dict_datasets: dict[str, xr.Dataset],
     name: str,
     data_version: str,
+    instrument_id: int,
 ) -> xr.Dataset:
     """
     Create dataset with defined datatypes for Extended Spin Data.
 
     Parameters
     ----------
-    hk_dataset : xarray.Dataset
-        Dataset containing l1a hk data.
-    rates_dataset : xarray.Dataset
-        Dataset containing l1a rates data.
-    de_dataset : xarray.Dataset
-        Dataset containing l1b de data.
+    dict_datasets : dict
+        Dictionary containing all the datasets.
     name : str
         Name of the dataset.
     data_version : str
         Version of the data.
+    instrument_id : int
+        Instrument ID.
 
     Returns
     -------
     extendedspin_dataset : xarray.Dataset
         Dataset containing the data.
     """
+    aux_dataset = dict_datasets[f"imap_ultra_l1a_{instrument_id}sensor-aux"]
+    de_dataset = dict_datasets[f"imap_ultra_l1b_{instrument_id}sensor-de"]
+
     extendedspin_dict = {}
     rates_qf, spin, energy_midpoints, n_sigma_per_energy = flag_spin(
-        de_dataset["event_times"].values,
+        de_dataset["spin"].values,
         de_dataset["energy"].values,
     )
-    spin_number = get_spin(de_dataset["event_times"].values)
     count_rates, _, counts, _ = get_energy_histogram(
-        spin_number, de_dataset["energy"].values
+        de_dataset["spin"].values, de_dataset["energy"].values
     )
     attitude_qf, spin_rates, spin_period, spin_starttime = flag_attitude(
-        de_dataset["event_times"].values
+        de_dataset["spin"].values, aux_dataset
     )
 
     # These will be the coordinates.