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