imap-processing 0.18.0__py3-none-any.whl → 0.19.2__py3-none-any.whl

imap_processing/swapi/l2/swapi_l2.py

@@ -12,7 +12,7 @@ from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 logger = logging.getLogger(__name__)
 
 
-TIME_PER_BIN = 0.167  # seconds
+SWAPI_LIVETIME = 0.145  # seconds
 
 
 def solve_full_sweep_energy(
@@ -66,7 +66,7 @@ def solve_full_sweep_energy(
 
     first_63_energies = []
 
-    for time, sweep_id in zip(data_time, sweep_table):
+    for time, sweep_id in zip(data_time, sweep_table, strict=False):
         # Find the sweep's ESA data for the given time and sweep_id
         subset = esa_table_df[
             (esa_table_df["timestamp"] <= time) & (esa_table_df["Sweep #"] == sweep_id)
@@ -159,14 +159,12 @@ def swapi_l2(
 
     To process science data to L2, we need to:
     - convert counts to rates. This is done by dividing the counts by the
-        TIME_PER_BIN time. TIME_PER_BIN is the exposure time per energy bin which is
-        obtained by dividing the time for one complete sweep
-        (12 s, coarse + fine sweep) by the total energy steps (72),
-        i.e., TIME_PER_BIN = 12/72 = 0.167 s. This will be constant.
+        SWAPI_LIVETIME time. LIVETIME is data acquisition time. It will
+        be constant, SWAPI_LIVETIME = 0.145 s.
 
     - update uncertainty. Calculate new uncertainty value using
-        SWP_PCEM_ERR data from level one and divide by TIME_PER_BIN. Eg.
-            SWP_PCEM_UNC = SWP_PCEM_ERR / TIME_PER_BIN
+        SWP_PCEM_ERR data from level one and divide by SWAPI_LIVETIME. Eg.
+            SWP_PCEM_UNC = SWP_PCEM_ERR / SWAPI_LIVETIME
         Do the same for SCEM and COIN data.
 
     Parameters
@@ -233,9 +231,9 @@ def swapi_l2(
     ]
 
     # convert counts to rate
-    l2_dataset["swp_pcem_rate"] = l1_dataset["swp_pcem_counts"] / TIME_PER_BIN
-    l2_dataset["swp_scem_rate"] = l1_dataset["swp_scem_counts"] / TIME_PER_BIN
-    l2_dataset["swp_coin_rate"] = l1_dataset["swp_coin_counts"] / TIME_PER_BIN
+    l2_dataset["swp_pcem_rate"] = l1_dataset["swp_pcem_counts"] / SWAPI_LIVETIME
+    l2_dataset["swp_scem_rate"] = l1_dataset["swp_scem_counts"] / SWAPI_LIVETIME
+    l2_dataset["swp_coin_rate"] = l1_dataset["swp_coin_counts"] / SWAPI_LIVETIME
     # update attrs
     l2_dataset["swp_pcem_rate"].attrs = cdf_manager.get_variable_attributes("pcem_rate")
     l2_dataset["swp_scem_rate"].attrs = cdf_manager.get_variable_attributes("scem_rate")
@@ -243,22 +241,22 @@ def swapi_l2(
 
     # update uncertainty
     l2_dataset["swp_pcem_rate_stat_uncert_plus"] = (
-        l1_dataset["swp_pcem_counts_stat_uncert_plus"] / TIME_PER_BIN
+        l1_dataset["swp_pcem_counts_stat_uncert_plus"] / SWAPI_LIVETIME
     )
     l2_dataset["swp_pcem_rate_stat_uncert_minus"] = (
-        l1_dataset["swp_pcem_counts_stat_uncert_minus"] / TIME_PER_BIN
+        l1_dataset["swp_pcem_counts_stat_uncert_minus"] / SWAPI_LIVETIME
     )
     l2_dataset["swp_scem_rate_stat_uncert_plus"] = (
-        l1_dataset["swp_scem_counts_stat_uncert_plus"] / TIME_PER_BIN
+        l1_dataset["swp_scem_counts_stat_uncert_plus"] / SWAPI_LIVETIME
     )
     l2_dataset["swp_scem_rate_stat_uncert_minus"] = (
-        l1_dataset["swp_scem_counts_stat_uncert_minus"] / TIME_PER_BIN
+        l1_dataset["swp_scem_counts_stat_uncert_minus"] / SWAPI_LIVETIME
     )
     l2_dataset["swp_coin_rate_stat_uncert_plus"] = (
-        l1_dataset["swp_coin_counts_stat_uncert_plus"] / TIME_PER_BIN
+        l1_dataset["swp_coin_counts_stat_uncert_plus"] / SWAPI_LIVETIME
     )
     l2_dataset["swp_coin_rate_stat_uncert_minus"] = (
-        l1_dataset["swp_coin_counts_stat_uncert_minus"] / TIME_PER_BIN
+        l1_dataset["swp_coin_counts_stat_uncert_minus"] / SWAPI_LIVETIME
     )
     # update attrs
     l2_dataset[

imap_processing/swe/l1b/swe_l1b.py

@@ -2,7 +2,6 @@
 
 import logging
 from pathlib import Path
-from typing import Union
 
 import numpy as np
 import numpy.typing as npt
@@ -51,7 +50,7 @@ def get_esa_dataframe(esa_table_number: int) -> pd.DataFrame:
 
 
 def deadtime_correction(
-    counts: np.ndarray, acq_duration: Union[int, npt.NDArray]
+    counts: np.ndarray, acq_duration: int | npt.NDArray
 ) -> npt.NDArray:
     """
     Calculate deadtime correction.

imap_processing/ultra/constants.py

@@ -51,7 +51,6 @@ class UltraConstants:
     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_PH_CTOF: float = (
         Z_DS - (2**0.5) * DF
@@ -79,53 +78,60 @@ class UltraConstants:
     CULLING_RPM_MIN = 2.0
     CULLING_RPM_MAX = 6.0
 
-    # Thresholds for culling based on counts.
+    # Thresholds for culling based on counts (keV).
     CULLING_ENERGY_BIN_EDGES: ClassVar[list] = [
+        3.0,
+        10.0,
+        20.0,
+        50.0,
+        300.0,
+        1e5,
+    ]
+    PSET_ENERGY_BIN_EDGES: ClassVar[list] = [
         3.385,
         4.13722222222222,
-        4.13722222222222,
         5.05660493827161,
-        5.05660493827161,
-        6.18029492455419,
         6.18029492455419,
         7.55369379667734,
-        7.55369379667734,
         9.23229241816119,
-        9.23229241816119,
-        11.2839129555303,
         11.2839129555303,
         13.7914491678704,
-        13.7914491678704,
-        16.8562156496194,
         16.8562156496194,
         20.6020413495348,
-        20.6020413495348,
-        25.1802727605426,
         25.1802727605426,
         30.775888929552,
-        30.775888929552,
         37.6149753583414,
-        37.6149753583414,
-        45.9738587713061,
         45.9738587713061,
         56.1902718315964,
-        56.1902718315964,
-        68.6769989052845,
         68.6769989052845,
         83.93855421757,
-        83.93855421757,
-        102.591566265919,
         102.591566265919,
         125.38969210279,
-        125.38969210279,
         153.254068125632,
-        153.254068125632,
-        187.310527709106,
         187.310527709106,
         228.93508942224,
-        228.93508942224,
-        279.809553738294,
         279.809553738294,
         341.989454569026,
         1e5,
     ]
+
+    # Valid event filter constants
+    # Note these appear similar to image params constants
+    # but they should be used only for the valid event filter.
+    ETOFOFF1_EVENTFILTER = 100
+    ETOFOFF2_EVENTFILTER = -50
+    ETOFSLOPE1_EVENTFILTER = 6667
+    ETOFSLOPE2_EVENTFILTER = 7500
+    ETOFMAX_EVENTFILTER = 90
+    ETOFMIN_EVENTFILTER = -400
+    TOFDIFFTPMIN_EVENTFILTER = 226
+    TOFDIFFTPMAX_EVENTFILTER = 266
+
+    TOFXE_SPECIES_GROUPS: ClassVar[dict[str, list[int]]] = {
+        "proton": [3],
+        "non_proton": [20, 28, 36],
+    }
+    TOFXPH_SPECIES_GROUPS: ClassVar[dict[str, list[int]]] = {
+        "proton": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
+        "non_proton": [20, 21, 22, 23, 24, 25, 26],
+    }

imap_processing/ultra/l0/ultra_utils.py

@@ -1,6 +1,6 @@
 """Contains data classes to support Ultra L0 processing."""
 
-from typing import NamedTuple, Union
+from typing import NamedTuple
 
 
 class PacketProperties(NamedTuple):
@@ -9,23 +9,21 @@ class PacketProperties(NamedTuple):
     apid: list  # List of APIDs
     logical_source: list  # List of logical sources
     addition_to_logical_desc: str  # Description of the logical source
-    width: Union[int, None]  # Width of binary data (could be None).
+    width: int | None  # Width of binary data (could be None).
     # Block, image_planes, pixel_window_rows, and pixel_window_columns are important for
     # decompressing the images and a description is available on page 171 of IMAP-Ultra
     # Flight Software Specification document (7523-9009_Rev_-.pdf).
-    block: Union[int, None]  # Number of values in each block (could be None).
-    len_array: Union[
-        int, None
-    ]  # Length of the array to be decompressed (could be None).
-    mantissa_bit_length: Union[int, None]  # used to determine the level of
+    block: int | None  # Number of values in each block (could be None).
+    len_array: int | None  # Length of the array to be decompressed (could be None).
+    mantissa_bit_length: int | None  # used to determine the level of
     # precision that can be recovered from compressed data (could be None).
-    image_planes: Union[int, None] = None
+    image_planes: int | None = None
     # number of images. See table 11 in the FSSD.
-    pixel_window_rows: Union[int, None] = None
+    pixel_window_rows: int | None = None
     # number of rows in each image. See table 49 in the FSSD.
-    pixel_window_columns: Union[int, None] = None
+    pixel_window_columns: int | None = None
     # number of columns in each image. See table 49 in the FSSD.
-    image_planes_per_packet: Union[int, None] = None
+    image_planes_per_packet: int | None = None
     # number of image planes in each packet. See table 52 in the FSSD.
 
 

imap_processing/ultra/l1a/ultra_l1a.py

@@ -1,7 +1,6 @@
 """Generate ULTRA L1a CDFs."""
 
 import logging
-from typing import Optional
 
 import xarray as xr
 
@@ -44,7 +43,7 @@ logger = logging.getLogger(__name__)
 
 
 def ultra_l1a(  # noqa: PLR0912
-    packet_file: str, apid_input: Optional[int] = None
+    packet_file: str, apid_input: int | None = None
 ) -> list[xr.Dataset]:
     """
     Will process ULTRA L0 data into L1A CDF files at output_filepath.

imap_processing/ultra/l1b/badtimes.py

@@ -7,13 +7,14 @@ from numpy.typing import NDArray
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset, extract_data_dict
 
 FILLVAL_UINT16 = 65535
+FILLVAL_FLOAT32 = -1.0e31
 FILLVAL_FLOAT64 = -1.0e31
 FILLVAL_UINT32 = 4294967295
 
 
 def calculate_badtimes(
     extendedspin_dataset: xr.Dataset,
-    cullingmask_spins: NDArray,
+    goodtimes_spins: NDArray,
     name: str,
 ) -> xr.Dataset:
     """
@@ -23,7 +24,7 @@ def calculate_badtimes(
     ----------
     extendedspin_dataset : xarray.Dataset
         Dataset containing the data.
-    cullingmask_spins : np.ndarray
+    goodtimes_spins : np.ndarray
         Dataset containing the culled data.
     name : str
         Name of the dataset.
@@ -33,11 +34,9 @@ def calculate_badtimes(
     badtimes_dataset : xarray.Dataset
         Dataset containing the extendedspin data that has been culled.
     """
+    n_bins = extendedspin_dataset.dims["energy_bin_geometric_mean"]
     culled_spins = np.setdiff1d(
-        extendedspin_dataset["spin_number"].values, cullingmask_spins
-    )
-    extendedspin_dataset = extendedspin_dataset.assign_coords(
-        epoch=("spin_number", extendedspin_dataset["epoch"].values)
+        extendedspin_dataset["spin_number"].values, goodtimes_spins
     )
     filtered_dataset = extendedspin_dataset.sel(spin_number=culled_spins)
 
@@ -48,9 +47,6 @@ def calculate_badtimes(
     if badtimes_dataset["spin_number"].size == 0:
         badtimes_dataset = badtimes_dataset.drop_dims("spin_number")
         badtimes_dataset = badtimes_dataset.expand_dims(spin_number=[FILLVAL_UINT32])
-        badtimes_dataset = badtimes_dataset.assign_coords(
-            epoch=("spin_number", [extendedspin_dataset["epoch"].values[0]])
-        )
         badtimes_dataset["spin_start_time"] = xr.DataArray(
             np.array([FILLVAL_FLOAT64], dtype="float64"), dims=["spin_number"]
         )
@@ -60,16 +56,44 @@ def calculate_badtimes(
         badtimes_dataset["spin_rate"] = xr.DataArray(
             np.array([FILLVAL_FLOAT64], dtype="float64"), dims=["spin_number"]
         )
+        badtimes_dataset["start_pulses_per_spin"] = xr.DataArray(
+            np.array([FILLVAL_FLOAT32], dtype="float32"),
+            dims=["spin_number"],
+        )
+        badtimes_dataset["stop_pulses_per_spin"] = xr.DataArray(
+            np.array([FILLVAL_FLOAT32], dtype="float32"),
+            dims=["spin_number"],
+        )
+        badtimes_dataset["coin_pulses_per_spin"] = xr.DataArray(
+            np.array([FILLVAL_FLOAT32], dtype="float32"),
+            dims=["spin_number"],
+        )
+        badtimes_dataset["rejected_events_per_spin"] = xr.DataArray(
+            np.array([FILLVAL_UINT32], dtype="uint32"),
+            dims=["spin_number"],
+        )
         badtimes_dataset["quality_attitude"] = xr.DataArray(
             np.array([FILLVAL_UINT16], dtype="uint16"), dims=["spin_number"]
         )
+        badtimes_dataset["quality_hk"] = xr.DataArray(
+            np.array([FILLVAL_UINT16], dtype="uint16"),
+            dims=["spin_number"],
+        )
+        badtimes_dataset["quality_instruments"] = xr.DataArray(
+            np.array([FILLVAL_UINT16], dtype="uint16"),
+            dims=["spin_number"],
+        )
         badtimes_dataset["quality_ena_rates"] = (
             ("energy_bin_geometric_mean", "spin_number"),
-            np.full((3, 1), FILLVAL_UINT16, dtype="uint16"),
+            np.full((n_bins, 1), FILLVAL_UINT16, dtype="uint16"),
         )
         badtimes_dataset["ena_rates"] = (
             ("energy_bin_geometric_mean", "spin_number"),
-            np.full((3, 1), FILLVAL_FLOAT64, dtype="float64"),
+            np.full((n_bins, 1), FILLVAL_FLOAT64, dtype="float64"),
+        )
+        badtimes_dataset["ena_rates_threshold"] = (
+            ("energy_bin_geometric_mean", "spin_number"),
+            np.full((n_bins, 1), FILLVAL_FLOAT32, dtype="float32"),
         )
 
     return badtimes_dataset

imap_processing/ultra/l1b/de.py

@@ -4,12 +4,16 @@ import numpy as np
 import xarray as xr
 
 from imap_processing.cdf.utils import parse_filename_like
-from imap_processing.quality_flags import ImapDEUltraFlags
+from imap_processing.quality_flags import (
+    ImapDEOutliersUltraFlags,
+    ImapDEScatteringUltraFlags,
+)
 from imap_processing.spice.geometry import SpiceFrame
 from imap_processing.ultra.l1b.lookup_utils import get_geometric_factor
 from imap_processing.ultra.l1b.ultra_l1b_annotated import (
     get_annotated_particle_velocity,
 )
+from imap_processing.ultra.l1b.ultra_l1b_culling import flag_scattering
 from imap_processing.ultra.l1b.ultra_l1b_extended import (
     StopType,
     determine_ebin_pulse_height,
@@ -32,10 +36,13 @@ from imap_processing.ultra.l1b.ultra_l1b_extended import (
     get_spin_number,
     get_ssd_back_position_and_tof_offset,
     get_ssd_tof,
+    is_back_tof_valid,
+    is_coin_ph_valid,
 )
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
 FILLVAL_UINT8 = 255
+FILLVAL_UINT32 = 4294967295
 FILLVAL_FLOAT32 = -1.0e31
 
 
@@ -76,7 +83,6 @@ def calculate_de(
         "event_type",
         "de_event_met",
         "phase_angle",
-        "spin",
     ]
     dataset_keys = [
         "coin_type",
@@ -84,11 +90,13 @@ def calculate_de(
         "stop_type",
         "shcoarse",
         "phase_angle",
-        "spin",
     ]
 
     de_dict.update(
-        {key: de_dataset[dataset_key] for key, dataset_key in zip(keys, dataset_keys)}
+        {
+            key: de_dataset[dataset_key]
+            for key, dataset_key in zip(keys, dataset_keys, strict=False)
+        }
     )
 
     valid_mask = de_dataset["start_type"].data != FILLVAL_UINT8
@@ -114,9 +122,11 @@ def calculate_de(
     tof = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
     etof = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
     ctof = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
+    tof_energy = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
     magnitude_v = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
     energy = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
     e_bin = np.full(len(de_dataset["epoch"]), FILLVAL_UINT8, dtype=np.uint8)
+    e_bin_l1a = np.full(len(de_dataset["epoch"]), FILLVAL_UINT8, dtype=np.uint8)
     species_bin = np.full(len(de_dataset["epoch"]), FILLVAL_UINT8, dtype=np.uint8)
     t2 = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
     event_times = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float32)
@@ -125,10 +135,19 @@ def calculate_de(
     sc_dps_velocity = np.full(shape, FILLVAL_FLOAT32, dtype=np.float32)
     helio_velocity = np.full(shape, FILLVAL_FLOAT32, dtype=np.float32)
     spin_starts = np.full(len(de_dataset["epoch"]), FILLVAL_FLOAT32, dtype=np.float64)
+    velocities = np.full(shape, FILLVAL_FLOAT32, dtype=np.float32)
+    v_hat = np.full(shape, FILLVAL_FLOAT32, dtype=np.float32)
+    r_hat = np.full(shape, FILLVAL_FLOAT32, dtype=np.float32)
 
     start_type = np.full(len(de_dataset["epoch"]), FILLVAL_UINT8, dtype=np.uint8)
     quality_flags = np.full(
-        de_dataset["epoch"].shape, ImapDEUltraFlags.NONE.value, dtype=np.uint16
+        de_dataset["epoch"].shape, ImapDEOutliersUltraFlags.NONE.value, dtype=np.uint16
+    )
+
+    scattering_quality_flags = np.full(
+        de_dataset["epoch"].shape,
+        ImapDEScatteringUltraFlags.NONE.value,
+        dtype=np.uint16,
     )
 
     xf[valid_indices] = get_front_x_position(
@@ -149,9 +168,13 @@ def calculate_de(
     )
 
     # Pulse height
-    tof[ph_indices], t2[ph_indices], xb[ph_indices], yb[ph_indices] = (
-        get_ph_tof_and_back_positions(de_dataset, xf, f"ultra{sensor}", ancillary_files)
+    ph_result = get_ph_tof_and_back_positions(
+        de_dataset, xf, f"ultra{sensor}", ancillary_files
     )
+    tof[ph_indices] = ph_result.tof
+    t2[ph_indices] = ph_result.t2
+    xb[ph_indices] = ph_result.xb
+    yb[ph_indices] = ph_result.yb
     d[ph_indices], yf[ph_indices] = get_front_y_position(
         de_dataset["start_type"].data[ph_indices], yb[ph_indices], ancillary_files
     )
@@ -174,16 +197,39 @@ def calculate_de(
         (xb[ph_indices], yb[ph_indices]),
         d[ph_indices],
     )
-    e_bin[ph_indices] = determine_ebin_pulse_height(
-        energy[ph_indices], tof[ph_indices], r[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}",
         ancillary_files,
     )
+    backtofvalid = is_back_tof_valid(
+        de_dataset,
+        xf,
+        f"ultra{sensor}",
+        ancillary_files,
+    )
+    coinphvalid = is_coin_ph_valid(
+        etof[ph_indices],
+        xc[ph_indices],
+        xb[ph_indices],
+        de_dataset["stop_north_tdc"][ph_indices].values,
+        de_dataset["stop_south_tdc"][ph_indices].values,
+        de_dataset["stop_east_tdc"][ph_indices].values,
+        de_dataset["stop_west_tdc"][ph_indices].values,
+        f"ultra{sensor}",
+        ancillary_files,
+        quality_flags[ph_indices],
+    )
+    e_bin[ph_indices] = determine_ebin_pulse_height(
+        energy[ph_indices],
+        tof[ph_indices],
+        r[ph_indices],
+        backtofvalid,
+        coinphvalid,
+        ancillary_files,
+    )
+    species_bin[ph_indices] = determine_species(e_bin[ph_indices], "PH")
     ctof[ph_indices], magnitude_v[ph_indices] = get_ctof(
         tof[ph_indices], r[ph_indices], "PH"
     )
@@ -211,11 +257,13 @@ def calculate_de(
         d[ssd_indices],
     )
     e_bin[ssd_indices] = determine_ebin_ssd(
-        energy[ssd_indices], tof[ssd_indices], r[ssd_indices]
-    )
-    species_bin[ssd_indices] = determine_species(
-        tof[ssd_indices], r[ssd_indices], "SSD"
+        energy[ssd_indices],
+        tof[ssd_indices],
+        r[ssd_indices],
+        f"ultra{sensor}",
+        ancillary_files,
     )
+    species_bin[ssd_indices] = determine_species(e_bin[ssd_indices], "SSD")
     ctof[ssd_indices], magnitude_v[ssd_indices] = get_ctof(
         tof[ssd_indices], r[ssd_indices], "SSD"
     )
@@ -237,19 +285,26 @@ def calculate_de(
     de_dict["phi"] = phi
     de_dict["theta"] = theta
 
-    v, vhat, r = get_de_velocity(
-        (de_dict["x_front"], de_dict["y_front"]),
-        (de_dict["x_back"], de_dict["y_back"]),
-        de_dict["front_back_distance"],
-        de_dict["tof_start_stop"],
+    velocities[valid_indices], v_hat[valid_indices], r_hat[valid_indices] = (
+        get_de_velocity(
+            (de_dict["x_front"][valid_indices], de_dict["y_front"][valid_indices]),
+            (de_dict["x_back"][valid_indices], de_dict["y_back"][valid_indices]),
+            de_dict["front_back_distance"][valid_indices],
+            de_dict["tof_start_stop"][valid_indices],
+        )
     )
-    de_dict["direct_event_velocity"] = v.astype(np.float32)
-    de_dict["direct_event_unit_velocity"] = vhat.astype(np.float32)
-    de_dict["direct_event_unit_position"] = r.astype(np.float32)
+    de_dict["direct_event_unit_velocity"] = v_hat.astype(np.float32)
+    de_dict["direct_event_unit_position"] = r_hat.astype(np.float32)
 
-    de_dict["tof_energy"] = get_de_energy_kev(v, species_bin)
+    tof_energy[valid_indices] = get_de_energy_kev(
+        velocities[valid_indices], species_bin[valid_indices]
+    )
+    de_dict["tof_energy"] = tof_energy
     de_dict["energy"] = energy
-    de_dict["ebin"] = e_bin
+    de_dict["computed_ebin"] = e_bin
+    valid_ebin = de_dataset["bin"].values != FILLVAL_UINT32
+    e_bin_l1a[valid_ebin] = de_dataset["bin"].values[valid_ebin]
+    de_dict["ebin"] = e_bin_l1a
     de_dict["species"] = species_bin
 
     # Annotated Events.
@@ -264,7 +319,7 @@ def calculate_de(
             helio_velocity[valid_events],
         ) = get_annotated_particle_velocity(
             event_times[valid_events],
-            de_dict["direct_event_velocity"][valid_events],
+            velocities.astype(np.float32)[valid_events],
             ultra_frame,
             SpiceFrame.IMAP_DPS,
             SpiceFrame.IMAP_SPACECRAFT,
@@ -289,20 +344,29 @@ def calculate_de(
         de_dict["tof_energy"], de_dict["phi"], de_dict["theta"], ancillary_files
     )
     de_dict["geometric_factor_blades"] = get_geometric_factor(
-        ancillary_files,
-        "l1b-sensor-gf-blades",
         de_dict["phi"],
         de_dict["theta"],
         quality_flags,
+        ancillary_files,
+        "l1b-sensor-gf-blades",
     )
     de_dict["geometric_factor_noblades"] = get_geometric_factor(
-        ancillary_files,
-        "l1b-sensor-gf-noblades",
         de_dict["phi"],
         de_dict["theta"],
         quality_flags,
+        ancillary_files,
+        "l1b-sensor-gf-noblades",
+    )
+    de_dict["quality_outliers"] = quality_flags
+    flag_scattering(
+        de_dict["tof_energy"],
+        de_dict["theta"],
+        de_dict["phi"],
+        ancillary_files,
+        sensor,
+        scattering_quality_flags,
     )
-    de_dict["quality_fov"] = quality_flags
+    de_dict["quality_scattering"] = scattering_quality_flags
 
     dataset = create_dataset(de_dict, name, "l1b")
 

imap_processing/ultra/l1b/extendedspin.py

@@ -2,8 +2,10 @@
 
 import numpy as np
 import xarray as xr
+from numpy.typing import NDArray
 
 from imap_processing.ultra.l1b.ultra_l1b_culling import (
+    count_rejected_events_per_spin,
     flag_attitude,
     flag_hk,
     flag_imap_instruments,
@@ -14,6 +16,7 @@ from imap_processing.ultra.l1b.ultra_l1b_culling import (
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
 FILLVAL_UINT16 = 65535
+FILLVAL_FLOAT32 = -1.0e31
 
 
 def calculate_extendedspin(
@@ -43,7 +46,7 @@ def calculate_extendedspin(
     de_dataset = dict_datasets[f"imap_ultra_l1b_{instrument_id}sensor-de"]
 
     extendedspin_dict = {}
-    rates_qf, spin, energy_midpoints, n_sigma_per_energy = flag_rates(
+    rates_qf, spin, energy_bin_geometric_mean, n_sigma_per_energy = flag_rates(
         de_dataset["spin"].values,
         de_dataset["energy"].values,
     )
@@ -57,34 +60,46 @@ def calculate_extendedspin(
     hk_qf = flag_hk(de_dataset["spin"].values)
     inst_qf = flag_imap_instruments(de_dataset["spin"].values)
 
-    # Get the first epoch for each spin.
-    mask = xr.DataArray(np.isin(de_dataset["spin"], spin), dims="epoch")
-    filtered_dataset = de_dataset.where(mask, drop=True)
-    _, first_indices = np.unique(filtered_dataset["spin"].values, return_index=True)
-    first_epochs = filtered_dataset["epoch"].values[first_indices]
-
     # Get the number of pulses per spin.
-    start_per_spin, stop_per_spin, coin_per_spin = get_pulses_per_spin(rates_dataset)
+    pulses = get_pulses_per_spin(rates_dataset)
 
+    # Track rejected events in each spin based on
+    # quality flags in de l1b data.
+    rejected_counts = count_rejected_events_per_spin(
+        de_dataset["spin"].values,
+        de_dataset["quality_scattering"].values,
+        de_dataset["quality_outliers"].values,
+    )
     # These will be the coordinates.
-    extendedspin_dict["epoch"] = first_epochs
     extendedspin_dict["spin_number"] = spin
-    extendedspin_dict["energy_bin_geometric_mean"] = energy_midpoints
+    extendedspin_dict["energy_bin_geometric_mean"] = energy_bin_geometric_mean
 
     extendedspin_dict["ena_rates"] = count_rates
     extendedspin_dict["ena_rates_threshold"] = n_sigma_per_energy
     extendedspin_dict["spin_start_time"] = spin_starttime
     extendedspin_dict["spin_period"] = spin_period
     extendedspin_dict["spin_rate"] = spin_rates
+
+    # Get index of pulses.unique_spins corresponding to each spin.
+    idx: NDArray[np.intp] = np.searchsorted(pulses.unique_spins, spin)
+
+    # Validate that the spin values match
+    valid = (idx < pulses.unique_spins.size) & (pulses.unique_spins[idx] == spin)
+
+    start_per_spin = np.full(len(spin), FILLVAL_FLOAT32, dtype=np.float32)
+    stop_per_spin = np.full(len(spin), FILLVAL_FLOAT32, dtype=np.float32)
+    coin_per_spin = np.full(len(spin), FILLVAL_FLOAT32, dtype=np.float32)
+
+    # Fill only the valid ones
+    start_per_spin[valid] = pulses.start_per_spin[idx[valid]]
+    stop_per_spin[valid] = pulses.stop_per_spin[idx[valid]]
+    coin_per_spin[valid] = pulses.coin_per_spin[idx[valid]]
+
+    # account for rates spins which are not in the direct event spins
     extendedspin_dict["start_pulses_per_spin"] = start_per_spin
     extendedspin_dict["stop_pulses_per_spin"] = stop_per_spin
     extendedspin_dict["coin_pulses_per_spin"] = coin_per_spin
-    # TODO: this will be used to track rejected events in each
-    #  spin based on quality flags in de l1b data.
-    extendedspin_dict["rejected_events_per_spin"] = np.full_like(
-        spin, FILLVAL_UINT16, dtype=np.uint16
-    )
-
+    extendedspin_dict["rejected_events_per_spin"] = rejected_counts
     extendedspin_dict["quality_attitude"] = attitude_qf
     extendedspin_dict["quality_ena_rates"] = rates_qf
     extendedspin_dict["quality_hk"] = hk_qf