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

imap_processing/ultra/l1c/ultra_l1c.py

@@ -6,7 +6,7 @@ from imap_processing.ultra.l1c.histogram import calculate_histogram
 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 (
@@ -44,7 +41,6 @@ def ultra_l1c(data_dict: dict, data_version: str) -> list[xr.Dataset]:
             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",
-            data_version,
         )
         # TODO: add calculate_helio_pset here
         output_datasets = [spacecraft_pset]

imap_processing/ultra/l1c/ultra_l1c_pset_bins.py

@@ -3,14 +3,14 @@
 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
 
@@ -70,7 +70,7 @@ def get_spacecraft_histogram(
         Array of energy bin edges.
     nside : int, optional
         The nside parameter of the Healpix tessellation.
-        Default is 32.
+        Default is 128.
     nested : bool, optional
         Whether the Healpix tessellation is nested. Default is False.
 
@@ -174,54 +174,57 @@ def get_spacecraft_exposure_times(constant_exposure: pandas.DataFrame) -> NDArra
 
 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, with angles in degrees.
+    # Get energy midpoints.
     _, energy_midpoints, _ = build_energy_bins()
-    az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
-        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, az_grid, el_grid), axis=-1)
-    cartesian_coords = spherical_to_cartesian(spherical_coords)
-    cartesian = cartesian_coords.reshape(-1, 3, order="F").T
+    # 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
 
-    # Spacecraft velocity in the pointing (DPS) frame wrt heliosphere.
+    # 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.
@@ -234,41 +237,35 @@ 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
         )
-        # 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
+        # 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]
 
-        # 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)
+        # Convert azimuth/elevation directions to HEALPix pixel indices.
+        hpix_idx = hp.ang2pix(nside, az, el, nest=nested, lonlat=True)
 
-        # 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")
+        # Accumulate exposure values into HEALPix pixels for this energy bin.
+        helio_exposure[:, i] = np.bincount(
+            hpix_idx, weights=exposure_flat, minlength=npix
+        )
 
-    return exposure_3d
+    return helio_exposure
 
 
 def get_spacecraft_sensitivity(
     efficiencies: pandas.DataFrame,
     geometric_function: pandas.DataFrame,
-) -> pandas.DataFrame:
+) -> tuple[pandas.DataFrame, NDArray, NDArray, NDArray]:
     """
-    Compute sensitivity.
+    Compute sensitivity as efficiency * geometric factor.
 
     Parameters
     ----------
@@ -281,19 +278,69 @@ def get_spacecraft_sensitivity(
     -------
     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 = efficiencies.columns.difference(
-        ["Right Ascension (deg)", "Declination (deg)"]
-    )
+    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
     )
 
-    # Add "Right Ascension (deg)" and "Declination (deg)" to the result
-    sensitivity.insert(
-        0, "Right Ascension (deg)", efficiencies["Right Ascension (deg)"]
+    right_ascension = efficiencies["Right Ascension (deg)"]
+    declination = efficiencies["Declination (deg)"]
+
+    energy_vals = np.array([float(col.replace("keV", "")) for col in energy_columns])
+
+    return sensitivity, energy_vals, right_ascension, declination
+
+
+def grid_sensitivity(
+    efficiencies: pandas.DataFrame,
+    geometric_function: pandas.DataFrame,
+    energy: float,
+) -> NDArray:
+    """
+    Grid the sensitivity.
+
+    Parameters
+    ----------
+    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
+    -------
+    interpolated_sensitivity : np.ndarray
+        Sensitivity with dimensions (HEALPIX pixel_number, 1).
+    """
+    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,
     )
-    sensitivity.insert(1, "Declination (deg)", efficiencies["Declination (deg)"])
 
-    return sensitivity
+    # Interpolate to energy
+    interpolated = interp_func(energy)
+
+    return interpolated

imap_processing/ultra/l2/ultra_l2.py

@@ -0,0 +1,299 @@
+"""Calculate ULTRA Level 2 (L2) ENA Map Product."""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import xarray as xr
+
+from imap_processing.ena_maps import ena_maps
+from imap_processing.ena_maps.utils.coordinates import CoordNames
+
+logger = logging.getLogger(__name__)
+logger.info("Importing ultra_l2 module")
+
+# Default properties for the Ultra L2 map
+DEFAULT_ULTRA_L2_MAP_STRUCTURE: ena_maps.RectangularSkyMap | ena_maps.HealpixSkyMap = (
+    ena_maps.AbstractSkyMap.from_dict(
+        {
+            "sky_tiling_type": "HEALPIX",
+            "spice_reference_frame": "ECLIPJ2000",
+            "values_to_push_project": [
+                "counts",
+                "exposure_factor",
+                "sensitivity",
+                "background_rates",
+            ],
+            "nside": 32,
+            "nested": False,
+        }
+    )
+)
+
+# Set some default Healpix parameters - these must be defined, even if also
+# present in the DEFAULT_ULTRA_L2_MAP_STRUCTURE, because we always make a Healpix map
+# regardless of the output map type
+DEFAULT_L2_HEALPIX_NSIDE = 32
+DEFAULT_L2_HEALPIX_NESTED = False
+
+
+# These variables must always be present in each L1C dataset
+REQUIRED_L1C_VARIABLES = [
+    "counts",
+    "exposure_factor",
+    "sensitivity",
+    "background_rates",
+]
+
+# These variables are projected to the map as the mean of pointing set pixels value,
+# weighted by that pointing set pixel's exposure and solid angle
+VARIABLES_TO_WEIGHT_BY_POINTING_SET_EXPOSURE_TIMES_SOLID_ANGLE = [
+    "sensitivity",
+    "background_rates",
+    "observation_time",
+]
+
+# These variables are dropped after they are used to calculate flux and flux uncertainty
+# They will not be present in the final map
+VARIABLES_TO_DROP_AFTER_FLUX_CALCULATION = [
+    "counts",
+    "background_rates",
+    "pointing_set_exposure_times_solid_angle",
+    "num_pointing_set_pixel_members",
+    "corrected_count_rate",
+]
+
+
+def generate_ultra_healpix_skymap(
+    ultra_l1c_psets: list[str | xr.Dataset],
+    output_map_structure: (
+        ena_maps.RectangularSkyMap | ena_maps.HealpixSkyMap
+    ) = DEFAULT_ULTRA_L2_MAP_STRUCTURE,
+) -> ena_maps.HealpixSkyMap:
+    """
+    Generate a Healpix skymap from ULTRA L1C pointing sets.
+
+    This function combines IMAP Ultra L1C pointing sets into a single L2 HealpixSkyMap.
+    It handles the projection of values from pointing sets to the map, applies necessary
+    weighting and background subtraction, and calculates flux and flux uncertainty.
+
+    Parameters
+    ----------
+    ultra_l1c_psets : list[str | xr.Dataset]
+        List of paths to ULTRA L1C pointing set files or xarray Datasets containing
+        pointing set data.
+    output_map_structure : ena_maps.RectangularSkyMap | ena_maps.HealpixSkyMap, optional
+        Empty SkyMap structure providing the properties of the map to be generated.
+        Defaults to DEFAULT_ULTRA_L2_MAP_STRUCTURE defined in this module.
+
+    Returns
+    -------
+    ena_maps.HealpixSkyMap
+        HealpixSkyMap object containing the combined data from all pointing sets,
+        with calculated flux and flux uncertainty values.
+
+    Notes
+    -----
+    The structure of this function goes as follows:
+    1. Initialize the HealpixSkyMap object with the specified properties.
+    2. Iterate over the input pointing sets and read them into UltraPointingSet objects.
+    3. For each pointing set, weight certain variables by exposure and solid angle of
+    the pointing set pixels.
+    4. Project the pointing set values to the map using the push method.
+    5. Perform subsequent processing for weighted quantities at the SkyMap level
+    (e.g., divide weighted quantities by their summed weights to
+    get their weighted mean)
+    6. Calculate corrected count rate with background subtraction applied.
+    7. Calculate flux and flux uncertainty.
+    8. Drop unnecessary variables from the map.
+    """
+    if output_map_structure.tiling_type is ena_maps.SkyTilingType.HEALPIX:
+        map_nside, map_nested = (
+            output_map_structure.nside,
+            output_map_structure.nested,
+        )
+    else:
+        map_nside, map_nested = (DEFAULT_L2_HEALPIX_NSIDE, DEFAULT_L2_HEALPIX_NESTED)
+
+    # Initialize the HealpixSkyMap object
+    skymap = ena_maps.HealpixSkyMap(
+        nside=map_nside,
+        nested=map_nested,
+        spice_frame=output_map_structure.spice_reference_frame,
+    )
+
+    # Add additional data variables to the map
+    output_map_structure.values_to_push_project.extend(
+        [
+            "observation_time",
+            "pointing_set_exposure_times_solid_angle",
+            "num_pointing_set_pixel_members",
+        ]
+    )
+
+    # Get full list of variables to push to the map: all requested variables plus
+    # any which are required for L2 processing
+    value_keys_to_push_project = list(
+        set(output_map_structure.values_to_push_project + REQUIRED_L1C_VARIABLES)
+    )
+
+    for ultra_l1c_pset in ultra_l1c_psets:
+        pointing_set = ena_maps.UltraPointingSet(ultra_l1c_pset)
+        logger.info(
+            f"Projecting a PointingSet with {pointing_set.num_points} pixels "
+            f"at epoch:{pointing_set.epoch}\n"
+            f"These values will be projected: {value_keys_to_push_project}"
+        )
+
+        pointing_set.data["num_pointing_set_pixel_members"] = xr.DataArray(
+            np.ones(pointing_set.num_points, dtype=int),
+            dims=(CoordNames.HEALPIX_INDEX.value),
+        )
+        pointing_set.data["observation_time"] = xr.DataArray(
+            np.full(pointing_set.num_points, pointing_set.epoch),
+            dims=(CoordNames.HEALPIX_INDEX.value),
+        )
+        # Add solid_angle * exposure of pointing set as data_var
+        # so this quantity is projected to map pixels for use in weighted averaging
+        pointing_set.data["pointing_set_exposure_times_solid_angle"] = (
+            pointing_set.data["exposure_factor"] * pointing_set.solid_angle
+        )
+
+        # Initial processing for weighted quantities at PSET level
+        # Weight the values by exposure and solid angle
+        pointing_set.data[
+            VARIABLES_TO_WEIGHT_BY_POINTING_SET_EXPOSURE_TIMES_SOLID_ANGLE
+        ] *= pointing_set.data["pointing_set_exposure_times_solid_angle"]
+
+        skymap.project_pset_values_to_map(
+            pointing_set=pointing_set,
+            value_keys=value_keys_to_push_project,
+            index_match_method=ena_maps.IndexMatchMethod.PUSH,
+        )
+
+    # Subsequent processing for weighted quantities at SkyMap level
+    skymap.data_1d[VARIABLES_TO_WEIGHT_BY_POINTING_SET_EXPOSURE_TIMES_SOLID_ANGLE] /= (
+        skymap.data_1d["pointing_set_exposure_times_solid_angle"]
+    )
+
+    # TODO: Ask Ultra team about this - I think this is a decent
+    # but imperfect approximation for meaning the exposure:
+    # (dividing by the 1/(number of PSETs)) to fix the exposure being mean-ed over
+    # all pixels in all PSETs which feed into a map superpixel,
+    # rather than being mean-ed over pixels in a PSET and summed over PSETs
+    skymap.data_1d["exposure_factor"] /= skymap.data_1d[
+        "num_pointing_set_pixel_members"
+    ] / len(ultra_l1c_psets)
+
+    # TODO: Ask Ultra team about background rates - I think they should increase when
+    # binned to larger pixels, as I've done here, but that was never explicitly stated
+    skymap.data_1d["background_rates"] *= skymap.solid_angle / pointing_set.solid_angle
+
+    # Get the energy bin widths from a PointingSet (they will all be the same)
+    delta_energy = pointing_set.data["energy_bin_delta"]
+
+    # Core calculations of flux and flux uncertainty for L2
+    # Exposure time may contain 0s, producing NaNs in the corrected count rate and flux.
+    # These NaNs are not incorrect, so we temporarily ignore numpy div by 0 warnings.
+    with np.errstate(divide="ignore"):
+        # Get corrected count rate with background subtraction applied
+        skymap.data_1d["corrected_count_rate"] = (
+            skymap.data_1d["counts"].astype(float) / skymap.data_1d["exposure_factor"]
+        ) - skymap.data_1d["background_rates"]
+
+        # Calculate flux = corrected_counts / (sensitivity * solid_angle * delta_energy)
+        skymap.data_1d["flux"] = skymap.data_1d["corrected_count_rate"] / (
+            skymap.data_1d["sensitivity"] * skymap.solid_angle * delta_energy
+        )
+
+        skymap.data_1d["flux_uncertainty"] = (
+            skymap.data_1d["counts"].astype(float) ** 0.5
+        ) / (
+            skymap.data_1d["exposure_factor"]
+            * skymap.data_1d["sensitivity"]
+            * skymap.solid_angle
+            * delta_energy
+        )
+
+    # Drop the variables that are no longer needed
+    skymap.data_1d = skymap.data_1d.drop_vars(
+        VARIABLES_TO_DROP_AFTER_FLUX_CALCULATION,
+    )
+
+    return skymap
+
+
+def ultra_l2(
+    data_dict: dict[str, xr.Dataset | str],
+    data_version: str,
+    output_map_structure: (
+        ena_maps.RectangularSkyMap | ena_maps.HealpixSkyMap
+    ) = DEFAULT_ULTRA_L2_MAP_STRUCTURE,
+) -> list[xr.Dataset]:
+    """
+    Generate and format Ultra L2 ENA Map Product from L1C Products.
+
+    Parameters
+    ----------
+    data_dict : dict[str, xr.Dataset]
+        Dict mapping l1c product identifiers to paths/Datasets containing l1c psets.
+    data_version : str
+        Version of the data product being created.
+    output_map_structure : ena_maps.RectangularSkyMap | ena_maps.HealpixSkyMap, optional
+        Empty SkyMap structure providing the properties of the map to be generated.
+        Defaults to DEFAULT_ULTRA_L2_MAP_STRUCTURE defined in this module.
+
+    Returns
+    -------
+    list[xarray.Dataset,]
+        L2 output dataset containing map of the counts on the sky.
+        Wrapped in a list for consistency with other product levels.
+
+    Raises
+    ------
+    NotImplementedError
+        If asked to project to a rectangular map.
+        # TODO: This is coming shortly
+    """
+    l1c_products = data_dict.values()
+    num_l1c_products = len(l1c_products)
+    logger.info(f"Running ultra_l2 processing on {num_l1c_products} L1C products")
+
+    # Regardless of the output sky tiling type, we will directly
+    # project the PSET values into a healpix map. However, if we are outputting
+    # a Healpix map, we can go directly to map with desired nside, nested params
+    healpix_skymap = generate_ultra_healpix_skymap(
+        ultra_l1c_psets=list(l1c_products),
+        output_map_structure=output_map_structure,
+    )
+
+    # Output formatting for HEALPIX tiling
+    if output_map_structure.tiling_type is ena_maps.SkyTilingType.HEALPIX:
+        map_dataset = healpix_skymap.to_dataset()
+        # Add attributes related to the map
+        map_attrs = {
+            "HEALPix_nside": output_map_structure.nside,
+            "HEALPix_nest": output_map_structure.nested,
+            "Data_version": data_version,
+        }
+
+    # TODO: Implement conversion to Rectangular map
+    elif output_map_structure.tiling_type is ena_maps.SkyTilingType.RECTANGULAR:
+        map_attrs = {
+            "Spacing_degrees": output_map_structure.spacing_deg,
+            "Data_version": data_version,
+        }
+        raise NotImplementedError
+
+    # Always add the following attributes to the map
+    map_attrs.update(
+        {
+            "Sky_tiling_type": output_map_structure.tiling_type.value,
+            "Spice_reference_frame": output_map_structure.spice_reference_frame,
+        }
+    )
+
+    # Add the defined attributes to the map's global attrs
+    map_dataset.attrs.update(map_attrs)
+    return [map_dataset]