diff --git a/imap_processing/ena_maps/ena_maps.py b/imap_processing/ena_maps/ena_maps.py index eab9e0c3ee..53258d5a59 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, @@ -925,11 +925,27 @@ 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 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 @@ -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,15 +993,20 @@ 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." ) - 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.") + # 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 multiple spatial axes present # (i.e (az, el) for rectangular coordinate PSET), # stack them into a single coordinate to match the raveled indices @@ -990,18 +1016,30 @@ 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 + # 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 } ) @@ -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,186 @@ 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}'." + ) + # 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}." + ) + 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 + 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] + 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: + 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"({reason}): PSET size {pset_size} != map size {map_size}." + ) + return None + + pset_labels = pset_data.coords[dim].values + map_labels = map_var.coords[dim].values + # 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 + 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( + f"PSET coordinate '{dim}' of '{value_key}' has 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 + @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..318c2aeb68 100644 --- a/imap_processing/tests/ena_maps/test_ena_maps.py +++ b/imap_processing/tests/ena_maps/test_ena_maps.py @@ -2119,3 +2119,409 @@ 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) + + 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] + ) + 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="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): + 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="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): + 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"]) + + @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") + 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"]) + + 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] + ) + 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") + 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"])