Coverage for python/lsst/images/fields/_chebyshev.py: 80%

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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("ChebyshevField", "ChebyshevFieldSerializationModel") 

15 

16from collections.abc import Iterator 

17from typing import TYPE_CHECKING, Any, ClassVar, Literal, final 

18 

19import astropy.units 

20import numpy as np 

21import pydantic 

22 

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 

28 

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] 

36 

37 

38@final 

39class ChebyshevField(BaseField): 

40 """A 2-d Chebyshev polynomial over a rectangular region. 

41 

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 """ 

56 

57 def __init__( 

58 self, bounds: Bounds, coefficients: np.ndarray, *, unit: astropy.units.UnitBase | None = None 

59 ) -> None: 

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 

72 

73 def __eq__(self, other: object) -> bool: 

74 if type(other) is not ChebyshevField: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true

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 ) 

81 

82 __hash__ = None # type: ignore[assignment] 

83 

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. 

99 

100 Parameters 

101 ---------- 

102 bounds 

103 Bounding box over which the Chebyshev field is defined. 

104 data 

105 Data points to fit. If this is an `astropy.units.Quantity`, 

106 this sets the units of the field and the ``unit`` argument cannot 

107 also be provided. 

108 order 

109 Maximum order for the Chebyshev polynomial in both dimensions. 

110 y 

111 Y coordinates of the data points. Must have either the same 

112 shape as ``data`` (providing the coordinates for all points 

113 directly), or be a 1-d array with the same size as 

114 ``data.shape[0]`` (when ``data`` is a 2-d image and ``y`` provides 

115 the coordinates of the rows). 

116 x 

117 X coordinates of the data points. Must have either the same 

118 shape as ``data`` (providing the coordinates for all points 

119 directly), or be a 1-d array with the same size as 

120 ``data.shape[1]`` (when ``data`` is a 2-d image and ``x`` provides 

121 the coordinates of the columns). 

122 weight 

123 Weights to apply to the data points. Must have the same shape as 

124 ``data``. 

125 y_order 

126 Maximum order for the Chebyshev polynomial in ``y``. Requires 

127 ``x_order`` to also be provided. Incompatible with ``order``. 

128 x_order 

129 Maximum order for the Chebyshev polynomial in ``x``. Requires 

130 ``y_order`` to also be provided. Incompatible with ``order``. 

131 triangular 

132 If `True`, only fit for coefficients of ``T_p(y) T_q(x)`` where 

133 ``p + q <= max(y_order, x_order)``. 

134 unit 

135 Units of the returned field. 

136 """ 

137 match (order, x_order, y_order): 

138 case (int(), None, None): 

139 x_order = order 

140 y_order = order 

141 case (None, int(), int()): 141 ↛ 143line 141 didn't jump to line 143 because the pattern on line 141 always matched

142 pass 

143 case _: 

144 raise TypeError("Either 'order' (only) or both 'x_order' and 'y_order' must be provided.") 

145 if weight is not None and weight.shape != data.shape: 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true

146 raise ValueError(f"Shape of 'data' {data.shape} does not match 'weight' {weight.shape}.") 

147 if isinstance(data, astropy.units.Quantity): 

148 if unit is not None: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true

149 raise TypeError("If 'data' is a Quantity, 'unit' cannot be provided separately.") 

150 unit = data.unit 

151 data = data.to_value() 

152 result = ChebyshevField(bounds, np.zeros((y_order + 1, x_order + 1), dtype=np.float64), unit=unit) 

153 if len(data.shape) == 2 and len(x.shape) == 1 and len(y.shape) == 1: 

154 if data.shape != y.shape + x.shape: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true

155 raise ValueError( 

156 f"Shape of 2-d 'data' {data.shape} does not match 1-d 'y' {y.shape} and/or 'x' {x.shape}." 

157 ) 

158 matrix = result._make_grid_matrix(x=x, y=y, triangular=triangular) 

159 else: 

160 if data.shape != y.shape: 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true

161 raise ValueError(f"Shape of 'data' {data.shape} does not match 'y' {y.shape}.") 

162 if data.shape != x.shape: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

163 raise ValueError(f"Shape of 'data' {data.shape} does not match 'x' {x.shape}.") 

164 matrix = result._make_general_matrix(x=x, y=y, triangular=triangular) 

165 if weight is not None: 

166 weight = weight.ravel() # copies only if needed 

167 matrix *= weight[:, np.newaxis] 

168 data = data.flatten() # always copies 

169 data *= weight 

170 mask = np.logical_and(weight > 0, np.isfinite(data)) 

171 else: 

172 data = data.ravel() 

173 mask = np.isfinite(data) 

174 n_good = mask.sum() 

175 if n_good == 0: 175 ↛ 176line 175 didn't jump to line 176 because the condition on line 175 was never true

176 raise ValueError("No good data points.") 

177 if n_good < data.size: 

178 data = data[mask] 

179 matrix = matrix[mask, :] 

180 packed_coefficients, *_ = np.linalg.lstsq(matrix, data) 

181 result._coefficients.flags.writeable = True 

182 for i, pq in result._packing_indices(triangular): 

183 result._coefficients[pq.y, pq.x] = packed_coefficients[i] 

184 result._coefficients.flags.writeable = False 

185 return result 

186 

187 @property 

188 def bounds(self) -> Bounds: 

189 return self._bounds 

190 

191 @property 

192 def unit(self) -> astropy.units.UnitBase | None: 

193 return self._unit 

194 

195 @property 

196 def x_order(self) -> int: 

197 """Maximum polynomial order in the column dimension (`int`).""" 

198 return self._coefficients.shape[1] - 1 

199 

200 @property 

201 def y_order(self) -> int: 

202 """Maximum polynomial order in the row dimension (`int`).""" 

203 return self._coefficients.shape[0] - 1 

204 

205 @property 

206 def order(self) -> int: 

207 """Maximum polynomial order in either dimension (`int`).""" 

208 return max(self.x_order, self.y_order) 

209 

210 @property 

211 def coefficients(self) -> np.ndarray: 

212 """Coefficients for the 2-d Chebyshev polynomial of the first kind, 

213 as a 2-d matrix in which element ``[p, q]`` corresponds to the 

214 coefficient of ``T_p(y) T_q(x)``. 

215 """ 

216 return self._coefficients 

217 

218 @property 

219 def is_constant(self) -> bool: 

220 return self.x_order == 0 and self.y_order == 0 

221 

222 def evaluate( 

223 self, *, x: np.ndarray, y: np.ndarray, quantity: bool 

224 ) -> np.ndarray | astropy.units.Quantity: 

225 m = self._remap(x=x.copy(), y=y.copy()) 

226 # We swap x and y relative to Numpy's docs because that's how our 

227 # coefficients are ordered. 

228 v = np.polynomial.chebyshev.chebval2d(m.y, m.x, self._coefficients) 

229 if quantity: 

230 return astropy.units.Quantity(v, self.unit) 

231 return v 

232 

233 def render(self, bbox: Box | None = None, *, dtype: np.typing.DTypeLike | None = None) -> Image: 

234 if bbox is None: 

235 bbox = self.bounds.bbox 

236 m = self._remap( 

237 x=bbox.x.arange.astype(np.float64), 

238 y=bbox.y.arange.astype(np.float64), 

239 ) 

240 # We swap x and y relative to Numpy's docs because that's how our 

241 # coefficients and images are ordered. 

242 v = np.polynomial.chebyshev.chebgrid2d(m.y, m.x, self._coefficients) 

243 return Image(v, bbox=bbox, unit=self.unit, dtype=dtype) 

244 

245 def multiply_constant( 

246 self, factor: float | astropy.units.Quantity | astropy.units.UnitBase 

247 ) -> ChebyshevField: 

248 factor, unit = self._handle_factor_units(factor) 

249 return ChebyshevField(self.bounds, self.coefficients * factor, unit=unit) 

250 

251 def serialize(self, archive: OutputArchive[Any]) -> ChebyshevFieldSerializationModel: 

252 """Serialize the Chebyshev field to an output archive. 

253 

254 Parameters 

255 ---------- 

256 archive 

257 Archive to write to. 

258 """ 

259 return ChebyshevFieldSerializationModel( 

260 bounds=self.bounds.serialize(), 

261 coefficients=self.coefficients, 

262 unit=self.unit, 

263 ) 

264 

265 @staticmethod 

266 def _get_archive_tree_type( 

267 pointer_type: type[Any], 

268 ) -> type[ChebyshevFieldSerializationModel]: 

269 """Return the serialization model type for this object for an archive 

270 type that uses the given pointer type. 

271 """ 

272 return ChebyshevFieldSerializationModel 

273 

274 @staticmethod 

275 def from_legacy( 

276 legacy: LegacyChebyshevBoundedField, 

277 unit: astropy.units.UnitBase | None = None, 

278 bounds: Bounds | None = None, 

279 ) -> ChebyshevField: 

280 """Convert from a legacy `lsst.afw.math.ChebyshevBoundedField`. 

281 

282 Parameters 

283 ---------- 

284 legacy 

285 Legacy field to convert. 

286 unit 

287 The units of the returned field (`lsst.afw.math.BoundedField` 

288 objects do not know their units). 

289 bounds 

290 The bounds of the returned field, if they should be different from 

291 the bounding box of ``legacy``. 

292 """ 

293 bbox = Box.from_legacy(legacy.getBBox()) 

294 if bounds is not None: 294 ↛ 295line 294 didn't jump to line 295 because the condition on line 294 was never true

295 if bounds.bbox != bbox: 

296 raise ValueError( 

297 "Custom bounds when converting a ChebyshevBoundedField must not change the bbox." 

298 ) 

299 else: 

300 bounds = bbox 

301 return ChebyshevField(bounds=bounds, coefficients=legacy.getCoefficients(), unit=unit) 

302 

303 def to_legacy(self) -> LegacyChebyshevBoundedField: 

304 """Convert to a legacy `lsst.afw.math.ChebyshevBoundedField`.""" 

305 from lsst.afw.math import ChebyshevBoundedField as LegacyChebyshevBoundedField 

306 

307 return LegacyChebyshevBoundedField(self.bounds.bbox.to_legacy(), self.coefficients) 

308 

309 @staticmethod 

310 def from_legacy_background( 

311 legacy_background: LegacyBackground, 

312 bounds: Bounds | None = None, 

313 unit: astropy.units.UnitBase | None = None, 

314 ) -> ChebyshevField: 

315 """Convert from a legacy `lsst.afw.math.BackgroundMI` instance. 

316 

317 Parameters 

318 ---------- 

319 legacy_background 

320 Legacy background object to convert. 

321 bounds 

322 The bounds of the returned field, if they should be different from 

323 the bounding box of ``legacy_background``. 

324 unit 

325 The units of the returned field (`lsst.afw.math.Background` 

326 objects do not know their units). 

327 """ 

328 from lsst.afw.math import ApproximateControl 

329 

330 approx_control = legacy_background.getBackgroundControl().getApproximateControl() 

331 stats_image = legacy_background.getStatsImage() 

332 if approx_control.getStyle() != ApproximateControl.CHEBYSHEV: 

333 raise TypeError("Legacy background does not use Chebyshev approximation.") 

334 if approx_control.getWeighting(): 

335 weight = stats_image.variance.array ** (-0.5) 

336 else: 

337 weight = None 

338 x = legacy_background.getBinCentersX() 

339 y = legacy_background.getBinCentersY() 

340 bbox = Box.from_legacy(legacy_background.getImageBBox()) 

341 if bounds is not None: 

342 if bounds.bbox != bbox: 

343 raise ValueError( 

344 "Custom bounds when converting a Chebyshev background must not change the bbox." 

345 ) 

346 else: 

347 bounds = bbox 

348 return ChebyshevField.fit( 

349 bounds, 

350 stats_image.image.array, 

351 x=x, 

352 y=y, 

353 x_order=approx_control.getOrderX(), 

354 y_order=approx_control.getOrderY(), 

355 weight=weight, 

356 unit=unit, 

357 ) 

358 

359 def _remap(self, *, x: np.ndarray, y: np.ndarray) -> YX[np.ndarray]: 

360 x -= self._xt 

361 x *= self._xs 

362 y -= self._yt 

363 y *= self._ys 

364 return YX(y=y, x=x) 

365 

366 def _packing_indices(self, triangular: bool) -> Iterator[tuple[int, YX[int]]]: 

367 i = 0 

368 for p in range(self.y_order + 1): 

369 for q in range(self.x_order + 1): 

370 if not triangular or p + q <= self.order: 

371 yield i, YX(y=p, x=q) 

372 i += 1 

373 

374 def _make_grid_matrix(self, *, x: np.ndarray, y: np.ndarray, triangular: bool) -> np.ndarray: 

375 r = self._remap( 

376 x=np.asarray(x, dtype=np.float64, copy=True), 

377 y=np.asarray(y, dtype=np.float64, copy=True), 

378 ) 

379 yv = np.polynomial.chebyshev.chebvander(r.y, self.y_order) 

380 xv = np.polynomial.chebyshev.chebvander(r.x, self.x_order) 

381 indices = list(self._packing_indices(triangular)) 

382 tensor = np.zeros(r.y.shape + r.x.shape + (len(indices),), dtype=np.float64) 

383 for i, pq in indices: 

384 tensor[:, :, i] = np.multiply.outer(yv[:, pq.y], xv[:, pq.x]) 

385 return tensor.reshape(y.shape[0] * x.shape[0], len(indices)) 

386 

387 def _make_general_matrix(self, *, x: np.ndarray, y: np.ndarray, triangular: bool) -> np.ndarray: 

388 r = self._remap( 

389 x=np.asarray(x, dtype=np.float64, copy=True).ravel(), 

390 y=np.asarray(y, dtype=np.float64, copy=True).ravel(), 

391 ) 

392 yv = np.polynomial.chebyshev.chebvander(r.y, self.y_order) 

393 xv = np.polynomial.chebyshev.chebvander(r.x, self.x_order) 

394 indices = list(self._packing_indices(triangular)) 

395 matrix = np.zeros(r.y.shape + (len(indices),), dtype=np.float64) 

396 for i, pq in indices: 

397 matrix[:, i] = yv[:, pq.y] * xv[:, pq.x] 

398 return matrix 

399 

400 

401class ChebyshevFieldSerializationModel(ArchiveTree): 

402 """Serialization model for `ChebyshevField`.""" 

403 

404 SCHEMA_NAME: ClassVar[str] = "chebyshev_field" 

405 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

406 MIN_READ_VERSION: ClassVar[int] = 1 

407 PUBLIC_TYPE: ClassVar[type] = ChebyshevField 

408 

409 bounds: SerializableBounds = pydantic.Field( 

410 description=( 

411 "The region where this field can be evaluated. " 

412 "The bbox of this region is grown by half a pixel on all sides and then used to remap " 

413 "coordinates to [-1, 1]x[-1, 1], which is the natural domain of a 2-d Chebyshev polynomial." 

414 ) 

415 ) 

416 

417 coefficients: InlineArray = pydantic.Field( 

418 description=( 

419 "Coefficients for a 2-d Chebyshev polynomial of the first kind, as a 2-d matrix in which " 

420 "element [p, q] corresponds to the coefficient of T_p(y) T_q(x)." 

421 ) 

422 ) 

423 

424 unit: Unit | None = pydantic.Field(default=None, description="Units of the field.") 

425 

426 field_type: Literal["CHEBYSHEV"] = "CHEBYSHEV" 

427 

428 def deserialize(self, archive: InputArchive, **kwargs: Any) -> ChebyshevField: 

429 """Deserialize the Chebyshev field from an input archive. 

430 

431 Parameters 

432 ---------- 

433 archive 

434 Archive to read from. 

435 **kwargs 

436 Unsupported keyword arguments are accepted only to provide 

437 better error messages (raising 

438 `.serialization.InvalidParameterError`). 

439 """ 

440 if kwargs: 440 ↛ 441line 440 didn't jump to line 441 because the condition on line 440 was never true

441 raise InvalidParameterError(f"Unrecognized parameters for ChebyshevField: {set(kwargs.keys())}.") 

442 return ChebyshevField(self.bounds.deserialize(), self.coefficients, unit=self.unit)