From daa382984ee84f9df0c1d01be78472b2b8600ddc Mon Sep 17 00:00:00 2001 From: Tim Plummer Date: Fri, 6 Nov 2026 09:50:58 -0700 Subject: [PATCH 1/4] Combine PSET non-spatial dims by coordinate label when projecting to a map AbstractSkyMap.project_pset_values_to_map now adds each PSET's binned values into the map variable by non-spatial coordinate label instead of by array position, for both PUSH and PULL. PSETs may hold a subset of the map's labels, in any order. Mismatched dims, unknown labels, and size mismatches raise ValueError instead of silently broadcasting. The epoch dim is combined by position. Map variables can be pre-initialized with the full coords. Closes #3632 Co-Authored-By: Claude Opus 5.5 --- imap_processing/ena_maps/ena_maps.py | 226 +++++++++++- .../tests/ena_maps/test_ena_maps.py | 332 ++++++++++++++++++ 2 files changed, 550 insertions(+), 8 deletions(-) diff --git a/imap_processing/ena_maps/ena_maps.py b/imap_processing/ena_maps/ena_maps.py index eab9e0c3ee..c80206777d 100644 --- a/imap_processing/ena_maps/ena_maps.py +++ b/imap_processing/ena_maps/ena_maps.py @@ -925,6 +925,22 @@ def project_pset_values_to_map( # noqa: PLR0912 the map grid correspond to which pixels in the pointing set grid, and then binning the values at those indices from the pointing set to the map. + Non-spatial dimensions (e.g. energy step, calibration product) are combined + into the map by coordinate label, not by array position. A pointing set may + therefore contain a subset of the map variable's labels, in any order, and + its values are added only into the matching map labels. The ``epoch`` + dimension, and any dimension without a coordinate on either the map or the + pointing set, is combined by position and must have the same size in both. + + If a value key is not yet a variable in ``self.data_1d``, it is created, + filled with zeros, using the dims of the pointing set variable. Dims that + already exist in ``self.data_1d`` keep the map's size and coords; other + dims take the size and coords of this pointing set. To project pointing + sets that each cover only some labels (e.g. some ESA energy steps), a + caller can pre-initialize the map variable in ``self.data_1d`` with the + full non-spatial coords and the spatial pixel dim as the last dim (required), + before projecting any pointing set. + Parameters ---------- pointing_set : PointingSet @@ -945,8 +961,13 @@ def project_pset_values_to_map( # noqa: PLR0912 Raises ------ - ValueError + KeyError If a value key is not found in the pointing set. + ValueError + If a pointing set variable cannot be combined into the map variable: + its non-spatial dims differ from the map variable's, one of its + coordinate labels is not in the map variable's coords, or a dim + combined by position differs in size. """ if value_keys is None: value_keys = list(pointing_set.data.data_vars.keys()) @@ -972,11 +993,19 @@ def project_pset_values_to_map( # noqa: PLR0912 input_object=self, output_object=pointing_set, ) + # The map pixels that pull from a valid PSET pixel, and the PSET + # pixels they pull from, are the same for every value key. + valid_map_mask = pset_valid_mask[matched_indices_pull] + valid_matched_indices_pull = matched_indices_pull[valid_map_mask] else: raise NotImplementedError( "Only PUSH and PULL index matching methods are supported." ) + # Map positions of this PSET's non-spatial labels, resolved once per dim + # and shared by all value keys (map variables share coords per dim). + dim_positions: dict[str, NDArray | None] = {} + for value_key in value_keys: if value_key not in pointing_set.data.data_vars: raise ValueError(f"Value key {value_key} not found in pointing set.") @@ -990,10 +1019,19 @@ def project_pset_values_to_map( # noqa: PLR0912 if value_key not in self.data_1d.data_vars: # Initialize the map data array if it doesn't exist (values start at 0) - output_shape = (*raveled_pset_data.shape[:-1], self.num_points) + # Non-spatial dims that already exist in the map (e.g. from a + # pre-initialized variable) keep the map's size, so that the + # PSET's values are combined into the map's labels below. + output_dims = pointing_set.unwrapped_dims_dict[value_key] + output_shape = tuple( + self.num_points + if dim == CoordNames.GENERIC_PIXEL.value + else self.data_1d.sizes.get(dim, raveled_pset_data.sizes[dim]) + for dim in output_dims + ) self.data_1d[value_key] = xr.DataArray( np.zeros(output_shape), - dims=pointing_set.unwrapped_dims_dict[value_key], + dims=output_dims, ) # Make coordinates for the map data array if they don't exist @@ -1033,20 +1071,29 @@ def project_pset_values_to_map( # noqa: PLR0912 # dividing pointing_projected_values by some binned weights. # For unweighted means, we could use the number of pointing set pixels # that correspond to each map pixel as the weights. - self.data_1d[value_key] += pointing_projected_values + self._add_projected_values( + value_key=value_key, + projected_values=pointing_projected_values, + pset_data=data_bc, + map_pixel_mask=None, + dim_positions=dim_positions, + ) elif index_match_method is IndexMatchMethod.PULL: - valid_map_mask = pset_valid_mask[matched_indices_pull] # We know that there will only be one value per sky map pixel, # so we can use the matched indices directly pointing_projected_values = raveled_pset_data.values[ - ..., matched_indices_pull[valid_map_mask] + ..., valid_matched_indices_pull ] # TODO: we may need to allow for unweighted/weighted means here by # dividing pointing_projected_values by some binned weights. # For unweighted means, we could use the number of pointing set pixels # that correspond to each map pixel as the weights. - self.data_1d[value_key].values[..., valid_map_mask] += ( - pointing_projected_values + self._add_projected_values( + value_key=value_key, + projected_values=pointing_projected_values, + pset_data=raveled_pset_data, + map_pixel_mask=valid_map_mask, + dim_positions=dim_positions, ) # TODO: The max epoch needs to include the pset duration. Right now it @@ -1054,6 +1101,169 @@ def project_pset_values_to_map( # noqa: PLR0912 self.min_epoch = min(self.min_epoch, pointing_set.epoch) self.max_epoch = max(self.max_epoch, pointing_set.epoch) + def _add_projected_values( + self, + value_key: str, + projected_values: NDArray, + pset_data: xr.DataArray, + map_pixel_mask: NDArray | None, + dim_positions: dict[str, NDArray | None], + ) -> None: + """ + Add a PSET's projected values into a map variable, matching coords by label. + + Each non-spatial dimension of the PSET is combined into the map variable + by coordinate label, so a PSET may hold a subset of the map variable's + labels, in any order. Positions are used instead of labels for the + ``epoch`` dimension (each PSET has its own epoch value) and for any + dimension that lacks a coordinate on the map or the PSET; those + dimensions must have the same size in both. + + Parameters + ---------- + value_key : str + Name of the variable in ``self.data_1d`` to add the values into. + projected_values : NDArray + The PSET values binned to the map. The leading axes match the + non-spatial dimensions of ``pset_data`` and the final axis + corresponds to the map pixels selected by ``map_pixel_mask``. + pset_data : xr.DataArray + The raveled PSET data the values were projected from. Its dims (all + except the final, spatial dim) and coords label the leading axes + of ``projected_values``. + map_pixel_mask : NDArray or None + Boolean mask over the map pixels that the final axis of + ``projected_values`` corresponds to. None means all map pixels. + dim_positions : dict[str, NDArray | None] + Cache of resolved map positions per non-spatial dim for the current + PSET, updated in place. Map variables share one coordinate per dim, + so a dim only needs resolving once per PSET. A None entry means the + PSET matches the map exactly along that dim. + + Raises + ------ + ValueError + If the map variable's last dim is not the spatial pixel dim, or if + the PSET and map variable non-spatial dims differ. + """ + map_var = self.data_1d[value_key] + pixel_dim = CoordNames.GENERIC_PIXEL.value + pset_dims = tuple(pset_data.dims[:-1]) + map_dims = map_var.dims[:-1] + + if map_var.dims[-1] != pixel_dim: + raise ValueError( + f"Map variable '{value_key}' has dims {map_var.dims}; the last dim " + f"must be the spatial pixel dim '{pixel_dim}'." + ) + if set(pset_dims) != set(map_dims) or len(pset_dims) != len(map_dims): + raise ValueError( + f"Non-spatial dims of PSET variable '{value_key}' {pset_dims} do not " + f"match the non-spatial dims of the map variable {map_dims}." + ) + for dim in map_dims: + if dim not in dim_positions: + dim_positions[dim] = self._resolve_dim_positions( + value_key, dim, pset_data, map_var + ) + + # Put the PSET values in the map variable's dim order + if pset_dims != map_dims: + projected_values = np.moveaxis( + projected_values, + [pset_dims.index(dim) for dim in map_dims], + range(len(map_dims)), + ) + if map_pixel_mask is not None: + # Spread the values onto the full pixel axis, so the add below uses + # a slice for pixels. This is much faster than indexing the pixels. + full_pixel_values = np.zeros( + (*projected_values.shape[:-1], self.num_points), + dtype=projected_values.dtype, + ) + full_pixel_values[..., map_pixel_mask] = projected_values + projected_values = full_pixel_values + + if all(dim_positions[dim] is None for dim in map_dims): + # Fast path: the PSET matches the map's labels exactly + map_var.values += projected_values + else: + # Orthogonal (outer-product) index over the non-spatial dims + non_spatial_index = np.ix_( + *( + np.arange(map_var.sizes[dim]) + if dim_positions[dim] is None + else dim_positions[dim] + for dim in map_dims + ) + ) + map_var.values[(*non_spatial_index, slice(None))] += projected_values + + @staticmethod + def _resolve_dim_positions( + value_key: str, dim: str, pset_data: xr.DataArray, map_var: xr.DataArray + ) -> NDArray | None: + """ + Find the map positions of a PSET's labels along one non-spatial dim. + + Parameters + ---------- + value_key : str + Name of the variable being projected, used in error messages. + dim : str + The non-spatial dimension to resolve. + pset_data : xr.DataArray + The raveled PSET data being projected. + map_var : xr.DataArray + The map variable the PSET data is being added into. + + Returns + ------- + NDArray or None + Map positions of each PSET label along ``dim``, or None if the + PSET matches the map exactly along ``dim`` (or ``dim`` is combined + by position). + + Raises + ------ + ValueError + If a PSET coordinate label is not in the map variable's coords, if a + PSET coordinate has duplicate labels, or if a positionally combined + dimension differs in size. + """ + pset_size = pset_data.sizes[dim] + map_size = map_var.sizes[dim] + label_matched = ( + dim != CoordNames.TIME.value + and dim in pset_data.coords + and dim in map_var.coords + ) + if not label_matched: + if pset_size != map_size: + raise ValueError( + f"Cannot combine dim '{dim}' of '{value_key}' by position: " + f"PSET size {pset_size} != map size {map_size}." + ) + return None + + pset_labels = pset_data.coords[dim].values + map_labels = map_var.coords[dim].values + if np.array_equal(pset_labels, map_labels): + return None + if len(np.unique(pset_labels)) != pset_size: + raise ValueError( + f"PSET coordinate '{dim}' of '{value_key}' has duplicate labels: " + f"{pset_labels}." + ) + positions = map_var.indexes[dim].get_indexer(pset_labels) + if np.any(positions < 0): + raise ValueError( + f"PSET coordinate '{dim}' of '{value_key}' has labels " + f"{pset_labels[positions < 0]} not found in the map coordinate " + f"values {map_labels}." + ) + return positions + @classmethod def from_properties_json( cls, json_path: str | Path diff --git a/imap_processing/tests/ena_maps/test_ena_maps.py b/imap_processing/tests/ena_maps/test_ena_maps.py index 4575a584fd..c20ea1d3e1 100644 --- a/imap_processing/tests/ena_maps/test_ena_maps.py +++ b/imap_processing/tests/ena_maps/test_ena_maps.py @@ -2119,3 +2119,335 @@ def test_to_json_and_from_json(self): # add the values_to_pull_project key to the original json original_json["values_to_pull_project"] = [] assert original_json == output_json + + +# Full non-spatial labels of the map variables in the label-based projection tests +FULL_GAIN_CONFIG_IDS = [0, 1] +FULL_CAL_PRODS = [0, 1] +FULL_ESA_STEPS = list(range(1, 10)) +NON_SPATIAL_DIMS = ("epoch", "gain_config_id", "calibration_prod", "esa_energy_step") + + +def make_labeled_rect_pset( + gain_config_ids=FULL_GAIN_CONFIG_IDS, + cal_prods=FULL_CAL_PRODS, + esa_steps=FULL_ESA_STEPS, + epoch=0, + seed=0, + spacing_deg=30, +) -> ena_maps.RectangularPointingSet: + """Make a small RectangularPointingSet with labeled non-spatial dims.""" + sky_grid = spatial_utils.AzElSkyGrid(spacing_deg=spacing_deg) + longitude = np.sort(sky_grid.az_bin_midpoints) + latitude = np.sort(sky_grid.el_bin_midpoints) + shape = ( + 1, + len(gain_config_ids), + len(cal_prods), + len(esa_steps), + len(longitude), + len(latitude), + ) + rng = np.random.default_rng(seed) + values = rng.integers(1, 100, size=shape).astype(float) + ds = xr.Dataset( + { + "counts": ( + (*NON_SPATIAL_DIMS, "longitude", "latitude"), + values, + ), + "exposure_factor": ( + (*NON_SPATIAL_DIMS, "longitude", "latitude"), + values * 10, + ), + "epoch_delta": (("epoch",), np.array([1_000_000_000])), + }, + coords={ + "epoch": np.array([epoch], dtype=np.int64), + # uint8 PSET labels are combined into int64 map labels + "gain_config_id": np.array(gain_config_ids, dtype=np.uint8), + "calibration_prod": np.array(cal_prods, dtype=np.uint8), + "esa_energy_step": np.array(esa_steps, dtype=np.uint8), + "longitude": longitude, + "latitude": latitude, + }, + ) + return ena_maps.RectangularPointingSet(ds) + + +def pad_pset_to_full_labels( + pset: ena_maps.RectangularPointingSet, +) -> ena_maps.RectangularPointingSet: + """Zero-fill a PSET to the full, ordered labels (the reference projection).""" + padded = pset.data.reindex( + gain_config_id=FULL_GAIN_CONFIG_IDS, + calibration_prod=FULL_CAL_PRODS, + esa_energy_step=FULL_ESA_STEPS, + fill_value=0, + ) + return ena_maps.RectangularPointingSet(padded) + + +def make_sky_map(map_type: str) -> ena_maps.AbstractSkyMap: + if map_type == "rectangular": + return ena_maps.RectangularSkyMap( + spacing_deg=30, spice_frame=geometry.SpiceFrame.ECLIPJ2000 + ) + return ena_maps.HealpixSkyMap(nside=4, spice_frame=geometry.SpiceFrame.ECLIPJ2000) + + +def pre_initialize_map_var( + sky_map: ena_maps.AbstractSkyMap, value_key: str = "counts", epoch: int = 0 +): + """Create a zero-filled map variable with the full non-spatial coords.""" + coords = { + "epoch": np.array([epoch], dtype=np.int64), + "gain_config_id": np.array(FULL_GAIN_CONFIG_IDS, dtype=np.int64), + "calibration_prod": np.array(FULL_CAL_PRODS, dtype=np.int64), + "esa_energy_step": np.array(FULL_ESA_STEPS, dtype=np.int64), + } + sky_map.data_1d[value_key] = xr.DataArray( + np.zeros((*[len(c) for c in coords.values()], sky_map.num_points)), + dims=(*NON_SPATIAL_DIMS, CoordNames.GENERIC_PIXEL.value), + coords=coords, + ) + + +@mock.patch("imap_processing.spice.geometry.frame_transform_az_el") +class TestProjectByCoordinateLabel: + """Tests that PSET non-spatial dims are combined into a map by label.""" + + @staticmethod + def _identity(mock_frame_transform_az_el): + mock_frame_transform_az_el.side_effect = ( + lambda et, az_el, from_frame, to_frame, degrees: az_el + ) + + @pytest.mark.parametrize("map_type", ["rectangular", "healpix"]) + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + @pytest.mark.parametrize("use_mask", [False, True]) + def test_subset_and_reordered_labels_match_padded_reference( + self, mock_frame_transform_az_el, map_type, method, use_mask + ): + """PSET labels that are a reordered subset land only in matching slots.""" + self._identity(mock_frame_transform_az_el) + pset = make_labeled_rect_pset( + gain_config_ids=[1], cal_prods=[1, 0], esa_steps=[2, 5, 9] + ) + mask = None + if use_mask: + mask = np.zeros(pset.num_points, dtype=bool) + mask[::3] = True + + labeled_map = make_sky_map(map_type) + pre_initialize_map_var(labeled_map) + labeled_map.project_pset_values_to_map( + pset, ["counts"], index_match_method=method, pset_valid_mask=mask + ) + + reference_map = make_sky_map(map_type) + reference_map.project_pset_values_to_map( + pad_pset_to_full_labels(pset), + ["counts"], + index_match_method=method, + pset_valid_mask=mask, + ) + + result = labeled_map.data_1d["counts"] + assert result.sum() > 0 + np.testing.assert_array_equal( + result.values, reference_map.data_1d["counts"].values + ) + # Slots for labels not in the PSET are untouched + assert (result.sel(gain_config_id=0) == 0).all() + assert (result.sel(esa_energy_step=[1, 3, 4, 6, 7, 8]) == 0).all() + # Values land at the matching labels, not by position + np.testing.assert_array_equal( + result.sel(gain_config_id=1, calibration_prod=0, esa_energy_step=5), + reference_map.data_1d["counts"].sel( + gain_config_id=1, calibration_prod=0, esa_energy_step=5 + ), + ) + + @pytest.mark.parametrize("map_type", ["rectangular", "healpix"]) + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + def test_identical_coords_fast_path( + self, mock_frame_transform_az_el, map_type, method + ): + """PSETs with identical coords accumulate positionally across epochs.""" + self._identity(mock_frame_transform_az_el) + pset_a = make_labeled_rect_pset(epoch=0, seed=1) + pset_b = make_labeled_rect_pset(epoch=10**9, seed=2) + + sky_map = make_sky_map(map_type) + sky_map.project_pset_values_to_map(pset_a, ["counts"], method) + sky_map.project_pset_values_to_map(pset_b, ["counts"], method) + + single_maps = [] + for pset in (pset_a, pset_b): + single_map = make_sky_map(map_type) + single_map.project_pset_values_to_map(pset, ["counts"], method) + single_maps.append(single_map.data_1d["counts"].values) + + result = sky_map.data_1d["counts"] + assert result.dims == (*NON_SPATIAL_DIMS, CoordNames.GENERIC_PIXEL.value) + # The map keeps the first PSET's epoch label + np.testing.assert_array_equal(result["epoch"], [0]) + np.testing.assert_array_equal(result.values, single_maps[0] + single_maps[1]) + + @pytest.mark.parametrize("map_type", ["rectangular", "healpix"]) + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + def test_projection_order_independent( + self, mock_frame_transform_az_el, map_type, method + ): + """Projecting A then B gives the same map as B then A.""" + self._identity(mock_frame_transform_az_el) + pset_a = make_labeled_rect_pset( + gain_config_ids=[0], esa_steps=[2, 5, 9], epoch=0, seed=1 + ) + pset_b = make_labeled_rect_pset(gain_config_ids=[1], epoch=10**9, seed=2) + + results = [] + for psets in ((pset_a, pset_b), (pset_b, pset_a)): + sky_map = make_sky_map(map_type) + pre_initialize_map_var(sky_map) + for pset in psets: + sky_map.project_pset_values_to_map(pset, ["counts"], method) + results.append(sky_map.data_1d["counts"].values) + + np.testing.assert_array_equal(results[0], results[1]) + # Both gain configs received counts + assert results[0][:, 0].sum() > 0 + assert results[0][:, 1].sum() > 0 + + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + @pytest.mark.parametrize("subset", [False, True]) + def test_pset_dim_order_differs_from_map( + self, mock_frame_transform_az_el, method, subset + ): + """PSET dims in a different order than the map's are moved to map order.""" + self._identity(mock_frame_transform_az_el) + kwargs = {"gain_config_ids": [1], "esa_steps": [2, 5, 9]} if subset else {} + pset = make_labeled_rect_pset(**kwargs) + mask = np.zeros(pset.num_points, dtype=bool) + mask[::2] = True + + reference_map = make_sky_map("rectangular") + pre_initialize_map_var(reference_map) + reference_map.project_pset_values_to_map( + pset, ["counts"], method, pset_valid_mask=mask + ) + + pset.data["counts"] = pset.data["counts"].transpose( + "esa_energy_step", "epoch", "calibration_prod", "gain_config_id", ... + ) + sky_map = make_sky_map("rectangular") + pre_initialize_map_var(sky_map) + sky_map.project_pset_values_to_map( + pset, ["counts"], method, pset_valid_mask=mask + ) + + assert sky_map.data_1d["counts"].sum() > 0 + np.testing.assert_array_equal( + sky_map.data_1d["counts"].values, reference_map.data_1d["counts"].values + ) + + def test_new_variable_uses_existing_map_coords(self, mock_frame_transform_az_el): + """A variable created after a pre-initialized one takes the map's labels.""" + self._identity(mock_frame_transform_az_el) + pset = make_labeled_rect_pset(gain_config_ids=[1], esa_steps=[2, 5, 9]) + sky_map = make_sky_map("rectangular") + pre_initialize_map_var(sky_map, "counts") + + sky_map.project_pset_values_to_map(pset, ["counts", "exposure_factor"]) + + exposure = sky_map.data_1d["exposure_factor"] + assert exposure.sizes["gain_config_id"] == len(FULL_GAIN_CONFIG_IDS) + assert exposure.sizes["esa_energy_step"] == len(FULL_ESA_STEPS) + np.testing.assert_array_equal( + exposure.values, sky_map.data_1d["counts"].values * 10 + ) + + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + def test_label_not_in_map_raises(self, mock_frame_transform_az_el, method): + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + pre_initialize_map_var(sky_map) + pset = make_labeled_rect_pset(gain_config_ids=[3]) + with pytest.raises(ValueError, match="not found in the map coordinate"): + sky_map.project_pset_values_to_map(pset, ["counts"], method) + + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + def test_size_one_dim_without_labels_raises( + self, mock_frame_transform_az_el, method + ): + """A length-1 dim without coords no longer broadcasts into the map.""" + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + sky_map.project_pset_values_to_map(make_labeled_rect_pset(), ["counts"], method) + + pset = make_labeled_rect_pset(gain_config_ids=[1]) + pset.data = pset.data.drop_vars("gain_config_id") + with pytest.raises(ValueError, match="by position"): + sky_map.project_pset_values_to_map(pset, ["counts"], method) + + def test_epoch_size_mismatch_raises(self, mock_frame_transform_az_el): + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + sky_map.data_1d["counts"] = xr.DataArray( + np.zeros((2, 2, 2, 9, sky_map.num_points)), + dims=(*NON_SPATIAL_DIMS, CoordNames.GENERIC_PIXEL.value), + coords={"epoch": [0, 1]}, + ) + with pytest.raises(ValueError, match="dim 'epoch'"): + sky_map.project_pset_values_to_map(make_labeled_rect_pset(), ["counts"]) + + def test_missing_dim_raises(self, mock_frame_transform_az_el): + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + pre_initialize_map_var(sky_map) + pset = make_labeled_rect_pset(gain_config_ids=[1]) + pset.data = pset.data.squeeze("gain_config_id", drop=True) + with pytest.raises(ValueError, match="do not match the non-spatial dims"): + sky_map.project_pset_values_to_map(pset, ["counts"]) + + def test_duplicate_labels_raise(self, mock_frame_transform_az_el): + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + pre_initialize_map_var(sky_map) + pset = make_labeled_rect_pset(esa_steps=[2, 2]) + with pytest.raises(ValueError, match="duplicate labels"): + sky_map.project_pset_values_to_map(pset, ["counts"]) + + def test_pre_initialized_without_pixel_dim_raises(self, mock_frame_transform_az_el): + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + sky_map.data_1d["counts"] = xr.DataArray( + np.zeros((1, 2, 2, 9)), dims=NON_SPATIAL_DIMS + ) + with pytest.raises(ValueError, match="must be the spatial pixel dim"): + sky_map.project_pset_values_to_map(make_labeled_rect_pset(), ["counts"]) + + def test_pre_initialized_pixel_dim_not_last_raises( + self, mock_frame_transform_az_el + ): + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + pre_initialize_map_var(sky_map) + sky_map.data_1d["counts"] = sky_map.data_1d["counts"].transpose( + CoordNames.GENERIC_PIXEL.value, ... + ) + with pytest.raises(ValueError, match="must be the spatial pixel dim"): + sky_map.project_pset_values_to_map(make_labeled_rect_pset(), ["counts"]) From da83d8a87e194b38d3d22155fd382ceec85a242e Mon Sep 17 00:00:00 2001 From: Tim Plummer Date: Fri, 6 Nov 2026 10:47:46 -0700 Subject: [PATCH 2/4] Address review: clearer label-matching errors, remove dead check - Label-not-found error prints labels at full precision and notes that labels are matched exactly. - Positional size-mismatch error states why the dim is positional (no coordinate on map/PSET, or epoch). - Remove unreachable value_key check (KeyError check above covers it). - Fix value_keys type in project_pset_values_to_map docstring. Co-Authored-By: Claude Opus 5.5 --- imap_processing/ena_maps/ena_maps.py | 24 +++++++++++-------- .../tests/ena_maps/test_ena_maps.py | 22 +++++++++++++++-- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/imap_processing/ena_maps/ena_maps.py b/imap_processing/ena_maps/ena_maps.py index c80206777d..9144f4f46e 100644 --- a/imap_processing/ena_maps/ena_maps.py +++ b/imap_processing/ena_maps/ena_maps.py @@ -911,7 +911,7 @@ def num_points(self) -> int: """ return self.az_el_points.shape[0] - def project_pset_values_to_map( # noqa: PLR0912 + def project_pset_values_to_map( self, pointing_set: PointingSet, value_keys: list[str] | None = None, @@ -945,7 +945,7 @@ def project_pset_values_to_map( # noqa: PLR0912 ---------- pointing_set : PointingSet The pointing set containing the values to project to the map. - value_keys : list[tuple[str, IndexMatchMethod]] | None + value_keys : list[str] | None The keys of the values in the PointingSet to project to the map. Ex.: ["counts", "flux"] data_vars named each key must be present, and of the same dimensionality in @@ -1007,9 +1007,6 @@ def project_pset_values_to_map( # noqa: PLR0912 dim_positions: dict[str, NDArray | None] = {} for value_key in value_keys: - if value_key not in pointing_set.data.data_vars: - raise ValueError(f"Value key {value_key} not found in pointing set.") - # If multiple spatial axes present # (i.e (az, el) for rectangular coordinate PSET), # stack them into a single coordinate to match the raveled indices @@ -1240,9 +1237,14 @@ def _resolve_dim_positions( ) if not label_matched: if pset_size != map_size: + reason = ( + "the epoch dim is always combined by position" + if dim == CoordNames.TIME.value + else "no coordinate on the map variable or PSET" + ) raise ValueError( - f"Cannot combine dim '{dim}' of '{value_key}' by position: " - f"PSET size {pset_size} != map size {map_size}." + f"Cannot combine dim '{dim}' of '{value_key}' by position " + f"({reason}): PSET size {pset_size} != map size {map_size}." ) return None @@ -1253,14 +1255,16 @@ def _resolve_dim_positions( if len(np.unique(pset_labels)) != pset_size: raise ValueError( f"PSET coordinate '{dim}' of '{value_key}' has duplicate labels: " - f"{pset_labels}." + f"{pset_labels.tolist()}." ) positions = map_var.indexes[dim].get_indexer(pset_labels) if np.any(positions < 0): + # Print labels with tolist() so floats show at full precision raise ValueError( f"PSET coordinate '{dim}' of '{value_key}' has labels " - f"{pset_labels[positions < 0]} not found in the map coordinate " - f"values {map_labels}." + f"{pset_labels[positions < 0].tolist()} not found in the map " + f"coordinate values {map_labels.tolist()} (labels are matched " + f"exactly)." ) return positions diff --git a/imap_processing/tests/ena_maps/test_ena_maps.py b/imap_processing/tests/ena_maps/test_ena_maps.py index c20ea1d3e1..733c518137 100644 --- a/imap_processing/tests/ena_maps/test_ena_maps.py +++ b/imap_processing/tests/ena_maps/test_ena_maps.py @@ -2387,6 +2387,24 @@ def test_label_not_in_map_raises(self, mock_frame_transform_az_el, method): with pytest.raises(ValueError, match="not found in the map coordinate"): sky_map.project_pset_values_to_map(pset, ["counts"], method) + def test_float_label_mismatch_message_shows_full_precision( + self, mock_frame_transform_az_el + ): + """Near-equal float labels raise, and the message shows the difference.""" + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + sky_map.project_pset_values_to_map(make_labeled_rect_pset(), ["counts"]) + sky_map.data_1d = sky_map.data_1d.assign_coords( + esa_energy_step=np.arange(1, 10, dtype=float) + ) + pset = make_labeled_rect_pset() + pset.data = pset.data.assign_coords( + esa_energy_step=np.arange(1, 10, dtype=float) + 1e-12 + ) + with pytest.raises(ValueError, match="matched exactly") as excinfo: + sky_map.project_pset_values_to_map(pset, ["counts"]) + assert "1.000000000001" in str(excinfo.value) + @pytest.mark.parametrize( "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] ) @@ -2400,7 +2418,7 @@ def test_size_one_dim_without_labels_raises( pset = make_labeled_rect_pset(gain_config_ids=[1]) pset.data = pset.data.drop_vars("gain_config_id") - with pytest.raises(ValueError, match="by position"): + with pytest.raises(ValueError, match="no coordinate on the map variable"): sky_map.project_pset_values_to_map(pset, ["counts"], method) def test_epoch_size_mismatch_raises(self, mock_frame_transform_az_el): @@ -2411,7 +2429,7 @@ def test_epoch_size_mismatch_raises(self, mock_frame_transform_az_el): dims=(*NON_SPATIAL_DIMS, CoordNames.GENERIC_PIXEL.value), coords={"epoch": [0, 1]}, ) - with pytest.raises(ValueError, match="dim 'epoch'"): + with pytest.raises(ValueError, match="epoch dim is always combined"): sky_map.project_pset_values_to_map(make_labeled_rect_pset(), ["counts"]) def test_missing_dim_raises(self, mock_frame_transform_az_el): From 3d745388de44c98bb76dda95c5f62efb954dffa5 Mon Sep 17 00:00:00 2001 From: Tim Plummer Date: Fri, 6 Nov 2026 11:09:13 -0700 Subject: [PATCH 3/4] Address Copilot review: keep coordinate-less dims positional, catch duplicates - When creating a map variable, copy only real PSET coordinates. Dims without a coordinate no longer get artificial 0..N-1 map labels, so they stay combined by position. - Check for duplicate PSET labels before the identical-labels shortcut, so duplicates raise even when the map's labels came from the same PSET. Co-Authored-By: Claude Opus 5.5 --- imap_processing/ena_maps/ena_maps.py | 11 +++-- .../tests/ena_maps/test_ena_maps.py | 47 ++++++++++++++++++- 2 files changed, 53 insertions(+), 5 deletions(-) diff --git a/imap_processing/ena_maps/ena_maps.py b/imap_processing/ena_maps/ena_maps.py index 9144f4f46e..535acd1475 100644 --- a/imap_processing/ena_maps/ena_maps.py +++ b/imap_processing/ena_maps/ena_maps.py @@ -1031,12 +1031,15 @@ def project_pset_values_to_map( dims=output_dims, ) - # Make coordinates for the map data array if they don't exist + # Copy the PSET's coordinates for dims the map doesn't have coords + # for yet. Dims without a PSET coordinate stay coordinate-less, so + # they keep being combined by position. self.data_1d.coords.update( { dim: pointing_set.data[dim] for dim in self.data_1d[value_key].dims if dim not in self.data_1d.coords + and dim in pointing_set.data.coords } ) @@ -1250,13 +1253,15 @@ def _resolve_dim_positions( pset_labels = pset_data.coords[dim].values map_labels = map_var.coords[dim].values - if np.array_equal(pset_labels, map_labels): - return None + # Check before the equality shortcut, so duplicates are caught even when + # the map's labels were copied from this PSET. if len(np.unique(pset_labels)) != pset_size: raise ValueError( f"PSET coordinate '{dim}' of '{value_key}' has duplicate labels: " f"{pset_labels.tolist()}." ) + if np.array_equal(pset_labels, map_labels): + return None positions = map_var.indexes[dim].get_indexer(pset_labels) if np.any(positions < 0): # Print labels with tolist() so floats show at full precision diff --git a/imap_processing/tests/ena_maps/test_ena_maps.py b/imap_processing/tests/ena_maps/test_ena_maps.py index 733c518137..b98f95a138 100644 --- a/imap_processing/tests/ena_maps/test_ena_maps.py +++ b/imap_processing/tests/ena_maps/test_ena_maps.py @@ -2441,14 +2441,57 @@ def test_missing_dim_raises(self, mock_frame_transform_az_el): with pytest.raises(ValueError, match="do not match the non-spatial dims"): sky_map.project_pset_values_to_map(pset, ["counts"]) - def test_duplicate_labels_raise(self, mock_frame_transform_az_el): + @pytest.mark.parametrize("pre_initialize", [False, True]) + def test_duplicate_labels_raise(self, mock_frame_transform_az_el, pre_initialize): + """Duplicate labels raise, including when they match the map's labels.""" self._identity(mock_frame_transform_az_el) sky_map = make_sky_map("rectangular") - pre_initialize_map_var(sky_map) + if pre_initialize: + pre_initialize_map_var(sky_map) pset = make_labeled_rect_pset(esa_steps=[2, 2]) with pytest.raises(ValueError, match="duplicate labels"): sky_map.project_pset_values_to_map(pset, ["counts"]) + def test_duplicate_labels_matching_map_raise(self, mock_frame_transform_az_el): + """Identical duplicate labels on PSET and map don't take the fast path.""" + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + sky_map.data_1d["counts"] = xr.DataArray( + np.zeros((1, 2, 2, 2, sky_map.num_points)), + dims=(*NON_SPATIAL_DIMS, CoordNames.GENERIC_PIXEL.value), + coords={"esa_energy_step": [2, 2]}, + ) + pset = make_labeled_rect_pset(esa_steps=[2, 2]) + with pytest.raises(ValueError, match="duplicate labels"): + sky_map.project_pset_values_to_map(pset, ["counts"]) + + @pytest.mark.parametrize( + "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] + ) + def test_coordinate_less_dim_stays_positional( + self, mock_frame_transform_az_el, method + ): + """A map dim created from a PSET without coords gets no artificial labels.""" + self._identity(mock_frame_transform_az_el) + pset_a = make_labeled_rect_pset(seed=1) + pset_a.data = pset_a.data.drop_vars("calibration_prod") + # Labels that differ from the 0..N-1 positions of the first PSET + pset_b = make_labeled_rect_pset(cal_prods=[20, 10], seed=2) + + sky_map = make_sky_map("rectangular") + sky_map.project_pset_values_to_map(pset_a, ["counts"], method) + assert "calibration_prod" not in sky_map.data_1d.coords + sky_map.project_pset_values_to_map(pset_b, ["counts"], method) + + single_maps = [] + for pset in (pset_a, pset_b): + single_map = make_sky_map("rectangular") + single_map.project_pset_values_to_map(pset, ["counts"], method) + single_maps.append(single_map.data_1d["counts"].values) + np.testing.assert_array_equal( + sky_map.data_1d["counts"].values, single_maps[0] + single_maps[1] + ) + def test_pre_initialized_without_pixel_dim_raises(self, mock_frame_transform_az_el): self._identity(mock_frame_transform_az_el) sky_map = make_sky_map("rectangular") From 3a44cd6f51a44bb200944f58a568151301290f80 Mon Sep 17 00:00:00 2001 From: Tim Plummer Date: Mon, 9 Nov 2026 13:45:03 -0700 Subject: [PATCH 4/4] Address PR feedback --- imap_processing/ena_maps/ena_maps.py | 18 +++++++++++++----- .../tests/ena_maps/test_ena_maps.py | 13 +++++++++++++ 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/imap_processing/ena_maps/ena_maps.py b/imap_processing/ena_maps/ena_maps.py index 535acd1475..53258d5a59 100644 --- a/imap_processing/ena_maps/ena_maps.py +++ b/imap_processing/ena_maps/ena_maps.py @@ -1156,7 +1156,9 @@ def _add_projected_values( f"Map variable '{value_key}' has dims {map_var.dims}; the last dim " f"must be the spatial pixel dim '{pixel_dim}'." ) - if set(pset_dims) != set(map_dims) or len(pset_dims) != len(map_dims): + # Sorted comparison also catches repeated dim names, which xarray allows + # at construction (with a warning). + if sorted(pset_dims) != sorted(map_dims): raise ValueError( f"Non-spatial dims of PSET variable '{value_key}' {pset_dims} do not " f"match the non-spatial dims of the map variable {map_dims}." @@ -1227,9 +1229,9 @@ def _resolve_dim_positions( Raises ------ ValueError - If a PSET coordinate label is not in the map variable's coords, if a - PSET coordinate has duplicate labels, or if a positionally combined - dimension differs in size. + If a PSET coordinate label is not in the map variable's coords, if + the PSET or map coordinate has duplicate labels, or if a positionally + combined dimension differs in size. """ pset_size = pset_data.sizes[dim] map_size = map_var.sizes[dim] @@ -1262,7 +1264,13 @@ def _resolve_dim_positions( ) if np.array_equal(pset_labels, map_labels): return None - positions = map_var.indexes[dim].get_indexer(pset_labels) + map_index = map_var.indexes[dim] + if not map_index.is_unique: + raise ValueError( + f"Map coordinate '{dim}' of '{value_key}' has duplicate labels: " + f"{map_labels.tolist()}." + ) + positions = map_index.get_indexer(pset_labels) if np.any(positions < 0): # Print labels with tolist() so floats show at full precision raise ValueError( diff --git a/imap_processing/tests/ena_maps/test_ena_maps.py b/imap_processing/tests/ena_maps/test_ena_maps.py index b98f95a138..318c2aeb68 100644 --- a/imap_processing/tests/ena_maps/test_ena_maps.py +++ b/imap_processing/tests/ena_maps/test_ena_maps.py @@ -2465,6 +2465,19 @@ def test_duplicate_labels_matching_map_raise(self, mock_frame_transform_az_el): with pytest.raises(ValueError, match="duplicate labels"): sky_map.project_pset_values_to_map(pset, ["counts"]) + def test_duplicate_map_labels_raise(self, mock_frame_transform_az_el): + """Duplicate labels on a pre-initialized map raise a clear ValueError.""" + self._identity(mock_frame_transform_az_el) + sky_map = make_sky_map("rectangular") + sky_map.data_1d["counts"] = xr.DataArray( + np.zeros((1, 2, 2, 3, sky_map.num_points)), + dims=(*NON_SPATIAL_DIMS, CoordNames.GENERIC_PIXEL.value), + coords={"esa_energy_step": [1, 1, 2]}, + ) + pset = make_labeled_rect_pset(esa_steps=[2, 1]) + with pytest.raises(ValueError, match="Map coordinate 'esa_energy_step'"): + sky_map.project_pset_values_to_map(pset, ["counts"]) + @pytest.mark.parametrize( "method", [ena_maps.IndexMatchMethod.PUSH, ena_maps.IndexMatchMethod.PULL] )