Coverage for python/lsst/images/psfs/_piff.py: 67%
200 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-29 22:40 +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 stamp_size
58 Side length in pixels of the PSF image stamps drawn by the model.
59 """
61 def __init__(self, impl: piff.PSF, bounds: Bounds, stamp_size: int) -> None:
62 self._impl = impl
63 self._bounds = bounds
64 self._stamp_size = stamp_size
66 @property
67 def bounds(self) -> Bounds:
68 return self._bounds
70 @cached_property
71 def kernel_bbox(self) -> Box:
72 r = self._stamp_size // 2
73 return Box.factory[-r : r + 1, -r : r + 1]
75 def compute_kernel_image(self, *, x: float, y: float) -> Image:
76 if "colorValue" in self._impl.interp_property_names:
77 raise NotImplementedError("Chromatic PSFs are not yet supported.")
78 gs_image = self._impl.draw(x, y, stamp_size=self._stamp_size, center=True)
79 r = self._stamp_size // 2
80 result = Image(gs_image.array.copy(), yx0=YX(y=-r, x=-r))
81 result.array /= np.sum(result.array)
82 return result
84 def compute_stellar_image(self, *, x: float, y: float) -> Image:
85 if "colorValue" in self._impl.interp_property_names:
86 raise NotImplementedError("Chromatic PSFs are not yet supported.")
87 gs_image = self._impl.draw(x, y, stamp_size=self._stamp_size, center=None)
88 r = self._stamp_size // 2
89 result = Image(gs_image.array.copy(), yx0=YX(y=round_half_up(y) - r, x=round_half_up(x) - r))
90 result.array /= np.sum(result.array)
91 return result
93 def compute_stellar_bbox(self, *, x: float, y: float) -> Box:
94 r = self._stamp_size // 2
95 xi = round_half_up(x)
96 yi = round_half_up(y)
97 return Box.factory[yi - r : yi + r + 1, xi - r : xi + r + 1]
99 @property
100 def piff_psf(self) -> piff.PSF:
101 """The backing `piff.PSF` object.
103 This is an internal object that must not be modified in place.
104 """
105 return self._impl
107 @classmethod
108 def from_legacy(cls, legacy_psf: LegacyPiffPsf, bounds: Bounds) -> PiffWrapper:
109 return cls(impl=legacy_psf._piffResult, bounds=bounds, stamp_size=int(legacy_psf.width))
111 def to_legacy(self) -> LegacyPiffPsf:
112 """Convert to a legacy `lsst.meas.extensions.piff.piffPsf`."""
113 from lsst.meas.extensions.piff.piffPsf import PiffPsf as LegacyPiffPsf
115 return LegacyPiffPsf(self._stamp_size, self._stamp_size, self._impl)
117 def serialize(self, archive: serialization.OutputArchive[Any]) -> PiffSerializationModel:
118 """Serialize the PSF to an archive.
120 This method is intended to be usable as the callback function passed to
121 `.serialization.OutputArchive.serialize_direct` or
122 `.serialization.OutputArchive.serialize_pointer`.
124 Parameters
125 ----------
126 archive
127 Archive to write to.
128 """
129 from piff.config import PiffLogger
131 writer = _ArchivePiffWriter()
132 with self._without_stars():
133 self._impl._write(writer, "piff", PiffLogger(_LOG))
134 piff_model = writer.serialize(archive)
135 return PiffSerializationModel(
136 piff=piff_model,
137 stamp_size=self._stamp_size,
138 bounds=self._bounds.serialize(),
139 stars=[MinimalStar.from_star(s) for s in self._impl.stars],
140 )
142 @staticmethod
143 def _get_archive_tree_type(
144 pointer_type: type[pydantic.BaseModel],
145 ) -> type[PiffSerializationModel]:
146 """Return the serialization model type for this object for an archive
147 type that uses the given pointer type.
148 """
149 return PiffSerializationModel
151 @contextmanager
152 def _without_stars(self) -> Iterator[None]:
153 """Temporarily drop the embedded list of stars used to fit the PSF.
155 Notes
156 -----
157 By default Piff saves the list of stars (including postage stamps) used
158 to fit the PSF, which makes the serialized form much larger. But the
159 upstream Piff serialization code recognizes the case where that
160 ``stars`` attribute has been deleted and serializes everything else.
162 Unfortunately, to date, Rubin's pickle-based Piff serialization instead
163 just deletes the postage stamp image attributes from inside the Piff
164 ``stars`` list, which is not a state the Piff serialization code
165 handles gracefully. So for now we have to drop the full stars list
166 during serialization if it is present. We then save the star
167 positions separately.
168 """
169 if hasattr(self._impl, "stars"):
170 stars = self._impl.stars
171 try:
172 del self._impl.stars
173 yield
174 finally:
175 self._impl.stars = stars
176 else:
177 yield
180class MinimalStar(pydantic.BaseModel):
181 """A partial duck-alike for `piff.star.Star`, holding just the image
182 position and some booleans (enough to compute an 'average position' on the
183 legacy PSF).
184 """
186 image_pos: XY
187 is_flagged: bool
188 is_reserve: bool
190 @classmethod
191 def from_star(cls, star: MinimalStar | piff.star.Star) -> MinimalStar:
192 if type(star) is cls:
193 return star
194 return cls(
195 image_pos=XY(x=star.image_pos.x, y=star.image_pos.y),
196 is_flagged=star.is_flagged,
197 is_reserve=star.is_reserve,
198 )
201# Conventions on public visibility of the serialization types:
202#
203# - We lift and document the outermost Pydantic model type, since that needs to
204# be included directly in the Pydantic models of types that hold a PSF. This
205# type needs to be very clearly documented and named as a *serialization*
206# model, since there are many other kinds of models in play in this package.
207#
208# - We do not lift or document types used in that outermost model, but we do
209# not give them leading underscores, since they aren't really private.
210#
211# - Other utility types do get leading underscores.
214# Piff serialization uses a lot of dictionaries and lists restricted to these
215# basic types.
216type PiffScalar = int | float | str | bool | None
217type PiffValue = PiffScalar | list[PiffValue]
218type PiffDict = dict[str, PiffValue]
221class GalSimPixelScaleModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
222 """Model used to serialize `galsim.wcs.PixelScale` instances."""
224 scale: float
225 wcs_type: Literal["pixel_scale"] = "pixel_scale"
228# We expect this discriminated union to grow to include other trivial
229# pixel-to-pixel transforms that get embedded in PSFs. If we someday have to
230# store Piff objects that embed more sophisticated PSFs, we'll hook them into
231# the AST-based coordinate transform system instead, but as long as we're just
232# talking about simple offsets and scalings, that's a lot of extra complexity
233# for very little gain.
234type GalSimLocalWcsModel = Annotated[GalSimPixelScaleModel, pydantic.Field(discriminator="wcs_type")]
237class PiffTableModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
238 """Serialization model used to embed a reference to a binary-data table in
239 a Piff serialization's JSON-like data.
240 """
242 metadata: PiffDict
243 table: serialization.TableModel
246class PiffObjectModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
247 """General-purpose serialization model used for various Piff objects."""
249 structs: dict[str, PiffDict] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
250 tables: dict[str, PiffTableModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
251 wcs: dict[str, GalSimLocalWcsModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
252 objects: dict[str, PiffObjectModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
255class PiffSerializationModel(serialization.ArchiveTree):
256 """Serialization model for a Piff PSF."""
258 SCHEMA_NAME: ClassVar[str] = "piff_psf"
259 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
260 MIN_READ_VERSION: ClassVar[int] = 1
261 PUBLIC_TYPE: ClassVar[type] = PiffWrapper
263 piff: PiffObjectModel = pydantic.Field(description="The Piff PSF object itself.")
265 stars: list[MinimalStar] = pydantic.Field(
266 description="Minimal information about the stars that went into the PSF model."
267 )
269 stamp_size: int = pydantic.Field(
270 description="Width of the (square) images returned by this PSF's methods."
271 )
273 bounds: SerializableBounds = pydantic.Field(
274 description="The bounds object that represents the PSF's validity region."
275 )
277 def deserialize(self, archive: serialization.InputArchive[Any], **kwargs: Any) -> PiffWrapper:
278 """Deserialize the PSF from an archive.
280 This method is intended to be usable as the callback function passed to
281 `.serialization.InputArchive.deserialize_pointer`.
283 Parameters
284 ----------
285 archive
286 Archive to read from.
287 **kwargs
288 Unsupported keyword arguments are accepted only to provide
289 better error messages (raising
290 `.serialization.InvalidParameterError`).
291 """
292 if kwargs: 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true
293 raise serialization.InvalidParameterError(
294 f"Unrecognized parameters for PiffWrapper: {set(kwargs.keys())}."
295 )
296 try:
297 from piff import PSF
298 from piff.config import PiffLogger
299 except ImportError:
300 raise serialization.ArchiveReadError("Failed to import piff.") from None
302 reader = _ArchivePiffReader(self.piff, archive)
303 impl = PSF._read(reader, "piff", PiffLogger(_LOG))
304 impl.stars = self.stars
305 return PiffWrapper(impl, bounds=self.bounds.deserialize(), stamp_size=self.stamp_size)
308class _ArchivePiffWriter:
309 """An adapter from the Piff serialization interface to the
310 `.serialization.OutputArchive` class.
312 Notes
313 -----
314 Piff has its own simple serialization framework (contributed upstream by
315 Rubin DM) that maps everything to dictionaries, structured numpy arrays,
316 and a library of GalSim WCS objects, with the native implementation writing
317 standalone FITS files. That mostly maps nicely to the `lsst.images`
318 archive system, but we don't get to leverage any Pydantic validation or
319 JSON schema functionality since we only get opaque dictionaries from Piff.
321 See `piff.FitsWriter` for most method documentation; this class is designed
322 to mimic it exactly (the Piff authors prefer to just use duck-typing rather
323 than ABCs or protocols for interface definition).
324 """
326 def __init__(self, base_name: str = "") -> None:
327 self._base_name = base_name
328 self.structs: dict[str, PiffDict] = {}
329 self.tables: dict[str, tuple[np.ndarray, PiffDict]] = {}
330 self.wcs_models: dict[str, GalSimLocalWcsModel] = {}
331 self.writers: dict[str, _ArchivePiffWriter] = {}
333 def write_struct(self, name: str, struct: PiffDict) -> None:
334 self.structs[name] = {k: self._to_builtin(v) for k, v in struct.items()}
336 def write_table(self, name: str, array: np.ndarray, metadata: PiffDict | None = None) -> None:
337 self.tables[name] = (
338 array,
339 {k: self._to_builtin(v) for k, v in (metadata or {}).items()},
340 )
342 def write_wcs_map(
343 self, name: str, wcs_map: dict[int, galsim.wcs.BaseWCS], pointing: galsim.CelestialCoord | None
344 ) -> None:
345 import galsim.wcs
347 match wcs_map:
348 case {0: galsim.wcs.PixelScale() as wcs} if pointing is None:
349 self.wcs_models[name] = GalSimPixelScaleModel(scale=wcs.scale)
350 case _:
351 raise NotImplementedError("PSFs with complex embedded WCSs are not supported.")
353 @contextmanager
354 def nested(self, name: str) -> Iterator[_ArchivePiffWriter]:
355 nested = _ArchivePiffWriter(self.get_full_name(name))
356 yield nested
357 self.writers[name] = nested
359 def get_full_name(self, name: str) -> str:
360 return f"{self._base_name}/{name}"
362 def serialize(self, archive: serialization.OutputArchive[Any]) -> PiffObjectModel:
363 """Serialize to an archive.
365 This method is intended to be used as the callable passed to
366 `.serialization.OutputArchive.serialize_direct` and
367 `.serialization.OutputArchive.serialize_pointer`, after first passing
368 this writer to a Piff object's ``write`` or ``_write`` method.
370 Parameters
371 ----------
372 archive
373 Archive to write to.
374 """
375 model = PiffObjectModel()
376 for name, struct in self.structs.items():
377 model.structs[name] = struct
378 for name, (array, metadata) in self.tables.items(): 378 ↛ 379line 378 didn't jump to line 379 because the loop on line 378 never started
379 model.tables[name] = PiffTableModel(
380 metadata=metadata,
381 table=archive.add_structured_array(
382 array, name=name, update_header=lambda header: header.update(metadata)
383 ),
384 )
385 for name, wcs_model in self.wcs_models.items(): 385 ↛ 386line 385 didn't jump to line 386 because the loop on line 385 never started
386 model.wcs[name] = wcs_model
387 for name, writer in self.writers.items(): 387 ↛ 388line 387 didn't jump to line 388 because the loop on line 387 never started
388 model.objects[name] = archive.serialize_direct(name, writer.serialize)
389 return model
391 @staticmethod
392 def _to_builtin(val: Any) -> PiffValue:
393 match val:
394 case np.integer():
395 return int(val)
396 case np.floating():
397 return float(val)
398 case np.bool_():
399 return bool(val)
400 case np.str_(): 400 ↛ 401line 400 didn't jump to line 401 because the pattern on line 400 never matched
401 return str(val)
402 case tuple() | list():
403 return [_ArchivePiffWriter._to_builtin(item) for item in val]
404 return val
407class _ArchivePiffReader:
408 """An adapter from the Piff serialization interface to the
409 `.serialization.InputArchive` class.
411 See `ArchivePiffWriter` for additional notes.
412 """
414 def __init__(
415 self, object_model: PiffObjectModel, archive: serialization.InputArchive[Any], base_name: str = ""
416 ) -> None:
417 self._model = object_model
418 self._archive = archive
419 self._base_name = base_name
421 def read_struct(self, name: str) -> PiffDict | None:
422 return self._model.structs.get(name)
424 def read_table(self, name: str, metadata: PiffDict | None = None) -> np.ndarray | None:
425 table_model = self._model.tables.get(name)
426 if table_model is None:
427 return None
428 if metadata is not None: 428 ↛ 429line 428 didn't jump to line 429 because the condition on line 428 was never true
429 metadata.update(table_model.metadata)
430 return self._archive.get_structured_array(
431 table_model.table, strip_header=astropy.io.fits.Header.clear
432 )
434 def read_wcs_map(
435 self, name: str, logger: piff.config.LoggerWrapper
436 ) -> tuple[dict[int, galsim.wcs.BaseWCS] | None, galsim.CelestialCoord | None]:
437 import galsim.wcs
439 match self._model.wcs.get(name):
440 case GalSimPixelScaleModel(scale=scale): 440 ↛ 442line 440 didn't jump to line 442 because the pattern on line 440 always matched
441 return {0: galsim.wcs.PixelScale(scale)}, None
442 case None:
443 return None, None
444 case unexpected:
445 raise serialization.ArchiveReadError(
446 f"{self.get_full_name(name)} should be a WCS or WCS map, not {unexpected!r}."
447 )
449 @contextmanager
450 def nested(self, name: str) -> Iterator[_ArchivePiffReader]:
451 nested_model = self._model.objects[name]
452 yield _ArchivePiffReader(nested_model, self._archive, self.get_full_name(name))
454 def get_full_name(self, name: str) -> str:
455 return f"{self._base_name}/{name}"