Coverage for python / lsst / images / _transforms / _transform.py: 30%
164 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-05 02:32 -0800
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 apply_forward[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
177 """Apply the forward transform to one or more points.
179 Parameters
180 ----------
181 x : `numpy.ndarray` | `float`
182 ``x`` values of the points to transform.
183 y : `numpy.ndarray` | `float`
184 ``y`` values of the points to transform.
186 Returns
187 -------
188 `XY` [`numpy.ndarray` | `float`]
189 The transformed point or points.
190 """
191 return _standardize_xy(
192 _ast_apply(
193 self._ast_mapping.applyForward,
194 x=self._in_frame.standardize_x(x),
195 y=self._in_frame.standardize_y(y),
196 ),
197 self._out_frame,
198 )
200 def apply_inverse[T: np.ndarray | float](self, *, x: T, y: T) -> XY[T]:
201 """Apply the inverse transform to one or more points.
203 Parameters
204 ----------
205 x : `numpy.ndarray` | `float`
206 ``x`` values of the points to transform.
207 y : `numpy.ndarray` | `float`
208 ``y`` values of the points to transform.
210 Returns
211 -------
212 `XY` [`numpy.ndarray` | `float`]
213 The transformed point or points.
214 """
215 return _standardize_xy(
216 _ast_apply(
217 self._ast_mapping.applyInverse,
218 x=self._out_frame.standardize_x(x),
219 y=self._out_frame.standardize_y(y),
220 ),
221 self._in_frame,
222 )
224 def apply_forward_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
225 """Apply the forward transform to one or more unit-aware points.
227 Parameters
228 ----------
229 x
230 ``x`` values of the points to transform.
231 y
232 ``y`` values of the points to transform.
234 Returns
235 -------
236 `XY` [`astropy.units.Quantity`]
237 The transformed point or points.
238 """
239 xy = self.apply_forward(x=x.to_value(self._in_frame.unit), y=y.to_value(self._in_frame.unit))
240 return XY(xy.x * self._out_frame.unit, xy.y * self._out_frame.unit)
242 def apply_inverse_q(self, *, x: u.Quantity, y: u.Quantity) -> XY[u.Quantity]:
243 """Apply the inverse transform to one or more unit-aware points.
245 Parameters
246 ----------
247 x
248 ``x`` values of the points to transform.
249 y
250 ``y`` values of the points to transform.
252 Returns
253 -------
254 `XY` [`astropy.units.Quantity`]
255 The transformed point or points.
256 """
257 xy = self.apply_inverse(x=x.to_value(self._out_frame.unit), y=y.to_value(self._out_frame.unit))
258 return XY(xy.x * self._in_frame.unit, xy.y * self._in_frame.unit)
260 def decompose(self) -> list[Transform[Any, Any]]:
261 """Deconstruct a composed transform into its constituent parts.
263 Notes
264 -----
265 Most transforms will just return a single-element list holding
266 ``self``. Identity transform will return an empty list, and
267 transforms composed with `then` will return the original transforms.
268 Transforms constructed by `FrameSet` may or may not be decomposable.
269 """
270 if not self._components:
271 if self.in_frame == self._out_frame:
272 return []
273 else:
274 return [self]
275 else:
276 return list(self._components)
278 def inverted(self) -> Transform[O, I]:
279 """Return the inverse of this transform."""
280 return Transform[O, I](
281 self._out_frame,
282 self._in_frame,
283 self._ast_mapping.inverted(),
284 in_bounds=self.out_bounds,
285 out_bounds=self.in_bounds,
286 components=[t.inverted() for t in reversed(self._components)],
287 )
289 def then[F: Frame](self, next: Transform[O, F], remember_components: bool = True) -> Transform[I, F]:
290 """Compose two transforms into another.
292 Parameters
293 ----------
294 next
295 Another transform to apply after ``self``.
296 remember_components
297 If `True`, the returned composed transform will remember ``self``
298 and ``other`` so they can be returned by `decompose`.
299 """
300 if self._out_frame != next._in_frame:
301 raise TransformCompositionError(
302 "Cannot compose transforms that do not share a common intermediate frame: "
303 f"{self._out_frame} != {next._in_frame}."
304 )
305 components = self.decompose() + next.decompose() if remember_components else ()
306 return Transform(
307 self._in_frame,
308 next._out_frame,
309 self._ast_mapping.then(next._ast_mapping),
310 in_bounds=self.in_bounds,
311 out_bounds=next.out_bounds,
312 components=components,
313 )
315 def as_projection(self: Transform[I, SkyFrame]) -> Projection[I]:
316 """Return a `Projection` view of this transform.
318 This is only valid when `out_frame` is `~SkyFrame.ICRS`.
319 """
320 from ._projection import Projection
322 return Projection(self)
324 def as_fits_wcs(self, bbox: Box) -> astropy.wcs.WCS | None:
325 """Return a FITS WCS representation of this transform, if possible.
327 Parameters
328 ----------
329 bbox
330 Bounding box of the array the FITS WCS will describe. This
331 transform object is assumed to work on the same coordinate system
332 in which ``bbox`` is defined, while the FITS WCS will consider the
333 first row and column in that box to be ``(0, 0)`` (in Astropy
334 interfaces) or ``(1, 1)`` (in the FITS representation itself).
336 Notes
337 -----
338 This method assumes the transform maps pixel coordinates to world
339 coordinates.
341 Not all transforms can be represented exactly; when a FITS
342 represention is not possible, `None` is returned. When the returned
343 WCS is not `None`, it will have the same functional form, but it may
344 not evaluate identically due to small implementation differences in
345 the order of floating-point operations.
346 """
347 import astshim
349 ast_frame_set = self._get_ast_frame_set()
350 _prepend_ast_shift(ast_frame_set, x=1.0 - bbox.x.start, y=1.0 - bbox.y.start, ast_domain="GRID")
351 ast_stream = astshim.StringStream()
352 ast_fits_chan = astshim.FitsChan(
353 ast_stream, "Encoding=FITS-WCS, CDMatrix=1, FitsAxisOrder=<copy>, FitsTol=0.0001"
354 )
355 ast_fits_chan.setFitsI("NAXIS1", bbox.x.size)
356 ast_fits_chan.setFitsI("NAXIS2", bbox.y.size)
357 n_writes = ast_fits_chan.write(ast_frame_set)
358 if not n_writes:
359 return None
360 header = astropy.io.fits.Header(astropy.io.fits.Card.fromstring(c) for c in ast_fits_chan)
361 return astropy.wcs.WCS(header)
363 def serialize[P: pydantic.BaseModel](
364 self, archive: OutputArchive[P], *, use_frame_sets: bool = False
365 ) -> TransformSerializationModel[P]:
366 """Serialize a transform to an archive.
368 Parameters
369 ----------
370 archive
371 Archive to serialize to.
372 use_frame_sets
373 If `True`, decompose the transform and try to reference component
374 mappings that were already serialized into a `FrameSet` in the
375 archive. Note that if multiple transforms exist between a pair of
376 frames (e.g. a `Projection` and its FITS approximation), this may
377 cause the wrong one to be saved. When this option is used, the
378 frame set must be saved before the transform, and it must be
379 deserialized before the transform as well.
381 Returns
382 -------
383 `TransformSerializationModel`
384 Serialized form of the transform.
385 """
386 model = TransformSerializationModel[P]()
387 if use_frame_sets:
388 for link in self.decompose():
389 model.frames.append(link.in_frame.serialize())
390 model.bounds.append(link.in_bounds.serialize() if link.in_bounds is not None else None)
391 for frame_set, pointer in archive.iter_frame_sets():
392 if link.in_frame in frame_set and link.out_frame in frame_set:
393 model.mappings.append(pointer)
394 break
395 else:
396 model.mappings.append(MappingSerializationModel(ast=link._ast_mapping.show()))
397 else:
398 model.frames.append(self.in_frame.serialize())
399 model.bounds.append(self.in_bounds.serialize() if self.in_bounds is not None else None)
400 model.mappings.append(MappingSerializationModel(ast=self._ast_mapping.show()))
401 model.frames.append(self.out_frame.serialize())
402 model.bounds.append(self.out_bounds.serialize() if self.out_bounds is not None else None)
403 return model
405 @staticmethod
406 def deserialize[P: pydantic.BaseModel](
407 model: TransformSerializationModel[P], archive: InputArchive[P]
408 ) -> Transform[Any, Any]:
409 """Deserialize a transform from an archive.
411 Parameters
412 ----------
413 model
414 Seralized form of the transform.
415 archive
416 Archive to read from.
417 """
418 import astshim
420 if len(model.frames) != len(model.bounds):
421 raise ArchiveReadError(
422 f"Inconsistent lengths for 'frames' ({len(model.frames)}) and 'bounds' ({len(model.bounds)})."
423 )
424 if len(model.frames) != len(model.mappings) + 1:
425 raise ArchiveReadError(
426 f"Inconsistent lengths for 'frames' ({len(model.frames)}) and "
427 f"'mappings' ({len(model.mappings)}; should be one less)."
428 )
429 transform = Transform.identity(Frame.deserialize(model.frames[0]))
430 for n, mapping in enumerate(model.mappings):
431 match mapping:
432 case MappingSerializationModel(ast=serialized_mapping):
433 ast_mapping = astshim.Mapping.fromString(serialized_mapping)
434 in_bounds = model.bounds[n]
435 out_bounds = model.bounds[n + 1]
436 transform = transform.then(
437 Transform(
438 Frame.deserialize(model.frames[n]),
439 Frame.deserialize(model.frames[n + 1]),
440 ast_mapping,
441 Bounds.deserialize(in_bounds) if in_bounds is not None else None,
442 Bounds.deserialize(out_bounds) if out_bounds is not None else None,
443 )
444 )
445 case reference:
446 frame_set = archive.get_frame_set(reference)
447 transform = transform.then(
448 frame_set[Frame.deserialize(model.frames[n]), Frame.deserialize(model.frames[n + 1])]
449 )
450 return transform
452 @staticmethod
453 def _get_archive_tree_type[P: pydantic.BaseModel](
454 pointer_type: type[P],
455 ) -> type[TransformSerializationModel[P]]:
456 """Return the serialization model type for this object for an archive
457 type that uses the given pointer type.
458 """
459 return TransformSerializationModel[pointer_type] # type: ignore
461 @staticmethod
462 def from_legacy(
463 legacy: Any,
464 in_frame: I,
465 out_frame: O,
466 in_bounds: Bounds | None = None,
467 out_bounds: Bounds | None = None,
468 ) -> Transform[I, O]:
469 """Construct a transform from a legacy `lsst.afw.geom.Transform`.
471 Parameters
472 ----------
473 legacy : `lsst.afw.geom.Transform`
474 Legacy transform object.
475 in_frame
476 Coordinate frame for input points to the forward transform.
477 out_frame
478 Coordinate frame for output points from the forward transform.
479 in_bounds
480 The region that bounds valid input points.
481 out_bounds
482 The region that bounds valid output points.
483 """
484 return Transform(
485 in_frame,
486 out_frame,
487 legacy.getMapping(),
488 in_bounds=in_bounds,
489 out_bounds=out_bounds,
490 )
492 def _get_ast_frame_set(self) -> Any:
493 import astshim
495 ast_frame_set = astshim.FrameSet(_make_ast_frame(self._in_frame))
496 ast_frame_set.addFrame(astshim.FrameSet.BASE, self._ast_mapping, _make_ast_frame(self._out_frame))
497 return ast_frame_set
500def _ast_apply[T: np.ndarray | float](method: Any, *, x: T, y: T) -> XY[T]:
501 # TODO: add bounds argument and check inputs
502 # TODO: broadcast arrays with different shapes.
503 xy_in = np.vstack([x, y]).astype(np.float64)
504 xy_out = method(xy_in)
505 return XY(xy_out[0, :], xy_out[1, :])
508def _prepend_ast_shift(ast_frame_set: Any, x: float, y: float, ast_domain: str) -> None:
509 import astshim
511 ast_output_frame_id = ast_frame_set.current
512 ast_frame_set.addFrame(
513 astshim.FrameSet.BASE,
514 astshim.ShiftMap([x, y]),
515 astshim.Frame(2, f"Domain={ast_domain}"),
516 )
517 ast_frame_set.base = ast_frame_set.current
518 ast_frame_set.current = ast_output_frame_id
521def _make_ast_frame(frame: Frame) -> Any:
522 import astshim
524 if frame is SkyFrame.ICRS:
525 return astshim.SkyFrame("")
526 ast_frame = astshim.Frame(2, f"Ident={frame._ast_ident}")
527 if frame.unit is not None:
528 fits_unit = frame.unit.to_string(format="fits")
529 ast_frame.setUnit(1, fits_unit)
530 ast_frame.setUnit(2, fits_unit)
531 ast_frame.setLabel(1, "x")
532 ast_frame.setLabel(2, "y")
533 return ast_frame
536def _standardize_xy[T: np.ndarray | float](xy: XY[T], frame: Frame) -> XY[T]:
537 return XY(x=frame.standardize_x(xy.x), y=frame.standardize_y(xy.y))
540class MappingSerializationModel(ArchiveTree):
541 """Serialization model for an AST Mapping."""
543 ast: str = pydantic.Field(description="A serialized Starlink AST Mapping, using the AST native encoding.")
546class TransformSerializationModel[P: pydantic.BaseModel](ArchiveTree):
547 """Serialization model for coordinate transforms."""
549 frames: list[SerializableFrame] = pydantic.Field(
550 default_factory=list,
551 description=textwrap.dedent(
552 """
553 List of frames that this transform passes through.
555 All transforms include at least two frames (the endpoints). Others
556 intermediate frames may be included to facilitate data-sharing
557 between transforms.
558 """
559 ),
560 )
562 bounds: list[SerializableBounds | None] = pydantic.Field(
563 default_factory=list,
564 description=textwrap.dedent(
565 """
566 List of the bounds of the ``frames`` for this transform.
568 This always has the same number of elements as ``frames``.
569 """
570 ),
571 )
573 mappings: list[P | MappingSerializationModel] = pydantic.Field(
574 default_factory=list,
575 description=textwrap.dedent(
576 """
577 The actual mappings between frames, or archive pointers to
578 serialized FrameSet objects from which they can be obtained.
580 This always has one fewer element than ``frames``.
581 """
582 ),
583 )