imap-processing 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

imap_processing/ultra/l1b/de.py

@@ -4,12 +4,20 @@ import numpy as np
 import xarray as xr
 
 from imap_processing.cdf.utils import parse_filename_like
+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,
+    determine_ebin_ssd,
     determine_species,
     get_coincidence_positions,
     get_ctof,
@@ -25,8 +33,11 @@ from imap_processing.ultra.l1b.ultra_l1b_extended import (
     get_path_length,
     get_ph_tof_and_back_positions,
     get_phi_theta,
+    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
 
@@ -59,7 +70,10 @@ def calculate_de(
 
     # Define epoch and spin.
     de_dict["epoch"] = de_dataset["epoch"].data
-    de_dict["spin"] = de_dataset["spin"].data
+    spin_number = get_spin_number(
+        de_dataset["shcoarse"].values, de_dataset["spin"].values
+    )
+    de_dict["spin"] = spin_number
 
     # Add already populated fields.
     keys = [
@@ -80,7 +94,10 @@ def calculate_de(
     ]
 
     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
@@ -106,8 +123,10 @@ 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)
     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)
@@ -116,40 +135,56 @@ 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)
-    spin_period_sec = 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, 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(
         de_dataset["start_type"].data[valid_indices],
         de_dataset["start_pos_tdc"].data[valid_indices],
         f"ultra{sensor}",
+        ancillary_files,
     )
     start_type[valid_indices] = de_dataset["start_type"].data[valid_indices]
 
     (
         event_times[valid_indices],
         spin_starts[valid_indices],
-        spin_period_sec[valid_indices],
+        _,
     ) = get_eventtimes(
-        de_dataset["spin"].data[valid_indices],
+        de_dict["spin"][valid_indices],
         de_dataset["phase_angle"].data[valid_indices],
     )
 
     # 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}")
+    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]
+        de_dataset["start_type"].data[ph_indices], yb[ph_indices], ancillary_files
     )
-    energy[ph_indices] = get_energy_pulse_height(
+    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],
         f"ultra{sensor}",
+        ancillary_files,
+        quality_flags[ph_indices],
     )
     r[ph_indices] = get_path_length(
         (xf[ph_indices], yf[ph_indices]),
@@ -163,24 +198,48 @@ def calculate_de(
     )
     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}"
+        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],
+        f"ultra{sensor}",
+        ancillary_files,
+    )
+    e_bin[ph_indices] = determine_ebin_pulse_height(
+        energy[ph_indices],
+        tof[ph_indices],
+        r[ph_indices],
+        backtofvalid,
+        coinphvalid,
+        ancillary_files,
     )
     ctof[ph_indices], magnitude_v[ph_indices] = get_ctof(
         tof[ph_indices], r[ph_indices], "PH"
     )
 
     # SSD
-    tof[ssd_indices] = get_ssd_tof(de_dataset, xf, f"ultra{sensor}")
+    tof[ssd_indices] = get_ssd_tof(de_dataset, xf, f"ultra{sensor}", ancillary_files)
     yb[ssd_indices], _, ssd_number = get_ssd_back_position_and_tof_offset(
-        de_dataset, f"ultra{sensor}"
+        de_dataset, f"ultra{sensor}", ancillary_files
     )
     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]
+        de_dataset["start_type"].data[ssd_indices], yb[ssd_indices], ancillary_files
     )
-    energy[ssd_indices] = get_energy_ssd(de_dataset, ssd_number)
+    energy[ssd_indices] = get_energy_ssd(de_dataset, ssd_number, ancillary_files)
     r[ssd_indices] = get_path_length(
         (xf[ssd_indices], yf[ssd_indices]),
         (xb[ssd_indices], yb[ssd_indices]),
@@ -191,6 +250,13 @@ def calculate_de(
         (xb[ssd_indices], yb[ssd_indices]),
         d[ssd_indices],
     )
+    e_bin[ssd_indices] = determine_ebin_ssd(
+        energy[ssd_indices],
+        tof[ssd_indices],
+        r[ssd_indices],
+        f"ultra{sensor}",
+        ancillary_files,
+    )
     species_bin[ssd_indices] = determine_species(
         tof[ssd_indices], r[ssd_indices], "SSD"
     )
@@ -202,7 +268,6 @@ def calculate_de(
     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
@@ -216,18 +281,24 @@ 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_velocity"] = velocities.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["species"] = species_bin
 
     # Annotated Events.
@@ -261,10 +332,35 @@ def calculate_de(
         de_dict["tof_energy"],
         de_dict["phi"],
         de_dict["theta"],
+        ancillary_files,
     )
     de_dict["event_efficiency"] = get_efficiency(
         de_dict["tof_energy"], de_dict["phi"], de_dict["theta"], ancillary_files
     )
+    de_dict["geometric_factor_blades"] = get_geometric_factor(
+        de_dict["phi"],
+        de_dict["theta"],
+        quality_flags,
+        ancillary_files,
+        "l1b-sensor-gf-blades",
+    )
+    de_dict["geometric_factor_noblades"] = get_geometric_factor(
+        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_scattering"] = scattering_quality_flags
 
     dataset = create_dataset(de_dict, name, "l1b")
 

imap_processing/ultra/l1b/extendedspin.py

@@ -4,12 +4,18 @@ import numpy as np
 import xarray as xr
 
 from imap_processing.ultra.l1b.ultra_l1b_culling import (
+    count_rejected_events_per_spin,
     flag_attitude,
-    flag_spin,
+    flag_hk,
+    flag_imap_instruments,
+    flag_rates,
     get_energy_histogram,
+    get_pulses_per_spin,
 )
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
+FILLVAL_UINT16 = 65535
+
 
 def calculate_extendedspin(
     dict_datasets: dict[str, xr.Dataset],
@@ -34,10 +40,11 @@ def calculate_extendedspin(
         Dataset containing the data.
     """
     aux_dataset = dict_datasets[f"imap_ultra_l1a_{instrument_id}sensor-aux"]
+    rates_dataset = dict_datasets[f"imap_ultra_l1a_{instrument_id}sensor-rates"]
     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(
+    rates_qf, spin, energy_midpoints, n_sigma_per_energy = flag_rates(
         de_dataset["spin"].values,
         de_dataset["energy"].values,
     )
@@ -47,12 +54,26 @@ def calculate_extendedspin(
     attitude_qf, spin_rates, spin_period, spin_starttime = flag_attitude(
         de_dataset["spin"].values, aux_dataset
     )
+    # TODO: We will add to this later
+    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.
+    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
@@ -63,8 +84,14 @@ def calculate_extendedspin(
     extendedspin_dict["spin_start_time"] = spin_starttime
     extendedspin_dict["spin_period"] = spin_period
     extendedspin_dict["spin_rate"] = spin_rates
+    extendedspin_dict["start_pulses_per_spin"] = pulses.start_per_spin
+    extendedspin_dict["stop_pulses_per_spin"] = pulses.stop_per_spin
+    extendedspin_dict["coin_pulses_per_spin"] = pulses.coin_per_spin
+    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
+    extendedspin_dict["quality_instruments"] = inst_qf
 
     extendedspin_dataset = create_dataset(extendedspin_dict, name, "l1b")