ENH extra smoothing in prepsf moments

CHANGES.md

@@ -1,3 +1,14 @@
+## v2.1.0
+
+### new features
+
+    - Added `fwhm_smooth` keyword to pre-PSF moments routines to allow for extra
+      smoothing of the profile before the moments are measured.
+    - Added `use_pix_weight` keyword to pre-PSF moments routines to enable inverse
+      variance pixel weighting.
+    - Added caching of FFTs in metacal and pre-PSF moment rountines.
+
+
 ## v2.0.6
 
 ### bug fixes

ngmix/_version.py

@@ -1 +1 @@
-__version__ = '2.0.6'  # noqa
+__version__ = '2.1.0'  # noqa

ngmix/prepsfmom.py

@@ -37,12 +37,26 @@ class PrePSFMom(object):
     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.
+    fwhm_smooth : float, optional
+        If non-zero, this optional applies additional Gaussian smoothing to the
+        object before computing the moments. Typically a non-zero value results
+        in less shape noise.
+    use_pix_weight : bool, optional
+        If `True`, this option applies inverse pixel variance weighting in Fourier
+        space when computing the moments. Typically this optional results in poorer
+        performance, but for round PSFs, it results in similar performance as
+        `fwhm_smooth`.
     """
-    def __init__(self, fwhm, kernel, pad_factor=4, ap_rad=1.5):
+    def __init__(
+        self, fwhm, kernel, pad_factor=4, ap_rad=1.5,
+        fwhm_smooth=0, use_pix_weight=False
+    ):
         self.fwhm = fwhm
         self.pad_factor = pad_factor
         self.kernel = kernel
         self.ap_rad = ap_rad
+        self.fwhm_smooth = fwhm_smooth
+        self.use_pix_weight = use_pix_weight
         if self.kernel == "ksigma":
             self.kind = "ksigma"
         elif self.kernel in ["gauss", "pgauss"]:
@@ -133,13 +147,15 @@ def _meas(self, obs, psf_obs, return_kernels):
                 float(self.fwhm),
                 float(obs.jacobian.dvdrow), float(obs.jacobian.dvdcol),
                 float(obs.jacobian.dudrow), float(obs.jacobian.dudcol),
+                float(self.fwhm_smooth),
             )
         elif self.kernel in ["gauss", "pgauss"]:
             kernels = _gauss_kernels(
                 int(target_dim),
                 float(self.fwhm),
                 float(obs.jacobian.dvdrow), float(obs.jacobian.dvdcol),
                 float(obs.jacobian.dudrow), float(obs.jacobian.dudcol),
+                float(self.fwhm_smooth),
             )
         else:
             raise ValueError(
@@ -154,6 +170,7 @@ def _meas(self, obs, psf_obs, return_kernels):
         mom, mom_cov = _measure_moments_fft(
             kim, kpsf_im, tot_var, eff_pad_factor, kernels,
             im_row - psf_im_row, im_col - psf_im_col,
+            use_pix_weight=self.use_pix_weight,
         )
         res = make_mom_result(mom, mom_cov)
         if res['flags'] != 0:
@@ -193,9 +210,23 @@ class KSigmaMom(PrePSFMom):
     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.
+    fwhm_smooth : float, optional
+        If non-zero, this optional applies additional Gaussian smoothing to the
+        object before computing the moments. Typically a non-zero value results
+        in less shape noise.
+    use_pix_weight : bool, optional
+        If `True`, this option applies inverse pixel variance weighting in Fourier
+        space when computing the moments. Typically this optional results in poorer
+        performance, but for round PSFs, it results in similar performance as
+        `fwhm_smooth`.
     """
-    def __init__(self, fwhm, pad_factor=4, ap_rad=1.5):
-        super().__init__(fwhm, 'ksigma', pad_factor=pad_factor, ap_rad=ap_rad)
+    def __init__(
+        self, fwhm, pad_factor=4, ap_rad=1.5, fwhm_smooth=0, use_pix_weight=False,
+    ):
+        super().__init__(
+            fwhm, 'ksigma', pad_factor=pad_factor, ap_rad=ap_rad,
+            fwhm_smooth=fwhm_smooth, use_pix_weight=use_pix_weight,
+        )
 
 
 class PGaussMom(PrePSFMom):
@@ -217,16 +248,32 @@ class PGaussMom(PrePSFMom):
     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.
+    fwhm_smooth : float, optional
+        If non-zero, this optional applies additional Gaussian smoothing to the
+        object before computing the moments. Typically a non-zero value results
+        in less shape noise.
+    use_pix_weight : bool, optional
+        If `True`, this option applies inverse pixel variance weighting in Fourier
+        space when computing the moments. Typically this optional results in poorer
+        performance, but for round PSFs, it results in similar performance as
+        `fwhm_smooth`.
     """
-    def __init__(self, fwhm, pad_factor=4, ap_rad=1.5):
-        super().__init__(fwhm, 'pgauss', pad_factor=pad_factor, ap_rad=ap_rad)
+    def __init__(
+        self, fwhm, pad_factor=4, ap_rad=1.5, fwhm_smooth=0, use_pix_weight=False,
+    ):
+        super().__init__(
+            fwhm, 'pgauss', pad_factor=pad_factor, ap_rad=ap_rad,
+            fwhm_smooth=fwhm_smooth, use_pix_weight=use_pix_weight,
+        )
 
 
 # keep this here for API consistency
 PrePSFGaussMom = PGaussMom
 
 
-def _measure_moments_fft(kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, dcol):
+def _measure_moments_fft(
+    kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, dcol, use_pix_weight=False,
+):
     # we only need to do things where the kernel is non-zero
     # this saves a bunch of CPU cycles
     msk = kernels["msk"]
@@ -262,10 +309,15 @@ def _measure_moments_fft(kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, d
     fkp = kernels["fkp"]
     fkc = kernels["fkc"]
 
-    mf = np.sum((kim * fkf).real) * df2
-    mr = np.sum((kim * fkr).real) * df2
-    mp = np.sum((kim * fkp).real) * df2
-    mc = np.sum((kim * fkc).real) * df2
+    if use_pix_weight:
+        wgt = (kpsf_im * np.conj(kpsf_im)).real
+    else:
+        wgt = 1.0
+
+    mf = np.sum((kim * fkf * wgt).real) * df2
+    mr = np.sum((kim * fkr * wgt).real) * df2
+    mp = np.sum((kim * fkp * wgt).real) * df2
+    mc = np.sum((kim * fkc * wgt).real) * df2
 
     # build a covariance matrix of the moments
     # here we assume each Fourier mode is independent and sum the variances
@@ -279,10 +331,17 @@ def _measure_moments_fft(kim, kpsf_im, tot_var, eff_pad_factor, kernels, drow, d
     m_cov[1, 1] = 1
     tot_var *= eff_pad_factor**2
     tot_var_df4 = tot_var * df4
-    kerns = [fkp / kpsf_im, fkc / kpsf_im, fkr / kpsf_im, fkf / kpsf_im]
+    psf_kerns_fac = wgt / kpsf_im
+    kerns = [
+        fkp * psf_kerns_fac,
+        fkc * psf_kerns_fac,
+        fkr * psf_kerns_fac,
+        fkf * psf_kerns_fac,
+    ]
     conj_kerns = [np.conj(k) for k in kerns]
     for i in range(2, 6):
         for j in range(i, 6):
+            # subtract two since kernels start at second moments
             m_cov[i, j] = np.sum((kerns[i-2] * conj_kerns[j-2]).real) * tot_var_df4
             m_cov[j, i] = m_cov[i, j]
 
@@ -428,6 +487,7 @@ def _ksigma_kernels(
     dim,
     kernel_size,
     dvdrow, dvdcol, dudrow, dudcol,
+    fwhm_smooth,
 ):
     """This function builds a ksigma kernel in Fourier-space.
 
@@ -498,6 +558,16 @@ def _ksigma_kernels(
             "norm = %f (should be 1)!" % (kernel_size, nrm)
         )
 
+    # add smoothing after norm check above for kernel size
+    if fwhm_smooth > 0:
+        sigma_smooth = fwhm_to_sigma(fwhm_smooth)
+        chi2_2_smooth = sigma_smooth * sigma_smooth / 2 * fmag2
+        exp_val_smooth = np.zeros_like(karg)
+        msk_smooth = (chi2_2_smooth < FASTEXP_MAX_CHI2/2) & (chi2_2_smooth >= 0)
+        exp_val_smooth[msk_smooth] = fexp_arr(-chi2_2_smooth[msk_smooth])
+        fkf *= exp_val_smooth
+        knrm *= exp_val_smooth
+
     # the moment kernels take a bit more work
     # product by u^2 in real space is -dk^2/dku^2 in Fourier space
     # same holds for v and cross deriv is -dk^2/dkudkv
@@ -536,6 +606,7 @@ def _gauss_kernels(
     dim,
     kernel_size,
     dvdrow, dvdcol, dudrow, dudcol,
+    fwhm_smooth,
 ):
     """This function builds a Gaussian kernel in Fourier-space.
 
@@ -591,6 +662,16 @@ def _gauss_kernels(
             "norm = %f (should be 1)!" % (kernel_size, nrm)
         )
 
+    # add smoothing after norm check above for kernel size
+    if fwhm_smooth > 0:
+        sigma_smooth = fwhm_to_sigma(fwhm_smooth)
+        chi2_2_smooth = sigma_smooth * sigma_smooth / 2 * fmag2
+        exp_val_smooth = np.zeros_like(exp_val)
+        msk_smooth = (chi2_2_smooth < FASTEXP_MAX_CHI2/2) & (chi2_2_smooth >= 0)
+        exp_val_smooth[msk_smooth] = fexp_arr(-chi2_2_smooth[msk_smooth])
+        fkf *= exp_val_smooth
+        knrm *= exp_val_smooth
+
     # the moment kernels take a bit more work
     # product by u^2 in real space is -dk^2/dku^2 in Fourier space
     # same holds for v and cross deriv is -dk^2/dkudkv

ngmix/tests/test_metacal_cache.py

@@ -1,11 +1,15 @@
 import time
 import numpy as np
+
+from flaky import flaky
+
 import ngmix
 import ngmix.metacal.metacal
 from ._galsim_sims import _get_obs
 from ..metacal.metacal import _cached_galsim_stuff
 
 
+@flaky(max_runs=10)
 def test_metacal_cache():
     # first warm up numba
     rng = np.random.RandomState(seed=100)