Coverage for python/lsst/images/_transforms/_transform.py: 78%
195 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-27 08:32 +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 SerializableBounds
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 - `inverted` returns the inverse of a transform;
83 - `from_legacy` converts an `lsst.afw.geom.Transform` instance.
85 When applied to celestial coordinate systems, ``x=ra`` and ``y=dec``.
86 `SkyProjection` provides a more natural interface for pixel-to-sky
87 transforms.
89 `Transform` is conceptually immutable (the internal AST Mapping should
90 never be modified in-place after construction), and hence does not need to
91 be copied when any object that holds it is copied.
92 """
94 def __init__(
95 self,
96 in_frame: I,
97 out_frame: O,
98 ast_mapping: astshim.Mapping,
99 in_bounds: Bounds | None = None,
100 out_bounds: Bounds | None = None,
101 components: Iterable[Transform[Any, Any]] = (),
102 ) -> None:
103 self._in_frame = in_frame
104 self._out_frame = out_frame
105 self._ast_mapping = ast_mapping
106 self._in_bounds = in_bounds or getattr(in_frame, "bbox", None)
107 self._out_bounds = out_bounds or getattr(out_frame, "bbox", None)
108 self._components = list(components)
110 def __eq__(self, other: Any) -> bool:
111 if self is other:
112 # Short circuit for case where you are quickly checking
113 # that the image WCS and variance WCS are the same object.
114 return True
115 if not isinstance(other, Transform):
116 return NotImplemented
117 if self._ast_mapping != other._ast_mapping:
118 return False
119 if self._in_bounds != other._in_bounds:
120 return False
121 if self._out_bounds != other._out_bounds:
122 return False
123 if self._in_frame != other._in_frame:
124 return False
125 if self._out_frame != other._out_frame:
126 return False
127 if self._components != other._components:
128 return False
129 return True
131 @staticmethod
132 def from_fits_wcs(
133 fits_wcs: astropy.wcs.WCS,
134 in_frame: I,
135 out_frame: O,
136 in_bounds: Bounds | None = None,
137 out_bounds: Bounds | None = None,
138 x0: int = 0,
139 y0: int = 0,
140 ) -> Transform[I, O]:
141 """Construct a transform from a FITS WCS.
143 Parameters
144 ----------
145 fits_wcs
146 FITS WCS to convert.
147 in_frame
148 Coordinate frame for input points to the forward transform.
149 out_frame
150 Coordinate frame for output points from the forward transform.
151 in_bounds
152 The region that bounds valid input points.
153 out_bounds
154 The region that bounds valid output points.
155 x0
156 Logical coordinate of the first column in the array this WCS
157 relates to world coordinates.
158 y0
159 Logical coordinate of the first column in the array this WCS
160 relates to world coordinates.
162 Notes
163 -----
164 The ``x0`` and ``y0`` parameters reflect the fact that for FITS, the
165 first row and column are always labeled ``(1, 1)``, while in Astropy
166 and most other Python libraries, they are ``(0, 0)``. The `types` in
167 this package (e.g. `Image`, `Mask`) allow them to be any pair of
168 integers.
170 See Also
171 --------
172 SkyProjection.from_fits_wcs
173 """
174 ast_stream = astshim.StringStream(fits_wcs.to_header_string(relax=True))
175 ast_fits_chan = astshim.FitsChan(ast_stream, "Encoding=FITS-WCS, SipReplace=0, IWC=1")
176 ast_frame_set = ast_fits_chan.read()
177 _prepend_ast_shift(ast_frame_set, x=x0 - 1.0, y=y0 - 1.0, ast_domain="PIXEL")
178 return Transform(
179 in_frame,
180 out_frame,
181 ast_frame_set,
182 in_bounds=in_bounds,
183 out_bounds=out_bounds,
184 )
186 @staticmethod
187 def identity(frame: I) -> Transform[I, I]:
188 """Construct a trivial transform that maps a frame to itelf.
190 Parameters
191 ----------
192 frame
193 Frame used for both input and output points.
194 """
195 return Transform(frame, frame, astshim.UnitMap(2))
197 @property
198 def in_frame(self) -> I:
199 """Coordinate frame for input points."""
200 return self._in_frame
202 @property
203 def out_frame(self) -> O:
204 """Coordinate frame for output points."""
205 return self._out_frame
207 @property
208 def in_bounds(self) -> Bounds | None:
209 """The region that bounds valid input points (`Bounds` | `None`)."""
210 return self._in_bounds
212 @property
213 def out_bounds(self) -> Bounds | None:
214 """The region that bounds valid output points (`Bounds` | `None`)."""
215 return self._out_bounds
217 def show(self, simplified: bool = False, comments: bool = False) -> str:
218 """Return the AST native representation of the transform.
220 Parameters
221 ----------
222 simplified
223 Whether to ask AST to simplify the mapping before showing it.
224 This will make it much more likely that two equivalent transforms
225 have the same `show` result. If the internal mapping is actually
226 a frame set (as needed to round-trip legacy
227 `lsst.afw.geom.SkyWcs` objects), this will also just show the
228 mapping with no frame set information.
229 comments
230 Whether to include descriptive comments.
231 """
232 ast_mapping = self._ast_mapping
233 if simplified:
234 if isinstance(ast_mapping, astshim.FrameSet):
235 ast_mapping = ast_mapping.getMapping()
236 ast_mapping = ast_mapping.simplified()
237 return ast_mapping.show(comments)
239 def apply_forward[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
240 """Apply the forward transform to one or more points.
242 Parameters
243 ----------
244 x : `numpy.ndarray` | `float`
245 ``x`` values of the points to transform.
246 y : `numpy.ndarray` | `float`
247 ``y`` values of the points to transform.
249 Returns
250 -------
251 `XY` [`numpy.ndarray` | `float`]
252 The transformed point or points.
253 """
254 return _standardize_xy(
255 _ast_apply(
256 self._ast_mapping.applyForward,
257 x=self._in_frame.standardize_x(x),
258 y=self._in_frame.standardize_y(y),
259 ),
260 self._out_frame,
261 )
263 def apply_inverse[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
264 """Apply the inverse transform to one or more points.
266 Parameters
267 ----------
268 x : `numpy.ndarray` | `float`
269 ``x`` values of the points to transform.
270 y : `numpy.ndarray` | `float`
271 ``y`` values of the points to transform.
273 Returns
274 -------
275 `XY` [`numpy.ndarray` | `float`]
276 The transformed point or points.
277 """
278 return _standardize_xy(
279 _ast_apply(
280 self._ast_mapping.applyInverse,
281 x=self._out_frame.standardize_x(x),
282 y=self._out_frame.standardize_y(y),
283 ),
284 self._in_frame,
285 )
287 def apply_forward_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
288 """Apply the forward transform to one or more unit-aware points.
290 Parameters
291 ----------
292 x
293 ``x`` values of the points to transform.
294 y
295 ``y`` values of the points to transform.
297 Returns
298 -------
299 `XY` [`astropy.units.Quantity`]
300 The transformed point or points.
301 """
302 xy = self.apply_forward(x=x.to_value(self._in_frame.unit), y=y.to_value(self._in_frame.unit))
303 return XY(xy.x * self._out_frame.unit, xy.y * self._out_frame.unit)
305 def apply_inverse_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
306 """Apply the inverse transform to one or more unit-aware points.
308 Parameters
309 ----------
310 x
311 ``x`` values of the points to transform.
312 y
313 ``y`` values of the points to transform.
315 Returns
316 -------
317 `XY` [`astropy.units.Quantity`]
318 The transformed point or points.
319 """
320 xy = self.apply_inverse(x=x.to_value(self._out_frame.unit), y=y.to_value(self._out_frame.unit))
321 return XY(xy.x * self._in_frame.unit, xy.y * self._in_frame.unit)
323 def decompose(self) -> list[Transform[Any, Any]]:
324 """Deconstruct a composed transform into its constituent parts.
326 Notes
327 -----
328 Most transforms will just return a single-element list holding
329 ``self``. Identity transform will return an empty list, and
330 transforms composed with `then` will return the original transforms.
331 Transforms constructed by `FrameSet` may or may not be decomposable.
332 """
333 if not self._components: 333 ↛ 339line 333 didn't jump to line 339 because the condition on line 333 was always true
334 if self.in_frame == self._out_frame: 334 ↛ 337line 334 didn't jump to line 337 because the condition on line 334 was always true
335 return []
336 else:
337 return [self]
338 else:
339 return list(self._components)
341 def inverted(self) -> Transform[O, I]:
342 """Return the inverse of this transform."""
343 return Transform[O, I](
344 self._out_frame,
345 self._in_frame,
346 self._ast_mapping.inverted(),
347 in_bounds=self.out_bounds,
348 out_bounds=self.in_bounds,
349 components=[t.inverted() for t in reversed(self._components)],
350 )
352 def then[F: Frame](self, next: Transform[O, F], remember_components: bool = True) -> Transform[I, F]:
353 """Compose two transforms into another.
355 Parameters
356 ----------
357 next
358 Another transform to apply after ``self``.
359 remember_components
360 If `True`, the returned composed transform will remember ``self``
361 and ``other`` so they can be returned by `decompose`.
362 """
363 if self._out_frame != next._in_frame:
364 raise TransformCompositionError(
365 "Cannot compose transforms that do not share a common intermediate frame: "
366 f"{self._out_frame} != {next._in_frame}."
367 )
368 components = self.decompose() + next.decompose() if remember_components else ()
369 return Transform(
370 self._in_frame,
371 next._out_frame,
372 self._ast_mapping.then(next._ast_mapping),
373 in_bounds=self.in_bounds,
374 out_bounds=next.out_bounds,
375 components=components,
376 )
378 def as_fits_wcs(self, bbox: Box) -> astropy.wcs.WCS | None:
379 """Return a FITS WCS representation of this transform, if possible.
381 Parameters
382 ----------
383 bbox
384 Bounding box of the array the FITS WCS will describe. This
385 transform object is assumed to work on the same coordinate system
386 in which ``bbox`` is defined, while the FITS WCS will consider the
387 first row and column in that box to be ``(0, 0)`` (in Astropy
388 interfaces) or ``(1, 1)`` (in the FITS representation itself).
390 Notes
391 -----
392 This method assumes the transform maps pixel coordinates to world
393 coordinates.
395 Not all transforms can be represented exactly; when a FITS
396 represention is not possible, `None` is returned. When the returned
397 WCS is not `None`, it will have the same functional form, but it may
398 not evaluate identically due to small implementation differences in
399 the order of floating-point operations.
400 """
401 ast_frame_set = self._get_ast_frame_set()
402 _prepend_ast_shift(ast_frame_set, x=1.0 - bbox.x.start, y=1.0 - bbox.y.start, ast_domain="GRID")
403 ast_stream = astshim.StringStream()
404 ast_fits_chan = astshim.FitsChan(
405 ast_stream, "Encoding=FITS-WCS, CDMatrix=1, FitsAxisOrder=<copy>, FitsTol=0.0001"
406 )
407 ast_fits_chan.setFitsI("NAXIS1", bbox.x.size)
408 ast_fits_chan.setFitsI("NAXIS2", bbox.y.size)
409 n_writes = ast_fits_chan.write(ast_frame_set)
410 if not n_writes:
411 return None
412 header = astropy.io.fits.Header(astropy.io.fits.Card.fromstring(c) for c in ast_fits_chan)
413 return astropy.wcs.WCS(header)
415 def serialize[P: pydantic.BaseModel](
416 self, archive: OutputArchive[P], *, use_frame_sets: bool = False
417 ) -> TransformSerializationModel[P]:
418 """Serialize a transform to an archive.
420 Parameters
421 ----------
422 archive
423 Archive to serialize to.
424 use_frame_sets
425 If `True`, decompose the transform and try to reference component
426 mappings that were already serialized into a `FrameSet` in the
427 archive. Note that if multiple transforms exist between a pair of
428 frames (e.g. a `SkyProjection` and its FITS approximation), this
429 may cause the wrong one to be saved. When this option is used, the
430 frame set must be saved before the transform, and it must be
431 deserialized before the transform as well.
433 Returns
434 -------
435 `TransformSerializationModel`
436 Serialized form of the transform.
437 """
438 model = TransformSerializationModel[P]()
439 if use_frame_sets: 439 ↛ 440line 439 didn't jump to line 440 because the condition on line 439 was never true
440 for link in self.decompose():
441 model.frames.append(link.in_frame.serialize())
442 model.bounds.append(link.in_bounds.serialize() if link.in_bounds is not None else None)
443 for frame_set, pointer in archive.iter_frame_sets():
444 if link.in_frame in frame_set and link.out_frame in frame_set:
445 model.mappings.append(pointer)
446 break
447 else:
448 model.mappings.append(MappingSerializationModel(ast=link._ast_mapping.show()))
449 else:
450 model.frames.append(self.in_frame.serialize())
451 model.bounds.append(self.in_bounds.serialize() if self.in_bounds is not None else None)
452 model.mappings.append(MappingSerializationModel(ast=self._ast_mapping.show()))
453 model.frames.append(self.out_frame.serialize())
454 model.bounds.append(self.out_bounds.serialize() if self.out_bounds is not None else None)
455 return model
457 @staticmethod
458 def _get_archive_tree_type[P: pydantic.BaseModel](
459 pointer_type: type[P],
460 ) -> type[TransformSerializationModel[P]]:
461 """Return the serialization model type for this object for an archive
462 type that uses the given pointer type.
463 """
464 return TransformSerializationModel[pointer_type] # type: ignore
466 @staticmethod
467 def from_legacy(
468 legacy: LegacyTransform,
469 in_frame: I,
470 out_frame: O,
471 in_bounds: Bounds | None = None,
472 out_bounds: Bounds | None = None,
473 ) -> Transform[I, O]:
474 """Construct a transform from a legacy `lsst.afw.geom.Transform`.
476 Parameters
477 ----------
478 legacy : `lsst.afw.geom.Transform`
479 Legacy transform object.
480 in_frame
481 Coordinate frame for input points to the forward transform.
482 out_frame
483 Coordinate frame for output points from the forward transform.
484 in_bounds
485 The region that bounds valid input points.
486 out_bounds
487 The region that bounds valid output points.
488 """
489 return Transform(
490 in_frame,
491 out_frame,
492 legacy.getMapping(),
493 in_bounds=in_bounds,
494 out_bounds=out_bounds,
495 )
497 def to_legacy(self) -> LegacyTransform:
498 """Convert to a legacy `lsst.afw.geom.TransformPoint2ToPoint2`
499 instance.
500 """
501 from lsst.afw.geom import TransformPoint2ToPoint2 as LegacyTransform
503 return LegacyTransform(self._ast_mapping, False)
505 def _get_ast_frame_set(self) -> Any:
506 ast_frame_set = astshim.FrameSet(_make_ast_frame(self._in_frame))
507 ast_frame_set.addFrame(astshim.FrameSet.BASE, self._ast_mapping, _make_ast_frame(self._out_frame))
508 return ast_frame_set
511def _ast_apply[T: np.ndarray | float](method: Any, *, x: T, y: T) -> XY[T]:
512 # TODO: add bounds argument and check inputs
513 # TODO: broadcast arrays with different shapes.
514 xy_in = np.vstack([x, y]).astype(np.float64)
515 xy_out = method(xy_in)
516 return XY(xy_out[0, :], xy_out[1, :])
519def _prepend_ast_shift(ast_frame_set: Any, x: float, y: float, ast_domain: str) -> None:
520 ast_output_frame_id = ast_frame_set.current
521 ast_frame_set.addFrame(
522 astshim.FrameSet.BASE,
523 astshim.ShiftMap([x, y]),
524 astshim.Frame(2, f"Domain={ast_domain}"),
525 )
526 ast_frame_set.base = ast_frame_set.current
527 ast_frame_set.current = ast_output_frame_id
530def _make_ast_frame(frame: Frame) -> Any:
531 if frame is SkyFrame.ICRS:
532 return astshim.SkyFrame("")
533 ast_frame = astshim.Frame(2, f"Ident={frame._ast_ident}")
534 if frame.unit is not None: 534 ↛ 538line 534 didn't jump to line 538 because the condition on line 534 was always true
535 fits_unit = frame.unit.to_string(format="fits")
536 ast_frame.setUnit(1, fits_unit)
537 ast_frame.setUnit(2, fits_unit)
538 ast_frame.setLabel(1, "x")
539 ast_frame.setLabel(2, "y")
540 return ast_frame
543def _standardize_xy[T: np.ndarray | float](xy: XY[T], frame: Frame) -> XY[T]:
544 return XY(x=frame.standardize_x(xy.x), y=frame.standardize_y(xy.y))
547class MappingSerializationModel(pydantic.BaseModel):
548 """Serialization model for an AST Mapping."""
550 ast: str = pydantic.Field(description="A serialized Starlink AST Mapping, using the AST native encoding.")
553class TransformSerializationModel[P: pydantic.BaseModel](ArchiveTree):
554 """Serialization model for coordinate transforms."""
556 SCHEMA_NAME: ClassVar[str] = "transform"
557 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
558 MIN_READ_VERSION: ClassVar[int] = 1
559 PUBLIC_TYPE: ClassVar[type] = Transform
561 frames: list[SerializableFrame] = pydantic.Field(
562 default_factory=list,
563 description=textwrap.dedent(
564 """
565 List of frames that this transform passes through.
567 All transforms include at least two frames (the endpoints). Others
568 intermediate frames may be included to facilitate data-sharing
569 between transforms.
570 """
571 ),
572 )
574 bounds: list[SerializableBounds | None] = pydantic.Field(
575 default_factory=list,
576 description=textwrap.dedent(
577 """
578 List of the bounds of the ``frames`` for this transform.
580 This always has the same number of elements as ``frames``.
581 """
582 ),
583 )
585 mappings: list[P | MappingSerializationModel] = pydantic.Field(
586 default_factory=list,
587 description=textwrap.dedent(
588 """
589 The actual mappings between frames, or archive pointers to
590 serialized FrameSet objects from which they can be obtained.
592 This always has one fewer element than ``frames``.
593 """
594 ),
595 )
597 def deserialize(self, archive: InputArchive[P], **kwargs: Any) -> Transform[Any, Any]:
598 """Deserialize a transform from an archive.
600 Parameters
601 ----------
602 archive
603 Archive to read from.
604 **kwargs
605 Unsupported keyword arguments are accepted only to provide better
606 error messages (raising `serialization.InvalidParameterError`).
607 """
608 if kwargs: 608 ↛ 609line 608 didn't jump to line 609 because the condition on line 608 was never true
609 raise InvalidParameterError(f"Unrecognized parameters for Transform: {set(kwargs.keys())}.")
610 if len(self.frames) != len(self.bounds): 610 ↛ 611line 610 didn't jump to line 611 because the condition on line 610 was never true
611 raise ArchiveReadError(
612 f"Inconsistent lengths for 'frames' ({len(self.frames)}) and 'bounds' ({len(self.bounds)})."
613 )
614 if len(self.frames) != len(self.mappings) + 1: 614 ↛ 615line 614 didn't jump to line 615 because the condition on line 614 was never true
615 raise ArchiveReadError(
616 f"Inconsistent lengths for 'frames' ({len(self.frames)}) and "
617 f"'mappings' ({len(self.mappings)}; should be one less)."
618 )
619 # We can't just compose onto an identity Transform if we want to
620 # preserve the FrameSet-ness of any of these mappings.
621 transform: Transform | None = None
622 for n, mapping in enumerate(self.mappings):
623 match mapping:
624 case MappingSerializationModel(ast=serialized_mapping): 624 ↛ 635line 624 didn't jump to line 635 because the pattern on line 624 always matched
625 ast_mapping = astshim.Mapping.fromString(serialized_mapping)
626 in_bounds = self.bounds[n]
627 out_bounds = self.bounds[n + 1]
628 new_transform = Transform(
629 self.frames[n].deserialize(),
630 self.frames[n + 1].deserialize(),
631 ast_mapping,
632 in_bounds.deserialize() if in_bounds is not None else None,
633 out_bounds.deserialize() if out_bounds is not None else None,
634 )
635 case reference:
636 frame_set = archive.get_frame_set(reference)
637 new_transform = frame_set[self.frames[n].deserialize(), self.frames[n + 1].deserialize()]
638 if transform is None: 638 ↛ 641line 638 didn't jump to line 641 because the condition on line 638 was always true
639 transform = new_transform
640 else:
641 transform = transform.then(new_transform)
642 if transform is None: 642 ↛ 643line 642 didn't jump to line 643 because the condition on line 642 was never true
643 transform = Transform.identity(self.frames[0].deserialize())
644 return transform