Coverage for python / lsst / images / fields / _chebyshev.py: 24%
169 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-07 08:34 +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__ = ("ChebyshevField", "ChebyshevFieldSerializationModel")
16from collections.abc import Iterator
17from typing import TYPE_CHECKING, Any, Literal, final
19import astropy.units
20import numpy as np
21import pydantic
23from .._concrete_bounds import SerializableBounds, deserialize_bounds
24from .._geom import YX, Bounds, Box
25from .._image import Image
26from ..serialization import ArchiveTree, InlineArray, InputArchive, OutputArchive, Unit
27from ._base import BaseField
29if TYPE_CHECKING:
30 try:
31 from lsst.afw.math import BackgroundMI as LegacyBackground
32 from lsst.afw.math import ChebyshevBoundedField as LegacyChebyshevBoundedField
33 except ImportError:
34 type LegacyBackground = Any # type: ignore[no-redef]
35 type LegacyChebyshevBoundedField = Any # type: ignore[no-redef]
38@final
39class ChebyshevField(BaseField):
40 """A 2-d Chebyshev polynomial over a rectangular region.
42 Parameters
43 ----------
44 bounds
45 The region where this field can be evaluated. The ``bbox`` of this
46 region is grown by half a pixel on all sides and then used to remap
47 coordinates to ``[-1, 1]x[-1, 1]``, which is the natural domain of a
48 2-d Chebyshev polynomial.
49 coefficients
50 Coefficients for the 2-d Chebyshev polynomial of the first kind, as a
51 2-d matrix in which element ``[p, q]`` corresponds to the coefficient
52 of ``T_p(y) T_q(x)``. Will be set to read-only in place.
53 unit
54 Units of the field.
55 """
57 def __init__(
58 self, bounds: Bounds, coefficients: np.ndarray, *, unit: astropy.units.UnitBase | None = None
59 ):
60 self._bounds = bounds
61 self._coefficients = coefficients
62 self._coefficients.flags.writeable = False
63 self._unit = unit
64 # Compute the scaling and translation that map points in the bbox
65 # (including an extra 0.5 on all sides, since the bbox is int-based)
66 # to [-1, 1].
67 bbox = bounds.bbox
68 self._xs = 2.0 / bbox.x.size
69 self._xt = bbox.x.min + 0.5 * bbox.x.size - 0.5
70 self._ys = 2.0 / bbox.y.size
71 self._yt = bbox.y.min + 0.5 * bbox.y.size - 0.5
73 @staticmethod
74 def fit(
75 bounds: Bounds,
76 data: np.ndarray | astropy.units.Quantity,
77 order: int | None = None,
78 *,
79 y: np.ndarray,
80 x: np.ndarray,
81 weight: np.ndarray | None = None,
82 y_order: int | None = None,
83 x_order: int | None = None,
84 triangular: bool = True,
85 unit: astropy.units.UnitBase | None = None,
86 ) -> ChebyshevField:
87 """Fit a Chebyshev field to data points using linear least squares.
89 Parameters
90 ----------
91 data
92 Data points to fit. If this is an `astropy.units.Quantity`,
93 this sets the units of the field and the ``unit`` argument cannot
94 also be provided.
95 order
96 Maximum order for the Chebyshev polynomial in both dimensions.
97 y
98 Y coordinates of the data points. Must have either the same
99 shape as ``data`` (providing the coordinates for all points
100 directly), or be a 1-d array with the same size as
101 ``data.shape[0]`` (when ``data`` is a 2-d image and ``y`` provides
102 the coordinates of the rows).
103 x
104 X coordinates of the data points. Must have either the same
105 shape as ``data`` (providing the coordinates for all points
106 directly), or be a 1-d array with the same size as
107 ``data.shape[1]`` (when ``data`` is a 2-d image and ``x`` provides
108 the coordinates of the columns).
109 weight
110 Weights to apply to the data points. Must have the same shape as
111 ``data``.
112 y_order
113 Maximum order for the Chebyshev polynomial in ``y``. Requires
114 ``x_order`` to also be provided. Incompatible with ``order``.
115 x_order
116 Maximum order for the Chebyshev polynomial in ``x``. Requires
117 ``y_order`` to also be provided. Incompatible with ``order``.
118 triangular
119 If `True`, only fit for coefficients of ``T_p(y) T_q(x)`` where
120 ``p + q <= max(y_order, x_order)``.
121 unit
122 Units of the returned field.
123 """
124 match (order, x_order, y_order):
125 case (int(), None, None):
126 x_order = order
127 y_order = order
128 case (None, int(), int()):
129 pass
130 case _:
131 raise TypeError("Either 'order' (only) or both 'x_order' and 'y_order' must be provided.")
132 if weight is not None and weight.shape != data.shape:
133 raise ValueError(f"Shape of 'data' {data.shape} does not match 'weight' {weight.shape}.")
134 if isinstance(data, astropy.units.Quantity):
135 if unit is not None:
136 raise TypeError("If 'data' is a Quantity, 'unit' cannot be provided separately.")
137 unit = data.unit
138 data = data.to_value()
139 result = ChebyshevField(bounds, np.zeros((y_order + 1, x_order + 1), dtype=np.float64), unit=unit)
140 if len(data.shape) == 2 and len(x.shape) == 1 and len(y.shape) == 1:
141 if data.shape != y.shape + x.shape:
142 raise ValueError(
143 f"Shape of 2-d 'data' {data.shape} does not match 1-d 'y' {y.shape} and/or 'x' {x.shape}."
144 )
145 matrix = result._make_grid_matrix(x=x, y=y, triangular=triangular)
146 else:
147 if data.shape != y.shape:
148 raise ValueError(f"Shape of 'data' {data.shape} does not match 'y' {y.shape}.")
149 if data.shape != x.shape:
150 raise ValueError(f"Shape of 'data' {data.shape} does not match 'x' {x.shape}.")
151 matrix = result._make_general_matrix(x=x, y=y, triangular=triangular)
152 if weight is not None:
153 weight = weight.ravel() # copies only if needed
154 matrix *= weight[:, np.newaxis]
155 data = data.flatten() # always copies
156 data *= weight
157 mask = np.logical_and(weight > 0, np.isfinite(data))
158 else:
159 data = data.ravel()
160 mask = np.isfinite(data)
161 n_good = mask.sum()
162 if n_good == 0:
163 raise ValueError("No good data points.")
164 if n_good < data.size:
165 data = data[mask]
166 matrix = matrix[mask, :]
167 packed_coefficients, *_ = np.linalg.lstsq(matrix, data)
168 result._coefficients.flags.writeable = True
169 for i, pq in result._packing_indices(triangular):
170 result._coefficients[pq.y, pq.x] = packed_coefficients[i]
171 result._coefficients.flags.writeable = False
172 return result
174 @property
175 def bounds(self) -> Bounds:
176 return self._bounds
178 @property
179 def unit(self) -> astropy.units.UnitBase | None:
180 return self._unit
182 @property
183 def x_order(self) -> int:
184 """Maximum polynomial order in the column dimension (`int`)."""
185 return self._coefficients.shape[1] - 1
187 @property
188 def y_order(self) -> int:
189 """Maximum polynomial order in the row dimension (`int`)."""
190 return self._coefficients.shape[0] - 1
192 @property
193 def order(self) -> int:
194 """Maximum polynomial order in either dimension (`int`)."""
195 return max(self.x_order, self.y_order)
197 @property
198 def coefficients(self) -> np.ndarray:
199 """Coefficients for the 2-d Chebyshev polynomial of the first kind,
200 as a 2-d matrix in which element ``[p, q]`` corresponds to the
201 coefficient of ``T_p(y) T_q(x)``.
202 """
203 return self._coefficients
205 def evaluate(
206 self, *, x: np.ndarray, y: np.ndarray, quantity: bool
207 ) -> np.ndarray | astropy.units.Quantity:
208 m = self._remap(x=x.copy(), y=y.copy())
209 # We swap x and y relative to Numpy's docs because that's how our
210 # coefficients are ordered.
211 v = np.polynomial.chebyshev.chebval2d(m.y, m.x, self._coefficients)
212 if quantity:
213 return astropy.units.Quantity(v, self.unit)
214 return v
216 def render(self, bbox: Box | None = None, *, dtype: np.typing.DTypeLike | None = None) -> Image:
217 if bbox is None:
218 bbox = self.bounds.bbox
219 m = self._remap(
220 x=bbox.x.arange.astype(np.float64),
221 y=bbox.y.arange.astype(np.float64),
222 )
223 # We swap x and y relative to Numpy's docs because that's how our
224 # coefficients and images are ordered.
225 v = np.polynomial.chebyshev.chebgrid2d(m.y, m.x, self._coefficients)
226 return Image(v, bbox=bbox, unit=self.unit, dtype=dtype)
228 def multiply_constant(
229 self, factor: float | astropy.units.Quantity | astropy.units.UnitBase
230 ) -> ChebyshevField:
231 factor, unit = self._handle_factor_units(factor)
232 return ChebyshevField(self.bounds, self.coefficients * factor, unit=unit)
234 def serialize(self, archive: OutputArchive[Any]) -> ChebyshevFieldSerializationModel:
235 """Serialize the Chebyshev field to an output archive."""
236 return ChebyshevFieldSerializationModel(
237 bounds=self.bounds.serialize(),
238 coefficients=self.coefficients,
239 unit=self.unit,
240 )
242 @staticmethod
243 def deserialize(model: ChebyshevFieldSerializationModel, archive: InputArchive[Any]) -> ChebyshevField:
244 """Deserialize the Chebyshev field from an input archive."""
245 return ChebyshevField(deserialize_bounds(model.bounds), model.coefficients, unit=model.unit)
247 @staticmethod
248 def _get_archive_tree_type(
249 pointer_type: type[Any],
250 ) -> type[ChebyshevFieldSerializationModel]:
251 """Return the serialization model type for this object for an archive
252 type that uses the given pointer type.
253 """
254 return ChebyshevFieldSerializationModel
256 @staticmethod
257 def from_legacy(
258 legacy: LegacyChebyshevBoundedField, unit: astropy.units.UnitBase | None = None
259 ) -> ChebyshevField:
260 """Convert from a legacy `lsst.afw.math.ChebyshevBoundedField`."""
261 return ChebyshevField(
262 bounds=Box.from_legacy(legacy.getBBox()), coefficients=legacy.getCoefficients(), unit=unit
263 )
265 def to_legacy(self) -> LegacyChebyshevBoundedField:
266 """Convert to a legacy `lsst.afw.math.ChebyshevBoundedField`."""
267 from lsst.afw.math import ChebyshevBoundedField as LegacyChebyshevBoundedField
269 return LegacyChebyshevBoundedField(self.bounds.bbox.to_legacy(), self.coefficients)
271 @staticmethod
272 def from_legacy_background(
273 legacy_background: LegacyBackground,
274 unit: astropy.units.UnitBase | None = None,
275 ) -> ChebyshevField:
276 """Convert from a legacy `lsst.afw.math.BackgroundMI` instance."""
277 from lsst.afw.math import ApproximateControl
279 approx_control = legacy_background.getBackgroundControl().getApproximateControl()
280 stats_image = legacy_background.getStatsImage()
281 if approx_control.getStyle() != ApproximateControl.CHEBYSHEV:
282 raise TypeError("Legacy background does not use Chebyshev approximation.")
283 if approx_control.getWeighting():
284 weight = stats_image.variance.array ** (-0.5)
285 else:
286 weight = None
287 x = legacy_background.getBinCentersX()
288 y = legacy_background.getBinCentersY()
289 return ChebyshevField.fit(
290 Box.from_legacy(legacy_background.getImageBBox()),
291 stats_image.image.array,
292 x=x,
293 y=y,
294 x_order=approx_control.getOrderX(),
295 y_order=approx_control.getOrderY(),
296 weight=weight,
297 unit=unit,
298 )
300 def _remap(self, *, x: np.ndarray, y: np.ndarray) -> YX[np.ndarray]:
301 x -= self._xt
302 x *= self._xs
303 y -= self._yt
304 y *= self._ys
305 return YX(y=y, x=x)
307 def _packing_indices(self, triangular: bool) -> Iterator[tuple[int, YX[int]]]:
308 i = 0
309 for p in range(self.y_order + 1):
310 for q in range(self.x_order + 1):
311 if not triangular or p + q <= self.order:
312 yield i, YX(y=p, x=q)
313 i += 1
315 def _make_grid_matrix(self, *, x: np.ndarray, y: np.ndarray, triangular: bool) -> np.ndarray:
316 r = self._remap(
317 x=np.asarray(x, dtype=np.float64, copy=True),
318 y=np.asarray(y, dtype=np.float64, copy=True),
319 )
320 yv = np.polynomial.chebyshev.chebvander(r.y, self.y_order)
321 xv = np.polynomial.chebyshev.chebvander(r.x, self.x_order)
322 indices = list(self._packing_indices(triangular))
323 tensor = np.zeros(r.y.shape + r.x.shape + (len(indices),), dtype=np.float64)
324 for i, pq in indices:
325 tensor[:, :, i] = np.multiply.outer(yv[:, pq.y], xv[:, pq.x])
326 return tensor.reshape(y.shape[0] * x.shape[0], len(indices))
328 def _make_general_matrix(self, *, x: np.ndarray, y: np.ndarray, triangular: bool) -> np.ndarray:
329 r = self._remap(
330 x=np.asarray(x, dtype=np.float64, copy=True).ravel(),
331 y=np.asarray(y, dtype=np.float64, copy=True).ravel(),
332 )
333 yv = np.polynomial.chebyshev.chebvander(r.y, self.y_order)
334 xv = np.polynomial.chebyshev.chebvander(r.x, self.x_order)
335 indices = list(self._packing_indices(triangular))
336 matrix = np.zeros(r.y.shape + (len(indices),), dtype=np.float64)
337 for i, pq in indices:
338 matrix[:, i] = yv[:, pq.y] * xv[:, pq.x]
339 return matrix
342class ChebyshevFieldSerializationModel(ArchiveTree):
343 """Serialization model for `ChebyshevField`."""
345 bounds: SerializableBounds = pydantic.Field(
346 description=(
347 "The region where this field can be evaluated. "
348 "The bbox of this region is grown by half a pixel on all sides and then used to remap "
349 "coordinates to [-1, 1]x[-1, 1], which is the natural domain of a 2-d Chebyshev polynomial."
350 )
351 )
353 coefficients: InlineArray = pydantic.Field(
354 description=(
355 "Coefficients for a 2-d Chebyshev polynomial of the first kind, as a 2-d matrix in which "
356 "element [p, q] corresponds to the coefficient of T_p(y) T_q(x)."
357 )
358 )
360 unit: Unit | None = pydantic.Field(default=None, description="Units of the field.")
362 field_type: Literal["CHEBYSHEV"] = "CHEBYSHEV"
364 def finish_deserialize(self, archive: InputArchive) -> ChebyshevField:
365 return ChebyshevField.deserialize(self, archive)