Coverage for metadetect / lsst / tests / test_lsst_masking.py: 9%

1import numpy as np 

2from metadetect.masking import get_ap_range 

3from metadetect.lsst import masking 

4from metadetect.lsst import vis 

5from metadetect.lsst import util 

6from metadetect.lsst.metadetect import run_metadetect 

7import lsst.geom as geom 

8import pytest 

9from lsst.utils import getPackageDir 

10 

11try: 

12 getPackageDir('descwl_shear_sims') 

13 getPackageDir('descwl_coadd') 

14 skip_tests_on_simulations = False 

15except LookupError: 

16 skip_tests_on_simulations = True 

17 

18 

19def make_lsst_sim( 

20 rng, dim, mag=22, hlr=0.5, bands=['r', 'i', 'z'], layout='grid', 

21): 

22 import descwl_shear_sims 

23 

24 coadd_dim = dim 

25 se_dim = dim 

26 

27 galaxy_catalog = descwl_shear_sims.galaxies.FixedGalaxyCatalog( 

28 rng=rng, 

29 coadd_dim=coadd_dim, 

30 buff=0, 

31 layout=layout, 

32 mag=mag, 

33 hlr=hlr, 

34 ) 

35 

36 psf = descwl_shear_sims.psfs.make_fixed_psf(psf_type='gauss') 

37 

38 sim_data = descwl_shear_sims.make_sim( 

39 rng=rng, 

40 galaxy_catalog=galaxy_catalog, 

41 coadd_dim=coadd_dim, 

42 se_dim=se_dim, 

43 g1=0.00, 

44 g2=0.00, 

45 psf=psf, 

46 bands=bands, 

47 ) 

48 return sim_data 

49 

50 

51def do_coadding(rng, sim_data, nowarp): 

52 from descwl_coadd.coadd import make_coadd 

53 from descwl_coadd.coadd_nowarp import make_coadd_nowarp 

54 

55 bands = list(sim_data['band_data'].keys()) 

56 

57 if nowarp: 

58 coadd_data_list = [ 

59 make_coadd_nowarp( 

60 exp=sim_data['band_data'][band][0], 

61 psf_dims=sim_data['psf_dims'], 

62 rng=rng, 

63 remove_poisson=False, 

64 ) 

65 for band in bands 

66 ] 

67 else: 

68 coadd_data_list = [ 

69 make_coadd( 

70 exps=sim_data['band_data'][band], 

71 psf_dims=sim_data['psf_dims'], 

72 rng=rng, 

73 coadd_wcs=sim_data['coadd_wcs'], 

74 coadd_bbox=sim_data['coadd_bbox'], 

75 remove_poisson=False, 

76 ) 

77 for band in bands 

78 ] 

79 

80 return util.extract_multiband_coadd_data(coadd_data_list) 

81 

82 

83@pytest.mark.skipif( 

84 skip_tests_on_simulations, 

85 reason='descwl_shear_sims and/or descwl_coadd not available' 

86) 

87def test_apply_apodized_bright_masks(show=False): 

88 ntrial = 5 

89 nmasked = 2 

90 

91 seed = 101 

92 dim = 200 

93 rng = np.random.RandomState(seed=seed) 

94 

95 for itrial in range(ntrial): 

96 sim_data = make_lsst_sim(rng=rng, dim=dim) 

97 data = do_coadding(rng=rng, sim_data=sim_data, nowarp=True) 

98 

99 bands = data['mbexp'].bands 

100 

101 dtype = [('ra', 'f8'), ('dec', 'f8'), ('radius_pixels', 'f8')] 

102 

103 wcs = data['mbexp'][bands[0]].getWcs() 

104 

105 buff = 21 

106 

107 exp = data['mbexp'][bands[0]] 

108 dims = exp.image.array.shape 

109 dim = dims[0] 

110 

111 x = rng.uniform(low=buff, high=dim-buff, size=nmasked) 

112 y = rng.uniform(low=buff, high=dim-buff, size=nmasked) 

113 radius = rng.uniform(low=10, high=20, size=nmasked) 

114 ra, dec = wcs.pixelToSkyArray(x=x, y=y, degrees=True) 

115 bright_info = np.zeros(nmasked, dtype=dtype) 

116 bright_info['ra'] = ra 

117 bright_info['dec'] = dec 

118 bright_info['radius_pixels'] = radius 

119 

120 xp, yp = wcs.skyToPixelArray(ra=ra, dec=dec, degrees=True) 

121 assert np.allclose(x, xp) 

122 assert np.allclose(y, yp) 

123 

124 noise_images = [ 

125 nexp.image.array.copy() for nexp in data['noise_mbexp'] 

126 ] 

127 

128 # works in place 

129 masking.apply_apodized_bright_masks_mbexp( 

130 mbexp=data['mbexp'], 

131 noise_mbexp=data['noise_mbexp'], 

132 mfrac_mbexp=data['mfrac_mbexp'], 

133 bright_info=bright_info, 

134 ormasks=data['ormasks'], 

135 ) 

136 

137 if show: 

138 vis.show_multi_mbexp(data['mbexp']) 

139 

140 bright = exp.mask.getPlaneBitMask('BRIGHT') 

141 bright_expanded = exp.mask.getPlaneBitMask('BRIGHT_EXPANDED') 

142 

143 ygrid, xgrid = np.mgrid[0:dim, 0:dim] 

144 for iband, band in enumerate(data['mbexp'].bands): 

145 exp_orig = sim_data['band_data'][band][0] 

146 exp = data['mbexp'][band] 

147 assert exp_orig is not exp 

148 nexp = data['noise_mbexp'][band] 

149 ormask = data['ormasks'][iband] 

150 mfrac = data['mfrac_mbexp'][band] 

151 

152 noise_image = noise_images[iband] 

153 

154 for xx, yy, rr in zip(x, y, radius): 

155 ix = int(np.floor(xx+0.5)) 

156 iy = int(np.floor(yy+0.5)) 

157 assert exp.image.array[iy, ix] == 0 

158 assert nexp.image.array[iy, ix] == 0 

159 assert exp.mask.array[iy, ix] == (bright | bright_expanded) 

160 assert nexp.mask.array[iy, ix] == (bright | bright_expanded) 

161 

162 # within mask the bit is set and data are modified 

163 r2 = (ygrid - yy)**2 + (xgrid - xx)**2 

164 

165 # shrink radius a little to deal with some roundoff errors 

166 w = np.where(r2 < (rr-0.1)**2) 

167 assert np.all(exp.mask.array[w] & bright != 0) 

168 assert np.all(nexp.mask.array[w] & bright != 0) 

169 assert np.all(ormask[w] & bright != 0) 

170 

171 assert np.all(exp.variance.array[w] == np.inf) 

172 assert np.all(nexp.variance.array[w] == np.inf) 

173 assert np.all(mfrac.image.array[w] == 1) 

174 

175 assert np.all(exp.image.array[w] != exp_orig.image.array[w]) 

176 assert np.all(nexp.image.array[w] != noise_image[w]) 

177 

178 # mask is expanded and bit is set, but data are not modified in 

179 # the expanded area 

180 w = np.where(r2 < (rr + masking.EXPAND_RAD)**2) 

181 assert np.all(exp.mask.array[w] & bright_expanded != 0) 

182 assert np.all(nexp.mask.array[w] & bright_expanded != 0) 

183 assert np.all(ormask[w] & bright_expanded != 0) 

184 

185 w = np.where( 

186 (exp.mask.array & bright == 0) & 

187 (exp.mask.array & bright_expanded != 0) 

188 ) 

189 assert np.all(exp.image.array[w] == exp_orig.image.array[w]) 

190 assert np.all(nexp.image.array[w] == noise_image[w]) 

191 assert np.all(exp.variance.array[w] != np.inf) 

192 assert np.all(nexp.variance.array[w] != np.inf) 

193 

194 

195def extract_cell_mbexp(mbexp, cell_size, start_x, start_y): 

196 from metadetect.lsst.util import get_mbexp, copy_mbexp 

197 

198 bbox_begin = mbexp.getBBox().getBegin() 

199 

200 new_begin = geom.Point2I( 

201 x=bbox_begin.getX() + start_x, 

202 y=bbox_begin.getY() + start_y, 

203 ) 

204 extent = geom.Extent2I(cell_size) 

205 new_bbox = geom.Box2I( 

206 new_begin, 

207 extent, 

208 ) 

209 

210 subexps = [] 

211 for band in mbexp.bands: 

212 exp = mbexp[band] 

213 # we need to make a copy of it 

214 subexp = exp[new_bbox] 

215 

216 assert np.all( 

217 exp.image.array[ 

218 start_y:start_y+cell_size, 

219 start_x:start_x+cell_size 

220 ] == subexp.image.array[:, :] 

221 ) 

222 

223 subexps.append(subexp) 

224 

225 return copy_mbexp(get_mbexp(subexps)) 

226 

227 

228@pytest.mark.skipif( 

229 skip_tests_on_simulations, 

230 reason='descwl_shear_sims and/or descwl_coadd not available' 

231) 

232def test_apply_apodized_bright_masks_subexp(show=False): 

233 """ 

234 test that this works for sub exposures 

235 """ 

236 seed = 101 

237 dim = 200 

238 rng = np.random.RandomState(seed=seed) 

239 

240 sim_data = make_lsst_sim(rng=rng, dim=dim) 

241 data = do_coadding(rng=rng, sim_data=sim_data, nowarp=True) 

242 

243 bands = data['mbexp'].bands 

244 

245 dtype = [('ra', 'f8'), ('dec', 'f8'), ('radius_pixels', 'f8')] 

246 

247 wcs = data['mbexp'][bands[0]].getWcs() 

248 

249 exp = data['mbexp'][bands[0]] 

250 xy0 = exp.getXY0() 

251 

252 # in bigger coordinate system 

253 x = xy0.x + 100 

254 y = xy0.y + 100 

255 radius = 20 

256 

257 spt = wcs.pixelToSky(x, y) 

258 ra = spt.getRa().asDegrees() 

259 dec = spt.getDec().asDegrees() 

260 

261 bright_info = np.zeros(1, dtype=dtype) 

262 bright_info['ra'] = ra 

263 bright_info['dec'] = dec 

264 bright_info['radius_pixels'] = radius 

265 

266 xp, yp = wcs.skyToPixelArray(ra=ra, dec=dec, degrees=True) 

267 assert np.allclose(x, xp) 

268 assert np.allclose(y, yp) 

269 

270 if show: 

271 vis.show_multi_mbexp(data['mbexp']) 

272 

273 cell_size = 50 

274 start_x = x - xy0.x - 25 

275 start_y = y - xy0.y - 25 

276 sub_mbexp = extract_cell_mbexp( 

277 mbexp=data['mbexp'], 

278 cell_size=cell_size, 

279 start_x=start_x, 

280 start_y=start_y, 

281 ) 

282 nsub_mbexp = extract_cell_mbexp( 

283 mbexp=data['noise_mbexp'], 

284 cell_size=cell_size, 

285 start_x=start_x, 

286 start_y=start_y, 

287 ) 

288 mfsub_mbexp = extract_cell_mbexp( 

289 mbexp=data['mfrac_mbexp'], 

290 cell_size=cell_size, 

291 start_x=start_x, 

292 start_y=start_y, 

293 ) 

294 ormasks = [ 

295 ormask[ 

296 start_y:start_y+cell_size, 

297 start_x:start_x+cell_size, 

298 ] for ormask in data['ormasks'] 

299 ] 

300 

301 masking.apply_apodized_bright_masks_mbexp( 

302 mbexp=sub_mbexp, 

303 noise_mbexp=nsub_mbexp, 

304 mfrac_mbexp=mfsub_mbexp, 

305 bright_info=bright_info, 

306 ormasks=ormasks, 

307 ) 

308 

309 bright = exp.mask.getPlaneBitMask('BRIGHT') 

310 bright_expanded = exp.mask.getPlaneBitMask('BRIGHT_EXPANDED') 

311 

312 ygrid, xgrid = np.mgrid[0:cell_size, 0:cell_size] 

313 

314 for iband, band in enumerate(data['mbexp'].bands): 

315 exp = sub_mbexp[band] 

316 nexp = nsub_mbexp[band] 

317 ormask = ormasks[iband] 

318 mfrac = mfsub_mbexp[band] 

319 

320 txy0 = exp.getXY0() 

321 

322 ix = int(np.floor(x+0.5)) - txy0.x 

323 iy = int(np.floor(y+0.5)) - txy0.y 

324 

325 assert exp.image.array[iy, ix] == 0 

326 assert nexp.image.array[iy, ix] == 0 

327 assert exp.mask.array[iy, ix] == (bright | bright_expanded) 

328 assert nexp.mask.array[iy, ix] == (bright | bright_expanded) 

329 

330 # within mask the bit is set and data are modified 

331 r2 = (ygrid - iy)**2 + (xgrid - ix)**2 

332 

333 # shrink radius a little to deal with some roundoff errors 

334 w = np.where(r2 < (radius-0.1)**2) 

335 

336 assert np.all(exp.mask.array[w] & bright != 0) 

337 assert np.all(nexp.mask.array[w] & bright != 0) 

338 assert np.all(ormask[w] & bright != 0) 

339 

340 assert np.all(exp.variance.array[w] == np.inf) 

341 assert np.all(nexp.variance.array[w] == np.inf) 

342 assert np.all(mfrac.image.array[w] == 1) 

343 

344 # mask is expanded and bit is set, but data are not modified in 

345 # the expanded area 

346 w = np.where(r2 < (radius + masking.EXPAND_RAD)**2) 

347 assert np.all(exp.mask.array[w] & bright_expanded != 0) 

348 assert np.all(nexp.mask.array[w] & bright_expanded != 0) 

349 assert np.all(ormask[w] & bright_expanded != 0) 

350 

351 w = np.where( 

352 (exp.mask.array & bright == 0) & 

353 (exp.mask.array & bright_expanded != 0) 

354 ) 

355 assert np.all(exp.variance.array[w] != np.inf) 

356 assert np.all(nexp.variance.array[w] != np.inf) 

357 

358 

359@pytest.mark.skipif( 

360 skip_tests_on_simulations, 

361 reason='descwl_shear_sims and/or descwl_coadd not available' 

362) 

363def test_apply_apodized_edge_masks(show=False): 

364 seed = 101 

365 dim = 250 

366 rng = np.random.RandomState(seed=seed) 

367 

368 sim_data = make_lsst_sim(rng=rng, dim=dim) 

369 data = do_coadding(rng=rng, sim_data=sim_data, nowarp=True) 

370 

371 noise_images = [ 

372 nexp.image.array.copy() for nexp in data['noise_mbexp'] 

373 ] 

374 

375 # works in place 

376 masking.apply_apodized_edge_masks_mbexp( 

377 mbexp=data['mbexp'], 

378 noise_mbexp=data['noise_mbexp'], 

379 mfrac_mbexp=data['mfrac_mbexp'], 

380 ormasks=data['ormasks'], 

381 ) 

382 

383 if show: 

384 vis.show_multi_mbexp(data['mbexp']) 

385 

386 band0 = data['mbexp'].bands[0] 

387 edge = data['mbexp'][band0].mask.getPlaneBitMask('APODIZED_EDGE') 

388 

389 ygrid, xgrid = np.mgrid[0:dim, 0:dim] 

390 

391 ap_range = get_ap_range(masking.AP_RAD) 

392 w = np.where( 

393 (xgrid < ap_range) | 

394 (ygrid < ap_range) | 

395 (xgrid > (dim - ap_range - 1)) | 

396 (ygrid > (dim - ap_range - 1)) 

397 ) 

398 for iband, band in enumerate(data['mbexp'].bands): 

399 exp_orig = sim_data['band_data'][band][0] 

400 exp = data['mbexp'][band] 

401 assert exp_orig is not exp 

402 nexp = data['noise_mbexp'][band] 

403 mfrac = data['mfrac_mbexp'][band] 

404 ormask = data['ormasks'][iband] 

405 

406 noise_image = noise_images[iband] 

407 

408 assert np.all(exp.mask.array[w] & edge != 0) 

409 assert np.all(nexp.mask.array[w] & edge != 0) 

410 assert np.all(ormask[w] & edge != 0) 

411 

412 assert np.all(exp.variance.array[w] == np.inf) 

413 assert np.all(nexp.variance.array[w] == np.inf) 

414 assert np.all(mfrac.image.array[w] == 1) 

415 

416 assert np.all(exp.image.array[w] != exp_orig.image.array[w]) 

417 assert np.all(nexp.image.array[w] != noise_image[w]) 

418 

419 

420@pytest.mark.skipif( 

421 skip_tests_on_simulations, 

422 reason='descwl_shear_sims and/or descwl_coadd not available' 

423) 

424def test_apply_apodized_bright_masks_metadetect(show=False): 

425 """ 

426 look for detected objects with expanded mask set 

427 """ 

428 ntrial = 5 

429 nmasked = 2 

430 

431 seed = 101 

432 dim = 200 

433 rng = np.random.RandomState(seed=seed) 

434 

435 found = False 

436 for itrial in range(ntrial): 

437 sim_data = make_lsst_sim( 

438 rng=rng, dim=dim, layout='random', bands=['i'], 

439 ) 

440 

441 data = do_coadding(rng=rng, sim_data=sim_data, nowarp=True) 

442 

443 bands = data['mbexp'].bands 

444 

445 dtype = [('ra', 'f8'), ('dec', 'f8'), ('radius_pixels', 'f8')] 

446 

447 wcs = data['mbexp'][bands[0]].getWcs() 

448 

449 buff = 21 

450 

451 exp = data['mbexp'][bands[0]] 

452 dims = exp.image.array.shape 

453 dim = dims[0] 

454 

455 x = rng.uniform(low=buff, high=dim-buff, size=nmasked) 

456 y = rng.uniform(low=buff, high=dim-buff, size=nmasked) 

457 radius = rng.uniform(low=10, high=20, size=nmasked) 

458 ra, dec = wcs.pixelToSkyArray(x=x, y=y, degrees=True) 

459 bright_info = np.zeros(nmasked, dtype=dtype) 

460 bright_info['ra'] = ra 

461 bright_info['dec'] = dec 

462 bright_info['radius_pixels'] = radius 

463 

464 # works in place 

465 masking.apply_apodized_bright_masks_mbexp( 

466 mbexp=data['mbexp'], 

467 noise_mbexp=data['noise_mbexp'], 

468 mfrac_mbexp=data['mfrac_mbexp'], 

469 bright_info=bright_info, 

470 ormasks=data['ormasks'], 

471 ) 

472 

473 res = run_metadetect(config=None, rng=rng, **data) 

474 

475 if show: 

476 vis.show_multi_mbexp(data['mbexp'], sources=res['noshear']) 

477 

478 # bright = exp.mask.getPlaneBitMask('BRIGHT') 

479 bright_expanded = exp.mask.getPlaneBitMask('BRIGHT_EXPANDED') 

480 

481 if np.any(res['noshear']['bmask'] & bright_expanded != 0): 

482 found = True 

483 break 

484 

485 assert found 

486 

487 

488if __name__ == '__main__': 

489 # test_apply_apodized_bright_masks(show=True) 

490 test_apply_apodized_bright_masks_subexp(show=True) 

491 # test_apply_apodized_edge_masks(show=True) 

492 # test_apply_apodized_bright_masks_metadetect(show=True)