How to use TreeCorr - 10 common examples

hdu=-1)
    assert cat8 == cat7  # -1 is allowed and means the last one.

    cat9 = treecorr.Catalog(fname, allow_xyz=True,
                            x_col='x', y_col='y', z_col='z',
                            ra_col='ra', dec_col='dec', r_col='r',
                            ra_units='rad', dec_units='rad',
                            w_col='w', wpos_col='wpos', flag_col='flag',
                            k_col='k', g1_col='g1', g2_col='g2',
                            x_hdu=1, y_hdu=1, z_hdu=1,
                            ra_hdu=2, dec_hdu=1, r_hdu=2,
                            w_hdu=1, wpos_hdu=2, flag_hdu=1,
                            k_hdu=1, g1_hdu=1, g2_hdu=2)
    assert cat9 == cat1

    cat10 = treecorr.Catalog(fname, allow_xyz=True,
                             x_col='x', y_col='y', z_col='z',
                             ra_col='ra', dec_col='dec', r_col='r',
                             ra_units='rad', dec_units='rad',
                             w_col='w', wpos_col='wpos', flag_col='flag',
                             k_col='k', g1_col='g1', g2_col='g2',
                             x_hdu=3, y_hdu=3, z_hdu=3,
                             ra_hdu=4, dec_hdu=4, r_hdu=4,
                             w_hdu=5, wpos_hdu=5, flag_hdu=5,
                             k_hdu=6, g1_hdu=6, g2_hdu=6)
    assert cat10 == cat1

    # Not all columns in given hdu
    with assert_raises(ValueError):
        treecorr.Catalog(fname, allow_xyz=True,
                         x_col='x', y_col='y', z_col='z',
                         ra_col='ra', dec_col='dec', r_col='r',

# We aren't intentionally making a constant term, but there will be some C signal due to
        # the finite number of pairs being rendered.  So let's figure out how much there is.
        gC = gQ * np.conj(g)
        np.add.at(true_gCr, index[mask], gC[mask].real)
        np.add.at(true_gCi, index[mask], -gC[mask].imag)

    true_gQ /= true_npairs
    true_gCr /= true_npairs
    true_gCi /= true_npairs
    print('true_gQ = ',true_gQ)
    print('true_gCr = ',true_gCr)
    print('true_gCi = ',true_gCi)

    # Start with brute force.
    # With only 100 lenses, this still runs very fast.
    lens_cat = treecorr.Catalog(x=xl, y=yl, z=zl, g1=gl.real, g2=gl.imag)
    source_cat = treecorr.Catalog(x=xs, y=ys, z=zs, g1=g1, g2=g2)
    gg0 = treecorr.GGCorrelation(bin_size=bin_size, min_sep=min_sep, max_sep=max_sep, verbose=1,
                                 metric='Rlens', brute=True)
    t0 = time.time()
    gg0.process(lens_cat, source_cat)
    t1 = time.time()

    Rlens = gg0.meanr
    theory_gQ = gamma0 * np.exp(-0.5*Rlens**2/R0**2)

    print('Results with brute force:')
    print('time = ',t1-t0)
    print('gg.npairs = ',gg0.npairs)
    print('true_npairs = ',true_npairs)
    np.testing.assert_array_equal(gg0.npairs, true_npairs)
    print('gg.xim = ',gg0.xim)

patch_centers=cat7.patch_centers)
    np.testing.assert_array_equal(cat8.patch, cat8b.patch)
    np.testing.assert_array_equal(cat8.patch_centers, cat8b.patch_centers)

    # Check flat
    cat9 = treecorr.Catalog(x=x, y=y, npatch=npatch)
    cen_file2 = os.path.join('output','test_cat_centers.txt')
    cat9.write_patch_centers(cen_file2)
    centers9 = cat9.read_patch_centers(cen_file2)
    np.testing.assert_allclose(centers9, cat9.patch_centers)

    cat10 = treecorr.Catalog(x=x, y=y, patch_centers=cen_file2)
    np.testing.assert_array_equal(cat10.patch, cat9.patch)
    np.testing.assert_array_equal(cat10.patch_centers, cat9.patch_centers)

    cat11 = treecorr.Catalog(x=x, y=y, patch=cat9.patch)
    cat12 = treecorr.Catalog(x=x, y=y, patch_centers=cat11.patch_centers)
    print('n diff = ',np.sum(cat12.patch != cat11.patch))
    assert np.sum(cat12.patch != cat11.patch) < 10

    cat13 = treecorr.Catalog(x=x, y=y, w=w, patch=cat9.patch)
    cat14 = treecorr.Catalog(x=x, y=y, w=w, patch_centers=cat13.patch_centers)
    print('n diff = ',np.sum(cat14.patch != cat13.patch))
    assert np.sum(cat14.patch != cat13.patch) < 200
    np.testing.assert_array_equal(cat14.patch_centers, cat13.patch_centers)

    # The patch centers from the patch sub-catalogs should match.
    cen13 = [c.patch_centers[0] for c in cat13.patches]
    np.testing.assert_array_equal(cen13, cat13.patch_centers)

    # Using the full patch centers, you can also just load a single patch.
    for i in range(npatch):

x = (rng.random_sample(ngal)-0.5) * L
    y = (rng.random_sample(ngal)-0.5) * L
    r2 = (x**2 + y**2)/s**2
    kappa = A * np.exp(-r2/2.)

    min_sep = 11.
    max_sep = 15.
    nbins = 3
    min_u = 0.7
    max_u = 1.0
    nubins = 3
    min_v = 0.1
    max_v = 0.3
    nvbins = 2

    cat = treecorr.Catalog(x=x, y=y, k=kappa, x_units='arcmin', y_units='arcmin')
    kkk = treecorr.KKKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins,
                                  min_u=min_u, max_u=max_u, min_v=min_v, max_v=max_v,
                                  nubins=nubins, nvbins=nvbins,
                                  sep_units='arcmin', verbose=1)
    kkk.process(cat)

    # log() != , but it should be close:
    print('meanlogd1 - log(meand1) = ',kkk.meanlogd1 - np.log(kkk.meand1))
    print('meanlogd2 - log(meand2) = ',kkk.meanlogd2 - np.log(kkk.meand2))
    print('meanlogd3 - log(meand3) = ',kkk.meanlogd3 - np.log(kkk.meand3))
    print('meand3 / meand2 = ',kkk.meand3 / kkk.meand2)
    print('meanu = ',kkk.meanu)
    print('max diff = ',np.max(np.abs(kkk.meand3/kkk.meand2 -kkk.meanu)))
    print('max rel diff = ',np.max(np.abs((kkk.meand3/kkk.meand2 -kkk.meanu)/kkk.meanu)))
    print('(meand1 - meand2)/meand3 = ',(kkk.meand1-kkk.meand2) / kkk.meand3)
    print('meanv = ',kkk.meanv)

s1 = hp.heap().size if hp else 0
    print('NG1: ',s1, t1-t0, s1-s0)
    gk_cat1.unload()
    ng2 = treecorr.NGCorrelation(bin_size=0.5, min_sep=1., max_sep=30., sep_units='arcmin')
    t0 = time.time()
    s0 = hp.heap().size if hp else 0
    ng2.process(gk_cat2, gk_cat2, low_mem=True)
    t1 = time.time()
    s1 = hp.heap().size if hp else 0
    print('NG2: ',s1, t1-t0, s1-s0)
    gk_cat2.unload()
    np.testing.assert_allclose(ng1.xi, ng2.xi)
    np.testing.assert_allclose(ng1.weight, ng2.weight)

    # KK with r_col now to test that that works properly.
    gk_cat1 = treecorr.Catalog(file_name,
                               ra_col='ra', dec_col='dec', ra_units='rad', dec_units='rad',
                               r_col='r', g1_col='g1', g2_col='g2', k_col='k',
                               patch_centers=patch_centers)
    gk_cat2 = treecorr.Catalog(file_name,
                               ra_col='ra', dec_col='dec', ra_units='rad', dec_units='rad',
                               r_col='r', g1_col='g1', g2_col='g2', k_col='k',
                               patch_centers=patch_centers, save_patch_dir='output')

    kk1 = treecorr.KKCorrelation(bin_size=0.5, min_sep=1., max_sep=20.)
    t0 = time.time()
    s0 = hp.heap().size if hp else 0
    kk1.process(gk_cat1)
    t1 = time.time()
    s1 = hp.heap().size if hp else 0
    print('KK1: ',s1, t1-t0, s1-s0)
    gk_cat1.unload()

def test_direct_spherical():
    # Repeat in spherical coords

    ngal = 50
    s = 10.
    rng = np.random.RandomState(8675309)
    x = rng.normal(0,s, (ngal,) )
    y = rng.normal(0,s, (ngal,) ) + 200  # Put everything at large y, so small angle on sky
    z = rng.normal(0,s, (ngal,) )
    w = rng.random_sample(ngal)

    ra, dec = coord.CelestialCoord.xyz_to_radec(x,y,z)

    cat = treecorr.Catalog(ra=ra, dec=dec, ra_units='rad', dec_units='rad', w=w)

    min_sep = 1.
    bin_size = 0.2
    nrbins = 10
    nubins = 5
    nvbins = 5
    max_sep = min_sep * np.exp(nrbins * bin_size)
    ddd = treecorr.NNNCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
                                  sep_units='deg', brute=True)
    ddd.process(cat, num_threads=2)

    r = np.sqrt(x**2 + y**2 + z**2)
    x /= r;  y /= r;  z /= r
    north_pole = coord.CelestialCoord(0*coord.radians, 90*coord.degrees)

    true_ntri = np.zeros((nrbins, nubins, 2*nvbins), dtype=int)

np.testing.assert_allclose(cat4.r[use], cat1.r)
    np.testing.assert_allclose(cat4.x[use], cat1.x)
    np.testing.assert_allclose(cat4.y[use], cat1.y)
    np.testing.assert_allclose(cat4.z[use], cat1.z)

    cat5 = treecorr.Catalog(fname,
                            x_col='x', y_col='y', z_col='z',
                            w_col='w', wpos_col='wpos', flag_col='flag',
                            hdu=5)
    np.testing.assert_array_equal(cat5.x, cat1.x)
    np.testing.assert_array_equal(cat5.y, cat1.y)
    np.testing.assert_array_equal(cat5.z, cat1.z)
    np.testing.assert_array_equal(cat5.w, cat1.w)
    np.testing.assert_array_equal(cat5.wpos, cat1.wpos)

    cat6 = treecorr.Catalog(fname,
                            x_col='x', y_col='y', z_col='z',
                            k_col='k', g1_col='g1', g2_col='g2',
                            hdu=6)
    np.testing.assert_array_equal(cat6.x[use], cat1.x)
    np.testing.assert_array_equal(cat6.y[use], cat1.y)
    np.testing.assert_array_equal(cat6.z[use], cat1.z)
    np.testing.assert_array_equal(cat6.k[use], cat1.k)
    np.testing.assert_array_equal(cat6.g1[use], cat1.g1)
    np.testing.assert_array_equal(cat6.g2[use], cat1.g2)

    cat7 = treecorr.Catalog(fname, allow_xyz=True,
                            x_col='x', y_col='y', z_col='z',
                            ra_col='ra', dec_col='dec', r_col='r',
                            ra_units='rad', dec_units='rad',
                            w_col='w', wpos_col='wpos', flag_col='flag',
                            k_col='k', g1_col='g1', g2_col='g2',

x1 = rng.normal(0,s, (ngal,) )
    y1 = rng.normal(0,s, (ngal,) )
    w1 = rng.random_sample(ngal)
    k1 = rng.normal(5,1, (ngal,) )

    x2 = rng.normal(0,s, (ngal,) )
    y2 = rng.normal(0,s, (ngal,) )
    w2 = rng.random_sample(ngal)
    g12 = rng.normal(0,0.2, (ngal,) )
    g22 = rng.normal(0,0.2, (ngal,) )

    w1 = np.ones_like(w1)
    w2 = np.ones_like(w2)

    cat1 = treecorr.Catalog(x=x1, y=y1, w=w1, k=k1)
    cat2 = treecorr.Catalog(x=x2, y=y2, w=w2, g1=g12, g2=g22)

    min_sep = 5.
    max_sep = 50.
    nbins = 10
    bin_size = np.log(max_sep/min_sep) / nbins
    kg = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
    with assert_warns(FutureWarning):
        kg.process_pairwise(cat1, cat2)
    kg.finalize(cat1.vark, cat2.varg)

    true_npairs = np.zeros(nbins, dtype=int)
    true_weight = np.zeros(nbins, dtype=float)
    true_xi = np.zeros(nbins, dtype=complex)

    rsq = (x1-x2)**2 + (y1-y2)**2
    r = np.sqrt(rsq)

xl = (rng.random_sample(nlens)-0.5) * L
    yl = (rng.random_sample(nlens)-0.5) * L
    kl = rng.normal(0.23, 0.05, (nlens,) )
    xs = (rng.random_sample(nsource)-0.5) * L
    ys = (rng.random_sample(nsource)-0.5) * L
    g1 = np.zeros( (nsource,) )
    g2 = np.zeros( (nsource,) )
    for x,y,k in zip(xl,yl,kl):
        dx = xs-x
        dy = ys-y
        r2 = dx**2 + dy**2
        gammat = gamma0 * np.exp(-0.5*r2/r0**2) / k
        g1 += -gammat * (dx**2-dy**2)/r2
        g2 += -gammat * (2.*dx*dy)/r2

    lens_cat = treecorr.Catalog(x=xl, y=yl, k=kl, x_units='arcmin', y_units='arcmin')
    source_cat = treecorr.Catalog(x=xs, y=ys, g1=g1, g2=g2, x_units='arcmin', y_units='arcmin')
    kg = treecorr.KGCorrelation(bin_size=0.1, min_sep=1., max_sep=20., sep_units='arcmin',
                                verbose=1)
    kg.process(lens_cat, source_cat)

    r = kg.meanr
    true_gt = gamma0 * np.exp(-0.5*r**2/r0**2)

    print('kg.xi = ',kg.xi)
    print('kg.xi_im = ',kg.xi_im)
    print('true_gammat = ',true_gt)
    print('ratio = ',kg.xi / true_gt)
    print('diff = ',kg.xi - true_gt)
    print('max diff = ',max(abs(kg.xi - true_gt)))
    np.testing.assert_allclose(kg.xi, true_gt, rtol=0.1)
    np.testing.assert_allclose(kg.xi_im, 0., atol=1.e-2)

# The edge bins may differ slightly from the binning approximations (bin_slop and such),
    # but the differences should be very small.  (When Erika reported the problem, the differences
    # were a few percent, which ended up making a bit difference in the correlation function.)
    np.testing.assert_allclose(dd1.npairs, dd2.npairs[2:-2], rtol=1.e-6)

    if __name__ == '__main__':
        # If we're running from the command line, go ahead and finish the calculation
        # This catalog has 10^6 objects, which takes quite a while.  I should really investigate
        # how to speed up the Rperp distance calculation.  Probably by having a faster over-
        # and under-estimate first, and then only do the full calculation when it seems like we
        # will actually need it.
        # Anyway, until then, let's not take forever by using last_row=200000
        get_from_wiki('nn_perp_rand.dat')
        rcat = treecorr.Catalog('data/nn_perp_rand.dat', config, last_row=200000)

        rr1 = treecorr.NNCorrelation(config)
        rr1.process(rcat, metric='OldRperp')
        rr2 = treecorr.NNCorrelation(config, min_sep=lower_min_sep, nbins=more_nbins)
        rr2.process(rcat, metric='OldRperp')
        print('rr1 npairs = ',rr1.npairs)
        print('rr2 npairs = ',rr2.npairs[2:-2])
        np.testing.assert_allclose(rr1.npairs, rr2.npairs[2:-2], rtol=1.e-6)

        dr1 = treecorr.NNCorrelation(config)
        dr1.process(dcat, rcat, metric='OldRperp')
        dr2 = treecorr.NNCorrelation(config, min_sep=lower_min_sep, nbins=more_nbins)
        dr2.process(dcat, rcat, metric='OldRperp')
        print('dr1 npairs = ',dr1.npairs)
        print('dr2 npairs = ',dr2.npairs[2:-2])
        np.testing.assert_allclose(dr1.npairs, dr2.npairs[2:-2], rtol=1.e-6)

        xi1, varxi1 = dd1.calculateXi(rr1, dr1)