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