Coverage for python / lsst / images / psfs / _piff.py: 37%
175 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-05 23:13 -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
14from lsst.images.serialization import ArchiveReadError
16__all__ = ("PiffSerializationModel", "PiffWrapper")
18import operator
19from collections.abc import Iterator
20from contextlib import contextmanager
21from functools import cached_property
22from logging import getLogger
23from typing import TYPE_CHECKING, Annotated, Any, Literal
25import astropy.io.fits
26import numpy as np
27import pydantic
29from .. import serialization
30from .._geom import Bounds, Box, SerializableBounds
31from .._image import Image
32from ._base import PointSpreadFunction
34if TYPE_CHECKING:
35 import galsim.wcs
36 import piff.config
38_LOG = getLogger(__name__)
41class PiffWrapper(PointSpreadFunction):
42 """A PSF model backed by the Piff library.
44 Parameters
45 ----------
46 impl
47 The Piff PSF object to wrap.
48 bounds
49 The pixel-coordinate region where the model can safely be evaluated.
50 """
52 def __init__(self, impl: piff.PSF, bounds: Bounds, stamp_size: int):
53 self._impl = impl
54 self._bounds = bounds
55 self._stamp_size = stamp_size
57 @property
58 def bounds(self) -> Bounds:
59 return self._bounds
61 @cached_property
62 def kernel_bbox(self) -> Box:
63 r = self._stamp_size // 2
64 return Box.factory[-r : r + 1, -r : r + 1]
66 def compute_kernel_image(self, *, x: float, y: float) -> Image:
67 if "colorValue" in self._impl.interp_property_names:
68 raise NotImplementedError("Chromatic PSFs are not yet supported.")
69 gs_image = self._impl.draw(x, y, stamp_size=self._stamp_size, center=True)
70 r = self._stamp_size // 2
71 result = Image(gs_image.array.copy(), start=(-r, -r))
72 result.array /= np.sum(result.array)
73 return result
75 def compute_stellar_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=None)
79 r = self._stamp_size // 2
80 result = Image(gs_image.array.copy(), start=(round(y) - r, round(x) - r))
81 result.array /= np.sum(result.array)
82 return result
84 def compute_stellar_bbox(self, *, x: float, y: float) -> Box:
85 r = self._stamp_size // 2
86 xi = round(x)
87 yi = round(y)
88 return Box.factory[yi - r : yi + r + 1, xi - r : xi + r + 1]
90 @property
91 def piff_psf(self) -> Any:
92 """The backing `piff.PSF` object.
94 This is an internal object that must not be modified in place.
95 """
96 return self._impl
98 @classmethod
99 def from_legacy(cls, legacy_psf: Any, bounds: Bounds) -> PiffWrapper:
100 return cls(impl=legacy_psf._piffResult, bounds=bounds, stamp_size=legacy_psf.width)
102 def serialize(self, archive: serialization.OutputArchive[Any]) -> PiffSerializationModel:
103 """Serialize the PSF to an archive.
105 This method is intended to be usable as the callback function passed to
106 `.serialization.OutputArchive.serialize_direct` or
107 `.serialization.OutputArchive.serialize_pointer`.
108 """
109 from piff.config import PiffLogger
111 writer = _ArchivePiffWriter()
112 with self._without_stars():
113 self._impl._write(writer, "piff", PiffLogger(_LOG))
114 piff_model = writer.serialize(archive)
115 return PiffSerializationModel(
116 piff=piff_model,
117 stamp_size=self._stamp_size,
118 bounds=self._bounds.serialize(),
119 )
121 @classmethod
122 def deserialize(
123 cls, model: PiffSerializationModel, archive: serialization.InputArchive[Any]
124 ) -> PiffWrapper:
125 """Deserialize the PSF from an archive.
127 This method is intended to be usable as the callback function passed to
128 `.serialization.InputArchive.deserialize_pointer`.
129 """
130 try:
131 from piff import PSF
132 from piff.config import PiffLogger
133 except ImportError:
134 raise ArchiveReadError("Failed to import piff.") from None
136 reader = _ArchivePiffReader(model.piff, archive)
137 impl = PSF._read(reader, "piff", PiffLogger(_LOG))
138 return cls(impl, bounds=Bounds.deserialize(model.bounds), stamp_size=model.stamp_size)
140 @staticmethod
141 def _get_archive_tree_type(
142 pointer_type: type[pydantic.BaseModel],
143 ) -> type[PiffSerializationModel]:
144 """Return the serialization model type for this object for an archive
145 type that uses the given pointer type.
146 """
147 return PiffSerializationModel
149 @contextmanager
150 def _without_stars(self) -> Iterator[None]:
151 """Temporarily drop the embedded list of stars used to fit the PSF.
153 Notes
154 -----
155 By default Piff saves the list of stars (including postage stamps) used
156 to fit the PSF, which makes the serialized form much larger. But the
157 upstream Piff serialization code recognizes the case where that
158 ``stars`` attribute has been deleted and serializes everything else.
160 Unfortunately, to date, Rubin's pickle-based Piff serialization instead
161 just deletes the postage stamp image attributes from inside the Piff
162 ``stars`` list, which is not a state the Piff serialization code
163 handles gracefully. So for now we have to drop the full stars list
164 during serialization if it is present.
165 """
166 if hasattr(self._impl, "stars"):
167 stars = self._impl.stars
168 try:
169 del self._impl.stars
170 yield
171 finally:
172 self._impl.stars = stars
173 else:
174 yield
177# Conventions on public visibility of the serialization types:
178#
179# - We lift and document the outermost Pydantic model type, since that needs to
180# be included directly in the Pydantic models of types that hold a PSF. This
181# type needs to be very clearly documented and named as a *serialization*
182# model, since there are many other kinds of models in play in this package.
183#
184# - We do not lift or document types used in that outermost model, but we do
185# not give them leading underscores, since they aren't really private.
186#
187# - Other utility types do get leading underscores.
190# Piff serialization uses a lot of dictionaries and lists restricted to these
191# basic types.
192type PiffScalar = int | float | str | bool | None
193type PiffDict = dict[str, PiffScalar | list[PiffScalar]]
196class GalSimPixelScaleModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
197 """Model used to serialize `galsim.wcs.PixelScale` instances."""
199 scale: float
200 wcs_type: Literal["pixel_scale"] = "pixel_scale"
203# We expect this discriminated union to grow to include other trivial
204# pixel-to-pixel transforms that get embedded in PSFs. If we someday have to
205# store Piff objects that embed more sophisticated PSFs, we'll hook them into
206# the AST-based coordinate transform system instead, but as long as we're just
207# talking about simple offsets and scalings, that's a lot of extra complexity
208# for very little gain.
209type GalSimLocalWcsModel = Annotated[GalSimPixelScaleModel, pydantic.Field(discriminator="wcs_type")]
212class PiffTableModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
213 """Serialization model used to embed a reference to a binary-data table in
214 a Piff serialization's JSON-like data.
215 """
217 metadata: PiffDict
218 table: serialization.TableModel
221class PiffObjectModel(pydantic.BaseModel, ser_json_inf_nan="constants"):
222 """General-purpose serialization model used for various Piff objects."""
224 structs: dict[str, PiffDict] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
225 tables: dict[str, PiffTableModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
226 wcs: dict[str, GalSimLocalWcsModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
227 objects: dict[str, PiffObjectModel] = pydantic.Field(default_factory=dict, exclude_if=operator.not_)
230class PiffSerializationModel(serialization.ArchiveTree):
231 """Serialization model for a Piff PSF."""
233 piff: PiffObjectModel = pydantic.Field(description="The Piff PSF object itself.")
235 stamp_size: int = pydantic.Field(
236 description="Width of the (square) images returned by this PSF's methods."
237 )
239 bounds: SerializableBounds = pydantic.Field(
240 description="The bounds object that represents the PSF's validity region."
241 )
244class _ArchivePiffWriter:
245 """An adapter from the Piff serialization interface to the
246 `.serialization.OutputArchive` class.
248 Notes
249 -----
250 Piff has its own simple serialization framework (contributed upstream by
251 Rubin DM) that maps everything to dictionaries, structured numpy arrays,
252 and a library of GalSim WCS objects, with the native implementation writing
253 standalone FITS files. That mostly maps nicely to the `lsst.images`
254 archive system, but we don't get to leverage any Pydantic validation or
255 JSON schema functionality since we only get opaque dictionaries from Piff.
257 See `piff.FitsWriter` for most method documentation; this class is designed
258 to mimic it exactly (the Piff authors prefer to just use duck-typing rather
259 than ABCs or protocols for interface definition).
260 """
262 def __init__(self, base_name: str = ""):
263 self._base_name = base_name
264 self.structs: dict[str, PiffDict] = {}
265 self.tables: dict[str, tuple[np.ndarray, PiffDict]] = {}
266 self.wcs_models: dict[str, GalSimLocalWcsModel] = {}
267 self.writers: dict[str, _ArchivePiffWriter] = {}
269 def write_struct(self, name: str, struct: PiffDict) -> None:
270 self.structs[name] = {k: self._to_builtin(v) for k, v in struct.items()}
272 def write_table(self, name: str, array: np.ndarray, metadata: PiffDict | None = None) -> None:
273 self.tables[name] = (array, metadata or {})
275 def write_wcs_map(
276 self, name: str, wcs_map: dict[int, galsim.wcs.BaseWCS], pointing: galsim.CelestialCoord | None
277 ) -> None:
278 import galsim.wcs
280 match wcs_map:
281 case {0: galsim.wcs.PixelScale() as wcs} if pointing is None:
282 self.wcs_models[name] = GalSimPixelScaleModel(scale=wcs.scale)
283 case _:
284 raise NotImplementedError("PSFs with complex embedded WCSs are not supported.")
286 @contextmanager
287 def nested(self, name: str) -> Iterator[_ArchivePiffWriter]:
288 nested = _ArchivePiffWriter(self.get_full_name(name))
289 yield nested
290 self.writers[name] = nested
292 def get_full_name(self, name: str) -> str:
293 return f"{self._base_name}/{name}"
295 def serialize(self, archive: serialization.OutputArchive[Any]) -> PiffObjectModel:
296 """Serialize to an archive.
298 This method is intended to be used as the callable passed to
299 `.serialization.OutputArchive.serialize_direct` and
300 `.serialization.OutputArchive.serialize_pointer`, after first passing
301 this writer to a Piff object's ``write`` or ``_write`` method.
302 """
303 model = PiffObjectModel()
304 for name, struct in self.structs.items():
305 model.structs[name] = struct
306 for name, (array, metadata) in self.tables.items():
307 model.tables[name] = PiffTableModel(
308 metadata=metadata,
309 table=archive.add_structured_array(
310 array, name=name, update_header=lambda header: header.update(metadata)
311 ),
312 )
313 for name, wcs_model in self.wcs_models.items():
314 model.wcs[name] = wcs_model
315 for name, writer in self.writers.items():
316 model.objects[name] = archive.serialize_direct(name, writer.serialize)
317 return model
319 @staticmethod
320 def _to_builtin(val: Any) -> PiffScalar:
321 match val:
322 case np.integer():
323 return int(val)
324 case np.floating():
325 return float(val)
326 case np.str_():
327 return str(val)
328 return val
331class _ArchivePiffReader:
332 """An adapter from the Piff serialization interface to the
333 `.serialization.InputArchive` class.
335 See `ArchivePiffWriter` for additional notes.
336 """
338 def __init__(
339 self, object_model: PiffObjectModel, archive: serialization.InputArchive[Any], base_name: str = ""
340 ):
341 self._model = object_model
342 self._archive = archive
343 self._base_name = base_name
345 def read_struct(self, name: str) -> PiffDict | None:
346 return self._model.structs.get(name)
348 def read_table(self, name: str, metadata: PiffDict | None = None) -> np.ndarray | None:
349 table_model = self._model.tables.get(name)
350 if table_model is None:
351 return None
352 if metadata is not None:
353 metadata.update(table_model.metadata)
354 return self._archive.get_structured_array(
355 table_model.table, strip_header=astropy.io.fits.Header.clear
356 )
358 def read_wcs_map(
359 self, name: str, logger: piff.config.LoggerWrapper
360 ) -> tuple[dict[int, galsim.wcs.BaseWCS] | None, galsim.CelestialCoord | None]:
361 import galsim.wcs
363 match self._model.wcs.get(name):
364 case GalSimPixelScaleModel(scale=scale):
365 return {0: galsim.wcs.PixelScale(scale)}, None
366 case None:
367 return None, None
368 case unexpected:
369 raise serialization.ArchiveReadError(
370 f"{self.get_full_name(name)} should be a WCS or WCS map, not {unexpected!r}."
371 )
373 @contextmanager
374 def nested(self, name: str) -> Iterator[_ArchivePiffReader]:
375 nested_model = self._model.objects[name]
376 yield _ArchivePiffReader(nested_model, self._archive, self.get_full_name(name))
378 def get_full_name(self, name: str) -> str:
379 return f"{self._base_name}/{name}"