ENH(twopoint): function for binning two-point data

heracles/twopoint.py

@@ -373,65 +373,94 @@ def pixelate_mms_healpix(mms, nside, *, inplace=False):
     return out
 
 
-def binned_cls(cls, bins, *, weights=None, out=None):
-    """compute binned angular power spectra"""
+def bin2pt(arr, bins, name, *, weights=None):
+    """Compute binned two-point data."""
 
     def norm(a, b):
         """divide a by b if a is nonzero"""
-        return np.divide(a, b, where=(a != 0), out=np.zeros_like(a))
+        out = np.zeros(np.broadcast(a, b).shape)
+        return np.divide(a, b, where=(a != 0), out=out)
+
+    # flatten list of bins
+    bins = np.reshape(bins, -1)
+    m = bins.size
+
+    # shape of the data
+    n, *ds = np.shape(arr)
+    ell = np.arange(n)
+
+    # weights from string or given array
+    if weights is None:
+        w = np.ones(n)
+    elif isinstance(weights, str):
+        if weights == "l(l+1)":
+            w = ell * (ell + 1)
+        elif weights == "2l+1":
+            w = 2 * ell + 1
+        else:
+            msg = f"unknown weights string: {weights}"
+            raise ValueError(msg)
+    else:
+        w = np.asanyarray(weights)[:n]
+
+    # create the structured output array
+    # if input data is multi-dimensional, then so will the `name` column be
+    binned = np.empty(
+        m - 1,
+        [
+            ("L", float),
+            (name, float, ds) if ds else (name, float),
+            ("LMIN", float),
+            ("LMAX", float),
+            ("W", float),
+        ],
+    )
+
+    # get the bin index for each ell
+    i = np.digitize(ell, bins)
+
+    assert i.size == ell.size
+
+    # get the binned weights
+    wb = np.bincount(i, weights=w, minlength=m)[1:m]
+
+    # bin data in ell
+    binned["L"] = norm(np.bincount(i, w * ell, m)[1:m], wb)
+    for j in np.ndindex(*ds):
+        x = (slice(None), *j)
+        binned[name][x] = norm(np.bincount(i, w * arr[x], m)[1:m], wb)
+
+    # add bin edges
+    binned["LMIN"] = bins[:-1]
+    binned["LMAX"] = bins[1:]
+
+    # add weights
+    binned["W"] = wb
+
+    # all done
+    return binned
+
 
-    m = len(bins)
+def binned_cls(cls, bins, *, weights=None, out=None):
+    """compute binned angular power spectra"""
 
     if out is None:
         out = TocDict()
 
     for key, cl in cls.items():
-        ell = np.arange(len(cl))
-
-        if weights is None:
-            w = np.ones_like(cl)
-        elif isinstance(weights, str):
-            if weights == "l(l+1)":
-                w = ell * (ell + 1)
-            elif weights == "2l+1":
-                w = 2 * ell + 1
-            else:
-                msg = f"unknown weights string: {weights}"
-                raise ValueError(msg)
-        else:
-            w = weights[: len(cl)]
-
-        # create the output data
-        binned = np.empty(
-            m - 1,
-            [
-                ("L", float),
-                ("CL", float),
-                ("LMIN", float),
-                ("LMAX", float),
-                ("W", float),
-            ],
-        )
-
-        # get the bin index for each ell
-        i = np.digitize(ell, bins)
-
-        # get the binned weights
-        wb = np.bincount(i, weights=w, minlength=m)[1:m]
-
-        # bin everything
-        binned["L"] = norm(np.bincount(i, weights=w * ell, minlength=m)[1:m], wb)
-        binned["CL"] = norm(np.bincount(i, weights=w * cl, minlength=m)[1:m], wb)
-
-        # add bin edges
-        binned["LMIN"] = bins[:-1]
-        binned["LMAX"] = bins[1:]
-
-        # add weights
-        binned["W"] = wb
-
-        # store result
-        out[key] = binned
+        out[key] = bin2pt(cl, bins, "CL", weights=weights)
+
+    return out
+
+
+def binned_mms(mms, bins, *, weights=None, out=None):
+    """compute binned mixing matrices"""
+
+    if out is None:
+        out = TocDict()
+
+    for key, mm in mms.items():
+        out[key] = bin2pt(mm, bins, "MM", weights=weights)
 
     return out
 

tests/test_twopoint.py

@@ -221,6 +221,50 @@ def test_binned_cls(weights):
         np.testing.assert_array_almost_equal(result[key]["W"], binned_w)
 
 
+@pytest.mark.parametrize("weights", [None, "l(l+1)", "2l+1", "<rand>"])
+@pytest.mark.parametrize("ndim", [1, 2, 3])
+def test_bin2pt(ndim, weights):
+    from heracles.twopoint import bin2pt
+
+    data = np.random.randn(*[21, 31, 41][:ndim])
+
+    bins = [2, 5, 10, 15, 20]
+
+    if weights == "<rand>":
+        weights_ = np.random.rand(51)
+    else:
+        weights_ = weights
+
+    result = bin2pt(data, bins, "XY", weights=weights_)
+
+    ell = np.arange(len(data))
+
+    if weights is None:
+        w = np.ones_like(ell)
+    elif weights == "l(l+1)":
+        w = ell * (ell + 1)
+    elif weights == "2l+1":
+        w = 2 * ell + 1
+    else:
+        w = weights_[: len(ell)]
+
+    binned_ell = np.empty(len(bins) - 1)
+    binned_xy = np.empty((len(bins) - 1, *data.shape[1:]))
+    binned_w = np.empty(len(bins) - 1)
+    for i, (a, b) in enumerate(zip(bins[:-1], bins[1:])):
+        inbin = (a <= ell) & (ell < b)
+        binned_ell[i] = np.average(ell[inbin], weights=w[inbin])
+        for j in np.ndindex(*binned_xy.shape[1:]):
+            binned_xy[(i, *j)] = np.average(data[(inbin, *j)], weights=w[inbin])
+        binned_w[i] = w[inbin].sum()
+
+    np.testing.assert_array_almost_equal(result["L"], binned_ell)
+    np.testing.assert_array_almost_equal(result["XY"], binned_xy)
+    np.testing.assert_array_equal(result["LMIN"], bins[:-1])
+    np.testing.assert_array_equal(result["LMAX"], bins[1:])
+    np.testing.assert_array_almost_equal(result["W"], binned_w)
+
+
 @pytest.mark.parametrize("full", [False, True])
 def test_random_bias(full):
     from heracles.twopoint import random_bias