Coverage for tests/test_transformDiaSourceCatalog.py: 98%

152 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-11-08 01:50 -0800

1# This file is part of ap_association 

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 os 

23import unittest 

24from unittest import mock 

25 

26import numpy as np 

27 

28from lsst.ap.association.transformDiaSourceCatalog import (TransformDiaSourceCatalogConfig, 

29 TransformDiaSourceCatalogTask) 

30from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

31import lsst.daf.base as dafBase 

32import lsst.afw.image as afwImage 

33import lsst.geom as geom 

34import lsst.meas.base.tests as measTests 

35from lsst.pipe.base import Struct 

36import lsst.utils.tests 

37 

38from lsst.ap.association.transformDiaSourceCatalog import UnpackApdbFlags 

39 

40TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

41 

42 

43class TestTransformDiaSourceCatalogTask(unittest.TestCase): 

44 def setUp(self): 

45 # Create an instance of random generator with fixed seed. 

46 rng = np.random.default_rng(1234) 

47 

48 self.nSources = 10 

49 # Default PSF size (psfDim in makeEmptyExposure) in TestDataset results 

50 # in an 18 pixel wide source box. 

51 self.bboxSize = 18 

52 self.yLoc = 100 

53 self.bbox = geom.Box2I(geom.Point2I(0, 0), 

54 geom.Extent2I(1024, 1153)) 

55 dataset = measTests.TestDataset(self.bbox) 

56 for srcIdx in range(self.nSources-1): 

57 # Place sources at (index, yLoc), so we can distinguish them later. 

58 dataset.addSource(100000.0, geom.Point2D(srcIdx, self.yLoc)) 

59 # Ensure the last source has no peak `significance` field. 

60 dataset.addSource(100000.0, geom.Point2D(srcIdx+1, self.yLoc), setPeakSignificance=False) 

61 schema = dataset.makeMinimalSchema() 

62 schema.addField("base_PixelFlags_flag", type="Flag") 

63 schema.addField("base_PixelFlags_flag_offimage", type="Flag") 

64 schema.addField("sky_source", type="Flag", doc="Sky objects.") 

65 schema.addField("reliability", doc="real/bogus score of this source", type=float) 

66 self.exposure, self.inputCatalog = dataset.realize(10.0, schema, randomSeed=1234) 

67 # Create schemas for use in initializing the TransformDiaSourceCatalog task. 

68 self.initInputs = {"diaSourceSchema": Struct(schema=schema)} 

69 self.initInputsBadFlags = {"diaSourceSchema": Struct(schema=dataset.makeMinimalSchema())} 

70 

71 self.inputCatalog["reliability"] = rng.random(self.nSources, dtype=np.float32) 

72 

73 self.expId = 4321 

74 self.date = dafBase.DateTime(nsecs=1400000000 * 10**9) 

75 detector = DetectorWrapper(id=23, bbox=self.exposure.getBBox()).detector 

76 visit = afwImage.VisitInfo( 

77 id=self.expId, 

78 exposureTime=200., 

79 date=self.date) 

80 self.exposure.info.id = self.expId + (10000*detector.getId()) 

81 self.exposure.setDetector(detector) 

82 self.exposure.info.setVisitInfo(visit) 

83 self.band = 'g' 

84 self.exposure.setFilter(afwImage.FilterLabel(band=self.band, physical='g.MP9401')) 

85 scale = 2 

86 scaleErr = 1 

87 self.photoCalib = afwImage.PhotoCalib(scale, scaleErr) 

88 self.exposure.setPhotoCalib(self.photoCalib) 

89 

90 self.config = TransformDiaSourceCatalogConfig() 

91 self.config.flagMap = os.path.join(TESTDIR, "data", "test-flag-map.yaml") 

92 self.config.functorFile = os.path.join(TESTDIR, 

93 "data", 

94 "testDiaSource.yaml") 

95 

96 def test_run(self): 

97 """Test output dataFrame is created and values are correctly inserted 

98 from the exposure. 

99 """ 

100 transformTask = TransformDiaSourceCatalogTask(initInputs=self.initInputs, 

101 config=self.config) 

102 result = transformTask.run(self.inputCatalog, 

103 self.exposure, 

104 self.band) 

105 

106 self.assertEqual(len(result.diaSourceTable), len(self.inputCatalog)) 

107 np.testing.assert_array_equal(result.diaSourceTable["bboxSize"], [self.bboxSize]*self.nSources) 

108 np.testing.assert_array_equal(result.diaSourceTable["visit"], 

109 [self.exposure.visitInfo.id]*self.nSources) 

110 np.testing.assert_array_equal(result.diaSourceTable["detector"], 

111 [self.exposure.detector.getId()]*self.nSources) 

112 np.testing.assert_array_equal(result.diaSourceTable["band"], [self.band]*self.nSources) 

113 np.testing.assert_array_equal(result.diaSourceTable["midpointMjdTai"], 

114 [self.date.get(system=dafBase.DateTime.MJD)]*self.nSources) 

115 np.testing.assert_array_equal(result.diaSourceTable["exposureTime"], 

116 [self.exposure.visitInfo.exposureTime]*self.nSources) 

117 np.testing.assert_array_equal(result.diaSourceTable["diaObjectId"], [0]*self.nSources) 

118 np.testing.assert_array_equal(result.diaSourceTable["x"], np.arange(self.nSources)) 

119 # The final snr value should be NaN because it doesn't have a peak significance field. 

120 expect_snr = [397.887353515625]*9 

121 expect_snr.append(np.nan) 

122 # Have to use allclose because assert_array_equal doesn't support equal_nan. 

123 np.testing.assert_allclose(result.diaSourceTable["snr"], expect_snr, equal_nan=True, rtol=0) 

124 np.testing.assert_array_equal(result.diaSourceTable["templateInstrument"], 

125 ["UNKNOWN"]*self.nSources) 

126 np.testing.assert_array_equal(result.diaSourceTable["mixedInstruments"], [False]*self.nSources) 

127 np.testing.assert_array_equal(result.diaSourceTable["templateVersion"], [0]*self.nSources) 

128 

129 def test_run_template_instrument(self): 

130 """Test that the template instrument columns follow the instruments 

131 of the science image and the template. 

132 """ 

133 self.config.doUseSchema = True 

134 transformTask = TransformDiaSourceCatalogTask(initInputs=self.initInputs, 

135 config=self.config) 

136 # instrument, templateInstrument, expected column values 

137 cases = [("LSSTCam", "LSSTCam", ("LSSTCam", False, 0)), 

138 ("LSSTCam", "LSSTComCam", ("LSSTComCam", True, 1)), 

139 ("LSSTCam", None, ("UNKNOWN", False, 0)), 

140 ("LSSTCam", "UNKNOWN", ("UNKNOWN", False, 0)), 

141 (None, "LSSTComCam", ("LSSTComCam", False, 0))] 

142 for instrument, templateInstrument, expected in cases: 

143 with self.subTest(instrument=instrument, templateInstrument=templateInstrument): 

144 result = transformTask.run(self.inputCatalog, self.exposure, self.band, 

145 instrument=instrument, templateInstrument=templateInstrument) 

146 table = result.diaSourceTable 

147 for name, value in zip(("templateInstrument", "mixedInstruments", "templateVersion"), 

148 expected): 

149 np.testing.assert_array_equal(table[name], [value]*self.nSources, err_msg=name) 

150 

151 def test_connections_image_type(self): 

152 """Test that input_image_type sets the storage class of diffIm.""" 

153 for image_type, storageClass in [("legacy", "ExposureF"), ("future", "DifferenceImage")]: 

154 with self.subTest(image_type=image_type): 

155 self.config.input_image_type = image_type 

156 connections = self.config.connections.ConnectionsClass(config=self.config) 

157 self.assertEqual(connections.diffIm.storageClass, storageClass) 

158 

159 def test_run_quantum_future(self): 

160 """Test that runQuantum passes run() the legacy form of a 

161 future-mode difference image, the science instrument, and the 

162 instrument recorded for the template, if any. 

163 """ 

164 self.config.input_image_type = "future" 

165 task = TransformDiaSourceCatalogTask(initInputs=self.initInputs, config=self.config) 

166 for recorded in (True, False): 

167 with self.subTest(recorded=recorded): 

168 diffIm = mock.Mock() 

169 diffIm.to_legacy.return_value = self.exposure 

170 if recorded: 

171 diffIm.templates = [mock.Mock(instrument="LSSTComCam")] 

172 else: 

173 # Accessing the templates of an image without them raises. 

174 del diffIm.templates 

175 butlerQC = mock.Mock() 

176 butlerQC.get.return_value = {"diaSourceCat": self.inputCatalog, "diffIm": diffIm} 

177 butlerQC.quantum.dataId = {"band": self.band, "instrument": "LSSTCam"} 

178 with mock.patch.object(task, "run") as run: 

179 task.runQuantum(butlerQC, mock.Mock(), mock.Mock()) 

180 kwargs = run.call_args.kwargs 

181 self.assertIs(kwargs["diffIm"], self.exposure) 

182 self.assertEqual(kwargs["band"], self.band) 

183 self.assertEqual(kwargs["instrument"], "LSSTCam") 

184 self.assertEqual(kwargs.get("templateInstrument"), "LSSTComCam" if recorded else None) 

185 

186 def test_run_with_apdb_schema(self): 

187 """Test output dataFrame is created and values are correctly inserted 

188 from the exposure. 

189 """ 

190 self.config.doUseSchema = True 

191 transformTask = TransformDiaSourceCatalogTask(initInputs=self.initInputs, 

192 config=self.config) 

193 result = transformTask.run(self.inputCatalog, 

194 self.exposure, 

195 self.band) 

196 

197 self.assertEqual(len(result.diaSourceTable), len(self.inputCatalog)) 

198 np.testing.assert_array_equal(result.diaSourceTable["bboxSize"], [self.bboxSize]*self.nSources) 

199 np.testing.assert_array_equal(result.diaSourceTable["visit"], 

200 [self.exposure.visitInfo.id]*self.nSources) 

201 np.testing.assert_array_equal(result.diaSourceTable["detector"], 

202 [self.exposure.detector.getId()]*self.nSources) 

203 np.testing.assert_array_equal(result.diaSourceTable["band"], [self.band]*self.nSources) 

204 np.testing.assert_array_equal(result.diaSourceTable["midpointMjdTai"], 

205 [self.date.get(system=dafBase.DateTime.MJD)]*self.nSources) 

206 np.testing.assert_array_equal(result.diaSourceTable["exposureTime"], 

207 [self.exposure.visitInfo.exposureTime]*self.nSources) 

208 np.testing.assert_array_equal(result.diaSourceTable["diaObjectId"], [0]*self.nSources) 

209 np.testing.assert_array_equal(result.diaSourceTable["x"], np.arange(self.nSources)) 

210 # The final snr value should be NaN because it doesn't have a peak significance field. 

211 expect_snr = [397.887353515625]*9 

212 expect_snr.append(np.nan) 

213 # Downcast to single precision to match `sdm_schemas` 

214 expect_snr = np.array(expect_snr, dtype=np.float32) 

215 # Have to use allclose because assert_array_equal doesn't support equal_nan. 

216 np.testing.assert_allclose(result.diaSourceTable["snr"], expect_snr, equal_nan=True, rtol=0) 

217 

218 def test_run_dia_source_wrong_flags(self): 

219 """Test that the proper errors are thrown when requesting flag columns 

220 that are not in the input schema. 

221 """ 

222 with self.assertRaises(KeyError): 

223 TransformDiaSourceCatalogTask(initInputs=self.initInputsBadFlags) 

224 

225 def test_computeBBoxSize(self): 

226 transform = TransformDiaSourceCatalogTask(initInputs=self.initInputs, 

227 config=self.config) 

228 boxSizes = transform.computeBBoxSizes(self.inputCatalog) 

229 

230 for size in boxSizes: 

231 self.assertEqual(size, self.bboxSize) 

232 self.assertEqual(len(boxSizes), self.nSources) 

233 

234 def test_flag_existence_check(self): 

235 unpacker = UnpackApdbFlags(self.config.flagMap, "DiaSource") 

236 

237 self.assertTrue(unpacker.flagExists('base_PixelFlags_flag')) 

238 self.assertFalse(unpacker.flagExists('')) 

239 with self.assertRaisesRegex(ValueError, 'column doesNotExist not in flag map'): 

240 unpacker.flagExists('base_PixelFlags_flag', columnName='doesNotExist') 

241 

242 def test_flag_bitmask(self): 

243 """Test that we get the expected bitmask back from supplied flag names. 

244 """ 

245 unpacker = UnpackApdbFlags(self.config.flagMap, "DiaSource") 

246 

247 with self.assertRaisesRegex(ValueError, "flag '' not included"): 

248 unpacker.makeFlagBitMask(['']) 

249 with self.assertRaisesRegex(ValueError, 'column doesNotExist not in flag map'): 

250 unpacker.makeFlagBitMask(['base_PixelFlags_flag'], columnName='doesNotExist') 

251 self.assertEqual(unpacker.makeFlagBitMask(['base_PixelFlags_flag']), np.uint64(1)) 

252 self.assertEqual(unpacker.makeFlagBitMask(['base_PixelFlags_flag_offimage']), np.uint64(4)) 

253 self.assertEqual(unpacker.makeFlagBitMask(['base_PixelFlags_flag', 

254 'base_PixelFlags_flag_offimage']), 

255 np.uint64(5)) 

256 

257 

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

259 pass 

260 

261 

262def setup_module(module): 

263 lsst.utils.tests.init() 

264 

265 

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

267 lsst.utils.tests.init() 

268 unittest.main()