Coverage for python/lsst/images/cells/_aperture_corrections.py: 42%
126 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-16 15:24 -0700
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 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: LegacyStichedApertureCorrection, 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 for cell_index in bounds.cell_indices():
177 array_index = cell_index - bounds.subgrid_start
178 array[array_index.i, array_index.j] = legacy.gc[cell_index.to_legacy()]
179 return CellField(bounds, array)
181 def to_legacy_aperture_correction(self) -> LegacyStichedApertureCorrection:
182 """Convert to a legacy
183 `lsst.cell_coadds.StitchedApertureCorrection`.
184 """
185 from lsst.cell_coadds import GridContainer, StitchedApertureCorrection
187 grid = self.bounds.grid.to_legacy()
188 gc = GridContainer[float](grid.shape)
189 for cell_index in self.bounds.cell_indices():
190 gc[cell_index.to_legacy()] = self.value_in_cell(cell_index)
191 return StitchedApertureCorrection(grid, gc)
194class CellApertureCorrectionMapSerializationModel(ArchiveTree):
195 """A serialization model for a `~collections.abc.Mapping` of `CellField`,
196 which is used to represent aperture corrections for cell-based coadds.
197 """
199 SCHEMA_NAME: ClassVar[str] = "cell_aperture_correction_map"
200 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
201 MIN_READ_VERSION: ClassVar[int] = 1
202 PUBLIC_TYPE: ClassVar[type] = dict
204 table: TableModel = pydantic.Field(
205 description="Table with one row for each cell and different photometry algorithms in columns."
206 )
207 bounds: CellGridBounds = pydantic.Field(
208 description=(
209 "Description of the cell grid and any missing cells. Array entries for "
210 "missing cells should be NaN."
211 ),
212 )
214 @staticmethod
215 def serialize(
216 aperture_correction_map: Mapping[str, CellField], archive: OutputArchive[Any]
217 ) -> CellApertureCorrectionMapSerializationModel | None:
218 if not aperture_correction_map:
219 return None
220 bounds = next(iter(aperture_correction_map.values())).bounds
221 if not all(field.bounds == bounds for field in aperture_correction_map.values()):
222 raise ValueError("Cell aperture corrections do not have consistent bounds.")
223 if any(field.unit is not None for field in aperture_correction_map.values()):
224 raise ValueError("Aperture corrections should be dimensionless.")
225 table = astropy.table.Table(
226 rows=[cell_index.as_tuple() for cell_index in bounds.cell_indices()], names=["cell_i", "cell_j"]
227 )
228 good_cell_mask = np.ones(bounds.subgrid_size.as_tuple(), dtype=bool)
229 for cell_index in bounds.missing:
230 array_index = cell_index - bounds.subgrid_start
231 good_cell_mask[array_index.i, array_index.j] = False
232 for name, field in aperture_correction_map.items():
233 table.add_column(field._array[good_cell_mask], name=name, copy=False)
234 return CellApertureCorrectionMapSerializationModel(
235 table=archive.add_table(table, name="table"), bounds=bounds
236 )
238 def deserialize(self, archive: InputArchive[Any], **kwargs: Any) -> dict[str, CellField]:
239 if kwargs: 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true
240 raise InvalidParameterError(
241 f"Unrecognized parameters for cell aperture correction map: {set(kwargs.keys())}."
242 )
243 good_cell_mask = np.zeros(self.bounds.subgrid_size.as_tuple(), dtype=bool)
244 table = archive.get_table(self.table)
245 for tbl_ij, cell_index in zip(
246 table["cell_i", "cell_j"].iterrows(), self.bounds.cell_indices(), strict=True
247 ):
248 if cell_index.as_tuple() != tbl_ij: 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true
249 raise ArchiveReadError(
250 "Inconsistency between serialized aperture correction bounds and table."
251 )
252 array_index = cell_index - self.bounds.subgrid_start
253 good_cell_mask[array_index.i, array_index.j] = True
254 result: dict[str, CellField] = {}
255 for name, column in table.columns.items():
256 if name in ("cell_i", "cell_j"):
257 continue
258 array = np.full(self.bounds.subgrid_size.as_tuple(), np.nan, dtype=np.float64)
259 array[good_cell_mask] = column
260 result[name] = CellField(self.bounds, array)
261 return result