diff --git a/requirements.txt b/requirements.txt index 1de92f3..abd017b 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/test_hyp_search.py b/tests/test_hyp_search.py index 0419258..e3e7e86 100644 --- a/tests/test_hyp_search.py +++ b/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. diff --git a/treegp/__init__.py b/treegp/__init__.py index 5ac47d3..189d2cf 100644 --- a/treegp/__init__.py +++ b/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", ] diff --git a/treegp/utils.py b/treegp/utils.py index 04e9826..1253832 100644 --- a/treegp/utils.py +++ b/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