Coverage for metadetect / masking.py: 8%
169 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-27 01:48 -0800
1from numba import njit
2import numpy as np
3from .interpolate import interpolate_image_at_mask
6@njit
7def get_ap_range(ap_rad):
8 """
9 Get the range over which the the apodization kernel drops to zero
11 int(6*ap_rad + 0.5)
13 Parameters
14 ----------
15 ap_rad: float
16 The apodization radius over which the kernel goes to zero
18 Returns
19 -------
20 ap range as an integer
21 """
22 return int(6*ap_rad + 0.5)
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.
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)
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
60@njit
61def _build_square_apodization_mask(ap_rad, ap_mask):
62 ap_range = get_ap_range(ap_rad)
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)
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)
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.
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 """
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 )
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
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):
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 )
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:
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
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
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
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
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 )
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
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.
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.
321 Returns
322 -------
323 bmask: np.ndarray
324 The bit mask.
325 """
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 )
337 if symmetrize:
338 bmask |= np.rot90(bmask)
340 return bmask
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.
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.
371 Returns
372 -------
373 ap_mask: np.ndarray
374 The apodization mask.
375 """
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 )
387 if symmetrize:
388 ap_mask *= np.rot90(ap_mask)
390 return ap_mask
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
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
408 Returns
409 -------
410 True if it intersects, otherwise False
411 """
413 low_row = -radius_pixels
414 high_row = nrows + radius_pixels - 1
415 low_col = -radius_pixels
416 high_col = ncols + radius_pixels - 1
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
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
446@njit
447def _do_apodization_mask(*, rows, cols, radius_pixels, ap_mask, ap_rad):
448 """low-level code to make the apodization mask
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]
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
473 if not _intersects(y, x, rad, ny, nx):
474 continue
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
484 return nmasked
487@njit
488def _do_mask_foreground(*, rows, cols, radius_pixels, bmask, flag):
489 """
490 low level code to mask foreground objects
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
508 for ifg in range(rows.size):
509 row = rows[ifg]
510 col = cols[ifg]
512 rad = radius_pixels[ifg]
513 rad2 = rad * rad
515 if not _intersects(row, col, rad, nrows, ncols):
516 continue
518 for irow in range(nrows):
519 rowdiff2 = (row - irow)**2
520 for icol in range(ncols):
522 r2 = rowdiff2 + (col - icol)**2
523 if r2 < rad2:
524 bmask[irow, icol] |= flag
525 nmasked += 1
527 return nmasked