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, ClassVar 

18 

19import numpy as np 

20import pydantic 

21 

22from .. import serialization 

23from .._concrete_bounds import BoundsSerializationModel 

24from .._geom import Bounds, Box 

25from .._image import Image 

26from ._base import PointSpreadFunction 

27 

28 

29class LegacyPointSpreadFunction(PointSpreadFunction): 

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

31 

32 Parameters 

33 ---------- 

34 impl 

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

36 bounds 

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

38 

39 Notes 

40 ----- 

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

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

43 """ 

44 

45 def __init__(self, impl: Any, bounds: Bounds) -> None: 

46 self._impl = impl 

47 self._bounds = bounds 

48 

49 @property 

50 def bounds(self) -> Bounds: 

51 return self._bounds 

52 

53 @cached_property 

54 def kernel_bbox(self) -> Box: 

55 from lsst.geom import Box2I, Point2D 

56 

57 biggest = Box2I() 

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

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

60 return Box.from_legacy(biggest) 

61 

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

63 from lsst.geom import Point2D 

64 

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

66 if result.bbox != self.kernel_bbox: 

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

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

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

70 result = padded 

71 return result 

72 

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

74 from lsst.geom import Point2D 

75 

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

77 

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

79 from lsst.geom import Point2D 

80 

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

82 

83 @property 

84 def legacy_psf(self) -> Any: 

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

86 return self._impl 

87 

88 @classmethod 

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

90 from lsst.meas.extensions.psfex import PsfexPsf 

91 

92 if isinstance(legacy_psf, PsfexPsf): 

93 return PSFExWrapper(legacy_psf, bounds) 

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

95 

96 

97class PSFExWrapper(LegacyPointSpreadFunction): 

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

99 

100 Parameters 

101 ---------- 

102 impl 

103 A `lsst.meas.extensions.psfex.PsfexPsf` instance. 

104 bounds 

105 The pixel-coordinate region where the model can safely be 

106 evaluated. 

107 """ 

108 

109 def __init__(self, impl: Any, bounds: Bounds) -> None: 

110 from lsst.meas.extensions.psfex import PsfexPsf 

111 

112 if not isinstance(impl, PsfexPsf): 

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

114 super().__init__(impl, bounds) 

115 

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

117 """Serialize the PSF to an archive. 

118 

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

120 `.serialization.OutputArchive.serialize_direct` or 

121 `.serialization.OutputArchive.serialize_pointer`. 

122 

123 Parameters 

124 ---------- 

125 archive 

126 Archive to write to. 

127 """ 

128 data = self._impl.getSerializationData() 

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

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

131 return PSFExSerializationModel( 

132 average_x=data.average_x, 

133 average_y=data.average_y, 

134 pixel_step=data.pixel_step, 

135 group=data.group, 

136 degree=data.degree, 

137 basis=data.basis, 

138 coeff=data.coeff, 

139 parameters=array_ref, 

140 context=data.context, 

141 bounds=self.bounds.serialize(), 

142 ) 

143 

144 @staticmethod 

145 def _get_archive_tree_type( 

146 pointer_type: type[pydantic.BaseModel], 

147 ) -> type[PSFExSerializationModel]: 

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

149 type that uses the given pointer type. 

150 """ 

151 return PSFExSerializationModel 

152 

153 

154class PSFExSerializationModel(serialization.ArchiveTree): 

155 """Serialization model for PSFEx PSFs.""" 

156 

157 SCHEMA_NAME: ClassVar[str] = "psfex_psf" 

158 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

159 MIN_READ_VERSION: ClassVar[int] = 1 

160 PUBLIC_TYPE: ClassVar[type] = PSFExWrapper 

161 

162 average_x: float = pydantic.Field( 

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

164 ) 

165 

166 average_y: float = pydantic.Field( 

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

168 ) 

169 

170 pixel_step: float = pydantic.Field( 

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

172 ) 

173 

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

175 default_factory=lambda: [0, 0], 

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

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

178 ) 

179 

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

181 

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

183 

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

185 

186 parameters: serialization.ArrayReferenceModel | serialization.InlineArrayModel = pydantic.Field( 

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

188 ) 

189 

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

191 

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

193 

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

195 

196 def deserialize(self, archive: serialization.InputArchive[Any], **kwargs: Any) -> PSFExWrapper: 

197 """Deserialize the PSF from an archive. 

198 

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

200 `.serialization.InputArchive.deserialize_pointer`. 

201 

202 Parameters 

203 ---------- 

204 archive 

205 Archive to read from. 

206 **kwargs 

207 Unsupported keyword arguments are accepted only to provide 

208 better error messages (raising 

209 `.serialization.InvalidParameterError`). 

210 """ 

211 if kwargs: 

212 raise serialization.InvalidParameterError( 

213 f"Unrecognized parameters for PsfExWrapper: {set(kwargs.keys())}." 

214 ) 

215 try: 

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

217 except ImportError: 

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

219 

220 parameters = archive.get_array(self.parameters).astype(np.float32) 

221 data = PsfexPsfSerializationData() 

222 data.average_x = self.average_x 

223 data.average_y = self.average_y 

224 data.pixel_step = self.pixel_step 

225 data.group = self.group 

226 data.degree = self.degree 

227 data.basis = self.basis 

228 data.coeff = self.coeff 

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

230 data.comp = parameters.flatten() 

231 data.context = self.context 

232 legacy_psf = PsfexPsf.fromSerializationData(data) 

233 return PSFExWrapper(legacy_psf, self.bounds.deserialize())