imap-processing 0.19.3__py3-none-any.whl → 0.19.4__py3-none-any.whl

imap_processing/lo/l1b/lo_l1b.py

@@ -3,9 +3,9 @@
 import logging
 from dataclasses import Field
 from pathlib import Path
-from typing import Any
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
@@ -16,7 +16,11 @@ from imap_processing.lo.l1b.tof_conversions import (
     TOF2_CONV,
     TOF3_CONV,
 )
-from imap_processing.spice.geometry import SpiceFrame, instrument_pointing
+from imap_processing.spice.geometry import (
+    SpiceFrame,
+    cartesian_to_latitudinal,
+    instrument_pointing,
+)
 from imap_processing.spice.repoint import get_pointing_times
 from imap_processing.spice.spin import get_spin_number
 from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_et
@@ -68,11 +72,6 @@ def lo_l1b(sci_dependencies: dict, anc_dependencies: list) -> list[Path]:
         avg_spin_durations_per_cycle = get_avg_spin_durations_per_cycle(
             acq_start, acq_end
         )
-        # get spin angle (0 - 360 degrees) for each DE
-        spin_angle = get_spin_angle(l1a_de)
-        # calculate and set the spin bin based on the spin angle
-        # spin bins are 0 - 60 bins
-        l1b_de = set_spin_bin(l1b_de, spin_angle)
         # set the spin cycle for each direct event
         l1b_de = set_spin_cycle(pointing_start_met, l1a_de, l1b_de)
         # get spin start times for each event
@@ -100,13 +99,13 @@ def lo_l1b(sci_dependencies: dict, anc_dependencies: list) -> list[Path]:
         l1b_de = convert_tofs_to_eu(l1a_de, l1b_de, attr_mgr_l1a, attr_mgr_l1b)
         # set the species for each direct event
         l1b_de = identify_species(l1b_de)
-        # set the badtimes
-        l1b_de = set_bad_times(l1b_de)
         # set the pointing direction for each direct event
         l1b_de = set_pointing_direction(l1b_de)
         # calculate and set the pointing bin based on the spin phase
         # pointing bin is 3600 x 40 bins
         l1b_de = set_pointing_bin(l1b_de)
+        # set the badtimes
+        l1b_de = set_bad_times(l1b_de, anc_dependencies)
 
     return [l1b_de]
 
@@ -275,54 +274,6 @@ def get_avg_spin_durations_per_cycle(
     return avg_spin_durations_per_cycle
 
 
-def get_spin_angle(l1a_de: xr.Dataset) -> np.ndarray[np.float64] | Any:
-    """
-    Get the spin angle (0 - 360 degrees) for each DE.
-
-    Parameters
-    ----------
-    l1a_de : xarray.Dataset
-        The L1A DE dataset.
-
-    Returns
-    -------
-    spin_angle : np.ndarray
-        The spin angle for each DE.
-    """
-    de_times = l1a_de["de_time"].values
-    # DE Time is 12 bit DN. The max possible value is 4096
-    spin_angle = np.array(de_times / 4096 * 360, dtype=np.float64)
-    return spin_angle
-
-
-def set_spin_bin(l1b_de: xr.Dataset, spin_angle: np.ndarray) -> xr.Dataset:
-    """
-    Set the spin bin (0 - 60 bins) for each Direct Event where each bin is 6 degrees.
-
-    Parameters
-    ----------
-    l1b_de : xarray.Dataset
-        The L1B Direct Event dataset.
-    spin_angle : np.ndarray
-        The spin angle (0-360 degrees) for each Direct Event.
-
-    Returns
-    -------
-    l1b_de : xarray.Dataset
-        The L1B DE dataset with the spin bin added.
-    """
-    # Get the spin bin for each DE
-    # Spin bins are 0 - 60 where each bin is 6 degrees
-    spin_bin = (spin_angle // 6).astype(int)
-    l1b_de["spin_bin"] = xr.DataArray(
-        spin_bin,
-        dims=["epoch"],
-        # TODO: Add spin angle to YAML file
-        # attrs=attr_mgr.get_variable_attributes("spin_bin"),
-    )
-    return l1b_de
-
-
 def set_spin_cycle(
     pointing_start_met: float, l1a_de: xr.Dataset, l1b_de: xr.Dataset
 ) -> xr.Dataset:
@@ -734,7 +685,7 @@ def identify_species(l1b_de: xr.Dataset) -> xr.Dataset:
     return l1b_de
 
 
-def set_bad_times(l1b_de: xr.Dataset) -> xr.Dataset:
+def set_bad_times(l1b_de: xr.Dataset, anc_dependencies: list) -> xr.Dataset:
     """
     Set the bad times for each direct event.
 
@@ -742,18 +693,27 @@ def set_bad_times(l1b_de: xr.Dataset) -> xr.Dataset:
     ----------
     l1b_de : xarray.Dataset
         The L1B DE dataset.
+    anc_dependencies : list
+        List of ancillary file paths.
 
     Returns
     -------
     l1b_de : xarray.Dataset
         The L1B DE dataset with the bad times added.
     """
-    # Initialize all times as not bad for now
-    # TODO: Update to set badtimes based on criteria that
-    #  will be defined in the algorithm document
+    badtimes_df = lo_ancillary.read_ancillary_file(
+        next(str(s) for s in anc_dependencies if "bad-times" in str(s))
+    )
+
+    esa_steps = l1b_de["esa_step"].values
+    epochs = l1b_de["epoch"].values
+    spin_bins = l1b_de["spin_bin"].values
+
+    badtimes = set_bad_or_goodtimes(badtimes_df, epochs, esa_steps, spin_bins)
+
     # 1 = badtime, 0 = not badtime
     l1b_de["badtimes"] = xr.DataArray(
-        np.zeros(len(l1b_de["epoch"]), dtype=int),
+        badtimes,
         dims=["epoch"],
         # TODO: Add to yaml
         # attrs=attr_mgr.get_variable_attributes("bad_times"),
@@ -762,6 +722,65 @@ def set_bad_times(l1b_de: xr.Dataset) -> xr.Dataset:
     return l1b_de
 
 
+def set_bad_or_goodtimes(
+    times_df: pd.DataFrame,
+    epochs: np.ndarray,
+    esa_steps: np.ndarray,
+    spin_bins: np.ndarray,
+) -> np.ndarray:
+    """
+    Find the good/bad time flags for each epoch based on the provided times DataFrame.
+
+    Parameters
+    ----------
+    times_df : pd.DataFrame
+        Good or Bad times dataframe containing time ranges and corresponding flags.
+    epochs : np.ndarray
+        Array of epochs in TTJ2000ns format.
+    esa_steps : np.ndarray
+        Array of ESA steps corresponding to each epoch.
+    spin_bins : np.ndarray
+        Array of spin bins corresponding to each epoch.
+
+    Returns
+    -------
+    time_flags : np.ndarray
+        Array of time good or bad time flags for each epoch.
+    """
+    if "BadTime_start" in times_df.columns and "BadTime_end" in times_df.columns:
+        times_start = met_to_ttj2000ns(times_df["BadTime_start"])
+        times_end = met_to_ttj2000ns(times_df["BadTime_end"])
+    elif "GoodTime_start" in times_df.columns and "GoodTime_end" in times_df.columns:
+        times_start = met_to_ttj2000ns(times_df["GoodTime_start"])
+        times_end = met_to_ttj2000ns(times_df["GoodTime_end"])
+    else:
+        raise ValueError("DataFrame must contain either BadTime or GoodTime columns.")
+
+    # Create masks for time and bin ranges using broadcasting
+    # the bin_start and bin_end are 6 degree bins and need to be converted to
+    # 0.1 degree bins to align with the spin_bins, so multiply by 60
+    time_mask = (epochs[:, None] >= times_start) & (epochs[:, None] <= times_end)
+    bin_mask = (spin_bins[:, None] >= times_df["bin_start"].values * 60) & (
+        spin_bins[:, None] <= times_df["bin_end"].values * 60
+    )
+
+    # Combined mask for epochs that fall within the time and bin ranges
+    combined_mask = time_mask & bin_mask
+
+    # Get the time flags for each epoch's esa_step from matching rows
+    time_flags = np.zeros(len(epochs), dtype=int)
+    for epoch_idx in range(len(epochs)):
+        matching_rows = np.where(combined_mask[epoch_idx])[0]
+        if len(matching_rows) > 0:
+            # Use the first matching row
+            row_idx = matching_rows[0]
+            esa_step = esa_steps[epoch_idx]
+            if f"E-Step{esa_step}" in times_df.columns:
+                time_flags[epoch_idx] = times_df[f"E-Step{esa_step}"].iloc[row_idx]
+
+    return time_flags
+
+
 def set_pointing_direction(l1b_de: xr.Dataset) -> xr.Dataset:
     """
     Set the pointing direction for each direct event.
@@ -782,22 +801,31 @@ def set_pointing_direction(l1b_de: xr.Dataset) -> xr.Dataset:
     """
     # Get the pointing bin for each DE
     et = ttj2000ns_to_et(l1b_de["epoch"])
-
-    direction = instrument_pointing(et, SpiceFrame.IMAP_LO_BASE, SpiceFrame.IMAP_DPS)
+    # get the direction in HAE coordinates
+    direction = instrument_pointing(
+        et, SpiceFrame.IMAP_LO_BASE, SpiceFrame.IMAP_HAE, cartesian=True
+    )
     # TODO: Need to ask Lo what to do if a latitude is outside of the
     # +/-2 degree range. Is that possible?
-    l1b_de["direction_lon"] = xr.DataArray(
+    l1b_de["hae_x"] = xr.DataArray(
         direction[:, 0],
         dims=["epoch"],
         # TODO: Add direction_lon to YAML file
-        # attrs=attr_mgr.get_variable_attributes("direction_lon"),
+        # attrs=attr_mgr.get_variable_attributes("hae_x"),
     )
 
-    l1b_de["direction_lat"] = xr.DataArray(
+    l1b_de["hae_y"] = xr.DataArray(
         direction[:, 1],
         dims=["epoch"],
         # TODO: Add direction_lat to YAML file
-        # attrs=attr_mgr.get_variable_attributes("direction_lat"),
+        # attrs=attr_mgr.get_variable_attributes("hae_y"),
+    )
+
+    l1b_de["hae_z"] = xr.DataArray(
+        direction[:, 2],
+        dims=["epoch"],
+        # TODO: Add direction_lat to YAML file
+        # attrs=attr_mgr.get_variable_attributes("hae_z"),
     )
 
     return l1b_de
@@ -807,7 +835,7 @@ def set_pointing_bin(l1b_de: xr.Dataset) -> xr.Dataset:
     """
     Set the pointing bin for each direct event.
 
-    The pointing bins are defined as 3600 bins for longitude and 40 bins for latitude.
+    The pointing bins are defined as 3600 bins for spin and 40 bins for off angle.
     Each bin is 0.1 degrees. The bins are defined as follows:
     Longitude bins: -180 to 180 degrees
     Latitude bins: -2 to 2 degrees
@@ -822,10 +850,16 @@ def set_pointing_bin(l1b_de: xr.Dataset) -> xr.Dataset:
     l1b_de : xarray.Dataset
         The L1B DE dataset with the pointing bins added.
     """
-    # First column: latitudes
-    lats = l1b_de["direction_lat"]
-    # Second column: longitudes
-    lons = l1b_de["direction_lon"]
+    x = l1b_de["hae_x"]
+    y = l1b_de["hae_y"]
+    z = l1b_de["hae_z"]
+    # convert the pointing direction to latitudinal coordinates
+    direction = cartesian_to_latitudinal(np.column_stack((x, y, z)))
+    # first column: radius (Not needed)
+    # second column: longitude
+    lons = direction[:, 1]
+    # third column: latitude
+    lats = direction[:, 2]
 
     # Define bin edges
     # 3600 bins, 0.1° each
@@ -838,18 +872,18 @@ def set_pointing_bin(l1b_de: xr.Dataset) -> xr.Dataset:
     lon_bins = np.digitize(lons, lon_bins) - 1
     lat_bins = np.digitize(lats, lat_bins) - 1
 
-    l1b_de["pointing_bin_lon"] = xr.DataArray(
+    l1b_de["spin_bin"] = xr.DataArray(
         lon_bins,
         dims=["epoch"],
         # TODO: Add pointing_bin_lon to YAML file
-        # attrs=attr_mgr.get_variable_attributes("pointing_bin_lon"),
+        # attrs=attr_mgr.get_variable_attributes("spin_bin"),
     )
 
-    l1b_de["pointing_bin_lat"] = xr.DataArray(
+    l1b_de["off_angle_bin"] = xr.DataArray(
         lat_bins,
         dims=["epoch"],
         # TODO: Add point_bin_lat to YAML file
-        # attrs=attr_mgr.get_variable_attributes("pointing_bin_lat"),
+        # attrs=attr_mgr.get_variable_attributes("spin_bin"),
     )
 
     return l1b_de

imap_processing/lo/l1c/lo_l1c.py

@@ -10,6 +10,7 @@ from scipy.stats import binned_statistic_dd
 
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.lo import lo_ancillary
+from imap_processing.lo.l1b.lo_l1b import set_bad_or_goodtimes
 from imap_processing.spice.repoint import get_pointing_times
 from imap_processing.spice.spin import get_spin_number
 from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_met
@@ -68,7 +69,7 @@ def lo_l1c(sci_dependencies: dict, anc_dependencies: list) -> list[xr.Dataset]:
         logical_source = "imap_lo_l1c_pset"
         l1b_de = sci_dependencies["imap_lo_l1b_de"]
         l1b_goodtimes_only = filter_goodtimes(l1b_de, anc_dependencies)
-
+        # TODO: Need to handle case where no good times are found
         # Set the pointing start and end times based on the first epoch
         pointing_start_met, pointing_end_met = get_pointing_times(
             ttj2000ns_to_met(l1b_goodtimes_only["epoch"][0].item())
@@ -200,19 +201,17 @@ def filter_goodtimes(l1b_de: xr.Dataset, anc_dependencies: list) -> xr.Dataset:
         next(str(s) for s in anc_dependencies if "good-times" in str(s))
     )
 
-    # convert goodtimes from MET to TTJ2000
-    goodtimes_start = met_to_ttj2000ns(goodtimes_table_df["GoodTime_start"])
-    goodtimes_end = met_to_ttj2000ns(goodtimes_table_df["GoodTime_end"])
-
-    # Create a mask for epochs within any of the start/end time ranges
-    goodtimes_mask = np.zeros_like(l1b_de["epoch"], dtype=bool)
+    esa_steps = l1b_de["esa_step"].values
+    epochs = l1b_de["epoch"].values
+    spin_bins = l1b_de["spin_bin"].values
 
-    # Iterate over the good times and create a mask
-    for start, end in zip(goodtimes_start, goodtimes_end, strict=False):
-        goodtimes_mask |= (l1b_de["epoch"] >= start) & (l1b_de["epoch"] < end)
+    # Get array of bools for each epoch 1 = good time, 0 not good time
+    goodtimes_mask = set_bad_or_goodtimes(
+        goodtimes_table_df, epochs, esa_steps, spin_bins
+    )
 
     # Filter the dataset using the mask
-    filtered_epochs = l1b_de.sel(epoch=goodtimes_mask)
+    filtered_epochs = l1b_de.sel(epoch=goodtimes_mask.astype(bool))
 
     return filtered_epochs
 

imap_processing/lo/l2/lo_l2.py

@@ -853,8 +853,8 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
     """
     logger.info("Applying bootstrap corrections")
 
-    # Table 3 bootstrap terms h_i,k
-    bootstrap_factor = np.array(
+    # Table 3 bootstrap terms h_i,k - convert to xarray for better dimension handling
+    bootstrap_factor_array = np.array(
         [
             [0, 0.03, 0.01, 0, 0, 0, 0, 0],
             [0, 0, 0.05, 0.02, 0.01, 0, 0, 0],
@@ -865,6 +865,15 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
             [0, 0, 0, 0, 0, 0, 0, 0.75],
         ]
     )
+    # Create xarray DataArray with named dimensions for proper broadcasting
+    bootstrap_factor = xr.DataArray(
+        bootstrap_factor_array,
+        dims=["energy_i", "energy_k"],
+        coords={
+            "energy_i": list(range(7)),
+            "energy_k": list(range(8)),  # Include virtual channel 7 (index 7)
+        },
+    )
 
     # Equation 14
     bg_intensity = dataset["bg_rates"] / (
@@ -910,6 +919,12 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
     dataset["bootstrap_intensity_sys_err"] = xr.zeros_like(dataset["ena_intensity"])
 
     for i in range(6, -1, -1):
+        # Create views for the current energy channel to avoid repeated indexing
+        bootstrap_intensity_i = dataset["bootstrap_intensity"][0, i, ...]
+        bootstrap_intensity_var_i = dataset["bootstrap_intensity_var"][0, i, ...]
+        j_c_prime_i = j_c_prime[0, i, ...]
+        j_c_prime_var_i = j_c_prime_var[0, i, ...]
+
         # Initialize the variable with the non-summation term and virtual
         # channel energy subtraction first, then iterate through the other
         # channels which can be looked up via indexing
@@ -917,31 +932,37 @@ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
         # included the k=8 term here.
         # NOTE: The paper uses 1-based indexing and we use 0-based indexing
         #       so there is an off-by-one difference in the indices.
-        dataset["bootstrap_intensity"][0, i, ...] = (
-            j_c_prime[0, i, ...] - bootstrap_factor[i, 7] * j_8_b[0, ...]
+        bootstrap_intensity_i[:] = (
+            j_c_prime_i - bootstrap_factor.sel(energy_i=i, energy_k=7) * j_8_b[0, ...]
         )
         # NOTE: We will square root at the end to get the uncertainty, but
         #       all equations are with variances
-        dataset["bootstrap_intensity_var"][0, i, ...] = j_c_prime_var[0, i, ...]
-
-        for k in range(i + 1, 7):
-            logger.debug(
-                f"Subtracting bootstrap factor h_{i},{k} * J_{k}_b from J_{i}_b"
-            )
-            # Subtraction terms from equations 18-23
-            dataset["bootstrap_intensity"][0, i, ...] -= (
-                bootstrap_factor[i, k] * dataset["bootstrap_intensity"][0, k, ...]
-            )
-
-            # Summation terms from equations 25-30
-            dataset["bootstrap_intensity_var"][0, i, ...] += (
-                bootstrap_factor[i, k] ** 2
-            ) * dataset["bootstrap_intensity_var"][0, k, ...]
+        bootstrap_intensity_var_i[:] = j_c_prime_var_i
+
+        # Vectorized summation using xarray's built-in broadcasting
+        # Select the relevant k indices for summation (k = i+1 to 6)
+        k_indices = list(range(i + 1, 7))
+
+        # Get bootstrap factors for this i and the relevant k values
+        # Rename energy_k dimension to energy for alignment with intensity
+        bootstrap_factors_k = bootstrap_factor.sel(
+            energy_i=i, energy_k=k_indices
+        ).rename({"energy_k": "energy"})
+
+        # Get intensity slices - these will have an 'energy' dimension still
+        intensity_k = dataset["bootstrap_intensity"][0, k_indices, ...]
+        intensity_var_k = dataset["bootstrap_intensity_var"][0, k_indices, ...]
+
+        # Subtraction terms from equations 18-23 (xarray vectorized)
+        bootstrap_intensity_i -= (bootstrap_factors_k * intensity_k).sum(dim="energy")
+
+        # Summation terms from equations 25-30 (xarray vectorized)
+        bootstrap_intensity_var_i += (bootstrap_factors_k**2 * intensity_var_k).sum(
+            dim="energy"
+        )
 
         # Again zero any bootstrap fluxes that are negative
-        dataset["bootstrap_intensity"][0, i, ...].values[
-            dataset["bootstrap_intensity"][0, i, ...] < 0
-        ] = 0.0
+        bootstrap_intensity_i.values[bootstrap_intensity_i < 0] = 0.0
 
     # Equation 31 - systematic error propagation for bootstrap intensity
     # Handle division by zero: only compute where j_c_prime > 0

imap_processing/mag/l1c/interpolation_methods.py

@@ -229,11 +229,13 @@ def cic_filter(
     cic1 = cic1 / decimation_factor
     cic2 = np.convolve(cic1, cic1)
     delay = (len(cic2) - 1) // 2
+
     input_filtered = input_timestamps
+    vectors_filtered = lfilter(cic2, 1, input_vectors, axis=0)
     if delay != 0:
         input_filtered = input_timestamps[:-delay]
+        vectors_filtered = vectors_filtered[delay:]
 
-    vectors_filtered = lfilter(cic2, 1, input_vectors, axis=0)[delay:]
     return input_filtered, vectors_filtered
 
 
@@ -270,6 +272,12 @@ def linear_filtered(
         Interpolated vectors of shape (m, 3) where m is equal to the number of output
         timestamps. Contains x, y, z components of the vector.
     """
+    if input_vectors.shape[0] != input_timestamps.shape[0]:
+        raise ValueError(
+            "Input vectors and input timestamps must have the same length. "
+            f"Got {input_vectors.shape[0]} and {input_timestamps.shape[0]}"
+        )
+
     input_filtered, vectors_filtered = cic_filter(
         input_vectors, input_timestamps, output_timestamps, input_rate, output_rate
     )