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

imap_processing/ultra/l1b/de.py

@@ -3,6 +3,27 @@
 import numpy as np
 import xarray as xr
 
+from imap_processing.cdf.utils import parse_filename_like
+from imap_processing.spice.geometry import SpiceFrame
+from imap_processing.ultra.l1b.ultra_l1b_annotated import (
+    get_annotated_particle_velocity,
+)
+from imap_processing.ultra.l1b.ultra_l1b_extended import (
+    StopType,
+    determine_species_pulse_height,
+    determine_species_ssd,
+    get_coincidence_positions,
+    get_ctof,
+    get_energy_pulse_height,
+    get_energy_ssd,
+    get_front_x_position,
+    get_front_y_position,
+    get_path_length,
+    get_ph_tof_and_back_positions,
+    get_ssd_back_position_and_tof_offset,
+    get_ssd_tof,
+    get_unit_vector,
+)
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
 
@@ -13,46 +34,169 @@ def calculate_de(de_dataset: xr.Dataset, name: str) -> xr.Dataset:
     Parameters
     ----------
     de_dataset : xarray.Dataset
-        Dataset containing direct event data.
+        L1a dataset containing direct event data.
     name : str
-        Name of the dataset.
+        Name of the l1a dataset.
 
     Returns
     -------
     dataset : xarray.Dataset
-        Dataset containing the data.
+        L1b de dataset.
     """
     de_dict = {}
+    sensor = parse_filename_like(name)["sensor"][0:2]
 
-    # Placeholder for calculations
-    epoch = de_dataset.coords["epoch"].values
+    # Instantiate arrays
+    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)
+    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)
+    energy = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
+    # TODO: uint8 fills with zeros instead of nans.
+    #  Confirm with Ultra team what fill values and dtype we want.
+    species_bin = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.uint8)
+    t2 = np.full(len(de_dataset["epoch"]), np.nan, dtype=np.float32)
 
+    # Drop events with invalid start type.
+    de_dataset = de_dataset.where(
+        de_dataset["START_TYPE"] != np.iinfo(np.int64).min, drop=True
+    )
+    # Define epoch.
     de_dict["epoch"] = de_dataset["epoch"]
-    de_dict["x_front"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["y_front"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["x_back"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["y_back"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["x_coin"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["tof_start_stop"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["tof_stop_coin"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["tof_corrected"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["eventtype"] = np.zeros(len(epoch), dtype=np.uint64)
-    de_dict["vx_ultra"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vy_ultra"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vz_ultra"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["energy"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["species"] = np.zeros(len(epoch), dtype=np.uint64)
-    de_dict["event_efficiency"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vx_sc"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vy_sc"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vz_sc"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vx_dps_sc"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vy_dps_sc"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vz_dps_sc"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vx_dps_helio"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vy_dps_helio"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["vz_dps_helio"] = np.zeros(len(epoch), dtype=np.float64)
-    de_dict["eventtimes"] = np.zeros(len(epoch), dtype=np.float64)
+
+    xf = get_front_x_position(
+        de_dataset["START_TYPE"].data,
+        de_dataset["START_POS_TDC"].data,
+    )
+
+    # 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}")
+    )
+    d[ph_indices], yf[ph_indices] = get_front_y_position(
+        de_dataset["START_TYPE"].data[ph_indices], yb[ph_indices]
+    )
+    energy[ph_indices] = get_energy_pulse_height(
+        de_dataset["STOP_TYPE"].data[ph_indices],
+        de_dataset["ENERGY_PH"].data[ph_indices],
+        xb[ph_indices],
+        yb[ph_indices],
+    )
+    r[ph_indices] = get_path_length(
+        (xf[ph_indices], yf[ph_indices]),
+        (xb[ph_indices], yb[ph_indices]),
+        d[ph_indices],
+    )
+    species_bin[ph_indices] = determine_species_pulse_height(
+        energy[ph_indices], tof[ph_indices], r[ph_indices]
+    )
+    etof[ph_indices], xc[ph_indices] = get_coincidence_positions(
+        de_dataset.isel(epoch=ph_indices), t2[ph_indices], f"ultra{sensor}"
+    )
+    ctof[ph_indices] = get_ctof(tof[ph_indices], r[ph_indices], "PH")
+
+    # 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))
+    xb[ssd_indices] = np.zeros(len(ssd_indices))
+    etof[ssd_indices] = np.zeros(len(ssd_indices))
+    d[ssd_indices], yf[ssd_indices] = get_front_y_position(
+        de_dataset["START_TYPE"].data[ssd_indices], yb[ssd_indices]
+    )
+    energy[ssd_indices] = get_energy_ssd(de_dataset, ssd_number)
+    r[ssd_indices] = get_path_length(
+        (xf[ssd_indices], yf[ssd_indices]),
+        (xb[ssd_indices], yb[ssd_indices]),
+        d[ssd_indices],
+    )
+    species_bin[ssd_indices] = determine_species_ssd(
+        energy[ssd_indices],
+        tof[ssd_indices],
+        r[ssd_indices],
+    )
+    ctof[ssd_indices] = get_ctof(tof[ssd_indices], r[ssd_indices], "SSD")
+
+    # Combine ph_yb and ssd_yb along with their indices
+    de_dict["x_front"] = xf.astype(np.float32)
+    de_dict["y_front"] = yf
+    de_dict["x_back"] = xb
+    de_dict["y_back"] = yb
+    de_dict["x_coin"] = xc
+    de_dict["tof_start_stop"] = tof
+    de_dict["tof_stop_coin"] = etof
+    de_dict["tof_corrected"] = ctof
+    de_dict["front_back_distance"] = d
+    de_dict["path_length"] = r
+
+    keys = [
+        "coincidence_type",
+        "start_type",
+        "event_type",
+        "de_event_met",
+    ]
+    dataset_keys = ["COIN_TYPE", "START_TYPE", "STOP_TYPE", "SHCOARSE"]
+
+    de_dict.update(
+        {key: de_dataset[dataset_key] for key, dataset_key in zip(keys, dataset_keys)}
+    )
+
+    vx_ultra, vy_ultra, vz_ultra = get_unit_vector(
+        (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"],
+    )
+
+    de_dict["vx_ultra"] = vx_ultra.astype(np.float32)
+    de_dict["vy_ultra"] = vy_ultra.astype(np.float32)
+    de_dict["vz_ultra"] = vz_ultra.astype(np.float32)
+    de_dict["energy"] = energy
+    de_dict["species"] = species_bin
+
+    # Annotated Events.
+    position = np.stack(
+        (de_dict["vx_ultra"], de_dict["vy_ultra"], de_dict["vz_ultra"]), axis=-1
+    )
+
+    ultra_frame = getattr(SpiceFrame, f"IMAP_ULTRA_{sensor}")
+    sc_velocity, sc_dps_velocity, helio_velocity = get_annotated_particle_velocity(
+        de_dataset.data_vars["EVENTTIMES"],
+        position,
+        ultra_frame,
+        SpiceFrame.IMAP_DPS,
+        SpiceFrame.IMAP_SPACECRAFT,
+    )
+
+    de_dict["vx_sc"], de_dict["vy_sc"], de_dict["vz_sc"] = (
+        sc_velocity[:, 0],
+        sc_velocity[:, 1],
+        sc_velocity[:, 2],
+    )
+    de_dict["vx_dps_sc"], de_dict["vy_dps_sc"], de_dict["vz_dps_sc"] = (
+        sc_dps_velocity[:, 0],
+        sc_dps_velocity[:, 1],
+        sc_dps_velocity[:, 2],
+    )
+    de_dict["vx_dps_helio"], de_dict["vy_dps_helio"], de_dict["vz_dps_helio"] = (
+        helio_velocity[:, 0],
+        helio_velocity[:, 1],
+        helio_velocity[:, 2],
+    )
+
+    # TODO: TBD.
+    de_dict["event_efficiency"] = np.full(
+        len(de_dataset["epoch"]), np.nan, dtype=np.float32
+    )
 
     dataset = create_dataset(de_dict, name, "l1b")
 

imap_processing/ultra/l1b/ultra_l1b_annotated.py

@@ -1,7 +1,5 @@
 """Calculates Annotated Events for ULTRA L1b."""
 
-import typing
-
 import numpy as np
 
 from imap_processing.spice.geometry import (
@@ -9,17 +7,15 @@ from imap_processing.spice.geometry import (
     frame_transform,
     imap_state,
 )
-from imap_processing.spice.kernels import ensure_spice
 
 
-@ensure_spice
-@typing.no_type_check
-def get_particle_velocity(
+def get_annotated_particle_velocity(
     time: np.ndarray,
     instrument_velocity: np.ndarray,
     instrument_frame: SpiceFrame,
     pointing_frame: SpiceFrame,
-) -> tuple[np.ndarray, np.ndarray]:
+    spacecraft_frame: SpiceFrame,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Get the particle velocity in the pointing (DPS) frame wrt the spacecraft.
 
@@ -33,11 +29,15 @@ def get_particle_velocity(
         Instrument frame.
     pointing_frame : SpiceFrame
         Pointing frame.
+    spacecraft_frame : SpiceFrame
+        Spacecraft frame.
 
     Returns
     -------
     particle_velocity_spacecraft : np.ndarray
         Particle velocity in the spacecraft frame.
+    particle_velocity_dps_spacecraft : np.ndarray
+        Particle velocity in DPS frame at rest WRT spacecraft .
     particle_velocity_heliosphere : np.ndarray
         Particle velocity in the heliosphere frame.
 
@@ -45,8 +45,16 @@ def get_particle_velocity(
     ----------
     https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy
     """
-    # Particle velocity in the pointing (DPS) frame wrt spacecraft.
+    # Particle velocity in the spacecraft frame.
     particle_velocity_spacecraft = frame_transform(
+        et=time,
+        position=instrument_velocity,
+        from_frame=instrument_frame,
+        to_frame=spacecraft_frame,
+    )
+
+    # Particle velocity in the pointing (DPS) frame wrt spacecraft.
+    particle_velocity_dps_spacecraft = frame_transform(
         et=time,
         position=instrument_velocity,
         from_frame=instrument_frame,
@@ -61,6 +69,12 @@ def get_particle_velocity(
 
     # Apply Compton-Getting.
     # Particle velocity in the DPS frame wrt to the heliosphere
-    particle_velocity_heliosphere = spacecraft_velocity + particle_velocity_spacecraft
+    particle_velocity_heliosphere = (
+        spacecraft_velocity + particle_velocity_dps_spacecraft
+    )
 
-    return particle_velocity_spacecraft, particle_velocity_heliosphere
+    return (
+        particle_velocity_spacecraft,
+        particle_velocity_dps_spacecraft,
+        particle_velocity_heliosphere,
+    )

imap_processing/ultra/l1b/ultra_l1b_extended.py

@@ -1,5 +1,6 @@
 """Calculates Extended Raw Events for ULTRA L1b."""
 
+# TODO: Come back and add in FSW logic.
 import logging
 from enum import Enum
 from typing import ClassVar
@@ -252,7 +253,9 @@ def get_ph_tof_and_back_positions(
     return tof, t2, xb, yb
 
 
-def get_path_length(front_position: tuple, back_position: tuple, d: float) -> float:
+def get_path_length(
+    front_position: tuple, back_position: tuple, d: np.ndarray
+) -> NDArray:
     """
     Calculate the path length.
 
@@ -262,15 +265,15 @@ def get_path_length(front_position: tuple, back_position: tuple, d: float) -> fl
         Front position (xf,yf) (hundredths of a millimeter).
     back_position : tuple of floats
         Back position (xb,yb) (hundredths of a millimeter).
-    d : float
+    d : np.ndarray
         Distance from slit to foil (hundredths of a millimeter).
 
     Returns
     -------
-    path_length : float
+    path_length : np.ndarray
         Path length (r) (hundredths of a millimeter).
     """
-    path_length: float = np.sqrt(
+    path_length = np.sqrt(
         (front_position[0] - back_position[0]) ** 2
         + (front_position[1] - back_position[1]) ** 2
         + (d) ** 2
@@ -434,9 +437,9 @@ def get_coincidence_positions(
     return etof, xc_array * 100
 
 
-def get_particle_velocity(
-    front_position: tuple[float, float],
-    back_position: tuple[float, float],
+def get_unit_vector(
+    front_position: tuple[NDArray, NDArray],
+    back_position: tuple[NDArray, NDArray],
     d: np.ndarray,
     tof: np.ndarray,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -486,9 +489,9 @@ def get_particle_velocity(
     vhat_y = -v_y / magnitude_v
     vhat_z = -v_z / magnitude_v
 
-    vhat_x[tof < 0] = np.iinfo(np.int64).min  # used as fillvals
-    vhat_y[tof < 0] = np.iinfo(np.int64).min
-    vhat_z[tof < 0] = np.iinfo(np.int64).min
+    vhat_x[tof < 0] = np.nan  # used as fillvals
+    vhat_y[tof < 0] = np.nan
+    vhat_z[tof < 0] = np.nan
 
     return vhat_x, vhat_y, vhat_z
 
@@ -645,7 +648,7 @@ def get_energy_ssd(de_dataset: xarray.Dataset, ssd: np.ndarray) -> NDArray[np.fl
     return energy_norm
 
 
-def get_ctof(tof: np.ndarray, path_length: np.ndarray) -> NDArray:
+def get_ctof(tof: np.ndarray, path_length: np.ndarray, type: str) -> NDArray:
     """
     Calculate the corrected TOF.
 
@@ -662,14 +665,18 @@ def get_ctof(tof: np.ndarray, path_length: np.ndarray) -> NDArray:
         Time of flight (tenths of a nanosecond).
     path_length : np.ndarray
         Path length (r) (hundredths of a millimeter).
+    type : str
+        Type of event, either "ph" or "ssd".
 
     Returns
     -------
     ctof : np.ndarray
         Corrected TOF (tenths of a ns).
     """
+    dmin_ctof = getattr(UltraConstants, f"DMIN_{type}_CTOF")
+
     # Multiply times 100 to convert to hundredths of a millimeter.
-    ctof = tof * UltraConstants.DMIN * 100 / path_length
+    ctof = tof * dmin_ctof * 100 / path_length
 
     return ctof
 
@@ -706,7 +713,7 @@ def determine_species_pulse_height(
         Species bin.
     """
     # PH event TOF normalization to Z axis
-    ctof = get_ctof(tof, path_length)
+    ctof = get_ctof(tof, path_length, "PH")
     # TODO: need lookup tables
     # placeholder
     bin = np.zeros(len(ctof))
@@ -749,7 +756,7 @@ def determine_species_ssd(
         Species bin.
     """
     # SSD event TOF normalization to Z axis
-    ctof = get_ctof(tof, path_length)
+    ctof = get_ctof(tof, path_length, "SSD")
 
     bin = np.zeros(len(ctof))  # placeholder