Coverage for examples/designdoc.py: 0%

365 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-11 03:20 -0800

1import pylab as plt 

2import numpy as np 

3import matplotlib 

4matplotlib.use('Agg') 

5 

6import lsst.afw.image as afwImage # noqa E402 

7import lsst.afw.detection as afwDet # noqa E402 

8import lsst.log # noqa E402 

9 

10from .utils import footprintToImage, getExtent, get_sigma1, getFamilies, plotDeblendFamily, readCatalog # noqa E402 

11from .suprime import getSuprimeDataref # noqa E402 

12 

13 

14def main(): 

15 ''' 

16 Runs the deblender and creates plots for the "design document", 

17 doc/design.tex. See the file NOTES for how to get set up to the 

18 point where you can actually run this on data. 

19 ''' 

20 

21 from optparse import OptionParser 

22 parser = OptionParser() 

23 parser.add_option('--root', dest='root', help='Root directory for Subaru data') 

24 parser.add_option('--outroot', '-o', dest='outroot', help='Output root directory for Subaru data') 

25 parser.add_option('--sources', help='Read a FITS table of sources') 

26 parser.add_option('--calexp', help='Read a FITS calexp') 

27 parser.add_option('--psf', help='Read a FITS PSF') 

28 

29 parser.add_option('--drill', '-D', dest='drill', action='append', type=str, default=[], 

30 help='Drill down on individual source IDs') 

31 parser.add_option('--drillxy', dest='drillxy', action='append', type=str, default=[], 

32 help='Drill down on individual source positions, eg 132,46;54,67') 

33 parser.add_option('--visit', dest='visit', type=int, default=108792, help='Suprimecam visit id') 

34 parser.add_option('--ccd', dest='ccd', type=int, default=5, help='Suprimecam CCD number') 

35 parser.add_option('--prefix', dest='prefix', default='design-', help='plot filename prefix') 

36 parser.add_option('--suffix', dest='suffix', default=None, help='plot filename suffix (default: ".png")') 

37 parser.add_option( 

38 '--pat', dest='pat', 

39 help='Plot filename pattern: eg, "design-%(pid)04i-%(name).png"; overrides --prefix and --suffix') 

40 parser.add_option('--pdf', dest='pdf', action='store_true', default=False, help='save in PDF format?') 

41 parser.add_option('-v', dest='verbose', action='store_true') 

42 parser.add_option('--figw', dest='figw', type=float, help='Figure window width (inches)', 

43 default=4.) 

44 parser.add_option('--figh', dest='figh', type=float, help='Figure window height (inches)', 

45 default=4.) 

46 parser.add_option('--order', dest='order', type=str, help='Child order: eg 3,0,1,2') 

47 

48 parser.add_option('--sdss', dest='sec', action='store_const', const='sdss', 

49 help='Produce plots for the SDSS section.') 

50 parser.add_option('--mono', dest='sec', action='store_const', const='mono', 

51 help='Produce plots for the "monotonic" section.') 

52 parser.add_option('--median', dest='sec', action='store_const', const='median', 

53 help='Produce plots for the "median filter" section.') 

54 parser.add_option('--ramp', dest='sec', action='store_const', const='ramp', 

55 help='Produce plots for the "ramp edges" section.') 

56 parser.add_option('--ramp2', dest='sec', action='store_const', const='ramp2', 

57 help='Produce plots for the "ramp edges + stray flux" section.') 

58 parser.add_option('--patch', dest='sec', action='store_const', const='patch', 

59 help='Produce plots for the "patch edges" section.') 

60 

61 opt, args = parser.parse_args() 

62 

63 # Logging 

64 if opt.verbose: 

65 lsst.log.setLevel('', lsst.log.DEBUG) 

66 else: 

67 lsst.log.setLevel('', lsst.log.INFO) 

68 

69 if opt.sec is None: 

70 opt.sec = 'sdss' 

71 if opt.pdf: 

72 if opt.suffix is None: 

73 opt.suffix = '' 

74 opt.suffix += '.pdf' 

75 if not opt.suffix: 

76 opt.suffix = '.png' 

77 

78 if opt.pat: 

79 plotpattern = opt.pat 

80 else: 

81 plotpattern = opt.prefix + '%(pid)04i-%(name)s' + opt.suffix 

82 

83 if opt.order is not None: 

84 opt.order = [int(x) for x in opt.order.split(',')] 

85 invorder = np.zeros(len(opt.order)) 

86 invorder[opt.order] = np.arange(len(opt.order)) 

87 

88 def mapchild(i): 

89 if opt.order is None: 

90 return i 

91 return invorder[i] 

92 

93 def savefig(pid, figname): 

94 fn = plotpattern % dict(pid=pid, name=figname) 

95 plt.savefig(fn) 

96 

97 # Load data using the butler, if desired 

98 dr = None 

99 if opt.sources is None or opt.calexp is None: 

100 print('Creating DataRef...') 

101 dr = getSuprimeDataref(opt.visit, opt.ccd, rootdir=opt.root, 

102 outrootdir=opt.outroot) 

103 print('Got', dr) 

104 

105 # Which parent ids / deblend families are we going to plot? 

106 keepids = None 

107 if len(opt.drill): 

108 keepids = [] 

109 for d in opt.drill: 

110 for dd in d.split(','): 

111 keepids.append(int(dd)) 

112 print('Keeping parent ids', keepids) 

113 

114 keepxys = None 

115 if len(opt.drillxy): 

116 keepxys = [] 

117 for d in opt.drillxy: 

118 for dd in d.split(';'): 

119 xy = dd.split(',') 

120 assert(len(xy) == 2) 

121 keepxys.append((int(xy[0]), int(xy[1]))) 

122 print('Keeping parents at xy', keepxys) 

123 

124 # Read from butler or local file 

125 cat = readCatalog(opt.sources, None, dataref=dr, keepids=keepids, 

126 keepxys=keepxys, patargs=dict(visit=opt.visit, ccd=opt.ccd)) 

127 print('Got', len(cat), 'sources') 

128 

129 # Load data from butler or local files 

130 if opt.calexp is not None: 

131 print('Reading exposure from', opt.calexp) 

132 exposure = afwImage.ExposureF(opt.calexp) 

133 else: 

134 exposure = dr.get('calexp') 

135 print('Exposure', exposure) 

136 mi = exposure.getMaskedImage() 

137 

138 if opt.psf is not None: 

139 print('Reading PSF from', opt.psf) 

140 psf = afwDet.Psf.readFits(opt.psf) 

141 print('Got', psf) 

142 elif dr: 

143 psf = dr.get('psf') 

144 else: 

145 psf = exposure.getPsf() 

146 

147 sigma1 = get_sigma1(mi) 

148 

149 fams = getFamilies(cat) 

150 print(len(fams), 'deblend families') 

151 

152 if False: 

153 for j, (parent, children) in enumerate(fams): 

154 print('parent', parent) 

155 print('children', children) 

156 plotDeblendFamily(mi, parent, children, cat, sigma1, ellipses=False) 

157 fn = '%04i.png' % parent.getId() 

158 plt.savefig(fn) 

159 print('wrote', fn) 

160 

161 def nlmap(X): 

162 return np.arcsinh(X / (3.*sigma1)) 

163 

164 def myimshow(im, **kwargs): 

165 kwargs = kwargs.copy() 

166 mn = kwargs.get('vmin', -5*sigma1) 

167 kwargs['vmin'] = nlmap(mn) 

168 mx = kwargs.get('vmax', 100*sigma1) 

169 kwargs['vmax'] = nlmap(mx) 

170 plt.imshow(nlmap(im), **kwargs) 

171 plt.figure(figsize=(opt.figw, opt.figh)) 

172 plt.subplot(1, 1, 1) 

173 plt.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.99, 

174 wspace=0.05, hspace=0.1) 

175 

176 # Make plots for each deblend family. 

177 

178 for j, (parent, children) in enumerate(fams): 

179 print('parent', parent.getId()) 

180 print('children', [ch.getId() for ch in children]) 

181 print('parent x,y', parent.getX(), parent.getY()) 

182 

183 pid = parent.getId() 

184 fp = parent.getFootprint() 

185 bb = fp.getBBox() 

186 pim = footprintToImage(parent.getFootprint(), mi).getArray() 

187 pext = getExtent(bb) 

188 imargs = dict(interpolation='nearest', origin='lower', vmax=pim.max() * 0.95, vmin=-3.*sigma1) 

189 pksty = dict(linestyle='None', marker='+', color='r', mew=3, ms=20, alpha=0.6) 

190 

191 plt.clf() 

192 myimshow(afwImage.ImageF(mi.getImage(), bb).getArray(), **imargs) 

193 plt.gray() 

194 plt.xticks([]) 

195 plt.yticks([]) 

196 savefig(pid, 'image') 

197 

198 # Parent footprint 

199 plt.clf() 

200 myimshow(pim, extent=pext, **imargs) 

201 plt.gray() 

202 pks = fp.getPeaks() 

203 plt.plot([pk.getIx() for pk in pks], [pk.getIy() for pk in pks], **pksty) 

204 plt.xticks([]) 

205 plt.yticks([]) 

206 plt.axis(pext) 

207 savefig(pid, 'parent') 

208 

209 from lsst.meas.deblender.baseline import deblend 

210 

211 xc = int((bb.getMinX() + bb.getMaxX()) / 2.) 

212 yc = int((bb.getMinY() + bb.getMaxY()) / 2.) 

213 if hasattr(psf, 'getFwhm'): 

214 psf_fwhm = psf.getFwhm(xc, yc) 

215 else: 

216 psf_fwhm = psf.computeShape(psf.getAveragePosition()).getDeterminantRadius() * 2.35 

217 

218 # Each section of the design doc runs the deblender with different args. 

219 

220 kwargs = dict(sigma1=sigma1, verbose=opt.verbose, 

221 getTemplateSum=True) 

222 

223 basic = kwargs.copy() 

224 basic.update(fit_psfs=False, 

225 median_smooth_template=False, 

226 monotonic_template=False, 

227 lstsq_weight_templates=False, 

228 assignStrayFlux=False, 

229 rampFluxAtEdge=False, 

230 patchEdges=False) 

231 

232 if opt.sec == 'sdss': 

233 # SDSS intro 

234 kwargs = basic 

235 kwargs.update(lstsq_weight_templates=True) 

236 

237 elif opt.sec == 'mono': 

238 kwargs = basic 

239 kwargs.update(lstsq_weight_templates=True, 

240 monotonic_template=True) 

241 elif opt.sec == 'median': 

242 kwargs = basic 

243 kwargs.update(lstsq_weight_templates=True, 

244 median_smooth_template=True, 

245 monotonic_template=True) 

246 elif opt.sec == 'ramp': 

247 kwargs = basic 

248 kwargs.update(median_smooth_template=True, 

249 monotonic_template=True, 

250 rampFluxAtEdge=True) 

251 

252 elif opt.sec == 'ramp2': 

253 kwargs = basic 

254 kwargs.update(median_smooth_template=True, 

255 monotonic_template=True, 

256 rampFluxAtEdge=True, 

257 assignStrayFlux=True) 

258 

259 elif opt.sec == 'patch': 

260 kwargs = basic 

261 kwargs.update(median_smooth_template=True, 

262 monotonic_template=True, 

263 patchEdges=True) 

264 

265 else: 

266 raise 'Unknown section: "%s"' % opt.sec 

267 

268 print('Running deblender with kwargs:', kwargs) 

269 res = deblend(fp, mi, psf, psf_fwhm, **kwargs) 

270 # print('got result with', [x for x in dir(res) if not x.startswith('__')]) 

271 # for pk in res.peaks: 

272 # print('got peak with', [x for x in dir(pk) if not x.startswith('__')]) 

273 # print(' deblend as psf?', pk.deblend_as_psf) 

274 

275 # Find bounding-box of all templates. 

276 tbb = fp.getBBox() 

277 for pkres, pk in zip(res.peaks, pks): 

278 tbb.include(pkres.template_foot.getBBox()) 

279 print('Bounding-box of all templates:', tbb) 

280 

281 # Sum-of-templates plot 

282 tsum = np.zeros((tbb.getHeight(), tbb.getWidth())) 

283 tx0, ty0 = tbb.getMinX(), tbb.getMinY() 

284 

285 # Sum-of-deblended children plot(s) 

286 # "heavy" bbox == template bbox. 

287 hsum = np.zeros((tbb.getHeight(), tbb.getWidth())) 

288 hsum2 = np.zeros((tbb.getHeight(), tbb.getWidth())) 

289 

290 # Sum of templates from the deblender itself 

291 plt.clf() 

292 t = res.templateSum 

293 myimshow(t.getArray(), extent=getExtent(t.getBBox()), **imargs) 

294 plt.gray() 

295 plt.xticks([]) 

296 plt.yticks([]) 

297 savefig(pid, 'tsum1') 

298 

299 # Make plots for each deblended child (peak) 

300 

301 k = 0 

302 for pkres, pk in zip(res.peaks, pks): 

303 

304 heavy = pkres.get_flux_portion() 

305 if heavy is None: 

306 print('Child has no HeavyFootprint -- skipping') 

307 continue 

308 

309 kk = mapchild(k) 

310 

311 w = pkres.template_weight 

312 

313 cfp = pkres.template_foot 

314 cbb = cfp.getBBox() 

315 cext = getExtent(cbb) 

316 

317 # Template image 

318 tim = pkres.template_mimg.getImage() 

319 timext = cext 

320 tim = tim.getArray() 

321 

322 (x0, x1, y0, y1) = timext 

323 print('tim ext', timext) 

324 tsum[y0-ty0:y1-ty0, x0-tx0:x1-tx0] += tim 

325 

326 # "Heavy" image -- flux assigned to child 

327 him = footprintToImage(heavy).getArray() 

328 hext = getExtent(heavy.getBBox()) 

329 

330 (x0, x1, y0, y1) = hext 

331 hsum[y0-ty0:y1-ty0, x0-tx0:x1-tx0] += him 

332 

333 # "Heavy" without stray flux 

334 h2 = pkres.get_flux_portion(strayFlux=False) 

335 him2 = footprintToImage(h2).getArray() 

336 hext2 = getExtent(h2.getBBox()) 

337 (x0, x1, y0, y1) = hext2 

338 hsum2[y0-ty0:y1-ty0, x0-tx0:x1-tx0] += him2 

339 

340 if opt.sec == 'median': 

341 try: 

342 med = pkres.median_filtered_template 

343 except Exception: 

344 med = pkres.orig_template 

345 

346 for im, nm in [(pkres.orig_template, 'symm'), (med, 'med')]: 

347 # print('im:', im) 

348 plt.clf() 

349 myimshow(im.getArray(), extent=cext, **imargs) 

350 plt.gray() 

351 plt.xticks([]) 

352 plt.yticks([]) 

353 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

354 plt.axis(pext) 

355 savefig(pid, nm + '%i' % (kk)) 

356 

357 # Template 

358 plt.clf() 

359 myimshow(pkres.template_mimg.getImage().getArray() / w, extent=cext, **imargs) 

360 plt.gray() 

361 plt.xticks([]) 

362 plt.yticks([]) 

363 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

364 plt.axis(pext) 

365 savefig(pid, 't%i' % (kk)) 

366 

367 # Weighted template 

368 plt.clf() 

369 myimshow(tim, extent=cext, **imargs) 

370 plt.gray() 

371 plt.xticks([]) 

372 plt.yticks([]) 

373 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

374 plt.axis(pext) 

375 savefig(pid, 'tw%i' % (kk)) 

376 

377 # "Heavy" 

378 plt.clf() 

379 myimshow(him, extent=hext, **imargs) 

380 plt.gray() 

381 plt.xticks([]) 

382 plt.yticks([]) 

383 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

384 plt.axis(pext) 

385 savefig(pid, 'h%i' % (kk)) 

386 

387 # Original symmetric template 

388 plt.clf() 

389 t = pkres.orig_template 

390 foot = pkres.orig_foot 

391 myimshow(t.getArray(), extent=getExtent(foot.getBBox()), **imargs) 

392 plt.gray() 

393 plt.xticks([]) 

394 plt.yticks([]) 

395 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

396 plt.axis(pext) 

397 savefig(pid, 'o%i' % (kk)) 

398 

399 if opt.sec == 'patch' and pkres.patched: 

400 pass 

401 

402 if opt.sec in ['ramp', 'ramp2'] and pkres.has_ramped_template: 

403 

404 # Ramped template 

405 plt.clf() 

406 t = pkres.ramped_template 

407 myimshow(t.getArray(), extent=getExtent(t.getBBox()), 

408 **imargs) 

409 plt.gray() 

410 plt.xticks([]) 

411 plt.yticks([]) 

412 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

413 plt.axis(pext) 

414 savefig(pid, 'r%i' % (kk)) 

415 

416 # Median-filtered template 

417 plt.clf() 

418 t = pkres.median_filtered_template 

419 myimshow(t.getArray(), extent=getExtent(t.getBBox()), 

420 **imargs) 

421 plt.gray() 

422 plt.xticks([]) 

423 plt.yticks([]) 

424 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

425 plt.axis(pext) 

426 savefig(pid, 'med%i' % (kk)) 

427 

428 # Assigned flux 

429 plt.clf() 

430 t = pkres.portion_mimg.getImage() 

431 myimshow(t.getArray(), extent=getExtent(t.getBBox()), 

432 **imargs) 

433 plt.gray() 

434 plt.xticks([]) 

435 plt.yticks([]) 

436 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

437 plt.axis(pext) 

438 savefig(pid, 'p%i' % (kk)) 

439 

440 if opt.sec == 'ramp2': 

441 # stray flux 

442 if pkres.stray_flux is not None: 

443 s = pkres.stray_flux 

444 strayim = footprintToImage(s).getArray() 

445 strayext = getExtent(s.getBBox()) 

446 

447 plt.clf() 

448 myimshow(strayim, extent=strayext, **imargs) 

449 plt.gray() 

450 plt.xticks([]) 

451 plt.yticks([]) 

452 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

453 plt.axis(pext) 

454 savefig(pid, 's%i' % (kk)) 

455 

456 # Assigned flux, omitting stray flux. 

457 plt.clf() 

458 myimshow(him2, extent=hext2, **imargs) 

459 plt.gray() 

460 plt.xticks([]) 

461 plt.yticks([]) 

462 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

463 plt.axis(pext) 

464 savefig(pid, 'hb%i' % (kk)) 

465 

466 k += 1 

467 

468 # sum of templates 

469 plt.clf() 

470 myimshow(tsum, extent=getExtent(tbb), **imargs) 

471 plt.gray() 

472 plt.xticks([]) 

473 plt.yticks([]) 

474 plt.plot([pk.getIx() for pk in pks], [pk.getIy() for pk in pks], **pksty) 

475 plt.axis(pext) 

476 savefig(pid, 'tsum') 

477 

478 # sum of assigned flux 

479 plt.clf() 

480 myimshow(hsum, extent=getExtent(tbb), **imargs) 

481 plt.gray() 

482 plt.xticks([]) 

483 plt.yticks([]) 

484 plt.plot([pk.getIx() for pk in pks], [pk.getIy() for pk in pks], **pksty) 

485 plt.axis(pext) 

486 savefig(pid, 'hsum') 

487 

488 plt.clf() 

489 myimshow(hsum2, extent=getExtent(tbb), **imargs) 

490 plt.gray() 

491 plt.xticks([]) 

492 plt.yticks([]) 

493 plt.plot([pk.getIx() for pk in pks], [pk.getIy() for pk in pks], **pksty) 

494 plt.axis(pext) 

495 savefig(pid, 'hsum2') 

496 

497 k = 0 

498 for pkres, pk in zip(res.peaks, pks): 

499 heavy = pkres.get_flux_portion() 

500 if heavy is None: 

501 continue 

502 

503 print('Template footprint:', pkres.template_foot.getBBox()) 

504 print('Template img:', pkres.template_mimg.getBBox()) 

505 print('Heavy footprint:', heavy.getBBox()) 

506 

507 cfp = pkres.template_foot 

508 cbb = cfp.getBBox() 

509 cext = getExtent(cbb) 

510 tim = pkres.template_mimg.getImage().getArray() 

511 (x0, x1, y0, y1) = cext 

512 

513 frac = tim / tsum[y0-ty0:y1-ty0, x0-tx0:x1-tx0] 

514 

515 msk = afwImage.ImageF(cbb.getWidth(), cbb.getHeight()) 

516 msk.setXY0(cbb.getMinX(), cbb.getMinY()) 

517 afwDet.setImageFromFootprint(msk, cfp, 1.) 

518 msk = msk.getArray() 

519 frac[msk == 0.] = np.nan 

520 

521 # Fraction of flux assigned to this child. 

522 plt.clf() 

523 plt.imshow(frac, extent=cext, interpolation='nearest', origin='lower', vmin=0, vmax=1) 

524 # plt.plot([x0,x0,x1,x1,x0], [y0,y1,y1,y0,y0], 'k-') 

525 plt.gray() 

526 plt.xticks([]) 

527 plt.yticks([]) 

528 plt.plot([pk.getIx()], [pk.getIy()], **pksty) 

529 plt.gca().set_axis_bgcolor((0.9, 0.9, 0.5)) 

530 plt.axis(pext) 

531 savefig(pid, 'f%i' % (mapchild(k))) 

532 

533 k += 1 

534 

535 

536if __name__ == '__main__': 

537 main()