Add optional higher order moments

CHANGES.md

@@ -1,5 +1,17 @@
 ## v2.3.2 (not yet released)
 
+
+## New features
+
+    - Added optional calculation of higher order moments when calling
+        gmix.get_weighted_moments(obs, with_higher_order=True)
+      This is also available with GaussMom(fwhm, with_higher_order=True)
+
+      Currently the sums are calculated and are stored in a larger "sums" and
+      "sums_cov" arrays in the results.  No normalized moments are returned.
+
+      Names and indices for moments can be found in moments.MOMENTS_NAME_MAP
+
 ## Bug Fixes
 
     - Fixed bug when moments are used in guesser and size is bad.

ngmix/gaussmom.py

@@ -12,12 +12,17 @@ class GaussMom(object):
     ----------
     fwhm: float
         The FWHM of the Gaussian weight function.
+    with_higher_order: bool, optional
+        If set to True, return higher order moments in the sums/sums_cov
+        arrays.  See ngmix.moments.MOMENTS_NAME_MAP for a map between
+        name and index.
     """
 
     kind = "wmom"
 
-    def __init__(self, fwhm):
+    def __init__(self, fwhm, with_higher_order=False):
         self.fwhm = fwhm
+        self.with_higher_order = with_higher_order
         self._set_mompars()
 
     def go(self, obs):
@@ -46,7 +51,9 @@ def _measure_moments(self, obs):
         measure weighted moments
         """
 
-        res = self.weight.get_weighted_moments(obs=obs, maxrad=1.e9)
+        res = self.weight.get_weighted_moments(
+            obs=obs, with_higher_order=self.with_higher_order,
+        )
 
         if res['flags'] != 0:
             return res

ngmix/gmix/gmix.py

@@ -638,7 +638,7 @@ def fill_fdiff(self, obs, fdiff, start=0):
             gm, obs._pixels, fdiff, start,
         )
 
-    def get_weighted_moments(self, obs, maxrad):
+    def get_weighted_moments(self, obs, maxrad=None, with_higher_order=False):
         """
         Get weighted moments using this mixture as the weight, including
         e1,e2,T,s2n etc.  If you just want the raw moments use
@@ -647,21 +647,30 @@ def get_weighted_moments(self, obs, maxrad):
         If you want the expected fluxes, you should set the flux to the inverse
         of the normalization which is 2*pi*sqrt(det)
 
-        parameters
+        Parameters
         ----------
         obs: Observation
             The Observation to compare with. See ngmix.observation.Observation
             The Observation must have a weight map set
+        maxrad: float, optional
+            If sent, limit moments to within the specified maximum radius
+        with_higher_order: bool, optional
+            If set to True, return higher order moments in the sums/sums_cov
+            arrays.  See ngmix.moments.MOMENTS_NAME_MAP for a map between
+            name and index.
 
-        returns:
-            result array with basic sums as well as summary statistics
-            such as e1,e2,T,s2n etc.
+        Returns
+        -------
+        result array with basic sums as well as summary statistics
+        such as e1,e2,T,s2n etc.
         """
 
-        res = self.get_weighted_sums(obs, maxrad)
+        res = self.get_weighted_sums(
+            obs, maxrad=maxrad, with_higher_order=with_higher_order,
+        )
         return get_weighted_moments_stats(res)
 
-    def get_weighted_sums(self, obs, maxrad, res=None):
+    def get_weighted_sums(self, obs, maxrad=None, with_higher_order=False, res=None):
         """
         Get weighted moments using this mixture as the weight.  To
         get more summary statistics use get_weighted_moments or
@@ -672,25 +681,43 @@ def get_weighted_sums(self, obs, maxrad, res=None):
         obs: Observation
             The Observation to compare with. See ngmix.observation.Observation
             The Observation must have a weight map set
+        maxrad: float, optional
+            If sent, limit moments to within the specified maximum radius
+        with_higher_order: bool, optional
+            If set to True, return higher order moments in the sums/sums_cov
+            arrays.  See ngmix.moments.MOMENTS_NAME_MAP for a map between
+            name and index.
         res: result array, optional
             If sent, sums will be added to the array rather than making
             a new one
+
+        Returns
+        -------
+        result array with sums
         """
         from . import gmix_nb
 
         self.set_norms_if_needed()
 
+        if maxrad is None:
+            T = self.get_T()
+            sigma = np.sqrt(T/2)
+            maxrad = 100 * sigma
+
         if res is None:
-            dt = np.dtype(_moments_result_dtype, align=True)
+            dt0 = get_moments_result_dtype(with_higher_order=with_higher_order)
+            dt = np.dtype(dt0, align=True)
             resarray = np.zeros(1, dtype=dt)
             res = resarray[0]
 
         wt_gm = self.get_data()
 
         # this will add to the sums
-        gmix_nb.get_weighted_sums(
-            wt_gm, obs.pixels, res, maxrad,
-        )
+        if with_higher_order:
+            gmix_nb.get_higher_order_weighted_sums(wt_gm, obs.pixels, res, maxrad)
+        else:
+            gmix_nb.get_weighted_sums(wt_gm, obs.pixels, res, maxrad)
+
         return res
 
     def get_model_s2n_sum(self, obs):
@@ -1251,14 +1278,20 @@ def get_weighted_moments_stats(ares):
     return res
 
 
-_moments_result_dtype = [
-    ('flags', 'i4'),
-    ('npix', 'i4'),
-    ('wsum', 'f8'),
-
-    ('sums', 'f8', 6),
-    ('sums_cov', 'f8', (6, 6)),
-    ('pars', 'f8', 6),
-    # temporary
-    ('F', 'f8', 6),
-]
+def get_moments_result_dtype(with_higher_order=False):
+    if with_higher_order:
+        nmom = 17
+    else:
+        nmom = 6
+
+    return [
+        ('flags', 'i4'),
+        ('npix', 'i4'),
+        ('wsum', 'f8'),
+
+        ('sums', 'f8', nmom),
+        ('sums_cov', 'f8', (nmom, nmom)),
+        ('pars', 'f8', nmom),
+        # temporary
+        ('F', 'f8', nmom),
+    ]

ngmix/gmix/gmix_nb.py

@@ -725,6 +725,93 @@ def get_weighted_sums(wt, pixels, res, maxrad):
                     res["sums_cov"][i, j] += w2 * var * F[i] * F[j]
 
 
+@njit
+def get_higher_order_weighted_sums(wt, pixels, res, maxrad):
+    """
+    Do sums for calculating the weighted moments.
+
+    Parameters
+    ----------
+    wt: array
+        The gaussian mixture with dtype ngmix.gmix.gmix._gauss2d_dtype
+    pixels: array
+        Array of pixels
+    res: array
+        The result array
+    maxrad: float
+        Maximum radius in u, v coordinates
+    """
+
+    maxrad2 = maxrad ** 2
+
+    vcen = wt["row"][0]
+    ucen = wt["col"][0]
+
+    F = res["F"]
+    nmom = F.size
+
+    n_pixels = pixels.size
+    for i_pixel in range(n_pixels):
+
+        pixel = pixels[i_pixel]
+
+        # v and u are really dv and du
+        v = pixel["v"] - vcen
+        u = pixel["u"] - ucen
+
+        r2 = u * u + v * v
+        if r2 < maxrad2:
+
+            weight = gmix_eval_pixel(wt, pixel)
+            var = 1.0 / (pixel["ierr"] * pixel["ierr"])
+
+            wdata = weight * pixel["val"]
+            w2 = weight * weight
+
+            u2 = u ** 2
+            v2 = v ** 2
+            vu = v * u
+            u4 = u2 ** 2
+            v4 = v2 ** 2
+
+            r4 = r2 * r2
+            r6 = r4 * r2
+            r8 = r6 * r2
+
+            # the order for first 6 matches get_weighted_sums
+            F[0] = pixel["v"]
+            F[1] = pixel["u"]
+            F[2] = u2 - v2
+            F[3] = 2 * vu
+            F[4] = r2
+            F[5] = 1.0
+
+            # third
+            F[6] = u * r2
+            F[7] = v * r2
+            F[8] = u * (u2 - 3 * v2)
+            F[9] = v * (3 * u2 - v2)
+
+            # fourth
+            F[10] = r4
+            F[11] = r2 * (u2 - v2)
+            F[12] = r2 * 2 * u * v
+            F[13] = u4 - 6 * u2 * v2 + v4
+            F[14] = (u2 - v2) * 4 * u * v
+
+            # Additional pure radial momentes 6, 8
+            F[15] = r6
+            F[16] = r8
+
+            res["wsum"] += weight
+            res["npix"] += 1
+
+            for i in range(nmom):
+                res["sums"][i] += wdata * F[i]
+                for j in range(nmom):
+                    res["sums_cov"][i, j] += w2 * var * F[i] * F[j]
+
+
 @njit
 def get_loglike(gmix, pixels):
     """

ngmix/moments.py

@@ -6,12 +6,58 @@
 from .util import get_ratio_error
 
 MOMENTS_NAME_MAP = {
+    # v
     "Mv": 0,
+    # u
     "Mu": 1,
+    # u^2 - v^2
     "M1": 2,
+    # v * u
     "M2": 3,
+    # v^2 + u^2
     "MT": 4,
+    # 1 (flux)
     "MF": 5,
+
+    # notation from piff.util.calculate_moments
+    # third order
+
+    # these are same as above but with the alternative notation
+    "M00": 5,  # same as MF
+    "M10": 1,  # same as Mu
+    "M01": 0,  # same as Mv
+    "M11": 4,  # same as MT
+    "M20": 2,  # same as M1
+    "M02": 3,  # same as M2
+
+    # u * r^2
+    "M21": 6,
+    # v * r^2
+    "M12": 7,
+    # u * (u^2 - 3 * v^2)
+    "M30": 8,
+    # v * (3 * u^2 - v^2)
+    "M03": 9,
+
+    # fourth order
+    # r^4
+    "M22": 10,
+    # r^2 * (u^ - v^)
+    "M31": 11,
+    # r^2 * 2 * u * v
+    "M13": 12,
+    # u^4 - 6 * u^2 * v^2 + v^4
+    "M40": 13,
+    # (u^2 - v^2) * 4 * u * v
+    "M14": 14,
+
+    # 6th order
+    # r^6
+    "M33": 15,
+
+    # 8th order
+    # r^8
+    "M44": 16,
 }
 
 
@@ -368,14 +414,14 @@ def make_mom_result(sums, sums_cov, sums_norm=None):
         A dictionary of results.
     """
     # for safety...
-    if len(sums) != 6:
+    if len(sums) != 6 and len(sums) != 17:
         raise ValueError(
-            "You must pass exactly 6 unnormalized moments in the order "
-            "[Mv, Mu, M1, M2, MT, MF] for ngmix.moments.make_mom_result."
+            "You must pass exactly 6 or 17 unnormalized moments in the order "
+            "[Mv, Mu, M1, M2, MT, MF, ...] for ngmix.moments.make_mom_result."
         )
-    if sums_cov.shape != (6, 6):
+    if sums_cov.shape != (6, 6) and sums_cov.shape != (17, 17):
         raise ValueError(
-            "You must pass a 6x6 matrix for ngmix.moments.make_mom_result."
+            "You must pass a 6x6 or 17x17 matrix for ngmix.moments.make_mom_result."
         )
 
     mv_ind = MOMENTS_NAME_MAP["Mv"]