Coverage for python/lsst/images/psfs/_piff.py: 67%
200 statements
« prev ^ index » next coverage.py v7.14.2, created at 2026-06-21 08:57 +0000
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__ = ("PiffSerializationModel", "PiffWrapper")
16import operator
17from collections.abc import Iterator
18from contextlib import contextmanager
19from functools import cached_property
20from logging import getLogger
21from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal
23import astropy.io.fits
24import numpy as np
25import pydantic
27from .. import serialization
28from .._concrete_bounds import SerializableBounds
29from .._geom import XY, YX, Bounds, Box
30from .._image import Image
31from ..utils import round_half_up
32from ._base import PointSpreadFunction
34if TYPE_CHECKING:
35 import galsim.wcs
36 import piff
37 import piff.config
39 try:
40 from lsst.meas.extensions.piff.piffPsf import PiffPsf as LegacyPiffPsf
41 except ImportError:
42 type LegacyPiffPsf = Any # type: ignore[no-redef]
45_LOG = getLogger(__name__)
48class PiffWrapper(PointSpreadFunction):
49 """A PSF model backed by the Piff library.
51 Parameters
52 ----------
53 impl
54 The Piff PSF object to wrap.
55 bounds
56 The pixel-coordinate region where the model can safely be evaluated.
57 """
59 def __init__(self, impl: piff.PSF, bounds: Bounds, stamp_size: int):
60 self._impl = impl
61 self._bounds = bounds
62 self._stamp_size = stamp_size
64 @property
65 def bounds(self) -> Bounds:
66 return self._bounds
68 @cached_property
69 def kernel_bbox(self) -> Box:
70 r = self._stamp_size // 2
71 return Box.factory[-r : r + 1, -r : r + 1]
73 def compute_kernel_image(self, *, x: float, y: float) -> Image:
74 if "colorValue" in self._impl.interp_property_names:
75 raise NotImplementedError("Chromatic PSFs are not yet supported.")
76 gs_image = self._impl.draw(x, y, stamp_size=self._stamp_size, center=True)
77 r = self._stamp_size // 2
78 result = Image(gs_image.array.copy(), yx0=YX(y=-r, x=-r))
79 result.array /= np.sum(result.array)
80 return result
82 def compute_stellar_image(self, *, x: float, y: float) -> Image:
83 if "colorValue" in self._impl.interp_property_names:
84 raise NotImplementedError("Chromatic PSFs are not yet supported.")
85 gs_image = self._impl.draw(x, y, stamp_size=self._stamp_size, center=None)
86 r = self._stamp_size // 2
87 result = Image(gs_image.array.copy(), yx0=YX(y=round_half_up(y) - r, x=round_half_up(x) - r))
88 result.array /= np.sum(result.array)
89 return result
91 def compute_stellar_bbox(self, *, x: float, y: float) -> Box:
92 r = self._stamp_size // 2
93 xi = round_half_up(x)
94 yi = round_half_up(y)
95 return Box.factory[yi - r : yi + r + 1, xi - r : xi + r + 1]
97 @property
98 def piff_psf(self) -> piff.PSF:
99 """The backing `piff.PSF` object.
101 This is an internal object that must not be modified in place.
102 """
103 return self._impl
105 @classmethod
106 def from_legacy(cls, legacy_psf: LegacyPiffPsf, bounds: Bounds) -> PiffWrapper:
107 return cls(impl=legacy_psf._piffResult, bounds=bounds, stamp_size=int(legacy_psf.width))
109 def to_legacy(self) -> LegacyPiffPsf:
110 """Convert to a legacy `lsst.meas.extensions.piff.piffPsf`."""
111 from lsst.meas.extensions.piff.piffPsf import PiffPsf as LegacyPiffPsf
113 return LegacyPiffPsf(self._stamp_size, self._stamp_size, self._impl)
115 def serialize(self, archive: serialization.OutputArchive[Any]) -> PiffSerializationModel:
116 """Serialize the PSF to an archive.
118 This method is intended to be usable as the callback function passed to
119 `.serialization.OutputArchive.serialize_direct` or
120 `.serialization.OutputArchive.serialize_pointer`.
121 """
122 from piff.config import PiffLogger
124 writer = _ArchivePiffWriter()
125 with self._without_stars():
126 self._impl._write(writer, "piff", PiffLogger(_LOG))
127 piff_model = writer.serialize(archive)
128 return PiffSerializationModel(
129 piff=piff_model,
130 stamp_size=self._stamp_size,
131 bounds=self._bounds.serialize(),
132 stars=[MinimalStar.from_star(s) for s in self._impl.stars],
133 )
135 @staticmethod
136 def _get_archive_tree_type(
137 pointer_type: type[pydantic.BaseModel],
138 ) -> type[PiffSerializationModel]:
139 """Return the serialization model type for this object for an archive
140 type that uses the given pointer type.
141 """
142 return PiffSerializationModel
144 @contextmanager
145 def _without_stars(self) -> Iterator[None]:
146 """Temporarily drop the embedded list of stars used to fit the PSF.
148 Notes
149 -----
150 By default Piff saves the list of stars (including postage stamps) used
151 to fit the PSF, which makes the serialized form much larger. But the
152 upstream Piff serialization code recognizes the case where that
153 ``stars`` attribute has been deleted and serializes everything else.
155 Unfortunately, to date, Rubin's pickle-based Piff serialization instead
156 just deletes the postage stamp image attributes from inside the Piff
157 ``stars`` list, which is not a state the Piff serialization code
158 handles gracefully. So for now we have to drop the full stars list
159 during serialization if it is present. We then save the star
160 positions separately.
161 """
162 if hasattr(self._impl, "stars"):
163 stars = self._impl.stars
164 try:
165 del self._impl.stars
166 yield
167 finally:
168 self._impl.stars = stars
169 else:
170 yield
173class MinimalStar(pydantic.BaseModel):
174 """A partial duck-alike for `piff.star.Star`, holding just the image
175 position and some booleans (enough to compute an 'average position' on the
176 legacy PSF).
177 """
179 image_pos: XY
180 is_flagged: bool
181 is_reserve: bool
183 @classmethod
184 def from_star(cls, star: MinimalStar | piff.star.Star) -> MinimalStar:
185 if type(star) is cls:
186 return star
187 return cls(
188 image_pos=XY(x=star.image_pos.x, y=star.image_pos.y),
189 is_flagged=star.is_flagged,
190 is_reserve=star.is_reserve,
191 )
194# Conventions on public visibility of the serialization types:
195#
196# - We lift and document the outermost Pydantic model type, since that needs to
197# be included directly in the Pydantic models of types that hold a PSF. This
198# type needs to be very clearly documented and named as a *serialization*
199# model, since there are many other kinds of models in play in this package.
200#
201# - We do not lift or document types used in that outermost model, but we do
202# not give them leading underscores, since they aren't really private.
203#
204# - Other utility types do get leading underscores.
207# Piff serialization uses a lot of dictionaries and lists restricted to these
208# basic types.
209type PiffScalar = int | float | str | bool | None
210type PiffValue = PiffScalar | list[PiffValue]
211type PiffDict = dict[str, PiffValue]
214class GalSimPixelScaleModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
215 """Model used to serialize `galsim.wcs.PixelScale` instances."""
217 scale: float
218 wcs_type: Literal["pixel_scale"] = "pixel_scale"
221# We expect this discriminated union to grow to include other trivial
222# pixel-to-pixel transforms that get embedded in PSFs. If we someday have to
223# store Piff objects that embed more sophisticated PSFs, we'll hook them into
224# the AST-based coordinate transform system instead, but as long as we're just
225# talking about simple offsets and scalings, that's a lot of extra complexity
226# for very little gain.
227type GalSimLocalWcsModel = Annotated[GalSimPixelScaleModel, pydantic.Field(discriminator="wcs_type")]
230class PiffTableModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
231 """Serialization model used to embed a reference to a binary-data table in
232 a Piff serialization's JSON-like data.
233 """
235 metadata: PiffDict
236 table: serialization.TableModel
239class PiffObjectModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
240 """General-purpose serialization model used for various Piff objects."""
242 structs: dict[str, PiffDict] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
243 tables: dict[str, PiffTableModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
244 wcs: dict[str, GalSimLocalWcsModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
245 objects: dict[str, PiffObjectModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
248class PiffSerializationModel(serialization.ArchiveTree):
249 """Serialization model for a Piff PSF."""
251 SCHEMA_NAME: ClassVar[str] = "piff_psf"
252 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
253 MIN_READ_VERSION: ClassVar[int] = 1
254 PUBLIC_TYPE: ClassVar[type] = PiffWrapper
256 piff: PiffObjectModel = pydantic.Field(description="The Piff PSF object itself.")
258 stars: list[MinimalStar] = pydantic.Field(
259 description="Minimal information about the stars that went into the PSF model."
260 )
262 stamp_size: int = pydantic.Field(
263 description="Width of the (square) images returned by this PSF's methods."
264 )
266 bounds: SerializableBounds = pydantic.Field(
267 description="The bounds object that represents the PSF's validity region."
268 )
270 def deserialize(self, archive: serialization.InputArchive[Any], **kwargs: Any) -> PiffWrapper:
271 """Deserialize the PSF from an archive.
273 This method is intended to be usable as the callback function passed to
274 `.serialization.InputArchive.deserialize_pointer`.
275 """
276 if kwargs: 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true
277 raise serialization.InvalidParameterError(
278 f"Unrecognized parameters for PiffWrapper: {set(kwargs.keys())}."
279 )
280 try:
281 from piff import PSF
282 from piff.config import PiffLogger
283 except ImportError:
284 raise serialization.ArchiveReadError("Failed to import piff.") from None
286 reader = _ArchivePiffReader(self.piff, archive)
287 impl = PSF._read(reader, "piff", PiffLogger(_LOG))
288 impl.stars = self.stars
289 return PiffWrapper(impl, bounds=self.bounds.deserialize(), stamp_size=self.stamp_size)
292class _ArchivePiffWriter:
293 """An adapter from the Piff serialization interface to the
294 `.serialization.OutputArchive` class.
296 Notes
297 -----
298 Piff has its own simple serialization framework (contributed upstream by
299 Rubin DM) that maps everything to dictionaries, structured numpy arrays,
300 and a library of GalSim WCS objects, with the native implementation writing
301 standalone FITS files. That mostly maps nicely to the `lsst.images`
302 archive system, but we don't get to leverage any Pydantic validation or
303 JSON schema functionality since we only get opaque dictionaries from Piff.
305 See `piff.FitsWriter` for most method documentation; this class is designed
306 to mimic it exactly (the Piff authors prefer to just use duck-typing rather
307 than ABCs or protocols for interface definition).
308 """
310 def __init__(self, base_name: str = ""):
311 self._base_name = base_name
312 self.structs: dict[str, PiffDict] = {}
313 self.tables: dict[str, tuple[np.ndarray, PiffDict]] = {}
314 self.wcs_models: dict[str, GalSimLocalWcsModel] = {}
315 self.writers: dict[str, _ArchivePiffWriter] = {}
317 def write_struct(self, name: str, struct: PiffDict) -> None:
318 self.structs[name] = {k: self._to_builtin(v) for k, v in struct.items()}
320 def write_table(self, name: str, array: np.ndarray, metadata: PiffDict | None = None) -> None:
321 self.tables[name] = (
322 array,
323 {k: self._to_builtin(v) for k, v in (metadata or {}).items()},
324 )
326 def write_wcs_map(
327 self, name: str, wcs_map: dict[int, galsim.wcs.BaseWCS], pointing: galsim.CelestialCoord | None
328 ) -> None:
329 import galsim.wcs
331 match wcs_map:
332 case {0: galsim.wcs.PixelScale() as wcs} if pointing is None:
333 self.wcs_models[name] = GalSimPixelScaleModel(scale=wcs.scale)
334 case _:
335 raise NotImplementedError("PSFs with complex embedded WCSs are not supported.")
337 @contextmanager
338 def nested(self, name: str) -> Iterator[_ArchivePiffWriter]:
339 nested = _ArchivePiffWriter(self.get_full_name(name))
340 yield nested
341 self.writers[name] = nested
343 def get_full_name(self, name: str) -> str:
344 return f"{self._base_name}/{name}"
346 def serialize(self, archive: serialization.OutputArchive[Any]) -> PiffObjectModel:
347 """Serialize to an archive.
349 This method is intended to be used as the callable passed to
350 `.serialization.OutputArchive.serialize_direct` and
351 `.serialization.OutputArchive.serialize_pointer`, after first passing
352 this writer to a Piff object's ``write`` or ``_write`` method.
353 """
354 model = PiffObjectModel()
355 for name, struct in self.structs.items():
356 model.structs[name] = struct
357 for name, (array, metadata) in self.tables.items(): 357 ↛ 358line 357 didn't jump to line 358 because the loop on line 357 never started
358 model.tables[name] = PiffTableModel(
359 metadata=metadata,
360 table=archive.add_structured_array(
361 array, name=name, update_header=lambda header: header.update(metadata)
362 ),
363 )
364 for name, wcs_model in self.wcs_models.items(): 364 ↛ 365line 364 didn't jump to line 365 because the loop on line 364 never started
365 model.wcs[name] = wcs_model
366 for name, writer in self.writers.items(): 366 ↛ 367line 366 didn't jump to line 367 because the loop on line 366 never started
367 model.objects[name] = archive.serialize_direct(name, writer.serialize)
368 return model
370 @staticmethod
371 def _to_builtin(val: Any) -> PiffValue:
372 match val:
373 case np.integer():
374 return int(val)
375 case np.floating():
376 return float(val)
377 case np.bool_():
378 return bool(val)
379 case np.str_(): 379 ↛ 380line 379 didn't jump to line 380 because the pattern on line 379 never matched
380 return str(val)
381 case tuple() | list():
382 return [_ArchivePiffWriter._to_builtin(item) for item in val]
383 return val
386class _ArchivePiffReader:
387 """An adapter from the Piff serialization interface to the
388 `.serialization.InputArchive` class.
390 See `ArchivePiffWriter` for additional notes.
391 """
393 def __init__(
394 self, object_model: PiffObjectModel, archive: serialization.InputArchive[Any], base_name: str = ""
395 ):
396 self._model = object_model
397 self._archive = archive
398 self._base_name = base_name
400 def read_struct(self, name: str) -> PiffDict | None:
401 return self._model.structs.get(name)
403 def read_table(self, name: str, metadata: PiffDict | None = None) -> np.ndarray | None:
404 table_model = self._model.tables.get(name)
405 if table_model is None:
406 return None
407 if metadata is not None: 407 ↛ 408line 407 didn't jump to line 408 because the condition on line 407 was never true
408 metadata.update(table_model.metadata)
409 return self._archive.get_structured_array(
410 table_model.table, strip_header=astropy.io.fits.Header.clear
411 )
413 def read_wcs_map(
414 self, name: str, logger: piff.config.LoggerWrapper
415 ) -> tuple[dict[int, galsim.wcs.BaseWCS] | None, galsim.CelestialCoord | None]:
416 import galsim.wcs
418 match self._model.wcs.get(name):
419 case GalSimPixelScaleModel(scale=scale): 419 ↛ 421line 419 didn't jump to line 421 because the pattern on line 419 always matched
420 return {0: galsim.wcs.PixelScale(scale)}, None
421 case None:
422 return None, None
423 case unexpected:
424 raise serialization.ArchiveReadError(
425 f"{self.get_full_name(name)} should be a WCS or WCS map, not {unexpected!r}."
426 )
428 @contextmanager
429 def nested(self, name: str) -> Iterator[_ArchivePiffReader]:
430 nested_model = self._model.objects[name]
431 yield _ArchivePiffReader(nested_model, self._archive, self.get_full_name(name))
433 def get_full_name(self, name: str) -> str:
434 return f"{self._base_name}/{name}"