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 .._concrete_bounds import SerializableBounds 

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): 

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 def __init__(self, impl: Any, bounds: Bounds): 

101 from lsst.meas.extensions.psfex import PsfexPsf 

102 

103 if not isinstance(impl, PsfexPsf): 

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

105 super().__init__(impl, bounds) 

106 

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

108 """Serialize the PSF to an archive. 

109 

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

111 `.serialization.OutputArchive.serialize_direct` or 

112 `.serialization.OutputArchive.serialize_pointer`. 

113 """ 

114 data = self._impl.getSerializationData() 

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

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

117 return PSFExSerializationModel( 

118 average_x=data.average_x, 

119 average_y=data.average_y, 

120 pixel_step=data.pixel_step, 

121 group=data.group, 

122 degree=data.degree, 

123 basis=data.basis, 

124 coeff=data.coeff, 

125 parameters=array_ref, 

126 context=data.context, 

127 bounds=self.bounds.serialize(), 

128 ) 

129 

130 @staticmethod 

131 def _get_archive_tree_type( 

132 pointer_type: type[pydantic.BaseModel], 

133 ) -> type[PSFExSerializationModel]: 

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

135 type that uses the given pointer type. 

136 """ 

137 return PSFExSerializationModel 

138 

139 

140class PSFExSerializationModel(serialization.ArchiveTree): 

141 """Serialization model for PSFEx PSFs.""" 

142 

143 average_x: float = pydantic.Field( 

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

145 ) 

146 

147 average_y: float = pydantic.Field( 

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

149 ) 

150 

151 pixel_step: float = pydantic.Field( 

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

153 ) 

154 

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

156 default_factory=lambda: [0, 0], 

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

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

159 ) 

160 

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

162 

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

164 

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

166 

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

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

169 ) 

170 

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

172 

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

174 

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

176 

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

178 """Deserialize the PSF from an archive. 

179 

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

181 `.serialization.InputArchive.deserialize_pointer`. 

182 """ 

183 try: 

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

185 except ImportError: 

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

187 

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

189 data = PsfexPsfSerializationData() 

190 data.average_x = self.average_x 

191 data.average_y = self.average_y 

192 data.pixel_step = self.pixel_step 

193 data.group = self.group 

194 data.degree = self.degree 

195 data.basis = self.basis 

196 data.coeff = self.coeff 

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

198 data.comp = parameters.flatten() 

199 data.context = self.context 

200 legacy_psf = PsfexPsf.fromSerializationData(data) 

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