Coverage for tests/test_transformDiaSourceCatalog.py: 98%

131 statements  

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

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 

125 def test_connections_image_type(self): 

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

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

128 with self.subTest(image_type=image_type): 

129 self.config.input_image_type = image_type 

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

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

132 

133 def test_run_quantum_future(self): 

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

135 future-mode difference image. 

136 """ 

137 self.config.input_image_type = "future" 

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

139 diffIm = mock.Mock() 

140 diffIm.to_legacy.return_value = self.exposure 

141 butlerQC = mock.Mock() 

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

143 butlerQC.quantum.dataId = {"band": self.band} 

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

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

146 self.assertIs(run.call_args.kwargs["diffIm"], self.exposure) 

147 self.assertEqual(run.call_args.kwargs["band"], self.band) 

148 

149 def test_run_with_apdb_schema(self): 

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

151 from the exposure. 

152 """ 

153 self.config.doUseSchema = True 

154 transformTask = TransformDiaSourceCatalogTask(initInputs=self.initInputs, 

155 config=self.config) 

156 result = transformTask.run(self.inputCatalog, 

157 self.exposure, 

158 self.band) 

159 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

174 expect_snr = [397.887353515625]*9 

175 expect_snr.append(np.nan) 

176 # Downcast to single precision to match `sdm_schemas` 

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

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

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

180 

181 def test_run_dia_source_wrong_flags(self): 

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

183 that are not in the input schema. 

184 """ 

185 with self.assertRaises(KeyError): 

186 TransformDiaSourceCatalogTask(initInputs=self.initInputsBadFlags) 

187 

188 def test_computeBBoxSize(self): 

189 transform = TransformDiaSourceCatalogTask(initInputs=self.initInputs, 

190 config=self.config) 

191 boxSizes = transform.computeBBoxSizes(self.inputCatalog) 

192 

193 for size in boxSizes: 

194 self.assertEqual(size, self.bboxSize) 

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

196 

197 def test_flag_existence_check(self): 

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

199 

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

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

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

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

204 

205 def test_flag_bitmask(self): 

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

207 """ 

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

209 

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

211 unpacker.makeFlagBitMask(['']) 

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

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

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

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

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

217 'base_PixelFlags_flag_offimage']), 

218 np.uint64(5)) 

219 

220 

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

222 pass 

223 

224 

225def setup_module(module): 

226 lsst.utils.tests.init() 

227 

228 

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

230 lsst.utils.tests.init() 

231 unittest.main()