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.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 @classmethod 

131 def deserialize( 

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

133 ) -> PSFExWrapper: 

134 """Deserialize the PSF from an archive. 

135 

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

137 `.serialization.InputArchive.deserialize_pointer`. 

138 """ 

139 try: 

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

141 except ImportError: 

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

143 

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

145 data = PsfexPsfSerializationData() 

146 data.average_x = model.average_x 

147 data.average_y = model.average_y 

148 data.pixel_step = model.pixel_step 

149 data.group = model.group 

150 data.degree = model.degree 

151 data.basis = model.basis 

152 data.coeff = model.coeff 

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

154 data.comp = parameters.flatten() 

155 data.context = model.context 

156 legacy_psf = PsfexPsf.fromSerializationData(data) 

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

158 

159 @staticmethod 

160 def _get_archive_tree_type( 

161 pointer_type: type[pydantic.BaseModel], 

162 ) -> type[PSFExSerializationModel]: 

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

164 type that uses the given pointer type. 

165 """ 

166 return PSFExSerializationModel 

167 

168 

169class PSFExSerializationModel(serialization.ArchiveTree): 

170 """Serialization model for PSFEx PSFs.""" 

171 

172 average_x: float = pydantic.Field( 

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

174 ) 

175 

176 average_y: float = pydantic.Field( 

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

178 ) 

179 

180 pixel_step: float = pydantic.Field( 

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

182 ) 

183 

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

185 default_factory=lambda: [0, 0], 

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

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

188 ) 

189 

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

191 

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

193 

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

195 

196 parameters: serialization.ArrayReferenceModel = pydantic.Field( 

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

198 ) 

199 

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

201 

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

203 

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