ENH deconvolve pixel in prepsf moments

CHANGES.md

@@ -4,6 +4,7 @@
 
    - Caching in pre-psf moments and metacal is now optional
      with an API to turn it on. Default is off.
+   - Pixel in pre-PSF moments is not deconvolved when doing PSFs.
 
 ## v2.2.1
 

ngmix/prepsfmom.py

@@ -132,8 +132,9 @@ def _meas(self, obs, psf_obs, return_kernels):
                 0,  # we do not apodize PSF stamps since it should not be needed
             )
         else:
-            # delta function in k-space
-            kpsf_im = np.ones_like(kim, dtype=np.complex128)
+            # pixel in real-space
+            kpsf_im = _pixel_fft(kim.shape[0])
+
             psf_im_row = 0.0
             psf_im_col = 0.0
 
@@ -484,6 +485,17 @@ def _zero_pad_and_compute_fft_impl(im, cen_row, cen_col, target_dim, ap_rad):
     return kpim, pad_cen_row, pad_cen_col
 
 
+@functools.lru_cache(maxsize=128)
+def _pixel_fft(dim):
+    # pixel in real-space
+    f = fft.fftfreq(dim)
+    f = np.sinc(f)
+    fx = f.reshape(1, -1)
+    fy = f.reshape(-1, 1)
+    kpsf_im = fx * fy
+    return kpsf_im
+
+
 # see https://stackoverflow.com/a/52332109 for how this works
 @functools.lru_cache(maxsize=128)
 def _zero_pad_and_compute_fft_cached_impl(
@@ -521,10 +533,14 @@ def _deconvolve_im_psf_inplace(kim, kpsf_im, max_amp, min_psf_frac=1e-5):
     """
     min_amp = min_psf_frac * max_amp
     abs_kpsf_im = np.abs(kpsf_im)
-    msk = abs_kpsf_im <= min_amp
+    msk = (abs_kpsf_im <= min_amp) & (abs_kpsf_im != 0)
     if np.any(msk):
         kpsf_im[msk] = kpsf_im[msk] / abs_kpsf_im[msk] * min_amp
 
+    msk = (abs_kpsf_im <= min_amp) & (abs_kpsf_im == 0)
+    if np.any(msk):
+        kpsf_im[msk] = min_amp
+
     kim /= kpsf_im
     return kim, kpsf_im, msk
 

ngmix/tests/test_prepsfmom.py

@@ -382,7 +382,7 @@ def test_prepsfmom_gauss(
         nx=image_size,
         ny=image_size,
         wcs=gs_wcs,
-        method='no_pixel').array
+    ).array
     obs = Observation(
         image=im_true,
         jacobian=jac,
@@ -508,7 +508,7 @@ def test_prepsfmom_mn_cov_psf(
         nx=image_size,
         ny=image_size,
         wcs=gs_wcs,
-        method='no_pixel').array
+    ).array
     obs = Observation(
         image=im_true,
         jacobian=jac,
@@ -645,7 +645,7 @@ def _run_sim_fwhm_smooth(fwhm_smooth):
             nx=image_size,
             ny=image_size,
             wcs=gs_wcs,
-            method='no_pixel').array
+        ).array
         obs = Observation(
             image=im_true,
             jacobian=jac,
@@ -937,7 +937,7 @@ def test_prepsfmom_gauss_true_flux(
         nx=image_size,
         ny=image_size,
         wcs=gs_wcs,
-        method='no_pixel').array
+    ).array
     obs = Observation(
         image=im_true,
         jacobian=jac,
@@ -1031,25 +1031,37 @@ def test_prepsfmom_comp_to_gaussmom_simple(
         ny=image_size,
         wcs=gs_wcs,
     ).array
+    im_true_nopixel = gal.drawImage(
+        nx=image_size,
+        ny=image_size,
+        wcs=gs_wcs,
+        method="no_pixel",
+    ).array
+
     obs = Observation(
         image=im_true,
         jacobian=jac,
     )
+    obs_nopixel = Observation(
+        image=im_true_nopixel,
+        jacobian=jac,
+    )
+
     res = PGaussMom(
         fwhm=mom_fwhm, pad_factor=pad_factor,
     ).go(
         obs=obs, no_psf=True, return_kernels=True,
     )
 
     from ngmix.gaussmom import GaussMom
-    res_gmom = GaussMom(fwhm=mom_fwhm).go(obs=obs)
+    res_gmom = GaussMom(fwhm=mom_fwhm).go(obs=obs_nopixel)
 
     for k in sorted(res):
         if k in res_gmom:
             print("%s:" % k, res[k], res_gmom[k])
 
     for k in ["flux", "flux_err", "T", "T_err", "e", "e_cov"]:
-        assert_allclose(res[k], res_gmom[k], atol=0, rtol=1e-2)
+        assert_allclose(res[k], res_gmom[k], atol=0, rtol=2.5e-2)
 
 
 @pytest.mark.parametrize('pixel_scale', [0.25, 0.125])
@@ -1110,9 +1122,16 @@ def test_prepsfmom_comp_to_gaussmom_fwhm_smooth(
             nx=image_size,
             ny=image_size,
             wcs=gs_wcs,
+            method="no_pixel",
         ).array
     else:
-        im_true_smooth = im_true
+        im_true_smooth = gal.drawImage(
+            nx=image_size,
+            ny=image_size,
+            wcs=gs_wcs,
+            method="no_pixel",
+        ).array
+
     obs_smooth = Observation(
         image=im_true_smooth,
         jacobian=jac,
@@ -1124,8 +1143,8 @@ def test_prepsfmom_comp_to_gaussmom_fwhm_smooth(
             print("%s:" % k, res[k], res_gmom[k])
 
     assert_allclose(res["flux"], res_gmom["flux"], atol=0, rtol=5e-4)
-    assert_allclose(res["T"], res_gmom["T"], atol=0, rtol=1e-3)
-    assert_allclose(res["e"], res_gmom["e"], atol=0, rtol=1e-3)
+    assert_allclose(res["T"], res_gmom["T"], atol=0, rtol=2e-3)
+    assert_allclose(res["e"], res_gmom["e"], atol=0, rtol=3e-3)
     # the errors do not match - this is because the underlying noise model is
     # different - the pure gaussian moments weight map is an error on the convolved
     # profile whereas the pre-PSF case uses error propagation through the
@@ -1209,8 +1228,33 @@ def _sim_apodize(flux_factor, ap_rad):
     ap_mask = np.ones_like(im)
     if ap_rad > 0:
         _build_square_apodization_mask(ap_rad, ap_mask)
+
+    # get true flux
+    im_nopixel = gal.drawImage(
+        nx=image_size,
+        ny=image_size,
+        wcs=gs_wcs,
+        method="no_pixel",
+    ).array
+
+    im_nopixel += 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="no_pixel",
+    ).array
+
     obs_ap = Observation(
-        image=im * ap_mask,
+        image=im_nopixel * ap_mask,
         jacobian=jac,
     )