ENH added regularization function for moments

CHANGES.md

@@ -1,3 +1,10 @@
+## unreleased
+
+### new features
+
+   - Added function to regularize moments results.
+
+
 ## v2.1.0
 
 ### bug fixes

ngmix/moments.py

@@ -489,3 +489,69 @@ def make_mom_result(sums, sums_cov, sums_norm=None):
     res["T_flagstr"] = ngmix.flags.get_flags_str(res["T_flags"])
 
     return res
+
+
+def regularize_mom_shapes(res, fwhm_reg):
+    """Apply regularization to the shapes computed from moments sums.
+
+    This routine transforms the shapes as
+
+        e_{1,2} = M_{1,2}/(T + T_reg)
+
+    where T_reg is the T value equivalent to fwhm_reg, T is the original T value
+    from the moments, and M_{1,2} are the moments for shapes e_{1,2}.
+
+    This form of regularization is equivalent, for Gaussians, to convolving the Gaussian
+    with an isotropic Gaussian smoothing kernel of size fwhm_reg.
+
+    Parameters
+    ----------
+    res : dict
+        The original moments result before regularization.
+    fwhm_reg : float
+        The regularization FWHM value. Typically this should be of order the size of
+        the PSF for a prePSF moment.
+
+    Returns
+    -------
+    res_reg : dict
+        The regularized moments result. The size and flux are unchanged.
+    """
+    if fwhm_reg > 0:
+        raw_mom = res["sums"]
+        raw_mom_cov = res["sums_cov"]
+
+        T_reg = fwhm_to_T(fwhm_reg)
+
+        # the moments are not normalized and are sums, so convert T_reg to a sum using
+        # the flux sum first via T_reg -> T_reg * raw_mom[5]
+        amat = np.eye(6)
+        amat[4, 5] = T_reg
+
+        # the pre-PSF fitters do not fill out the centroid moments so hack around that
+        raw_mom_orig = raw_mom.copy()
+        if np.isnan(raw_mom_orig[0]):
+            raw_mom[0] = 0
+        if np.isnan(raw_mom_orig[1]):
+            raw_mom[1] = 0
+        reg_mom = np.dot(amat, raw_mom)
+        if np.isnan(raw_mom_orig[0]):
+            raw_mom[0] = np.nan
+            reg_mom[0] = np.nan
+        if np.isnan(raw_mom_orig[1]):
+            raw_mom[1] = np.nan
+            reg_mom[1] = np.nan
+
+        reg_mom_cov = np.dot(amat, np.dot(raw_mom_cov, amat.T))
+        momres = make_mom_result(reg_mom, reg_mom_cov)
+
+        # use old T
+        for col in ["T", "T_err", "T_flags", "T_flagstr"]:
+            momres[col] = res[col]
+
+        momres["flags"] |= res["flags"]
+        momres["flagstr"] = ngmix.flags.get_flags_str(momres["flags"])
+
+        return momres
+    else:
+        return res

ngmix/tests/test_moments.py

@@ -1,6 +1,10 @@
 import numpy as np
 
 import pytest
+import galsim
+from ..gaussmom import GaussMom
+from ..jacobian import Jacobian
+from ..observation import Observation
 
 from ..gexceptions import GMixRangeError
 
@@ -22,6 +26,7 @@
     mom2g,
     e2mom,
     g2mom,
+    regularize_mom_shapes,
 )
 
 
@@ -84,3 +89,70 @@ def test_moments_moms_to_e1e2_raises():
 
     with pytest.raises(GMixRangeError):
         moms_to_e1e2(np.array([0.1]), np.array([0.2]), np.array([-0.1]))
+
+
+@pytest.mark.parametrize("fwhm_reg", [0, 0.8])
+@pytest.mark.parametrize("has_nan", [True, False])
+def test_regularize_mom_shapes(fwhm_reg, has_nan):
+    rng = np.random.RandomState(seed=100)
+
+    fwhm = 0.9
+    image_size = 107
+    cen = (image_size - 1)/2
+    gs_wcs = galsim.ShearWCS(
+        0.125, galsim.Shear(g1=0, g2=0)).jacobian()
+
+    obj = galsim.Gaussian(
+        fwhm=fwhm
+    ).shear(
+        g1=-0.1, g2=0.3
+    ).withFlux(
+        400)
+    im = obj.drawImage(
+        nx=image_size,
+        ny=image_size,
+        wcs=gs_wcs,
+        method='no_pixel').array
+    noise = np.sqrt(np.sum(im**2)) / 1e2
+    wgt = np.ones_like(im) / noise**2
+
+    fitter = GaussMom(fwhm=1.2)
+
+    # get true flux
+    jac = Jacobian(
+        y=cen, x=cen,
+        dudx=gs_wcs.dudx, dudy=gs_wcs.dudy,
+        dvdx=gs_wcs.dvdx, dvdy=gs_wcs.dvdy)
+    obs = Observation(
+        image=im,
+        jacobian=jac,
+        weight=wgt,
+    )
+    res = fitter.go(obs=obs)
+
+    if has_nan:
+        res["sums"][0] = np.nan
+        res["sums"][1] = np.nan
+
+    res_reg = regularize_mom_shapes(res, fwhm_reg)
+
+    if has_nan:
+        assert np.isnan(res_reg["sums"][0])
+        assert np.isnan(res_reg["sums"][1])
+    assert np.all(np.isfinite(res_reg["sums"][2:]))
+
+    T_reg = fwhm_to_T(fwhm_reg)
+
+    if not has_nan:
+        assert np.allclose(res["sums"][[0, 1]], res_reg["sums"][[0, 1]])
+    assert np.allclose(res["sums"][4] + T_reg * res["sums"][5], res_reg["sums"][4])
+    if fwhm_reg > 0:
+        assert not np.allclose(res["sums"][4], res_reg["sums"][4])
+    assert np.allclose(res["sums"][[2, 3, 5]], res_reg["sums"][[2, 3, 5]])
+    for col in ["flux", "flux_err", "T", "T_err", "T_flags", "s2n"]:
+        assert np.allclose(res[col], res_reg[col])
+    for col in ["e1", "e2", "e", "e_err", "e_cov"]:
+        if fwhm_reg > 0:
+            assert not np.allclose(res[col], res_reg[col])
+        else:
+            assert np.allclose(res[col], res_reg[col])