Coverage for python/lsst/images/cells/_aperture_corrections.py: 54%
126 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 08:38 +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 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: 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true
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()): 221 ↛ 222line 221 didn't jump to line 222 because the condition on line 221 was never true
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()): 223 ↛ 224line 223 didn't jump to line 224 because the condition on line 223 was never true
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