Coverage for shear_meas_test / test_shear_meas.py: 8%

255 statements  

« prev     ^ index     » next       coverage.py v7.13.2, created at 2026-02-08 17:04 -0800

1import time 

2import copy 

3import numpy as np 

4import ngmix 

5import galsim 

6import metadetect 

7from esutil.pbar import PBar 

8import joblib 

9 

10import pytest 

11 

12 

13TEST_METADETECT_CONFIG = { 

14 "model": "wmom", 

15 

16 'weight': { 

17 'fwhm': 1.2, # arcsec 

18 }, 

19 

20 'metacal': { 

21 'psf': 'fitgauss', 

22 'types': ['noshear', '1p', '1m', '2p', '2m'], 

23 }, 

24 

25 'sx': { 

26 # in sky sigma 

27 # DETECT_THRESH 

28 'detect_thresh': 0.8, 

29 

30 # Minimum contrast parameter for deblending 

31 # DEBLEND_MINCONT 

32 'deblend_cont': 0.00001, 

33 

34 # minimum number of pixels above threshold 

35 # DETECT_MINAREA: 6 

36 'minarea': 4, 

37 

38 'filter_type': 'conv', 

39 

40 # 7x7 convolution mask of a gaussian PSF with FWHM = 3.0 pixels. 

41 'filter_kernel': [ 

42 [0.004963, 0.021388, 0.051328, 0.068707, 0.051328, 0.021388, 0.004963], # noqa 

43 [0.021388, 0.092163, 0.221178, 0.296069, 0.221178, 0.092163, 0.021388], # noqa 

44 [0.051328, 0.221178, 0.530797, 0.710525, 0.530797, 0.221178, 0.051328], # noqa 

45 [0.068707, 0.296069, 0.710525, 0.951108, 0.710525, 0.296069, 0.068707], # noqa 

46 [0.051328, 0.221178, 0.530797, 0.710525, 0.530797, 0.221178, 0.051328], # noqa 

47 [0.021388, 0.092163, 0.221178, 0.296069, 0.221178, 0.092163, 0.021388], # noqa 

48 [0.004963, 0.021388, 0.051328, 0.068707, 0.051328, 0.021388, 0.004963], # noqa 

49 ] 

50 }, 

51 

52 'meds': { 

53 'min_box_size': 32, 

54 'max_box_size': 32, 

55 

56 'box_type': 'iso_radius', 

57 

58 'rad_min': 4, 

59 'rad_fac': 2, 

60 'box_padding': 2, 

61 }, 

62 

63 # check for an edge hit 

64 'bmask_flags': 2**30, 

65 

66 'nodet_flags': 2**0, 

67} 

68 

69 

70def make_sim( 

71 *, 

72 seed, 

73 g1, 

74 g2, 

75 dim=251, 

76 buff=34, 

77 scale=0.25, 

78 dens=100, 

79 ngrid=7, 

80 snr=1e6, 

81): 

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

83 

84 half_loc = (dim-buff*2)*scale/2 

85 

86 if ngrid is None: 

87 area_arcmin2 = ((dim - buff*2)*scale/60)**2 

88 nobj = int(dens * area_arcmin2) 

89 x = rng.uniform(low=-half_loc, high=half_loc, size=nobj) 

90 y = rng.uniform(low=-half_loc, high=half_loc, size=nobj) 

91 else: 

92 half_ngrid = (ngrid-1)/2 

93 x, y = np.meshgrid(np.arange(ngrid), np.arange(ngrid)) 

94 x = (x.ravel() - half_ngrid)/half_ngrid * half_loc 

95 y = (y.ravel() - half_ngrid)/half_ngrid * half_loc 

96 nobj = x.shape[0] 

97 

98 cen = (dim-1)/2 

99 psf_dim = 53 

100 psf_cen = (psf_dim-1)/2 

101 

102 psf = galsim.Gaussian(fwhm=0.9) 

103 gals = [] 

104 for ind in range(nobj): 

105 u, v = rng.uniform(low=-scale, high=scale, size=2) 

106 u += x[ind] 

107 v += y[ind] 

108 gals.append(galsim.Exponential(half_light_radius=0.5).shift(u, v)) 

109 gals = galsim.Add(gals) 

110 gals = gals.shear(g1=g1, g2=g2) 

111 gals = galsim.Convolve([gals, psf]) 

112 

113 im = gals.drawImage(nx=dim, ny=dim, scale=scale).array 

114 psf_im = psf.drawImage(nx=psf_dim, ny=psf_dim, scale=scale).array 

115 

116 nse = ( 

117 np.sqrt(np.sum( 

118 galsim.Convolve([ 

119 psf, 

120 galsim.Exponential(half_light_radius=0.5), 

121 ]).drawImage(scale=0.25).array**2) 

122 ) 

123 / snr 

124 ) 

125 

126 im += rng.normal(size=im.shape, scale=nse) 

127 wgt = np.ones_like(im) / nse**2 

128 jac = ngmix.DiagonalJacobian(scale=scale, row=cen, col=cen) 

129 psf_jac = ngmix.DiagonalJacobian(scale=scale, row=psf_cen, col=psf_cen) 

130 

131 obs = ngmix.Observation( 

132 image=im, 

133 weight=wgt, 

134 jacobian=jac, 

135 ormask=np.zeros_like(im, dtype=np.int32), 

136 bmask=np.zeros_like(im, dtype=np.int32), 

137 psf=ngmix.Observation( 

138 image=psf_im, 

139 jacobian=psf_jac, 

140 ), 

141 ) 

142 mbobs = ngmix.MultiBandObsList() 

143 obslist = ngmix.ObsList() 

144 obslist.append(obs) 

145 mbobs.append(obslist) 

146 return mbobs 

147 

148 

149def make_sim_color( 

150 *, 

151 seed, 

152 g1, 

153 g2, 

154 dim=251, 

155 buff=34, 

156 scale=0.25, 

157 dens=100, 

158 ngrid=7, 

159 snr=1e6, 

160): 

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

162 

163 half_loc = (dim-buff*2)*scale/2 

164 

165 if ngrid is None: 

166 area_arcmin2 = ((dim - buff*2)*scale/60)**2 

167 nobj = int(dens * area_arcmin2) 

168 x = rng.uniform(low=-half_loc, high=half_loc, size=nobj) 

169 y = rng.uniform(low=-half_loc, high=half_loc, size=nobj) 

170 else: 

171 half_ngrid = (ngrid-1)/2 

172 x, y = np.meshgrid(np.arange(ngrid), np.arange(ngrid)) 

173 x = (x.ravel() - half_ngrid)/half_ngrid * half_loc 

174 y = (y.ravel() - half_ngrid)/half_ngrid * half_loc 

175 nobj = x.shape[0] 

176 

177 cen = (dim-1)/2 

178 psf_dim = 53 

179 psf_cen = (psf_dim-1)/2 

180 us = rng.uniform(low=-scale, high=scale, size=nobj) 

181 vs = rng.uniform(low=-scale, high=scale, size=nobj) 

182 colors = rng.randint(low=0, high=2, size=nobj) 

183 

184 mbobs = ngmix.MultiBandObsList() 

185 base_psf = galsim.Gaussian(fwhm=0.9) 

186 # this dict holds the PSFs as a function of color (c0 or c1) and band (index 0 or 1) 

187 psfs = { 

188 "c0": [galsim.Gaussian(fwhm=0.83), galsim.Gaussian(fwhm=0.96)], 

189 "c1": [galsim.Gaussian(fwhm=0.88), galsim.Gaussian(fwhm=1.03)], 

190 } 

191 for band in range(2): 

192 # we render all objects with the same color at once and so loop through the 

193 # colors here 

194 for color in range(2): 

195 gals = [] 

196 psf = psfs["c%d" % color][band] 

197 for ind in range(nobj): 

198 # we only render objects at this color 

199 if colors[ind] != color: 

200 continue 

201 

202 # the object should be brighter in one band than the other 

203 # brighter in band 0 is c0 and thew opposite is c1 

204 if colors[ind] == band: 

205 flux = 0.8 

206 else: 

207 flux = 0.2 

208 

209 u = us[ind] + x[ind] 

210 v = vs[ind] + y[ind] 

211 gals.append( 

212 flux * galsim.Exponential(half_light_radius=0.5).shift(u, v) 

213 ) 

214 gals = galsim.Add(gals) 

215 gals = gals.shear(g1=g1, g2=g2) 

216 gals = galsim.Convolve([gals, psf]) 

217 

218 if color == 0: 

219 im = gals.drawImage(nx=dim, ny=dim, scale=scale).array 

220 else: 

221 im += gals.drawImage(nx=dim, ny=dim, scale=scale).array 

222 

223 psf_im = base_psf.drawImage(nx=psf_dim, ny=psf_dim, scale=scale).array 

224 

225 nse = ( 

226 np.sqrt(np.sum( 

227 galsim.Convolve([ 

228 base_psf, 

229 galsim.Exponential(half_light_radius=0.5), 

230 ]).drawImage(scale=0.25).array**2) 

231 ) 

232 / snr 

233 ) 

234 

235 im += rng.normal(size=im.shape, scale=nse) 

236 wgt = np.ones_like(im) / nse**2 

237 jac = ngmix.DiagonalJacobian(scale=scale, row=cen, col=cen) 

238 psf_jac = ngmix.DiagonalJacobian(scale=scale, row=psf_cen, col=psf_cen) 

239 

240 obs = ngmix.Observation( 

241 image=im, 

242 weight=wgt, 

243 jacobian=jac, 

244 ormask=np.zeros_like(im, dtype=np.int32), 

245 bmask=np.zeros_like(im, dtype=np.int32), 

246 psf=ngmix.Observation( 

247 image=psf_im, 

248 jacobian=psf_jac, 

249 ), 

250 ) 

251 obslist = ngmix.ObsList() 

252 obslist.append(obs) 

253 mbobs.append(obslist) 

254 

255 # these are the observations with the correct PSFs for a given color 

256 color_dep_mbobs = { 

257 "c0": mbobs.copy(), 

258 "c1": mbobs.copy(), 

259 } 

260 for color in range(2): 

261 for band in range(2): 

262 psf = psfs["c%d" % color][band] 

263 psf_im = psf.drawImage(nx=psf_dim, ny=psf_dim, scale=scale).array 

264 psf_jac = ngmix.DiagonalJacobian(scale=scale, row=psf_cen, col=psf_cen) 

265 

266 psf = ngmix.Observation( 

267 image=psf_im, 

268 jacobian=psf_jac, 

269 ) 

270 color_dep_mbobs["c%d" % color][band][0].psf = psf 

271 

272 return mbobs, color_dep_mbobs 

273 

274 

275def _shear_cuts(arr, model): 

276 if model == "wmom": 

277 tmin = 1.2 

278 else: 

279 tmin = 0.5 

280 msk = ( 

281 (arr[f'{model}_flags'] == 0) 

282 & (arr[f'{model}_s2n'] > 10) 

283 & (arr[f'{model}_T_ratio'] > tmin) 

284 ) 

285 return msk 

286 

287 

288def _meas_shear_data(res, model): 

289 msk = _shear_cuts(res['noshear'], model) 

290 g1 = np.mean(res['noshear'][f'{model}_g'][msk, 0]) 

291 g2 = np.mean(res['noshear'][f'{model}_g'][msk, 1]) 

292 

293 msk = _shear_cuts(res['1p'], model) 

294 g1_1p = np.mean(res['1p'][f'{model}_g'][msk, 0]) 

295 msk = _shear_cuts(res['1m'], model) 

296 g1_1m = np.mean(res['1m'][f'{model}_g'][msk, 0]) 

297 R11 = (g1_1p - g1_1m) / 0.02 

298 

299 msk = _shear_cuts(res['2p'], model) 

300 g2_2p = np.mean(res['2p'][f'{model}_g'][msk, 1]) 

301 msk = _shear_cuts(res['2m'], model) 

302 g2_2m = np.mean(res['2m'][f'{model}_g'][msk, 1]) 

303 R22 = (g2_2p - g2_2m) / 0.02 

304 

305 dt = [ 

306 ('g1', 'f8'), 

307 ('g2', 'f8'), 

308 ('R11', 'f8'), 

309 ('R22', 'f8')] 

310 return np.array([(g1, g2, R11, R22)], dtype=dt) 

311 

312 

313def _bootstrap_stat(d1, d2, func, seed, nboot=500): 

314 dim = d1.shape[0] 

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

316 stats = [] 

317 for _ in range(nboot): 

318 ind = rng.choice(dim, size=dim, replace=True) 

319 stats.append(func(d1[ind], d2[ind])) 

320 return stats 

321 

322 

323def meas_m_c_cancel(pres, mres): 

324 x = np.mean(pres['g1'] - mres['g1'])/2 

325 y = np.mean(pres['R11'] + mres['R11'])/2 

326 m = x/y/0.02 - 1 

327 

328 x = np.mean(pres['g2'] + mres['g2'])/2 

329 y = np.mean(pres['R22'] + mres['R22'])/2 

330 c = x/y 

331 

332 return m, c 

333 

334 

335def boostrap_m_c(pres, mres): 

336 m, c = meas_m_c_cancel(pres, mres) 

337 bdata = _bootstrap_stat(pres, mres, meas_m_c_cancel, 14324, nboot=500) 

338 merr, cerr = np.std(bdata, axis=0) 

339 return m, merr, c, cerr 

340 

341 

342def run_sim(seed, mdet_seed, model, **kwargs): 

343 mbobs_p = make_sim(seed=seed, g1=0.02, g2=0.0, **kwargs) 

344 cfg = copy.deepcopy(TEST_METADETECT_CONFIG) 

345 cfg["model"] = model 

346 _pres = metadetect.do_metadetect( 

347 copy.deepcopy(cfg), 

348 mbobs_p, 

349 np.random.RandomState(seed=mdet_seed) 

350 ) 

351 if _pres is None: 

352 return None 

353 

354 mbobs_m = make_sim(seed=seed, g1=-0.02, g2=0.0, **kwargs) 

355 _mres = metadetect.do_metadetect( 

356 copy.deepcopy(cfg), 

357 mbobs_m, 

358 np.random.RandomState(seed=mdet_seed) 

359 ) 

360 if _mres is None: 

361 return None 

362 

363 return _meas_shear_data(_pres, model), _meas_shear_data(_mres, model) 

364 

365 

366def run_sim_color(seed, mdet_seed, model, **kwargs): 

367 # if an object is bright in band 0 it is c0, otherwise it is c1 

368 def _color_key_func(fluxes): 

369 if fluxes[0] > fluxes[1]: 

370 return "c0" 

371 else: 

372 return "c1" 

373 

374 mbobs_p, color_dep_mbobs_p = make_sim_color(seed=seed, g1=0.02, g2=0.0, **kwargs) 

375 cfg = copy.deepcopy(TEST_METADETECT_CONFIG) 

376 cfg["model"] = model 

377 _pres = metadetect.do_metadetect( 

378 copy.deepcopy(cfg), 

379 mbobs_p, 

380 np.random.RandomState(seed=mdet_seed), 

381 color_dep_mbobs=color_dep_mbobs_p, 

382 color_key_func=_color_key_func, 

383 ) 

384 if _pres is None: 

385 return None 

386 

387 mbobs_m, color_dep_mbobs_m = make_sim_color(seed=seed, g1=-0.02, g2=0.0, **kwargs) 

388 _mres = metadetect.do_metadetect( 

389 copy.deepcopy(cfg), 

390 mbobs_m, 

391 np.random.RandomState(seed=mdet_seed), 

392 color_dep_mbobs=color_dep_mbobs_m, 

393 color_key_func=_color_key_func, 

394 ) 

395 if _mres is None: 

396 return None 

397 

398 return _meas_shear_data(_pres, model), _meas_shear_data(_mres, model) 

399 

400 

401@pytest.mark.parametrize( 

402 'model,snr,ngrid,ntrial', [ 

403 ("wmom", 1e6, 7, 64), 

404 ("pgauss", 1e6, 7, 64), 

405 ] 

406) 

407def test_shear_meas_color(model, snr, ngrid, ntrial): 

408 pytest.importorskip("sxdes") 

409 

410 nsub = max(ntrial // 100, 8) 

411 nitr = ntrial // nsub 

412 rng = np.random.RandomState(seed=116) 

413 seeds = rng.randint(low=1, high=2**29, size=ntrial) 

414 mdet_seeds = rng.randint(low=1, high=2**29, size=ntrial) 

415 

416 tm0 = time.time() 

417 

418 print("") 

419 

420 pres = [] 

421 mres = [] 

422 loc = 0 

423 with joblib.Parallel(n_jobs=-1, verbose=100, backend='loky') as par: 

424 for itr in PBar(range(nitr)): 

425 jobs = [ 

426 joblib.delayed(run_sim_color)( 

427 seeds[loc+i], mdet_seeds[loc+i], model, snr=snr, ngrid=ngrid, 

428 ) 

429 for i in range(nsub) 

430 ] 

431 print("\n", end="", flush=True) 

432 outputs = par(jobs) 

433 

434 for out in outputs: 

435 if out is None: 

436 continue 

437 pres.append(out[0]) 

438 mres.append(out[1]) 

439 loc += nsub 

440 

441 m, merr, c, cerr = boostrap_m_c( 

442 np.concatenate(pres), 

443 np.concatenate(mres), 

444 ) 

445 print( 

446 ( 

447 "\n" 

448 "nsims: %d\n" 

449 "m [1e-3, 3sigma]: %s +/- %s\n" 

450 "c [1e-5, 3sigma]: %s +/- %s\n" 

451 "\n" 

452 ) % ( 

453 len(pres), 

454 m/1e-3, 

455 3*merr/1e-3, 

456 c/1e-5, 

457 3*cerr/1e-5, 

458 ), 

459 flush=True, 

460 ) 

461 

462 total_time = time.time()-tm0 

463 print("time per:", total_time/ntrial, flush=True) 

464 

465 pres = np.concatenate(pres) 

466 mres = np.concatenate(mres) 

467 m, merr, c, cerr = boostrap_m_c(pres, mres) 

468 

469 print( 

470 ( 

471 "m [1e-3, 3sigma]: %s +/- %s" 

472 "\nc [1e-5, 3sigma]: %s +/- %s" 

473 ) % ( 

474 m/1e-3, 

475 3*merr/1e-3, 

476 c/1e-5, 

477 3*cerr/1e-5, 

478 ), 

479 flush=True, 

480 ) 

481 

482 assert np.abs(m) < max(1e-3, 3*merr) 

483 assert np.abs(c) < 3*cerr 

484 

485 

486@pytest.mark.parametrize( 

487 'model,snr,ngrid,ntrial', [ 

488 ("wmom", 1e6, 7, 64), 

489 ("ksigma", 1e6, 7, 64), 

490 ("pgauss", 1e6, 7, 64), 

491 ("am", 1e6, 7, 64), 

492 ("gauss", 1e6, 7, 64), 

493 # this test takes ~3 hours on github actions so we only run it once and not 

494 # with color 

495 ("wmom", 1e6, None, 9500), 

496 ] 

497) 

498def test_shear_meas_simple(model, snr, ngrid, ntrial): 

499 pytest.importorskip("sxdes") 

500 

501 nsub = max(ntrial // 128, 8) 

502 nitr = ntrial // nsub 

503 rng = np.random.RandomState(seed=116) 

504 seeds = rng.randint(low=1, high=2**29, size=ntrial) 

505 mdet_seeds = rng.randint(low=1, high=2**29, size=ntrial) 

506 

507 tm0 = time.time() 

508 

509 print("") 

510 

511 pres = [] 

512 mres = [] 

513 loc = 0 

514 with joblib.Parallel(n_jobs=-1, verbose=100, backend='loky') as par: 

515 for itr in PBar(range(nitr)): 

516 jobs = [ 

517 joblib.delayed(run_sim)( 

518 seeds[loc+i], mdet_seeds[loc+i], model, snr=snr, ngrid=ngrid, 

519 ) 

520 for i in range(nsub) 

521 ] 

522 print("\n", end="", flush=True) 

523 outputs = par(jobs) 

524 

525 for out in outputs: 

526 if out is None: 

527 continue 

528 pres.append(out[0]) 

529 mres.append(out[1]) 

530 loc += nsub 

531 

532 m, merr, c, cerr = boostrap_m_c( 

533 np.concatenate(pres), 

534 np.concatenate(mres), 

535 ) 

536 print( 

537 ( 

538 "\n" 

539 "nsims: %d\n" 

540 "m [1e-3, 3sigma]: %s +/- %s\n" 

541 "c [1e-5, 3sigma]: %s +/- %s\n" 

542 "\n" 

543 ) % ( 

544 len(pres), 

545 m/1e-3, 

546 3*merr/1e-3, 

547 c/1e-5, 

548 3*cerr/1e-5, 

549 ), 

550 flush=True, 

551 ) 

552 

553 total_time = time.time()-tm0 

554 print("time per:", total_time/ntrial, flush=True) 

555 

556 pres = np.concatenate(pres) 

557 mres = np.concatenate(mres) 

558 m, merr, c, cerr = boostrap_m_c(pres, mres) 

559 

560 print( 

561 ( 

562 "m [1e-3, 3sigma]: %s +/- %s" 

563 "\nc [1e-5, 3sigma]: %s +/- %s" 

564 ) % ( 

565 m/1e-3, 

566 3*merr/1e-3, 

567 c/1e-5, 

568 3*cerr/1e-5, 

569 ), 

570 flush=True, 

571 ) 

572 

573 assert np.abs(m) < max(1e-3, 3*merr) 

574 assert np.abs(c) < 3*cerr 

575 

576 

577@pytest.mark.parametrize( 

578 'model,snr,ngrid', [ 

579 ("wmom", 1e6, 7), 

580 ("ksigma", 1e6, 7), 

581 ("pgauss", 1e6, 7), 

582 ] 

583) 

584def test_shear_meas_timing(model, snr, ngrid): 

585 pytest.importorskip("sxdes") 

586 

587 rng = np.random.RandomState(seed=116) 

588 seeds = rng.randint(low=1, high=2**29, size=1) 

589 mdet_seeds = rng.randint(low=1, high=2**29, size=1) 

590 

591 run_sim( 

592 seeds[0], mdet_seeds[0], model, snr=snr, ngrid=ngrid, 

593 )