Coverage for python/lsst/cell_coadds/_stitched_psf.py: 38%

1# This file is part of cell_coadds. 

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# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24import pickle 

25from functools import partial 

26from typing import TYPE_CHECKING, Any, ClassVar 

27 

28import numpy as np 

29 

30import lsst.geom as geom 

31from lsst.afw.image import ImageD 

32from lsst.afw.typehandling import StorableHelperFactory 

33from lsst.meas.algorithms import ImagePsf 

34 

35from ._grid_container import GridContainer 

36from ._uniform_grid import UniformGrid 

37 

38if TYPE_CHECKING: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true

39 from lsst.afw.image import Color 

40 

41__all__ = ("StitchedPsf",) 

42 

43 

44class StitchedPsf(ImagePsf): 

45 """A piecewise PSF implementation backed by a 2-d grid of images.""" 

46 

47 # We need to ensure a C++ StorableHelperFactory is constructed and 

48 # available before any unpersists of this class. Placing this "private" 

49 # class attribute here accomplishes that. 

50 _factory: ClassVar[type[StorableHelperFactory]] = StorableHelperFactory("lsst.cell_coadds", "StitchedPsf") 

51 

52 def __init__(self, images: GridContainer[ImageD], grid: UniformGrid) -> None: 

53 self._validate_args(images, grid) 

54 

55 super().__init__() 

56 self._images = images 

57 self._grid = grid 

58 self._averagePosition = None 

59 

60 @staticmethod 

61 def _validate_args(images: GridContainer[ImageD], grid: UniformGrid) -> None: 

62 """Validate the images and grid. 

63 

64 Parameters 

65 ---------- 

66 images : `GridContainer` 

67 The images to validate. 

68 grid : `UniformGrid` 

69 The grid to validate. 

70 

71 Raises 

72 ------ 

73 ValueError 

74 Raised if the images and grid are incompatible. 

75 """ 

76 min_x = min(index.x for index in images.indices()) 

77 min_y = min(index.y for index in images.indices()) 

78 max_x = max(index.x for index in images.indices()) 

79 max_y = max(index.y for index in images.indices()) 

80 

81 if ((max_x - min_x + 1) > grid.shape.x) or ((max_y - min_y + 1) > grid.shape.y): 

82 raise ValueError("Images do not fit on grid.") 

83 

84 @property 

85 def images(self) -> GridContainer[ImageD]: 

86 """The images that make up this PSF.""" 

87 return self._images 

88 

89 @property 

90 def grid(self) -> UniformGrid: 

91 """The grid on which the images are placed.""" 

92 return self._grid 

93 

94 def getAveragePosition(self) -> geom.Point2D: 

95 """Get a position where PSF can be evaluated on a patch. 

96 

97 This defaults to the center of the patch bounding box, unless there are 

98 no inputs there. In that case, it switches to find an arbitrary cell, 

99 typically at a corner that has inputs and returns the center position 

100 of the cell. 

101 """ 

102 if self._averagePosition is None: 

103 center = self._grid.bbox.getCenter() 

104 if self.grid.index(geom.Point2I(center)) not in self._images: 

105 arbitrary_index = next(iter(self._images)) 

106 bbox = self._grid.bbox_of(arbitrary_index) 

107 center = bbox.getCenter() 

108 

109 self._averagePosition = center 

110 

111 return self._averagePosition 

112 

113 # The _do* methods make use of the ImagePsf trampoline. 

114 def _doComputeBBox(self, position: geom.Point2D | geom.Point2I, color: Color = None) -> geom.Box2I: 

115 return self._images[self._grid.index(geom.Point2I(position))].getBBox() 

116 

117 def _doComputeKernelImage(self, position: geom.Point2D | geom.Point2I, color: Color = None) -> ImageD: 

118 return self._images[self._grid.index(geom.Point2I(position))] 

119 

120 def clone(self) -> StitchedPsf: 

121 """Return a deep copy of this object.""" 

122 return StitchedPsf(self.images, self.grid) 

123 

124 def __deepcopy__(self, memo: dict[int, Any] | None = None) -> StitchedPsf: 

125 """Return a deep copy of this object.""" 

126 return StitchedPsf(self.images, self.grid) 

127 

128 def __eq__(self, other: object) -> bool: 

129 if not isinstance(other, StitchedPsf): 

130 return False 

131 

132 if not (self.grid == other.grid): 

133 return False 

134 

135 for index in self.images.indices(): 

136 if ( 

137 not (self.images[index].array.shape == other.images[index].array.shape) 

138 or not np.equal(self.images[index].array, other.images[index].array).all() 

139 ): 

140 return False 

141 

142 return True 

143 

144 @staticmethod 

145 def _callback(image: ImageD, bbox: geom.Box2I) -> ImageD: 

146 if image.getBBox().contains(bbox): 

147 return image[bbox] 

148 else: 

149 # Make a new image big enough to fit current bbox and new bbox, 

150 # copy current image into it, then subset that for the returned 

151 # PSF. 

152 bigger_image = ImageD(bbox=bbox.expandedTo(image.getBBox()), initialValue=0.0) 

153 bigger_image[image.getBBox()] = image 

154 return bigger_image[bbox] 

155 

156 def resized(self, width: int, height: int) -> StitchedPsf: 

157 if not (width % 2 == 1 and width > 0): 

158 raise ValueError("resized width must be a positive odd integer; got {width}.") 

159 if not (height % 2 == 1 and height > 0): 

160 raise ValueError("resized height must be a positive odd integer; got {height}.") 

161 

162 bbox = geom.Box2I(geom.Point2I(-(width // 2), -(height // 2)), geom.Extent2I(width, height)) 

163 gc = self._images.rebuild_transformed(transform=partial(self._callback, bbox=bbox)) 

164 return StitchedPsf(gc, self.grid) 

165 

166 @staticmethod 

167 def isPersistable() -> bool: 

168 return True 

169 

170 @staticmethod 

171 def _getPersistenceName() -> str: 

172 return "StitchedPsf" 

173 

174 @staticmethod 

175 def _getPythonModule() -> str: 

176 return __name__ 

177 

178 # The get/set state methods are needed to support pickle. 

179 def __getstate__(self) -> dict: 

180 return {"images": self.images, "grid": self.grid} 

181 

182 def __setstate__(self, state: dict) -> None: 

183 StitchedPsf.__init__(self, state["images"], state["grid"]) 

184 

185 def _write(self) -> bytes: 

186 return pickle.dumps((self._images, self._grid)) 

187 

188 @staticmethod 

189 def _read(pkl: bytes) -> StitchedPsf: 

190 return StitchedPsf(*pickle.loads(pkl)) 

191 

192 def writeFits(self, name: str) -> None: 

193 """Persist the PSF as a FITS file.""" 

194 raise NotImplementedError("FITS persistence not implemented for StitchedPsf.") 

195 

196 def readFits(self, name: str) -> None: 

197 """Persist the PSF as a FITS file.""" 

198 raise NotImplementedError("FITS persistence not implemented for StitchedPsf.")