Index Mapping Util for Rectangular grids

imap_processing/ena_maps/map_utils.py

@@ -0,0 +1,168 @@
+"""Utilities for generating ENA maps."""
+
+from __future__ import annotations
+
+import logging
+import typing
+
+import numpy as np
+from numpy.typing import NDArray
+
+from imap_processing.spice import geometry
+from imap_processing.ultra.utils import spatial_utils
+
+logger = logging.getLogger(__name__)
+
+
+# Ignore linting rule to allow for 6 unrelated args
+# Also need to manually specify allowed str literals for order parameter
+def match_rect_indices_frame_to_frame(  # noqa: PLR0913
+    input_frame: geometry.SpiceFrame,
+    projection_frame: geometry.SpiceFrame,
+    event_time: float,
+    input_frame_spacing_deg: float,
+    projection_frame_spacing_deg: float,
+    order: typing.Literal["C"] | typing.Literal["F"] = "F",
+) -> tuple[NDArray, NDArray, NDArray, NDArray, NDArray, NDArray, NDArray, NDArray]:
+    """
+    Match flattened rectangular grid indices between two reference frames.
+
+    Parameters
+    ----------
+    input_frame : geometry.SpiceFrame
+        The frame in which the input grid is defined.
+    projection_frame : geometry.SpiceFrame
+        The frame to which the input grid will be projected.
+    event_time : float
+        The time at which to project the grid.
+    input_frame_spacing_deg : float, optional
+        The spacing of the input frame grid in degrees,
+        by default DEFAULT_SPACING_DEG.
+    projection_frame_spacing_deg : float, optional
+        The spacing of the projection frame grid in degrees,
+        by default DEFAULT_SPACING_DEG.
+    order : str, optional
+        The order of unraveling/re-raveling the grid points,
+        by default "F" for column-major Fortran ordering.
+        See numpy.ravel documentation for more information.
+
+    Returns
+    -------
+    tuple[NDArray]
+        Tuple of the following arrays:
+        - Ordered 1D array of pixel indices covering the entire input grid.
+        Guaranteed to have one index for each point on the input grid,
+        (meaning it will cover the entire input grid).
+        Size is number of pixels on input grid
+        `= (az_grid_input.size * el_grid_input.size)`.
+        - Ordered 1D array of indices in the projection frame
+        corresponding to each index in the input grid.
+        Not guaranteed to cover the entire projection grid.
+        Size is the same as the input indices
+        = `(az_grid_input.size * el_grid_input.size)`.
+        - Azimuth values of the input grid points in the input frame, raveled.
+        Same size as the input indices.
+        - Elevation values of the input grid points in the input frame, raveled.
+        Same size as the input indices.
+        - Azimuth values of the input grid points in the projection frame, raveled.
+        Same size as the input indices.
+        - Elevation values of the input grid points in the projection frame, raveled.
+        Same size as the input indices.
+        - Azimuth indices of the input grid points in the projection frame az grid.
+        Same size as the input indices.
+        - Elevation indices of the input grid points in the projection frame el grid.
+        Same size as the input indices.
+    """
+    if input_frame_spacing_deg > projection_frame_spacing_deg:
+        logger.warning(
+            "Input frame has a larger spacing than the projection frame."
+            f"\nReceived: input = {input_frame_spacing_deg} degrees "
+            f"and {projection_frame_spacing_deg} degrees."
+        )
+    # Build the azimuth, elevation grid in the inertial frame
+    az_grid_in, el_grid_in = spatial_utils.build_az_el_grid(
+        spacing=input_frame_spacing_deg,
+        input_degrees=True,
+        output_degrees=False,
+        centered_azimuth=False,  # (0, 2pi rad = 0, 360 deg)
+        centered_elevation=True,  # (-pi/2, pi/2 rad = -90, 90 deg)
+    )[2:4]
+
+    # Unwrap the grid to a 1D array for same interface as tessellation
+    az_grid_input_raveled = az_grid_in.ravel(order=order)
+    el_grid_input_raveled = el_grid_in.ravel(order=order)
+
+    # Get the flattened indices of the grid points in the input frame (Fortran order)
+    # TODO: Discuss with Nick/Tim/Laura: Should this just be an arange?
+    flat_indices_in = np.ravel(
+        np.arange(az_grid_input_raveled.size).reshape(az_grid_in.shape),
+        order=order,
+    )
+
+    radii_in = geometry.spherical_to_cartesian(
+        np.stack(
+            (
+                np.ones_like(az_grid_input_raveled),
+                az_grid_input_raveled,
+                el_grid_input_raveled,
+            ),
+            axis=-1,
+        )
+    )
+
+    # Project the grid points from the input frame to the projection frame
+    # radii_proj are cartesian (x,y,z) radii vectors in the projection frame
+    # corresponding to the grid points in the input frame
+    radii_proj = geometry.frame_transform(
+        et=event_time,
+        position=radii_in,
+        from_frame=input_frame,
+        to_frame=projection_frame,
+    )
+
+    # Convert the (x,y,z) vectors to spherical coordinates in the projection frame
+    # Then extract the azimuth, elevation angles in the projection frame. Ignore radius.
+    input_grid_in_proj_spherical_coord = geometry.cartesian_to_spherical(
+        radii_proj, degrees=False
+    )
+
+    input_az_in_proj_az = input_grid_in_proj_spherical_coord[:, 1]
+    input_el_in_proj_el = input_grid_in_proj_spherical_coord[:, 2]
+
+    # Create bin edges for azimuth (0 to 2pi) and elevation (-pi/2 to pi/2)
+    proj_frame_az_bin_edges, proj_frame_el_bin_edges = spatial_utils.build_az_el_grid(
+        spacing=projection_frame_spacing_deg,
+        input_degrees=True,
+        output_degrees=False,
+        centered_azimuth=False,  # (0, 2pi rad = 0, 360 deg)
+        centered_elevation=True,  # (-pi/2, pi/2 rad = -90, 90 deg)
+    )[4:6]
+
+    # Use digitize to find indices (-1 since digitize returns 1-based indices)
+    input_az_in_proj_az_indices = (
+        np.digitize(input_az_in_proj_az, proj_frame_az_bin_edges) - 1
+    )
+    input_el_in_proj_el_indices = (
+        np.digitize(input_el_in_proj_el, proj_frame_el_bin_edges[::-1]) - 1
+    )
+
+    # NOTE: This method of matching indices 1:1 is rectangular grid-focused
+    # It will not necessarily work with a tessellation (e.g. Healpix)
+    flat_indices_proj = np.ravel_multi_index(
+        multi_index=(input_az_in_proj_az_indices, input_el_in_proj_el_indices),
+        dims=(
+            int(360 // projection_frame_spacing_deg),
+            int(180 // projection_frame_spacing_deg),
+        ),
+        order=order,
+    )
+    return (
+        flat_indices_in,
+        flat_indices_proj,
+        az_grid_input_raveled,
+        el_grid_input_raveled,
+        input_az_in_proj_az,
+        input_el_in_proj_el,
+        input_az_in_proj_az_indices,
+        input_el_in_proj_el_indices,
+    )

imap_processing/tests/ena_maps/test_map_utils.py

@@ -0,0 +1,103 @@
+"""Tests coverage for ENA Map Mapping Util functions."""
+
+from unittest import mock
+
+import numpy as np
+import pytest
+
+from imap_processing.ena_maps import map_utils
+from imap_processing.spice import geometry
+from imap_processing.ultra.utils import spatial_utils
+
+
+class TestENAMapMappingUtils:
+    @mock.patch("imap_processing.ena_maps.map_utils.geometry.frame_transform")
+    def test_match_rect_indices_frame_to_frame_same_frame(self, mock_frame_transform):
+        # Mock frame transform to return the input position vectors (no transform)
+        mock_frame_transform.side_effect = (
+            lambda et, position, from_frame, to_frame: position
+        )
+        et = -1
+        spacing_deg = 1
+        (
+            flat_indices_in,
+            flat_indices_proj,
+            az_grid_input_raveled,
+            el_grid_input_raveled,
+            input_az_in_proj_az,
+            input_el_in_proj_el,
+        ) = map_utils.match_rect_indices_frame_to_frame(
+            input_frame=geometry.SpiceFrame.IMAP_DPS,
+            projection_frame=geometry.SpiceFrame.IMAP_DPS,
+            event_time=et,
+            input_frame_spacing_deg=spacing_deg,
+            projection_frame_spacing_deg=spacing_deg,
+        )[:6]
+
+        # The two arrays should be the same if the input and projection
+        # frames are the same and the spacing is the same.
+        np.testing.assert_array_equal(flat_indices_in, flat_indices_proj)
+
+        # The input and projection azimuth and elevation grids should be the same
+        np.testing.assert_allclose(az_grid_input_raveled, input_az_in_proj_az)
+        np.testing.assert_allclose(el_grid_input_raveled, input_el_in_proj_el)
+
+        # The projection azimuth and elevation, when reshaped to a 2D grid,
+        # should be the same as the input grid created with build_az_el_grid.
+        np.testing.assert_allclose(
+            input_az_in_proj_az.reshape(
+                (180 // spacing_deg, 360 // spacing_deg), order="F"
+            ),
+            spatial_utils.build_az_el_grid(spacing=np.deg2rad(spacing_deg))[2],
+        )
+        np.testing.assert_allclose(
+            input_el_in_proj_el.reshape(
+                (180 // spacing_deg, 360 // spacing_deg), order="F"
+            ),
+            spatial_utils.build_az_el_grid(spacing=np.deg2rad(spacing_deg))[3],
+        )
+
+    @mock.patch("imap_processing.ena_maps.map_utils.geometry.frame_transform")
+    @pytest.mark.parametrize("input_spacing_deg", [1 / 4, 1 / 3, 1 / 2, 1])
+    def test_match_rect_indices_frame_to_frame_different_spacings(
+        self, mock_frame_transform, input_spacing_deg
+    ):
+        # Mock frame transform to return the input position vectors (no transform)
+        mock_frame_transform.side_effect = (
+            lambda et, position, from_frame, to_frame: position
+        )
+
+        et = -1
+        proj_spacing_deg = 1
+
+        (
+            flat_indices_in,
+            flat_indices_proj,
+        ) = map_utils.match_rect_indices_frame_to_frame(
+            input_frame=geometry.SpiceFrame.IMAP_DPS,
+            projection_frame=geometry.SpiceFrame.IMAP_DPS,
+            event_time=et,
+            input_frame_spacing_deg=input_spacing_deg,
+            projection_frame_spacing_deg=proj_spacing_deg,
+        )[:2]
+
+        # The input indices should still be the same length as the projection indices
+        assert len(flat_indices_in) == len(flat_indices_proj)
+
+        # However, the min and max of the input indices should go from:
+        # 0 to (180/spacing_deg of the input)*(360/spacing_deg of the input)
+        # and the min and max of the projection indices should go from:
+        # 0 to (180/spacing_deg of the projection)*(360/spacing_deg of the projection)
+        assert flat_indices_in.min() == 0
+        assert flat_indices_in.max() == (
+            (180 / (input_spacing_deg)) * (360 / (input_spacing_deg)) - 1
+        )
+        assert flat_indices_proj.min() == 0
+        assert (
+            flat_indices_proj.max()
+            == ((180 / proj_spacing_deg) * (360 / proj_spacing_deg)) - 1
+        )
+
+        # There should be (proj_spacing_deg/input_spacing_deg)**2
+        # instances of each projection index, and proj_spacing_deg = 1.
+        assert np.all(np.bincount(flat_indices_proj) == 1 / (input_spacing_deg**2))

imap_processing/tests/ultra/unit/test_spatial_utils.py

@@ -69,12 +69,14 @@ def test_build_solid_angle_map_invalid_spacing(spacing, match_str):
 @pytest.mark.parametrize("spacing", valid_spacings)
 def test_build_az_el_grid(spacing):
     """Test build_az_el_grid function."""
-    az_range, el_range, az_grid, el_grid = spatial_utils.build_az_el_grid(
-        spacing=spacing,
-        input_degrees=True,
-        output_degrees=True,
-        centered_azimuth=False,
-        centered_elevation=True,
+    (az_range, el_range, az_grid, el_grid, az_bin_edges, el_bin_edges) = (
+        spatial_utils.build_az_el_grid(
+            spacing=spacing,
+            input_degrees=True,
+            output_degrees=True,
+            centered_azimuth=False,
+            centered_elevation=True,
+        )
     )
 
     # Size checks
@@ -92,6 +94,12 @@ def test_build_az_el_grid(spacing):
     npt.assert_allclose(az_range, expected_az_range, atol=1e-12)
     npt.assert_allclose(el_range, expected_el_range, atol=1e-12)
 
+    # Check bin edges
+    expected_az_bin_edges = np.arange(0, 360 + spacing, spacing)
+    expected_el_bin_edges = np.arange(-90, 90 + spacing, spacing)
+    npt.assert_allclose(az_bin_edges, expected_az_bin_edges, atol=1e-11)
+    npt.assert_allclose(el_bin_edges, expected_el_bin_edges, atol=1e-11)
+
 
 def test_rewrap_even_spaced_el_az_grid_1d():
     """Test rewrap_even_spaced_el_az_grid function, without extra axis."""

imap_processing/ultra/utils/spatial_utils.py

@@ -95,9 +95,9 @@ def build_az_el_grid(
     output_degrees: bool = False,
     centered_azimuth: bool = False,
     centered_elevation: bool = True,
-) -> tuple[NDArray, NDArray, NDArray, NDArray]:
+) -> tuple[NDArray, NDArray, NDArray, NDArray, NDArray, NDArray]:
     """
-    Build a 2D grid of azimuth and elevation angles.
+    Build a 2D grid of azimuth and elevation angles, and their 1D bin edges.
 
     Azimuth and Elevation values represent the center of each grid cell,
     so the grid is offset by half the spacing.
@@ -123,7 +123,7 @@ def build_az_el_grid(
 
     Returns
     -------
-    tuple[NDArray, NDArray, NDArray, NDArray]
+    tuple[NDArray, NDArray, NDArray, NDArray, NDArray, NDArray]
         - The evenly spaced, 1D range of azimuth angles
         e.g.(0, 0.5, 1, ..., 359.5) deg.
         - The evenly spaced, 1D range of elevation angles
@@ -132,6 +132,12 @@ def build_az_el_grid(
         This grid will be constant along the elevation (0th) axis.
         - The 2D grid of elevation angles (elevations for each azimuth).
         This grid will be constant along the azimuth (1st) axis.
+        - The 1D bin edges for azimuth angles.
+        e.g. if spacing=1 deg:
+        az_bin_edges = [0, 1, 2, ..., 359, 360] deg.
+        - The 1D bin edges for elevation angles.
+        e.g. if spacing=1 deg:
+        el_bin_edges = [-90, -89, -88, ..., 89, 90] deg.
 
     Raises
     ------
@@ -147,6 +153,14 @@ def build_az_el_grid(
     if not np.isclose((np.pi / spacing) % 1, 0):
         raise ValueError("Spacing must divide evenly into pi radians.")
 
+    # Create the bin edges for azimuth and elevation.
+    # E.g. for spacing=1, az_bin_edges = [0, 1, 2, ..., 359, 360] deg.
+    el_bin_edges = np.linspace(-np.pi / 2, np.pi / 2, int(np.pi / spacing) + 1)
+    az_bin_edges = np.linspace(0, 2 * np.pi, int(2 * np.pi / spacing) + 1)
+
+    # Create the 2D grid of azimuth and elevation angles at center of each bin.
+    # These ranges are offset by half the spacing and are
+    # one element shorter than the bin edges.
     el_range = np.arange(spacing / 2, np.pi, spacing)
     az_range = np.arange(spacing / 2, 2 * np.pi, spacing)
     if centered_azimuth:
@@ -165,8 +179,10 @@ def build_az_el_grid(
         el_range = np.rad2deg(el_range)
         az_grid = np.rad2deg(az_grid)
         el_grid = np.rad2deg(el_grid)
+        az_bin_edges = np.rad2deg(az_bin_edges)
+        el_bin_edges = np.rad2deg(el_bin_edges)
 
-    return az_range, el_range, az_grid, el_grid
+    return az_range, el_range, az_grid, el_grid, az_bin_edges, el_bin_edges
 
 
 @typing.no_type_check