Coverage for tests/test_dipole.py: 13%

272 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-04-14 02:51 -0700

1# This file is part of ip_diffim. 

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 unittest 

23 

24import numpy as np 

25 

26import lsst.utils.tests 

27import lsst.daf.base as dafBase 

28import lsst.afw.image as afwImage 

29import lsst.afw.table as afwTable 

30import lsst.afw.math as afwMath 

31import lsst.geom as geom 

32import lsst.meas.algorithms as measAlg 

33import lsst.ip.diffim as ipDiffim 

34 

35display = False 

36try: 

37 display 

38except NameError: 

39 display = False 

40else: 

41 import lsst.afw.display as afwDisplay 

42 afwDisplay.setDefaultMaskTransparency(75) 

43 

44sigma2fwhm = 2.*np.sqrt(2.*np.log(2.)) 

45 

46 

47def makePluginAndCat(alg, name, control, metadata=False, centroid=None): 

48 schema = afwTable.SourceTable.makeMinimalSchema() 

49 if centroid: 

50 schema.addField(centroid + "_x", type=float) 

51 schema.addField(centroid + "_y", type=float) 

52 schema.addField(centroid + "_flag", type='Flag') 

53 schema.getAliasMap().set("slot_Centroid", centroid) 

54 if metadata: 

55 plugin = alg(control, name, schema, dafBase.PropertySet()) 

56 else: 

57 plugin = alg(control, name, schema) 

58 cat = afwTable.SourceCatalog(schema) 

59 return plugin, cat 

60 

61 

62def createDipole(w, h, xc, yc, scaling=100.0, fracOffset=1.2): 

63 # Make random noise image: set image plane to normal distribution 

64 image = afwImage.MaskedImageF(w, h) 

65 image.set(0) 

66 array = image.getImage().getArray() 

67 array[:, :] = np.random.randn(w, h) 

68 # Set variance to 1.0 

69 var = image.getVariance() 

70 var.set(1.0) 

71 

72 if display: 

73 afwDisplay.Display(frame=1).mtv(image, title="Original image") 

74 afwDisplay.Display(frame=2).mtv(image.getVariance(), title="Original variance") 

75 

76 # Create Psf for dipole creation and measurement 

77 psfSize = 17 

78 psf = measAlg.DoubleGaussianPsf(psfSize, psfSize, 2.0, 3.5, 0.1) 

79 pos = psf.getAveragePosition() 

80 psfFwhmPix = sigma2fwhm*psf.computeShape(pos).getDeterminantRadius() 

81 psfim = psf.computeImage(pos).convertF() 

82 psfim *= scaling/psf.computePeak(pos) 

83 psfw, psfh = psfim.getDimensions() 

84 psfSum = np.sum(psfim.getArray()) 

85 

86 # Create the dipole, offset by fracOffset of the Psf FWHM (pixels) 

87 offset = fracOffset*psfFwhmPix//2 

88 array = image.getImage().getArray() 

89 xp = int(xc - psfw//2 + offset) 

90 yp = int(yc - psfh//2 + offset) 

91 array[yp:yp + psfh, xp:xp + psfw] += psfim.getArray() 

92 

93 xn = int(xc - psfw//2 - offset) 

94 yn = int(yc - psfh//2 - offset) 

95 array[yn:yn + psfh, xn:xn + psfw] -= psfim.getArray() 

96 

97 if display: 

98 afwDisplay.Display(frame=3).mtv(image, title="With dipole") 

99 

100 # Create an exposure, detect positive and negative peaks separately 

101 exp = afwImage.makeExposure(image) 

102 exp.setPsf(psf) 

103 config = measAlg.SourceDetectionConfig() 

104 config.thresholdPolarity = "both" 

105 config.reEstimateBackground = False 

106 schema = afwTable.SourceTable.makeMinimalSchema() 

107 task = measAlg.SourceDetectionTask(schema, config=config) 

108 table = afwTable.SourceTable.make(schema) 

109 results = task.run(table, exp) 

110 if display: 

111 afwDisplay.Display(frame=4).mtv(image, title="Detection plane") 

112 

113 # Merge them together 

114 assert(len(results.sources) == 2) 

115 fpSet = results.positive 

116 fpSet.merge(results.negative, 0, 0, False) 

117 sources = afwTable.SourceCatalog(table) 

118 fpSet.makeSources(sources) 

119 assert(len(sources) == 1) 

120 s = sources[0] 

121 assert(len(s.getFootprint().getPeaks()) == 2) 

122 

123 return psf, psfSum, exp, s 

124 

125 

126class DipoleAlgorithmTest(lsst.utils.tests.TestCase): 

127 """ A test case for dipole algorithms""" 

128 

129 def setUp(self): 

130 np.random.seed(666) 

131 self.w, self.h = 100, 100 # size of image 

132 self.xc, self.yc = 50, 50 # location of center of dipole 

133 

134 def testNaiveDipoleCentroid(self): 

135 control = ipDiffim.DipoleCentroidControl() 

136 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

137 plugin, cat = makePluginAndCat(ipDiffim.NaiveDipoleCentroid, "test", control, centroid="centroid") 

138 source = cat.addNew() 

139 source.set("centroid_x", 50) 

140 source.set("centroid_y", 50) 

141 source.setFootprint(s.getFootprint()) 

142 plugin.measure(source, exposure) 

143 for key in ("_pos_x", "_pos_y", "_pos_xErr", "_pos_yErr", "_pos_flag", 

144 "_neg_x", "_neg_y", "_neg_xErr", "_neg_yErr", "_neg_flag"): 

145 try: 

146 source.get("test" + key) 

147 except Exception: 

148 self.fail() 

149 

150 def testNaiveDipoleFluxControl(self): 

151 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

152 control = ipDiffim.DipoleFluxControl() 

153 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

154 plugin, cat = makePluginAndCat(ipDiffim.NaiveDipoleFlux, "test", control, centroid="centroid") 

155 source = cat.addNew() 

156 source.set("centroid_x", 50) 

157 source.set("centroid_y", 50) 

158 source.setFootprint(s.getFootprint()) 

159 plugin.measure(source, exposure) 

160 for key in ("_pos_instFlux", "_pos_instFluxErr", "_pos_flag", "_npos", 

161 "_neg_instFlux", "_neg_instFluxErr", "_neg_flag", "_nneg"): 

162 try: 

163 source.get("test" + key) 

164 except Exception: 

165 self.fail() 

166 

167 def testPsfDipoleFluxControl(self): 

168 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

169 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

170 control = ipDiffim.PsfDipoleFluxControl() 

171 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

172 plugin, cat = makePluginAndCat(ipDiffim.PsfDipoleFlux, "test", control, centroid="centroid") 

173 source = cat.addNew() 

174 source.set("centroid_x", 50) 

175 source.set("centroid_y", 50) 

176 source.setFootprint(s.getFootprint()) 

177 plugin.measure(source, exposure) 

178 for key in ("_pos_instFlux", "_pos_instFluxErr", "_pos_flag", 

179 "_neg_instFlux", "_neg_instFluxErr", "_neg_flag"): 

180 try: 

181 source.get("test" + key) 

182 except Exception: 

183 self.fail() 

184 

185 def testAll(self): 

186 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

187 self.measureDipole(s, exposure) 

188 

189 def _makeModel(self, exposure, psf, fp, negCenter, posCenter): 

190 

191 negPsf = psf.computeImage(negCenter).convertF() 

192 posPsf = psf.computeImage(posCenter).convertF() 

193 negPeak = psf.computePeak(negCenter) 

194 posPeak = psf.computePeak(posCenter) 

195 negPsf /= negPeak 

196 posPsf /= posPeak 

197 

198 model = afwImage.ImageF(fp.getBBox()) 

199 negModel = afwImage.ImageF(fp.getBBox()) 

200 posModel = afwImage.ImageF(fp.getBBox()) 

201 

202 # The center of the Psf should be at negCenter, posCenter 

203 negPsfBBox = negPsf.getBBox() 

204 posPsfBBox = posPsf.getBBox() 

205 modelBBox = model.getBBox() 

206 

207 # Portion of the negative Psf that overlaps the montage 

208 negOverlapBBox = geom.Box2I(negPsfBBox) 

209 negOverlapBBox.clip(modelBBox) 

210 self.assertFalse(negOverlapBBox.isEmpty()) 

211 

212 # Portion of the positivePsf that overlaps the montage 

213 posOverlapBBox = geom.Box2I(posPsfBBox) 

214 posOverlapBBox.clip(modelBBox) 

215 self.assertFalse(posOverlapBBox.isEmpty()) 

216 

217 negPsfSubim = type(negPsf)(negPsf, negOverlapBBox) 

218 modelSubim = type(model)(model, negOverlapBBox) 

219 negModelSubim = type(negModel)(negModel, negOverlapBBox) 

220 modelSubim += negPsfSubim # just for debugging 

221 negModelSubim += negPsfSubim # for fitting 

222 

223 posPsfSubim = type(posPsf)(posPsf, posOverlapBBox) 

224 modelSubim = type(model)(model, posOverlapBBox) 

225 posModelSubim = type(posModel)(posModel, posOverlapBBox) 

226 modelSubim += posPsfSubim 

227 posModelSubim += posPsfSubim 

228 

229 data = afwImage.ImageF(exposure.getMaskedImage().getImage(), fp.getBBox()) 

230 var = afwImage.ImageF(exposure.getMaskedImage().getVariance(), fp.getBBox()) 

231 matrixNorm = 1./np.sqrt(np.median(var.getArray())) 

232 

233 if display: 

234 afwDisplay.Display(frame=5).mtv(model, title="Unfitted model") 

235 afwDisplay.Display(frame=6).mtv(data, title="Data") 

236 

237 posPsfSum = np.sum(posPsf.getArray()) 

238 negPsfSum = np.sum(negPsf.getArray()) 

239 

240 M = np.array((np.ravel(negModel.getArray()), np.ravel(posModel.getArray()))).T.astype(np.float64) 

241 B = np.array((np.ravel(data.getArray()))).astype(np.float64) 

242 M *= matrixNorm 

243 B *= matrixNorm 

244 

245 # Numpy solution 

246 fneg0, fpos0 = np.linalg.lstsq(M, B, rcond=-1)[0] 

247 

248 # Afw solution 

249 lsq = afwMath.LeastSquares.fromDesignMatrix(M, B, afwMath.LeastSquares.DIRECT_SVD) 

250 fneg, fpos = lsq.getSolution() 

251 

252 # Should be exaxtly the same as each other 

253 self.assertAlmostEqual(1e-2*fneg0, 1e-2*fneg) 

254 self.assertAlmostEqual(1e-2*fpos0, 1e-2*fpos) 

255 

256 # Recreate model 

257 fitted = afwImage.ImageF(fp.getBBox()) 

258 negFit = type(negPsf)(negPsf, negOverlapBBox, afwImage.PARENT, True) 

259 negFit *= float(fneg) 

260 posFit = type(posPsf)(posPsf, posOverlapBBox, afwImage.PARENT, True) 

261 posFit *= float(fpos) 

262 

263 fitSubim = type(fitted)(fitted, negOverlapBBox) 

264 fitSubim += negFit 

265 fitSubim = type(fitted)(fitted, posOverlapBBox) 

266 fitSubim += posFit 

267 if display: 

268 afwDisplay.Display(frame=7).mtv(fitted, title="Fitted model") 

269 

270 fitted -= data 

271 

272 if display: 

273 afwDisplay.Display(frame=8).mtv(fitted, title="Residuals") 

274 

275 fitted *= fitted 

276 fitted /= var 

277 

278 if display: 

279 afwDisplay.Display(frame=9).mtv(fitted, title="Chi2") 

280 

281 return fneg, negPsfSum, fpos, posPsfSum, fitted 

282 

283 def testPsfDipoleFit(self, scaling=100.): 

284 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc, scaling=scaling) 

285 source = self.measureDipole(s, exposure) 

286 # Recreate the simultaneous joint Psf fit in python 

287 fp = source.getFootprint() 

288 peaks = fp.getPeaks() 

289 speaks = [(p.getPeakValue(), p) for p in peaks] 

290 speaks.sort() 

291 dpeaks = [speaks[0][1], speaks[-1][1]] 

292 

293 negCenter = geom.Point2D(dpeaks[0].getFx(), dpeaks[0].getFy()) 

294 posCenter = geom.Point2D(dpeaks[1].getFx(), dpeaks[1].getFy()) 

295 

296 fneg, negPsfSum, fpos, posPsfSum, residIm = self._makeModel(exposure, psf, fp, negCenter, posCenter) 

297 

298 # Should be close to the same as the inputs; as fracOffset 

299 # gets smaller this will be worse. This works for scaling = 

300 # 100. 

301 self.assertAlmostEqual(1e-2*scaling, -1e-2*fneg, 2) 

302 self.assertAlmostEqual(1e-2*scaling, 1e-2*fpos, 2) 

303 

304 # Now compare the LeastSquares results fitted here to the C++ 

305 # implementation: Since total flux is returned, and this is of 

306 # order 1e4 for this default test, scale back down so that 

307 # assertAlmostEqual behaves reasonably (the comparison to 2 

308 # places means to 0.01). Also note that PsfDipoleFlux returns 

309 # the total flux, while here we are just fitting for the 

310 # scaling of the Psf. Therefore the comparison is 

311 # fneg*negPsfSum to flux.dipole.psf.neg. 

312 self.assertAlmostEqual(1e-4*fneg*negPsfSum, 

313 1e-4*source.get("ip_diffim_PsfDipoleFlux_neg_instFlux"), 

314 2) 

315 self.assertAlmostEqual(1e-4*fpos*posPsfSum, 

316 1e-4*source.get("ip_diffim_PsfDipoleFlux_pos_instFlux"), 

317 2) 

318 

319 self.assertGreater(source.get("ip_diffim_PsfDipoleFlux_pos_instFluxErr"), 0.0) 

320 self.assertGreater(source.get("ip_diffim_PsfDipoleFlux_neg_instFluxErr"), 0.0) 

321 self.assertFalse(source.get("ip_diffim_PsfDipoleFlux_neg_flag")) 

322 self.assertFalse(source.get("ip_diffim_PsfDipoleFlux_pos_flag")) 

323 

324 self.assertAlmostEqual(source.get("ip_diffim_PsfDipoleFlux_centroid_x"), 50.0, 1) 

325 self.assertAlmostEqual(source.get("ip_diffim_PsfDipoleFlux_centroid_y"), 50.0, 1) 

326 self.assertAlmostEqual(source.get("ip_diffim_PsfDipoleFlux_neg_centroid_x"), negCenter[0], 1) 

327 self.assertAlmostEqual(source.get("ip_diffim_PsfDipoleFlux_neg_centroid_y"), negCenter[1], 1) 

328 self.assertAlmostEqual(source.get("ip_diffim_PsfDipoleFlux_pos_centroid_x"), posCenter[0], 1) 

329 self.assertAlmostEqual(source.get("ip_diffim_PsfDipoleFlux_pos_centroid_y"), posCenter[1], 1) 

330 self.assertFalse(source.get("ip_diffim_PsfDipoleFlux_neg_flag")) 

331 self.assertFalse(source.get("ip_diffim_PsfDipoleFlux_pos_flag")) 

332 

333 self.assertGreater(source.get("ip_diffim_PsfDipoleFlux_chi2dof"), 0.0) 

334 

335 def measureDipole(self, s, exp): 

336 msConfig = ipDiffim.DipoleMeasurementConfig() 

337 schema = afwTable.SourceTable.makeMinimalSchema() 

338 schema.addField("centroid_x", type=float) 

339 schema.addField("centroid_y", type=float) 

340 schema.addField("centroid_flag", type='Flag') 

341 task = ipDiffim.DipoleMeasurementTask(schema, config=msConfig) 

342 measCat = afwTable.SourceCatalog(schema) 

343 measCat.defineCentroid("centroid") 

344 source = measCat.addNew() 

345 source.set("centroid_x", self.xc) 

346 source.set("centroid_y", self.yc) 

347 source.setFootprint(s.getFootprint()) 

348 # Then run the default SFM task. Results not checked 

349 task.run(measCat, exp) 

350 return measCat[0] 

351 

352 def testDipoleAnalysis(self): 

353 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

354 source = self.measureDipole(s, exposure) 

355 dpAnalysis = ipDiffim.DipoleAnalysis() 

356 dpAnalysis(source) 

357 

358 def testDipoleDeblender(self): 

359 psf, psfSum, exposure, s = createDipole(self.w, self.h, self.xc, self.yc) 

360 source = self.measureDipole(s, exposure) 

361 dpDeblender = ipDiffim.DipoleDeblender() 

362 dpDeblender(source, exposure) 

363 

364 

365class DipoleMeasurementTaskTest(lsst.utils.tests.TestCase): 

366 """A test case for the DipoleMeasurementTask. Essentially just 

367 test the classification flag since the invididual algorithms are 

368 tested above""" 

369 

370 def setUp(self): 

371 np.random.seed(666) 

372 self.config = ipDiffim.DipoleMeasurementConfig() 

373 

374 def tearDown(self): 

375 del self.config 

376 

377 def testMeasure(self): 

378 schema = afwTable.SourceTable.makeMinimalSchema() 

379 task = ipDiffim.DipoleMeasurementTask(schema, config=self.config) 

380 table = afwTable.SourceTable.make(schema) 

381 sources = afwTable.SourceCatalog(table) 

382 source = sources.addNew() 

383 # make fake image 

384 psf, psfSum, exposure, s = createDipole(100, 100, 50, 50) 

385 

386 # set it in source with the appropriate schema 

387 source.setFootprint(s.getFootprint()) 

388 task.run(sources, exposure) 

389 self.assertEqual(source.get("ip_diffim_ClassificationDipole_value"), 1.0) 

390 

391 

392class TestMemory(lsst.utils.tests.MemoryTestCase): 

393 pass 

394 

395 

396def setup_module(module): 

397 lsst.utils.tests.init() 

398 

399 

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

401 lsst.utils.tests.init() 

402 unittest.main()