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

imap_processing/ultra/l1c/helio_pset.py

@@ -1,27 +1,47 @@
 """Calculate Pointing Set Grids."""
 
+import logging
+
+import astropy_healpix.healpy as hp
 import numpy as np
-import pandas as pd
 import xarray as xr
 
-from imap_processing.spice.time import sct_to_et
+from imap_processing.quality_flags import ImapPSETUltraFlags
+from imap_processing.spice.repoint import get_pointing_times
+from imap_processing.spice.time import (
+    et_to_met,
+    met_to_ttj2000ns,
+    ttj2000ns_to_et,
+)
+from imap_processing.ultra.l1b.ultra_l1b_culling import get_de_rejection_mask
+from imap_processing.ultra.l1c.l1c_lookup_utils import (
+    calculate_fwhm_spun_scattering,
+    get_spacecraft_pointing_lookup_tables,
+)
+from imap_processing.ultra.l1c.ultra_l1c_culling import compute_culling_mask
 from imap_processing.ultra.l1c.ultra_l1c_pset_bins import (
     build_energy_bins,
-    get_helio_background_rates,
-    get_helio_exposure_times,
-    get_helio_sensitivity,
+    get_efficiencies_and_geometric_function,
+    get_energy_delta_minus_plus,
+    get_helio_adjusted_data,
+    get_spacecraft_exposure_times,
     get_spacecraft_histogram,
 )
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
+logger = logging.getLogger(__name__)
+
 
 def calculate_helio_pset(
     de_dataset: xr.Dataset,
-    extendedspin_dataset: xr.Dataset,
-    cullingmask_dataset: xr.Dataset,
+    goodtimes_dataset: xr.Dataset,
+    rates_dataset: xr.Dataset,
+    params_dataset: xr.Dataset,
     name: str,
     ancillary_files: dict,
-) -> xr.Dataset:
+    instrument_id: int,
+    species_id: list,
+) -> xr.Dataset | None:
     """
     Create dictionary with defined datatype for Pointing Set Grid Data.
 
@@ -29,14 +49,20 @@ def calculate_helio_pset(
     ----------
     de_dataset : xarray.Dataset
         Dataset containing de data.
-    extendedspin_dataset : xarray.Dataset
-        Dataset containing extendedspin data.
-    cullingmask_dataset : xarray.Dataset
-        Dataset containing cullingmask data.
+    goodtimes_dataset : xarray.Dataset
+        Dataset containing goodtimes data.
+    rates_dataset : xarray.Dataset
+        Dataset containing image rates data.
+    params_dataset : xarray.Dataset
+        Dataset containing image parameters data.
     name : str
         Name of the dataset.
     ancillary_files : dict
         Ancillary files.
+    instrument_id : int
+        Instrument ID, either 45 or 90.
+    species_id : List
+        Species ID.
 
     Returns
     -------
@@ -44,56 +70,132 @@ def calculate_helio_pset(
         Dataset containing the data.
     """
     pset_dict: dict[str, np.ndarray] = {}
+    # Select only the species we are interested in.
+    indices = np.where(np.isin(de_dataset["e_bin"].values, species_id))[0]
+    species_dataset = de_dataset.isel(epoch=indices)
+
+    rejected = get_de_rejection_mask(
+        species_dataset["quality_scattering"].values,
+        species_dataset["quality_outliers"].values,
+    )
+    species_dataset = species_dataset.isel(epoch=~rejected)
 
     v_mag_helio_spacecraft = np.linalg.norm(
-        de_dataset["velocity_dps_helio"].values, axis=1
+        species_dataset["velocity_dps_helio"].values, axis=1
     )
     vhat_dps_helio = (
-        de_dataset["velocity_dps_helio"].values / v_mag_helio_spacecraft[:, np.newaxis]
+        species_dataset["velocity_dps_helio"].values
+        / v_mag_helio_spacecraft[:, np.newaxis]
     )
     intervals, _, energy_bin_geometric_means = build_energy_bins()
+    # Get lookup table for FOR indices by spin phase step
+    (
+        for_indices_by_spin_phase,
+        theta_vals,
+        phi_vals,
+        ra_and_dec,
+        boundary_scale_factors,
+    ) = get_spacecraft_pointing_lookup_tables(ancillary_files, instrument_id)
+
+    logger.info("calculating spun FWHM scattering values.")
+    pixels_below_scattering, scattering_theta, scattering_phi, scattering_thresholds = (
+        calculate_fwhm_spun_scattering(
+            for_indices_by_spin_phase,
+            theta_vals,
+            phi_vals,
+            ancillary_files,
+            instrument_id,
+        )
+    )
+
+    nside = hp.npix2nside(for_indices_by_spin_phase.shape[0])
     counts, latitude, longitude, n_pix = get_spacecraft_histogram(
         vhat_dps_helio,
-        de_dataset["energy_heliosphere"].values,
+        species_dataset["energy_heliosphere"].values,
         intervals,
-        nside=128,
+        nside=nside,
+    )
+    helio_pset_quality_flags = np.full(
+        n_pix, ImapPSETUltraFlags.NONE.value, dtype=np.uint16
     )
-
     healpix = np.arange(n_pix)
 
-    # calculate background rates
-    background_rates = get_helio_background_rates()
-
-    efficiencies = ancillary_files["l1c-90sensor-efficiencies"]
-    geometric_function = ancillary_files["l1c-90sensor-gf"]
+    logger.info("Calculating spacecraft exposure times with deadtime correction.")
+    exposure_time, deadtime_ratios = get_spacecraft_exposure_times(
+        rates_dataset,
+        params_dataset,
+        pixels_below_scattering,
+        boundary_scale_factors,
+        n_pix=n_pix,
+    )
+    logger.info("Calculating spun efficiencies and geometric function.")
+    # calculate efficiency and geometric function as a function of energy
+    efficiencies, geometric_function = get_efficiencies_and_geometric_function(
+        pixels_below_scattering,
+        boundary_scale_factors,
+        theta_vals,
+        phi_vals,
+        n_pix,
+        ancillary_files,
+    )
+    # Get midpoint timestamp for pointing.
+    pointing_start, pointing_stop = get_pointing_times(
+        et_to_met(species_dataset["event_times"].data[0])
+    )
+    mid_time = ttj2000ns_to_et(met_to_ttj2000ns((pointing_start + pointing_stop) / 2))
 
-    df_efficiencies = pd.read_csv(efficiencies)
-    df_geometric_function = pd.read_csv(geometric_function)
-    mid_time = sct_to_et(np.median(de_dataset["event_times"].data))
-    sensitivity = get_helio_sensitivity(
+    logger.info("Adjusting data for helio frame.")
+    exposure_time, efficiency, geometric_function = get_helio_adjusted_data(
         mid_time,
-        df_efficiencies,
-        df_geometric_function,
+        exposure_time,
+        geometric_function,
+        efficiencies,
+        ra_and_dec[:, 0],
+        ra_and_dec[:, 1],
+        nside=nside,
     )
+    sensitivity = efficiencies * geometric_function
 
-    # Calculate exposure
-    constant_exposure = ancillary_files["l1c-90sensor-dps-exposure"]
-    df_exposure = pd.read_csv(constant_exposure)
-    exposure_pointing = get_helio_exposure_times(mid_time, df_exposure)
+    start: float = np.min(species_dataset["event_times"].values)
+    end: float = np.max(species_dataset["event_times"].values)
 
-    # For ISTP, epoch should be the center of the time bin.
-    pset_dict["epoch"] = de_dataset.epoch.data[:1].astype(np.int64)
+    # Time bins in 30 minute intervals
+    time_bins = np.arange(start, end + 1800, 1800)
+
+    # Compute mask for culling the Earth
+    compute_culling_mask(
+        time_bins,
+        6378.1,  # Earth radius
+        helio_pset_quality_flags,
+        nside=nside,
+    )
+    pointing_start = met_to_ttj2000ns(pointing_start)
+    # Epoch should be the start of the pointing
+    pset_dict["epoch"] = np.atleast_1d(pointing_start).astype(np.int64)
     pset_dict["counts"] = counts[np.newaxis, ...]
     pset_dict["latitude"] = latitude[np.newaxis, ...]
     pset_dict["longitude"] = longitude[np.newaxis, ...]
     pset_dict["energy_bin_geometric_mean"] = energy_bin_geometric_means
-    pset_dict["background_rates"] = background_rates[np.newaxis, ...]
-    pset_dict["helio_exposure_factor"] = exposure_pointing[np.newaxis, ...]
+    pset_dict["helio_exposure_factor"] = exposure_time
     pset_dict["pixel_index"] = healpix
     pset_dict["energy_bin_delta"] = np.diff(intervals, axis=1).squeeze()[
         np.newaxis, ...
     ]
-    pset_dict["sensitivity"] = sensitivity[np.newaxis, ...]
+    pset_dict["sensitivity"] = sensitivity
+    pset_dict["efficiency"] = efficiencies
+    pset_dict["geometric_function"] = geometric_function
+    pset_dict["dead_time_ratio"] = deadtime_ratios
+    pset_dict["spin_phase_step"] = np.arange(len(deadtime_ratios))
+    pset_dict["quality_flags"] = helio_pset_quality_flags[np.newaxis, ...]
+
+    pset_dict["scatter_theta"] = scattering_theta
+    pset_dict["scatter_phi"] = scattering_phi
+    pset_dict["scatter_threshold"] = scattering_thresholds
+
+    # Add the energy delta plus/minus to the dataset
+    energy_delta_minus, energy_delta_plus = get_energy_delta_minus_plus()
+    pset_dict["energy_delta_minus"] = energy_delta_minus
+    pset_dict["energy_delta_plus"] = energy_delta_plus
 
     dataset = create_dataset(pset_dict, name, "l1c")
 

imap_processing/ultra/l1c/l1c_lookup_utils.py

@@ -0,0 +1,289 @@
+"""Contains tools for lookup tables for l1c."""
+
+import logging
+
+import numpy as np
+import pandas as pd
+from numpy._typing import NDArray
+
+from imap_processing.ultra.l1b.lookup_utils import (
+    get_scattering_coefficients,
+    get_scattering_thresholds,
+    load_scattering_lookup_tables,
+)
+from imap_processing.ultra.l1c.ultra_l1c_pset_bins import build_energy_bins
+
+logger = logging.getLogger(__name__)
+
+
+def mask_below_fwhm_scattering_threshold(
+    theta_coeffs: np.ndarray,
+    phi_coeffs: np.ndarray,
+    energy: np.ndarray,
+    scattering_thresholds: np.ndarray,
+) -> tuple[NDArray, NDArray, NDArray]:
+    """
+    Determine indices of theta and phi values below the FWHM scattering threshold.
+
+    For each phi and theta, calculate the FWHM using the formula:
+    FWHM = A*E^g
+    If Phi FWHM or Theta FWHM > the scattering requirements from the table above,
+    mask the instrument frame pixel.
+
+    Parameters
+    ----------
+    theta_coeffs : NDArray
+        Coefficients for theta FWHM calculation (a and g) for each pixel.
+    phi_coeffs : NDArray
+        Coefficients for phi FWHM calculation (a and g) for each pixel.
+    energy : NDArray
+        Energy corresponding to each theta and phi val in keV.
+    scattering_thresholds : dict
+        Scattering thresholds corresponding to each energy.
+
+    Returns
+    -------
+    scattering_mask : numpy.ndarray
+        Boolean array indicating indices below the scattering threshold.
+    fwhm_theta : numpy.ndarray
+        Calculated FWHM values for theta.
+    fwhm_phi : numpy.ndarray
+        Calculated FWHM values for phi.
+    """
+    # Calculate FWHM for all pixels and all energies
+    fwhm_theta = theta_coeffs[..., 0:1] * (
+        energy ** theta_coeffs[..., 1:2]
+    )  # (npix, energy.shape[1])
+    fwhm_phi = phi_coeffs[..., 0:1] * (
+        energy ** phi_coeffs[..., 1:2]
+    )  # (npix, energy.shape[1])
+
+    thresholds = scattering_thresholds[np.newaxis, :]  # (1, energy.shape[1])
+
+    # Combine conditions for both theta and phi.
+    # shape = (npix, energy.shape[1])
+    scattering_mask = np.logical_and(fwhm_theta <= thresholds, fwhm_phi <= thresholds)
+    return scattering_mask, fwhm_theta, fwhm_phi
+
+
+def calculate_fwhm_spun_scattering(
+    for_indices_by_spin_phase: np.ndarray,
+    theta_vals: np.ndarray,
+    phi_vals: np.ndarray,
+    ancillary_files: dict,
+    instrument_id: int,
+) -> tuple[list, NDArray, NDArray, NDArray]:
+    """
+    Calculate FWHM scattering values for each pixel, energy bin, and spin phase step.
+
+    This function also calculates a mask for pixels that are below the FWHM threshold.
+
+    Parameters
+    ----------
+    for_indices_by_spin_phase : np.ndarray
+        A 2D boolean array where cols are spin phase steps are rows are HEALPix pixels.
+        True indicates pixels that are within the Field of Regard (FOR) at that
+        spin phase.
+    theta_vals : np.ndarray
+        A 2D array of theta values for each HEALPix pixel at each spin phase step.
+    phi_vals : np.ndarray
+         A 2D array of phi values for each HEALPix pixel at each spin phase step.
+    ancillary_files : dict
+        Dictionary containing ancillary files.
+    instrument_id : int,
+        Instrument ID, either 45 or 90.
+
+    Returns
+    -------
+    pixels_below_scattering : list
+        A Nested list of arrays indicating pixels within the scattering threshold.
+        The outer list indicates spin phase steps, the middle list indicates energy
+        bins, and the inner arrays contain indices indicating pixels that are below
+        the FWHM scattering threshold.
+    scattering_fwhm_theta : NDArray
+        Calculated FWHM scatting values for theta at each energy bin and averaged
+        over spin phase.
+    scattering_fwhm_phi : NDArray
+        Calculated FWHM scatting values for theta at each energy bin and averaged
+        over spin phase.
+    scattering_thresholds_for_energy_mean : NDArray
+        Scattering thresholds corresponding to each energy bin.
+    """
+    # Load scattering coefficient lookup table
+    scattering_luts = load_scattering_lookup_tables(ancillary_files, instrument_id)
+    pixels_below_scattering = []
+    # Get energy bin geometric means
+    energy_bin_geometric_means = build_energy_bins()[2]
+    # Load scattering thresholds for the energy bin geometric means
+    scattering_thresholds_for_energy_mean = get_scattering_thresholds_for_energy(
+        energy_bin_geometric_means, ancillary_files
+    )
+    # Initialize arrays to accumulate FWHM values for averaging
+    fwhm_theta_sum = np.zeros(
+        (len(energy_bin_geometric_means), for_indices_by_spin_phase.shape[0])
+    )
+    fwhm_phi_sum = np.zeros_like(fwhm_theta_sum)
+    sample_count = np.zeros_like(fwhm_theta_sum)
+
+    steps = for_indices_by_spin_phase.shape[1]
+    energies = energy_bin_geometric_means[np.newaxis, :]
+    # The "for_indices_by_spin_phase" lookup table contains the boolean values of each
+    # pixel at each spin phase step, indicating whether the pixel is inside the FOR.
+    # It starts at Spin-phase = 0, and increments in fine steps (1 ms), spinning the
+    # spacecraft in the despun frame. At each iteration, query for the pixels in the
+    # FOR, and calculate whether the FWHM value is below the threshold at the energy.
+    for i in range(steps):
+        # Calculate spin phase for the current iteration
+        for_inds = for_indices_by_spin_phase[:, i]
+
+        # Skip if no pixels in FOR
+        if not np.any(for_inds):
+            logger.info(f"No pixels found in FOR at spin phase step {i}")
+            pixels_below_scattering.append(
+                [
+                    np.array([], dtype=int)
+                    for _ in range(len(energy_bin_geometric_means))
+                ]
+            )
+            continue
+        # Using the lookup table, get the indices of the pixels inside the FOR at
+        # the current spin phase step.
+        theta = theta_vals[for_inds, i]
+        phi = phi_vals[for_inds, i]
+        theta_coeffs, phi_coeffs = get_scattering_coefficients(
+            theta, phi, lookup_tables=scattering_luts
+        )
+        # Get a mask for pixels below the FWHM scattering threshold
+        scattering_mask, fwhm_theta, fwhm_phi = mask_below_fwhm_scattering_threshold(
+            theta_coeffs,
+            phi_coeffs,
+            energies,
+            scattering_thresholds=scattering_thresholds_for_energy_mean,
+        )
+        # Extract pixel indices for each energy
+        for_pixel_indices = np.where(for_inds)[0]
+        pixels_below_scattering_for_energy = []
+
+        for energy_idx in range(len(energy_bin_geometric_means)):
+            valid_pixels = scattering_mask[:, energy_idx]
+            pixels_below_scattering_for_energy.append(for_pixel_indices[valid_pixels])
+
+        pixels_below_scattering.append(pixels_below_scattering_for_energy)
+        # Accumulate FWHM values for averaging
+        fwhm_theta_sum[:, for_inds] += fwhm_theta.T
+        fwhm_phi_sum[:, for_inds] += fwhm_phi.T
+        sample_count[:, for_inds] += 1
+
+    fwhm_phi_avg = np.divide(fwhm_theta_sum, sample_count, where=sample_count != 0)
+    fwhm_theta_avg = np.divide(fwhm_theta_sum, sample_count, where=sample_count != 0)
+    return (
+        pixels_below_scattering,
+        fwhm_theta_avg,
+        fwhm_phi_avg,
+        scattering_thresholds_for_energy_mean,
+    )
+
+
+def get_spacecraft_pointing_lookup_tables(
+    ancillary_files: dict, instrument_id: int
+) -> tuple[NDArray, NDArray, NDArray, NDArray, NDArray]:
+    """
+    Get indices of pixels in the nominal FOR as a function of spin phase.
+
+    This function also returns the theta / phi values in the instrument frame per spin
+    phase, right ascension / declination values in the SC frame, and boundary scale
+    factors for each pixel at each spin phase.
+
+    Parameters
+    ----------
+    ancillary_files : dict[Path]
+        Ancillary files.
+    instrument_id : int
+        Instrument ID, either 45 or 90.
+
+    Returns
+    -------
+    for_indices_by_spin_phase : NDArray
+        A 2D boolean array of shape (npix, n_spin_phase_steps).
+        True indicates pixels that are within the Field of Regard (FOR) at that
+        spin phase.
+    theta_vals : NDArray
+        A 2D array of theta values for each HEALPix pixel at each spin phase step.
+    phi_vals : NDArray
+         A 2D array of phi values for each HEALPix pixel at each spin phase step.
+    ra_and_dec : NDArray
+        A 2D array of right ascension and declination values for each HEALPix pixel.
+    boundary_scale_factors : NDArray
+        A 2D array of boundary scale factors for each HEALPix pixel at each spin phase
+        step.
+    """
+    theta_descriptor = f"l1c-{instrument_id}sensor-sc-pointing-theta"
+    phi_descriptor = f"l1c-{instrument_id}sensor-sc-pointing-phi"
+    index_descriptor = f"l1c-{instrument_id}sensor-sc-pointing-index"
+    bsf_descriptor = f"l1c-{instrument_id}sensor-sc-pointing-bsf"
+
+    theta_vals = pd.read_csv(
+        ancillary_files[theta_descriptor], header=None, skiprows=1
+    ).to_numpy(dtype=float)[:, 2:]
+    phi_vals = pd.read_csv(
+        ancillary_files[phi_descriptor], header=None, skiprows=1
+    ).to_numpy(dtype=float)[:, 2:]
+    index_grid = pd.read_csv(
+        ancillary_files[index_descriptor], header=None, skiprows=1
+    ).to_numpy(dtype=float)
+    boundary_scale_factors = pd.read_csv(
+        ancillary_files[bsf_descriptor], header=None, skiprows=1
+    ).to_numpy(dtype=float)[:, 2:]
+
+    ra_and_dec = index_grid[:, :2]  # Shape (npix, 2)
+    # This array indicates whether each pixel is in the nominal FOR at each spin phase
+    # step (15000 steps for a full rotation with 1 ms resolution).
+    for_indices_by_spin_phase = np.nan_to_num(index_grid[:, 2:], nan=0).astype(
+        bool
+    )  # Shape (npix, 15000)
+    return (
+        for_indices_by_spin_phase,
+        theta_vals,
+        phi_vals,
+        ra_and_dec,
+        boundary_scale_factors,
+    )
+
+
+def get_scattering_thresholds_for_energy(
+    energy: np.ndarray, ancillary_files: dict
+) -> np.ndarray:
+    """
+    Find the scattering thresholds for each energy bin.
+
+    Parameters
+    ----------
+    energy : np.ndarray
+        Array of energy values in keV.
+    ancillary_files : dict
+        Dictionary containing ancillary files.
+
+    Returns
+    -------
+    np.ndarray
+        Array of scattering thresholds for each energy bin.
+    """
+    scattering_thresholds = get_scattering_thresholds(ancillary_files)
+    # Get thresholds for all energies
+    thresholds = []
+    for e in energy:
+        try:
+            threshold = next(
+                threshold
+                for energy_range, threshold in scattering_thresholds.items()
+                if energy_range[0] <= e < energy_range[1]
+            )
+        except StopIteration:
+            logger.warning(
+                f"Energy {e} keV is out of bounds for scattering thresholds. Using"
+                f" zero for as threshold."
+            )
+
+            threshold = 0
+        thresholds.append(threshold)
+    return np.array(thresholds)