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