Coverage for python/lsst/images/_transforms/_transform.py: 53%

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__ = ( 

15 "Transform", 

16 "TransformCompositionError", 

17 "TransformSerializationModel", 

18) 

19 

20import enum 

21import textwrap 

22from collections.abc import Iterable 

23from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, assert_type, cast, final, overload 

24 

25import astropy.io.fits.header 

26import astropy.units as u 

27import numpy as np 

28import numpy.typing as npt 

29import pydantic 

30 

31from .._concrete_bounds import BoundsSerializationModel 

32from .._geom import XY, YX, Bounds, Box 

33from ..describe import DescribableMixin, DescribeOptions, FieldRole, Report, ReportField 

34from ..serialization import ArchiveReadError, ArchiveTree, InputArchive, InvalidParameterError, OutputArchive 

35from . import _ast as astshim 

36from ._frames import Frame, SerializableFrame, SkyFrame 

37 

38if TYPE_CHECKING: 

39 try: 

40 from lsst.afw.geom import TransformPoint2ToPoint2 as LegacyTransform 

41 except ImportError: 

42 type LegacyTransform = Any # type: ignore[no-redef] 

43 

44# These pre-python-3.12 declaration are needed by Sphinx (probably the 

45# autodoc-typehints plugin. 

46I = TypeVar("I", bound=Frame) # noqa: E741 

47O = TypeVar("O", bound=Frame) # noqa: E741 

48P = TypeVar("P", bound=pydantic.BaseModel) 

49 

50 

51class TransformCompositionError(RuntimeError): 

52 """Exception raised when two transforms cannot be composed.""" 

53 

54 

55def _frame_label(frame: Frame) -> str: 

56 """Return a short name identifying a coordinate frame. 

57 

58 Most frames are pydantic models whose ``str`` spells out every field, 

59 which is far too much for a one-line summary, so the type name stands in 

60 for them. The sky frames are an enum, where the member is the name. 

61 """ 

62 if isinstance(frame, enum.Enum): 

63 return str(frame.value) 

64 return type(frame).__name__ 

65 

66 

67@final 

68class Transform[I: Frame, O: Frame](DescribableMixin): 

69 """A transform that maps two coordinate frames. 

70 

71 Parameters 

72 ---------- 

73 in_frame 

74 Input coordinate frame. 

75 out_frame 

76 Output coordinate frame. 

77 ast_mapping 

78 AST mapping that implements the transform. 

79 in_bounds 

80 Bounds of the input frame, defaulting to the input frame's 

81 bounding box. 

82 out_bounds 

83 Bounds of the output frame, defaulting to the output frame's 

84 bounding box. 

85 components 

86 Component transforms that this transform was composed from. 

87 

88 Notes 

89 ----- 

90 The `Transform` class constructor is considered a private implementation 

91 detail. Instead of using this, various factory methods are available: 

92 

93 - `from_fits_wcs` constructs a transform from a FITS WCS, as represented 

94 `astropy.wcs.WCS`; 

95 - `then` composes two transforms; 

96 - `identity` constructs a trivial transform that does nothing; 

97 - `affine` contructs an affine transform from a 2x2 or 3x3 matrix; 

98 - `inverted` returns the inverse of a transform; 

99 - `from_legacy` converts an `lsst.afw.geom.Transform` instance. 

100 

101 When applied to celestial coordinate systems, ``x=ra`` and ``y=dec``. 

102 `SkyProjection` provides a more natural interface for pixel-to-sky 

103 transforms. 

104 

105 `Transform` is conceptually immutable (the internal AST Mapping should 

106 never be modified in-place after construction), and hence does not need to 

107 be copied when any object that holds it is copied. 

108 """ 

109 

110 def __init__( 

111 self, 

112 in_frame: I, 

113 out_frame: O, 

114 ast_mapping: astshim.Mapping, 

115 in_bounds: Bounds | None = None, 

116 out_bounds: Bounds | None = None, 

117 components: Iterable[Transform[Any, Any]] = (), 

118 ) -> None: 

119 self._in_frame = in_frame 

120 self._out_frame = out_frame 

121 self._ast_mapping = ast_mapping 

122 self._in_bounds = in_bounds or getattr(in_frame, "bbox", None) 

123 self._out_bounds = out_bounds or getattr(out_frame, "bbox", None) 

124 self._components = list(components) 

125 

126 def __eq__(self, other: Any) -> bool: 

127 if self is other: 

128 # Short circuit for case where you are quickly checking 

129 # that the image WCS and variance WCS are the same object. 

130 return True 

131 if not isinstance(other, Transform): 

132 return NotImplemented 

133 if self._ast_mapping != other._ast_mapping: 

134 return False 

135 if self._in_bounds != other._in_bounds: 

136 return False 

137 if self._out_bounds != other._out_bounds: 

138 return False 

139 if self._in_frame != other._in_frame: 

140 return False 

141 if self._out_frame != other._out_frame: 

142 return False 

143 if self._components != other._components: 

144 return False 

145 return True 

146 

147 @staticmethod 

148 def from_fits_wcs( 

149 fits_wcs: astropy.wcs.WCS, 

150 in_frame: I, 

151 out_frame: O, 

152 in_bounds: Bounds | None = None, 

153 out_bounds: Bounds | None = None, 

154 x0: int = 0, 

155 y0: int = 0, 

156 ) -> Transform[I, O]: 

157 """Construct a transform from a FITS WCS. 

158 

159 Parameters 

160 ---------- 

161 fits_wcs 

162 FITS WCS to convert. 

163 in_frame 

164 Coordinate frame for input points to the forward transform. 

165 out_frame 

166 Coordinate frame for output points from the forward transform. 

167 in_bounds 

168 The region that bounds valid input points. 

169 out_bounds 

170 The region that bounds valid output points. 

171 x0 

172 Logical coordinate of the first column in the array this WCS 

173 relates to world coordinates. 

174 y0 

175 Logical coordinate of the first column in the array this WCS 

176 relates to world coordinates. 

177 

178 Notes 

179 ----- 

180 The ``x0`` and ``y0`` parameters reflect the fact that for FITS, the 

181 first row and column are always labeled ``(1, 1)``, while in Astropy 

182 and most other Python libraries, they are ``(0, 0)``. The `types` in 

183 this package (e.g. `Image`, `Mask`) allow them to be any pair of 

184 integers. 

185 

186 See Also 

187 -------- 

188 SkyProjection.from_fits_wcs 

189 """ 

190 ast_stream = astshim.StringStream(fits_wcs.to_header_string(relax=True)) 

191 ast_fits_chan = astshim.FitsChan(ast_stream, "Encoding=FITS-WCS, SipReplace=0, IWC=1") 

192 ast_frame_set = ast_fits_chan.read() 

193 _prepend_ast_shift(ast_frame_set, x=x0 - 1.0, y=y0 - 1.0, ast_domain="PIXEL") 

194 return Transform( 

195 in_frame, 

196 out_frame, 

197 ast_frame_set, 

198 in_bounds=in_bounds, 

199 out_bounds=out_bounds, 

200 ) 

201 

202 @staticmethod 

203 def identity(frame: I) -> Transform[I, I]: 

204 """Construct a trivial transform that maps a frame to itelf. 

205 

206 Parameters 

207 ---------- 

208 frame 

209 Frame used for both input and output points. 

210 """ 

211 return Transform(frame, frame, astshim.UnitMap(2)) 

212 

213 @staticmethod 

214 def affine(in_frame: I, out_frame: O, matrix: np.ndarray) -> Transform[I, O]: 

215 """Construct an affine transform from a matrix. 

216 

217 Parameters 

218 ---------- 

219 in_frame 

220 Coordinate frame for input points to the forward transform. 

221 out_frame 

222 Coordinate frame for output points from the forward transform. 

223 matrix 

224 Matrix of coefficients, either a 2x2 linear transform or a 3x3 

225 augmented affine transform, with a shift embedded in the third 

226 column and ``[0, 0, 1]`` the third row. 

227 """ 

228 if matrix.shape == (2, 2): 

229 return Transform(in_frame, out_frame, astshim.MatrixMap(matrix.copy())) 

230 elif matrix.shape == (3, 3): 230 ↛ 237line 230 didn't jump to line 237 because the condition on line 230 was always true

231 linear = astshim.MatrixMap(matrix[:2, :2].copy()) 

232 shift = astshim.ShiftMap(matrix[:2, 2]) 

233 if not np.array_equal(matrix[2, :], np.array([0.0, 0.0, 1.0])): 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true

234 raise ValueError("3x3 affine transform array must have [0, 0, 1] in its last row.") 

235 return Transform(in_frame, out_frame, linear.then(shift)) 

236 else: 

237 raise ValueError("Affine transform array must be 2x2 or 3x3.") 

238 

239 @property 

240 def in_frame(self) -> I: 

241 """Coordinate frame for input points.""" 

242 return self._in_frame 

243 

244 @property 

245 def out_frame(self) -> O: 

246 """Coordinate frame for output points.""" 

247 return self._out_frame 

248 

249 @property 

250 def in_bounds(self) -> Bounds | None: 

251 """The region that bounds valid input points (`Bounds` | `None`).""" 

252 return self._in_bounds 

253 

254 @property 

255 def out_bounds(self) -> Bounds | None: 

256 """The region that bounds valid output points (`Bounds` | `None`).""" 

257 return self._out_bounds 

258 

259 def show(self, simplified: bool = False, comments: bool = False) -> str: 

260 """Return the AST native representation of the transform. 

261 

262 Parameters 

263 ---------- 

264 simplified 

265 Whether to ask AST to simplify the mapping before showing it. 

266 This will make it much more likely that two equivalent transforms 

267 have the same `show` result. If the internal mapping is actually 

268 a frame set (as needed to round-trip legacy 

269 `lsst.afw.geom.SkyWcs` objects), this will also just show the 

270 mapping with no frame set information. 

271 comments 

272 Whether to include descriptive comments. 

273 """ 

274 ast_mapping = self._ast_mapping 

275 if simplified: 

276 if isinstance(ast_mapping, astshim.FrameSet): 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true

277 ast_mapping = ast_mapping.getMapping() 

278 ast_mapping = ast_mapping.simplified() 

279 return ast_mapping.show(comments) 

280 

281 def _describe(self, options: DescribeOptions = DescribeOptions(), /) -> Report: 

282 """Return a `Report` describing this transform. 

283 

284 Parameters 

285 ---------- 

286 options : `DescribeOptions`, optional 

287 Rendering options. `DescribeOptions.brief` reports only the input 

288 and output frames, skipping the bounds and the (potentially very 

289 long) AST mapping dump. 

290 """ 

291 fields = [ 

292 ReportField(label="in_frame", value=self.in_frame, role=FieldRole.DERIVED), 

293 ReportField(label="out_frame", value=self.out_frame, role=FieldRole.DERIVED), 

294 ] 

295 if not options.brief: 295 ↛ 305line 295 didn't jump to line 305 because the condition on line 295 was always true

296 fields += [ 

297 ReportField(label="in_bounds", value=self.in_bounds, role=FieldRole.DERIVED), 

298 ReportField(label="out_bounds", value=self.out_bounds, role=FieldRole.DERIVED), 

299 ReportField( 

300 label="mapping", 

301 value=self.show(simplified=True), 

302 role=FieldRole.DERIVED, 

303 ), 

304 ] 

305 return Report( 

306 type_name="Transform", 

307 summary=f"{_frame_label(self.in_frame)} → {_frame_label(self.out_frame)}", 

308 fields=fields, 

309 ) 

310 

311 @overload 

312 def apply_forward(self, point: XY[int | float] | YX[int | float], /) -> XY[float]: ... 

313 

314 @overload 

315 def apply_forward(self, point: XY[npt.ArrayLike] | YX[npt.ArrayLike], /) -> XY[np.ndarray]: ... 

316 

317 @overload 

318 def apply_forward(self, /, *, x: int | float, y: int | float) -> XY[float]: ... 

319 

320 @overload 

321 def apply_forward(self, /, *, x: npt.ArrayLike, y: npt.ArrayLike) -> XY[np.ndarray]: ... 

322 

323 def apply_forward( 

324 self, point: XY[Any] | YX[Any] | None = None, /, *, x: Any = None, y: Any = None 

325 ) -> XY[float] | XY[np.ndarray]: 

326 """Apply the forward transform to one or more points. 

327 

328 Parameters 

329 ---------- 

330 point 

331 An `XY` or `YX` coordinate pair to transform. Mutually exclusive 

332 with ``x`` and ``y``. 

333 x : `float` | array-like 

334 ``x`` values of the points to transform, as a scalar or any 

335 array-like. Results are broadcast against ``y``. 

336 Mutually exclusive with ``point``. 

337 y : `float` | array-like 

338 ``y`` values of the points to transform, as a scalar or any 

339 array-like. Results are broadcast against ``x``. 

340 Mutually exclusive with ``point``. 

341 

342 Returns 

343 ------- 

344 `XY` [`float` | `numpy.ndarray`] 

345 The transformed point or points. A scalar input pair returns 

346 `XY` of `float`; array-like inputs return `XY` of 

347 `numpy.ndarray` with the broadcast shape of ``x`` and ``y``. 

348 """ 

349 match point: 

350 case None: 

351 if x is None or y is None: 

352 raise TypeError("Pass either a point or both x= and y= to 'apply_forward'.") 

353 case XY() | YX(): 353 ↛ 357line 353 didn't jump to line 357 because the pattern on line 353 always matched

354 if x is not None or y is not None: 

355 raise TypeError("'apply_forward' point argument is mutually exclusive with x= and y=.") 

356 x, y = point.x, point.y 

357 case _: 

358 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.") 

359 return _standardize_xy( 

360 _ast_apply( 

361 self._ast_mapping.applyForward, 

362 x=self._in_frame.standardize_x(x), 

363 y=self._in_frame.standardize_y(y), 

364 ), 

365 self._out_frame, 

366 ) 

367 

368 @overload 

369 def apply_inverse(self, point: XY[int | float] | YX[int | float], /) -> XY[float]: ... 

370 

371 @overload 

372 def apply_inverse(self, point: XY[npt.ArrayLike] | YX[npt.ArrayLike], /) -> XY[np.ndarray]: ... 

373 

374 @overload 

375 def apply_inverse(self, /, *, x: int | float, y: int | float) -> XY[float]: ... 

376 

377 @overload 

378 def apply_inverse(self, /, *, x: npt.ArrayLike, y: npt.ArrayLike) -> XY[np.ndarray]: ... 

379 

380 def apply_inverse( 

381 self, point: XY[Any] | YX[Any] | None = None, /, *, x: Any = None, y: Any = None 

382 ) -> XY[float] | XY[np.ndarray]: 

383 """Apply the inverse transform to one or more points. 

384 

385 Parameters 

386 ---------- 

387 point 

388 An `XY` or `YX` coordinate pair to transform. Mutually exclusive 

389 with ``x`` and ``y``. 

390 x : `float` | array-like 

391 ``x`` values of the points to transform, as a scalar or any 

392 array-like. Results are broadcast against ``y``. 

393 Mutually exclusive with ``point``. 

394 y : `float` | array-like 

395 ``y`` values of the points to transform, as a scalar or any 

396 array-like. Results are broadcast against ``x``. 

397 Mutually exclusive with ``point``. 

398 

399 Returns 

400 ------- 

401 `XY` [`float` | `numpy.ndarray`] 

402 The transformed point or points. A scalar input pair returns 

403 `XY` of `float`; array-like inputs return `XY` of 

404 `numpy.ndarray` with the broadcast shape of ``x`` and ``y``. 

405 """ 

406 match point: 

407 case None: 

408 if x is None or y is None: 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true

409 raise TypeError("Pass either a point or both x= and y= to 'apply_inverse'.") 

410 case XY() | YX(): 410 ↛ 414line 410 didn't jump to line 414 because the pattern on line 410 always matched

411 if x is not None or y is not None: 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true

412 raise TypeError("'apply_inverse' point argument is mutually exclusive with x= and y=.") 

413 x, y = point.x, point.y 

414 case _: 

415 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.") 

416 return _standardize_xy( 

417 _ast_apply( 

418 self._ast_mapping.applyInverse, 

419 x=self._out_frame.standardize_x(x), 

420 y=self._out_frame.standardize_y(y), 

421 ), 

422 self._in_frame, 

423 ) 

424 

425 @overload 

426 def apply_forward_q(self, point: XY[u.Quantity] | YX[u.Quantity], /) -> XY[u.Quantity]: ... 

427 

428 @overload 

429 def apply_forward_q(self, /, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]: ... 

430 

431 def apply_forward_q( 

432 self, point: XY[u.Quantity] | YX[u.Quantity] | None = None, /, *, x: Any = None, y: Any = None 

433 ) -> XY[u.Quantity]: 

434 """Apply the forward transform to one or more unit-aware points. 

435 

436 Parameters 

437 ---------- 

438 point 

439 An `XY` or `YX` coordinate pair of `~astropy.units.Quantity` to 

440 transform. Mutually exclusive with ``x`` and ``y``. 

441 x 

442 ``x`` values of the points to transform. 

443 Mutually exclusive with ``point``. 

444 y 

445 ``y`` values of the points to transform. 

446 Mutually exclusive with ``point``. 

447 

448 Returns 

449 ------- 

450 `XY` [`astropy.units.Quantity`] 

451 The transformed point or points. 

452 """ 

453 match point: 

454 case None: 

455 if x is None or y is None: 455 ↛ 456line 455 didn't jump to line 456 because the condition on line 455 was never true

456 raise TypeError("Pass either a point or both x= and y= to 'apply_forward_q'.") 

457 case XY() | YX(): 457 ↛ 461line 457 didn't jump to line 461 because the pattern on line 457 always matched

458 if x is not None or y is not None: 458 ↛ 459line 458 didn't jump to line 459 because the condition on line 458 was never true

459 raise TypeError("'apply_forward_q' point argument is mutually exclusive with x= and y=.") 

460 x, y = point.x, point.y 

461 case _: 

462 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.") 

463 xy = self.apply_forward(x=x.to_value(self._in_frame.unit), y=y.to_value(self._in_frame.unit)) 

464 return XY(xy.x * self._out_frame.unit, xy.y * self._out_frame.unit) 

465 

466 @overload 

467 def apply_inverse_q(self, point: XY[u.Quantity] | YX[u.Quantity], /) -> XY[u.Quantity]: ... 

468 

469 @overload 

470 def apply_inverse_q(self, /, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]: ... 

471 

472 def apply_inverse_q( 

473 self, point: XY[u.Quantity] | YX[u.Quantity] | None = None, /, *, x: Any = None, y: Any = None 

474 ) -> XY[u.Quantity]: 

475 """Apply the inverse transform to one or more unit-aware points. 

476 

477 Parameters 

478 ---------- 

479 point 

480 An `XY` or `YX` coordinate pair of `~astropy.units.Quantity` to 

481 transform. Mutually exclusive with ``x`` and ``y``. 

482 x 

483 ``x`` values of the points to transform. 

484 Mutually exclusive with ``point``. 

485 y 

486 ``y`` values of the points to transform. 

487 Mutually exclusive with ``point``. 

488 

489 Returns 

490 ------- 

491 `XY` [`astropy.units.Quantity`] 

492 The transformed point or points. 

493 """ 

494 match point: 

495 case None: 

496 if x is None or y is None: 496 ↛ 497line 496 didn't jump to line 497 because the condition on line 496 was never true

497 raise TypeError("Pass either a point or both x= and y= to 'apply_inverse_q'.") 

498 case XY() | YX(): 498 ↛ 502line 498 didn't jump to line 502 because the pattern on line 498 always matched

499 if x is not None or y is not None: 499 ↛ 500line 499 didn't jump to line 500 because the condition on line 499 was never true

500 raise TypeError("'apply_inverse_q' point argument is mutually exclusive with x= and y=.") 

501 x, y = point.x, point.y 

502 case _: 

503 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.") 

504 xy = self.apply_inverse(x=x.to_value(self._out_frame.unit), y=y.to_value(self._out_frame.unit)) 

505 return XY(xy.x * self._in_frame.unit, xy.y * self._in_frame.unit) 

506 

507 def decompose(self) -> list[Transform[Any, Any]]: 

508 """Deconstruct a composed transform into its constituent parts. 

509 

510 Notes 

511 ----- 

512 Most transforms will just return a single-element list holding 

513 ``self``. Identity transform will return an empty list, and 

514 transforms composed with `then` will return the original transforms. 

515 Transforms constructed by `FrameSet` may or may not be decomposable. 

516 """ 

517 if not self._components: 517 ↛ 523line 517 didn't jump to line 523 because the condition on line 517 was always true

518 if self.in_frame == self._out_frame: 518 ↛ 521line 518 didn't jump to line 521 because the condition on line 518 was always true

519 return [] 

520 else: 

521 return [self] 

522 else: 

523 return list(self._components) 

524 

525 def inverted(self) -> Transform[O, I]: 

526 """Return the inverse of this transform.""" 

527 return Transform[O, I]( 

528 self._out_frame, 

529 self._in_frame, 

530 self._ast_mapping.inverted(), 

531 in_bounds=self.out_bounds, 

532 out_bounds=self.in_bounds, 

533 components=[t.inverted() for t in reversed(self._components)], 

534 ) 

535 

536 def then[F: Frame](self, next: Transform[O, F], remember_components: bool = True) -> Transform[I, F]: 

537 """Compose two transforms into another. 

538 

539 Parameters 

540 ---------- 

541 next 

542 Another transform to apply after ``self``. 

543 remember_components 

544 If `True`, the returned composed transform will remember ``self`` 

545 and ``other`` so they can be returned by `decompose`. 

546 """ 

547 if self._out_frame != next._in_frame: 

548 raise TransformCompositionError( 

549 "Cannot compose transforms that do not share a common intermediate frame: " 

550 f"{self._out_frame} != {next._in_frame}." 

551 ) 

552 components = self.decompose() + next.decompose() if remember_components else () 

553 return Transform( 

554 self._in_frame, 

555 next._out_frame, 

556 self._ast_mapping.then(next._ast_mapping), 

557 in_bounds=self.in_bounds, 

558 out_bounds=next.out_bounds, 

559 components=components, 

560 ) 

561 

562 def as_fits_wcs(self, bbox: Box) -> astropy.wcs.WCS | None: 

563 """Return a FITS WCS representation of this transform, if possible. 

564 

565 Parameters 

566 ---------- 

567 bbox 

568 Bounding box of the array the FITS WCS will describe. This 

569 transform object is assumed to work on the same coordinate system 

570 in which ``bbox`` is defined, while the FITS WCS will consider the 

571 first row and column in that box to be ``(0, 0)`` (in Astropy 

572 interfaces) or ``(1, 1)`` (in the FITS representation itself). 

573 

574 Notes 

575 ----- 

576 This method assumes the transform maps pixel coordinates to world 

577 coordinates. 

578 

579 Not all transforms can be represented exactly; when a FITS 

580 represention is not possible, `None` is returned. When the returned 

581 WCS is not `None`, it will have the same functional form, but it may 

582 not evaluate identically due to small implementation differences in 

583 the order of floating-point operations. 

584 """ 

585 ast_frame_set = self._get_ast_frame_set() 

586 _prepend_ast_shift(ast_frame_set, x=1.0 - bbox.x.start, y=1.0 - bbox.y.start, ast_domain="GRID") 

587 ast_stream = astshim.StringStream() 

588 ast_fits_chan = astshim.FitsChan( 

589 ast_stream, "Encoding=FITS-WCS, CDMatrix=1, FitsAxisOrder=<copy>, FitsTol=0.0001" 

590 ) 

591 ast_fits_chan.setFitsI("NAXIS1", bbox.x.size) 

592 ast_fits_chan.setFitsI("NAXIS2", bbox.y.size) 

593 n_writes = ast_fits_chan.write(ast_frame_set) 

594 if not n_writes: 

595 return None 

596 header = astropy.io.fits.Header(astropy.io.fits.Card.fromstring(c) for c in ast_fits_chan) 

597 return astropy.wcs.WCS(header) 

598 

599 def serialize[P: pydantic.BaseModel]( 

600 self, archive: OutputArchive[P], *, use_frame_sets: bool = False 

601 ) -> TransformSerializationModel[P]: 

602 """Serialize a transform to an archive. 

603 

604 Parameters 

605 ---------- 

606 archive 

607 Archive to serialize to. 

608 use_frame_sets 

609 If `True`, decompose the transform and try to reference component 

610 mappings that were already serialized into a `FrameSet` in the 

611 archive. Note that if multiple transforms exist between a pair of 

612 frames (e.g. a `SkyProjection` and its FITS approximation), this 

613 may cause the wrong one to be saved. When this option is used, the 

614 frame set must be saved before the transform, and it must be 

615 deserialized before the transform as well. 

616 

617 Returns 

618 ------- 

619 `TransformSerializationModel` 

620 Serialized form of the transform. 

621 """ 

622 model = TransformSerializationModel[P]() 

623 if use_frame_sets: 623 ↛ 624line 623 didn't jump to line 624 because the condition on line 623 was never true

624 for link in self.decompose(): 

625 model.frames.append(link.in_frame.serialize()) 

626 model.bounds.append(link.in_bounds.serialize() if link.in_bounds is not None else None) 

627 for frame_set, pointer in archive.iter_frame_sets(): 

628 if link.in_frame in frame_set and link.out_frame in frame_set: 

629 model.mappings.append(pointer) 

630 break 

631 else: 

632 model.mappings.append(MappingSerializationModel(ast=link._ast_mapping.show())) 

633 else: 

634 model.frames.append(self.in_frame.serialize()) 

635 model.bounds.append(self.in_bounds.serialize() if self.in_bounds is not None else None) 

636 model.mappings.append(MappingSerializationModel(ast=self._ast_mapping.show())) 

637 model.frames.append(self.out_frame.serialize()) 

638 model.bounds.append(self.out_bounds.serialize() if self.out_bounds is not None else None) 

639 return model 

640 

641 @staticmethod 

642 def _get_archive_tree_type[P: pydantic.BaseModel]( 

643 pointer_type: type[P], 

644 ) -> type[TransformSerializationModel[P]]: 

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

646 type that uses the given pointer type. 

647 """ 

648 return TransformSerializationModel[pointer_type] # type: ignore 

649 

650 @staticmethod 

651 def from_legacy( 

652 legacy: LegacyTransform, 

653 in_frame: I, 

654 out_frame: O, 

655 in_bounds: Bounds | None = None, 

656 out_bounds: Bounds | None = None, 

657 ) -> Transform[I, O]: 

658 """Construct a transform from a legacy `lsst.afw.geom.Transform`. 

659 

660 Parameters 

661 ---------- 

662 legacy : `lsst.afw.geom.Transform` 

663 Legacy transform object. 

664 in_frame 

665 Coordinate frame for input points to the forward transform. 

666 out_frame 

667 Coordinate frame for output points from the forward transform. 

668 in_bounds 

669 The region that bounds valid input points. 

670 out_bounds 

671 The region that bounds valid output points. 

672 """ 

673 return Transform( 

674 in_frame, 

675 out_frame, 

676 legacy.getMapping(), 

677 in_bounds=in_bounds, 

678 out_bounds=out_bounds, 

679 ) 

680 

681 def to_legacy(self) -> LegacyTransform: 

682 """Convert to a legacy `lsst.afw.geom.TransformPoint2ToPoint2` 

683 instance. 

684 """ 

685 from lsst.afw.geom import TransformPoint2ToPoint2 as LegacyTransform 

686 

687 return LegacyTransform(self._ast_mapping, False) 

688 

689 def _get_ast_frame_set(self) -> Any: 

690 ast_frame_set = astshim.FrameSet(_make_ast_frame(self._in_frame)) 

691 ast_frame_set.addFrame(astshim.FrameSet.BASE, self._ast_mapping, _make_ast_frame(self._out_frame)) 

692 return ast_frame_set 

693 

694 

695def _ast_apply(method: Any, *, x: Any, y: Any) -> XY[float] | XY[np.ndarray]: 

696 # TODO: add bounds argument and check inputs 

697 xa = np.asarray(x) 

698 ya = np.asarray(y) 

699 broadcast_shape = np.broadcast(xa, ya).shape 

700 scalar = not broadcast_shape 

701 xb, yb = np.broadcast_arrays(xa, ya) 

702 xy_in = np.vstack([xb.ravel(), yb.ravel()]).astype(np.float64) 

703 xy_out = method(xy_in) 

704 if scalar: 

705 return XY(float(xy_out[0, 0]), float(xy_out[1, 0])) 

706 return XY(xy_out[0].reshape(broadcast_shape), xy_out[1].reshape(broadcast_shape)) 

707 

708 

709def _prepend_ast_shift(ast_frame_set: Any, x: float, y: float, ast_domain: str) -> None: 

710 ast_output_frame_id = ast_frame_set.current 

711 ast_frame_set.addFrame( 

712 astshim.FrameSet.BASE, 

713 astshim.ShiftMap([x, y]), 

714 astshim.Frame(2, f"Domain={ast_domain}"), 

715 ) 

716 ast_frame_set.base = ast_frame_set.current 

717 ast_frame_set.current = ast_output_frame_id 

718 

719 

720def _make_ast_frame(frame: Frame) -> Any: 

721 if frame is SkyFrame.ICRS: 

722 return astshim.SkyFrame("") 

723 ast_frame = astshim.Frame(2, f"Ident={frame._ast_ident}") 

724 if frame.unit is not None: 724 ↛ 728line 724 didn't jump to line 728 because the condition on line 724 was always true

725 fits_unit = frame.unit.to_string(format="fits") 

726 ast_frame.setUnit(1, fits_unit) 

727 ast_frame.setUnit(2, fits_unit) 

728 ast_frame.setLabel(1, "x") 

729 ast_frame.setLabel(2, "y") 

730 return ast_frame 

731 

732 

733def _standardize_xy(xy: XY[Any], frame: Frame) -> XY[Any]: 

734 return XY(x=frame.standardize_x(xy.x), y=frame.standardize_y(xy.y)) 

735 

736 

737class MappingSerializationModel(pydantic.BaseModel): 

738 """Serialization model for an AST Mapping.""" 

739 

740 ast: str = pydantic.Field(description="A serialized Starlink AST Mapping, using the AST native encoding.") 

741 

742 

743class TransformSerializationModel[P: pydantic.BaseModel](ArchiveTree): 

744 """Serialization model for coordinate transforms.""" 

745 

746 SCHEMA_NAME: ClassVar[str] = "transform" 

747 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

748 MIN_READ_VERSION: ClassVar[int] = 1 

749 PUBLIC_TYPE: ClassVar[type] = Transform 

750 

751 frames: list[SerializableFrame] = pydantic.Field( 

752 default_factory=list, 

753 description=textwrap.dedent( 

754 """ 

755 List of frames that this transform passes through. 

756 

757 All transforms include at least two frames (the endpoints). Others 

758 intermediate frames may be included to facilitate data-sharing 

759 between transforms. 

760 """ 

761 ), 

762 ) 

763 

764 bounds: list[BoundsSerializationModel | None] = pydantic.Field( 

765 default_factory=list, 

766 description=textwrap.dedent( 

767 """ 

768 List of the bounds of the ``frames`` for this transform. 

769 

770 This always has the same number of elements as ``frames``. 

771 """ 

772 ), 

773 ) 

774 

775 mappings: list[P | MappingSerializationModel] = pydantic.Field( 

776 default_factory=list, 

777 description=textwrap.dedent( 

778 """ 

779 The actual mappings between frames, or archive pointers to 

780 serialized FrameSet objects from which they can be obtained. 

781 

782 This always has one fewer element than ``frames``. 

783 """ 

784 ), 

785 ) 

786 

787 def deserialize(self, archive: InputArchive[P], **kwargs: Any) -> Transform[Any, Any]: 

788 """Deserialize a transform from an archive. 

789 

790 Parameters 

791 ---------- 

792 archive 

793 Archive to read from. 

794 **kwargs 

795 Unsupported keyword arguments are accepted only to provide better 

796 error messages (raising `serialization.InvalidParameterError`). 

797 """ 

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

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

800 if len(self.frames) != len(self.bounds): 800 ↛ 801line 800 didn't jump to line 801 because the condition on line 800 was never true

801 raise ArchiveReadError( 

802 f"Inconsistent lengths for 'frames' ({len(self.frames)}) and 'bounds' ({len(self.bounds)})." 

803 ) 

804 if len(self.frames) != len(self.mappings) + 1: 804 ↛ 805line 804 didn't jump to line 805 because the condition on line 804 was never true

805 raise ArchiveReadError( 

806 f"Inconsistent lengths for 'frames' ({len(self.frames)}) and " 

807 f"'mappings' ({len(self.mappings)}; should be one less)." 

808 ) 

809 # We can't just compose onto an identity Transform if we want to 

810 # preserve the FrameSet-ness of any of these mappings. 

811 transform: Transform | None = None 

812 for n, mapping in enumerate(self.mappings): 

813 match mapping: 

814 case MappingSerializationModel(ast=serialized_mapping): 814 ↛ 825line 814 didn't jump to line 825 because the pattern on line 814 always matched