Coverage for tests/test_templateInputs.py: 98%

162 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-13 10:09 +0000

1# This file is part of analysis_ap. 

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 types 

23import unittest 

24import uuid 

25 

26import astropy.table 

27import numpy as np 

28 

29import lsst.afw.geom as afwGeom 

30import lsst.afw.image as afwImage 

31import lsst.afw.table as afwTable 

32import lsst.geom 

33import lsst.utils.tests 

34from lsst.analysis.ap.templateInputs import ( 

35 _CellCoaddReader, _combine_weights, _LegacyCoaddReader, get_template_inputs, summarize_template_inputs, 

36) 

37from lsst.daf.base import DateTime 

38from lsst.images import YX, Box, DifferenceImageTemplateInfo, Polygon, TractFrame 

39from lsst.images.cells import CellGrid, CoaddProvenance 

40from lsst.images.tests import make_random_sky_projection 

41 

42 

43def _make_wcs(ra=150.0, dec=2.0): 

44 """Return a TAN WCS with 0.2 arcsec pixels and its reference pixel at 

45 the origin. 

46 """ 

47 return afwGeom.makeSkyWcs( 

48 crpix=lsst.geom.Point2D(0.0, 0.0), 

49 crval=lsst.geom.SpherePoint(ra, dec, lsst.geom.degrees), 

50 cdMatrix=afwGeom.makeCdMatrix(scale=0.2*lsst.geom.arcseconds), 

51 ) 

52 

53 

54def _make_ccds(wcs, rows): 

55 """Return a legacy coadd ``ccds`` catalog. 

56 

57 Parameters 

58 ---------- 

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

60 WCS shared by every input. 

61 rows : `list` [`tuple`] 

62 ``(visit, ccd, weight, bbox, valid_box, mjd)`` for each input. 

63 """ 

64 schema = afwTable.ExposureTable.makeMinimalSchema() 

65 schema.addField("visit", type="L", doc="") 

66 schema.addField("ccd", type="I", doc="") 

67 schema.addField("weight", type="D", doc="") 

68 ccds = afwTable.ExposureCatalog(schema) 

69 for visit, ccd, weight, bbox, valid_box, mjd in rows: 

70 record = ccds.addNew() 

71 record["visit"] = visit 

72 record["ccd"] = ccd 

73 record["weight"] = weight 

74 record.setWcs(wcs) 

75 record.setBBox(bbox) 

76 record.setValidPolygon(afwGeom.Polygon(lsst.geom.Box2D(valid_box))) 

77 record.setVisitInfo(afwImage.VisitInfo(exposureTime=30.0, 

78 date=DateTime(mjd, DateTime.MJD, DateTime.TAI))) 

79 return ccds 

80 

81 

82def _sky(wcs, x, y): 

83 """Return the (ra, dec) arrays, in degrees, of pixel positions.""" 

84 points = [wcs.pixelToSky(lsst.geom.Point2D(px, py)) for px, py in zip(x, y)] 

85 return (np.array([p.getRa().asDegrees() for p in points]), 

86 np.array([p.getDec().asDegrees() for p in points])) 

87 

88 

89class CombineWeightsTestCase(lsst.utils.tests.TestCase): 

90 """Test the patch and tract weighting of overlapping coadds.""" 

91 

92 def _inputs(self, weights): 

93 return astropy.table.Table({"coadd_weight": np.array(weights, dtype=float)}) 

94 

95 def testSingleCoadd(self): 

96 fractions = _combine_weights([("a", 1, self._inputs([1.0, 3.0]))], {}) 

97 self.assertEqual(fractions, {"a": 0.25}) 

98 

99 def testPatchOverlap(self): 

100 # Patch weights go as variance**-0.5: 1/1 and 1/2, so 2/3 and 1/3. 

101 coadds = [("a", 1, self._inputs([2.0])), ("b", 1, self._inputs([1.0, 1.0]))] 

102 fractions = _combine_weights(coadds, {"a": 1.0, "b": 4.0}) 

103 self.assertAlmostEqual(fractions["a"]*2.0, 2.0/3.0) 

104 self.assertAlmostEqual(fractions["b"]*2.0, 1.0/3.0) 

105 

106 def testTractOverlap(self): 

107 # Tract 1 averages variances 1 and 4 with weights 2/3 and 1/3, giving 

108 # 2; tract 2 has variance 8. Tract weights 2**-0.5 and 8**-0.5 are in 

109 # the ratio 2:1. 

110 coadds = [("a", 1, self._inputs([1.0])), ("b", 1, self._inputs([1.0])), 

111 ("c", 2, self._inputs([1.0]))] 

112 fractions = _combine_weights(coadds, {"a": 1.0, "b": 4.0, "c": 8.0}) 

113 self.assertAlmostEqual(fractions["a"], 2.0/3.0*2.0/3.0) 

114 self.assertAlmostEqual(fractions["b"], 2.0/3.0*1.0/3.0) 

115 self.assertAlmostEqual(fractions["c"], 1.0/3.0) 

116 self.assertAlmostEqual(sum(fractions.values()), 1.0) 

117 

118 def testUnusableCoadds(self): 

119 # A coadd with no inputs, or with no valid variance, is left out the 

120 # way GetTemplateTask leaves out its pixels. 

121 coadds = [("a", 1, self._inputs([1.0])), ("b", 1, self._inputs([1.0])), 

122 ("c", 1, self._inputs([]))] 

123 fractions = _combine_weights(coadds, {"a": 1.0, "b": np.nan, "c": 1.0}) 

124 self.assertEqual(fractions, {"a": 1.0}) 

125 

126 

127class LegacyCoaddReaderTestCase(lsst.utils.tests.TestCase): 

128 """Test finding inputs in a legacy ``coaddInputs`` catalog.""" 

129 

130 def testFindInputs(self): 

131 wcs = _make_wcs() 

132 bbox = lsst.geom.Box2I(lsst.geom.Point2I(0, 0), lsst.geom.Extent2I(100, 100)) 

133 valid = lsst.geom.Box2I(lsst.geom.Point2I(0, 0), lsst.geom.Extent2I(50, 100)) 

134 ccds = _make_ccds(wcs, [(1, 10, 2.0, bbox, bbox, 57000.0), 

135 (2, 11, 3.0, bbox, valid, 57001.0)]) 

136 ra, dec = _sky(wcs, [20.0, 80.0, 150.0], [50.0, 50.0, 50.0]) 

137 inputs = _LegacyCoaddReader(ccds).find_inputs(ra, dec) 

138 found = {(row["point"], row["input_visit"]) for row in inputs} 

139 # Point 1 is outside the valid polygon of visit 2; point 2 is 

140 # outside every input. 

141 self.assertEqual(found, {(0, 1), (0, 2), (1, 1)}) 

142 row = inputs[(inputs["point"] == 0) & (inputs["input_visit"] == 2)][0] 

143 self.assertEqual(row["input_detector"], 11) 

144 self.assertEqual(row["coadd_weight"], 3.0) 

145 self.assertAlmostEqual(row["mid_mjd_tai"], 57001.0) 

146 self.assertEqual(row["exposure_time"], 30.0) 

147 

148 

149class CellCoaddReaderTestCase(lsst.utils.tests.TestCase): 

150 """Test finding inputs in a cell coadd's provenance.""" 

151 

152 def testFindInputs(self): 

153 bbox = Box.from_shape((200, 200)) 

154 frame = TractFrame(skymap="test", tract=1, bbox=bbox) 

155 projection = make_random_sky_projection(np.random.default_rng(5), frame, bbox) 

156 grid = CellGrid(bbox=bbox, cell_shape=YX(100, 100)) 

157 inputs = CoaddProvenance.make_empty_input_table(2) 

158 inputs["instrument"] = "Cam" 

159 inputs["physical_filter"] = "r" 

160 inputs["visit"] = [1, 2] 

161 inputs["detector"] = [5, 6] 

162 inputs["day_obs"] = 20250101 

163 inputs["polygon"] = [Polygon.from_box(bbox), Polygon.from_box(Box.from_shape((200, 50)))] 

164 contributions = CoaddProvenance.make_empty_contribution_table(3) 

165 contributions["cell_i"] = [0, 0, 1] 

166 contributions["cell_j"] = [0, 0, 0] 

167 contributions["instrument"] = "Cam" 

168 contributions["visit"] = [1, 2, 1] 

169 contributions["detector"] = [5, 6, 5] 

170 contributions["weight"] = [2.0, 3.0, 4.0] 

171 times = {("Cam", 1): (57000.0, 30.0), ("Cam", 2): (57001.0, 15.0)} 

172 reader = _CellCoaddReader(CoaddProvenance(inputs, contributions), projection, grid, times) 

173 # Point 0 is in cell (0, 0) inside both polygons; point 1 is in the 

174 # same cell outside the polygon of visit 2; point 2 is in cell (1, 0), 

175 # where only visit 1 contributes, with a weight of its own. 

176 sky = projection.pixel_to_sky(x=np.array([20.0, 80.0, 20.0]), y=np.array([20.0, 20.0, 150.0])) 

177 found = reader.find_inputs(sky.ra.deg, sky.dec.deg) 

178 self.assertEqual({(row["point"], row["input_visit"], row["coadd_weight"]) for row in found}, 

179 {(0, 1, 2.0), (0, 2, 3.0), (1, 1, 2.0), (2, 1, 4.0)}) 

180 row = found[found["input_visit"] == 2][0] 

181 self.assertEqual((row["mid_mjd_tai"], row["exposure_time"]), (57001.0, 15.0)) 

182 

183 

184class _FakeButler: 

185 """Butler that serves one template and legacy coadd inputs. 

186 

187 Parameters 

188 ---------- 

189 templates : `list` [`lsst.images.DifferenceImageTemplateInfo`] 

190 Templates of every science image. 

191 ccds : `dict` [`uuid.UUID`, `lsst.afw.table.ExposureCatalog`] 

192 Coadd ``ccds`` catalogs keyed by dataset ID. 

193 variances : `dict` [`uuid.UUID`, `float`] 

194 Constant variance of each coadd. 

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

196 WCS of every coadd. 

197 template_components : `set` [`str`] 

198 Components of the template storage class. 

199 """ 

200 

201 def __init__(self, templates, ccds, variances, wcs, template_components=("templates",)): 

202 self._templates = templates 

203 self._ccds = ccds 

204 self._variances = variances 

205 self._wcs = wcs 

206 self._template_components = set(template_components) 

207 

208 def get_dataset_type(self, name): 

209 storage_class = types.SimpleNamespace(allComponents=lambda: self._template_components) 

210 return types.SimpleNamespace(storageClass=storage_class, storageClass_name="Test") 

211 

212 def get_dataset(self, dataset_id): 

213 storage_class = types.SimpleNamespace(allComponents=lambda: {"coaddInputs": None}) 

214 dataset_type = types.SimpleNamespace(storageClass=storage_class, storageClass_name="ExposureF") 

215 return types.SimpleNamespace( 

216 id=dataset_id, datasetType=dataset_type, 

217 makeComponentRef=lambda component: (dataset_id, component), 

218 ) 

219 

220 def get(self, ref, data_id=None, parameters=None): 

221 if isinstance(ref, str): 

222 return self._templates 

223 dataset_id, component = ref 

224 match component: 

225 case "coaddInputs": 

226 return types.SimpleNamespace(ccds=self._ccds[dataset_id]) 

227 case "wcs": 

228 return self._wcs 

229 case "bbox": 

230 return lsst.geom.Box2I(lsst.geom.Point2I(0, 0), lsst.geom.Extent2I(100, 100)) 

231 case "variance": 231 ↛ exitline 231 didn't return from function 'get' because the pattern on line 231 always matched

232 image = afwImage.ImageF(parameters["bbox"]) 

233 image.array[:, :] = self._variances[dataset_id] 

234 return image 

235 

236 

237class GetTemplateInputsTestCase(lsst.utils.tests.TestCase): 

238 """Test the full chain from diaSources to template inputs.""" 

239 

240 def setUp(self): 

241 self.wcs = _make_wcs() 

242 bbox = lsst.geom.Box2I(lsst.geom.Point2I(0, 0), lsst.geom.Extent2I(100, 100)) 

243 self.ids = [uuid.uuid4(), uuid.uuid4()] 

244 # Two patches of one tract cover x < 60 and x >= 40 of the science 

245 # image, which here shares the coadd pixel grid. 

246 self.templates = [ 

247 DifferenceImageTemplateInfo( 

248 skymap="test", tract=1, patch=patch, dataset_id=dataset_id, dataset_run="run", 

249 bounds=Polygon.from_box(Box.from_shape((100, 60), start=(0, x0))), 

250 psf_shape_xx=1.0, psf_shape_yy=1.0, psf_shape_xy=0.0, psf_shape_flag=False, 

251 ) 

252 for patch, dataset_id, x0 in [(0, self.ids[0], 0), (1, self.ids[1], 40)] 

253 ] 

254 self.ccds = { 

255 self.ids[0]: _make_ccds(self.wcs, [(1, 10, 1.0, bbox, bbox, 57000.0), 

256 (2, 10, 3.0, bbox, bbox, 57002.0)]), 

257 self.ids[1]: _make_ccds(self.wcs, [(1, 10, 1.0, bbox, bbox, 57000.0)]), 

258 } 

259 x = np.array([20.0, 50.0]) 

260 y = np.array([50.0, 50.0]) 

261 ra, dec = _sky(self.wcs, x, y) 

262 self.sources = astropy.table.Table({"diaSourceId": [100, 101], "visit": [7, 7], "detector": [3, 3], 

263 "ra": ra, "dec": dec, "x": x, "y": y}) 

264 

265 def testInputs(self): 

266 butler = _FakeButler(self.templates, self.ccds, {self.ids[0]: 1.0, self.ids[1]: 4.0}, self.wcs) 

267 inputs = get_template_inputs(butler, self.sources) 

268 # Source 100 is in patch 0 only. Source 101 is in both, which are 

269 # weighted 2/3 and 1/3; visit 1 reaches it through both. 

270 first = inputs[inputs["diaSourceId"] == 100] 

271 self.assertFloatsAlmostEqual(np.array(first["weight"]), np.array([0.25, 0.75])) 

272 second = inputs[inputs["diaSourceId"] == 101] 

273 self.assertEqual(list(second["patch"]), [0, 0, 1]) 

274 self.assertFloatsAlmostEqual(np.array(second["weight"]), np.array([1/6, 1/2, 1/3]), rtol=1e-12) 

275 summary = summarize_template_inputs(inputs) 

276 self.assertEqual(list(summary["n_inputs"]), [2, 2]) 

277 self.assertFloatsAlmostEqual(np.array(summary["mean_mjd_tai"]), np.array([57001.5, 57001.0])) 

278 

279 def testLegacyTemplateRejected(self): 

280 butler = _FakeButler(self.templates, self.ccds, {}, self.wcs, template_components=()) 

281 with self.assertRaisesRegex(TypeError, "legacy Exposures are not supported"): 

282 get_template_inputs(butler, self.sources) 

283 

284 

285class MemoryTester(lsst.utils.tests.MemoryTestCase): 

286 pass 

287 

288 

289def setup_module(module): 

290 lsst.utils.tests.init() 

291 

292 

293if __name__ == "__main__": 293 ↛ 294line 293 didn't jump to line 294 because the condition on line 293 was never true

294 lsst.utils.tests.init() 

295 unittest.main()