Coverage for python/lsst/images/cells/_aperture_corrections.py: 69%
144 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-14 09:21 +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 LegacyStitchedApertureCorrection
41 except ImportError:
42 type LegacyApCorrMap = Any # type: ignore[no-redef]
43 type LegacyStitchedApertureCorrection = 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 bounds
53 Description of the cell grid and any missing cells. Array entries for
54 missing cells should be NaN.
55 array
56 A 2-d array of cell values with shape
57 ``bounds.subgrid_size.as_tuple()``.
58 unit
59 Units of the field values, or `None` if dimensionless.
61 Notes
62 -----
63 `CellField` is not directly serializable and is not included in the
64 ``Field`` union type alias as a result. A `~collections.abc.Mapping` of
65 `CellField` is instead serializable via
66 `CellApertureCorrectionMapSerializationModel`.
67 """
69 def __init__(
70 self, bounds: CellGridBounds, array: np.ndarray, unit: astropy.units.UnitBase | None = None
71 ) -> None:
72 self._array = array
73 self._bounds = bounds
74 self._unit = unit
75 if self._array.shape != self._bounds.subgrid_size.as_tuple(): 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true
76 raise ValueError(
77 f"Array shape ({self._array.shape}) differs from subgrid size ({self._bounds.subgrid_size})."
78 )
80 __hash__ = None # type: ignore[assignment]
82 @property
83 def bounds(self) -> CellGridBounds:
84 return self._bounds
86 @property
87 def unit(self) -> astropy.units.UnitBase | None:
88 return self._unit
90 @property
91 def is_constant(self) -> bool:
92 indices = iter(self._bounds.cell_indices())
93 try:
94 first = self.value_in_cell(next(indices))
95 except StopIteration:
96 return True
97 for other_index in indices:
98 if self.value_in_cell(other_index) != first:
99 return False
100 return True
102 def value_in_cell(self, key: CellIJ) -> float:
103 """Return the value of the field in the cell with the given index.
105 Parameters
106 ----------
107 key
108 Index of the cell to evaluate.
109 """
110 if key in self._bounds.missing:
111 raise BoundsError(f"Cell {key} is missing for this field.")
112 index = key - self._bounds.subgrid_start
113 try:
114 return self._array[index.i, index.j]
115 except IndexError:
116 raise BoundsError(f"Cell {key} is out of bounds for this field.") from None
118 def quantity_in_cell(self, key: CellIJ) -> astropy.units.Quantity:
119 """Return the quantity (value with units) of the field in the cell
120 with the given index.
122 Parameters
123 ----------
124 key
125 Index of the cell to evaluate.
126 """
127 return astropy.units.Quantity(self.value_in_cell(key), self._unit)
129 def _evaluate(
130 self, *, x: np.ndarray, y: np.ndarray, quantity: bool
131 ) -> np.ndarray | astropy.units.Quantity:
132 # This implementation is optimized for the case where there are many
133 # more evaluation points than cells. We could switch to an
134 # implementation that zip-broadcast-iterates over x and y when that is
135 # not the case, but that feels like a premature optimization right now.
136 result = np.full(np.broadcast_shapes(y.shape, x.shape), np.nan, dtype=np.float64)
137 for cell_index in self._bounds.cell_indices():
138 cell_bbox = self._bounds.grid.bbox_of(cell_index)
139 result[cell_bbox.contains(x=x, y=y)] = self.value_in_cell(cell_index)
140 if quantity:
141 return astropy.units.Quantity(result, self._unit)
142 return result
144 def render(self, bbox: Box | None = None, *, dtype: np.typing.DTypeLike | None = None) -> Image:
145 if bbox is None:
146 bbox = self._bounds.bbox
147 bounds = self._bounds
148 else:
149 bounds = self._bounds[bbox]
150 result = Image(np.nan, bbox=bbox, dtype=dtype, unit=self._unit)
151 for cell_index in bounds.cell_indices():
152 cell_bbox = self._bounds.grid.bbox_of(cell_index).intersection(bbox)
153 result[cell_bbox].array = self.value_in_cell(cell_index)
154 return result
156 def _multiply_constant(
157 self, factor: float | astropy.units.Quantity | astropy.units.UnitBase
158 ) -> CellField:
159 factor, unit = self._handle_factor_units(factor)
160 return CellField(self._bounds, self._array * factor, unit=unit)
162 @staticmethod
163 def from_legacy_aperture_correction(
164 legacy: LegacyStitchedApertureCorrection, bounds: CellGridBounds
165 ) -> CellField:
166 """Convert from a legacy `lsst.cell_coadds.StitchedApertureCorrection`.
168 Parameters
169 ----------
170 legacy
171 Legacy field to convert.
172 bounds
173 The grid and bounds of the returned field.
174 """
175 array = np.full(bounds.subgrid_size.as_tuple(), np.nan, dtype=np.float64)
176 new_missing: set[CellIJ] = set()
177 for cell_index in bounds.cell_indices():
178 array_index = cell_index - bounds.subgrid_start
179 try:
180 value = legacy.gc[cell_index.to_legacy()]
181 except KeyError:
182 value = np.nan
183 new_missing.add(cell_index)
184 array[array_index.i, array_index.j] = value
185 if new_missing:
186 bounds = CellGridBounds(
187 grid=bounds.grid, missing=frozenset(bounds.missing | new_missing), bbox=bounds.bbox
188 )
189 return CellField(bounds, array)
191 def to_legacy_aperture_correction(self) -> LegacyStitchedApertureCorrection:
192 """Convert to a legacy
193 `lsst.cell_coadds.StitchedApertureCorrection`.
194 """
195 from lsst.cell_coadds import GridContainer, StitchedApertureCorrection
197 grid = self.bounds.grid.to_legacy()
198 gc = GridContainer[float](grid.shape)
199 for cell_index in self.bounds.cell_indices():
200 gc[cell_index.to_legacy()] = self.value_in_cell(cell_index)
201 return StitchedApertureCorrection(grid, gc)
204class CellApertureCorrectionMapSerializationModel(ArchiveTree):
205 """A serialization model for a `~collections.abc.Mapping` of `CellField`,
206 which is used to represent aperture corrections for cell-based coadds.
207 """
209 SCHEMA_NAME: ClassVar[str] = "cell_aperture_correction_map"
210 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
211 MIN_READ_VERSION: ClassVar[int] = 1
212 PUBLIC_TYPE: ClassVar[type] = dict
214 table: TableModel = pydantic.Field(
215 description="Table with one row for each cell and different photometry algorithms in columns."
216 )
217 bounds: CellGridBounds = pydantic.Field(
218 description=(
219 "Description of the cell grid and any missing cells. Array entries for "
220 "missing cells should be NaN."
221 ),
222 )
224 @staticmethod
225 def serialize(
226 aperture_correction_map: Mapping[str, CellField], archive: OutputArchive[Any]
227 ) -> CellApertureCorrectionMapSerializationModel | None:
228 if not aperture_correction_map: 228 ↛ 229line 228 didn't jump to line 229 because the condition on line 228 was never true
229 return None
230 bounds = next(iter(aperture_correction_map.values())).bounds
231 if not all(field.bounds == bounds for field in aperture_correction_map.values()):
232 # In rare cases some field can have more missing cells from the
233 # others. This isn't worth a full denormalization into full
234 # per-field bounds (especially since that's a schema change).
235 # Instead, the shared, serialized 'bounds' only marks a cell as
236 # missing if it's absent from all fields (as is the case when
237 # there's no data), and we restore the per-field missing cells in
238 # deserialized.
239 always_missing: set[CellIJ] = set(bounds.missing)
240 for field in aperture_correction_map.values():
241 if field.bounds.grid != bounds.grid:
242 raise ValueError("Cell aperture corrections do not have a consistent grid.")
243 always_missing &= field.bounds.missing
244 bounds = CellGridBounds(grid=bounds.grid, missing=frozenset(always_missing), bbox=bounds.bbox)
245 if any(field.unit is not None for field in aperture_correction_map.values()): 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true
246 raise ValueError("Aperture corrections should be dimensionless.")
247 table = astropy.table.Table(
248 rows=[cell_index.as_tuple() for cell_index in bounds.cell_indices()], names=["cell_i", "cell_j"]
249 )
250 good_cell_mask = np.ones(bounds.subgrid_size.as_tuple(), dtype=bool)
251 for cell_index in bounds.missing:
252 array_index = cell_index - bounds.subgrid_start
253 good_cell_mask[array_index.i, array_index.j] = False
254 for name, field in aperture_correction_map.items():
255 table.add_column(field._array[good_cell_mask], name=name, copy=False)
256 return CellApertureCorrectionMapSerializationModel(
257 table=archive.add_table(table, name="table"), bounds=bounds
258 )
260 def deserialize(self, archive: InputArchive[Any], **kwargs: Any) -> dict[str, CellField]:
261 if kwargs: 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true
262 raise InvalidParameterError(
263 f"Unrecognized parameters for cell aperture correction map: {set(kwargs.keys())}."
264 )
265 good_cell_mask = np.zeros(self.bounds.subgrid_size.as_tuple(), dtype=bool)
266 table = archive.get_table(self.table)
267 for tbl_ij, cell_index in zip(
268 table["cell_i", "cell_j"].iterrows(), self.bounds.cell_indices(), strict=True
269 ):
270 if cell_index.as_tuple() != tbl_ij: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 raise ArchiveReadError(
272 "Inconsistency between serialized aperture correction bounds and table."
273 )
274 array_index = cell_index - self.bounds.subgrid_start
275 good_cell_mask[array_index.i, array_index.j] = True
276 result: dict[str, CellField] = {}
277 for name, column in table.columns.items():
278 if name in ("cell_i", "cell_j"):
279 continue
280 extra_missing_mask = np.isnan(column)
281 # If there are any NaN entries in this column, make a custom
282 # CellGridBounds with additional missing cells for this field.
283 bounds = self.bounds
284 if extra_missing_mask.any():
285 extra_missing = frozenset(
286 {CellIJ(i=i, j=j) for i, j in table["cell_i", "cell_j"][extra_missing_mask].iterrows()}
287 )
288 bounds = CellGridBounds(
289 grid=self.bounds.grid, missing=self.bounds.missing | extra_missing, bbox=self.bounds.bbox
290 )
291 array = np.full(self.bounds.subgrid_size.as_tuple(), np.nan, dtype=np.float64)
292 array[good_cell_mask] = column
293 result[name] = CellField(bounds, array)
294 return result