Coverage for python/lsst/images/_transforms/_transform.py: 78%
206 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-04 17:17 +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__ = (
15 "Transform",
16 "TransformCompositionError",
17 "TransformSerializationModel",
18)
20import textwrap
21from collections.abc import Iterable
22from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, final
24import astropy.io.fits.header
25import astropy.units as u
26import numpy as np
27import pydantic
29from .._concrete_bounds import BoundsSerializationModel
30from .._geom import XY, Bounds, Box
31from ..serialization import ArchiveReadError, ArchiveTree, InputArchive, InvalidParameterError, OutputArchive
32from . import _ast as astshim
33from ._frames import Frame, SerializableFrame, SkyFrame
35if TYPE_CHECKING:
36 try:
37 from lsst.afw.geom import TransformPoint2ToPoint2 as LegacyTransform
38 except ImportError:
39 type LegacyTransform = Any # type: ignore[no-redef]
41# These pre-python-3.12 declaration are needed by Sphinx (probably the
42# autodoc-typehints plugin.
43I = TypeVar("I", bound=Frame) # noqa: E741
44O = TypeVar("O", bound=Frame) # noqa: E741
45P = TypeVar("P", bound=pydantic.BaseModel)
48class TransformCompositionError(RuntimeError):
49 """Exception raised when two transforms cannot be composed."""
52@final
53class Transform[I: Frame, O: Frame]:
54 """A transform that maps two coordinate frames.
56 Parameters
57 ----------
58 in_frame
59 Input coordinate frame.
60 out_frame
61 Output coordinate frame.
62 ast_mapping
63 AST mapping that implements the transform.
64 in_bounds
65 Bounds of the input frame, defaulting to the input frame's
66 bounding box.
67 out_bounds
68 Bounds of the output frame, defaulting to the output frame's
69 bounding box.
70 components
71 Component transforms that this transform was composed from.
73 Notes
74 -----
75 The `Transform` class constructor is considered a private implementation
76 detail. Instead of using this, various factory methods are available:
78 - `from_fits_wcs` constructs a transform from a FITS WCS, as represented
79 `astropy.wcs.WCS`;
80 - `then` composes two transforms;
81 - `identity` constructs a trivial transform that does nothing;
82 - `affine` contructs an affine transform from a 2x2 or 3x3 matrix;
83 - `inverted` returns the inverse of a transform;
84 - `from_legacy` converts an `lsst.afw.geom.Transform` instance.
86 When applied to celestial coordinate systems, ``x=ra`` and ``y=dec``.
87 `SkyProjection` provides a more natural interface for pixel-to-sky
88 transforms.
90 `Transform` is conceptually immutable (the internal AST Mapping should
91 never be modified in-place after construction), and hence does not need to
92 be copied when any object that holds it is copied.
93 """
95 def __init__(
96 self,
97 in_frame: I,
98 out_frame: O,
99 ast_mapping: astshim.Mapping,
100 in_bounds: Bounds | None = None,
101 out_bounds: Bounds | None = None,
102 components: Iterable[Transform[Any, Any]] = (),
103 ) -> None:
104 self._in_frame = in_frame
105 self._out_frame = out_frame
106 self._ast_mapping = ast_mapping
107 self._in_bounds = in_bounds or getattr(in_frame, "bbox", None)
108 self._out_bounds = out_bounds or getattr(out_frame, "bbox", None)
109 self._components = list(components)
111 def __eq__(self, other: Any) -> bool:
112 if self is other:
113 # Short circuit for case where you are quickly checking
114 # that the image WCS and variance WCS are the same object.
115 return True
116 if not isinstance(other, Transform):
117 return NotImplemented
118 if self._ast_mapping != other._ast_mapping:
119 return False
120 if self._in_bounds != other._in_bounds:
121 return False
122 if self._out_bounds != other._out_bounds:
123 return False
124 if self._in_frame != other._in_frame:
125 return False
126 if self._out_frame != other._out_frame:
127 return False
128 if self._components != other._components:
129 return False
130 return True
132 @staticmethod
133 def from_fits_wcs(
134 fits_wcs: astropy.wcs.WCS,
135 in_frame: I,
136 out_frame: O,
137 in_bounds: Bounds | None = None,
138 out_bounds: Bounds | None = None,
139 x0: int = 0,
140 y0: int = 0,
141 ) -> Transform[I, O]:
142 """Construct a transform from a FITS WCS.
144 Parameters
145 ----------
146 fits_wcs
147 FITS WCS to convert.
148 in_frame
149 Coordinate frame for input points to the forward transform.
150 out_frame
151 Coordinate frame for output points from the forward transform.
152 in_bounds
153 The region that bounds valid input points.
154 out_bounds
155 The region that bounds valid output points.
156 x0
157 Logical coordinate of the first column in the array this WCS
158 relates to world coordinates.
159 y0
160 Logical coordinate of the first column in the array this WCS
161 relates to world coordinates.
163 Notes
164 -----
165 The ``x0`` and ``y0`` parameters reflect the fact that for FITS, the
166 first row and column are always labeled ``(1, 1)``, while in Astropy
167 and most other Python libraries, they are ``(0, 0)``. The `types` in
168 this package (e.g. `Image`, `Mask`) allow them to be any pair of
169 integers.
171 See Also
172 --------
173 SkyProjection.from_fits_wcs
174 """
175 ast_stream = astshim.StringStream(fits_wcs.to_header_string(relax=True))
176 ast_fits_chan = astshim.FitsChan(ast_stream, "Encoding=FITS-WCS, SipReplace=0, IWC=1")
177 ast_frame_set = ast_fits_chan.read()
178 _prepend_ast_shift(ast_frame_set, x=x0 - 1.0, y=y0 - 1.0, ast_domain="PIXEL")
179 return Transform(
180 in_frame,
181 out_frame,
182 ast_frame_set,
183 in_bounds=in_bounds,
184 out_bounds=out_bounds,
185 )
187 @staticmethod
188 def identity(frame: I) -> Transform[I, I]:
189 """Construct a trivial transform that maps a frame to itelf.
191 Parameters
192 ----------
193 frame
194 Frame used for both input and output points.
195 """
196 return Transform(frame, frame, astshim.UnitMap(2))
198 @staticmethod
199 def affine(in_frame: I, out_frame: O, matrix: np.ndarray) -> Transform[I, O]:
200 """Construct an affine transform from a matrix.
202 Parameters
203 ----------
204 in_frame
205 Coordinate frame for input points to the forward transform.
206 out_frame
207 Coordinate frame for output points from the forward transform.
208 matrix
209 Matrix of coefficients, either a 2x2 linear transform or a 3x3
210 augmented affine transform, with a shift embedded in the third
211 column and ``[0, 0, 1]`` the third row.
212 """
213 if matrix.shape == (2, 2):
214 return Transform(in_frame, out_frame, astshim.MatrixMap(matrix.copy()))
215 elif matrix.shape == (3, 3): 215 ↛ 222line 215 didn't jump to line 222 because the condition on line 215 was always true
216 linear = astshim.MatrixMap(matrix[:2, :2].copy())
217 shift = astshim.ShiftMap(matrix[:2, 2])
218 if not np.array_equal(matrix[2, :], np.array([0.0, 0.0, 1.0])): 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true
219 raise ValueError("3x3 affine transform array must have [0, 0, 1] in its last row.")
220 return Transform(in_frame, out_frame, linear.then(shift))
221 else:
222 raise ValueError("Affine transform array must be 2x2 or 3x3.")
224 @property
225 def in_frame(self) -> I:
226 """Coordinate frame for input points."""
227 return self._in_frame
229 @property
230 def out_frame(self) -> O:
231 """Coordinate frame for output points."""
232 return self._out_frame
234 @property
235 def in_bounds(self) -> Bounds | None:
236 """The region that bounds valid input points (`Bounds` | `None`)."""
237 return self._in_bounds
239 @property
240 def out_bounds(self) -> Bounds | None:
241 """The region that bounds valid output points (`Bounds` | `None`)."""
242 return self._out_bounds
244 def show(self, simplified: bool = False, comments: bool = False) -> str:
245 """Return the AST native representation of the transform.
247 Parameters
248 ----------
249 simplified
250 Whether to ask AST to simplify the mapping before showing it.
251 This will make it much more likely that two equivalent transforms
252 have the same `show` result. If the internal mapping is actually
253 a frame set (as needed to round-trip legacy
254 `lsst.afw.geom.SkyWcs` objects), this will also just show the
255 mapping with no frame set information.
256 comments
257 Whether to include descriptive comments.
258 """
259 ast_mapping = self._ast_mapping
260 if simplified:
261 if isinstance(ast_mapping, astshim.FrameSet):
262 ast_mapping = ast_mapping.getMapping()
263 ast_mapping = ast_mapping.simplified()
264 return ast_mapping.show(comments)
266 def apply_forward[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
267 """Apply the forward transform to one or more points.
269 Parameters
270 ----------
271 x : `numpy.ndarray` | `float`
272 ``x`` values of the points to transform.
273 y : `numpy.ndarray` | `float`
274 ``y`` values of the points to transform.
276 Returns
277 -------
278 `XY` [`numpy.ndarray` | `float`]
279 The transformed point or points.
280 """
281 return _standardize_xy(
282 _ast_apply(
283 self._ast_mapping.applyForward,
284 x=self._in_frame.standardize_x(x),
285 y=self._in_frame.standardize_y(y),
286 ),
287 self._out_frame,
288 )
290 def apply_inverse[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
291 """Apply the inverse transform to one or more points.
293 Parameters
294 ----------
295 x : `numpy.ndarray` | `float`
296 ``x`` values of the points to transform.
297 y : `numpy.ndarray` | `float`
298 ``y`` values of the points to transform.
300 Returns
301 -------
302 `XY` [`numpy.ndarray` | `float`]
303 The transformed point or points.
304 """
305 return _standardize_xy(
306 _ast_apply(
307 self._ast_mapping.applyInverse,
308 x=self._out_frame.standardize_x(x),
309 y=self._out_frame.standardize_y(y),
310 ),
311 self._in_frame,
312 )
314 def apply_forward_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
315 """Apply the forward transform to one or more unit-aware points.
317 Parameters
318 ----------
319 x
320 ``x`` values of the points to transform.
321 y
322 ``y`` values of the points to transform.
324 Returns
325 -------
326 `XY` [`astropy.units.Quantity`]
327 The transformed point or points.
328 """
329 xy = self.apply_forward(x=x.to_value(self._in_frame.unit), y=y.to_value(self._in_frame.unit))
330 return XY(xy.x * self._out_frame.unit, xy.y * self._out_frame.unit)
332 def apply_inverse_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
333 """Apply the inverse transform to one or more unit-aware points.
335 Parameters
336 ----------
337 x
338 ``x`` values of the points to transform.
339 y
340 ``y`` values of the points to transform.
342 Returns
343 -------
344 `XY` [`astropy.units.Quantity`]
345 The transformed point or points.
346 """
347 xy = self.apply_inverse(x=x.to_value(self._out_frame.unit), y=y.to_value(self._out_frame.unit))
348 return XY(xy.x * self._in_frame.unit, xy.y * self._in_frame.unit)
350 def decompose(self) -> list[Transform[Any, Any]]:
351 """Deconstruct a composed transform into its constituent parts.
353 Notes
354 -----
355 Most transforms will just return a single-element list holding
356 ``self``. Identity transform will return an empty list, and
357 transforms composed with `then` will return the original transforms.
358 Transforms constructed by `FrameSet` may or may not be decomposable.
359 """
360 if not self._components: 360 ↛ 366line 360 didn't jump to line 366 because the condition on line 360 was always true
361 if self.in_frame == self._out_frame: 361 ↛ 364line 361 didn't jump to line 364 because the condition on line 361 was always true
362 return []
363 else:
364 return [self]
365 else:
366 return list(self._components)
368 def inverted(self) -> Transform[O, I]:
369 """Return the inverse of this transform."""
370 return Transform[O, I](
371 self._out_frame,
372 self._in_frame,
373 self._ast_mapping.inverted(),
374 in_bounds=self.out_bounds,
375 out_bounds=self.in_bounds,
376 components=[t.inverted() for t in reversed(self._components)],
377 )
379 def then[F: Frame](self, next: Transform[O, F], remember_components: bool = True) -> Transform[I, F]:
380 """Compose two transforms into another.
382 Parameters
383 ----------
384 next
385 Another transform to apply after ``self``.
386 remember_components
387 If `True`, the returned composed transform will remember ``self``
388 and ``other`` so they can be returned by `decompose`.
389 """
390 if self._out_frame != next._in_frame:
391 raise TransformCompositionError(
392 "Cannot compose transforms that do not share a common intermediate frame: "
393 f"{self._out_frame} != {next._in_frame}."
394 )
395 components = self.decompose() + next.decompose() if remember_components else ()
396 return Transform(
397 self._in_frame,
398 next._out_frame,
399 self._ast_mapping.then(next._ast_mapping),
400 in_bounds=self.in_bounds,
401 out_bounds=next.out_bounds,
402 components=components,
403 )
405 def as_fits_wcs(self, bbox: Box) -> astropy.wcs.WCS | None:
406 """Return a FITS WCS representation of this transform, if possible.
408 Parameters
409 ----------
410 bbox
411 Bounding box of the array the FITS WCS will describe. This
412 transform object is assumed to work on the same coordinate system
413 in which ``bbox`` is defined, while the FITS WCS will consider the
414 first row and column in that box to be ``(0, 0)`` (in Astropy
415 interfaces) or ``(1, 1)`` (in the FITS representation itself).
417 Notes
418 -----
419 This method assumes the transform maps pixel coordinates to world
420 coordinates.
422 Not all transforms can be represented exactly; when a FITS
423 represention is not possible, `None` is returned. When the returned
424 WCS is not `None`, it will have the same functional form, but it may
425 not evaluate identically due to small implementation differences in
426 the order of floating-point operations.
427 """
428 ast_frame_set = self._get_ast_frame_set()
429 _prepend_ast_shift(ast_frame_set, x=1.0 - bbox.x.start, y=1.0 - bbox.y.start, ast_domain="GRID")
430 ast_stream = astshim.StringStream()
431 ast_fits_chan = astshim.FitsChan(
432 ast_stream, "Encoding=FITS-WCS, CDMatrix=1, FitsAxisOrder=<copy>, FitsTol=0.0001"
433 )
434 ast_fits_chan.setFitsI("NAXIS1", bbox.x.size)
435 ast_fits_chan.setFitsI("NAXIS2", bbox.y.size)
436 n_writes = ast_fits_chan.write(ast_frame_set)
437 if not n_writes:
438 return None
439 header = astropy.io.fits.Header(astropy.io.fits.Card.fromstring(c) for c in ast_fits_chan)
440 return astropy.wcs.WCS(header)
442 def serialize[P: pydantic.BaseModel](
443 self, archive: OutputArchive[P], *, use_frame_sets: bool = False
444 ) -> TransformSerializationModel[P]:
445 """Serialize a transform to an archive.
447 Parameters
448 ----------
449 archive
450 Archive to serialize to.
451 use_frame_sets
452 If `True`, decompose the transform and try to reference component
453 mappings that were already serialized into a `FrameSet` in the
454 archive. Note that if multiple transforms exist between a pair of
455 frames (e.g. a `SkyProjection` and its FITS approximation), this
456 may cause the wrong one to be saved. When this option is used, the
457 frame set must be saved before the transform, and it must be
458 deserialized before the transform as well.
460 Returns
461 -------
462 `TransformSerializationModel`
463 Serialized form of the transform.
464 """
465 model = TransformSerializationModel[P]()
466 if use_frame_sets: 466 ↛ 467line 466 didn't jump to line 467 because the condition on line 466 was never true
467 for link in self.decompose():
468 model.frames.append(link.in_frame.serialize())
469 model.bounds.append(link.in_bounds.serialize() if link.in_bounds is not None else None)
470 for frame_set, pointer in archive.iter_frame_sets():
471 if link.in_frame in frame_set and link.out_frame in frame_set:
472 model.mappings.append(pointer)
473 break
474 else:
475 model.mappings.append(MappingSerializationModel(ast=link._ast_mapping.show()))
476 else:
477 model.frames.append(self.in_frame.serialize())
478 model.bounds.append(self.in_bounds.serialize() if self.in_bounds is not None else None)
479 model.mappings.append(MappingSerializationModel(ast=self._ast_mapping.show()))
480 model.frames.append(self.out_frame.serialize())
481 model.bounds.append(self.out_bounds.serialize() if self.out_bounds is not None else None)
482 return model
484 @staticmethod
485 def _get_archive_tree_type[P: pydantic.BaseModel](
486 pointer_type: type[P],
487 ) -> type[TransformSerializationModel[P]]:
488 """Return the serialization model type for this object for an archive
489 type that uses the given pointer type.
490 """
491 return TransformSerializationModel[pointer_type] # type: ignore
493 @staticmethod
494 def from_legacy(
495 legacy: LegacyTransform,
496 in_frame: I,
497 out_frame: O,
498 in_bounds: Bounds | None = None,
499 out_bounds: Bounds | None = None,
500 ) -> Transform[I, O]:
501 """Construct a transform from a legacy `lsst.afw.geom.Transform`.
503 Parameters
504 ----------
505 legacy : `lsst.afw.geom.Transform`
506 Legacy transform object.
507 in_frame
508 Coordinate frame for input points to the forward transform.
509 out_frame
510 Coordinate frame for output points from the forward transform.
511 in_bounds
512 The region that bounds valid input points.
513 out_bounds
514 The region that bounds valid output points.
515 """
516 return Transform(
517 in_frame,
518 out_frame,
519 legacy.getMapping(),
520 in_bounds=in_bounds,
521 out_bounds=out_bounds,
522 )
524 def to_legacy(self) -> LegacyTransform:
525 """Convert to a legacy `lsst.afw.geom.TransformPoint2ToPoint2`
526 instance.
527 """
528 from lsst.afw.geom import TransformPoint2ToPoint2 as LegacyTransform
530 return LegacyTransform(self._ast_mapping, False)
532 def _get_ast_frame_set(self) -> Any:
533 ast_frame_set = astshim.FrameSet(_make_ast_frame(self._in_frame))
534 ast_frame_set.addFrame(astshim.FrameSet.BASE, self._ast_mapping, _make_ast_frame(self._out_frame))
535 return ast_frame_set
538def _ast_apply[T: np.ndarray | float](method: Any, *, x: T, y: T) -> XY[T]:
539 # TODO: add bounds argument and check inputs
540 # TODO: broadcast arrays with different shapes.
541 xy_in = np.vstack([x, y]).astype(np.float64)
542 xy_out = method(xy_in)
543 return XY(xy_out[0, :], xy_out[1, :])
546def _prepend_ast_shift(ast_frame_set: Any, x: float, y: float, ast_domain: str) -> None:
547 ast_output_frame_id = ast_frame_set.current
548 ast_frame_set.addFrame(
549 astshim.FrameSet.BASE,
550 astshim.ShiftMap([x, y]),
551 astshim.Frame(2, f"Domain={ast_domain}"),
552 )
553 ast_frame_set.base = ast_frame_set.current
554 ast_frame_set.current = ast_output_frame_id
557def _make_ast_frame(frame: Frame) -> Any:
558 if frame is SkyFrame.ICRS:
559 return astshim.SkyFrame("")
560 ast_frame = astshim.Frame(2, f"Ident={frame._ast_ident}")
561 if frame.unit is not None: 561 ↛ 565line 561 didn't jump to line 565 because the condition on line 561 was always true
562 fits_unit = frame.unit.to_string(format="fits")
563 ast_frame.setUnit(1, fits_unit)
564 ast_frame.setUnit(2, fits_unit)
565 ast_frame.setLabel(1, "x")
566 ast_frame.setLabel(2, "y")
567 return ast_frame
570def _standardize_xy[T: np.ndarray | float](xy: XY[T], frame: Frame) -> XY[T]:
571 return XY(x=frame.standardize_x(xy.x), y=frame.standardize_y(xy.y))
574class MappingSerializationModel(pydantic.BaseModel):
575 """Serialization model for an AST Mapping."""
577 ast: str = pydantic.Field(description="A serialized Starlink AST Mapping, using the AST native encoding.")
580class TransformSerializationModel[P: pydantic.BaseModel](ArchiveTree):
581 """Serialization model for coordinate transforms."""
583 SCHEMA_NAME: ClassVar[str] = "transform"
584 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
585 MIN_READ_VERSION: ClassVar[int] = 1
586 PUBLIC_TYPE: ClassVar[type] = Transform
588 frames: list[SerializableFrame] = pydantic.Field(
589 default_factory=list,
590 description=textwrap.dedent(
591 """
592 List of frames that this transform passes through.
594 All transforms include at least two frames (the endpoints). Others
595 intermediate frames may be included to facilitate data-sharing
596 between transforms.
597 """
598 ),
599 )
601 bounds: list[BoundsSerializationModel | None] = pydantic.Field(
602 default_factory=list,
603 description=textwrap.dedent(
604 """
605 List of the bounds of the ``frames`` for this transform.
607 This always has the same number of elements as ``frames``.
608 """
609 ),
610 )
612 mappings: list[P | MappingSerializationModel] = pydantic.Field(
613 default_factory=list,
614 description=textwrap.dedent(
615 """
616 The actual mappings between frames, or archive pointers to
617 serialized FrameSet objects from which they can be obtained.
619 This always has one fewer element than ``frames``.
620 """
621 ),
622 )
624 def deserialize(self, archive: InputArchive[P], **kwargs: Any) -> Transform[Any, Any]:
625 """Deserialize a transform from an archive.
627 Parameters
628 ----------
629 archive
630 Archive to read from.
631 **kwargs
632 Unsupported keyword arguments are accepted only to provide better
633 error messages (raising `serialization.InvalidParameterError`).
634 """
635 if kwargs: 635 ↛ 636line 635 didn't jump to line 636 because the condition on line 635 was never true
636 raise InvalidParameterError(f"Unrecognized parameters for Transform: {set(kwargs.keys())}.")
637 if len(self.frames) != len(self.bounds): 637 ↛ 638line 637 didn't jump to line 638 because the condition on line 637 was never true
638 raise ArchiveReadError(
639 f"Inconsistent lengths for 'frames' ({len(self.frames)}) and 'bounds' ({len(self.bounds)})."
640 )
641 if len(self.frames) != len(self.mappings) + 1: 641 ↛ 642line 641 didn't jump to line 642 because the condition on line 641 was never true
642 raise ArchiveReadError(
643 f"Inconsistent lengths for 'frames' ({len(self.frames)}) and "
644 f"'mappings' ({len(self.mappings)}; should be one less)."
645 )
646 # We can't just compose onto an identity Transform if we want to
647 # preserve the FrameSet-ness of any of these mappings.
648 transform: Transform | None = None
649 for n, mapping in enumerate(self.mappings):
650 match mapping:
651 case MappingSerializationModel(ast=serialized_mapping): 651 ↛ 662line 651 didn't jump to line 662 because the pattern on line 651 always matched
652 ast_mapping = astshim.Mapping.fromString(serialized_mapping)
653 in_bounds = self.bounds[n]
654 out_bounds = self.bounds[n + 1]
655 new_transform = Transform(
656 self.frames[n].deserialize(),
657 self.frames[n + 1].deserialize(),
658 ast_mapping,
659 in_bounds.deserialize() if in_bounds is not None else None,
660 out_bounds.deserialize() if out_bounds is not None else None,
661 )
662 case reference:
663 frame_set = archive.get_frame_set(reference)
664 new_transform = frame_set[self.frames[n].deserialize(), self.frames[n + 1].deserialize()]
665 if transform is None: 665 ↛ 668line 665 didn't jump to line 668 because the condition on line 665 was always true
666 transform = new_transform
667 else:
668 transform = transform.then(new_transform)
669 if transform is None: 669 ↛ 670line 669 didn't jump to line 670 because the condition on line 669 was never true
670 transform = Transform.identity(self.frames[0].deserialize())
671 return transform