Comp E/B mode using TreeCorr

requirements.txt

@@ -2,8 +2,7 @@ setuptools>=65.5.1
 numpy>=1.16.5
 scipy
 matplotlib
-
-iminuit>2  # iminuit changed API in v2.0.  We use the 2.x API.
-treecorr>=4.2
+iminuit>2
+treecorr>=5.0
 fitsio>=0.9.12
 scikit-learn>=0.18

tests/test_hyp_search.py

@@ -134,6 +134,7 @@ def test_hyperparameter_search_2d():
         if opt == "anisotropic":
             try:
                 e_mode, b_mode, logr = treegp.comp_eb(x[:, 0], x[:, 1], y, y)
+                e_mode, b_mode, logr = treegp.comp_eb_treecorr(x[:, 0], x[:, 1], y, y)
             except:
                 raise ValueError("Failed to compute E/B decomposition")
         # test if found hyperparameters are close the true hyperparameters.

treegp/__init__.py

@@ -18,7 +18,7 @@
 
 from .meanify import meanify
 
-from .utils import comp_eb
+from .utils import comp_eb, comp_eb_treecorr
 
 __all__ = [
     "__version__",
@@ -32,4 +32,5 @@
     "eval_kernel",
     "meanify",
     "comp_eb",
+    "comp_eb_treecorr",
 ]

treegp/utils.py

@@ -1,7 +1,5 @@
 import numpy as np
-
-# TO DO: E/B mode computation using numpy should use
-# treecorr instead.
+import treecorr
 
 
 def vcorr(x, y, dx, dy, rmin=5.0 / 3600.0, rmax=1.5, dlogr=0.05, maxpts=30000):
@@ -31,13 +29,11 @@ def vcorr(x, y, dx, dy, rmin=5.0 / 3600.0, rmax=1.5, dlogr=0.05, maxpts=30000):
     if len(x) > maxpts:
         # Subsample array to get desired number of points
         rate = float(maxpts) / len(x)
-        print("Subsampling rate {:5.3f}%".format(rate * 100.0))
         use = np.random.random(len(x)) <= rate
         x = x[use]
         y = y[use]
         dx = dx[use]
         dy = dy[use]
-    print("Length ", len(x))
     # Get index arrays that make all unique pairs
     i1, i2 = np.triu_indices(len(x))
     # Omit self-pairs
@@ -109,3 +105,51 @@ def comp_eb(u, v, du, dv, **kwargs):
     xib = xiB(logr, xiplus, ximinus)
     xie = xiplus - xib
     return xie, xib, logr
+
+
+class compEbTreecorr:
+
+    def __init__(self, x, y, dx, dy, rmin=5.0 / 3600.0, rmax=1.5, dlogr=0.05):
+
+        self.catalog = treecorr.Catalog(x=x, y=y, v1=dx, v2=dy)
+        self.vvCorr = treecorr.VVCorrelation(min_sep=rmin, max_sep=rmax, bin_size=dlogr)
+
+    def vcorr(self):
+
+        self.vvCorr.process(self.catalog)
+
+    def xiB(self, logr, xiplus, ximinus):
+        """
+        Return estimate of pure B-mode correlation function
+        """
+        # Integral of d(log r) ximinus(r) from r to infty:
+        dlogr = np.zeros_like(logr)
+        dlogr[1:-1] = 0.5 * (logr[2:] - logr[:-2])
+        tmp = np.array(ximinus) * dlogr
+        integral = np.cumsum(tmp[::-1])[::-1]
+        return 0.5 * (xiplus - ximinus) + integral
+
+    def comp_eb(self):
+        self.vcorr()
+        xib = self.xiB(self.vvCorr.logr, self.vvCorr.xip, self.vvCorr.xim)
+        xie = self.vvCorr.xip - xib
+        return xie, xib, self.vvCorr.logr
+
+
+def comp_eb_treecorr(u, v, du, dv, **kwargs):
+    """
+    Compute E/B decomposition of astrometric error correlation function
+
+    Parameters
+    ----------
+    u, v : array_like. positions of objects.
+    du, dv : array_like. astrometric shift.
+
+    returns
+    -------
+    xie, xib, logr : array_like. E-mode, B-mode,
+    and log of binned distance separation in 2-point correlation function.
+    """
+    cebt = compEbTreecorr(u, v, du, dv, **kwargs)
+    xie, xib, logr = cebt.comp_eb()
+    return xie, xib, logr