Coverage for metadetect / masking.py: 8%

169 statements  

« prev     ^ index     » next       coverage.py v7.13.2, created at 2026-02-11 01:42 -0800

1from numba import njit 

2import numpy as np 

3from .interpolate import interpolate_image_at_mask 

4 

5 

6@njit 

7def get_ap_range(ap_rad): 

8 """ 

9 Get the range over which the the apodization kernel drops to zero 

10 

11 int(6*ap_rad + 0.5) 

12 

13 Parameters 

14 ---------- 

15 ap_rad: float 

16 The apodization radius over which the kernel goes to zero 

17 

18 Returns 

19 ------- 

20 ap range as an integer 

21 """ 

22 return int(6*ap_rad + 0.5) 

23 

24 

25def apply_apodization_corrections(*, mbobs, ap_rad, mask_bit_val): 

26 """Apply an apodization mask around the edge of the images in an mbobs to 

27 prevent FFT artifacts. 

28 

29 Parameters 

30 ---------- 

31 mbobs: ngmix.MultiBandObsList 

32 The observations to mask 

33 mask_bit_val: int 

34 The bit to set in the bit mask for areas that are apodized. 

35 ap_rad: float 

36 When apodizing, the scale of the kernel. The total kernel goes from 0 to 1 

37 over 6*ap_rad. 

38 """ 

39 ap_mask = np.ones_like(mbobs[0][0].image) 

40 _build_square_apodization_mask(ap_rad, ap_mask) 

41 

42 msk = np.where(ap_mask < 1) 

43 if msk[0].size > 0: 

44 for obslist in mbobs: 

45 for obs in obslist: 

46 # the pixels list will be reset upon exiting 

47 with obs.writeable(): 

48 obs.image[msk] *= ap_mask[msk] 

49 obs.noise[msk] *= ap_mask[msk] 

50 obs.bmask[msk] |= mask_bit_val 

51 if hasattr(obs, "mfrac"): 

52 obs.mfrac[msk] = 1.0 

53 if msk[0].size == obs.image.size: 

54 obs.ignore_zero_weight = False 

55 obs.weight[msk] = 0.0 

56 if np.all(obs.weight == 0): 

57 obs.ignore_zero_weight = False 

58 

59 

60@njit 

61def _build_square_apodization_mask(ap_rad, ap_mask): 

62 ap_range = get_ap_range(ap_rad) 

63 

64 ny, nx = ap_mask.shape 

65 for y in range(min(ap_range+1, ny)): 

66 for x in range(nx): 

67 ap_mask[y, x] *= _ap_kern_kern(y, ap_range, ap_rad) 

68 ap_mask[ny-1 - y, x] *= _ap_kern_kern(y, ap_range, ap_rad) 

69 

70 for y in range(ny): 

71 for x in range(min(ap_range+1, nx)): 

72 ap_mask[y, x] *= _ap_kern_kern(x, ap_range, ap_rad) 

73 ap_mask[y, nx - 1 - x] *= _ap_kern_kern(x, ap_range, ap_rad) 

74 

75 

76def apply_foreground_masking_corrections( 

77 *, mbobs, xm, ym, rm, method, mask_expand_rad, 

78 mask_bit_val, expand_mask_bit_val, interp_bit_val, 

79 symmetrize, ap_rad, iso_buff, rng, 

80): 

81 """ 

82 Apply corrections for masks of large foreground objects like local galaxies 

83 and stars. 

84 

85 Parameters 

86 ---------- 

87 mbobs: ngmix.MultiBandObsList 

88 The observations to mask 

89 xm: np.ndarray 

90 The x/column location of the mask holes in zero-indexed pixels. 

91 ym: np.ndarray 

92 The y/row location of the mask holes in zero-indexed pixels. 

93 rm: np.ndarray 

94 The radii of the mask holes in pixels. 

95 method: str 

96 The method to use for masking corrections. It should be one of 'interp', 

97 'interp-noise', or 'apodize'. 

98 mask_bit_val: int 

99 The bit to set in the bit mask for areas inside the mask holes. 

100 expand_mask_bit_val: int 

101 The bit to set in the bit mask for areas inside the expanded mask holes. 

102 interp_bit_val: int 

103 The bit to set for areas in the mask that are interpolated. 

104 symmetrize: bool 

105 If True, the mask holes will be symmetrized via a 90 degree rotation. 

106 ap_rad: float 

107 When apodizing, the scale of the kernel. The total kernel goes from 0 to 1 

108 over 6*ap_rad. 

109 iso_buff: float 

110 When using 'interp-noise', the number of pixels away from a good pixel a 

111 given pixel must be to be noise interpolated. 

112 rng: np.random.RandomState 

113 An RNG to use when doing 'interp-noise'. 

114 """ 

115 

116 if method == 'interp': 

117 _apply_mask_interp( 

118 mbobs=mbobs, 

119 xm=xm, 

120 ym=ym, 

121 rm=rm, 

122 symmetrize=symmetrize, 

123 mask_bit_val=mask_bit_val, 

124 interp_bit_val=interp_bit_val, 

125 fill_isolated_with_noise=False, 

126 iso_buff=iso_buff, 

127 rng=rng, 

128 ) 

129 elif method == 'interp-noise': 

130 _apply_mask_interp( 

131 mbobs=mbobs, 

132 xm=xm, 

133 ym=ym, 

134 rm=rm, 

135 symmetrize=symmetrize, 

136 mask_bit_val=mask_bit_val, 

137 interp_bit_val=interp_bit_val, 

138 fill_isolated_with_noise=True, 

139 iso_buff=iso_buff, 

140 rng=rng, 

141 ) 

142 elif method == 'apodize': 

143 _apply_mask_apodize( 

144 mbobs=mbobs, 

145 xm=xm, 

146 ym=ym, 

147 rm=rm, 

148 symmetrize=symmetrize, 

149 ap_rad=ap_rad, 

150 mask_bit_val=mask_bit_val, 

151 ) 

152 else: 

153 raise RuntimeError( 

154 "Can only do one of 'interp', 'interp-noise' or 'apodize' for " 

155 "handling foreground masks (got %s)!" % method 

156 ) 

157 

158 if mask_expand_rad > 0: 

159 expanded_bmask = make_foreground_bmask( 

160 xm=xm, 

161 ym=ym, 

162 rm=rm + mask_expand_rad, 

163 dims=mbobs[0][0].image.shape, 

164 symmetrize=symmetrize, 

165 mask_bit_val=expand_mask_bit_val, 

166 ) 

167 for obslist in mbobs: 

168 for obs in obslist: 

169 # the pixels list will be reset upon exiting 

170 with obs.writeable(): 

171 obs.bmask |= expanded_bmask 

172 

173 

174def _apply_mask_interp( 

175 *, 

176 mbobs, 

177 xm, 

178 ym, 

179 rm, 

180 symmetrize, 

181 mask_bit_val, 

182 interp_bit_val, 

183 fill_isolated_with_noise, 

184 iso_buff, 

185 rng, 

186): 

187 

188 # masking is same for all, just take the first 

189 fg_bmask = make_foreground_bmask( 

190 xm=xm, 

191 ym=ym, 

192 rm=rm, 

193 dims=mbobs[0][0].image.shape, 

194 symmetrize=symmetrize, 

195 mask_bit_val=mask_bit_val, 

196 ) 

197 

198 # now modify the masks, weight maps, and interpolate in all 

199 # bands 

200 bad_logic = fg_bmask != 0 

201 wbad = np.where(bad_logic) 

202 if wbad[0].size > 0: 

203 

204 for obslist in mbobs: 

205 for obs in obslist: 

206 # the pixels list will be reset upon exiting 

207 with obs.writeable(): 

208 obs.bmask |= fg_bmask 

209 

210 if hasattr(obs, "mfrac"): 

211 obs.mfrac[wbad] = 1.0 

212 if np.all(bad_logic): 

213 obs.ignore_zero_weight = False 

214 obs.weight[wbad] = 0.0 

215 if np.all(obs.weight == 0): 

216 obs.ignore_zero_weight = False 

217 

218 if not np.all(bad_logic): 

219 wmsk = obs.weight > 0 

220 wgt = np.median(obs.weight[wmsk]) 

221 interp_image = interpolate_image_at_mask( 

222 image=obs.image, 

223 bad_msk=bad_logic, 

224 maxfrac=1.0, 

225 iso_buff=iso_buff, 

226 fill_isolated_with_noise=fill_isolated_with_noise, 

227 rng=rng, 

228 weight=wgt, 

229 ) 

230 interp_noise = interpolate_image_at_mask( 

231 image=obs.noise, 

232 bad_msk=bad_logic, 

233 maxfrac=1.0, 

234 iso_buff=iso_buff, 

235 fill_isolated_with_noise=fill_isolated_with_noise, 

236 rng=rng, 

237 weight=wgt, 

238 ) 

239 else: 

240 interp_image = None 

241 interp_noise = None 

242 

243 if interp_image is None or interp_noise is None: 

244 obs.bmask |= mask_bit_val 

245 if hasattr(obs, "mfrac"): 

246 obs.mfrac[:, :] = 1.0 

247 obs.ignore_zero_weight = False 

248 obs.weight[:, :] = 0.0 

249 else: 

250 obs.image = interp_image 

251 obs.noise = interp_noise 

252 obs.bmask[wbad] |= interp_bit_val 

253 

254 

255def _apply_mask_apodize( 

256 *, 

257 mbobs, 

258 xm, 

259 ym, 

260 rm, 

261 symmetrize, 

262 ap_rad, 

263 mask_bit_val, 

264): 

265 obs0 = mbobs[0][0] 

266 ap_mask = make_foreground_apodization_mask( 

267 xm=xm, 

268 ym=ym, 

269 rm=rm, 

270 dims=obs0.image.shape, 

271 symmetrize=symmetrize, 

272 ap_rad=ap_rad, 

273 ) 

274 

275 msk = ap_mask < 1 

276 if np.any(msk): 

277 for obslist in mbobs: 

278 for obs in obslist: 

279 # the pixels list will be reset upon exiting 

280 with obs.writeable(): 

281 obs.image *= ap_mask 

282 obs.noise *= ap_mask 

283 obs.bmask[msk] |= mask_bit_val 

284 if hasattr(obs, "mfrac"): 

285 obs.mfrac[msk] = 1.0 

286 if np.all(msk): 

287 obs.ignore_zero_weight = False 

288 obs.weight[msk] = 0.0 

289 if np.all(obs.weight == 0): 

290 obs.ignore_zero_weight = False 

291 

292 

293def make_foreground_bmask( 

294 *, 

295 xm, 

296 ym, 

297 rm, 

298 dims, 

299 symmetrize, 

300 mask_bit_val, 

301): 

302 """ 

303 Make a bit mask marking the locations of holes at (xm,ym) with radii rm 

304 with bit mask_bit_val. 

305 

306 Parameters 

307 ---------- 

308 xm: np.ndarray 

309 The x/column location of the mask holes in zero-indexed pixels. 

310 ym: np.ndarray 

311 The y/row location of the mask holes in zero-indexed pixels. 

312 rm: np.ndarray 

313 The radii of the mask holes in pixels. 

314 dims: tuple of ints 

315 The dimensions of the mask. 

316 symmetrize: bool 

317 If True, the mask holes will be symmetrized via a 90 degree rotation. 

318 mask_bit_val: int 

319 The bit to set in the bit mask for areas inside the mask holes. 

320 

321 Returns 

322 ------- 

323 bmask: np.ndarray 

324 The bit mask. 

325 """ 

326 

327 # must be native byte order for numba 

328 bmask = np.zeros(dims, dtype='i4') 

329 _do_mask_foreground( 

330 rows=ym.astype('f8'), 

331 cols=xm.astype('f8'), 

332 radius_pixels=rm.astype('f8'), 

333 bmask=bmask, 

334 flag=mask_bit_val, 

335 ) 

336 

337 if symmetrize: 

338 bmask |= np.rot90(bmask) 

339 

340 return bmask 

341 

342 

343def make_foreground_apodization_mask( 

344 *, 

345 xm, 

346 ym, 

347 rm, 

348 dims, 

349 symmetrize, 

350 ap_rad, 

351): 

352 """ 

353 Make foreground apodization mask for mask holes at (xm,ym) with radius rm. 

354 

355 Parameters 

356 ---------- 

357 xm: np.ndarray 

358 The x/column location of the mask holes in zero-indexed pixels. 

359 ym: np.ndarray 

360 The y/row location of the mask holes in zero-indexed pixels. 

361 rm: np.ndarray 

362 The radii of the mask holes in pixels. 

363 dims: tuple of ints 

364 The dimensions of the mask. 

365 symmetrize: bool 

366 If True, the mask holes will be symmetrized via a 90 degree rotation. 

367 ap_rad: float 

368 When apodizing, the scale of the kernel. The total kernel goes from 0 to 1 

369 over 6*ap_rad. 

370 

371 Returns 

372 ------- 

373 ap_mask: np.ndarray 

374 The apodization mask. 

375 """ 

376 

377 # must be native byte order for numba 

378 ap_mask = np.ones(dims, dtype='f8') 

379 _do_apodization_mask( 

380 rows=ym.astype('f8'), 

381 cols=xm.astype('f8'), 

382 radius_pixels=rm.astype('f8'), 

383 ap_mask=ap_mask, 

384 ap_rad=ap_rad, 

385 ) 

386 

387 if symmetrize: 

388 ap_mask *= np.rot90(ap_mask) 

389 

390 return ap_mask 

391 

392 

393@njit 

394def _intersects(row, col, radius_pixels, nrows, ncols): 

395 """ 

396 low level routine to check if the mask intersects the image. 

397 For simplicty just check the bounding rectangle 

398 

399 Parameters 

400 ---------- 

401 row, col: float 

402 The row and column of the mask center 

403 radius_pixels: float 

404 The radius for the mask 

405 nrows, ncols: int 

406 Shape of the image 

407 

408 Returns 

409 ------- 

410 True if it intersects, otherwise False 

411 """ 

412 

413 low_row = -radius_pixels 

414 high_row = nrows + radius_pixels - 1 

415 low_col = -radius_pixels 

416 high_col = ncols + radius_pixels - 1 

417 

418 if ( 

419 row > low_row and row < high_row and 

420 col > low_col and col < high_col 

421 ): 

422 return True 

423 else: 

424 return False 

425 

426 

427@njit 

428def _ap_kern_kern(x, m, h): 

429 # cumulative triweight kernel 

430 y = (x - m) / h + 3 

431 if y < -3: 

432 return 0 

433 elif y > 3: 

434 return 1 

435 else: 

436 val = ( 

437 -5 * y ** 7 / 69984 

438 + 7 * y ** 5 / 2592 

439 - 35 * y ** 3 / 864 

440 + 35 * y / 96 

441 + 1 / 2 

442 ) 

443 return val 

444 

445 

446@njit 

447def _do_apodization_mask(*, rows, cols, radius_pixels, ap_mask, ap_rad): 

448 """low-level code to make the apodization mask 

449 

450 Parameters 

451 ---------- 

452 rows, cols: arrays 

453 Arrays of rows/cols of mask locations in the "local" pixel frame of the 

454 slice, not the overall pixels of the big coadd. These positions may be 

455 off the image. 

456 radius_pixels: array 

457 The radius for each mask. 

458 ap_mask: array 

459 The array to fill with the apodization fraction. 

460 ap_rad: float 

461 The scale in pixels for the apodization transition from 1 to 0. 

462 """ 

463 ny, nx = ap_mask.shape 

464 ns = cols.shape[0] 

465 

466 nmasked = 0 

467 for i in range(ns): 

468 x = cols[i] 

469 y = rows[i] 

470 rad = radius_pixels[i] 

471 rad2 = rad**2 

472 

473 if not _intersects(y, x, rad, ny, nx): 

474 continue 

475 

476 for _y in range(ny): 

477 dy2 = (_y - y)**2 

478 for _x in range(nx): 

479 dr2 = (_x - x)**2 + dy2 

480 if dr2 < rad2: 

481 ap_mask[_y, _x] *= _ap_kern_kern(np.sqrt(dr2), rad, ap_rad) 

482 nmasked += 1 

483 

484 return nmasked 

485 

486 

487@njit 

488def _do_mask_foreground(*, rows, cols, radius_pixels, bmask, flag): 

489 """ 

490 low level code to mask foreground objects 

491 

492 Parameters 

493 ---------- 

494 rows, cols: arrays 

495 Arrays of rows/cols of mask locations in the "local" pixel frame of the 

496 slice, not the overall pixels of the big coadd. These positions may be 

497 off the image. 

498 radius_pixels: array 

499 The radius for each mask. 

500 bmask: array 

501 The bmask to modify. 

502 flag: int 

503 The flag value to "or" into the bmask. 

504 """ 

505 nrows, ncols = bmask.shape 

506 nmasked = 0 

507 

508 for ifg in range(rows.size): 

509 row = rows[ifg] 

510 col = cols[ifg] 

511 

512 rad = radius_pixels[ifg] 

513 rad2 = rad * rad 

514 

515 if not _intersects(row, col, rad, nrows, ncols): 

516 continue 

517 

518 for irow in range(nrows): 

519 rowdiff2 = (row - irow)**2 

520 for icol in range(ncols): 

521 

522 r2 = rowdiff2 + (col - icol)**2 

523 if r2 < rad2: 

524 bmask[irow, icol] |= flag 

525 nmasked += 1 

526 

527 return nmasked