Coverage for tests/test_diaPipe.py: 20%

155 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-05-02 04:31 -0700

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 tempfile 

23import unittest 

24from unittest.mock import patch, Mock, MagicMock, DEFAULT 

25import warnings 

26 

27import numpy as np 

28import pandas as pd 

29 

30import lsst.afw.image as afwImage 

31import lsst.afw.table as afwTable 

32import lsst.dax.apdb as daxApdb 

33import lsst.pex.config as pexConfig 

34import lsst.utils.tests 

35from lsst.pipe.base.testUtils import assertValidOutput 

36 

37from lsst.ap.association import DiaPipelineTask 

38from utils_tests import makeExposure, makeDiaObjects 

39 

40 

41def _makeMockDataFrame(): 

42 """Create a new mock of a DataFrame. 

43 

44 Returns 

45 ------- 

46 mock : `unittest.mock.Mock` 

47 A mock guaranteed to accept all operations used by `pandas.DataFrame`. 

48 """ 

49 with warnings.catch_warnings(): 

50 # spec triggers deprecation warnings on DataFrame, but will 

51 # automatically adapt to any removals. 

52 warnings.simplefilter("ignore", category=DeprecationWarning) 

53 return MagicMock(spec=pd.DataFrame()) 

54 

55 

56class TestDiaPipelineTask(unittest.TestCase): 

57 

58 @classmethod 

59 def _makeDefaultConfig(cls, 

60 config_file, 

61 doPackageAlerts=False, 

62 doSolarSystemAssociation=False): 

63 config = DiaPipelineTask.ConfigClass() 

64 config.doConfigureApdb = False 

65 config.apdb_config_url = config_file 

66 config.doPackageAlerts = doPackageAlerts 

67 config.doSolarSystemAssociation = doSolarSystemAssociation 

68 return config 

69 

70 def setUp(self): 

71 # schemas are persisted in both Gen 2 and Gen 3 butler as prototypical catalogs 

72 srcSchema = afwTable.SourceTable.makeMinimalSchema() 

73 srcSchema.addField("base_PixelFlags_flag", type="Flag") 

74 srcSchema.addField("base_PixelFlags_flag_offimage", type="Flag") 

75 self.srcSchema = afwTable.SourceCatalog(srcSchema) 

76 

77 apdb_config = daxApdb.ApdbSql.init_database(db_url="sqlite://") 

78 self.config_file = tempfile.NamedTemporaryFile() 

79 self.addCleanup(self.config_file.close) 

80 apdb_config.save(self.config_file.name) 

81 

82 # TODO: remove on DM-43419 

83 def testConfigApdbNestedOk(self): 

84 config = DiaPipelineTask.ConfigClass() 

85 config.doConfigureApdb = True 

86 with self.assertWarns(FutureWarning): 

87 config.apdb.db_url = "sqlite://" 

88 config.freeze() 

89 config.validate() 

90 

91 # TODO: remove on DM-43419 

92 def testConfigApdbNestedInvalid(self): 

93 config = DiaPipelineTask.ConfigClass() 

94 config.doConfigureApdb = True 

95 # Don't set db_url 

96 config.freeze() 

97 with self.assertRaises(pexConfig.FieldValidationError): 

98 config.validate() 

99 

100 # TODO: remove on DM-43419 

101 def testConfigApdbFileOk(self): 

102 config = DiaPipelineTask.ConfigClass() 

103 config.doConfigureApdb = False 

104 config.apdb_config_url = "some/file/path.yaml" 

105 config.freeze() 

106 config.validate() 

107 

108 # TODO: remove on DM-43419 

109 def testConfigApdbFileInvalid(self): 

110 config = DiaPipelineTask.ConfigClass() 

111 config.doConfigureApdb = False 

112 # Don't set apdb_config_url 

113 config.freeze() 

114 with self.assertRaises(pexConfig.FieldValidationError): 

115 config.validate() 

116 

117 def testRun(self): 

118 """Test running while creating and packaging alerts. 

119 """ 

120 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=True) 

121 

122 def testRunWithSolarSystemAssociation(self): 

123 """Test running while creating and packaging alerts. 

124 """ 

125 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=True) 

126 

127 def testRunWithAlerts(self): 

128 """Test running while creating and packaging alerts. 

129 """ 

130 self._testRun(doPackageAlerts=True, doSolarSystemAssociation=False) 

131 

132 def testRunWithoutAlertsOrSolarSystem(self): 

133 """Test running without creating and packaging alerts. 

134 """ 

135 self._testRun(doPackageAlerts=False, doSolarSystemAssociation=False) 

136 

137 def _testRun(self, doPackageAlerts=False, doSolarSystemAssociation=False): 

138 """Test the normal workflow of each ap_pipe step. 

139 """ 

140 config = self._makeDefaultConfig( 

141 config_file=self.config_file.name, 

142 doPackageAlerts=doPackageAlerts, 

143 doSolarSystemAssociation=doSolarSystemAssociation) 

144 task = DiaPipelineTask(config=config) 

145 # Set DataFrame index testing to always return False. Mocks return 

146 # true for this check otherwise. 

147 task.testDataFrameIndex = lambda x: False 

148 diffIm = Mock(spec=afwImage.ExposureF) 

149 exposure = Mock(spec=afwImage.ExposureF) 

150 template = Mock(spec=afwImage.ExposureF) 

151 diaSrc = _makeMockDataFrame() 

152 ssObjects = _makeMockDataFrame() 

153 ccdExposureIdBits = 32 

154 

155 # Each of these subtasks should be called once during diaPipe 

156 # execution. We use mocks here to check they are being executed 

157 # appropriately. 

158 subtasksToMock = [ 

159 "diaCatalogLoader", 

160 "diaCalculation", 

161 "diaForcedSource", 

162 ] 

163 if doPackageAlerts: 

164 subtasksToMock.append("alertPackager") 

165 else: 

166 self.assertFalse(hasattr(task, "alertPackager")) 

167 

168 if not doSolarSystemAssociation: 

169 self.assertFalse(hasattr(task, "solarSystemAssociator")) 

170 

171 def concatMock(_data, **_kwargs): 

172 return _makeMockDataFrame() 

173 

174 # Mock out the run() methods of these two Tasks to ensure they 

175 # return data in the correct form. 

176 def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, diffIm): 

177 return lsst.pipe.base.Struct(nTotalSsObjects=42, 

178 nAssociatedSsObjects=30, 

179 ssoAssocDiaSources=_makeMockDataFrame(), 

180 unAssocDiaSources=_makeMockDataFrame()) 

181 

182 def associator_run(table, diaObjects, exposure_time=None): 

183 return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3, 

184 matchedDiaSources=_makeMockDataFrame(), 

185 unAssocDiaSources=_makeMockDataFrame(), 

186 longTrailedSources=None) 

187 

188 # apdb isn't a subtask, but still needs to be mocked out for correct 

189 # execution in the test environment. 

190 with patch.multiple(task, **{task: DEFAULT for task in subtasksToMock + ["apdb"]}), \ 

191 patch('lsst.ap.association.diaPipe.pd.concat', side_effect=concatMock), \ 

192 patch('lsst.ap.association.association.AssociationTask.run', 

193 side_effect=associator_run) as mainRun, \ 

194 patch('lsst.ap.association.ssoAssociation.SolarSystemAssociationTask.run', 

195 side_effect=solarSystemAssociator_run) as ssRun: 

196 

197 result = task.run(diaSrc, 

198 ssObjects, 

199 diffIm, 

200 exposure, 

201 template, 

202 ccdExposureIdBits, 

203 "g") 

204 for subtaskName in subtasksToMock: 

205 getattr(task, subtaskName).run.assert_called_once() 

206 assertValidOutput(task, result) 

207 # Exact type and contents of apdbMarker are undefined. 

208 self.assertIsInstance(result.apdbMarker, pexConfig.Config) 

209 meta = task.getFullMetadata() 

210 # Check that the expected metadata has been set. 

211 self.assertEqual(meta["diaPipe.numUpdatedDiaObjects"], 2) 

212 self.assertEqual(meta["diaPipe.numUnassociatedDiaObjects"], 3) 

213 # and that associators ran once or not at all. 

214 mainRun.assert_called_once() 

215 if doSolarSystemAssociation: 

216 ssRun.assert_called_once() 

217 else: 

218 ssRun.assert_not_called() 

219 

220 def test_createDiaObjects(self): 

221 """Test that creating new DiaObjects works as expected. 

222 """ 

223 nSources = 5 

224 diaSources = pd.DataFrame(data=[ 

225 {"ra": 0.04*idx, "dec": 0.04*idx, 

226 "diaSourceId": idx + 1 + nSources, "diaObjectId": 0, 

227 "ssObjectId": 0} 

228 for idx in range(nSources)]) 

229 

230 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False) 

231 task = DiaPipelineTask(config=config) 

232 result = task.createNewDiaObjects(diaSources) 

233 self.assertEqual(nSources, len(result.newDiaObjects)) 

234 self.assertTrue(np.all(np.equal( 

235 result.diaSources["diaObjectId"].to_numpy(), 

236 result.diaSources["diaSourceId"].to_numpy()))) 

237 self.assertTrue(np.all(np.equal( 

238 result.newDiaObjects["diaObjectId"].to_numpy(), 

239 result.diaSources["diaSourceId"].to_numpy()))) 

240 

241 def test_purgeDiaObjects(self): 

242 """Remove diaOjects that are outside an image's bounding box. 

243 """ 

244 

245 config = self._makeDefaultConfig(config_file=self.config_file.name, doPackageAlerts=False) 

246 task = DiaPipelineTask(config=config) 

247 exposure = makeExposure(False, False) 

248 nObj0 = 20 

249 

250 # Create diaObjects 

251 diaObjects = makeDiaObjects(nObj0, exposure) 

252 # Shrink the bounding box so that some of the diaObjects will be outside 

253 bbox = exposure.getBBox() 

254 size = np.minimum(bbox.getHeight(), bbox.getWidth()) 

255 bbox.grow(-size//4) 

256 exposureCut = exposure[bbox] 

257 sizeCut = np.minimum(bbox.getHeight(), bbox.getWidth()) 

258 buffer = 10 

259 bbox.grow(buffer) 

260 

261 def check_diaObjects(bbox, wcs, diaObjects): 

262 raVals = diaObjects.ra.to_numpy() 

263 decVals = diaObjects.dec.to_numpy() 

264 xVals, yVals = wcs.skyToPixelArray(raVals, decVals, degrees=True) 

265 selector = bbox.contains(xVals, yVals) 

266 return selector 

267 

268 selector0 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects) 

269 nIn0 = np.count_nonzero(selector0) 

270 nOut0 = np.count_nonzero(~selector0) 

271 self.assertEqual(nObj0, nIn0 + nOut0) 

272 

273 diaObjects1 = task.purgeDiaObjects(exposureCut.getBBox(), exposureCut.getWcs(), diaObjects, 

274 buffer=buffer) 

275 # Verify that the bounding box was not changed 

276 sizeCheck = np.minimum(exposureCut.getBBox().getHeight(), exposureCut.getBBox().getWidth()) 

277 self.assertEqual(sizeCut, sizeCheck) 

278 selector1 = check_diaObjects(bbox, exposureCut.getWcs(), diaObjects1) 

279 nIn1 = np.count_nonzero(selector1) 

280 nOut1 = np.count_nonzero(~selector1) 

281 nObj1 = len(diaObjects1) 

282 self.assertEqual(nObj1, nIn0) 

283 # Verify that not all diaObjects were removed 

284 self.assertGreater(nObj1, 0) 

285 # Check that some diaObjects were removed 

286 self.assertLess(nObj1, nObj0) 

287 # Verify that no objects outside the bounding box remain 

288 self.assertEqual(nOut1, 0) 

289 # Verify that no objects inside the bounding box were removed 

290 self.assertEqual(nIn1, nIn0) 

291 

292 

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

294 pass 

295 

296 

297def setup_module(module): 

298 lsst.utils.tests.init() 

299 

300 

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

302 lsst.utils.tests.init() 

303 unittest.main()