Corrfunc usage in reality

corrfunc_3dcf_2dcf.py

#%% imports
from Corrfunc.io import read_catalog
from Corrfunc.mocks.DDrppi_mocks import DDrppi_mocks
from Corrfunc.utils import convert_3d_counts_to_cf, convert_rp_pi_counts_to_wp
import numpy as  np
import matplotlib.pyplot as plt

#%% function
def calc_xi_and_wp(nd1: int, nd2: int,
                   nr1: int, nr2: int,
                   d1d2: np.ndarray, d1r2: np.ndarray,
                   d2r1: np.ndarray, r1r2: np.ndarray,
                   where: str = 'mid') -> Tuple[list, list]:
    """
    spatial_cf creates a spatial correlation function from pair counts using Landy-Szalay estimator. Spatial CF has a shape (number of rp bins)x(number of pi bins).  Spatial cf with indexing like [:,0] is a CF for a fixed pi between 0 and 1.
    It calculates rp and pi bins from the input d1d2 pair counts.
    It  also returns projected correlation function as an integral of spatial CF along Pi axis.


    rp and pi are from np.meshgrid function, so one may simply plot 2d cf as:
    rp, pi, cf = calc_LS(arguments)
    plot = plt.pcolormesh(rp, pi, cf, cmap='RdBu', shading='flat')
    plt.colorbar(plot)
    plt.contour(rp, pi, cf)


    Args:
        nd1 (int): number of data points in the first set
        nd2 (int): same as above for the second set
        nr1 (int): number of random points in the first set
        nr2 (int): number of random points in the second set
        d1d2 (np.ndarray): DDrppi_mocks array for cross-correlation between D1 and D2
        d1r2 (np.ndarray): same but for D1 R2
        d2r1 (np.ndarray): same but for D2 and R1
        r1r2 (np.ndarray): same but for R1 and R2
        where (str): defines what is an rp bin array for plotting. It may be left, right, mid for rmin, rmax, 0.5(rmin+rmax). Otherwise uses rpavg from d1d2 pair counts

    Returns:
        Tuple[list, list]: [rp (grid), pi (grid) and  spatial correlation function] , [rp and projected spatial correlation function].

    """

    piarr = np.unique(d1d2['pimax'])
    # assert linear 1 Mpc/h binning in pi direction
    assert np.all(np.diff(piarr) == np.ones(len(piarr)-1))

    if where == 'mid':
        rarr = (np.unique(d1d2['rmax']) + np.unique(d1d2['rmin'])) / 2
    elif where == 'left':
        rarr = np.unique(d1d2['rmin'])
    elif where == 'right':
        rarr = np.unique(d1d2['rmax'])
    else:
        print('rmax position is not recognized. Using ravg')
        rarr = np.unique(d1d2['rpavg'])

    grid_pi, grid_rp = np.meshgrid(piarr, rarr)

    xi = convert_3d_counts_to_cf(ND1=nd1, ND2=nd2,
                                 NR1=nr1, NR2=nr2,
                                 D1D2=d1d2, D1R2=d1r2,
                                 D2R1=d2r1, R1R2=r1r2,
                                 estimator='LS')

    xi_reshape = np.reshape(xi, (-1, int(piarr[-1])))

    assert xi_reshape.shape == (len(rarr), len(piarr)), 'shape mismatch!'

    # because the firs bin is from 0 to min(piarr)
    piarr = np.insert(piarr, 0, 0)
    dpi = np.diff(piarr)

    wp = 2*np.average(xi_reshape, axis=1, weights=dpi) * \
        np.sum(dpi)  # this is an integral over pi direction

    return [grid_rp, grid_pi, xi_reshape], [rarr, wp]


#%% test on mock from Corrfunc package
  
galaxy_catalog = "<corrfunc path here>/Corrfunc_mock/Mr19_mock_northonly.rdcz.ff"

RA, DEC, CZ = read_catalog(galaxy_catalog)
N = len(RA)

random_catalog = "<corrfunc path here>/Corrfunc_mock/Mr19_randoms_northonly.rdcz.ff"
rand_RA, rand_DEC, rand_CZ = read_catalog(random_catalog)
rand_N = len(rand_RA)

nbins = 25
bins = np.logspace(np.log10(0.01), np.log10(2), nbins + 1)
pimax = 10
cosmology = 1
nthreads = 4



DD_counts = DDrppi_mocks(1, cosmology, pimax=pimax, binfile=bins,
                         RA1=RA, DEC1=DEC, CZ1=CZ,
                         verbose=True,
                         output_rpavg=True,
                         nthreads=nthreads)
print('done DD')
autocorr = 0
DR_counts = DDrppi_mocks(autocorr, cosmology, nthreads, pimax, bins,
                         RA, DEC, CZ,
                         RA2=rand_RA, DEC2=rand_DEC, CZ2=rand_CZ, verbose=True)
print('done DR')

autocorr = 1
RR_counts = DDrppi_mocks(autocorr, cosmology, nthreads, pimax, bins,
                         rand_RA, rand_DEC, rand_CZ, verbose=True)
print('done RR')


# %% calc CFs and compare with Corrfunc

cf_corrfunc = convert_3d_counts_to_cf(N, N, rand_N, rand_N,
                                      DD_counts, DR_counts,
                                      DR_counts, RR_counts)

my_cf = calc_xi_and_wp(N, N, rand_N, rand_N,
                       DD_counts, DR_counts,
                       DR_counts, RR_counts)
a, b = my_cf
rp, pi, cf = a
rpp, wp = b
plt.figure()
plot = plt.pcolormesh(rp, pi, cf, cmap='RdBu', shading='flat')
plt.colorbar(plot)
plt.contour(rp, pi, cf)
plt.xscale('log')
plt.xlabel('rp, Mpc/h')
plt.ylabel('pi, Mpc/h')


plt.figure()
plt.loglog(rpp, wp)
plt.xlabel('rp')
plt.ylabel('wp(rp)')

wp_corrfunc = convert_rp_pi_counts_to_wp(N, N, rand_N, rand_N,
                                         DD_counts, DR_counts,
                                         DR_counts, RR_counts, nbins, pimax)

plt.plot(rpp, wp_corrfunc, 's-', label ='Corrfunc convert_rp_pi_counts_to_wp')


assert np.allclose(cf_corrfunc.reshape(len(rpp), pimax), cf)
assert np.allclose(wp_corrfunc, wp)