diff --git a/CHANGES.md b/CHANGES.md index ee42f45f..86e202e1 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -5,6 +5,7 @@ - implemented check for equality, copy() method, copy, and deepcopy for Gmix objects, jacobians, and Observation/ObsList/MultiBandObsList - added get_data() method for jacobian objects + - added apodization to the pre-PSF moments to help prevent FFT artifacts ## v2.0.4 diff --git a/ngmix/prepsfmom.py b/ngmix/prepsfmom.py index 96112296..4cf8ea2e 100644 --- a/ngmix/prepsfmom.py +++ b/ngmix/prepsfmom.py @@ -2,6 +2,7 @@ import numpy as np import scipy.fft as fft +from numba import njit from ngmix.observation import Observation from ngmix.moments import fwhm_to_sigma, make_mom_result @@ -32,11 +33,15 @@ class PrePSFMom(object): are aliases for the same thing. pad_factor : int, optional The factor by which to pad the FFTs used for the image. Default is 4. + ap_rad : float, optional + The apodization radius for the stamp in pixels. The default of 1.5 is likely + fine for most ground based surveys. """ - def __init__(self, fwhm, kernel, pad_factor=4): + def __init__(self, fwhm, kernel, pad_factor=4, ap_rad=1.5): self.fwhm = fwhm self.pad_factor = pad_factor self.kernel = kernel + self.ap_rad = ap_rad def go(self, obs, return_kernels=False, no_psf=False): """Measure the pre-PSF ksigma moments. @@ -73,6 +78,7 @@ def _meas(self, obs, psf_obs, return_kernels): # pad image, psf and weight map, get FFTs, apply cen_phases kim, im_row, im_col = _zero_pad_and_compute_fft( obs.image, obs.jacobian.row0, obs.jacobian.col0, target_dim, + self.ap_rad, ) fft_dim = kim.shape[0] @@ -81,6 +87,7 @@ def _meas(self, obs, psf_obs, return_kernels): psf_obs.image, psf_obs.jacobian.row0, psf_obs.jacobian.col0, target_dim, + 0, # we do not apodize PSF stamps since it should not be needed ) else: # delta function in k-space @@ -174,12 +181,15 @@ class KSigmaMom(PrePSFMom): arcseconds. pad_factor : int, optional The factor by which to pad the FFTs used for the image. Default is 4. + ap_rad : float, optional + The apodization radius for the stamp in pixels. The default of 1.5 is likely + fine for most ground based surveys. """ kind = "ksigma" - def __init__(self, fwhm, pad_factor=4): - super().__init__(fwhm, 'ksigma', pad_factor=pad_factor) + def __init__(self, fwhm, pad_factor=4, ap_rad=1.5): + super().__init__(fwhm, 'ksigma', pad_factor=pad_factor, ap_rad=ap_rad) class PGaussMom(PrePSFMom): @@ -198,12 +208,15 @@ class PGaussMom(PrePSFMom): arcseconds. pad_factor : int, optional The factor by which to pad the FFTs used for the image. Default is 4. + ap_rad : float, optional + The apodization radius for the stamp in pixels. The default of 1.5 is likely + fine for most ground based surveys. """ kind = "pgauss" - def __init__(self, fwhm, pad_factor=4): - super().__init__(fwhm, 'pgauss', pad_factor=pad_factor) + def __init__(self, fwhm, pad_factor=4, ap_rad=1.5): + super().__init__(fwhm, 'pgauss', pad_factor=pad_factor, ap_rad=ap_rad) # keep this here for API consistency @@ -275,6 +288,41 @@ def _measure_moments_fft(kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, d return mom, m_cov +@njit +def _ap_kern_kern(x, m, h): + # cumulative triweight kernel + y = (x - m) / h + 3 + if y < -3: + return 0 + elif y > 3: + return 1 + else: + val = ( + -5 * y ** 7 / 69984 + + 7 * y ** 5 / 2592 + - 35 * y ** 3 / 864 + + 35 * y / 96 + + 1 / 2 + ) + return val + + +@njit +def _build_square_apodization_mask(ap_rad, ap_mask): + ap_range = int(6*ap_rad + 0.5) + + ny, nx = ap_mask.shape + for y in range(min(ap_range+1, ny)): + for x in range(nx): + ap_mask[y, x] *= _ap_kern_kern(y, ap_range, ap_rad) + ap_mask[ny-1 - y, x] *= _ap_kern_kern(y, ap_range, ap_rad) + + for y in range(ny): + for x in range(min(ap_range+1, nx)): + ap_mask[y, x] *= _ap_kern_kern(x, ap_range, ap_rad) + ap_mask[y, nx - 1 - x] *= _ap_kern_kern(x, ap_range, ap_rad) + + def _zero_pad_image(im, target_dim): """zero pad an image, returning it and the offsets before and after the original image""" @@ -315,11 +363,16 @@ def _compute_cen_phase_shift(cen_row, cen_col, dim, msk=None): return np.cos(kcen) + 1j*np.sin(kcen) -def _zero_pad_and_compute_fft(im, cen_row, cen_col, target_dim): +def _zero_pad_and_compute_fft(im, cen_row, cen_col, target_dim, ap_rad): """zero pad and compute the FFT Returns the fft, cen_row in the padded image, and cen_col in the padded image. """ + if ap_rad > 0: + ap_mask = np.ones_like(im) + _build_square_apodization_mask(ap_rad, ap_mask) + im = im * ap_mask + pim, pad_width_before, _ = _zero_pad_image(im, target_dim) pad_cen_row = cen_row + pad_width_before pad_cen_col = cen_col + pad_width_before diff --git a/ngmix/tests/test_prepsfmom.py b/ngmix/tests/test_prepsfmom.py index 83db7a90..5221a870 100644 --- a/ngmix/tests/test_prepsfmom.py +++ b/ngmix/tests/test_prepsfmom.py @@ -2,7 +2,7 @@ import numpy as np import pytest -from ngmix.prepsfmom import KSigmaMom, PGaussMom +from ngmix.prepsfmom import KSigmaMom, PGaussMom, _build_square_apodization_mask from ngmix import Jacobian from ngmix import Observation from ngmix.moments import make_mom_result @@ -663,7 +663,7 @@ def test_prepsfmom_gauss_true_flux( 2, 0.5, ]) @pytest.mark.parametrize('image_size', [ - 53, + 107, ]) @pytest.mark.parametrize('pad_factor', [ 3.5, 4, @@ -721,3 +721,109 @@ def test_prepsfmom_comp_to_gaussmom( for k in ["flux", "flux_err", "T", "T_err", "e", "e_cov"]: assert np.allclose(res[k], res_gmom[k], atol=0, rtol=1e-2) + + +def _sim_apodize(flux_factor, ap_rad): + """ + we are simulating an object at the center with a bright object right on the + edge of the stamp. + + We then apply apodization to the image and measure the same Gaussian moment + with either the Fourier-space code or the real-space one. + + We compare the case with zero apodization to non-zero in the test below + and assert that with apodization the results from Fourier-space match the + real-space results better. + """ + rng = np.random.RandomState(seed=100) + image_size = 53 + pixel_scale = 0.25 + fwhm = 0.9 + mom_fwhm = 2.0 + pad_factor = 4 + + cen = (image_size - 1)/2 + gs_wcs = galsim.ShearWCS( + pixel_scale, galsim.Shear(g1=-0, g2=0.0)).jacobian() + scale = np.sqrt(gs_wcs.pixelArea()) + shift = rng.uniform(low=-scale/2, high=scale/2, size=2) + xy = gs_wcs.toImage(galsim.PositionD(shift)) + + jac = Jacobian( + y=cen + xy.y, x=cen + xy.x, + dudx=gs_wcs.dudx, dudy=gs_wcs.dudy, + dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy) + + gal = galsim.Gaussian( + fwhm=fwhm + ).shear( + g1=-0.1, g2=0.2 + ).withFlux( + 400 + ).shift( + dx=shift[0], dy=shift[1] + ) + + # get true flux + im_true = gal.drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs, + ).array + + im = im_true.copy() + im += galsim.Exponential( + half_light_radius=fwhm + ).shear( + g1=-0.5, g2=0.2 + ).shift( + cen*pixel_scale, + 0, + ).withFlux( + 400*flux_factor + ).drawImage( + nx=image_size, + ny=image_size, + wcs=gs_wcs, + method="real_space", + ).array + + obs = Observation( + image=im, + jacobian=jac, + ) + res = PGaussMom(fwhm=mom_fwhm, pad_factor=pad_factor, ap_rad=ap_rad).go( + obs=obs, no_psf=True, return_kernels=True, + ) + + ap_mask = np.ones_like(im) + if ap_rad > 0: + _build_square_apodization_mask(ap_rad, ap_mask) + obs_ap = Observation( + image=im * ap_mask, + jacobian=jac, + ) + + from ngmix.gaussmom import GaussMom + res_gmom = GaussMom(fwhm=mom_fwhm).go(obs=obs_ap) + + return res, res_gmom + + +@pytest.mark.parametrize("flux_factor", [1e2, 1e3, 1e5]) +def test_prepsfmom_apodize(flux_factor): + res, res_geom = _sim_apodize(flux_factor, 1.5) + ap_diffs = np.array([ + np.abs(res[k] - res_geom[k]) + for k in ["e1", "e2", "T", "flux"] + ]) + print("apodized:", ap_diffs) + + res, res_geom = _sim_apodize(flux_factor, 0) + zero_diffs = np.array([ + np.abs(res[k] - res_geom[k]) + for k in ["e1", "e2", "T", "flux"] + ]) + print("non-apodized:", zero_diffs) + + assert np.all(zero_diffs > ap_diffs)