imap-processing 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

imap_processing/ena_maps/utils/corrections.py

@@ -1,10 +1,40 @@
 """L2 corrections common to multiple IMAP ENA instruments."""
 
+import logging
 from pathlib import Path
+from typing import TypeVar
 
 import numpy as np
 import pandas as pd
+import xarray as xr
 from numpy.polynomial import Polynomial
+from scipy.constants import electron_volt, erg, proton_mass
+
+from imap_processing.ena_maps.ena_maps import (
+    LoHiBasePointingSet,
+)
+from imap_processing.ena_maps.utils.coordinates import CoordNames
+from imap_processing.spice import geometry
+from imap_processing.spice.time import ttj2000ns_to_et
+
+logger = logging.getLogger(__name__)
+
+# Tell ruff to ignore ambiguous Greek letters in formulas in this file
+# ruff: noqa: RUF003
+
+# Create a TypeVar to represent the specific class being passed in
+# Bound to LoHiBasePointingSet, meaning it must be LoHiBasePointingSet
+# or a subclass of it
+LoHiBasePsetSubclass = TypeVar("LoHiBasePsetSubclass", bound=LoHiBasePointingSet)
+
+# Physical constants for Compton-Getting correction
+# Units: electron_volt = [J / eV]
+#        erg = [J / erg]
+# To get [erg / eV], => electron_volt [J / eV] / erg [J / erg] = erg_per_ev [erg / eV]
+ERG_PER_EV = electron_volt / erg  # erg per eV - unit conversion factor
+# Units: proton_mass = [kg]
+# Here, we convert proton_mass to grams
+PROTON_MASS_GRAMS = proton_mass * 1e3  # proton mass in grams
 
 
 class PowerLawFluxCorrector:
@@ -289,3 +319,423 @@ class PowerLawFluxCorrector:
             )
 
         return corrected_flux, corrected_flux_stat_unc
+
+
+def _add_spacecraft_velocity_to_pset(
+    pset: LoHiBasePsetSubclass,
+) -> LoHiBasePsetSubclass:
+    """
+    Calculate and add spacecraft velocity data to pointing set.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object to be updated.
+
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with spacecraft velocity data added.
+
+    Notes
+    -----
+    Adds the following DataArrays to pset.data:
+    - "sc_velocity": Spacecraft velocity vector (km/s) with dims ["x_y_z"]
+    - "sc_direction_vector": Spacecraft velocity unit vector with dims ["x_y_z"]
+    """
+    # Compute ephemeris time (J2000 seconds) of PSET midpoint time
+    # TODO: Use the Pointing midpoint time. Epoch should be start time
+    #     but use it until we can make Lo and Hi PSETs have a consistent
+    #     variable to hold the midpoint time.
+    et = ttj2000ns_to_et(pset.data["epoch"].values[0])
+    # Get spacecraft state in HAE frame
+    sc_state = geometry.imap_state(et, ref_frame=geometry.SpiceFrame.IMAP_HAE)
+    sc_velocity_vector = sc_state[3:6]
+
+    # Store spacecraft velocity as DataArray
+    pset.data["sc_velocity"] = xr.DataArray(
+        sc_velocity_vector, dims=[CoordNames.CARTESIAN_VECTOR.value]
+    )
+
+    # Calculate spacecraft speed and direction
+    sc_velocity_km_per_sec = np.linalg.norm(
+        pset.data["sc_velocity"], axis=-1, keepdims=True
+    )
+    pset.data["sc_direction_vector"] = pset.data["sc_velocity"] / sc_velocity_km_per_sec
+
+    return pset
+
+
+def _add_cartesian_look_direction(pset: LoHiBasePsetSubclass) -> LoHiBasePsetSubclass:
+    """
+    Calculate and add look direction vectors to pointing set.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object to be updated.
+
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with look direction vectors added.
+
+    Notes
+    -----
+    Adds the following DataArray to pset.data:
+    - "look_direction": Cartesian unit vectors with dims [...spatial_dims, "x_y_z"]
+    """
+    longitudes = pset.data["hae_longitude"]
+    latitudes = pset.data["hae_latitude"]
+
+    # Stack spherical coordinates (r=1 for unit vectors, azimuth, elevation)
+    spherical_coords = np.stack(
+        [
+            np.ones_like(longitudes),  # r = 1 for unit vectors
+            longitudes,  # azimuth = longitude
+            latitudes,  # elevation = latitude
+        ],
+        axis=-1,
+    )
+
+    # Convert to Cartesian coordinates and store as DataArray
+    pset.data["look_direction"] = xr.DataArray(
+        geometry.spherical_to_cartesian(spherical_coords),
+        dims=[*longitudes.dims, CoordNames.CARTESIAN_VECTOR.value],
+    )
+
+    return pset
+
+
+def _calculate_compton_getting_transform(
+    pset: LoHiBasePsetSubclass,
+    energy_hf: xr.DataArray,
+) -> LoHiBasePsetSubclass:
+    """
+    Apply Compton-Getting transformation to compute ENA source directions.
+
+    This implements the Compton-Getting velocity transformation to correct
+    for the motion of the spacecraft through the heliosphere. The transformation
+    accounts for the Doppler shift of ENA energies and the aberration of
+    arrival directions.
+
+    All calculations are performed using xarray DataArrays to preserve
+    dimension information throughout the computation.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object with sc_velocity, sc_direction_vector, and
+        look_direction already added.
+    energy_hf : xr.DataArray
+        ENA energies in the heliosphere frame in eV.
+
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Pointing set object with Compton-Getting related variables added and
+        updated az_el_points.
+
+    Notes
+    -----
+    The algorithm is based on the "Appendix A. The IMAP-Lo Mapping Algorithms"
+    document.
+    Adds the following DataArrays to pset.data:
+    - "energy_sc": ENA energies in spacecraft frame (eV)
+    - "energy_hf": ENA energies in the heliosphere frame (eV)
+    - "ram_mask": Mask indicating whether ENA source direction is from the ram
+      direction.
+    Updates the following DataArrays in pset.data:
+    - "hae_longitude": ENA source longitudes in heliosphere frame (degrees)
+    - "hae_latitude": ENA source latitudes in heliosphere frame (degrees)
+    """
+    # Store heliosphere frame energies
+    pset.data["energy_hf"] = energy_hf
+
+    # Calculate spacecraft speed
+    sc_velocity_km_per_sec = np.linalg.norm(
+        pset.data["sc_velocity"], axis=-1, keepdims=True
+    )
+
+    # Calculate dot product between look directions and spacecraft direction vector
+    # Use Einstein summation for efficient vectorized dot product
+    dot_product = xr.DataArray(
+        np.einsum(
+            "...i,...i->...",
+            pset.data["look_direction"],
+            pset.data["sc_direction_vector"],
+        ),
+        dims=pset.data["look_direction"].dims[:-1],
+    )
+
+    # Calculate the kinetic energy of a hydrogen ENA traveling at spacecraft velocity
+    # E_u = (1/2) * m * U_sc^2 (convert km/s to cm/s with 1.0e5 factor)
+    energy_u = (
+        0.5 * PROTON_MASS_GRAMS * (sc_velocity_km_per_sec * 1e5) ** 2 / ERG_PER_EV
+    )
+
+    # Note: Tim thinks that this approach seems backwards. Here, we are assuming
+    #     that ENAs are observed in the heliosphere frame at the ESA energy levels.
+    #     We then calculate the velocity that said ENAs would have in the spacecraft
+    #     frame as well as the CG corrected energy level in the spacecraft frame.
+    #     We then use this velocity to calculate and the velocity of the spacecraft
+    #     to do the vector math which determines the ENA source direction in the
+    #     heliosphere frame.
+    #     The ENAs are in fact observed in the spacecraft frame at a known energy
+    #     level in the spacecraft frame. Why don't we use that energy level to
+    #     calculate the source direction in the spacecraft frame and then do the
+    #     vector math to find the source direction in the heliosphere frame? We
+    #     would also need to calculate the CG corrected ENA energy in the heliosphere
+    #     frame and keep track of that when binning.
+
+    # Calculate y values for each energy level (Equation 61)
+    # y_k = sqrt(E^h_k / E^u)
+    y = np.sqrt(pset.data["energy_hf"] / energy_u)
+
+    # Velocity magnitude factor calculation (Equation 62)
+    # x_k = (êₛ · û_sc) + sqrt(y² + (êₛ · û_sc)² - 1)
+    x = dot_product + np.sqrt(y**2 + dot_product**2 - 1)
+    # Get the dimensions in the right order so that spatial is last
+    x = x.transpose(dot_product.dims[0], y.dims[0], dot_product.dims[1])
+
+    # Calculate ENA speed in the spacecraft frame
+    # |v⃗_sc| = x_k * U_sc
+    velocity_sc = x * sc_velocity_km_per_sec
+
+    # Calculate the kinetic energy in the spacecraft frame
+    # E_sc = (1/2) * M_p * v_sc² (convert km/s to cm/s with 1.0e5 factor)
+    pset.data["energy_sc"] = (
+        0.5 * PROTON_MASS_GRAMS * (velocity_sc * 1e5) ** 2 / ERG_PER_EV
+    )
+
+    # Calculate the velocity vector in the spacecraft frame
+    # v⃗_sc = |v_sc| * êₛ (velocity direction follows look direction)
+    velocity_vector_sc = velocity_sc * pset.data["look_direction"]
+
+    # Calculate the ENA velocity vector in the heliosphere frame
+    # v⃗_helio = v⃗_sc - U⃗_sc (simple velocity addition)
+    velocity_vector_helio = velocity_vector_sc - pset.data["sc_velocity"]
+
+    # Convert to spherical coordinates to get ENA source directions
+    ena_source_direction_helio = geometry.cartesian_to_spherical(
+        velocity_vector_helio.data
+    )
+
+    # Update the PSET hae_longitude and hae_latitude variables with the new
+    # energy-dependent values.
+    pset.data["hae_longitude"] = (
+        pset.data["energy_sc"].dims,
+        ena_source_direction_helio[..., 1],
+    )
+    pset.data["hae_latitude"] = (
+        pset.data["energy_sc"].dims,
+        ena_source_direction_helio[..., 2],
+    )
+
+    # For ram/anti-ram filtering we can use the sign of the scalar projection
+    # of the ENA source direction onto the spacecraft velocity vector.
+    # ram_mask = (v⃗_helio · û_sc) >= 0
+    ram_mask = (
+        np.einsum(
+            "...i,...i->...", velocity_vector_helio, pset.data["sc_direction_vector"]
+        )
+        >= 0
+    )
+    pset.data["ram_mask"] = xr.DataArray(
+        ram_mask,
+        dims=velocity_vector_helio.dims[:-1],
+    )
+
+    return pset
+
+
+def apply_compton_getting_correction(
+    pset: LoHiBasePsetSubclass,
+    energy_hf: xr.DataArray,
+) -> LoHiBasePsetSubclass:
+    """
+    Apply Compton-Getting correction to a pointing set and update coordinates.
+
+    This function performs the Compton-Getting velocity transformation to correct
+    ENA observations for the motion of the spacecraft through the heliosphere.
+    The corrected coordinates represent the true source directions of the ENAs
+    in the heliosphere frame.
+
+    The pointing set is modified in-place: new variables are added to the dataset
+    for the corrected coordinates and energies, and the az_el_points attribute
+    is updated to use the corrected coordinates for binning.
+
+    All calculations are performed using xarray DataArrays to preserve dimension
+    information throughout the computation.
+
+    Parameters
+    ----------
+    pset : LoHiBasePointingSet
+        Pointing set object containing HAE longitude/latitude coordinates.
+    energy_hf : xr.DataArray
+        ENA energies in the heliosphere frame in eV. Must be 1D with an
+        energy dimension.
+
+    Returns
+    -------
+    pset : LoHiBasePointingSet
+        Updated pointing set object with Compton-Getting related variables added.
+
+    Notes
+    -----
+    This function adds the following variables to the pointing set dataset:
+    - "sc_velocity": Spacecraft velocity vector (km/s)
+    - "sc_direction_vector": Spacecraft velocity unit vector
+    - "look_direction": Cartesian unit vectors of observation directions
+    - "energy_hf": ENA energies in heliosphere frame (eV)
+    - "energy_sc": ENA energies in spacecraft frame (eV)
+    This function modifies the following variables in the pointing set dataset:
+    - "hae_longitude": ENA source longitudes in heliosphere frame (degrees)
+    - "hae_latitude": ENA source latitudes in heliosphere frame (degrees)
+
+    The az_el_points attribute is updated to use the corrected coordinates,
+    which will be used for subsequent binning operations.
+    """
+    # Step 1: Add spacecraft velocity and direction to pset
+    pset = _add_spacecraft_velocity_to_pset(pset)
+
+    # Step 2: Calculate and add look direction vectors to pset
+    pset = _add_cartesian_look_direction(pset)
+
+    # Step 3: Apply Compton-Getting transformation
+    pset = _calculate_compton_getting_transform(pset, energy_hf)
+
+    # Step 4: Update az_el_points to use the corrected coordinates
+    pset.update_az_el_points()
+
+    return pset
+
+
+def interpolate_map_flux_to_helio_frame(
+    map_ds: xr.Dataset,
+    esa_energies_ev: xr.DataArray,
+    helio_energies_ev: xr.DataArray,
+) -> xr.Dataset:
+    """
+    Interpolate flux from spacecraft frame to heliocentric frame energies.
+
+    This implements the Compton-Getting interpolation step that transforms
+    flux measurements from the spacecraft frame to the heliocentric frame.
+    The algorithm follows these steps:
+    1. For each spatial pixel and energy step, get the spacecraft energy
+    2. Find bounding ESA energy channels for interpolation
+    3. Perform power-law interpolation between bounding channels to spacecraft energy
+    4. Apply energy scaling transformation to heliocentric frame
+
+    Parameters
+    ----------
+    map_ds : xarray.Dataset
+        Map dataset with `energy_sc` data variable containing the spacecraft
+        frame energies for each spatial pixel and ESA energy step.
+    esa_energies_ev : xarray.DataArray
+        The ESA nominal central energies (in eV).
+    helio_energies_ev : xarray.DataArray
+        The heliocentric frame energies to interpolate to (in eV).
+        In practice, these are the same as esa_energies_ev.
+
+    Returns
+    -------
+    map_ds : xarray.Dataset
+        Updated map dataset with interpolated heliocentric frame fluxes.
+    """
+    logger.info("Performing Compton-Getting interpolation to heliocentric frame")
+
+    # Work with xarray DataArrays to handle arbitrary spatial dimensions
+    energy_sc = map_ds["energy_sc"]
+    intensity = map_ds["ena_intensity"]
+    stat_unc = map_ds["ena_intensity_stat_uncert"]
+    sys_err = map_ds["ena_intensity_sys_err"]
+
+    # Step 1: Find bounding ESA energy indices for each position
+    # Use np.searchsorted on flattened array, then reshape back
+    esa_energy_vals = esa_energies_ev.values
+    energy_sc_flat = energy_sc.values.ravel()
+
+    # Find right bound index for each element (vectorized)
+    right_idx_flat = np.searchsorted(esa_energy_vals, energy_sc_flat, side="right")
+    right_idx_flat = np.clip(right_idx_flat, 1, len(esa_energy_vals) - 1)
+    left_idx_flat = right_idx_flat - 1
+
+    # Reshape indices back to match energy_sc shape
+    right_idx = right_idx_flat.reshape(energy_sc.shape)
+    left_idx = left_idx_flat.reshape(energy_sc.shape)
+
+    # Create DataArrays for indices with same dims as energy_sc
+    # Note: we need to avoid coordinate name conflicts when using isel()
+    # The energy dimension should be present in dims but not as a coordinate
+    # since we're using these as indices into the energy dimension
+    # Create coordinates dict without the energy coordinate
+    coords_without_energy = {k: v for k, v in energy_sc.coords.items() if k != "energy"}
+
+    right_idx_da = xr.DataArray(
+        right_idx, dims=energy_sc.dims, coords=coords_without_energy
+    )
+    left_idx_da = xr.DataArray(
+        left_idx, dims=energy_sc.dims, coords=coords_without_energy
+    )
+
+    # Step 2: Extract flux values at bounding energy channels
+    # Use xarray's advanced indexing to get fluxes at left and right indices
+    flux_left = intensity.isel({"energy": left_idx_da})
+    flux_right = intensity.isel({"energy": right_idx_da})
+    stat_unc_left = stat_unc.isel({"energy": left_idx_da})
+    stat_unc_right = stat_unc.isel({"energy": right_idx_da})
+    sys_err_left = sys_err.isel({"energy": left_idx_da})
+
+    # Get energy values at boundaries - select from esa_energies_ev using indices
+    energy_left = esa_energies_ev.isel({"energy": left_idx_da})
+    energy_right = esa_energies_ev.isel({"energy": right_idx_da})
+
+    # Step 3: Perform power-law interpolation to spacecraft energy
+    # slope = log(f_right/f_left) / log(e_right/e_left)
+    # flux_sc = f_left * (energy_sc / e_left)^slope
+    with np.errstate(divide="ignore", invalid="ignore"):
+        # Calculate slope for power-law interpolation
+        slope = np.log(flux_right / flux_left) / np.log(energy_right / energy_left)
+
+        # Interpolate flux using power-law
+        flux_sc = flux_left * ((energy_sc / energy_left) ** slope)
+
+        # Interpolation factor for uncertainty propagation (Equations 75 & 76)
+        unc_factor = np.log(energy_sc / energy_left) / np.log(
+            energy_right / energy_left
+        )
+
+        # Statistical uncertainty propagation (Equation 75):
+        # δJ = J * sqrt((δJ_left/J_left)^2 * (1 + unc_factor^2) + (δJ_right/J_right)^2)
+        stat_unc_sc = flux_sc * np.sqrt(
+            (stat_unc_left / flux_left) ** 2 * (1.0 + unc_factor**2)
+            + (stat_unc_right / flux_right) ** 2
+        )
+
+        # Systematic uncertainty propagation (Equation 76):
+        # σJ^g = σJ^src_kref * (⟨E^s_kref⟩ / E^ESA_kref)^γ_kref * (E^h / ⟨E^s_kref⟩)
+        # Systematic error scales proportionally with flux during power-law
+        # interpolation
+        sys_err_sc = sys_err_left * ((energy_sc / energy_left) ** slope)
+
+    # Step 4: Energy scaling transformation (Liouville theorem)
+    # flux_helio = flux_sc * (helio_energy / energy_sc)
+    # Use xarray broadcasting - helio_energies_ev will broadcast along esa_energy_step
+    with np.errstate(divide="ignore", invalid="ignore"):
+        energy_ratio = helio_energies_ev / energy_sc
+        flux_helio = flux_sc * energy_ratio
+        stat_unc_helio = stat_unc_sc * energy_ratio
+        sys_err_helio = sys_err_sc * energy_ratio
+
+    # Set any location where the value is not finite to NaN (converts +/-inf to NaN)
+    flux_helio = flux_helio.where(np.isfinite(flux_helio), np.nan)
+    stat_unc_helio = stat_unc_helio.where(np.isfinite(stat_unc_helio), np.nan)
+    sys_err_helio = sys_err_helio.where(np.isfinite(sys_err_helio), np.nan)
+
+    # Update the dataset with interpolated values
+    map_ds["ena_intensity"] = flux_helio
+    map_ds["ena_intensity_stat_uncert"] = stat_unc_helio
+    map_ds["ena_intensity_sys_err"] = sys_err_helio
+
+    return map_ds

imap_processing/ena_maps/utils/map_utils.py

@@ -10,6 +10,89 @@ from numpy.typing import NDArray
 logger = logging.getLogger(__name__)
 
 
+def vectorized_bincount(
+    indices: NDArray, weights: NDArray | None = None, minlength: int = 0
+) -> NDArray:
+    """
+    Vectorized version of np.bincount for multi-dimensional arrays.
+
+    This function applies np.bincount across multi-dimensional input arrays by
+    adding offsets to the indices and flattening, then reshaping the result.
+    This approach allows broadcasting between indices and weights.
+
+    Parameters
+    ----------
+    indices : NDArray
+        Array of non-negative integers to be binned. Can be multi-dimensional.
+        If multi-dimensional, bincount is applied independently along each
+        leading dimension.
+    weights : NDArray, optional
+        Array of weights that is broadcastable with indices. If provided, each
+        weight is accumulated into its corresponding bin. If None (default),
+        each index contributes a count of 1.
+    minlength : int, optional
+        Minimum number of bins in the output array. Applied to each independent
+        bincount operation. Default is 0.
+
+    Returns
+    -------
+    NDArray
+        Array of binned values with the same leading dimensions as the input
+        arrays, and a final dimension of size minlength (or the maximum index + 1,
+        whichever is larger).
+
+    See Also
+    --------
+    numpy.bincount : The underlying function being vectorized.
+
+    Examples
+    --------
+    >>> indices = np.array([[0, 1, 1], [2, 2, 3]])
+    >>> vectorized_bincount(indices, minlength=4)
+    array([[1., 2., 0., 0.],
+           [0., 0., 2., 1.]])
+    """
+    # Handle 1D case directly
+    if indices.ndim == 1 and (weights is None or weights.ndim == 1):
+        return np.bincount(indices, weights=weights, minlength=minlength)
+
+    # For multi-dimensional arrays, broadcast indices and weights
+    if weights is not None:
+        indices_bc, weights_bc = np.broadcast_arrays(indices, weights)
+        weights_flat = weights_bc.ravel()
+    else:
+        indices_bc = indices
+        weights_flat = None
+
+    # Get the shape for reshaping output
+    non_spatial_shape = indices_bc.shape[:-1]
+    n_binsets = np.prod(non_spatial_shape)
+
+    # Determine actual minlength if not specified
+    if minlength == 0:
+        minlength = int(np.max(indices_bc)) + 1
+
+    # We want to flatten the multi-dimensional bincount problem into a 1D problem.
+    # This can be done by offsetting the indices for each element of each additional
+    # dimension by an integer multiple of the number of bins. Doing so gives
+    # each element in the additional dimensions its own set of 1D bins: index 0
+    # uses bins [0, minlength), index 1 uses bins [minlength, 2*minlength), etc.
+    offsets = np.arange(n_binsets).reshape(*non_spatial_shape, 1) * minlength
+    indices_flat = (indices_bc + offsets).ravel()
+
+    # Single bincount call with flattened data
+    binned_flat = np.bincount(
+        indices_flat, weights=weights_flat, minlength=n_binsets * minlength
+    )
+
+    # Reshape to separate each sample's bins
+    binned_values = binned_flat.reshape(n_binsets, -1)[:, :minlength].reshape(
+        *non_spatial_shape, minlength
+    )
+
+    return binned_values
+
+
 def bin_single_array_at_indices(
     value_array: NDArray,
     projection_grid_shape: tuple[int, ...],
@@ -25,7 +108,7 @@ def bin_single_array_at_indices(
     Parameters
     ----------
     value_array : NDArray
-        Array of values to bin. The final axis be the one and only spatial axis.
+        Array of values to bin. The final axis is the one and only spatial axis.
         If other axes are present, they will be binned independently
         along the spatial axis.
     projection_grid_shape : tuple[int, ...]
@@ -34,71 +117,89 @@ def bin_single_array_at_indices(
         or just (number of bins,) if the grid is 1D.
     projection_indices : NDArray
         Ordered indices for projection grid, corresponding to indices in input grid.
-        1 dimensional. May be non-unique, depending on the projection method.
+        Can be 1-dimensional or multi-dimensional. If multi-dimensional, must be
+        broadcastable with value_array. May contain non-unique indices, depending
+        on the projection method.
     input_indices : NDArray
         Ordered indices for input grid, corresponding to indices in projection grid.
         1 dimensional. May be non-unique, depending on the projection method.
-        If None (default), an arange of the same length as the
-        final axis of value_array is used.
+        If None (default), an numpy.arange of the same length as the final axis of
+        value_array is used.
     input_valid_mask : NDArray, optional
         Boolean mask array for valid values in input grid.
         If None, all pixels are considered valid. Default is None.
+        Must be broadcastable with value_array and projection_indices.
 
     Returns
     -------
     NDArray
-        Binned values on the projection grid.
+        Binned values on the projection grid. The output shape depends on the
+        input shapes after broadcasting:
+        - If value_array is 1D: returns 1D array of shape (num_projection_indices,)
+        - If value_array is multi-dimensional: returns array with shape
+          (*value_array.shape[:-1], num_projection_indices), where the leading
+          dimensions match value_array's non-spatial dimensions and the final
+          dimension contains the binned values for each projection grid position.
+        - If projection_indices is multi-dimensional and broadcasts with value_array,
+          the output shape will be (broadcasted_shape[:-1], num_projection_indices).
 
     Raises
     ------
     ValueError
-        If the input and projection indices are not 1D arrays
-        with the same number of elements.
-    NotImplementedError
-        If the input value_array has dimensionality less than 1.
+        If input_indices is not a 1D array, or if the arrays cannot be
+        broadcast together.
     """
+    # Set and check input_indices
     if input_indices is None:
         input_indices = np.arange(value_array.shape[-1])
-    if input_valid_mask is None:
-        input_valid_mask = np.ones(value_array.shape[-1], dtype=bool)
-
-    # Both sets of indices must be 1D with the same number of elements
-    if input_indices.ndim != 1 or projection_indices.ndim != 1:
+    # input_indices must be 1D
+    if input_indices.ndim != 1:
         raise ValueError(
-            "Indices must be 1D arrays. "
+            "input_indices must be a 1D array. "
             "If using a rectangular grid, the indices must be unwrapped."
         )
-    if input_indices.size != projection_indices.size:
-        raise ValueError(
-            "The number of input and projection indices must be the same. \n"
-            f"Received {input_indices.size} input indices and {projection_indices.size}"
-            " projection indices."
+
+    # Verify projection_indices is broadcastable with value_array
+    try:
+        broadcasted_shape = np.broadcast_shapes(
+            projection_indices.shape, value_array.shape
         )
+    except ValueError as e:
+        raise ValueError(
+            f"projection_indices shape {projection_indices.shape} must be "
+            f"broadcastable with value_array shape {value_array.shape}"
+        ) from e
 
-    input_valid_mask = np.asarray(input_valid_mask, dtype=bool)
-    mask_idx = input_valid_mask[input_indices]
+    # Set and check input_valid_mask
+    if input_valid_mask is None:
+        input_valid_mask = np.ones(value_array.shape[-1], dtype=bool)
+    else:
+        input_valid_mask = np.asarray(input_valid_mask, dtype=bool)
+    # Verify input_valid_mask is broadcastable with value_array
+    try:
+        np.broadcast_shapes(input_valid_mask.shape, value_array.shape)
+    except ValueError as e:
+        raise ValueError(
+            f"input_valid_mask shape {input_valid_mask.shape} must be "
+            f"broadcastable with value_array shape {value_array.shape}"
+        ) from e
 
-    num_projection_indices = np.prod(projection_grid_shape)
+    # Broadcast input_valid_mask to match value_array shape if needed
+    input_valid_mask_bc = np.broadcast_to(input_valid_mask, broadcasted_shape)
+
+    # Select values at input_indices positions along the spatial axis
+    values = value_array[..., input_indices]
+
+    # Apply mask: set invalid values to 0
+    values_masked = np.where(input_valid_mask_bc, values, 0)
+
+    num_projection_indices = int(np.prod(projection_grid_shape))
+
+    # Use vectorized_bincount to handle arbitrary dimensions
+    binned_values = vectorized_bincount(
+        projection_indices, weights=values_masked, minlength=num_projection_indices
+    )
 
-    # Only valid values are summed into bins.
-    if value_array.ndim == 1:
-        values = value_array[input_indices]
-        binned_values = np.bincount(
-            projection_indices[mask_idx],
-            weights=values[mask_idx],
-            minlength=num_projection_indices,
-        )
-    elif value_array.ndim >= 2:
-        # Apply bincount to each row independently
-        binned_values = np.apply_along_axis(
-            lambda x: np.bincount(
-                projection_indices[mask_idx],
-                weights=x[..., input_indices][mask_idx],
-                minlength=num_projection_indices,
-            ),
-            axis=-1,
-            arr=value_array,
-        )
     return binned_values