Coverage for python / lsst / images / _transforms / _transform.py: 29%
169 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-24 02: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, TypeVar, final
24import astropy.io.fits.header
25import astropy.units as u
26import numpy as np
27import pydantic
29from .._geom import XY, Bounds, Box, SerializableBounds
30from ..serialization import ArchiveReadError, ArchiveTree, InputArchive, OutputArchive
31from ._frames import Frame, SerializableFrame, SkyFrame
33if TYPE_CHECKING:
34 from ._projection import Projection
36# These pre-python-3.12 declaration are needed by Sphinx (probably the
37# autodoc-typehints plugin.
38I = TypeVar("I", bound=Frame) # noqa: E741
39O = TypeVar("O", bound=Frame) # noqa: E741
42class TransformCompositionError(RuntimeError):
43 """Exception raised when two transforms cannot be composed."""
46@final
47class Transform[I: Frame, O: Frame]:
48 """A transform that maps two coordinate frames.
50 Notes
51 -----
52 The `Transform` class constructor is considered a private implementation
53 detail. Instead of using this, various factory methods are available:
55 - `from_fits_wcs` constructs a transform from a FITS WCS, as represented
56 `astropy.wcs.WCS`;
57 - `then` composes two transforms;
58 - `identity` constructs a trivial transform that does nothing;
59 - `inverted` returns the inverse of a transform;
60 - `from_legacy` converts an `lsst.afw.geom.Transform` instance.
62 When applied to celestial coordinate systems, ``x=ra`` and ``y=dec``.
63 `Projection` provides a more natural interface for pixel-to-sky transforms.
65 `Transform` is conceptually immutable (the internal AST Mapping should
66 never be modified in-place after construction), and hence does not need to
67 be copied when any object that holds it is copied.
68 """
70 def __init__(
71 self,
72 in_frame: I,
73 out_frame: O,
74 ast_mapping: Any,
75 in_bounds: Bounds | None = None,
76 out_bounds: Bounds | None = None,
77 components: Iterable[Transform[Any, Any]] = (),
78 ):
79 self._in_frame = in_frame
80 self._out_frame = out_frame
81 self._ast_mapping = ast_mapping
82 self._in_bounds = in_bounds or getattr(in_frame, "bbox", None)
83 self._out_bounds = out_bounds or getattr(out_frame, "bbox", None)
84 self._components = list(components)
86 @staticmethod
87 def from_fits_wcs(
88 fits_wcs: astropy.wcs.WCS,
89 in_frame: I,
90 out_frame: O,
91 in_bounds: Bounds | None = None,
92 out_bounds: Bounds | None = None,
93 x0: int = 0,
94 y0: int = 0,
95 ) -> Transform[I, O]:
96 """Construct a transform from a FITS WCS.
98 Parameters
99 ----------
100 fits_wcs
101 FITS WCS to convert.
102 in_frame
103 Coordinate frame for input points to the forward transform.
104 out_frame
105 Coordinate frame for output points from the forward transform.
106 in_bounds
107 The region that bounds valid input points.
108 out_bounds
109 The region that bounds valid output points.
110 x0
111 Logical coordinate of the first column in the array this WCS
112 relates to world coordinates.
113 y0
114 Logical coordinate of the first column in the array this WCS
115 relates to world coordinates.
117 Notes
118 -----
119 The ``x0`` and ``y0`` parameters reflect the fact that for FITS, the
120 first row and column are always labeled ``(1, 1)``, while in Astropy
121 and most other Python libraries, they are ``(0, 0)``. The `types` in
122 this package (e.g. `Image`, `Mask`) allow them to be any pair of
123 integers.
125 See Also
126 --------
127 Projection.from_fits_wcs
128 """
129 import astshim
131 ast_stream = astshim.StringStream(fits_wcs.to_header_string(relax=True))
132 ast_fits_chan = astshim.FitsChan(ast_stream, "Encoding=FITS-WCS, SipReplace=0")
133 ast_frame_set = ast_fits_chan.read()
134 _prepend_ast_shift(ast_frame_set, x=x0 - 1.0, y=y0 - 1.0, ast_domain="PIXEL")
135 return Transform(
136 in_frame,
137 out_frame,
138 ast_frame_set.getMapping(),
139 in_bounds=in_bounds,
140 out_bounds=out_bounds,
141 )
143 @staticmethod
144 def identity(frame: I) -> Transform[I, I]:
145 """Construct a trivial transform that maps a frame to itelf.
147 Parameters
148 ----------
149 frame
150 Frame used for both input and output points.
151 """
152 import astshim
154 return Transform(frame, frame, astshim.UnitMap(2))
156 @property
157 def in_frame(self) -> I:
158 """Coordinate frame for input points."""
159 return self._in_frame
161 @property
162 def out_frame(self) -> O:
163 """Coordinate frame for output points."""
164 return self._out_frame
166 @property
167 def in_bounds(self) -> Bounds | None:
168 """The region that bounds valid input points (`Bounds` | `None`)."""
169 return self._in_bounds
171 @property
172 def out_bounds(self) -> Bounds | None:
173 """The region that bounds valid output points (`Bounds` | `None`)."""
174 return self._out_bounds
176 def show(self, simplified: bool = False, comments: bool = False) -> str:
177 """Return the AST native representation of the transform.
179 Parameters
180 ----------
181 simplified
182 Whether to ask AST to simplify the mapping before showing it.
183 This will make it much more likely that two equivalent transforms
184 have the same `show` result.
185 comments
186 Whether to include descriptive comments.
187 """
188 ast_mapping = self._ast_mapping
189 if simplified:
190 ast_mapping = ast_mapping.simplified()
191 return ast_mapping.show(comments)
193 def apply_forward[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
194 """Apply the forward transform to one or more points.
196 Parameters
197 ----------
198 x : `numpy.ndarray` | `float`
199 ``x`` values of the points to transform.
200 y : `numpy.ndarray` | `float`
201 ``y`` values of the points to transform.
203 Returns
204 -------
205 `XY` [`numpy.ndarray` | `float`]
206 The transformed point or points.
207 """
208 return _standardize_xy(
209 _ast_apply(
210 self._ast_mapping.applyForward,
211 x=self._in_frame.standardize_x(x),
212 y=self._in_frame.standardize_y(y),
213 ),
214 self._out_frame,
215 )
217 def apply_inverse[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
218 """Apply the inverse transform to one or more points.
220 Parameters
221 ----------
222 x : `numpy.ndarray` | `float`
223 ``x`` values of the points to transform.
224 y : `numpy.ndarray` | `float`
225 ``y`` values of the points to transform.
227 Returns
228 -------
229 `XY` [`numpy.ndarray` | `float`]
230 The transformed point or points.
231 """
232 return _standardize_xy(
233 _ast_apply(
234 self._ast_mapping.applyInverse,
235 x=self._out_frame.standardize_x(x),
236 y=self._out_frame.standardize_y(y),
237 ),
238 self._in_frame,
239 )
241 def apply_forward_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
242 """Apply the forward transform to one or more unit-aware points.
244 Parameters
245 ----------
246 x
247 ``x`` values of the points to transform.
248 y
249 ``y`` values of the points to transform.
251 Returns
252 -------
253 `XY` [`astropy.units.Quantity`]
254 The transformed point or points.
255 """
256 xy = self.apply_forward(x=x.to_value(self._in_frame.unit), y=y.to_value(self._in_frame.unit))
257 return XY(xy.x * self._out_frame.unit, xy.y * self._out_frame.unit)
259 def apply_inverse_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
260 """Apply the inverse transform to one or more unit-aware points.
262 Parameters
263 ----------
264 x
265 ``x`` values of the points to transform.
266 y
267 ``y`` values of the points to transform.
269 Returns
270 -------
271 `XY` [`astropy.units.Quantity`]
272 The transformed point or points.
273 """
274 xy = self.apply_inverse(x=x.to_value(self._out_frame.unit), y=y.to_value(self._out_frame.unit))
275 return XY(xy.x * self._in_frame.unit, xy.y * self._in_frame.unit)
277 def decompose(self) -> list[Transform[Any, Any]]:
278 """Deconstruct a composed transform into its constituent parts.
280 Notes
281 -----
282 Most transforms will just return a single-element list holding
283 ``self``. Identity transform will return an empty list, and
284 transforms composed with `then` will return the original transforms.
285 Transforms constructed by `FrameSet` may or may not be decomposable.
286 """
287 if not self._components:
288 if self.in_frame == self._out_frame:
289 return []
290 else:
291 return [self]
292 else:
293 return list(self._components)
295 def inverted(self) -> Transform[O, I]:
296 """Return the inverse of this transform."""
297 return Transform[O, I](
298 self._out_frame,
299 self._in_frame,
300 self._ast_mapping.inverted(),
301 in_bounds=self.out_bounds,
302 out_bounds=self.in_bounds,
303 components=[t.inverted() for t in reversed(self._components)],
304 )
306 def then[F: Frame](self, next: Transform[O, F], remember_components: bool = True) -> Transform[I, F]:
307 """Compose two transforms into another.
309 Parameters
310 ----------
311 next
312 Another transform to apply after ``self``.
313 remember_components
314 If `True`, the returned composed transform will remember ``self``
315 and ``other`` so they can be returned by `decompose`.
316 """
317 if self._out_frame != next._in_frame:
318 raise TransformCompositionError(
319 "Cannot compose transforms that do not share a common intermediate frame: "
320 f"{self._out_frame} != {next._in_frame}."
321 )
322 components = self.decompose() + next.decompose() if remember_components else ()
323 return Transform(
324 self._in_frame,
325 next._out_frame,
326 self._ast_mapping.then(next._ast_mapping),
327 in_bounds=self.in_bounds,
328 out_bounds=next.out_bounds,
329 components=components,
330 )
332 def as_projection(self: Transform[I, SkyFrame]) -> Projection[I]:
333 """Return a `Projection` view of this transform.
335 This is only valid when `out_frame` is `~SkyFrame.ICRS`.
336 """
337 from ._projection import Projection
339 return Projection(self)
341 def as_fits_wcs(self, bbox: Box) -> astropy.wcs.WCS | None:
342 """Return a FITS WCS representation of this transform, if possible.
344 Parameters
345 ----------
346 bbox
347 Bounding box of the array the FITS WCS will describe. This
348 transform object is assumed to work on the same coordinate system
349 in which ``bbox`` is defined, while the FITS WCS will consider the
350 first row and column in that box to be ``(0, 0)`` (in Astropy
351 interfaces) or ``(1, 1)`` (in the FITS representation itself).
353 Notes
354 -----
355 This method assumes the transform maps pixel coordinates to world
356 coordinates.
358 Not all transforms can be represented exactly; when a FITS
359 represention is not possible, `None` is returned. When the returned
360 WCS is not `None`, it will have the same functional form, but it may
361 not evaluate identically due to small implementation differences in
362 the order of floating-point operations.
363 """
364 import astshim
366 ast_frame_set = self._get_ast_frame_set()
367 _prepend_ast_shift(ast_frame_set, x=1.0 - bbox.x.start, y=1.0 - bbox.y.start, ast_domain="GRID")
368 ast_stream = astshim.StringStream()
369 ast_fits_chan = astshim.FitsChan(
370 ast_stream, "Encoding=FITS-WCS, CDMatrix=1, FitsAxisOrder=<copy>, FitsTol=0.0001"
371 )
372 ast_fits_chan.setFitsI("NAXIS1", bbox.x.size)
373 ast_fits_chan.setFitsI("NAXIS2", bbox.y.size)
374 n_writes = ast_fits_chan.write(ast_frame_set)
375 if not n_writes:
376 return None
377 header = astropy.io.fits.Header(astropy.io.fits.Card.fromstring(c) for c in ast_fits_chan)
378 return astropy.wcs.WCS(header)
380 def serialize[P: pydantic.BaseModel](
381 self, archive: OutputArchive[P], *, use_frame_sets: bool = False
382 ) -> TransformSerializationModel[P]:
383 """Serialize a transform to an archive.
385 Parameters
386 ----------
387 archive
388 Archive to serialize to.
389 use_frame_sets
390 If `True`, decompose the transform and try to reference component
391 mappings that were already serialized into a `FrameSet` in the
392 archive. Note that if multiple transforms exist between a pair of
393 frames (e.g. a `Projection` and its FITS approximation), this may
394 cause the wrong one to be saved. When this option is used, the
395 frame set must be saved before the transform, and it must be
396 deserialized before the transform as well.
398 Returns
399 -------
400 `TransformSerializationModel`
401 Serialized form of the transform.
402 """
403 model = TransformSerializationModel[P]()
404 if use_frame_sets:
405 for link in self.decompose():
406 model.frames.append(link.in_frame.serialize())
407 model.bounds.append(link.in_bounds.serialize() if link.in_bounds is not None else None)
408 for frame_set, pointer in archive.iter_frame_sets():
409 if link.in_frame in frame_set and link.out_frame in frame_set:
410 model.mappings.append(pointer)
411 break
412 else:
413 model.mappings.append(MappingSerializationModel(ast=link._ast_mapping.show()))
414 else:
415 model.frames.append(self.in_frame.serialize())
416 model.bounds.append(self.in_bounds.serialize() if self.in_bounds is not None else None)
417 model.mappings.append(MappingSerializationModel(ast=self._ast_mapping.show()))
418 model.frames.append(self.out_frame.serialize())
419 model.bounds.append(self.out_bounds.serialize() if self.out_bounds is not None else None)
420 return model
422 @staticmethod
423 def deserialize[P: pydantic.BaseModel](
424 model: TransformSerializationModel[P], archive: InputArchive[P]
425 ) -> Transform[Any, Any]:
426 """Deserialize a transform from an archive.
428 Parameters
429 ----------
430 model
431 Seralized form of the transform.
432 archive
433 Archive to read from.
434 """
435 import astshim
437 if len(model.frames) != len(model.bounds):
438 raise ArchiveReadError(
439 f"Inconsistent lengths for 'frames' ({len(model.frames)}) and 'bounds' ({len(model.bounds)})."
440 )
441 if len(model.frames) != len(model.mappings) + 1:
442 raise ArchiveReadError(
443 f"Inconsistent lengths for 'frames' ({len(model.frames)}) and "
444 f"'mappings' ({len(model.mappings)}; should be one less)."
445 )
446 transform = Transform.identity(Frame.deserialize(model.frames[0]))
447 for n, mapping in enumerate(model.mappings):
448 match mapping:
449 case MappingSerializationModel(ast=serialized_mapping):
450 ast_mapping = astshim.Mapping.fromString(serialized_mapping)
451 in_bounds = model.bounds[n]
452 out_bounds = model.bounds[n + 1]
453 transform = transform.then(
454 Transform(
455 Frame.deserialize(model.frames[n]),
456 Frame.deserialize(model.frames[n + 1]),
457 ast_mapping,
458 Bounds.deserialize(in_bounds) if in_bounds is not None else None,
459 Bounds.deserialize(out_bounds) if out_bounds is not None else None,
460 )
461 )
462 case reference:
463 frame_set = archive.get_frame_set(reference)
464 transform = transform.then(
465 frame_set[Frame.deserialize(model.frames[n]), Frame.deserialize(model.frames[n + 1])]
466 )
467 return transform
469 @staticmethod
470 def _get_archive_tree_type[P: pydantic.BaseModel](
471 pointer_type: type[P],
472 ) -> type[TransformSerializationModel[P]]:
473 """Return the serialization model type for this object for an archive
474 type that uses the given pointer type.
475 """
476 return TransformSerializationModel[pointer_type] # type: ignore
478 @staticmethod
479 def from_legacy(
480 legacy: Any,
481 in_frame: I,
482 out_frame: O,
483 in_bounds: Bounds | None = None,
484 out_bounds: Bounds | None = None,
485 ) -> Transform[I, O]:
486 """Construct a transform from a legacy `lsst.afw.geom.Transform`.
488 Parameters
489 ----------
490 legacy : `lsst.afw.geom.Transform`
491 Legacy transform object.
492 in_frame
493 Coordinate frame for input points to the forward transform.
494 out_frame
495 Coordinate frame for output points from the forward transform.
496 in_bounds
497 The region that bounds valid input points.
498 out_bounds
499 The region that bounds valid output points.
500 """
501 return Transform(
502 in_frame,
503 out_frame,
504 legacy.getMapping(),
505 in_bounds=in_bounds,
506 out_bounds=out_bounds,
507 )
509 def _get_ast_frame_set(self) -> Any:
510 import astshim
512 ast_frame_set = astshim.FrameSet(_make_ast_frame(self._in_frame))
513 ast_frame_set.addFrame(astshim.FrameSet.BASE, self._ast_mapping, _make_ast_frame(self._out_frame))
514 return ast_frame_set
517def _ast_apply[T: np.ndarray | float](method: Any, *, x: T, y: T) -> XY[T]:
518 # TODO: add bounds argument and check inputs
519 # TODO: broadcast arrays with different shapes.
520 xy_in = np.vstack([x, y]).astype(np.float64)
521 xy_out = method(xy_in)
522 return XY(xy_out[0, :], xy_out[1, :])
525def _prepend_ast_shift(ast_frame_set: Any, x: float, y: float, ast_domain: str) -> None:
526 import astshim
528 ast_output_frame_id = ast_frame_set.current
529 ast_frame_set.addFrame(
530 astshim.FrameSet.BASE,
531 astshim.ShiftMap([x, y]),
532 astshim.Frame(2, f"Domain={ast_domain}"),
533 )
534 ast_frame_set.base = ast_frame_set.current
535 ast_frame_set.current = ast_output_frame_id
538def _make_ast_frame(frame: Frame) -> Any:
539 import astshim
541 if frame is SkyFrame.ICRS:
542 return astshim.SkyFrame("")
543 ast_frame = astshim.Frame(2, f"Ident={frame._ast_ident}")
544 if frame.unit is not None:
545 fits_unit = frame.unit.to_string(format="fits")
546 ast_frame.setUnit(1, fits_unit)
547 ast_frame.setUnit(2, fits_unit)
548 ast_frame.setLabel(1, "x")
549 ast_frame.setLabel(2, "y")
550 return ast_frame
553def _standardize_xy[T: np.ndarray | float](xy: XY[T], frame: Frame) -> XY[T]:
554 return XY(x=frame.standardize_x(xy.x), y=frame.standardize_y(xy.y))
557class MappingSerializationModel(ArchiveTree):
558 """Serialization model for an AST Mapping."""
560 ast: str = pydantic.Field(description="A serialized Starlink AST Mapping, using the AST native encoding.")
563class TransformSerializationModel[P: pydantic.BaseModel](ArchiveTree):
564 """Serialization model for coordinate transforms."""
566 frames: list[SerializableFrame] = pydantic.Field(
567 default_factory=list,
568 description=textwrap.dedent(
569 """
570 List of frames that this transform passes through.
572 All transforms include at least two frames (the endpoints). Others
573 intermediate frames may be included to facilitate data-sharing
574 between transforms.
575 """
576 ),
577 )
579 bounds: list[SerializableBounds | None] = pydantic.Field(
580 default_factory=list,
581 description=textwrap.dedent(
582 """
583 List of the bounds of the ``frames`` for this transform.
585 This always has the same number of elements as ``frames``.
586 """
587 ),
588 )
590 mappings: list[P | MappingSerializationModel] = pydantic.Field(
591 default_factory=list,
592 description=textwrap.dedent(
593 """
594 The actual mappings between frames, or archive pointers to
595 serialized FrameSet objects from which they can be obtained.
597 This always has one fewer element than ``frames``.
598 """
599 ),
600 )