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