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

imap_processing/ultra/l1b/ultra_l1b_extended.py

@@ -2,6 +2,7 @@
 
 # TODO: Come back and add in FSW logic.
 import logging
+from collections import namedtuple
 from enum import Enum
 
 import numpy as np
@@ -11,11 +12,14 @@ from numpy import ndarray
 from numpy.typing import NDArray
 from scipy.interpolate import LinearNDInterpolator, RegularGridInterpolator
 
+from imap_processing.quality_flags import ImapDEOutliersUltraFlags
 from imap_processing.spice.spin import get_spin_data
+from imap_processing.spice.time import sct_to_et
 from imap_processing.ultra.constants import UltraConstants
 from imap_processing.ultra.l1b.lookup_utils import (
     get_angular_profiles,
     get_back_position,
+    get_ebins,
     get_energy_efficiencies,
     get_energy_norm,
     get_image_params,
@@ -26,6 +30,10 @@ from imap_processing.ultra.l1b.lookup_utils import (
 
 logger = logging.getLogger(__name__)
 
+FILLVAL_UINT8 = 255
+FILLVAL_FLOAT32 = -1.0e31
+FILLVAL_FLOAT64 = -1.0e31
+
 
 class StartType(Enum):
     """Start Type: 1=Left, 2=Right."""
@@ -50,6 +58,9 @@ class CoinType(Enum):
     Bottom = 2
 
 
+PHTOFResult = namedtuple("PHTOFResult", ["tof", "t2", "xb", "yb", "tofx", "tofy"])
+
+
 def get_front_x_position(
     start_type: ndarray, start_position_tdc: ndarray, sensor: str, ancillary_files: dict
 ) -> ndarray:
@@ -161,7 +172,7 @@ def get_front_y_position(
 
 def get_ph_tof_and_back_positions(
     de_dataset: xarray.Dataset, xf: np.ndarray, sensor: str, ancillary_files: dict
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+) -> PHTOFResult:
     """
     Calculate back xb, yb position and tof.
 
@@ -196,6 +207,10 @@ def get_ph_tof_and_back_positions(
         Back positions in x direction (hundredths of a millimeter).
     yb : np.array
         Back positions in y direction (hundredths of a millimeter).
+    tofx : np.array
+        X front position tof offset (tenths of a nanosecond).
+    tofy : np.array
+        Y front position tof offset (tenths of a nanosecond).
     """
     indices = np.nonzero(
         np.isin(de_dataset["stop_type"], [StopType.Top.value, StopType.Bottom.value])
@@ -278,7 +293,7 @@ def get_ph_tof_and_back_positions(
         stop_type_bottom
     ] / 10 * get_image_params("XFTTOF", sensor, ancillary_files)
 
-    return tof, t2, xb, yb
+    return PHTOFResult(tof=tof, t2=t2, xb=xb, yb=yb, tofx=tofx, tofy=tofy)
 
 
 def get_path_length(
@@ -530,16 +545,16 @@ def get_de_velocity(
     v_y = delta_v[:, 1] / tof * 1e3
     v_z = delta_v[:, 2] / tof * 1e3
 
-    v_x[tof < 0] = np.nan  # used as fillvals
-    v_y[tof < 0] = np.nan
-    v_z[tof < 0] = np.nan
+    v_x[tof < 0] = FILLVAL_FLOAT32  # used as fillvals
+    v_y[tof < 0] = FILLVAL_FLOAT32
+    v_z[tof < 0] = FILLVAL_FLOAT32
 
     velocities = np.vstack((v_x, v_y, v_z)).T
 
     v_hat = velocities / np.linalg.norm(velocities, axis=1)[:, None]
     r_hat = -v_hat
 
-    return velocities, v_hat, r_hat
+    return velocities, -v_hat, -r_hat
 
 
 def get_ssd_tof(
@@ -621,13 +636,17 @@ def get_de_energy_kev(v: np.ndarray, species: np.ndarray) -> NDArray:
     # Compute the sum of squares.
     v2 = np.sum(vv**2, axis=1)
 
-    index_hydrogen = np.where(species == 1)
-    energy = np.full_like(v2, np.nan)
+    # Only compute where species == 1 and v is valid
+    index_hydrogen = species == 1
+    valid_velocity = np.isfinite(v2)
+    valid_mask = index_hydrogen & valid_velocity
+
+    energy = np.full_like(v2, FILLVAL_FLOAT32)
 
     # TODO: we will calculate the energies of the different species here.
     # 1/2 mv^2 in Joules, convert to keV
-    energy[index_hydrogen] = (
-        0.5 * UltraConstants.MASS_H * v2[index_hydrogen] * UltraConstants.J_KEV
+    energy[valid_mask] = (
+        0.5 * UltraConstants.MASS_H * v2[valid_mask] * UltraConstants.J_KEV
     )
 
     return energy
@@ -820,25 +839,16 @@ def get_ctof(
     return ctof, magnitude_v
 
 
-def determine_species(tof: np.ndarray, path_length: np.ndarray, type: str) -> NDArray:
+def determine_species(e_bin: np.ndarray, type: str) -> NDArray:
     """
     Determine the species for pulse-height events.
 
-    Species is determined from the particle velocity.
-    For velocity, the particle TOF is normalized with respect
-    to a fixed distance dmin between the front and back detectors.
-    The normalized TOF is termed the corrected TOF (ctof).
-    Particle species are determined from ctof using thresholds.
-
-    Further description is available on pages 42-44 of
-    IMAP-Ultra Flight Software Specification document.
+    Species is determined using the computed e_bin.
 
     Parameters
     ----------
-    tof : np.ndarray
-        Time of flight of the SSD event (tenths of a nanosecond).
-    path_length : np.ndarray
-        Path length (r) (hundredths of a millimeter).
+    e_bin : np.ndarray
+        Computed e_bin.
     type : str
         Type of data (PH or SSD).
 
@@ -847,11 +857,17 @@ def determine_species(tof: np.ndarray, path_length: np.ndarray, type: str) -> ND
     species_bin : np.array
         Species bin.
     """
-    # Event TOF normalization to Z axis
-    ctof, _ = get_ctof(tof, path_length, type)
-    # Assign Species 1 ("H") to bins
-    # TODO: this is a placeholder for future species assignments.
-    species_bin = np.full(len(ctof), 1, dtype=np.uint8)
+    if type == "PH":
+        species_groups = UltraConstants.TOFXPH_SPECIES_GROUPS
+    if type == "SSD":
+        species_groups = UltraConstants.TOFXE_SPECIES_GROUPS
+
+    non_proton_bins = species_groups["non_proton"]
+    proton_bins = species_groups["proton"]
+
+    species_bin = np.full(e_bin.shape, fill_value=2, dtype=int)
+    species_bin[np.isin(e_bin, non_proton_bins)] = 0
+    species_bin[np.isin(e_bin, proton_bins)] = 1
 
     return species_bin
 
@@ -933,7 +949,7 @@ def get_spin_number(de_met: NDArray, de_spin: NDArray) -> NDArray:
     possible_spins = spin_numbers & 0xFF
 
     # Assign each group based on time.
-    for start, end in zip(spin_start_indices, spin_end_indices):
+    for start, end in zip(spin_start_indices, spin_end_indices, strict=False):
         # Now that we have the possible spins from the Universal Spin Table,
         # we match the times of those spins to the nearest times in the DE data.
         possible_times = spin_start_mets[possible_spins == de_spin_sorted[start]]
@@ -967,9 +983,9 @@ def get_eventtimes(
     Returns
     -------
     event_times : np.ndarray
-        Event times.
+        Event times in et.
     spin_starts : np.ndarray
-        Spin start times.
+        Spin start times in et.
     spin_period_sec : np.ndarray
         Spin period in seconds.
 
@@ -990,7 +1006,7 @@ def get_eventtimes(
 
     event_times = spin_starts + spin_period_sec * (phase_angle / 720)
 
-    return event_times, spin_starts, spin_period_sec
+    return sct_to_et(event_times), sct_to_et(spin_starts), spin_period_sec
 
 
 def interpolate_fwhm(
@@ -1030,8 +1046,11 @@ def interpolate_fwhm(
     )
 
     # Note: will return nan for those out-of-bounds inputs.
-    phi_interp = interp_phi((energy, phi_inst))
-    theta_interp = interp_theta((energy, theta_inst))
+    phi_vals = interp_phi((energy, phi_inst))
+    theta_vals = interp_theta((energy, theta_inst))
+
+    phi_interp = np.where(np.isnan(phi_vals), FILLVAL_FLOAT32, phi_vals)
+    theta_interp = np.where(np.isnan(theta_vals), FILLVAL_FLOAT32, theta_vals)
 
     return phi_interp, theta_interp
 
@@ -1069,8 +1088,8 @@ def get_fwhm(
     theta_interp : NDArray
         Interpolated theta FWHM values.
     """
-    phi_interp = np.full_like(phi_inst, np.nan, dtype=np.float64)
-    theta_interp = np.full_like(theta_inst, np.nan, dtype=np.float64)
+    phi_interp = np.full_like(phi_inst, FILLVAL_FLOAT64, dtype=np.float64)
+    theta_interp = np.full_like(theta_inst, FILLVAL_FLOAT64, dtype=np.float64)
     lt_table = get_angular_profiles("left", sensor, ancillary_files)
     rt_table = get_angular_profiles("right", sensor, ancillary_files)
 
@@ -1089,20 +1108,12 @@ def get_fwhm(
     return phi_interp, theta_interp
 
 
-def get_efficiency(
-    energy: NDArray, phi_inst: NDArray, theta_inst: NDArray, ancillary_files: dict
-) -> NDArray:
+def get_efficiency_interpolator(ancillary_files: dict) -> RegularGridInterpolator:
     """
-    Interpolate efficiency values for each event.
+    Return a callable function that interpolates efficiency values for each event.
 
     Parameters
     ----------
-    energy : NDArray
-        Energy values for each event.
-    phi_inst : NDArray
-        Instrument-frame azimuth angle for each event.
-    theta_inst : NDArray
-        Instrument-frame elevation angle for each event.
     ancillary_files : dict
         Ancillary files.
 
@@ -1127,14 +1138,54 @@ def get_efficiency(
         (theta_vals, phi_vals, energy_vals),
         efficiency_grid,
         bounds_error=False,
-        fill_value=np.nan,
+        fill_value=FILLVAL_FLOAT32,
     )
 
+    return interpolator
+
+
+def get_efficiency(
+    energy: NDArray,
+    phi_inst: NDArray,
+    theta_inst: NDArray,
+    ancillary_files: dict,
+    interpolator: RegularGridInterpolator = None,
+) -> np.ndarray:
+    """
+    Return interpolated efficiency values for each event.
+
+    Parameters
+    ----------
+    energy : NDArray
+        Energy values for each event.
+    phi_inst : NDArray
+        Instrument-frame azimuth angle for each event.
+    theta_inst : NDArray
+        Instrument-frame elevation angle for each event.
+    ancillary_files : dict
+        Ancillary files.
+    interpolator : RegularGridInterpolator, optional
+        Precomputed interpolator to use for efficiency lookup.
+        If None, a new interpolator will be created from the ancillary files.
+
+    Returns
+    -------
+    efficiency : NDArray
+        Interpolated efficiency values.
+    """
+    if not interpolator:
+        interpolator = get_efficiency_interpolator(ancillary_files)
+
     return interpolator((theta_inst, phi_inst, energy))
 
 
 def determine_ebin_pulse_height(
-    energy: np.ndarray, tof: np.ndarray, path_length: np.ndarray
+    energy: NDArray,
+    tof: NDArray,
+    path_length: NDArray,
+    backtofvalid: NDArray,
+    coinphvalid: NDArray,
+    ancillary_files: dict,
 ) -> NDArray:
     """
     Determine the species for pulse-height events.
@@ -1152,12 +1203,18 @@ def determine_ebin_pulse_height(
 
     Parameters
     ----------
-    energy : np.ndarray
+    energy : NDArray
         Energy from the PH event (keV).
-    tof : np.ndarray
+    tof : NDArray
         Time of flight of the PH event (tenths of a nanosecond).
-    path_length : np.ndarray
+    path_length : NDArray
         Path length (r) (hundredths of a millimeter).
+    backtofvalid : NDArray
+        Boolean array indicating if the back TOF is valid.
+    coinphvalid : NDArray
+        Boolean array indicating if the Coincidence PH is valid.
+    ancillary_files : dict
+        Ancillary files containing the lookup tables.
 
     Returns
     -------
@@ -1166,15 +1223,22 @@ def determine_ebin_pulse_height(
     """
     # PH event TOF normalization to Z axis
     ctof, _ = get_ctof(tof, path_length, type="PH")
-    # TODO: need lookup tables
-    # placeholder
-    ebin = np.full(len(ctof), 255, dtype=np.uint8)
 
-    return ebin
+    ebins = np.full(path_length.shape, FILLVAL_UINT8, dtype=np.uint8)
+    valid = backtofvalid & coinphvalid
+    ebins[valid] = get_ebins(
+        "l1b-tofxph", energy[valid], ctof[valid], ebins[valid], ancillary_files
+    )
+
+    return ebins
 
 
 def determine_ebin_ssd(
-    energy: np.ndarray, tof: np.ndarray, path_length: np.ndarray
+    energy: NDArray,
+    tof: NDArray,
+    path_length: NDArray,
+    sensor: str,
+    ancillary_files: dict,
 ) -> NDArray:
     """
     Determine the species for SSD events.
@@ -1194,29 +1258,207 @@ def determine_ebin_ssd(
 
     Parameters
     ----------
-    energy : np.ndarray
+    energy : NDArray
         Energy from the SSD event (keV).
-    tof : np.ndarray
+    tof : NDArray
         Time of flight of the SSD event (tenths of a nanosecond).
-    path_length : np.ndarray
+    path_length : NDArray
         Path length (r) (hundredths of a millimeter).
+    sensor : str
+        Sensor name: "ultra45" or "ultra90".
+    ancillary_files : dict
+        Ancillary files containing the lookup tables.
 
     Returns
     -------
-    bin : np.ndarray
+    bin : NDArray
         Species bin.
     """
     # SSD event TOF normalization to Z axis
     ctof, _ = get_ctof(tof, path_length, type="SSD")
 
-    ebin = np.full(len(ctof), 255, dtype=np.uint8)  # placeholder
+    ebins = np.full(path_length.shape, FILLVAL_UINT8, dtype=np.uint8)
+    steep_path_length = get_image_params("PathSteepThresh", sensor, ancillary_files)
+    medium_path_length = get_image_params("PathMediumThresh", sensor, ancillary_files)
+
+    steep_mask = path_length < steep_path_length
+    medium_mask = (path_length >= steep_path_length) & (
+        path_length < medium_path_length
+    )
+    flat_mask = path_length >= medium_path_length
+
+    ebins[steep_mask] = get_ebins(
+        f"l1b-{sensor[5::]}sensor-tofxesteep",
+        energy[steep_mask],
+        ctof[steep_mask],
+        ebins[steep_mask],
+        ancillary_files,
+    )
+    ebins[medium_mask] = get_ebins(
+        f"l1b-{sensor[5::]}sensor-tofxemedium",
+        energy[medium_mask],
+        ctof[medium_mask],
+        ebins[medium_mask],
+        ancillary_files,
+    )
+    ebins[flat_mask] = get_ebins(
+        f"l1b-{sensor[5::]}sensor-tofxeflat",
+        energy[flat_mask],
+        ctof[flat_mask],
+        ebins[flat_mask],
+        ancillary_files,
+    )
+
+    return ebins
+
+
+def is_back_tof_valid(
+    de_dataset: xarray.Dataset,
+    xf: NDArray,
+    sensor: str,
+    ancillary_files: dict,
+) -> NDArray:
+    """
+    Determine whether back TOF is valid based on stop type.
+
+    Parameters
+    ----------
+    de_dataset : xarray.Dataset
+        Data in xarray format.
+    xf : NDArray
+        X front position in (hundredths of a millimeter).
+        Has same length as de_dataset.
+    sensor : str
+        Sensor name: "ultra45" or "ultra90".
+    ancillary_files : dict
+        Ancillary files for lookup.
+
+    Returns
+    -------
+    valid_mask : NDArray
+        Boolean array indicating whether back TOF is valid.
+
+    Notes
+    -----
+    From page 33 of the IMAP-Ultra Flight Software Specification document.
+    """
+    _, _, _, _, tofx, tofy = get_ph_tof_and_back_positions(
+        de_dataset, xf, "ultra45", ancillary_files
+    )
+    diff = tofy - tofx
+
+    indices = np.nonzero(
+        np.isin(de_dataset["stop_type"], [StopType.Top.value, StopType.Bottom.value])
+    )[0]
+    de_ph = de_dataset.isel(epoch=indices)
+
+    top_mask = de_ph["stop_type"] == StopType.Top.value
+    bottom_mask = de_ph["stop_type"] == StopType.Bottom.value
+
+    valid = np.zeros_like(diff, dtype=bool)
+
+    diff_tp_min = get_image_params("TOFDiffTpMin", sensor, ancillary_files)
+    diff_tp_max = get_image_params("TOFDiffTpMax", sensor, ancillary_files)
+    diff_bt_min = get_image_params("TOFDiffBtMin", sensor, ancillary_files)
+    diff_bt_max = get_image_params("TOFDiffBtMax", sensor, ancillary_files)
+
+    valid[top_mask] = (diff[top_mask] >= diff_tp_min) & (diff[top_mask] <= diff_tp_max)
+    valid[bottom_mask] = (diff[bottom_mask] >= diff_bt_min) & (
+        diff[bottom_mask] <= diff_bt_max
+    )
+
+    return valid
+
+
+def is_coin_ph_valid(
+    etof: NDArray,
+    xc: NDArray,
+    xb: NDArray,
+    stop_north_tdc: NDArray,
+    stop_south_tdc: NDArray,
+    stop_east_tdc: NDArray,
+    stop_west_tdc: NDArray,
+    sensor: str,
+    ancillary_files: dict,
+    quality_flags: NDArray,
+) -> NDArray:
+    """
+    Determine event validity.
+
+    Parameters
+    ----------
+    etof : NDArray
+        Time for the electrons to travel back to the coincidence
+        anode (tenths of a nanosecond).
+    xc : NDArray
+        X coincidence position (hundredths of a millimeter).
+    xb : NDArray
+        Back positions in x direction (hundredths of a millimeter).
+    stop_north_tdc : NDArray
+        Stop North Time to Digital Converter.
+    stop_south_tdc : NDArray
+        Stop South Time to Digital Converter.
+    stop_east_tdc : NDArray
+        Stop East Time to Digital Converter.
+    stop_west_tdc : NDArray
+        Stop West Time to Digital Converter.
+    sensor : str
+        Sensor name: "ultra45" or "ultra90".
+    ancillary_files : dict
+        Ancillary files for lookup.
+    quality_flags : NDArray
+        Quality flag to set when there is an outlier.
+
+    Returns
+    -------
+    combined_mask : NDArray
+        Boolean array indicating whether back TOF is valid.
+
+    Notes
+    -----
+    From page 36 of the IMAP-Ultra Flight Software Specification document.
+    """
+    # Make certain etof is within range for tenths of a nanosecond.
+    etof_valid = (etof >= UltraConstants.ETOFMIN_EVENTFILTER) & (
+        etof <= UltraConstants.ETOFMAX_EVENTFILTER
+    )
+
+    # Hundredths of a mm.
+    diff_x = xc - xb
+
+    t1 = (
+        (etof - UltraConstants.ETOFOFF1_EVENTFILTER)
+        * UltraConstants.ETOFSLOPE1_EVENTFILTER
+        / 1024
+    )
+    t2 = (
+        (etof - UltraConstants.ETOFOFF2_EVENTFILTER)
+        * UltraConstants.ETOFSLOPE2_EVENTFILTER
+        / 1024
+    )
+
+    condition_1 = (diff_x >= t1) & (diff_x <= t2)
+    condition_2 = (diff_x >= -t2) & (diff_x <= -t1)
+
+    spatial_valid = condition_1 | condition_2
+
+    sp_n_norm = get_norm(stop_north_tdc, "SpN", sensor, ancillary_files)
+    sp_s_norm = get_norm(stop_south_tdc, "SpS", sensor, ancillary_files)
+    sp_e_norm = get_norm(stop_east_tdc, "SpE", sensor, ancillary_files)
+    sp_w_norm = get_norm(stop_west_tdc, "SpW", sensor, ancillary_files)
+
+    tofx = sp_n_norm + sp_s_norm
+    tofy = sp_e_norm + sp_w_norm
+
+    # Units in tenths of a nanosecond
+    delta_tof = tofy - tofx
+
+    delta_tof_mask = (delta_tof >= UltraConstants.TOFDIFFTPMIN_EVENTFILTER) & (
+        delta_tof <= UltraConstants.TOFDIFFTPMAX_EVENTFILTER
+    )
+
+    combined_mask = etof_valid & spatial_valid & delta_tof_mask
 
-    # TODO: get these lookup tables
-    # if r < get_image_params("PathSteepThresh"):
-    #     # bin = ExTOFSpeciesSteep[energy, ctof]
-    # elif r < get_image_params("PathMediumThresh"):
-    #     # bin = ExTOFSpeciesMedium[energy, ctof]
-    # else:
-    #     # bin = ExTOFSpeciesFlat[energy, ctof]
+    quality_flags[~combined_mask] |= ImapDEOutliersUltraFlags.COINPH.value
 
-    return ebin
+    return combined_mask