Coverage for python/lsst/images/convolution_kernels.py: 24%

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. 

11from __future__ import annotations 

12 

13__all__ = ( 

14 "ConvolutionKernel", 

15 "ConvolutionKernelSerializationModel", 

16 "ImageBasisConvolutionKernel", 

17 "ImageBasisConvolutionKernelSerializationModel", 

18) 

19 

20from abc import ABC, abstractmethod 

21from collections.abc import Iterable, Iterator, Sequence 

22from typing import TYPE_CHECKING, Any, ClassVar, Literal 

23 

24import numpy as np 

25import pydantic 

26 

27from ._geom import YX, Bounds, Box 

28from ._image import Image 

29from .describe import DescribableMixin, DescribeOptions, FieldRole, Report, ReportField 

30from .fields import ChebyshevField, Field, FieldSerializationModel 

31from .serialization import ( 

32 ArchiveTree, 

33 ArrayReferenceModel, 

34 InlineArrayModel, 

35 InputArchive, 

36 InvalidParameterError, 

37 OutputArchive, 

38) 

39 

40if TYPE_CHECKING: 

41 try: 

42 from lsst.afw.math import LinearCombinationKernel as LegacyLinearCombinationKernel 

43 except ImportError: 

44 type LegacyLinearCombinationKernel = Any # type: ignore[no-redef] 

45 

46 

47# This may become a union in the future. 

48type ConvolutionKernelSerializationModel = ImageBasisConvolutionKernelSerializationModel 

49 

50 

51class ConvolutionKernel(DescribableMixin, ABC): 

52 """An abstract base class for spatially-varying convolution kernels.""" 

53 

54 def _describe(self, options: DescribeOptions = DescribeOptions(), /) -> Report: 

55 """Return a `Report` describing this kernel. 

56 

57 Parameters 

58 ---------- 

59 options : `DescribeOptions`, optional 

60 Unused; accepted for interface compatibility. 

61 """ 

62 return Report( 

63 type_name=type(self).__name__, 

64 summary=f"{type(self).__name__} over {self.bounds}", 

65 fields=[ 

66 ReportField(label="bounds", value=self.bounds, role=FieldRole.DERIVED), 

67 ReportField(label="kernel_bbox", value=self.kernel_bbox, role=FieldRole.DERIVED), 

68 ], 

69 ) 

70 

71 @property 

72 @abstractmethod 

73 def bounds(self) -> Bounds: 

74 """The region where this convolution kernel is valid 

75 (`~lsst.images.Bounds`). 

76 """ 

77 raise NotImplementedError() 

78 

79 @property 

80 @abstractmethod 

81 def kernel_bbox(self) -> Box: 

82 """Bounding box of all images returned by `compute_kernel_image` 

83 (`~lsst.images.Box`). 

84 """ 

85 raise NotImplementedError() 

86 

87 @abstractmethod 

88 def compute_kernel_image(self, *, x: int, y: int) -> Image: 

89 """Evaluate the kernel at a point. 

90 

91 Parameters 

92 ---------- 

93 x 

94 Column position coordinate to evaluate at. 

95 y 

96 Row position coordinate to evaluate at. 

97 

98 Returns 

99 ------- 

100 Image 

101 An image of the kernel, centered on the center of the center pixel, 

102 which is defined to be ``(0, 0)`` by the image's origin. 

103 """ 

104 raise NotImplementedError() 

105 

106 @abstractmethod 

107 def serialize(self, archive: OutputArchive[Any]) -> ConvolutionKernelSerializationModel: 

108 """Serialize the kernel to an output archive. 

109 

110 Parameters 

111 ---------- 

112 archive 

113 Archive to write to. 

114 """ 

115 raise NotImplementedError() 

116 

117 

118class ImageBasisConvolutionKernel(ConvolutionKernel): 

119 """A convolution kernel formed by a linear combination of images 

120 multiplied by `~lsst.images.fields.BaseField` instances. 

121 

122 Parameters 

123 ---------- 

124 basis 

125 A 3-d array holding the kernel images each basis function, with shape 

126 ``(n, height, width)``. 

127 spatial 

128 Iterable of `.fields.BaseField` of length ``basis.shape[0]``, holding 

129 the spatial variation of each basis kernel. 

130 center_y 

131 Center of the basis kernels in the x dimension. Defaults to 

132 ``height//2``. 

133 center_x 

134 Center of the basis kernels in the x dimension. Defaults to 

135 ``width//2``. 

136 """ 

137 

138 def __init__( 

139 self, 

140 basis: np.ndarray, 

141 spatial: Iterable[Field], 

142 center_y: int | None = None, 

143 center_x: int | None = None, 

144 ): 

145 self._spatial = tuple(spatial) 

146 bounds: Bounds | None = None 

147 for field in self._spatial: 

148 if field.unit is not None: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true

149 raise ValueError("Kernel spatial fields should not have units.") 

150 if bounds is None: 

151 bounds = field.bounds 

152 else: 

153 bounds = bounds.intersection(field.bounds) 

154 if bounds is None: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true

155 raise ValueError("Must have at least one basis function.") 

156 self._bounds = bounds 

157 self._basis = basis 

158 if self._basis.ndim != 3: 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true

159 raise ValueError(f"Basis array must be 3-d; shape={self._basis.shape}.") 

160 if len(self._spatial) != self._basis.shape[0]: 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true

161 raise ValueError( 

162 f"Number of spatial fields ({len(self._spatial)}) " 

163 f"does not match basis array shape ({self._basis.shape})." 

164 ) 

165 if center_y is None: 

166 center_y = self._basis.shape[1] // 2 

167 if center_x is None: 

168 center_x = self._basis.shape[2] // 2 

169 self._kernel_bbox = Box.from_shape(self._basis.shape[1:], start=YX(y=-center_y, x=-center_x)) 

170 

171 @property 

172 def bounds(self) -> Bounds: 

173 return self._bounds 

174 

175 @property 

176 def kernel_bbox(self) -> Box: 

177 return self._kernel_bbox 

178 

179 @property 

180 def spatial(self) -> Sequence[Field]: 

181 """The spatial variation of each basis function 

182 (`~collections.abc.Sequence` [`~.fields.BaseField`]). 

183 """ 

184 return self._spatial 

185 

186 @property 

187 def basis(self) -> np.ndarray: 

188 """The kernel basis functions, as an array with shape ``(n, h, w)`` 

189 (`numpy.ndarray`). 

190 """ 

191 return self._basis 

192 

193 def __len__(self) -> int: 

194 return len(self._spatial) 

195 

196 def __iter__(self) -> Iterator[tuple[Image, Field]]: 

197 for field, array in zip(self._spatial, self._basis, strict=True): 

198 yield Image(array, bbox=self._kernel_bbox), field 

199 

200 def _describe(self, options: DescribeOptions = DescribeOptions(), /) -> Report: 

201 """Return a `Report` describing this kernel. 

202 

203 Parameters 

204 ---------- 

205 options : `DescribeOptions`, optional 

206 Unused; accepted for interface compatibility. 

207 

208 Notes 

209 ----- 

210 The basis functions are summarized rather than listed: a realistic 

211 kernel has dozens of them, and each is a `~.fields.BaseField` whose 

212 own report would bury the rest of the image it belongs to. The 

213 shape of the basis images is `kernel_bbox`, already reported. 

214 """ 

215 report = super()._describe(options) 

216 report.summary = f"{type(self).__name__} with {len(self._spatial)} basis images over {self.bounds}" 

217 report.fields.append( 

218 ReportField(label="basis images", value=len(self._spatial), role=FieldRole.DERIVED) 

219 ) 

220 report.fields.append( 

221 ReportField( 

222 label="spatial variation", 

223 value=", ".join(sorted({type(field).__name__ for field in self._spatial})), 

224 role=FieldRole.DERIVED, 

225 ) 

226 ) 

227 return report 

228 

229 def compute_kernel_image(self, *, x: int, y: int) -> Image: 

230 # TODO[DM-54965]: simplify this once BaseField.__call__ behaves more 

231 # like a real ufunc and can handle scalars directly. 

232 x_array = np.array([x], dtype=np.float64) 

233 y_array = np.array([y], dtype=np.float64) 

234 weights = np.array( 

235 [spatial_field(x=x_array, y=y_array)[0] for spatial_field in self._spatial], 

236 dtype=np.float64, 

237 ) 

238 return Image(np.tensordot(weights, self._basis, axes=(0, 0)), bbox=self._kernel_bbox) 

239 

240 def serialize(self, archive: OutputArchive[Any]) -> ImageBasisConvolutionKernelSerializationModel: 

241 """Serialize the kernel to an output archive. 

242 

243 Parameters 

244 ---------- 

245 archive 

246 Archive to write to. 

247 """ 

248 serialized_basis = archive.add_array(self._basis, name="basis") 

249 serialized_spatial = [archive.serialize_direct("spatial", f.serialize) for f in self._spatial] 

250 return ImageBasisConvolutionKernelSerializationModel( 

251 basis=serialized_basis, 

252 spatial=serialized_spatial, 

253 center_y=-self._kernel_bbox.y.min, 

254 center_x=-self._kernel_bbox.x.min, 

255 ) 

256 

257 @staticmethod 

258 def _get_archive_tree_type( 

259 pointer_type: type[Any], 

260 ) -> type[ImageBasisConvolutionKernelSerializationModel]: 

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

262 type that uses the given pointer type. 

263 """ 

264 return ImageBasisConvolutionKernelSerializationModel 

265 

266 @staticmethod 

267 def from_legacy(legacy_kernel: LegacyLinearCombinationKernel) -> ImageBasisConvolutionKernel: 

268 """Convert from a legacy `lsst.afw.math.LinearCombinationKernel`. 

269 

270 Parameters 

271 ---------- 

272 legacy_kernel 

273 The kernel to convert. Must use Chebyshev polynomials for its 

274 spatial variation and `lsst.afw.math.FixedKernel` objects with a 

275 consistent shape and center for its basis functions. 

276 """ 

277 from lsst.afw.math import FixedKernel as LegacyFixedKernel 

278 from lsst.afw.math import LinearCombinationKernel as LegacyLinearCombinationKernel 

279 

280 if not isinstance(legacy_kernel, LegacyLinearCombinationKernel): 

281 raise TypeError( 

282 f"Cannot convert {type(legacy_kernel).__name__} instance to an ImageBasisConvolutionKernel." 

283 ) 

284 dimensions = legacy_kernel.getDimensions() 

285 center = legacy_kernel.getCtr() 

286 basis = np.zeros((legacy_kernel.getNBasisKernels(), dimensions.y, dimensions.x), dtype=np.float64) 

287 for n, basis_kernel in enumerate(legacy_kernel.getKernelList()): 

288 if basis_kernel.getDimensions() != dimensions: 

289 raise ValueError("Cannot convert LinearCombinationKernel with different-size basis kernels.") 

290 if basis_kernel.getCtr() != center: 

291 raise ValueError( 

292 "Cannot convert LinearCombinationKernel with differently-centered basis kernels." 

293 ) 

294 if not isinstance(basis_kernel, LegacyFixedKernel): 

295 raise ValueError("Cannot convert LinearCombinationKernel with non-fixed basis kernels.") 

296 legacy_image_view = Image(basis[n, :, :], dtype=np.float64).to_legacy() 

297 basis_kernel.computeImage(legacy_image_view, doNormalize=False) 

298 spatial = [ChebyshevField.from_legacy_function2(f) for f in legacy_kernel.getSpatialFunctionList()] 

299 return ImageBasisConvolutionKernel(basis=basis, spatial=spatial, center_y=center.y, center_x=center.x) 

300 

301 def to_legacy(self) -> LegacyLinearCombinationKernel: 

302 """Convert to a legacy `lsst.afw.math.LinearCombinationKernel`. 

303 

304 This only works if all spatial variation is handled by 

305 `lsst.images.ChebyshevField`. 

306 """ 

307 from lsst.afw.math import FixedKernel as LegacyFixedKernel 

308 from lsst.afw.math import LinearCombinationKernel as LegacyLinearCombinationKernel 

309 from lsst.geom import Point2I as LegacyPoint2I 

310 

311 basis_kernels = [] 

312 spatial_functions = [] 

313 legacy_center = LegacyPoint2I(-self._kernel_bbox.x.min, -self._kernel_bbox.y.min) 

314 for image, field in self: 

315 legacy_image = image.to_legacy() 

316 legacy_image.setXY0(LegacyPoint2I()) 

317 basis_kernel = LegacyFixedKernel(legacy_image) 

318 basis_kernel.setCtr(legacy_center) 

319 basis_kernels.append(basis_kernel) 

320 if not isinstance(field, ChebyshevField): 

321 raise ValueError("Only Chebyshev spatial variation can be converted.") 

322 spatial_functions.append(field.to_legacy_function2()) 

323 result = LegacyLinearCombinationKernel(basis_kernels, spatial_functions) 

324 result.setCtr(legacy_center) 

325 return result 

326 

327 

328class ImageBasisConvolutionKernelSerializationModel(ArchiveTree): 

329 """The serialization model for `ImageBasisConvolutionKernel`.""" 

330 

331 SCHEMA_NAME: ClassVar[str] = "image_basis_convolution_kernel" 

332 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

333 MIN_READ_VERSION: ClassVar[int] = 1 

334 PUBLIC_TYPE: ClassVar[type] = ImageBasisConvolutionKernel 

335 

336 basis: ArrayReferenceModel | InlineArrayModel = pydantic.Field( 

337 description="The basis images, with shape (n, h, w)." 

338 ) 

339 spatial: list[FieldSerializationModel] = pydantic.Field( 

340 description="The spatial variation of each basis function." 

341 ) 

342 center_y: int = pydantic.Field(description="Center row of the kernel in the basis images.") 

343 center_x: int = pydantic.Field(description="Center column of the kernel in the basis images.") 

344 

345 kernel_type: Literal["IMAGE_BASIS"] = "IMAGE_BASIS" 

346 

347 def deserialize(self, archive: InputArchive[Any], **kwargs: Any) -> ImageBasisConvolutionKernel: 

348 if kwargs: 348 ↛ 349line 348 didn't jump to line 349 because the condition on line 348 was never true

349 raise InvalidParameterError(f"Unrecognized parameters for ChebyshevField: {set(kwargs.keys())}.") 

350 basis = archive.get_array(self.basis) 

351 spatial = [f.deserialize(archive) for f in self.spatial] 

352 return ImageBasisConvolutionKernel(basis, spatial, center_y=self.center_y, center_x=self.center_x)