Coverage for python/lsst/source/injection/utils/test_utils.py: 96%

1# This file is part of source_injection. 

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 

22import astropy.units as u 

23import numpy as np 

24 

25from lsst.afw.geom import makeCdMatrix, makeSkyWcs 

26from lsst.afw.image import ExposureF, makePhotoCalibFromCalibZeroPoint 

27from lsst.geom import Box2I, Extent2I, Point2D, Point2I, SpherePoint, degrees 

28from lsst.images import YX, Box, Image, Mask, SkyProjection, TractFrame 

29from lsst.images.cells import ( 

30 CellCoadd, 

31 CellGrid, 

32 CellGridBounds, 

33 CellIJ, 

34 CellPointSpreadFunction, 

35 CoaddProvenance, 

36 PatchDefinition, 

37) 

38from lsst.ip.isr.isrTask import IsrTask 

39from lsst.meas.algorithms.testUtils import plantSources 

40from lsst.pipe.base import Pipeline 

41from lsst.pipe.base.pipelineIR import ContractIR, LabeledSubset 

42from lsst.pipe.tasks.calibrate import CalibrateTask 

43from lsst.pipe.tasks.characterizeImage import CharacterizeImageTask 

44from lsst.source.injection import generate_injection_catalog 

45 

46 

47def make_test_cell_coadd( 

48 exposure: ExposureF, 

49 cell_shape: YX, 

50 psf_shape: tuple[int, int], 

51 missing: set[CellIJ] | None = None, 

52 image_unit: u.Unit = u.nJy, 

53 **kwargs, 

54) -> CellCoadd: 

55 """Make a cell coadd out of an exposure. 

56 

57 Parameters 

58 ---------- 

59 exposure : `lsst.afw.image.ExposureF` 

60 An exposure with a simple PSF. 

61 cell_shape : `lsst.images.YX` 

62 The shape of each cell. 

63 psf_shape : `tuple[int, int]` 

64 The shape of the PSF image array. 

65 missing : `set[CellIJ]` or `None` 

66 The set of cells that have invalid PSFs, which will be set to nan. 

67 image_unit : `astropy.units.Unit` 

68 The unit of the exposure's image plane. 

69 kwargs 

70 Additional keyword arguments to pass to the CellCoadd constructor. 

71 

72 Returns 

73 ------- 

74 cell_coadd : `lsst.images.cells.CellCoadd` 

75 The exposure as a CellCoadd. 

76 """ 

77 if missing is None: 77 ↛ 78line 77 didn't jump to line 78 because the condition on line 77 was never true

78 missing = set() 

79 

80 exposure_psf = exposure.psf 

81 box = Box.from_legacy(exposure.getBBox()) 

82 cell_grid = CellGrid(bbox=box, cell_shape=cell_shape) 

83 cell_grid_bounds = CellGridBounds(grid=cell_grid, bbox=box, missing=missing) 

84 

85 # Make the cell PSF grid by evaluating the PSF at the center of each cell 

86 cell_psfs = [] 

87 for cell_i in range(cell_grid.grid_size.i): 

88 row_psfs = [] 

89 for cell_j in range(cell_grid.grid_size.j): 

90 if (cell_ij := CellIJ(cell_i, cell_j)) in missing: 

91 psf = np.full(psf_shape, np.nan) 

92 else: 

93 center = cell_grid.bbox_of(cell_ij).to_legacy().getCenter() 

94 psf = exposure_psf.computeKernelImage(center).array 

95 if psf.shape != psf_shape: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true

96 raise ValueError(f"Exposure PSF image at {center=} has shape={psf.shape} != {psf_shape=}") 

97 row_psfs.append(psf) 

98 cell_psfs.append(row_psfs) 

99 

100 cell_psfs = CellPointSpreadFunction(bounds=cell_grid_bounds, array=np.array(cell_psfs)) 

101 # CellCoadds only take TractFrame for now 

102 sky_projection = SkyProjection.from_legacy( 

103 exposure.wcs, 

104 TractFrame(skymap="dummy", tract=0, bbox=box), 

105 pixel_bounds=box, 

106 ) 

107 coadd_contributions = CoaddProvenance.make_empty_contribution_table(n_rows=1) 

108 coadd_inputs = CoaddProvenance.make_empty_contribution_table(n_rows=1) 

109 patch = PatchDefinition( 

110 id=0, 

111 index=YX(0, 0), 

112 inner_bbox=box, 

113 cells=cell_grid, 

114 ) 

115 

116 cell_coadd = CellCoadd( 

117 image=Image.from_legacy(exposure.image, unit=image_unit), 

118 variance=Image.from_legacy(exposure.variance), 

119 mask=Mask.from_legacy(exposure.mask), 

120 patch=patch, 

121 provenance=CoaddProvenance(inputs=coadd_inputs, contributions=coadd_contributions), 

122 psf=cell_psfs, 

123 sky_projection=sky_projection, 

124 **kwargs, 

125 ) 

126 return cell_coadd 

127 

128 

129def make_test_exposure(): 

130 """Make a test exposure with a PSF attached and stars placed randomly. 

131 

132 This function generates a noisy synthetic image with Gaussian PSFs injected 

133 into the frame. 

134 The exposure is returned with a WCS, PhotoCalib and PSF attached. 

135 

136 Returns 

137 ------- 

138 exposure : `lsst.afw.image.Exposure` 

139 Exposure with calibs attached and stars placed randomly. 

140 """ 

141 # Inspired by meas_algorithms test_dynamicDetection.py. 

142 xy0 = Point2I(12345, 67890) # xy0 for image 

143 dims = Extent2I(2345, 2345) # Dimensions of image 

144 bbox = Box2I(xy0, dims) # Bounding box of image 

145 sigma = 3.21 # PSF sigma 

146 buffer = 4.0 # Buffer for star centers around edge 

147 n_sigma = 5.0 # Number of PSF sigmas for kernel 

148 sky = 12345.6 # Initial sky level 

149 num_stars = 100 # Number of stars 

150 noise = np.sqrt(sky) * np.pi * sigma**2 # Poisson noise per PSF 

151 faint = 1.0 * noise # Faintest level for star fluxes 

152 bright = 100.0 * noise # Brightest level for star fluxes 

153 star_bbox = Box2I(bbox) # Area on image in which we can put star centers 

154 star_bbox.grow(-int(buffer * sigma)) # Shrink star_bbox 

155 pixel_scale = 1.0e-5 * degrees # Pixel scale (1E-5 deg = 0.036 arcsec) 

156 

157 # Make an exposure with a PSF attached; place stars randomly. 

158 rng = np.random.default_rng(12345) 

159 stars = [ 

160 (xx, yy, ff, sigma) 

161 for xx, yy, ff in zip( 

162 rng.uniform(star_bbox.getMinX(), star_bbox.getMaxX(), num_stars), 

163 rng.uniform(star_bbox.getMinY(), star_bbox.getMaxY(), num_stars), 

164 np.linspace(faint, bright, num_stars), 

165 ) 

166 ] 

167 exposure = plantSources(bbox, 2 * int(n_sigma * sigma) + 1, sky, stars, True) 

168 

169 # Set WCS and PhotoCalib. 

170 exposure.setWcs( 

171 makeSkyWcs( 

172 crpix=Point2D(0, 0), 

173 crval=SpherePoint(0, 0, degrees), 

174 cdMatrix=makeCdMatrix(scale=pixel_scale), 

175 ) 

176 ) 

177 exposure.setPhotoCalib(makePhotoCalibFromCalibZeroPoint(1e10, 1e8)) 

178 return exposure 

179 

180 

181def make_test_injection_catalog(wcs, bbox): 

182 """Make a test source injection catalog. 

183 

184 This function generates a test source injection catalog consisting of 30 

185 star-like sources of varying magnitude. 

186 

187 Parameters 

188 ---------- 

189 wcs : `lsst.afw.geom.SkyWcs` 

190 WCS associated with the exposure. 

191 bbox : `lsst.geom.Box2I` 

192 Bounding box of the exposure. 

193 

194 Returns 

195 ------- 

196 injection_catalog : `astropy.table.Table` 

197 Source injection catalog. 

198 """ 

199 radec0 = wcs.pixelToSky(bbox.getBeginX(), bbox.getBeginY()) 

200 radec1 = wcs.pixelToSky(bbox.getEndX(), bbox.getEndY()) 

201 ra_lim = sorted([radec0.getRa().asDegrees(), radec1.getRa().asDegrees()]) 

202 dec_lim = sorted([radec0.getDec().asDegrees(), radec1.getDec().asDegrees()]) 

203 injection_catalog = generate_injection_catalog( 

204 ra_lim=ra_lim, 

205 dec_lim=dec_lim, 

206 wcs=wcs, 

207 number=10, 

208 source_type="DeltaFunction", 

209 mag=[10.0, 15.0, 20.0], 

210 ) 

211 return injection_catalog 

212 

213 

214def make_test_reference_pipeline(): 

215 """Make a test reference pipeline containing initial single-frame tasks.""" 

216 reference_pipeline = Pipeline("reference_pipeline") 

217 reference_pipeline.addTask(IsrTask, "isr") 

218 reference_pipeline.addTask(CharacterizeImageTask, "characterizeImage") 

219 reference_pipeline.addTask(CalibrateTask, "calibrate") 

220 reference_pipeline._pipelineIR.labeled_subsets["test_subset"] = LabeledSubset("test_subset", set(), None) 

221 reference_pipeline.addLabelToSubset("test_subset", "isr") 

222 reference_pipeline.addLabelToSubset("test_subset", "characterizeImage") 

223 reference_pipeline.addLabelToSubset("test_subset", "calibrate") 

224 isr_out = "isr.connections.ConnectionsClass(config=isr).outputExposure.name" 

225 char_inp = "characterizeImage.connections.ConnectionsClass(config=characterizeImage).exposure.name" 

226 char_out = "characterizeImage.connections.ConnectionsClass(config=characterizeImage).characterized.name" 

227 calib_inp = "calibrate.connections.ConnectionsClass(config=calibrate).exposure.name" 

228 # When injecting into a post_isr_image, contract1 should be violated. 

229 # The make_injection_pipeline utility should be robust to this. 

230 contract1 = ContractIR( 

231 contract=f"{isr_out} == {char_inp}", 

232 msg="isr output == characterizeImage input", 

233 ) 

234 contract2 = ContractIR( 

235 contract=f"{char_out} == {calib_inp}", 

236 msg="characterizeImage output == calibrate input", 

237 ) 

238 reference_pipeline._pipelineIR.contracts = [contract1, contract2] 

239 return reference_pipeline