imap-processing 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

imap_processing/ultra/l1c/histogram.py

@@ -6,9 +6,7 @@ import xarray as xr
 from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
 
 
-def calculate_histogram(
-    histogram_dataset: xr.Dataset, name: str, data_version: str
-) -> xr.Dataset:
+def calculate_histogram(histogram_dataset: xr.Dataset, name: str) -> xr.Dataset:
     """
     Create dictionary with defined datatype for Histogram Data.
 
@@ -18,8 +16,6 @@ def calculate_histogram(
         Dataset containing histogram data.
     name : str
         Name of dataset.
-    data_version : str
-        Version of the data.
 
     Returns
     -------
@@ -35,6 +31,6 @@ def calculate_histogram(
     histogram_dict["epoch"] = epoch
     histogram_dict["sid"] = np.zeros(len(epoch), dtype=np.uint8)
 
-    dataset = create_dataset(histogram_dict, name, "l1c", data_version)
+    dataset = create_dataset(histogram_dict, name, "l1c")
 
     return dataset

imap_processing/ultra/l1c/spacecraft_pset.py

@@ -0,0 +1,84 @@
+"""Calculate Pointing Set Grids."""
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+from imap_processing import imap_module_directory
+from imap_processing.ultra.l1c.ultra_l1c_pset_bins import (
+    build_energy_bins,
+    get_background_rates,
+    get_spacecraft_exposure_times,
+    get_spacecraft_histogram,
+)
+from imap_processing.ultra.utils.ultra_l1_utils import create_dataset
+
+# TODO: This is a placeholder for the API lookup table directory.
+TEST_PATH = imap_module_directory / "tests" / "ultra" / "data" / "l1"
+
+
+def calculate_spacecraft_pset(
+    de_dataset: xr.Dataset,
+    extendedspin_dataset: xr.Dataset,
+    cullingmask_dataset: xr.Dataset,
+    name: str,
+) -> xr.Dataset:
+    """
+    Create dictionary with defined datatype for Pointing Set Grid Data.
+
+    Parameters
+    ----------
+    de_dataset : xarray.Dataset
+        Dataset containing de data.
+    extendedspin_dataset : xarray.Dataset
+        Dataset containing extendedspin data.
+    cullingmask_dataset : xarray.Dataset
+        Dataset containing cullingmask data.
+    name : str
+        Name of the dataset.
+
+    Returns
+    -------
+    dataset : xarray.Dataset
+        Dataset containing the data.
+    """
+    pset_dict: dict[str, np.ndarray] = {}
+
+    v_mag_dps_spacecraft = np.linalg.norm(de_dataset["velocity_dps_sc"].values, axis=1)
+    vhat_dps_spacecraft = (
+        de_dataset["velocity_dps_sc"].values / v_mag_dps_spacecraft[:, np.newaxis]
+    )
+
+    intervals, _, energy_bin_geometric_means = build_energy_bins()
+    counts, latitude, longitude, n_pix = get_spacecraft_histogram(
+        vhat_dps_spacecraft,
+        de_dataset["energy_spacecraft"].values,
+        intervals,
+        nside=128,
+    )
+    healpix = np.arange(n_pix)
+
+    # calculate background rates
+    background_rates = get_background_rates()
+
+    # TODO: calculate sensitivity and interpolate based on energy.
+
+    # Calculate exposure
+    constant_exposure = TEST_PATH / "ultra_90_dps_exposure.csv"
+    df_exposure = pd.read_csv(constant_exposure)
+    exposure_pointing = get_spacecraft_exposure_times(df_exposure)
+
+    # For ISTP, epoch should be the center of the time bin.
+    pset_dict["epoch"] = de_dataset.epoch.data[0].astype(np.int64)
+    pset_dict["counts"] = counts
+    pset_dict["latitude"] = latitude
+    pset_dict["longitude"] = longitude
+    pset_dict["energy_bin_geometric_mean"] = energy_bin_geometric_means
+    pset_dict["background_rates"] = background_rates
+    pset_dict["exposure_factor"] = exposure_pointing
+    pset_dict["healpix"] = healpix
+    pset_dict["energy_bin_delta"] = np.diff(intervals, axis=1).squeeze()
+
+    dataset = create_dataset(pset_dict, name, "l1c")
+
+    return dataset

imap_processing/ultra/l1c/ultra_l1c.py

@@ -3,10 +3,10 @@
 import xarray as xr
 
 from imap_processing.ultra.l1c.histogram import calculate_histogram
-from imap_processing.ultra.l1c.pset import calculate_pset
+from imap_processing.ultra.l1c.spacecraft_pset import calculate_spacecraft_pset
 
 
-def ultra_l1c(data_dict: dict, data_version: str) -> list[xr.Dataset]:
+def ultra_l1c(data_dict: dict) -> list[xr.Dataset]:
     """
     Will process ULTRA L1A and L1B data into L1C CDF files at output_filepath.
 
@@ -14,8 +14,6 @@ def ultra_l1c(data_dict: dict, data_version: str) -> list[xr.Dataset]:
     ----------
     data_dict : dict
         The data itself and its dependent data.
-    data_version : str
-        Version of the data product being created.
 
     Returns
     -------
@@ -31,7 +29,6 @@ def ultra_l1c(data_dict: dict, data_version: str) -> list[xr.Dataset]:
         histogram_dataset = calculate_histogram(
             data_dict[f"imap_ultra_l1a_{instrument_id}sensor-histogram"],
             f"imap_ultra_l1c_{instrument_id}sensor-histogram",
-            data_version,
         )
         output_datasets = [histogram_dataset]
     elif (
@@ -39,12 +36,14 @@ def ultra_l1c(data_dict: dict, data_version: str) -> list[xr.Dataset]:
         and f"imap_ultra_l1b_{instrument_id}sensor-de" in data_dict
         and f"imap_ultra_l1b_{instrument_id}sensor-extendedspin" in data_dict
     ):
-        pset_dataset = calculate_pset(
+        spacecraft_pset = calculate_spacecraft_pset(
             data_dict[f"imap_ultra_l1b_{instrument_id}sensor-de"],
-            f"imap_ultra_l1c_{instrument_id}sensor-pset",
-            data_version,
+            data_dict[f"imap_ultra_l1b_{instrument_id}sensor-extendedspin"],
+            data_dict[f"imap_ultra_l1b_{instrument_id}sensor-cullingmask"],
+            f"imap_ultra_l1c_{instrument_id}sensor-spacecraftpset",
         )
-        output_datasets = [pset_dataset]
+        # TODO: add calculate_helio_pset here
+        output_datasets = [spacecraft_pset]
     else:
         raise ValueError("Data dictionary does not contain the expected keys.")
 

imap_processing/ultra/l1c/ultra_l1c_pset_bins.py

@@ -1,24 +1,23 @@
 """Module to create pointing sets."""
 
-from pathlib import Path
-
-import cdflib
+import astropy_healpix.healpy as hp
 import numpy as np
+import pandas
+import pandas as pd
 from numpy.typing import NDArray
+from scipy.interpolate import interp1d
 
-from imap_processing.ena_maps.utils.spatial_utils import build_spatial_bins
 from imap_processing.spice.geometry import (
     SpiceFrame,
     cartesian_to_spherical,
     imap_state,
-    spherical_to_cartesian,
 )
 from imap_processing.ultra.constants import UltraConstants
 
 # TODO: add species binning.
 
 
-def build_energy_bins() -> tuple[list[tuple[float, float]], np.ndarray]:
+def build_energy_bins() -> tuple[list[tuple[float, float]], np.ndarray, np.ndarray]:
     """
     Build energy bin boundaries.
 
@@ -28,6 +27,8 @@ def build_energy_bins() -> tuple[list[tuple[float, float]], np.ndarray]:
         Energy bins.
     energy_midpoints : np.ndarray
         Array of energy bin midpoints.
+    energy_bin_geometric_means : np.ndarray
+        Array of geometric means of energy bins.
     """
     # Calculate energy step
     energy_step = (1 + UltraConstants.ALPHA / 2) / (1 - UltraConstants.ALPHA / 2)
@@ -44,19 +45,20 @@ def build_energy_bins() -> tuple[list[tuple[float, float]], np.ndarray]:
         (float(energy_bin_edges[i]), float(energy_bin_edges[i + 1]))
         for i in range(len(energy_bin_edges) - 1)
     ]
+    energy_bin_geometric_means = np.sqrt(energy_bin_edges[:-1] * energy_bin_edges[1:])
 
-    return intervals, energy_midpoints
+    return intervals, energy_midpoints, energy_bin_geometric_means
 
 
-def get_histogram(
+def get_spacecraft_histogram(
     vhat: tuple[np.ndarray, np.ndarray, np.ndarray],
     energy: np.ndarray,
-    az_bin_edges: np.ndarray,
-    el_bin_edges: np.ndarray,
     energy_bin_edges: list[tuple[float, float]],
-) -> NDArray:
+    nside: int = 128,
+    nested: bool = False,
+) -> tuple[NDArray, NDArray, NDArray, NDArray]:
     """
-    Compute a 3D histogram of the particle data.
+    Compute a 3D histogram of the particle data using HEALPix binning.
 
     Parameters
     ----------
@@ -64,123 +66,165 @@ def get_histogram(
         The x,y,z-components of the unit velocity vector.
     energy : np.ndarray
         The particle energy.
-    az_bin_edges : np.ndarray
-        Array of azimuth bin boundary values.
-    el_bin_edges : np.ndarray
-        Array of elevation bin boundary values.
     energy_bin_edges : list[tuple[float, float]]
         Array of energy bin edges.
+    nside : int, optional
+        The nside parameter of the Healpix tessellation.
+        Default is 128.
+    nested : bool, optional
+        Whether the Healpix tessellation is nested. Default is False.
 
     Returns
     -------
     hist : np.ndarray
-        A 3D histogram array.
+        A 3D histogram array with shape (n_pix, n_energy_bins).
+    latitude : np.ndarray
+        Array of latitude values.
+    longitude : np.ndarray
+        Array of longitude values.
+    n_pix : int
+        Number of healpix pixels.
 
     Notes
     -----
-    The histogram will now work properly for overlapping energy bins, i.e.
+    The histogram will work properly for overlapping energy bins, i.e.
     the same energy value can fall into multiple bins if the intervals overlap.
+
+    azimuthal angle [0, 360], elevation angle [-90, 90]
     """
-    spherical_coords = cartesian_to_spherical(vhat)
+    # vhat = direction in which particle is traveling
+    # Make negative to see where it came from
+    spherical_coords = cartesian_to_spherical(-np.array(vhat), degrees=True)
     az, el = (
         spherical_coords[..., 1],
         spherical_coords[..., 2],
     )
 
-    # Initialize histogram
-    hist_total = np.zeros(
-        (len(az_bin_edges) - 1, len(el_bin_edges) - 1, len(energy_bin_edges))
-    )
+    # Compute number of HEALPix pixels that cover the sphere
+    n_pix = hp.nside2npix(nside)
+
+    # Calculate the corresponding longitude (az) latitude (el)
+    # center coordinates
+    longitude, latitude = hp.pix2ang(nside, np.arange(n_pix), lonlat=True)
+
+    # Get HEALPix pixel indices for each event
+    # HEALPix expects latitude in [-90, 90] so we don't need to change elevation
+    hpix_idx = hp.ang2pix(nside, az, el, nest=nested, lonlat=True)
+
+    # Initialize histogram: (n_energy_bins, n_HEALPix pixels)
+    hist = np.zeros((len(energy_bin_edges), n_pix))
 
+    # Bin data in energy & HEALPix space
     for i, (e_min, e_max) in enumerate(energy_bin_edges):
-        # Filter data for current energy bin.
         mask = (energy >= e_min) & (energy < e_max)
-        hist, _ = np.histogramdd(
-            sample=(az[mask], el[mask], energy[mask]),
-            bins=[az_bin_edges, el_bin_edges, [e_min, e_max]],
-        )
-        # Assign 2D histogram to current energy bin.
-        hist_total[:, :, i] = hist[:, :, 0]
+        # Only count the events that fall within the energy bin
+        hist[i, :] += np.bincount(hpix_idx[mask], minlength=n_pix).astype(np.float64)
 
-    return hist_total
+    return hist, latitude, longitude, n_pix
 
 
-def get_pointing_frame_exposure_times(
-    constant_exposure: Path, n_spins: int, sensor: str
+def get_background_rates(
+    nside: int = 128,
 ) -> NDArray:
     """
-    Compute a 2D array of the exposure.
+    Calculate background rates.
 
     Parameters
     ----------
-    constant_exposure : Path
-        Path to file containing constant exposure data.
-    n_spins : int
-        Number of spins per pointing.
-    sensor : str
-        Sensor (45 or 90).
+    nside : int, optional
+        The nside parameter of the Healpix tessellation (default is 128).
 
     Returns
     -------
-    exposure : np.ndarray
-        A 2D array with dimensions (az, el).
+    background_rates : np.ndarray
+        Array of background rates.
+
+    Notes
+    -----
+    This is a placeholder.
     """
-    with cdflib.CDF(constant_exposure) as cdf_file:
-        exposure = cdf_file.varget(f"dps_grid{sensor}") * n_spins
+    n_pix = hp.nside2npix(nside)
+    return np.zeros(n_pix)
 
-    return exposure
+
+def get_spacecraft_exposure_times(constant_exposure: pandas.DataFrame) -> NDArray:
+    """
+    Compute exposure times for HEALPix pixels.
+
+    Parameters
+    ----------
+    constant_exposure : pandas.DataFrame
+        Exposure data.
+
+    Returns
+    -------
+    exposure_pointing : np.ndarray
+        Total exposure times of pixels in a
+        Healpix tessellation of the sky
+        in the pointing (dps) frame.
+    """
+    # TODO: use the universal spin table and
+    #  universal pointing table here to determine actual number of spins
+    exposure_pointing = (
+        constant_exposure["Exposure Time"] * 5760
+    )  # 5760 spins per pointing (for now)
+
+    return exposure_pointing
 
 
 def get_helio_exposure_times(
     time: np.ndarray,
-    sc_exposure: np.ndarray,
+    df_exposure: pd.DataFrame,
+    nside: int = 128,
+    nested: bool = False,
 ) -> NDArray:
     """
-    Compute a 3D array of the exposure in the helio frame.
+    Compute a 2D (Healpix index, energy) array of exposure in the helio frame.
 
     Parameters
     ----------
     time : np.ndarray
-        Median time of pointing in J2000 seconds.
-    sc_exposure : np.ndarray
-        Spacecraft exposure.
+        Median time of pointing in et.
+    df_exposure : pd.DataFrame
+        Spacecraft exposure in healpix coordinates.
+    nside : int, optional
+        The nside parameter of the Healpix tessellation (default is 128).
+    nested : bool, optional
+        Whether the Healpix tessellation is nested (default is False).
 
     Returns
     -------
-    exposure_3d : np.ndarray
-        A 3D array with dimensions (az, el, energy).
+    helio_exposure : np.ndarray
+        A 2D array of shape (npix, n_energy_bins).
 
     Notes
     -----
     These calculations are performed once per pointing.
     """
-    # Get bins and midpoints, and convert from radians to degrees.
-    _, energy_midpoints = build_energy_bins()
-    az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
-        np.rad2deg(angle_radians) for angle_radians in (build_spatial_bins())
-    )
-
-    # Initialize the exposure grid.
-    exposure_3d = np.zeros(
-        (len(el_bin_midpoints), len(az_bin_midpoints), len(energy_midpoints))
-    )
-
-    # Create a 3D Cartesian grid from spherical coordinates
-    # using azimuth and elevation midpoints.
-    az_grid, el_grid = np.meshgrid(az_bin_midpoints, el_bin_midpoints[::-1])
-
-    # Radial distance.
-    r = np.ones(el_grid.shape)
-    spherical_coords = np.stack((r, np.radians(az_grid), np.radians(el_grid)), axis=-1)
-    cartesian_coords = spherical_to_cartesian(spherical_coords)
-    cartesian = cartesian_coords.reshape(-1, 3, order="F").T
-
-    # Spacecraft velocity in the pointing (DPS) frame wrt heliosphere.
+    # Get energy midpoints.
+    _, energy_midpoints, _ = build_energy_bins()
+    # Extract (RA/Dec) and exposure from the spacecraft frame.
+    ra = df_exposure["Right Ascension (deg)"].values
+    dec = df_exposure["Declination (deg)"].values
+    exposure_flat = df_exposure["Exposure Time"].values
+
+    # The Cartesian state vector representing the position and velocity of the
+    # IMAP spacecraft.
     state = imap_state(time, ref_frame=SpiceFrame.IMAP_DPS)
 
     # Extract the velocity part of the state vector
     spacecraft_velocity = state[3:6]
+    # Convert (RA, Dec) angles into 3D unit vectors.
+    # Each unit vector represents a direction in the sky where the spacecraft observed
+    # and accumulated exposure time.
+    unit_dirs = hp.ang2vec(ra, dec, lonlat=True).T  # Shape (N, 3)
+
+    # Initialize output array.
+    # Each row corresponds to a HEALPix pixel, and each column to an energy bin.
+    npix = hp.nside2npix(nside)
+    helio_exposure = np.zeros((npix, len(energy_midpoints)))
 
+    # Loop through energy bins and compute transformed exposure.
     for i, energy_midpoint in enumerate(energy_midpoints):
         # Convert the midpoint energy to a velocity (km/s).
         # Based on kinetic energy equation: E = 1/2 * m * v^2.
@@ -193,56 +237,110 @@ def get_helio_exposure_times(
         # to the velocity wrt heliosphere.
         # energy_velocity * cartesian -> apply the magnitude of the velocity
         # to every position on the grid in the despun grid.
-        helio_velocity = spacecraft_velocity.reshape(3, 1) + energy_velocity * cartesian
+        helio_velocity = spacecraft_velocity.reshape(1, 3) + energy_velocity * unit_dirs
 
         # Normalized vectors representing the direction of the heliocentric velocity.
-        helio_normalized = helio_velocity.T / np.linalg.norm(
-            helio_velocity.T, axis=1, keepdims=True
+        helio_normalized = helio_velocity / np.linalg.norm(
+            helio_velocity, axis=1, keepdims=True
+        )
+
+        # Convert Cartesian heliocentric vectors into spherical coordinates.
+        # Result: azimuth (longitude) and elevation (latitude) in degrees.
+        helio_spherical = cartesian_to_spherical(helio_normalized)
+        az, el = helio_spherical[:, 1], helio_spherical[:, 2]
+
+        # Convert azimuth/elevation directions to HEALPix pixel indices.
+        hpix_idx = hp.ang2pix(nside, az, el, nest=nested, lonlat=True)
+
+        # Accumulate exposure values into HEALPix pixels for this energy bin.
+        helio_exposure[:, i] = np.bincount(
+            hpix_idx, weights=exposure_flat, minlength=npix
         )
-        # Converts vectors from Cartesian coordinates (x, y, z)
-        # into spherical coordinates.
-        spherical_coords = cartesian_to_spherical(helio_normalized)
-        az, el = spherical_coords[..., 1], spherical_coords[..., 2]
 
-        # Assign values from sc_exposure directly to bins.
-        az_idx = np.digitize(az, az_bin_edges) - 1
-        el_idx = np.digitize(el, el_bin_edges[::-1]) - 1
+    return helio_exposure
+
+
+def get_spacecraft_sensitivity(
+    efficiencies: pandas.DataFrame,
+    geometric_function: pandas.DataFrame,
+) -> tuple[pandas.DataFrame, NDArray, NDArray, NDArray]:
+    """
+    Compute sensitivity as efficiency * geometric factor.
+
+    Parameters
+    ----------
+    efficiencies : pandas.DataFrame
+        Efficiencies at different energy levels.
+    geometric_function : pandas.DataFrame
+        Geometric function.
+
+    Returns
+    -------
+    pointing_sensitivity : pandas.DataFrame
+        Sensitivity with dimensions (HEALPIX pixel_number, energy).
+    energy_vals : NDArray
+        Energy values of dataframe.
+    right_ascension : NDArray
+        Right ascension (longitude/azimuth) values of dataframe (0 - 360 degrees).
+    declination : NDArray
+        Declination (latitude/elevation) values of dataframe (-90 to 90 degrees).
+    """
+    # Exclude "Right Ascension (deg)" and "Declination (deg)" from the multiplication
+    energy_columns = [
+        col
+        for col in efficiencies.columns
+        if col not in ["Right Ascension (deg)", "Declination (deg)"]
+    ]
+    sensitivity = efficiencies[energy_columns].mul(
+        geometric_function["Response (cm2-sr)"].values, axis=0
+    )
 
-        # Ensure az_idx and el_idx are within bounds.
-        az_idx = np.clip(az_idx, 0, len(az_bin_edges) - 2)
-        el_idx = np.clip(el_idx, 0, len(el_bin_edges) - 2)
+    right_ascension = efficiencies["Right Ascension (deg)"]
+    declination = efficiencies["Declination (deg)"]
 
-        # A 1D array of linear indices used to track the bin_id.
-        idx = el_idx + az_idx * az_grid.shape[0]
-        # Bins the transposed sc_exposure array.
-        binned_exposure = sc_exposure.T.flatten(order="F")[idx]
-        # Reshape the binned exposure.
-        exposure_3d[:, :, i] = binned_exposure.reshape(az_grid.shape, order="F")
+    energy_vals = np.array([float(col.replace("keV", "")) for col in energy_columns])
 
-    return exposure_3d
+    return sensitivity, energy_vals, right_ascension, declination
 
 
-def get_pointing_frame_sensitivity(
-    constant_sensitivity: Path, n_spins: int, sensor: str
+def grid_sensitivity(
+    efficiencies: pandas.DataFrame,
+    geometric_function: pandas.DataFrame,
+    energy: float,
 ) -> NDArray:
     """
-    Compute a 3D array of the sensitivity.
+    Grid the sensitivity.
 
     Parameters
     ----------
-    constant_sensitivity : Path
-        Path to file containing constant sensitivity data.
-    n_spins : int
-        Number of spins per pointing.
-    sensor : str
-        Sensor (45 or 90).
+    efficiencies : pandas.DataFrame
+        Efficiencies at different energy levels.
+    geometric_function : pandas.DataFrame
+        Geometric function.
+        energy : np.ndarray
+        The particle energy.
+    energy : float
+        Energy to which we are interpolating.
 
     Returns
     -------
-    sensitivity : np.ndarray
-        A 3D array with dimensions (az, el, energy).
+    interpolated_sensitivity : np.ndarray
+        Sensitivity with dimensions (HEALPIX pixel_number, 1).
     """
-    with cdflib.CDF(constant_sensitivity) as cdf_file:
-        sensitivity = cdf_file.varget(f"dps_sensitivity{sensor}") * n_spins
+    sensitivity, energy_vals, right_ascension, declination = get_spacecraft_sensitivity(
+        efficiencies, geometric_function
+    )
+
+    # Create interpolator over energy dimension for each pixel (axis=1)
+    interp_func = interp1d(
+        energy_vals,
+        sensitivity.values,
+        axis=1,
+        bounds_error=False,
+        fill_value=np.nan,
+    )
+
+    # Interpolate to energy
+    interpolated = interp_func(energy)
 
-    return sensitivity
+    return interpolated