Coverage for python/lsst/images/fields/_chebyshev.py: 22%
186 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-05-29 08:43 +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
24from .._geom import YX, Bounds, Box
25from .._image import Image
26from ..serialization import ArchiveTree, InlineArray, InputArchive, InvalidParameterError, 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 def __eq__(self, other: object) -> bool:
74 if type(other) is not ChebyshevField:
75 return NotImplemented
76 return (
77 self._bounds == other._bounds
78 and self._unit == other._unit
79 and np.array_equal(self._coefficients, other._coefficients, equal_nan=True)
80 )
82 __hash__ = None # type: ignore[assignment]
84 @staticmethod
85 def fit(
86 bounds: Bounds,
87 data: np.ndarray | astropy.units.Quantity,
88 order: int | None = None,
89 *,
90 y: np.ndarray,
91 x: np.ndarray,
92 weight: np.ndarray | None = None,
93 y_order: int | None = None,
94 x_order: int | None = None,
95 triangular: bool = True,
96 unit: astropy.units.UnitBase | None = None,
97 ) -> ChebyshevField:
98 """Fit a Chebyshev field to data points using linear least squares.
100 Parameters
101 ----------
102 data
103 Data points to fit. If this is an `astropy.units.Quantity`,
104 this sets the units of the field and the ``unit`` argument cannot
105 also be provided.
106 order
107 Maximum order for the Chebyshev polynomial in both dimensions.
108 y
109 Y coordinates of the data points. Must have either the same
110 shape as ``data`` (providing the coordinates for all points
111 directly), or be a 1-d array with the same size as
112 ``data.shape[0]`` (when ``data`` is a 2-d image and ``y`` provides
113 the coordinates of the rows).
114 x
115 X coordinates of the data points. Must have either the same
116 shape as ``data`` (providing the coordinates for all points
117 directly), or be a 1-d array with the same size as
118 ``data.shape[1]`` (when ``data`` is a 2-d image and ``x`` provides
119 the coordinates of the columns).
120 weight
121 Weights to apply to the data points. Must have the same shape as
122 ``data``.
123 y_order
124 Maximum order for the Chebyshev polynomial in ``y``. Requires
125 ``x_order`` to also be provided. Incompatible with ``order``.
126 x_order
127 Maximum order for the Chebyshev polynomial in ``x``. Requires
128 ``y_order`` to also be provided. Incompatible with ``order``.
129 triangular
130 If `True`, only fit for coefficients of ``T_p(y) T_q(x)`` where
131 ``p + q <= max(y_order, x_order)``.
132 unit
133 Units of the returned field.
134 """
135 match (order, x_order, y_order):
136 case (int(), None, None):
137 x_order = order
138 y_order = order
139 case (None, int(), int()):
140 pass
141 case _:
142 raise TypeError("Either 'order' (only) or both 'x_order' and 'y_order' must be provided.")
143 if weight is not None and weight.shape != data.shape:
144 raise ValueError(f"Shape of 'data' {data.shape} does not match 'weight' {weight.shape}.")
145 if isinstance(data, astropy.units.Quantity):
146 if unit is not None:
147 raise TypeError("If 'data' is a Quantity, 'unit' cannot be provided separately.")
148 unit = data.unit
149 data = data.to_value()
150 result = ChebyshevField(bounds, np.zeros((y_order + 1, x_order + 1), dtype=np.float64), unit=unit)
151 if len(data.shape) == 2 and len(x.shape) == 1 and len(y.shape) == 1:
152 if data.shape != y.shape + x.shape:
153 raise ValueError(
154 f"Shape of 2-d 'data' {data.shape} does not match 1-d 'y' {y.shape} and/or 'x' {x.shape}."
155 )
156 matrix = result._make_grid_matrix(x=x, y=y, triangular=triangular)
157 else:
158 if data.shape != y.shape:
159 raise ValueError(f"Shape of 'data' {data.shape} does not match 'y' {y.shape}.")
160 if data.shape != x.shape:
161 raise ValueError(f"Shape of 'data' {data.shape} does not match 'x' {x.shape}.")
162 matrix = result._make_general_matrix(x=x, y=y, triangular=triangular)
163 if weight is not None:
164 weight = weight.ravel() # copies only if needed
165 matrix *= weight[:, np.newaxis]
166 data = data.flatten() # always copies
167 data *= weight
168 mask = np.logical_and(weight > 0, np.isfinite(data))
169 else:
170 data = data.ravel()
171 mask = np.isfinite(data)
172 n_good = mask.sum()
173 if n_good == 0:
174 raise ValueError("No good data points.")
175 if n_good < data.size:
176 data = data[mask]
177 matrix = matrix[mask, :]
178 packed_coefficients, *_ = np.linalg.lstsq(matrix, data)
179 result._coefficients.flags.writeable = True
180 for i, pq in result._packing_indices(triangular):
181 result._coefficients[pq.y, pq.x] = packed_coefficients[i]
182 result._coefficients.flags.writeable = False
183 return result
185 @property
186 def bounds(self) -> Bounds:
187 return self._bounds
189 @property
190 def unit(self) -> astropy.units.UnitBase | None:
191 return self._unit
193 @property
194 def x_order(self) -> int:
195 """Maximum polynomial order in the column dimension (`int`)."""
196 return self._coefficients.shape[1] - 1
198 @property
199 def y_order(self) -> int:
200 """Maximum polynomial order in the row dimension (`int`)."""
201 return self._coefficients.shape[0] - 1
203 @property
204 def order(self) -> int:
205 """Maximum polynomial order in either dimension (`int`)."""
206 return max(self.x_order, self.y_order)
208 @property
209 def coefficients(self) -> np.ndarray:
210 """Coefficients for the 2-d Chebyshev polynomial of the first kind,
211 as a 2-d matrix in which element ``[p, q]`` corresponds to the
212 coefficient of ``T_p(y) T_q(x)``.
213 """
214 return self._coefficients
216 @property
217 def is_constant(self) -> bool:
218 return self.x_order == 0 and self.y_order == 0
220 def evaluate(
221 self, *, x: np.ndarray, y: np.ndarray, quantity: bool
222 ) -> np.ndarray | astropy.units.Quantity:
223 m = self._remap(x=x.copy(), y=y.copy())
224 # We swap x and y relative to Numpy's docs because that's how our
225 # coefficients are ordered.
226 v = np.polynomial.chebyshev.chebval2d(m.y, m.x, self._coefficients)
227 if quantity:
228 return astropy.units.Quantity(v, self.unit)
229 return v
231 def render(self, bbox: Box | None = None, *, dtype: np.typing.DTypeLike | None = None) -> Image:
232 if bbox is None:
233 bbox = self.bounds.bbox
234 m = self._remap(
235 x=bbox.x.arange.astype(np.float64),
236 y=bbox.y.arange.astype(np.float64),
237 )
238 # We swap x and y relative to Numpy's docs because that's how our
239 # coefficients and images are ordered.
240 v = np.polynomial.chebyshev.chebgrid2d(m.y, m.x, self._coefficients)
241 return Image(v, bbox=bbox, unit=self.unit, dtype=dtype)
243 def multiply_constant(
244 self, factor: float | astropy.units.Quantity | astropy.units.UnitBase
245 ) -> ChebyshevField:
246 factor, unit = self._handle_factor_units(factor)
247 return ChebyshevField(self.bounds, self.coefficients * factor, unit=unit)
249 def serialize(self, archive: OutputArchive[Any]) -> ChebyshevFieldSerializationModel:
250 """Serialize the Chebyshev field to an output archive."""
251 return ChebyshevFieldSerializationModel(
252 bounds=self.bounds.serialize(),
253 coefficients=self.coefficients,
254 unit=self.unit,
255 )
257 @staticmethod
258 def _get_archive_tree_type(
259 pointer_type: type[Any],
260 ) -> type[ChebyshevFieldSerializationModel]:
261 """Return the serialization model type for this object for an archive
262 type that uses the given pointer type.
263 """
264 return ChebyshevFieldSerializationModel
266 @staticmethod
267 def from_legacy(
268 legacy: LegacyChebyshevBoundedField,
269 unit: astropy.units.UnitBase | None = None,
270 bounds: Bounds | None = None,
271 ) -> ChebyshevField:
272 """Convert from a legacy `lsst.afw.math.ChebyshevBoundedField`.
274 Parameters
275 ----------
276 legacy
277 Legacy field to convert.
278 unit
279 The units of the returned field (`lsst.afw.math.BoundedField`
280 objects do not know their units).
281 bounds
282 The bounds of the returned field, if they should be different from
283 the bounding box of ``legacy``.
284 """
285 bbox = Box.from_legacy(legacy.getBBox())
286 if bounds is not None:
287 if bounds.bbox != bbox:
288 raise ValueError(
289 "Custom bounds when converting a ChebyshevBoundedField must not change the bbox."
290 )
291 else:
292 bounds = bbox
293 return ChebyshevField(bounds=bounds, coefficients=legacy.getCoefficients(), unit=unit)
295 def to_legacy(self) -> LegacyChebyshevBoundedField:
296 """Convert to a legacy `lsst.afw.math.ChebyshevBoundedField`."""
297 from lsst.afw.math import ChebyshevBoundedField as LegacyChebyshevBoundedField
299 return LegacyChebyshevBoundedField(self.bounds.bbox.to_legacy(), self.coefficients)
301 @staticmethod
302 def from_legacy_background(
303 legacy_background: LegacyBackground,
304 bounds: Bounds | None = None,
305 unit: astropy.units.UnitBase | None = None,
306 ) -> ChebyshevField:
307 """Convert from a legacy `lsst.afw.math.BackgroundMI` instance.
309 Parameters
310 ----------
311 legacy
312 Legacy background object to convert.
313 bounds
314 The bounds of the returned field, if they should be different from
315 the bounding box of ``legacy_background``.
316 unit
317 The units of the returned field (`lsst.afw.math.Background`
318 objects do not know their units).
319 """
320 from lsst.afw.math import ApproximateControl
322 approx_control = legacy_background.getBackgroundControl().getApproximateControl()
323 stats_image = legacy_background.getStatsImage()
324 if approx_control.getStyle() != ApproximateControl.CHEBYSHEV:
325 raise TypeError("Legacy background does not use Chebyshev approximation.")
326 if approx_control.getWeighting():
327 weight = stats_image.variance.array ** (-0.5)
328 else:
329 weight = None
330 x = legacy_background.getBinCentersX()
331 y = legacy_background.getBinCentersY()
332 bbox = Box.from_legacy(legacy_background.getImageBBox())
333 if bounds is not None:
334 if bounds.bbox != bbox:
335 raise ValueError(
336 "Custom bounds when converting a Chebyshev background must not change the bbox."
337 )
338 else:
339 bounds = bbox
340 return ChebyshevField.fit(
341 bounds,
342 stats_image.image.array,
343 x=x,
344 y=y,
345 x_order=approx_control.getOrderX(),
346 y_order=approx_control.getOrderY(),
347 weight=weight,
348 unit=unit,
349 )
351 def _remap(self, *, x: np.ndarray, y: np.ndarray) -> YX[np.ndarray]:
352 x -= self._xt
353 x *= self._xs
354 y -= self._yt
355 y *= self._ys
356 return YX(y=y, x=x)
358 def _packing_indices(self, triangular: bool) -> Iterator[tuple[int, YX[int]]]:
359 i = 0
360 for p in range(self.y_order + 1):
361 for q in range(self.x_order + 1):
362 if not triangular or p + q <= self.order:
363 yield i, YX(y=p, x=q)
364 i += 1
366 def _make_grid_matrix(self, *, x: np.ndarray, y: np.ndarray, triangular: bool) -> np.ndarray:
367 r = self._remap(
368 x=np.asarray(x, dtype=np.float64, copy=True),
369 y=np.asarray(y, dtype=np.float64, copy=True),
370 )
371 yv = np.polynomial.chebyshev.chebvander(r.y, self.y_order)
372 xv = np.polynomial.chebyshev.chebvander(r.x, self.x_order)
373 indices = list(self._packing_indices(triangular))
374 tensor = np.zeros(r.y.shape + r.x.shape + (len(indices),), dtype=np.float64)
375 for i, pq in indices:
376 tensor[:, :, i] = np.multiply.outer(yv[:, pq.y], xv[:, pq.x])
377 return tensor.reshape(y.shape[0] * x.shape[0], len(indices))
379 def _make_general_matrix(self, *, x: np.ndarray, y: np.ndarray, triangular: bool) -> np.ndarray:
380 r = self._remap(
381 x=np.asarray(x, dtype=np.float64, copy=True).ravel(),
382 y=np.asarray(y, dtype=np.float64, copy=True).ravel(),
383 )
384 yv = np.polynomial.chebyshev.chebvander(r.y, self.y_order)
385 xv = np.polynomial.chebyshev.chebvander(r.x, self.x_order)
386 indices = list(self._packing_indices(triangular))
387 matrix = np.zeros(r.y.shape + (len(indices),), dtype=np.float64)
388 for i, pq in indices:
389 matrix[:, i] = yv[:, pq.y] * xv[:, pq.x]
390 return matrix
393class ChebyshevFieldSerializationModel(ArchiveTree):
394 """Serialization model for `ChebyshevField`."""
396 bounds: SerializableBounds = pydantic.Field(
397 description=(
398 "The region where this field can be evaluated. "
399 "The bbox of this region is grown by half a pixel on all sides and then used to remap "
400 "coordinates to [-1, 1]x[-1, 1], which is the natural domain of a 2-d Chebyshev polynomial."
401 )
402 )
404 coefficients: InlineArray = pydantic.Field(
405 description=(
406 "Coefficients for a 2-d Chebyshev polynomial of the first kind, as a 2-d matrix in which "
407 "element [p, q] corresponds to the coefficient of T_p(y) T_q(x)."
408 )
409 )
411 unit: Unit | None = pydantic.Field(default=None, description="Units of the field.")
413 field_type: Literal["CHEBYSHEV"] = "CHEBYSHEV"
415 def deserialize(self, archive: InputArchive, **kwargs: Any) -> ChebyshevField:
416 """Deserialize the Chebyshev field from an input archive."""
417 if kwargs:
418 raise InvalidParameterError(f"Unrecognized parameters for ChebyshevField: {set(kwargs.keys())}.")
419 return ChebyshevField(self.bounds.deserialize(), self.coefficients, unit=self.unit)