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

imap_processing/tests/ena_maps/test_ena_maps.py

@@ -0,0 +1,309 @@
+"""Test classes and methods in ena_maps.py."""
+
+from __future__ import annotations
+
+from unittest import mock
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from imap_processing.ena_maps import ena_maps
+from imap_processing.spice import geometry
+from imap_processing.tests.ultra.test_data.mock_data import mock_l1c_pset_product
+
+
+@pytest.fixture()
+def l1c_pset_products():
+    """Make fake L1C Ultra PSET products for testing"""
+    l1c_spatial_bin_spacing_deg = 10
+    return {
+        "spacing": l1c_spatial_bin_spacing_deg,
+        "products": [
+            mock_l1c_pset_product(
+                spacing_deg=l1c_spatial_bin_spacing_deg,
+                stripe_center_lon=mid_longitude,
+                timestr=f"2025-09-{i + 1:02d}T12:00:00",
+                head=("45" if (i % 2 == 0) else "90"),
+            )
+            for i, mid_longitude in enumerate(
+                np.arange(
+                    0,
+                    360,
+                    45,
+                )
+            )
+        ],
+    }
+
+
+class TestUltraPointingSet:
+    @pytest.fixture(autouse=True)
+    def _setup_ultra_l1c_pset_products(self, l1c_pset_products):
+        """Setup fixture data as class attributes"""
+        self.l1c_spatial_bin_spacing_deg = l1c_pset_products["spacing"]
+        self.l1c_pset_products = l1c_pset_products["products"]
+
+    @pytest.mark.usefixtures("_setup_ultra_l1c_pset_products")
+    def test_instantiate(self):
+        """Test instantiation of UltraPointingSet"""
+        ultra_psets = [
+            ena_maps.UltraPointingSet(
+                spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
+                l1c_dataset=l1c_product,
+            )
+            for l1c_product in self.l1c_pset_products
+        ]
+
+        for ultra_pset in ultra_psets:
+            # Check tiling is rectangular
+            assert ultra_pset.tiling_type == ena_maps.SkyTilingType.RECTANGULAR
+
+            # Check that the reference frame is correctly set
+            assert ultra_pset.spice_reference_frame == geometry.SpiceFrame.IMAP_DPS
+
+            # Check the number of points is (360/0.5) * (180/0.5)
+            np.testing.assert_equal(
+                ultra_pset.num_points,
+                int(360 * 180 / (self.l1c_spatial_bin_spacing_deg**2)),
+            )
+
+            # Check the repr exists
+            assert "UltraPointingSet" in repr(ultra_pset)
+
+    @pytest.mark.usefixtures("_setup_ultra_l1c_pset_products")
+    def test_uneven_spacing_raises_error(self):
+        """Test that uneven spacing in az/el raises ValueError"""
+
+        # Create dataset with uneven az spacing
+        uneven_az_dataset = xr.Dataset()
+        uneven_az_dataset["epoch"] = 1
+        uneven_az_dataset["azimuth_bin_center"] = np.array([0, 5, 15, 20, 30])
+        uneven_az_dataset["elevation_bin_center"] = np.arange(5)
+
+        with pytest.raises(ValueError, match="Azimuth bin spacing is not uniform"):
+            ena_maps.UltraPointingSet(
+                spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
+                l1c_dataset=uneven_az_dataset,
+            )
+
+        uneven_az_dataset["azimuth_bin_center"] = np.arange(5)
+        uneven_az_dataset["elevation_bin_center"] = np.array([0, 5, 15, 20, 30])
+
+        with pytest.raises(ValueError, match="Elevation bin spacing is not uniform"):
+            ena_maps.UltraPointingSet(
+                spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
+                l1c_dataset=uneven_az_dataset,
+            )
+
+        # Even but not the same spacing between az and el
+        uneven_az_dataset["azimuth_bin_center"] = np.arange(5)
+        uneven_az_dataset["elevation_bin_center"] = np.arange(5) * 2
+
+        with pytest.raises(
+            ValueError, match="Azimuth and elevation bin spacing do not match:"
+        ):
+            ena_maps.UltraPointingSet(
+                spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
+                l1c_dataset=uneven_az_dataset,
+            )
+
+
+class TestRectangularSkyMap:
+    @pytest.fixture(autouse=True)
+    def _setup_ultra_l1c_pset_products(self, l1c_pset_products):
+        """Setup fixture data as class attributes"""
+        self.l1c_spatial_bin_spacing_deg = l1c_pset_products["spacing"]
+        self.l1c_pset_products = l1c_pset_products["products"]
+        self.ultra_psets = [
+            ena_maps.UltraPointingSet(
+                spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
+                l1c_dataset=l1c_product,
+            )
+            for l1c_product in self.l1c_pset_products
+        ]
+
+    def test_instantiate(self):
+        """Test instantiation of RectangularSkyMap"""
+        rm = ena_maps.RectangularSkyMap(
+            spacing_deg=2,
+            spice_frame=geometry.SpiceFrame.ECLIPJ2000,
+        )
+
+        # Check that the map is empty
+        assert rm.data_dict == {}
+
+        # Check that the reference frame is correctly set
+        assert rm.spice_reference_frame == geometry.SpiceFrame.ECLIPJ2000
+
+        # Check the number of points is (360/2) * (180/2)
+        np.testing.assert_equal(rm.num_points, int(360 * 180 / 4))
+
+        # Check the repr exists
+        assert "RectangularSkyMap" in repr(rm)
+
+    @pytest.mark.usefixtures("_setup_ultra_l1c_pset_products")
+    @mock.patch("imap_processing.spice.geometry.frame_transform_az_el")
+    def test_project_pset_values_to_map_push_method(self, mock_frame_transform_az_el):
+        """
+        Test projection of PSET values to Rect. Map w "push" index matching method.
+
+        If frame_transform_az_el is mocked to return the az and el unchanged, and the
+        map has the same spacing as the PSETs, then the map should have
+        the same values as the PSETs, summed.
+        """
+        index_matching_method = ena_maps.IndexMatchMethod.PUSH
+
+        pset_spacing_deg = self.ultra_psets[0].spacing_deg
+
+        # Mock frame_transform to return the az and el unchanged
+        mock_frame_transform_az_el.side_effect = (
+            lambda et, az_el, from_frame, to_frame, degrees: az_el
+        )
+
+        rectangular_map = ena_maps.RectangularSkyMap(
+            spacing_deg=pset_spacing_deg,
+            spice_frame=geometry.SpiceFrame.ECLIPJ2000,
+        )
+
+        # Project each PSET's values to the map
+        for ultra_pset in self.ultra_psets:
+            rectangular_map.project_pset_values_to_map(
+                ultra_pset,
+                value_keys=["counts", "exposure_time"],
+                index_match_method=index_matching_method,
+            )
+
+        # Check that the map has been updated
+        assert rectangular_map.data_dict != {}
+
+        # Check that the map has the same values as the PSETs, summed
+        simple_summed_pset_counts = np.sum(
+            [pset["counts"].values for pset in self.l1c_pset_products], axis=0
+        ).reshape(rectangular_map.data_dict["counts"].shape)
+
+        np.testing.assert_allclose(
+            rectangular_map.data_dict["counts"],
+            simple_summed_pset_counts,
+        )
+
+    @pytest.mark.usefixtures("_setup_ultra_l1c_pset_products")
+    @mock.patch("imap_processing.spice.geometry.frame_transform_az_el")
+    def test_project_pset_values_to_map_pull_method(self, mock_frame_transform_az_el):
+        """Test projection to Rect. Map fails w "pull" index matching method."""
+
+        index_matching_method = ena_maps.IndexMatchMethod.PULL
+        rectangular_map = ena_maps.RectangularSkyMap(
+            spacing_deg=10,
+            spice_frame=geometry.SpiceFrame.ECLIPJ2000,
+        )
+
+        with pytest.raises(NotImplementedError):
+            rectangular_map.project_pset_values_to_map(
+                self.ultra_psets[0],
+                value_keys=["counts", "exposure_time"],
+                index_match_method=index_matching_method,
+            )
+
+
+class TestIndexMatching:
+    @pytest.fixture(autouse=True)
+    def _setup_ultra_l1c_pset_products(self, l1c_pset_products):
+        """Setup fixture data as class attributes"""
+        self.l1c_spatial_bin_spacing_deg = l1c_pset_products["spacing"]
+        self.l1c_pset_products = l1c_pset_products["products"]
+
+    @pytest.mark.parametrize(
+        "map_spacing_deg",
+        [0.5, 1, 10],
+    )
+    @mock.patch("imap_processing.spice.geometry.frame_transform_az_el")
+    def test_match_coords_to_indices_pset_to_rect_map(
+        self, mock_frame_transform_az_el, map_spacing_deg
+    ):
+        # Mock frame_transform to return the az and el unchanged
+        mock_frame_transform_az_el.side_effect = (
+            lambda et, az_el, from_frame, to_frame, degrees: az_el
+        )
+
+        # Mock a PSET, overriding the az/el points
+        mock_pset_input_frame = ena_maps.UltraPointingSet(
+            l1c_dataset=self.l1c_pset_products[0],
+            spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
+        )
+        manual_az_el_coords = np.array(
+            [
+                [0, -90],  # always -> RectangularSkyMap pixel 0
+                [0.4999999, -90],
+                [180.5, -89.5],
+                [359.5, -89.5],
+                [0.5, 0],
+                [180.5, 0],
+                [359.5, 0],
+                [0.5, 89.5],
+                [180.5, 89.5],
+                [359.5, 89.5],
+                [359.999999, 89.99999],
+            ]
+        )
+        mock_pset_input_frame.az_el_points = np.deg2rad(manual_az_el_coords)
+
+        # Manually calculate the resulting 1D pixel indices for each az/el pair
+        # (num of pixels in an az row spanning 180 deg of elevation) * (current az row)
+        # + (pixel along in current az row)
+        expected_output_pixel = np.array(
+            [
+                (az // map_spacing_deg) * (180 // map_spacing_deg)
+                + ((90 + el) // map_spacing_deg)
+                for [az, el] in manual_az_el_coords
+            ]
+        )
+
+        # Mock the rectangular map and check the output values
+        mock_rect_map = ena_maps.RectangularSkyMap(
+            spacing_deg=map_spacing_deg,
+            spice_frame=geometry.SpiceFrame.ECLIPJ2000,
+        )
+        flat_indices_input_grid_output_frame = ena_maps.match_coords_to_indices(
+            mock_pset_input_frame, mock_rect_map
+        )
+        assert mock_rect_map.num_points == 360 * 180 / map_spacing_deg**2
+        assert len(flat_indices_input_grid_output_frame) == len(manual_az_el_coords)
+        np.testing.assert_equal(
+            flat_indices_input_grid_output_frame, expected_output_pixel
+        )
+
+        # Check that the map's az/el points at the matched indices
+        # are the same as the input az/el points to within the spacing of the map
+        matched_map_az_el = mock_rect_map.az_el_points[
+            flat_indices_input_grid_output_frame
+        ]
+        np.testing.assert_allclose(
+            matched_map_az_el[:, 0],
+            mock_pset_input_frame.az_el_points[:, 0],
+            atol=np.deg2rad(map_spacing_deg),
+        )
+
+    def test_match_coords_to_indices_pset_to_healpix_map_other_map(
+        self,
+    ):
+        mock_pset_input_frame = ena_maps.UltraPointingSet(
+            l1c_dataset=self.l1c_pset_products[0],
+            spice_reference_frame=geometry.SpiceFrame.ECLIPJ2000,
+        )
+
+        # Until implemented, just change the tiling on a RectangularSkyMap
+        mock_hp_map = ena_maps.RectangularSkyMap(
+            spacing_deg=2,
+            spice_frame=geometry.SpiceFrame.ECLIPJ2000,
+        )
+        mock_hp_map.tiling_type = ena_maps.SkyTilingType.HEALPIX
+
+        # Should raise NotImplementedError
+        with pytest.raises(NotImplementedError):
+            ena_maps.match_coords_to_indices(mock_pset_input_frame, mock_hp_map)
+
+        mock_other_map = mock_hp_map
+        mock_other_map.tiling_type = "INVALID"
+        with pytest.raises(ValueError, match="Tiling type of the output frame"):
+            ena_maps.match_coords_to_indices(mock_pset_input_frame, mock_other_map)

imap_processing/tests/ena_maps/test_map_utils.py

@@ -0,0 +1,286 @@
+import numpy as np
+import pytest
+
+from imap_processing.ena_maps.utils import map_utils
+
+
+class TestENAMapMappingUtils:
+    def test_bin_single_array_at_indices(
+        self,
+    ):
+        """Test coverage for bin_single_array_at_indices function w/ simple 1D input"""
+        value_array = np.array([1, 2, 3, 4, 5, 6])
+        input_indices = np.array([0, 1, 2, 2, 1, 0])
+        projection_indices = np.array([1, 2, 3, 1, 2, 3])
+        projection_grid_shape = (5,)
+        expected_projection_values = np.array([0, 4, 4, 4, 0])
+        projection_values = map_utils.bin_single_array_at_indices(
+            value_array,
+            input_indices=input_indices,
+            projection_indices=projection_indices,
+            projection_grid_shape=projection_grid_shape,
+        )
+        np.testing.assert_equal(projection_values, expected_projection_values)
+
+    def test_bin_single_array_at_indices_extra_axis(
+        self,
+    ):
+        """Test coverage for bin_single_array_at_indices function w/ simple 2D input,
+        Corresponding to an extra axis that is not spatially binned.
+        """
+        # Binning will occur along axis 0 (combining 1, 2, 3 and 4, 5, 6 separately)
+        value_array = np.array(
+            [
+                [1, 4],
+                [2, 5],
+                [3, 6],
+            ]
+        )
+        input_indices = np.array([0, 1, 2, 2])
+        projection_indices = np.array([1, 0, 1, 6])
+        projection_grid_shape = (7, 1)
+        expected_projection_values = np.array(
+            [
+                [2, 5],
+                [4, 10],
+                [0, 0],
+                [0, 0],
+                [0, 0],
+                [0, 0],
+                [3, 6],
+            ]
+        )
+        projection_values = map_utils.bin_single_array_at_indices(
+            value_array,
+            input_indices=input_indices,
+            projection_indices=projection_indices,
+            projection_grid_shape=projection_grid_shape,
+        )
+
+        np.testing.assert_equal(projection_values, expected_projection_values)
+
+    @pytest.mark.parametrize(
+        "projection_grid_shape", [(1, 1), (10, 10), (180, 360), (360, 720), (360, 180)]
+    )
+    @pytest.mark.parametrize(
+        "input_grid_shape", [(1, 1), (10, 10), (180, 360), (360, 720), (360, 180)]
+    )
+    def test_bin_single_array_at_indices_complex_2d(
+        self, projection_grid_shape, input_grid_shape
+    ):
+        """Test coverage for bin_single_array_at_indices function w/ complex 2D input,
+        Corresponding to an extra axis that is not spatially binned.
+        Parameterized across different input and projection grid shapes.
+        """
+        np.random.seed(0)
+        extra_axis_size = 11  # Another axis which is not spatially binned, e.g. energy
+        input_grid_size = np.prod(input_grid_shape)
+        projection_grid_size = np.prod(projection_grid_shape)
+        value_array = np.random.rand(input_grid_size, extra_axis_size)
+        input_indices = np.random.randint(0, input_grid_size, size=1000)
+        projection_indices = np.random.randint(0, projection_grid_size, size=1000)
+        projection_values = map_utils.bin_single_array_at_indices(
+            value_array,
+            input_indices=input_indices,
+            projection_indices=projection_indices,
+            projection_grid_shape=projection_grid_shape,
+        )
+
+        # Explicitly check that the shape of the output is the same as projection grid
+        np.testing.assert_equal(
+            projection_values.shape,
+            (
+                projection_grid_size,
+                extra_axis_size,
+            ),
+        )
+
+        # Create the expected projection values by summing the input values in a loop
+        # This is different from the binning function, which uses np.bincount
+        expected_projection_values = np.zeros((projection_grid_size, extra_axis_size))
+        for ii, ip in zip(input_indices, projection_indices):
+            expected_projection_values[ip, :] += value_array[ii, :]
+
+        np.testing.assert_allclose(projection_values, expected_projection_values)
+
+    @pytest.mark.parametrize("projection_grid_shape", [(1, 1), (10, 10), (180, 360)])
+    @pytest.mark.parametrize("input_grid_shape", [(1, 1), (10, 10), (180, 360)])
+    @pytest.mark.parametrize("num_extra_dims", [1, 2, 3, 5])
+    def test_bin_single_array_at_indices_complex_3d(
+        self, projection_grid_shape, input_grid_shape, num_extra_dims
+    ):
+        """Test coverage for bin_single_array_at_indices function w/ complex N-Dim input
+        Corresponding to 2 extra axes that are not spatially binned.
+        Parameterized across different input and projection grid shapes.
+        """
+        np.random.seed(0)
+        extra_axes_sizes = np.full(num_extra_dims, 3, dtype=int).tolist()
+        input_grid_size = np.prod(input_grid_shape)
+        projection_grid_size = np.prod(projection_grid_shape)
+        value_array = np.random.rand(input_grid_size, *extra_axes_sizes)
+        input_indices = np.random.randint(0, input_grid_size, size=1000)
+        projection_indices = np.random.randint(0, projection_grid_size, size=1000)
+        projection_values = map_utils.bin_single_array_at_indices(
+            value_array,
+            input_indices=input_indices,
+            projection_indices=projection_indices,
+            projection_grid_shape=projection_grid_shape,
+        )
+
+        # Explicitly check that the shape of the output is the same as projection grid
+        np.testing.assert_equal(
+            projection_values.shape,
+            (projection_grid_size, *extra_axes_sizes),
+        )
+
+        # Create the expected projection values by summing the input values in a loop
+        # This is different from the binning function, which uses np.bincount
+        expected_projection_values = np.zeros((projection_grid_size, *extra_axes_sizes))
+        for ii, ip in zip(input_indices, projection_indices):
+            expected_projection_values[ip, ...] += value_array[ii, ...]
+
+        np.testing.assert_allclose(projection_values, expected_projection_values)
+
+    # Parameterize by the size of the projection grid,
+    # which is not necessarily same size as input grid
+    @pytest.mark.parametrize("projection_grid_shape", [(1, 1), (10, 10), (360, 720)])
+    def test_bin_values_at_indices_collapse_to_idx_zero(self, projection_grid_shape):
+        """Test coverage for bin_values_at_indices function w/ dict of multiple
+        1D input value arrays and a single 2D input value array.
+        All input values are binned to the first index of the projection grid.
+        Parameterized across different projection grid shapes.
+        """
+        # 1D input values (2nd will be scalar multiple of 1st)
+        input_values_1d_1 = np.array([0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        scale_factor_1d = 1.5
+        input_values_1d_2 = input_values_1d_1 * scale_factor_1d
+
+        # 2D input values. The second axis (different cols) will be summed independently
+        input_values_2d = np.array(
+            [
+                [-0.5, 0, 0.5],
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ]
+        )
+        extra_axis_size_2d = input_values_2d.shape[1]
+
+        # 3D input values
+        input_values_3d = np.zeros((input_values_2d.shape[0], 3, 3))
+        input_values_3d[:2] = np.array(
+            [
+                [
+                    [1, 2, 3],
+                    [4, 5, 6],
+                    [7, 8, 9],
+                ],
+                [
+                    [10, 11, 12],
+                    [13, 14, 15],
+                    [16, 17, 18],
+                ],
+            ]
+        )
+        extra_axes_size_3d = input_values_3d.shape[1:]
+
+        # Set up the expected projection values
+        expected_projection_values_1d_1 = np.zeros(projection_grid_shape).ravel()
+        expected_projection_values_1d_1[0] = np.sum(input_values_1d_1)
+        expected_projection_values_1d_2 = (
+            expected_projection_values_1d_1 * scale_factor_1d
+        )
+        expected_projection_values_2d = np.zeros(
+            (np.prod(projection_grid_shape), extra_axis_size_2d)
+        )
+        expected_projection_values_2d[0, :] = np.array([11.5, 15, 18.5])
+        expected_projection_values_3d = np.zeros(
+            (np.prod(projection_grid_shape), *extra_axes_size_3d)
+        )
+        expected_projection_values_3d[0, :, :] = np.array(
+            [
+                [11, 13, 15],
+                [17, 19, 21],
+                [23, 25, 27],
+            ]
+        )
+
+        input_values_to_bin = {
+            "sum_variable_1d_1": input_values_1d_1,
+            "sum_variable_1d_2": input_values_1d_2,
+            "sum_variable_2d": np.array(input_values_2d),
+            "sum_variable_3d": np.array(input_values_3d),
+        }
+
+        # Set up indices
+        input_indices = np.arange(len(input_values_1d_1))
+        projection_indices = np.zeros_like(input_indices)
+
+        output_dict = map_utils.bin_values_at_indices(
+            projection_indices=projection_indices,
+            projection_grid_shape=projection_grid_shape,
+            input_values_to_bin=input_values_to_bin,
+            input_indices=input_indices,
+        )
+
+        np.testing.assert_equal(
+            output_dict["sum_variable_1d_1"], expected_projection_values_1d_1
+        )
+        np.testing.assert_equal(
+            output_dict["sum_variable_1d_2"], expected_projection_values_1d_2
+        )
+        np.testing.assert_equal(
+            output_dict["sum_variable_2d"], expected_projection_values_2d
+        )
+        np.testing.assert_equal(
+            output_dict["sum_variable_3d"], expected_projection_values_3d
+        )
+
+    def test_bin_values_at_indices_2d_indices_raises(self):
+        """2D indices are not supported for binning.
+        Test that ValueError is raised."""
+        input_values = np.array([1, 2, 3])
+        input_indices = np.array([[0, 1], [1, 2]])
+        projection_indices = np.array([0, 1, 2])
+        projection_grid_shape = (3,)
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                "Indices must be 1D arrays. If using a rectangular grid, "
+                "the indices must be unwrapped."
+            ),
+        ):
+            map_utils.bin_single_array_at_indices(
+                input_values,
+                input_indices=input_indices,
+                projection_indices=projection_indices,
+                projection_grid_shape=projection_grid_shape,
+            )
+
+    def test_bin_values_at_indices_mismatched_sizes_raises(self):
+        """Mismatched input and projection indices should raise an error.
+        Test that ValueError is raised."""
+        input_values = np.array([1, 2, 3])
+        input_indices = np.array([0, 1, 0, 1])
+        projection_indices = np.array([0, 1, 2])
+        projection_grid_shape = (3,)
+
+        with pytest.raises(
+            ValueError,
+            match=("The number of input and projection indices must be the same"),
+        ):
+            map_utils.bin_single_array_at_indices(
+                input_values,
+                input_indices=input_indices,
+                projection_indices=projection_indices,
+                projection_grid_shape=projection_grid_shape,
+            )