Coverage for python / lsst / images / psfs / _legacy.py: 42%

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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("LegacyPointSpreadFunction", "PSFExSerializationModel", "PSFExWrapper") 

15 

16from functools import cached_property 

17from typing import Any 

18 

19import numpy as np 

20import pydantic 

21 

22from .. import serialization 

23from .._geom import Bounds, Box, SerializableBounds 

24from .._image import Image 

25from ._base import PointSpreadFunction 

26 

27 

28class LegacyPointSpreadFunction(PointSpreadFunction): 

29 """A PSF model backed by a legacy `lsst.afw.detection.Psf` object. 

30 

31 Parameters 

32 ---------- 

33 impl 

34 An `lsst.afw.detection.Psf` instance. 

35 bounds 

36 The pixel-coordinate region where the model can safely be evaluated. 

37 

38 Notes 

39 ----- 

40 This wrapper is usable as-is on any `lsst.afw.detection.Psf` instance, 

41 but subclasses (e.g. `PSFExWrapper`) must be used for serialization. 

42 """ 

43 

44 def __init__(self, impl: Any, bounds: Bounds): 

45 self._impl = impl 

46 self._bounds = bounds 

47 

48 @property 

49 def bounds(self) -> Bounds: 

50 return self._bounds 

51 

52 @cached_property 

53 def kernel_bbox(self) -> Box: 

54 from lsst.geom import Box2I, Point2D 

55 

56 biggest = Box2I() 

57 for y, x in self._bounds.boundary(): 

58 biggest.include(self._impl.computeKernelBBox(Point2D(x, y))) 

59 return Box.from_legacy(biggest) 

60 

61 def compute_kernel_image(self, *, x: float, y: float) -> Image: 

62 from lsst.geom import Point2D 

63 

64 result = Image.from_legacy(self._impl.computeKernelImage(Point2D(x, y))) 

65 if result.bbox != self.kernel_bbox: 

66 # afw does not guarantee a consistent kernel_bbox, but we do now. 

67 padded = Image(0.0, bbox=self.kernel_bbox, dtype=np.float64) 

68 padded[self.kernel_bbox] = result[self.kernel_bbox] 

69 result = padded 

70 return result 

71 

72 def compute_stellar_image(self, *, x: float, y: float) -> Image: 

73 from lsst.geom import Point2D 

74 

75 return Image.from_legacy(self._impl.computeImage(Point2D(x, y))) 

76 

77 def compute_stellar_bbox(self, *, x: float, y: float) -> Box: 

78 from lsst.geom import Point2D 

79 

80 return Box.from_legacy(self._impl.computeImageBBox(Point2D(x, y))) 

81 

82 @property 

83 def legacy_psf(self) -> Any: 

84 """The backing `lsst.afw.detection.Psf` object.""" 

85 return self._impl 

86 

87 @classmethod 

88 def from_legacy(cls, legacy_psf: Any, bounds: Bounds) -> LegacyPointSpreadFunction: 

89 from lsst.meas.extensions.psfex import PsfexPsf 

90 

91 if isinstance(legacy_psf, PsfexPsf): 

92 return PSFExWrapper(legacy_psf, bounds) 

93 return cls(impl=legacy_psf, bounds=bounds) 

94 

95 

96class PSFExWrapper(LegacyPointSpreadFunction): 

97 """A specialization of LegacyPointSpreadFunction for the PSFEx backend.""" 

98 

99 def __init__(self, impl: Any, bounds: Bounds): 

100 from lsst.meas.extensions.psfex import PsfexPsf 

101 

102 if not isinstance(impl, PsfexPsf): 

103 raise TypeError(f"{impl!r} is not a PSFEx object.") 

104 super().__init__(impl, bounds) 

105 

106 def serialize(self, archive: serialization.OutputArchive[Any]) -> PSFExSerializationModel: 

107 """Serialize the PSF to an archive. 

108 

109 This method is intended to be usable as the callback function passed to 

110 `.serialization.OutputArchive.serialize_direct` or 

111 `.serialization.OutputArchive.serialize_pointer`. 

112 """ 

113 data = self._impl.getSerializationData() 

114 shape = tuple(reversed(data.size)) 

115 array_ref = archive.add_array(data.comp.reshape(*shape), name="parameters") 

116 return PSFExSerializationModel( 

117 average_x=data.average_x, 

118 average_y=data.average_y, 

119 pixel_step=data.pixel_step, 

120 group=data.group, 

121 degree=data.degree, 

122 basis=data.basis, 

123 coeff=data.coeff, 

124 parameters=array_ref, 

125 context=data.context, 

126 bounds=self.bounds.serialize(), 

127 ) 

128 

129 @classmethod 

130 def deserialize( 

131 cls, model: PSFExSerializationModel, archive: serialization.InputArchive[Any] 

132 ) -> PSFExWrapper: 

133 """Deserialize the PSF from an archive. 

134 

135 This method is intended to be usable as the callback function passed to 

136 `.serialization.InputArchive.deserialize_pointer`. 

137 """ 

138 try: 

139 from lsst.meas.extensions.psfex import PsfexPsf, PsfexPsfSerializationData 

140 except ImportError: 

141 raise serialization.ArchiveReadError("Failed to import lsst.meas.extensions.psfex.") from None 

142 

143 parameters = archive.get_array(model.parameters).astype(np.float32) 

144 data = PsfexPsfSerializationData() 

145 data.average_x = model.average_x 

146 data.average_y = model.average_y 

147 data.pixel_step = model.pixel_step 

148 data.group = model.group 

149 data.degree = model.degree 

150 data.basis = model.basis 

151 data.coeff = model.coeff 

152 data.size = list(reversed(parameters.shape)) 

153 data.comp = parameters.flatten() 

154 data.context = model.context 

155 legacy_psf = PsfexPsf.fromSerializationData(data) 

156 return cls(legacy_psf, Bounds.deserialize(model.bounds)) 

157 

158 @staticmethod 

159 def _get_archive_tree_type( 

160 pointer_type: type[pydantic.BaseModel], 

161 ) -> type[PSFExSerializationModel]: 

162 """Return the serialization model type for this object for an archive 

163 type that uses the given pointer type. 

164 """ 

165 return PSFExSerializationModel 

166 

167 

168class PSFExSerializationModel(serialization.ArchiveTree): 

169 """Serialization model for PSFEx PSFs.""" 

170 

171 average_x: float = pydantic.Field( 

172 description="Average X position of the stars used to build this PSF model." 

173 ) 

174 

175 average_y: float = pydantic.Field( 

176 description="Average Y position of the stars used to build this PSF model." 

177 ) 

178 

179 pixel_step: float = pydantic.Field( 

180 description="Size of a model pixel, as a fraction or multiple of the native pixel size." 

181 ) 

182 

183 group: list[int] = pydantic.Field( 

184 default_factory=lambda: [0, 0], 

185 exclude_if=lambda v: v == [0, 0], 

186 description="Number of model groups in each dimension.", 

187 ) 

188 

189 degree: list[int] = pydantic.Field(description="Polynomial degree for each model group.") 

190 

191 basis: list[float] = pydantic.Field(description="Basis function values.") 

192 

193 coeff: list[float] = pydantic.Field(description="Polynomial coefficients.") 

194 

195 parameters: serialization.ArrayReferenceModel = pydantic.Field( 

196 description="Reference to an array with the complete model parameters." 

197 ) 

198 

199 context: serialization.InlineArray = pydantic.Field(description="Internal PSFEx context array.") 

200 

201 bounds: SerializableBounds = pydantic.Field(description="Validity range for this PSF model.") 

202 

203 model_config = pydantic.ConfigDict(ser_json_inf_nan="constants")