Coverage for python/lsst/images/fields/

1# This file is part of lsst-images.

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.

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, ClassVar, Literal, final

19import astropy.units

20import numpy as np

21import pydantic

23from .._concrete_bounds import BoundsSerializationModel

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 Chebyshev1Function2D as LegacyChebyshev1Function2D

33 from lsst.afw.math import ChebyshevBoundedField as LegacyChebyshevBoundedField

34 except ImportError:

35 type LegacyBackground = Any # type: ignore[no-redef]

36 type LegacyChebyshevBoundedField = Any # type: ignore[no-redef]

37 type LegacyChebyshev1Function2D = Any # type: ignore[no-redef]

40@final

41class ChebyshevField(BaseField):

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

44 Parameters

45 ----------

46 bounds

47 The region where this field can be evaluated. The ``bbox`` of this

48 region is grown by half a pixel on all sides and then used to remap

49 coordinates to ``[-1, 1]x[-1, 1]``, which is the natural domain of a

50 2-d Chebyshev polynomial.

51 coefficients

52 Coefficients for the 2-d Chebyshev polynomial of the first kind, as a

53 2-d matrix in which element ``[p, q]`` corresponds to the coefficient

54 of ``T_p(y) T_q(x)``. Will be set to read-only in place.

55 unit

56 Units of the field.

57 """

59 def __init__(

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

61 ) -> None:

62 self._bounds = bounds

63 self._coefficients = coefficients

64 self._coefficients.flags.writeable = False

65 self._unit = unit

66 # Compute the scaling and translation that map points in the bbox

67 # (including an extra 0.5 on all sides, since the bbox is int-based)

68 # to [-1, 1].

69 bbox = bounds.bbox

70 self._xs = 2.0 / bbox.x.size

71 self._xt = bbox.x.min + 0.5 * bbox.x.size - 0.5

72 self._ys = 2.0 / bbox.y.size

73 self._yt = bbox.y.min + 0.5 * bbox.y.size - 0.5

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

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

77 return NotImplemented

78 return (

79 self._bounds == other._bounds

80 and self._unit == other._unit

81 and np.array_equal(self._coefficients, other._coefficients, equal_nan=True)

82 )

84 __hash__ = None # type: ignore[assignment]

86 @staticmethod

87 def fit(

88 bounds: Bounds,

89 data: np.ndarray | astropy.units.Quantity,

90 order: int | None = None,

91 *,

92 y: np.ndarray,

93 x: np.ndarray,

94 weight: np.ndarray | None = None,

95 y_order: int | None = None,

96 x_order: int | None = None,

97 triangular: bool = True,

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

99 ) -> ChebyshevField:

100 """Fit a Chebyshev field to data points using linear least squares.

101

102 Parameters

103 ----------

104 bounds

105 Bounding box over which the Chebyshev field is defined.

106 data

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

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

109 also be provided.

110 order

111 Maximum order for the Chebyshev polynomial in both dimensions.

112 y

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

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

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

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

117 the coordinates of the rows).

118 x

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

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

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

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

123 the coordinates of the columns).

124 weight

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

126 ``data``.

127 y_order

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

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

130 x_order

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

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

133 triangular

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

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

136 unit

137 Units of the returned field.

138 """

139 match (order, x_order, y_order):

140 case (int(), None, None):

141 x_order = order

142 y_order = order

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

144 pass

145 case _:

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

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

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

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

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

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

152 unit = data.unit

153 data = data.to_value()

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

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

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

157 raise ValueError(

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

159 )

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

161 else:

162 if data.shape != y.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 'y' {y.shape}.")

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

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

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

167 if weight is not None:

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

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

170 data = data.flatten() # always copies

171 data *= weight

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

173 else:

174 data = data.ravel()

175 mask = np.isfinite(data)

176 n_good = mask.sum()

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

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

179 if n_good < data.size:

180 data = data[mask]

181 matrix = matrix[mask, :]

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

183 result._coefficients.flags.writeable = True

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

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

186 result._coefficients.flags.writeable = False

187 return result

188

189 @property

190 def bounds(self) -> Bounds:

191 return self._bounds

192

193 @property

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

195 return self._unit

196

197 @property

198 def x_order(self) -> int:

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

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

201

202 @property

203 def y_order(self) -> int:

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

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

206

207 @property

208 def order(self) -> int:

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

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

211

212 @property

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

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

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

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

217 """

218 return self._coefficients

219

220 @property

221 def is_constant(self) -> bool:

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

223

224 def _evaluate(

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

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

227 y, x = np.broadcast_arrays(y, x)

228 m = self._remap(x=x.astype(np.float64, copy=True), y=y.astype(np.float64, copy=True))

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

230 # coefficients are ordered.

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

232 if quantity:

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

234 return v

235

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

237 if bbox is None:

238 bbox = self.bounds.bbox

239 m = self._remap(

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

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

242 )

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

244 # coefficients and images are ordered.

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

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

247

248 def _multiply_constant(

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

250 ) -> ChebyshevField:

251 factor, unit = self._handle_factor_units(factor)

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

253

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

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

256

257 Parameters

258 ----------

259 archive

260 Archive to write to.

261 """

262 return ChebyshevFieldSerializationModel(

263 bounds=self.bounds.serialize(),

264 coefficients=self.coefficients,

265 unit=self.unit,

266 )

267

268 @staticmethod

269 def _get_archive_tree_type(

270 pointer_type: type[Any],

271 ) -> type[ChebyshevFieldSerializationModel]:

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

273 type that uses the given pointer type.

274 """

275 return ChebyshevFieldSerializationModel

276

277 @staticmethod

278 def from_legacy(

279 legacy: LegacyChebyshevBoundedField,

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

281 bounds: Bounds | None = None,

282 ) -> ChebyshevField:

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

284

285 Parameters

286 ----------

287 legacy

288 Legacy field to convert.

289 unit

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

291 objects do not know their units).

292 bounds

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

294 the bounding box of ``legacy``.

295 """

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

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

298 if bounds.bbox != bbox:

299 raise ValueError(

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

301 )

302 else:

303 bounds = bbox

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

305

306 def to_legacy(self) -> LegacyChebyshevBoundedField:

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

308 from lsst.afw.math import ChebyshevBoundedField as LegacyChebyshevBoundedField

309

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

311

312 @staticmethod

313 def from_legacy_background(

314 legacy_background: LegacyBackground,

315 bounds: Bounds | None = None,

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

317 ) -> ChebyshevField:

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

319

320 Parameters

321 ----------

322 legacy_background

323 Legacy background object to convert.

324 bounds

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

326 the bounding box of ``legacy_background``.

327 unit

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

329 objects do not know their units).

330 """

331 from lsst.afw.math import ApproximateControl

332

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

334 stats_image = legacy_background.getStatsImage()

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

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

337 if approx_control.getWeighting():

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

339 else:

340 weight = None

341 x = legacy_background.getBinCentersX()

342 y = legacy_background.getBinCentersY()

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

344 if bounds is not None:

345 if bounds.bbox != bbox:

346 raise ValueError(

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

348 )

349 else:

350 bounds = bbox

351 return ChebyshevField.fit(

352 bounds,

353 stats_image.image.array,

354 x=x,

355 y=y,

356 x_order=approx_control.getOrderX(),

357 y_order=approx_control.getOrderY(),

358 weight=weight,

359 unit=unit,

360 )

361

362 @staticmethod

363 def from_legacy_function2(

364 legacy_function2: LegacyChebyshev1Function2D,

365 bounds: Bounds | None = None,

366 unit: astropy.units.Unit | None = None,

367 ) -> ChebyshevField:

368 """Convert from a legacy `lsst.afw.math.Chebyshev1Function2D`.

369

370 Parameters

371 ----------

372 legacy_function2

373 Legacy function object to convert.

374 bounds

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

376 the bounding box of ``legacy_background``.

377 unit

378 The units of the returned field.

379 """

380 xy_range = legacy_function2.getXYRange()

381 bbox = Box.factory[

382 _require_int(xy_range.y.min + 0.5) : _require_int(xy_range.y.max + 0.5),

383 _require_int(xy_range.x.min + 0.5) : _require_int(xy_range.x.max + 0.5),

384 ]

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

386 if bounds.bbox != bbox:

387 raise ValueError(

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

389 )

390 else:

391 bounds = bbox

392 order = legacy_function2.getOrder()

393 coefficients = np.zeros((order + 1, order + 1), dtype=np.float64)

394 for i, pq in ChebyshevField._legacy_function2_indices(order):

395 coefficients[pq.y, pq.x] = legacy_function2.getParameter(i)

396 return ChebyshevField(bbox, coefficients, unit=unit)

397

398 def to_legacy_function2(self) -> LegacyChebyshev1Function2D:

399 """Convert to a legacy `lsst.afw.math.Chebyshev1Function2D`."""

400 from lsst.afw.math import Chebyshev1Function2D as LegacyChebyshev1Function2D

401 from lsst.geom import Box2D as LegacyBox2D

402

403 order = max(self.y_order, self.x_order)

404 result = LegacyChebyshev1Function2D(order, LegacyBox2D(self.bounds.bbox.to_legacy()))

405 for i, pq in self._legacy_function2_indices(order):

406 result.setParameter(

407 i,

408 (

409 self._coefficients[pq.y, pq.x]

410 if pq.y < self._coefficients.shape[0] and pq.x < self._coefficients.shape[1]

411 else 0.0

412 ),

413 )

414 return result

415

416 @staticmethod

417 def _legacy_function2_indices(order: int) -> Iterator[tuple[int, YX[int]]]:

418 i = 0

419 for n in range(order + 1):

420 for p in range(0, n + 1):

421 q = n - p

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

423 i += 1

424

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

426 x -= self._xt

427 x *= self._xs

428 y -= self._yt

429 y *= self._ys

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

431

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

433 i = 0

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

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

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

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

438 i += 1

439

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

441 r = self._remap(

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

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

444 )

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

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

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

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

449 for i, pq in indices:

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

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

452

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

454 r = self._remap(

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

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

457 )

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

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

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

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

462 for i, pq in indices:

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

464 return matrix

465

466

467class ChebyshevFieldSerializationModel(ArchiveTree):

468 """Serialization model for `ChebyshevField`."""

469

470 SCHEMA_NAME: ClassVar[str] = "chebyshev_field"

471 SCHEMA_VERSION: ClassVar[str] = "1.0.0"

472 MIN_READ_VERSION: ClassVar[int] = 1

473 PUBLIC_TYPE: ClassVar[type] = ChebyshevField

474

475 bounds: BoundsSerializationModel = pydantic.Field(

476 description=(

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

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

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

480 )

481 )

482

483 coefficients: InlineArray = pydantic.Field(

484 description=(

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

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

487 )

488 )

489

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

491

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

493

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

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

496

497 Parameters

498 ----------

499 archive

500 Archive to read from.

501 **kwargs

502 Unsupported keyword arguments are accepted only to provide

503 better error messages (raising

504 `.serialization.InvalidParameterError`).

505 """

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

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

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

509

510

511def _require_int(v: float) -> int:

512 if (z := int(v)) == v: 512 ↛ 514line 512 didn't jump to line 514 because the condition on line 512 was always true

513 return z

514 raise ValueError("Legacy Chebyshev1Function2 XY range must be at half-integer positions.")

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

224 statements