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