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

130 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-01 02:06 -0700

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 .fields import ChebyshevField, Field, FieldSerializationModel 

30from .serialization import ( 

31 ArchiveTree, 

32 ArrayReferenceModel, 

33 InlineArrayModel, 

34 InputArchive, 

35 InvalidParameterError, 

36 OutputArchive, 

37) 

38 

39if TYPE_CHECKING: 

40 try: 

41 from lsst.afw.math import LinearCombinationKernel as LegacyLinearCombinationKernel 

42 except ImportError: 

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

44 

45 

46# This may become a union in the future. 

47type ConvolutionKernelSerializationModel = ImageBasisConvolutionKernelSerializationModel 

48 

49 

50class ConvolutionKernel(ABC): 

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

52 

53 @property 

54 @abstractmethod 

55 def bounds(self) -> Bounds: 

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

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

58 """ 

59 raise NotImplementedError() 

60 

61 @property 

62 @abstractmethod 

63 def kernel_bbox(self) -> Box: 

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

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

66 """ 

67 raise NotImplementedError() 

68 

69 @abstractmethod 

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

71 """Evaluate the kernel at a point. 

72 

73 Parameters 

74 ---------- 

75 x 

76 Column position coordinate to evaluate at. 

77 y 

78 Row position coordinate to evaluate at. 

79 

80 Returns 

81 ------- 

82 Image 

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

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

85 """ 

86 raise NotImplementedError() 

87 

88 @abstractmethod 

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

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

91 

92 Parameters 

93 ---------- 

94 archive 

95 Archive to write to. 

96 """ 

97 raise NotImplementedError() 

98 

99 

100class ImageBasisConvolutionKernel(ConvolutionKernel): 

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

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

103 

104 Parameters 

105 ---------- 

106 basis 

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

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

109 spatial 

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

111 the spatial variation of each basis kernel. 

112 center_y 

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

114 ``height//2``. 

115 center_x 

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

117 ``width//2``. 

118 """ 

119 

120 def __init__( 

121 self, 

122 basis: np.ndarray, 

123 spatial: Iterable[Field], 

124 center_y: int | None = None, 

125 center_x: int | None = None, 

126 ): 

127 self._spatial = tuple(spatial) 

128 bounds: Bounds | None = None 

129 for field in self._spatial: 

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

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

132 if bounds is None: 

133 bounds = field.bounds 

134 else: 

135 bounds = bounds.intersection(field.bounds) 

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

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

138 self._bounds = bounds 

139 self._basis = basis 

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

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

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

143 raise ValueError( 

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

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

146 ) 

147 if center_y is None: 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true

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

149 if center_x is None: 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true

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

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

152 

153 @property 

154 def bounds(self) -> Bounds: 

155 return self._bounds 

156 

157 @property 

158 def kernel_bbox(self) -> Box: 

159 return self._kernel_bbox 

160 

161 @property 

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

163 """The spatial variation of each basis function 

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

165 """ 

166 return self._spatial 

167 

168 @property 

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

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

171 (`numpy.ndarray`). 

172 """ 

173 return self._basis 

174 

175 def __len__(self) -> int: 

176 return len(self._spatial) 

177 

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

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

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

181 

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

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

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

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

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

187 weights = np.array( 

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

189 dtype=np.float64, 

190 ) 

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

192 

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

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

195 

196 Parameters 

197 ---------- 

198 archive 

199 Archive to write to. 

200 """ 

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

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

203 return ImageBasisConvolutionKernelSerializationModel( 

204 basis=serialized_basis, 

205 spatial=serialized_spatial, 

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

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

208 ) 

209 

210 @staticmethod 

211 def _get_archive_tree_type( 

212 pointer_type: type[Any], 

213 ) -> type[ImageBasisConvolutionKernelSerializationModel]: 

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

215 type that uses the given pointer type. 

216 """ 

217 return ImageBasisConvolutionKernelSerializationModel 

218 

219 @staticmethod 

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

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

222 

223 Parameters 

224 ---------- 

225 legacy_kernel 

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

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

228 consistent shape and center for its basis functions. 

229 """ 

230 from lsst.afw.math import FixedKernel as LegacyFixedKernel 

231 from lsst.afw.math import LinearCombinationKernel as LegacyLinearCombinationKernel 

232 

233 if not isinstance(legacy_kernel, LegacyLinearCombinationKernel): 

234 raise TypeError( 

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

236 ) 

237 dimensions = legacy_kernel.getDimensions() 

238 center = legacy_kernel.getCtr() 

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

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

241 if basis_kernel.getDimensions() != dimensions: 

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

243 if basis_kernel.getCtr() != center: 

244 raise ValueError( 

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

246 ) 

247 if not isinstance(basis_kernel, LegacyFixedKernel): 

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

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

250 basis_kernel.computeImage(legacy_image_view, doNormalize=False) 

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

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

253 

254 def to_legacy(self) -> LegacyLinearCombinationKernel: 

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

256 

257 This only works if all spatial variation is handled by 

258 `lsst.images.ChebyshevField`. 

259 """ 

260 from lsst.afw.math import FixedKernel as LegacyFixedKernel 

261 from lsst.afw.math import LinearCombinationKernel as LegacyLinearCombinationKernel 

262 from lsst.geom import Point2I as LegacyPoint2I 

263 

264 basis_kernels = [] 

265 spatial_functions = [] 

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

267 for image, field in self: 

268 legacy_image = image.to_legacy() 

269 legacy_image.setXY0(LegacyPoint2I()) 

270 basis_kernel = LegacyFixedKernel(legacy_image) 

271 basis_kernel.setCtr(legacy_center) 

272 basis_kernels.append(basis_kernel) 

273 if not isinstance(field, ChebyshevField): 

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

275 spatial_functions.append(field.to_legacy_function2()) 

276 result = LegacyLinearCombinationKernel(basis_kernels, spatial_functions) 

277 result.setCtr(legacy_center) 

278 return result 

279 

280 

281class ImageBasisConvolutionKernelSerializationModel(ArchiveTree): 

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

283 

284 SCHEMA_NAME: ClassVar[str] = "image_basis_convolution_kernel" 

285 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

286 MIN_READ_VERSION: ClassVar[int] = 1 

287 PUBLIC_TYPE: ClassVar[type] = ImageBasisConvolutionKernel 

288 

289 basis: ArrayReferenceModel | InlineArrayModel = pydantic.Field( 

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

291 ) 

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

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

294 ) 

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

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

297 

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

299 

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

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

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

303 basis = archive.get_array(self.basis) 

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

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