ENH apodize the pre-PSF moments

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
 

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

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)