Coverage for python/lsst/images/cells/_aperture_corrections.py: 25%
126 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-10 16:39 +0000
1# This file is part of lsst-images.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = ("CellApertureCorrectionMapSerializationModel", "CellField")
16from collections.abc import Mapping
17from typing import TYPE_CHECKING, Any, ClassVar, final
19import astropy.table
20import astropy.units
21import numpy as np
22import pydantic
24from .._cell_grid import CellGridBounds, CellIJ
25from .._geom import BoundsError, Box
26from .._image import Image
27from ..fields import BaseField
28from ..serialization import (
29 ArchiveReadError,
30 ArchiveTree,
31 InputArchive,
32 InvalidParameterError,
33 OutputArchive,
34 TableModel,
35)
37if TYPE_CHECKING:
38 try:
39 from lsst.afw.image import ApCorrMap as LegacyApCorrMap
40 from lsst.cell_coadds import StitchedApertureCorrection as LegacyStichedApertureCorrection
41 except ImportError:
42 type LegacyApCorrMap = Any # type: ignore[no-redef]
43 type LegacyStichedApertureCorrection = Any # type: ignore[no-redef]
46@final
47class CellField(BaseField):
48 """A piecewise 2-d function on a cell-coadd grid.
50 Parameters
51 ----------
52 array
53 A 2-d array of cell values with shape
54 ``bounds.subgrid_size.as_tuple()``.
55 bounds
56 Description of the cell grid and any missing cells. Array entries for
57 missing cells should be NaN.
59 Notes
60 -----
61 `CellField` is not directly serializable and is not included in the
62 ``Field`` union type alias as a result. A `~collections.abc.Mapping` of
63 `CellField` is instead serializable via
64 `CellApertureCorrectionMapSerializationModel`.
65 """
67 def __init__(
68 self, bounds: CellGridBounds, array: np.ndarray, unit: astropy.units.UnitBase | None = None
69 ) -> None:
70 self._array = array
71 self._bounds = bounds
72 self._unit = unit
73 if self._array.shape != self._bounds.subgrid_size.as_tuple():
74 raise ValueError(
75 f"Array shape ({self._array.shape}) differs from subgrid size ({self._bounds.subgrid_size})."
76 )
78 __hash__ = None # type: ignore[assignment]
80 @property
81 def bounds(self) -> CellGridBounds:
82 return self._bounds
84 @property
85 def unit(self) -> astropy.units.UnitBase | None:
86 return self._unit
88 @property
89 def is_constant(self) -> bool:
90 indices = iter(self._bounds.cell_indices())
91 try:
92 first = self.value_in_cell(next(indices))
93 except StopIteration:
94 return True
95 for other_index in indices:
96 if self.value_in_cell(other_index) != first:
97 return False
98 return True
100 def value_in_cell(self, key: CellIJ) -> float:
101 """Return the value of the field in the cell with the given index."""
102 if key in self._bounds.missing:
103 raise BoundsError(f"Cell {key} is missing for this field.")
104 index = key - self._bounds.subgrid_start
105 try:
106 return self._array[index.i, index.j]
107 except IndexError:
108 raise BoundsError(f"Cell {key} is out of bounds for this field.") from None
110 def quantity_in_cell(self, key: CellIJ) -> astropy.units.Quantity:
111 """Return the quantity (value with units) of the field in the cell
112 with the given index.
113 """
114 return astropy.units.Quantity(self.value_in_cell(key), self._unit)
116 def evaluate(
117 self, *, x: np.ndarray, y: np.ndarray, quantity: bool
118 ) -> np.ndarray | astropy.units.Quantity:
119 # This implementation is optimized for the case where there are many
120 # more evaluation points than cells. We could switch to an
121 # implementation that zip-broadcast-iterates over x and y when that is
122 # not the case, but that feels like a premature optimization right now.
123 result = np.full(np.broadcast_shapes(y.shape, x.shape), np.nan, dtype=np.float64)
124 for cell_index in self._bounds.cell_indices():
125 cell_bbox = self._bounds.grid.bbox_of(cell_index)
126 result[cell_bbox.contains(x=x, y=y)] = self.value_in_cell(cell_index)
127 if quantity:
128 return astropy.units.Quantity(result, self._unit)
129 return result
131 def render(self, bbox: Box | None = None, *, dtype: np.typing.DTypeLike | None = None) -> Image:
132 if bbox is None:
133 bbox = self._bounds.bbox
134 bounds = self._bounds
135 else:
136 bounds = self._bounds[bbox]
137 result = Image(np.nan, bbox=bbox, dtype=dtype, unit=self._unit)
138 for cell_index in bounds.cell_indices():
139 cell_bbox = self._bounds.grid.bbox_of(cell_index).intersection(bbox)
140 result[cell_bbox].array = self.value_in_cell(cell_index)
141 return result
143 def multiply_constant(self, factor: float | astropy.units.Quantity | astropy.units.UnitBase) -> CellField:
144 factor, unit = self._handle_factor_units(factor)
145 return CellField(self._bounds, self._array * factor, unit=unit)
147 @staticmethod
148 def from_legacy_aperture_correction(
149 legacy: LegacyStichedApertureCorrection, bounds: CellGridBounds
150 ) -> CellField:
151 """Convert from a legacy `lsst.cell_coadds.StitchedApertureCorrection`.
153 Parameters
154 ----------
155 legacy
156 Legacy field to convert.
157 bounds
158 The grid and bounds of the returned field.
159 """
160 array = np.full(bounds.subgrid_size.as_tuple(), np.nan, dtype=np.float64)
161 for cell_index in bounds.cell_indices():
162 array_index = cell_index - bounds.subgrid_start
163 array[array_index.i, array_index.j] = legacy.gc[cell_index.to_legacy()]
164 return CellField(bounds, array)
166 def to_legacy_aperture_correction(self) -> LegacyStichedApertureCorrection:
167 """Convert to a legacy
168 `lsst.cell_coadds.StitchedApertureCorrection`.
169 """
170 from lsst.cell_coadds import GridContainer, StitchedApertureCorrection
172 grid = self.bounds.grid.to_legacy()
173 gc = GridContainer[float](grid.shape)
174 for cell_index in self.bounds.cell_indices():
175 gc[cell_index.to_legacy()] = self.value_in_cell(cell_index)
176 return StitchedApertureCorrection(grid, gc)
179class CellApertureCorrectionMapSerializationModel(ArchiveTree):
180 """A serialization model for a `~collections.abc.Mapping` of `CellField`,
181 which is used to represent aperture corrections for cell-based coadds.
182 """
184 SCHEMA_NAME: ClassVar[str] = "cell_aperture_correction_map"
185 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
186 MIN_READ_VERSION: ClassVar[int] = 1
187 PUBLIC_TYPE: ClassVar[type] = dict
189 table: TableModel = pydantic.Field(
190 description="Table with one row for each cell and different photometry algorithms in columns."
191 )
192 bounds: CellGridBounds = pydantic.Field(
193 description=(
194 "Description of the cell grid and any missing cells. Array entries for "
195 "missing cells should be NaN."
196 ),
197 )
199 @staticmethod
200 def serialize(
201 aperture_correction_map: Mapping[str, CellField], archive: OutputArchive[Any]
202 ) -> CellApertureCorrectionMapSerializationModel | None:
203 if not aperture_correction_map:
204 return None
205 bounds = next(iter(aperture_correction_map.values())).bounds
206 if not all(field.bounds == bounds for field in aperture_correction_map.values()):
207 raise ValueError("Cell aperture corrections do not have consistent bounds.")
208 if any(field.unit is not None for field in aperture_correction_map.values()):
209 raise ValueError("Aperture corrections should be dimensionless.")
210 table = astropy.table.Table(
211 rows=[cell_index.as_tuple() for cell_index in bounds.cell_indices()], names=["cell_i", "cell_j"]
212 )
213 good_cell_mask = np.ones(bounds.subgrid_size.as_tuple(), dtype=bool)
214 for cell_index in bounds.missing:
215 array_index = cell_index - bounds.subgrid_start
216 good_cell_mask[array_index.i, array_index.j] = False
217 for name, field in aperture_correction_map.items():
218 table.add_column(field._array[good_cell_mask], name=name, copy=False)
219 return CellApertureCorrectionMapSerializationModel(
220 table=archive.add_table(table, name="table"), bounds=bounds
221 )
223 def deserialize(self, archive: InputArchive[Any], **kwargs: Any) -> dict[str, CellField]:
224 if kwargs:
225 raise InvalidParameterError(
226 f"Unrecognized parameters for cell aperture correction map: {set(kwargs.keys())}."
227 )
228 good_cell_mask = np.zeros(self.bounds.subgrid_size.as_tuple(), dtype=bool)
229 table = archive.get_table(self.table)
230 for tbl_ij, cell_index in zip(
231 table["cell_i", "cell_j"].iterrows(), self.bounds.cell_indices(), strict=True
232 ):
233 if cell_index.as_tuple() != tbl_ij:
234 raise ArchiveReadError(
235 "Inconsistency between serialized aperture correction bounds and table."
236 )
237 array_index = cell_index - self.bounds.subgrid_start
238 good_cell_mask[array_index.i, array_index.j] = True
239 result: dict[str, CellField] = {}
240 for name, column in table.columns.items():
241 if name in ("cell_i", "cell_j"):
242 continue
243 array = np.full(self.bounds.subgrid_size.as_tuple(), np.nan, dtype=np.float64)
244 array[good_cell_mask] = column
245 result[name] = CellField(self.bounds, array)
246 return result