def comp_2pcf(self, X, y, y_err):
"""
Estimate 2-point correlation function using TreeCorr.
:param X: The independent covariates. (n_samples, 2)
:param y: The dependent responses. (n_samples, n_targets)
:param y_err: Error of y. (n_samples, n_targets)
"""
if np.sum(y_err) == 0:
w = None
else:
w = 1./y_err**2
if self.anisotropic:
cat = treecorr.Catalog(x=X[:,0], y=X[:,1], k=(y-np.mean(y)), w=w)
kk = treecorr.KKCorrelation(min_sep=self.min_sep, max_sep=self.max_sep, nbins=self.nbins,
bin_type='TwoD', bin_slop=0)
kk.process(cat)
# Need a mask in the case of the 2D correlation function, to compute
# the covariance matrix using the bootstrap. The 2D correlation
# function is symmetric, so just half of the correlation function
# is useful to compute the covariance matrix. If it is not done,
# the covariance matrix is consequently not positive definite.
npixels = len(kk.xi)**2
mask = np.ones_like(kk.xi, dtype=bool)
mask = mask.reshape((int(np.sqrt(npixels)), int(np.sqrt(npixels))))
boolean_mask_odd = (len(mask)%2 == 0)
even_or_odd = len(mask)%2
nmask = int((len(mask)/2) + even_or_odd)
mask[nmask:,:] = False
mask[nmask-1][nmask:] = boolean_mask_odd
mask = mask.reshape(npixels)
distance = np.array([kk.dx.reshape(npixels), kk.dy.reshape(npixels)]).T
Coord = distance
xi = kk.xi.reshape(npixels)
else:
cat = treecorr.Catalog(x=X[:,0], y=X[:,1], k=(y-np.mean(y)), w=w)
kk = treecorr.KKCorrelation(min_sep=self.min_sep, max_sep=self.max_sep, nbins=self.nbins)
kk.process(cat)
distance = kk.meanr
mask = np.ones_like(kk.xi, dtype=bool)
Coord = np.array([distance,np.zeros_like(distance)]).T
xi = kk.xi
return xi, distance, Coord, mask
def test_constant():
# A fairly trivial test is to use a constant value of kappa everywhere.
ngal = 100000
A = 0.05
L = 100.
numpy.random.seed(8675309)
x = (numpy.random.random_sample(ngal) - 0.5) * L
y = (numpy.random.random_sample(ngal) - 0.5) * L
kappa = A * numpy.ones(ngal)
cat = treecorr.Catalog(x=x,
y=y,
k=kappa,
x_units='arcmin',
y_units='arcmin')
kk = treecorr.KKCorrelation(bin_size=0.1,
min_sep=0.1,
max_sep=10.,
sep_units='arcmin')
kk.process(cat)
print('kk.xi = ', kk.xi)
numpy.testing.assert_almost_equal(kk.xi, A**2, decimal=10)
# Now add some noise to the values. It should still work, but at slightly lower accuracy.
kappa += 0.001 * (numpy.random.random_sample(ngal) - 0.5)
cat = treecorr.Catalog(x=x,
y=y,
k=kappa,
x_units='arcmin',
y_units='arcmin')
kk.process(cat)
print('kk.xi = ', kk.xi)
numpy.testing.assert_almost_equal(kk.xi, A**2, decimal=6)
def fastXi(x, y, dx, dy, rrange=(5. / 3600., 1.5), nbins=100):
"""
Use treecorr to get xi of (dx*dx + dy*dy) (xiplus).
Returns arrays of mean radius in each (log-spaced) bin, mean xi.
"""
catx = tc.Catalog(x=x, y=y, k=dx)
caty = tc.Catalog(x=x, y=y, k=dy)
kk = tc.KKCorrelation(min_sep=rrange[0], max_sep=rrange[1], nbins=nbins)
kk.process(catx)
xx = np.array(kk.xi)
kk = tc.KKCorrelation(min_sep=rrange[0], max_sep=rrange[1], nbins=nbins)
kk.process(caty)
yy = np.array(kk.xi)
return np.exp(kk.meanlogr), xx + yy
def calc_shear_shear(self,i,j,verbose,num_threads):
m1,m2,mask = self.get_m(i)
if self.params['has_sheared']:
cat_i = treecorr.Catalog(g1=self.shape['e1'][mask]/m1[mask], g2=self.shape['e2'][mask]/m2[mask], w=self.weight[mask], ra=self.shape['ra'][mask], dec=self.shape['dec'][mask], ra_units='deg', dec_units='deg')
else:
cat_i = treecorr.Catalog(g1=self.shape['e1'][mask], g2=self.shape['e2'][mask], w=self.weight[mask], ra=self.shape['ra'][mask], dec=self.shape['dec'][mask], ra_units='deg', dec_units='deg')
biascat_i = treecorr.Catalog(k=np.sqrt(self.shape['m1'][mask]*self.shape['m2'][mask]), w=self.weight[mask], ra=self.shape['ra'][mask], dec=self.shape['dec'][mask], ra_units='deg', dec_units='deg')
m1,m2,mask = self.get_m(j)
if self.params['has_sheared']:
cat_j = treecorr.Catalog(g1=self.shape['e1'][mask]/m1[mask], g2=self.shape['e2'][mask]/m2[mask], w=self.weight[mask], ra=self.shape['ra'][mask], dec=self.shape['dec'][mask], ra_units='deg', dec_units='deg')
else:
cat_j = treecorr.Catalog(g1=self.shape['e1'][mask], g2=self.shape['e2'][mask], w=self.weight[mask], ra=self.shape['ra'][mask], dec=self.shape['dec'][mask], ra_units='deg', dec_units='deg')
biascat_j = treecorr.Catalog(k=np.sqrt(self.shape['m1'][mask]*self.shape['m2'][mask]), w=self.weight[mask], ra=self.shape['ra'][mask], dec=self.shape['dec'][mask], ra_units='deg', dec_units='deg')
gg = treecorr.GGCorrelation(nbins=self.params['tbins'], min_sep=self.params['tbounds'][0], max_sep=self.params['tbounds'][1], sep_units='arcmin', bin_slop=self.params['slop'], verbose=verbose,num_threads=num_threads)
gg.process(cat_i,cat_j)
if self.params['has_sheared']:
norm = 1.
else:
kk = treecorr.KKCorrelation(nbins=self.params['tbins'], min_sep=self.params['tbounds'][0], max_sep=self.params['tbounds'][1], sep_units='arcmin', bin_slop=self.params['slop'], verbose=verbose,num_threads=num_threads)
kk.process(biascat_i,biascat_j)
norm = kk.xi
theta=np.exp(gg.meanlogr)
xip = gg.xip/norm
xim = gg.xim/norm
xiperr = ximerr = np.sqrt(gg.varxi)/norm
return theta, xip, xim, xiperr, ximerr
def compute_2pt_pix(self):
kk = treecorr.KKCorrelation(config=self.config_2pt)
toc = time.time()
kk.process(self.cat)
tic = time.time()
print '2PCF took', tic - toc
stdout.flush()
self.kk = kk
def plot_cc(sys_name):
sys_autocorr = treecorr.KKCorrelation(nbins=10, min_sep=10, max_sep=100, sep_units='arcmin')
sys_autocorr.read(data_dir+'/auto_correlations/autocc_'+sys_name+'.fits')
print(sys_autocorr.xi)
ngal_autocorr = treecorr.KKCorrelation(nbins=10, min_sep=10, max_sep=100, sep_units='arcmin')
ngal_autocorr.read(data_dir+'/auto_correlations/autocc_ngal_norm.fits')
kk_no_weights = treecorr.KKCorrelation(nbins=10, min_sep=10, max_sep=100, sep_units = 'arcmin')
kk_no_weights.read(data_dir+'/cross_correlations/cc_ngal_'+sys_name+'.fits')
print(kk_no_weights.xi)
kk_weights = treecorr.KKCorrelation(nbins=10, min_sep=10, max_sep=100, sep_units = 'arcmin')
kk_weights.read(data_dir+'/cross_correlations/cc_ngal_'+sys_name+'_weights.fits')
print(kk_weights.xi)
#print(kk_no_weights.varxi)
#print(np.sqrt(kk_no_weights.varxi))
#Try error propagation
# f = <dg sys> / sqrt(<sys sys> <dg dg>)
std_f_no_weights = error_propagation(kk_no_weights.xi, sys_autocorr.xi, ngal_autocorr.xi, kk_no_weights.varxi, sys_autocorr.varxi, ngal_autocorr.varxi)
std_f_weights = error_propagation(kk_weights.xi, sys_autocorr.xi, ngal_autocorr.xi, kk_weights.varxi, sys_autocorr.varxi, ngal_autocorr.varxi)
print(std_f_no_weights)
print(std_f_weights)
plt.errorbar(x=np.arange(10,100,9), y=kk_no_weights.xi/np.sqrt(sys_autocorr.xi*ngal_autocorr.xi), yerr=std_f_no_weights, label='Without weights')
plt.errorbar(x=np.arange(10,100,9), y=kk_weights.xi/np.sqrt(sys_autocorr.xi*ngal_autocorr.xi), yerr=std_f_weights, label = 'With weights')
plt.title('Cross-Correlation NGAL with {}'.format(sys_name))
plt.legend()
plt.savefig(graph_dir+'/cross_correlation/cc_'+sys_name)
plt.show()
def calculate_autocc(df, sys_name):
sys_cat = treecorr.Catalog(ra=df['RA'], dec=df['DEC'], ra_units='deg', dec_units='deg', k=df[sys_name])
autokk = treecorr.KKCorrelation(nbins=10, min_sep=10, max_sep=100, sep_units='arcmin')
print('Start calculating autocorrelation of {}'.format(sys_name))
start = time.time()
autokk.process(sys_cat, num_threads=80) # For auto-correlation.
autokk.write(data_dir+'/auto_correlations/autocc_'+sys_name+'.fits')
end = time.time()
print('Done!')
print('Time: {}'.format(end-start))
def test_varxi():
# Test that varxi is correct (or close) based on actual variance of many runs.
# Signal doesn't matter much. Use the one from test_gg.
kappa0 = 0.03
r0 = 10.
L = 50. * r0
rng = np.random.RandomState(8675309)
# Note: to get a good estimate of var(xi), you need a lot of runs. The number of
# runs matters much more than the number of galaxies for getting this to pass.
if __name__ == '__main__':
ngal = 1000
nruns = 50000
tol_factor = 1
else:
ngal = 100
nruns = 5000
tol_factor = 5
all_kks = []
for run in range(nruns):
# In addition to the shape noise below, there is shot noise from the random x,y positions.
x = (rng.random_sample(ngal) - 0.5) * L
y = (rng.random_sample(ngal) - 0.5) * L
# Varied weights are hard, but at least check that non-unit weights work correctly.
w = np.ones_like(x) * 5
r2 = (x**2 + y**2) / r0**2
k = kappa0 * np.exp(-r2 / 2.)
k += rng.normal(0, 0.1, size=ngal)
cat = treecorr.Catalog(x=x, y=y, w=w, k=k)
kk = treecorr.KKCorrelation(bin_size=0.1, min_sep=10., nbins=24)
kk.process(cat)
all_kks.append(kk)
mean_xi = np.mean([kk.xi for kk in all_kks], axis=0)
var_xi = np.var([kk.xi for kk in all_kks], axis=0)
mean_varxi = np.mean([kk.varxi for kk in all_kks], axis=0)
print('mean_xi = ', mean_xi)
print('mean_varxi = ', mean_varxi)
print('var_xi = ', var_xi)
print('ratio = ', var_xi / mean_varxi)
print('max relerr for xi = ', np.max(np.abs(
(var_xi - mean_varxi) / var_xi)))
np.testing.assert_allclose(mean_varxi, var_xi, rtol=0.02 * tol_factor)
def __init__(self, dataname, runname, sample_type, n_jackknife):
self.sample_type = sample_type #redmagicHD, dark_matter
self.dataname = dataname
self.runname = runname
config_fname = output_path + self.runname + '.config'
try:
with open(config_fname) as f:
self.configdict = yaml.load(f.read())
except IOError:
print 'config file ' + config_fname + ' not found.'
raise
self.kk = treecorr.KKCorrelation(config=self.configdict['2PCF'])
self.kkk = treecorr.KKKCorrelation(config=self.configdict['3PCF'])
self.n_jackknife = n_jackknife
self.zetas = []
self.weights = []
self.xis = []
def compute_xi(ra, dec, e1, e2, w, s):
# Apply the mean sensitivity in each case
#means = numpy.sum(w*s) / numpy.sum(w)
#e1c = e1 / means
#e2c = e2 / means
# Build a TreeCorr catalog
cat = treecorr.Catalog(ra=ra,
dec=dec,
g1=e1,
g2=e2,
k=s,
w=w,
ra_units='deg',
dec_units='deg')
# Compute the correlation function
gg = treecorr.GGCorrelation(bin_size=0.05,
min_sep=0.1,
max_sep=500,
sep_units='arcmin',
output_dots=True,
verbose=2)
ss = treecorr.KKCorrelation(bin_size=0.05,
min_sep=0.1,
max_sep=500,
sep_units='arcmin',
output_dots=True,
verbose=2)
print 'Start corr2'
gg.process(cat)
ss.process(cat)
gg.xip /= ss.xi
gg.xim /= ss.xi
gg.varxi /= ss.xi**2
return gg
def calculate_cc(df, sys_name, weights=False):
if weights:
ngal_cat = treecorr.Catalog(ra=df['RA'], dec=df['DEC'], ra_units='deg', dec_units='deg', k=df['ngal_norm'], w=df['weights'])
else:
ngal_cat = treecorr.Catalog(ra=df['RA'], dec=df['DEC'], ra_units='deg', dec_units='deg', k=df['ngal_norm'])
sys_cat = treecorr.Catalog(ra=df['RA'], dec=df['DEC'], ra_units='deg', dec_units='deg', k=df[sys_name])
kk = treecorr.KKCorrelation(nbins=10, min_sep=10, max_sep=100, sep_units = 'arcmin')
print('Start calculating correlation...')
start = time.time()
kk.process(ngal_cat,sys_cat, num_threads=80) # Compute the cross-correlation function.
if weights:
kk.write(data_dir+'/cross_correlations/cc_ngal_'+sys_name+'_weights.fits')
else:
kk.write(data_dir+'/cross_correlations/cc_ngal_'+sys_name+'.fits')
end = time.time()
print('Done!')
print('Time: {}'.format(end-start))
def comp_xi(self):
cat_du = treecorr.Catalog(x=self.u,
y=self.v,
k=(self.du - np.mean(self.du)),
w=1. / self.du_err**2)
cat_dv = treecorr.Catalog(x=self.u,
y=self.v,
k=(self.dv - np.mean(self.dv)),
w=1. / self.dv_err**2)
kk = treecorr.KKCorrelation(min_sep=0,
max_sep=20. * 60.,
nbins=25,
bin_type='TwoD',
bin_slop=0)
kk.process(cat_du)
self.xi_dudu = copy.deepcopy(kk.xi)
kk.process(cat_dv)
self.xi_dvdv = copy.deepcopy(kk.xi)
kk.process(cat_du, cat_dv)
self.xi_dudv = copy.deepcopy(kk.xi)
npixels = len(kk.xi)**2
self.xi_sep = np.array(
[kk.dx.reshape(npixels),
kk.dy.reshape(npixels)]).T
self.dic_output['2pcf_stat'].update({
'xi_dudu': self.xi_dudu,
'xi_dudv': self.xi_dudv,
'xi_dvdv': self.xi_dvdv,
'xi_sep': self.xi_sep
})
def measure_rho(data,
max_sep,
max_mag,
tag=None,
use_xy=False,
prefix='piff',
alt_tt=False,
opt=None,
subtract_mean=False,
do_rho0=False):
"""Compute the rho statistics
"""
import treecorr
e1 = data['obs_e1']
e2 = data['obs_e2']
T = data['obs_T']
p_e1 = data[prefix + '_e1']
p_e2 = data[prefix + '_e2']
p_T = data[prefix + '_T']
m = data['mag']
if max_mag > 0:
e1 = e1[m < max_mag]
e2 = e2[m < max_mag]
T = T[m < max_mag]
p_e1 = p_e1[m < max_mag]
p_e2 = p_e2[m < max_mag]
p_T = p_T[m < max_mag]
q1 = e1 - p_e1
q2 = e2 - p_e2
dt = (T - p_T) / T
w1 = e1 * dt
w2 = e2 * dt
print('mean e = ', np.mean(e1), np.mean(e2))
print('std e = ', np.std(e1), np.std(e2))
print('mean T = ', np.mean(T))
print('std T = ', np.std(T))
print('mean de = ', np.mean(q1), np.mean(q2))
print('std de = ', np.std(q1), np.std(q2))
print('mean dT = ', np.mean(T - p_T))
print('std dT = ', np.std(T - p_T))
print('mean dT/T = ', np.mean(dt))
print('std dT/T = ', np.std(dt))
if subtract_mean:
e1 -= np.mean(e1)
e2 -= np.mean(e2)
q1 -= np.mean(q1)
q2 -= np.mean(q2)
w1 -= np.mean(w1)
w2 -= np.mean(w2)
dt -= np.mean(dt)
if use_xy:
x = data['fov_x']
y = data['fov_y']
if max_mag > 0:
x = x[m < max_mag]
y = y[m < max_mag]
print('x = ', x)
print('y = ', y)
ecat = treecorr.Catalog(x=x,
y=y,
x_units='arcsec',
y_units='arcsec',
g1=e1,
g2=e2)
qcat = treecorr.Catalog(x=x,
y=y,
x_units='arcsec',
y_units='arcsec',
g1=q1,
g2=q2)
wcat = treecorr.Catalog(x=x,
y=y,
x_units='arcsec',
y_units='arcsec',
g1=w1,
g2=w2,
k=dt)
else:
ra = data['ra']
dec = data['dec']
if max_mag > 0:
ra = ra[m < max_mag]
dec = dec[m < max_mag]
print('ra = ', ra)
print('dec = ', dec)
ecat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='deg',
dec_units='deg',
g1=e1,
g2=e2)
qcat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='deg',
dec_units='deg',
g1=q1,
g2=q2)
wcat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='deg',
dec_units='deg',
g1=w1,
g2=w2,
k=dt)
ecat.name = 'ecat'
qcat.name = 'qcat'
wcat.name = 'wcat'
if tag is not None:
for cat in [ecat, qcat, wcat]:
cat.name = tag + ":" + cat.name
bin_config = dict(
sep_units='arcmin',
bin_slop=0.1,
min_sep=0.5,
max_sep=max_sep,
bin_size=0.2,
)
if opt == 'lucas':
bin_config['min_sep'] = 2.5
bin_config['max_sep'] = 250.
bin_config['nbins'] = 20
del bin_config['bin_size']
if opt == 'fine_bin':
bin_config['min_sep'] = 0.1
bin_config['max_sep'] = 2000.
bin_config['bin_size'] = 0.01
pairs = [(qcat, qcat), (ecat, qcat), (wcat, wcat), (qcat, wcat),
(ecat, wcat)]
if do_rho0:
pairs.append((ecat, ecat))
results = []
for (cat1, cat2) in pairs:
print('Doing correlation of %s vs %s' % (cat1.name, cat2.name))
rho = treecorr.GGCorrelation(bin_config, verbose=2)
if cat1 is cat2:
rho.process(cat1)
else:
rho.process(cat1, cat2)
print('mean xi+ = ', rho.xip.mean())
print('mean xi- = ', rho.xim.mean())
results.append(rho)
if alt_tt:
print('Doing alt correlation of %s vs %s' % (dtcat.name, dtcat.name))
rho = treecorr.KKCorrelation(bin_config, verbose=2)
rho.process(dtcat)
results.append(rho)
return results
min_sep, max_sep = rmin, rmax
nbins = 20
dataK = tc.Catalog(k=cmb_map,
ra=cmb_ra,
dec=cmb_dec,
ra_units='deg',
dec_units='deg')
datag = tc.Catalog(k=g_map,
ra=g_ra,
dec=g_dec,
ra_units='deg',
dec_units='deg')
Kg = tc.KKCorrelation(nbins=nbins,
min_sep=min_sep,
max_sep=max_sep,
bin_slop=0.01,
verbose=0,
sep_units='arcmin')
Kg.process(dataK, datag, metric='Arc', num_threads=32)
xim = Kg.xi
rm = Kg.meanr
np.savetxt(
'../data/BASS/' + name_save + name_pre + m1 + m2 + 'z' + z1 + '-' + z2 +
'R' + Radius + 'r' + np.str(rmin) + '-' + np.str(rmax) + 'nbin' +
np.str(nbins) + 'nside' + np.str(nside) + 'nsidecmb' + np.str(nside_cmb) +
'nsidejack' + np.str(nside_jack) + 'wmap.dat',
np.vstack((Kg.xi, Kg.meanr, Kg.npairs, np.sqrt(Kg.varxi))))
#(5)jackknife error bar
flag_max = np.int32(flag_area.max())
xi = np.zeros((flag_max, 2, nbins))
logging.info(f"Calculating for NSIDE = {nside}")
rcscat = treecorr.Catalog("rcslens.fits", rcsconfig)
szcat = treecorr.Catalog(f"szmaps_masked{maskno}_{nside}.fits", szconfig)
m1cat = treecorr.Catalog("rcslens.fits", m1config)
kg = treecorr.KGCorrelation(corrconfig,logger=logging.getLogger())
kg.process(szcat, rcscat) # Calculate the cross-correlation
kg.write(f"crosscorr{maskno}_{nside}.result")
nk = treecorr.NKCorrelation(corrconfig,logger=logging.getLogger())
nk.process(szcat, m1cat)
nk.write(f"calib{maskno}_{nside}.result")
ny = treecorr.NKCorrelation(corrconfig,logger=logging.getLogger())
ny.process(rcscat, szcat)
ny.write(f"ycorr{maskno}_{nside}.result")
ng = treecorr.NGCorrelation(corrconfig,logger=logging.getLogger())
ng.process(rcscat, rcscat)
ng.write(f"shearcorr{maskno}_{nside}.result")
gg = treecorr.GGCorrelation(corrconfig,logger=logging.getLogger())
gg.process(rcscat, rcscat)
gg.write(f"shear_auto_corr{maskno}_{nside}.result")
kk = treecorr.KKCorrelation(corrconfig,logger=logging.getLogger())
kk.process(szcat, szcat)
kk.write(f"sz_auto_corr{maskno}_{nside}.result")
logging.info("DONE")
def test_periodic_ps():
# Make a kappa field with a known power spectrum on a periodic grid.
# This is heavily based on the GalSim PowerSpectrumRealizer class.
ngrid = 256
L = 100. # grid size
r0 = 10. # scale size of Gaussian in real space
k0 = 1./ (2.*np.pi * r0) # scale in Fourier space
# Set up grid in k space
kx_1d = np.fft.fftfreq(ngrid, L/ngrid)
ky_1d = np.fft.fftfreq(ngrid, L/ngrid)
kx, ky = np.meshgrid(kx_1d, ky_1d)
# And real space
x_1d = np.linspace(0, L, ngrid)
y_1d = np.linspace(0, L, ngrid)
x, y = np.meshgrid(x_1d, y_1d)
# Use a Gaussian for the power spectrum, so the correlation function is also Gaussian.
kz = kx + 1j * ky
ksq = kz*np.conj(kz)
pk = np.exp(-ksq/(2*k0**2))
# Multiply by a Gaussian random field
rng = np.random.RandomState(8675309)
r1 = rng.normal(0,1,ngrid**2).reshape(pk.shape)
r2 = rng.normal(0,1,ngrid**2).reshape(pk.shape)
pk *= (r1 + 1j*r2) / np.sqrt(2.)
pk[0,0] = 0.
# pk is now the FT for the kappa field.
# The gamma field is just rotated by exp(2ipsi)
ksq[0,0] = 1. # Not zero, so no division by 0.
exp2ipsi = kz*kz / ksq
pg = pk * exp2ipsi
# Take the inverse FFT to get the real-space fields
kappa = ngrid * np.fft.ifft2(pk)
gamma = ngrid * np.fft.ifft2(pg)
# Now measure the correlation function with TreeCorr using periodic boundary conditions.
cat = treecorr.Catalog(x=x.ravel(), y=y.ravel(), k=kappa.real.ravel(),
g1=gamma.real.ravel(), g2=gamma.imag.ravel())
kk = treecorr.KKCorrelation(min_sep=3, max_sep=10, nbins=20, period=ngrid)
kk.process(cat, metric='Periodic')
print('kk.xi = ',kk.xi)
# I admit that I didn't determine the amplitude of this from first principles.
# I think it has to scale as k0**2 (or 1/r0**2), but the factors of 2 and pi confuse me.
# The 1/pi here is empirical.
true_xi = np.exp(-kk.meanr**2/(2*r0**2)) * k0**2/np.pi
print('true_xi = ',true_xi)
print('ratio = ',kk.xi/true_xi)
np.testing.assert_allclose(kk.xi, true_xi, rtol=0.15)
gg = treecorr.GGCorrelation(min_sep=3, max_sep=10, nbins=20, period=ngrid)
gg.process(cat, metric='Periodic')
print('gg.xip = ',gg.xip)
print('gg.xim = ',gg.xim)
# I thought these were supposed to be equal, but xip is larger by a factor of sqrt(2).
# I'm sure this somehow makes sense because of the 2 complex components of the shear,
# but I didn't try to figure it out. I'm just going with it and assuming this is right.
true_xip = true_xi * np.sqrt(2)
print('true_xip = ',true_xip)
print('ratio = ',gg.xip/true_xip)
np.testing.assert_allclose(gg.xip, true_xip, rtol=0.15)
def test_direct_spherical():
# Repeat in spherical coords
ngal = 100
s = 10.
rng = np.random.RandomState(8675309)
x1 = rng.normal(0,s, (ngal,) )
y1 = rng.normal(0,s, (ngal,) ) + 200 # Put everything at large y, so small angle on sky
z1 = rng.normal(0,s, (ngal,) )
w1 = rng.random_sample(ngal)
k1 = rng.normal(10,1, (ngal,) )
x2 = rng.normal(0,s, (ngal,) )
y2 = rng.normal(0,s, (ngal,) ) + 200
z2 = rng.normal(0,s, (ngal,) )
w2 = rng.random_sample(ngal)
k2 = rng.normal(0,3, (ngal,) )
ra1, dec1 = coord.CelestialCoord.xyz_to_radec(x1,y1,z1)
ra2, dec2 = coord.CelestialCoord.xyz_to_radec(x2,y2,z2)
cat1 = treecorr.Catalog(ra=ra1, dec=dec1, ra_units='rad', dec_units='rad', w=w1, k=k1)
cat2 = treecorr.Catalog(ra=ra2, dec=dec2, ra_units='rad', dec_units='rad', w=w2, k=k2)
min_sep = 1.
max_sep = 10.
nbins = 50
bin_size = np.log(max_sep/min_sep) / nbins
kk = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins,
sep_units='deg', brute=True)
kk.process(cat1, cat2)
r1 = np.sqrt(x1**2 + y1**2 + z1**2)
r2 = np.sqrt(x2**2 + y2**2 + z2**2)
x1 /= r1; y1 /= r1; z1 /= r1
x2 /= r2; y2 /= r2; z2 /= r2
north_pole = coord.CelestialCoord(0*coord.radians, 90*coord.degrees)
true_npairs = np.zeros(nbins, dtype=int)
true_weight = np.zeros(nbins, dtype=float)
true_xi = np.zeros(nbins, dtype=float)
c1 = [coord.CelestialCoord(r*coord.radians, d*coord.radians) for (r,d) in zip(ra1, dec1)]
c2 = [coord.CelestialCoord(r*coord.radians, d*coord.radians) for (r,d) in zip(ra2, dec2)]
for i in range(ngal):
for j in range(ngal):
rsq = (x1[i]-x2[j])**2 + (y1[i]-y2[j])**2 + (z1[i]-z2[j])**2
r = np.sqrt(rsq)
r *= coord.radians / coord.degrees
logr = np.log(r)
index = np.floor(np.log(r/min_sep) / bin_size).astype(int)
if index < 0 or index >= nbins:
continue
ww = w1[i] * w2[j]
xi = ww * k1[i] * k2[j]
true_npairs[index] += 1
true_weight[index] += ww
true_xi[index] += xi
true_xi /= true_weight
print('true_npairs = ',true_npairs)
print('diff = ',kk.npairs - true_npairs)
np.testing.assert_array_equal(kk.npairs, true_npairs)
print('true_weight = ',true_weight)
print('diff = ',kk.weight - true_weight)
np.testing.assert_allclose(kk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
print('true_xi = ',true_xi)
print('kk.xi = ',kk.xi)
np.testing.assert_allclose(kk.xi, true_xi, rtol=1.e-4, atol=1.e-8)
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
# Check that running via the corr2 script works correctly.
config = treecorr.config.read_config('configs/kk_direct_spherical.yaml')
cat1.write(config['file_name'])
cat2.write(config['file_name2'])
treecorr.corr2(config)
data = fitsio.read(config['kk_file_name'])
np.testing.assert_allclose(data['r_nom'], kk.rnom)
np.testing.assert_allclose(data['npairs'], kk.npairs)
np.testing.assert_allclose(data['weight'], kk.weight)
np.testing.assert_allclose(data['xi'], kk.xi, rtol=1.e-3)
# Repeat with binslop = 0
# And don't do any top-level recursion so we actually test not going to the leaves.
kk = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins,
sep_units='deg', bin_slop=0, max_top=0)
kk.process(cat1, cat2)
np.testing.assert_array_equal(kk.npairs, true_npairs)
np.testing.assert_allclose(kk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kk.xi, true_xi, rtol=1.e-3, atol=1.e-6)
def xi_2pt(cata,
catb=None,
k=None,
ga=None,
gb=None,
corr='GG',
maska=None,
maskb=None,
wa=None,
wb=None,
ran=True,
mock=False,
erron=True,
jkmask=None,
label0='',
plot=False):
"""
This is a flexible convenience wrapper for interaction with treecorr to work on CatalogStore objects. Some basic examples are given in corr_tests() of the main testsuite.py. g1, g2 correctly by c1, c2 if ellipticities and cat.bs is true. Correction by sensitivity, 1+m applied if cat.bs=True. Weighting applied if cat.wt is true. Other config properties for treecorr stored in CatalogStore object. See catalog.py or config.py. Not all correlation types fully integrated or tested. For example, only one kappa value is currently possible. Will be updated in future as useful.
Use:
:cata, catb: CatalogStore - Must supply both cata, catb (can be same reference) if NG or NK correlation. Otherwise catb is optional.
:k: str - Array name in cata, catb to use for kappa correlation.
:ga, gb: str - Array names for g1, g2 treecorr inputs. If None assume e1, e2.
:corr: str - Type of correlation for treecorr.
:maska, maskb: [bool] - Masking array to apply to input catalogs.
:wa, wb: [float] - Additional weights to apply after cat.w is used. Combined as e.g., w=sqrt(cat.w*wa).
:ran: bool - Use randoms in correlation calculation. If True, assumes cat.ran_ra, cat.ran_dec exist.
:mock: bool - If mock catalog from sims. Used when calculating covariances from sims, not currently migrated from SV code.
:erron: bool - Calculate jackknife or sim cov errors. If False, uses treecorr error outputs. Not currently migrated from SV code. When implemented requires cat.calc_err in ('jk', 'mock').
:jkmask: [bool] - For jk, mock cov calculation loop over regions/sims.
:label0: str - Additional (optional) label string used in some outputs.
:plot: bool - Plot output?
Output (len cat.tbins):
:theta: [float] - Treecorr np.exp(meanlogr)
:out: ([float]x4) - Output of signal e.g., (xi+,xi-,xi+im,x-im). For correlations with only one xi output, (xi,0.,xi_im,0.).
:err: ([float]x4) - Same but for sqrt(var).
:chi2: ([float]x4) - Same but for chi^2 if using jk or sim covariance.
"""
maska = catalog.CatalogMethods.check_mask(cata.coadd, maska)
jkmask = catalog.CatalogMethods.check_mask(cata.coadd, jkmask)
maska0 = maska & jkmask
if wa is None:
wa = np.ones(len(cata.coadd))
e1, e2, w, ms = lin.linear_methods.get_lin_e_w_ms(cata,
xi=True,
mock=mock,
mask=maska0,
w1=wa)
if catb is None:
if corr not in ['GG', 'NN', 'KK']:
raise UseError(
'Must supply both cata,catb for NG,NK correlations.')
if ga is not None:
e1 = getattr(cata, ga + '1')[maska]
e2 = getattr(cata, ga + '2')[maska]
else:
ga = 'e'
if catb is None:
gb = ga
if (corr == 'GG') | ((catb != None) & (corr == 'KG')):
catxa = treecorr.Catalog(g1=e1,
g2=e2,
w=w,
ra=cata.ra[maska0],
dec=cata.dec[maska0],
ra_units='deg',
dec_units='deg')
catma = treecorr.Catalog(k=ms,
w=w,
ra=cata.ra[maska0],
dec=cata.dec[maska0],
ra_units='deg',
dec_units='deg')
elif (corr == 'NN') | ((catb != None) & (corr in ['NG', 'NK'])):
catxa = treecorr.Catalog(w=w,
ra=cata.ra[maska0],
dec=cata.dec[maska0],
ra_units='deg',
dec_units='deg')
if ran:
catra = treecorr.Catalog(w=w,
ra=cata.ran_ra[maska0],
dec=cata.ran_dec[maska0],
ra_units='deg',
dec_units='deg')
elif corr == 'KK':
if k is None:
raise UseError('Must specify k for KK correlation.')
if k not in dir(cata):
raise UseError('Unknown k field specified.')
catxa = treecorr.Catalog(k=getattr(cata, k)[maska0],
w=w,
ra=cata.ra[maska0],
dec=cata.dec[maska0],
ra_units='deg',
dec_units='deg')
if catb is not None:
maskb = catalog.CatalogMethods.check_mask(catb.coadd, maskb)
if wb is None:
wb = np.ones(len(catb.coadd))
e1, e2, w, ms = lin.linear_methods.get_lin_e_w_ms(catb,
xi=True,
mock=mock,
mask=maskb,
w1=wb)
if gb is not None:
e1 = getattr(cata, gb + '1')[maskb]
e2 = getattr(cata, gb + '2')[maskb]
else:
gb = 'e'
if corr in ['GG', 'NG', 'KG']:
catxb = treecorr.Catalog(g1=e1,
g2=e2,
w=w,
ra=catb.ra[maskb],
dec=catb.dec[maskb],
ra_units='deg',
dec_units='deg')
catmb = treecorr.Catalog(k=ms,
w=w,
ra=catb.ra[maskb],
dec=catb.dec[maskb],
ra_units='deg',
dec_units='deg')
elif corr == 'NN':
catxb = treecorr.Catalog(w=w,
ra=catb.ra[maskb],
dec=catb.dec[maskb],
ra_units='deg',
dec_units='deg')
if ran:
catrb = treecorr.Catalog(w=w,
ra=catb.ran_ra[maskb],
dec=catb.ran_dec[maskb],
ra_units='deg',
dec_units='deg')
elif corr in ['KK', 'NK']:
if k is None:
raise UseError('Must specify k for KK correlation.')
if k not in dir(catb):
raise UseError('Unknown k field specified.')
catxb = treecorr.Catalog(k=getattr(catb, k)[maskb],
w=w,
ra=catb.ra[maskb],
dec=catb.dec[maskb],
ra_units='deg',
dec_units='deg')
xim = None
xip_im = None
xim_im = None
ximerr = None
xiperr_im = None
ximerr_im = None
if corr == 'GG':
gg = treecorr.GGCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
kk = treecorr.KKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
if catb is None:
gg.process(catxa)
kk.process(catma)
else:
gg.process(catxa, catxb)
kk.process(catma, catmb)
xip = gg.xip / kk.xi
xim = gg.xim / kk.xi
xiperr = ximerr = np.sqrt(gg.varxi)
xip_im = gg.xip_im / kk.xi
xim_im = gg.xim_im / kk.xi
theta = np.exp(gg.meanlogr)
elif corr == 'NN':
nn = treecorr.NNCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
if ran:
nr = treecorr.NNCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
rr = treecorr.NNCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
if catb is None:
nn.process(catxa)
xip = nn.npairs
xiperr = np.sqrt(nn.npairs)
if ran:
nr.process(catxa, catra)
rr.process(catra)
xip, xiperr = nn.calculateXi(rr, nr)
xiperr = np.sqrt(xiperr)
else:
rn = treecorr.NNCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
nn.process(catxa, catxb)
xip = nn.npairs
xiperr = np.sqrt(nn.npairs)
if ran:
nr.process(catxa, catrb)
nr.process(catra, catxb)
rr.process(catra, catrb)
xip, xiperr = nn.calculateXi(rr, nr, rn)
xiperr = np.sqrt(xiperr)
theta = np.exp(nn.meanlogr)
elif corr == 'KK':
kk = treecorr.KKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
if catb is None:
kk.process(catxa)
else:
kk.process(catxa, catxb)
xip = kk.xi
xiperr = np.sqrt(kk.varxi)
theta = np.exp(kk.meanlogr)
elif corr == 'KG':
kg = treecorr.KGCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
kk = treecorr.KKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
kg.process(catxa, catxb)
kk.process(catxa, catmb)
xip = kg.xi / kk.xi
xiperr = np.sqrt(kg.varxi)
xip_im = kg.xi_im / kk.xi
theta = np.exp(kg.meanlogr)
elif corr == 'NG':
ng = treecorr.NGCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
nk = treecorr.NKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
ng.process(catxa, catxb)
nk.process(catxa, catmb)
xip = ng.xi / nk.xi
xiperr = np.sqrt(ng.varxi)
xip_im = ng.xi_im / nk.xi
if ran:
rg = treecorr.NGCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
rk = treecorr.NKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
rg.process(catra, catxb)
rk.process(catra, catmb)
xip, xip_im, xiperr = ng.calculateXi(rg)
tmpa, tmp = nk.calculateXi(rk)
if np.sum(tmpa) == 0:
tmpa = np.ones(len(xip))
xip /= tmpa
xiperr = np.sqrt(xiperr)
xip_im /= tmpa
theta = np.exp(ng.meanlogr)
elif corr == 'NK':
nk = treecorr.NKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
nk.process(catxa, catxb)
xip = nk.xi
xiperr = np.sqrt(nk.varxi)
if ran:
rk = treecorr.NKCorrelation(nbins=cata.tbins,
min_sep=cata.sep[0],
max_sep=cata.sep[1],
sep_units='arcmin',
bin_slop=cata.slop,
verbose=0)
rk.process(catra, catxb)
xip, xiperr = nk.calculateXi(rk)
xiperr = np.sqrt(xiperr)
theta = np.exp(nk.meanlogr)
out = [xip, xim, xip_im, xim_im]
err = [xiperr, ximerr, xiperr, ximerr]
chi2 = [0., 0., 0., 0.]
if erron:
kwargs = {
'catb': catb,
'k': k,
'corr': corr,
'maska': maska,
'maskb': maskb,
'wa': wa,
'wb': wb,
'ran': ran
}
if catb is None:
if corr in ['KK', 'NK', 'KG']:
label = 'xi_2pt_' + cata.name + '_' + k + '_' + corr + '_' + label0
else:
label = 'xi_2pt_' + cata.name + '_' + corr + '_' + label0
else:
if corr in ['KK', 'NK', 'KG']:
label = 'xi_2pt_' + cata.name + '-' + catb.name + '_' + k + '_' + corr + '_' + label0
else:
label = 'xi_2pt_' + cata.name + '-' + catb.name + '_' + corr + '_' + label0
if cata.calc_err == 'jk':
err, chi2 = jackknife_methods.jk(cata, xi_2pt.xi_2pt,
[xip, xim, xip_im, xim_im],
label, **kwargs)
elif cata.calc_err == 'mock':
ggperr, ggmerr, chi2p, chi2m, ceerr, cberr, cechi2, cbchi2 = BCC_Methods.jk_iter_xi(
cat,
ggp,
ggm,
ce,
cb,
mask,
w,
cosebi=cosebi,
parallel=parallel)
if plot:
fig.plot_methods.fig_create_xi(cata, corr, theta, out, err, k, ga,
gb)
return theta, out, err, chi2
def test_direct():
# If the catalogs are small enough, we can do a direct calculation to see if comes out right.
# This should exactly match the treecorr result if brute force.
ngal = 200
s = 10.
rng = np.random.RandomState(8675309)
x1 = rng.normal(0,s, (ngal,) )
y1 = rng.normal(0,s, (ngal,) )
w1 = rng.random_sample(ngal)
k1 = rng.normal(10,1, (ngal,) )
x2 = rng.normal(0,s, (ngal,) )
y2 = rng.normal(0,s, (ngal,) )
w2 = rng.random_sample(ngal)
k2 = rng.normal(0,3, (ngal,) )
cat1 = treecorr.Catalog(x=x1, y=y1, w=w1, k=k1)
cat2 = treecorr.Catalog(x=x2, y=y2, w=w2, k=k2)
min_sep = 1.
max_sep = 50.
nbins = 50
bin_size = np.log(max_sep/min_sep) / nbins
kk = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, brute=True)
kk.process(cat1, cat2)
true_npairs = np.zeros(nbins, dtype=int)
true_weight = np.zeros(nbins, dtype=float)
true_xi = np.zeros(nbins, dtype=float)
for i in range(ngal):
# It's hard to do all the pairs at once with numpy operations (although maybe possible).
# But we can at least do all the pairs for each entry in cat1 at once with arrays.
rsq = (x1[i]-x2)**2 + (y1[i]-y2)**2
r = np.sqrt(rsq)
logr = np.log(r)
expmialpha = ((x1[i]-x2) - 1j*(y1[i]-y2)) / r
ww = w1[i] * w2
xi = ww * k1[i] * k2
index = np.floor(np.log(r/min_sep) / bin_size).astype(int)
mask = (index >= 0) & (index < nbins)
np.add.at(true_npairs, index[mask], 1)
np.add.at(true_weight, index[mask], ww[mask])
np.add.at(true_xi, index[mask], xi[mask])
true_xi /= true_weight
print('true_npairs = ',true_npairs)
print('diff = ',kk.npairs - true_npairs)
np.testing.assert_array_equal(kk.npairs, true_npairs)
print('true_weight = ',true_weight)
print('diff = ',kk.weight - true_weight)
np.testing.assert_allclose(kk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
print('true_xi = ',true_xi)
print('kk.xi = ',kk.xi)
np.testing.assert_allclose(kk.xi, true_xi, rtol=1.e-4, atol=1.e-8)
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
# Check that running via the corr2 script works correctly.
config = treecorr.config.read_config('configs/kk_direct.yaml')
cat1.write(config['file_name'])
cat2.write(config['file_name2'])
treecorr.corr2(config)
data = fitsio.read(config['kk_file_name'])
np.testing.assert_allclose(data['r_nom'], kk.rnom)
np.testing.assert_allclose(data['npairs'], kk.npairs)
np.testing.assert_allclose(data['weight'], kk.weight)
np.testing.assert_allclose(data['xi'], kk.xi, rtol=1.e-3)
# Repeat with binslop = 0
# And don't do any top-level recursion so we actually test not going to the leaves.
kk = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, bin_slop=0,
max_top=0)
kk.process(cat1, cat2)
np.testing.assert_array_equal(kk.npairs, true_npairs)
np.testing.assert_allclose(kk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kk.xi, true_xi, rtol=1.e-4, atol=1.e-8)
# Check a few basic operations with a KKCorrelation object.
do_pickle(kk)
kk2 = kk.copy()
kk2 += kk
np.testing.assert_allclose(kk2.npairs, 2*kk.npairs)
np.testing.assert_allclose(kk2.weight, 2*kk.weight)
np.testing.assert_allclose(kk2.meanr, 2*kk.meanr)
np.testing.assert_allclose(kk2.meanlogr, 2*kk.meanlogr)
np.testing.assert_allclose(kk2.xi, 2*kk.xi)
kk2.clear()
kk2 += kk
np.testing.assert_allclose(kk2.npairs, kk.npairs)
np.testing.assert_allclose(kk2.weight, kk.weight)
np.testing.assert_allclose(kk2.meanr, kk.meanr)
np.testing.assert_allclose(kk2.meanlogr, kk.meanlogr)
np.testing.assert_allclose(kk2.xi, kk.xi)
ascii_name = 'output/kk_ascii.txt'
kk.write(ascii_name, precision=16)
kk3 = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
kk3.read(ascii_name)
np.testing.assert_allclose(kk3.npairs, kk.npairs)
np.testing.assert_allclose(kk3.weight, kk.weight)
np.testing.assert_allclose(kk3.meanr, kk.meanr)
np.testing.assert_allclose(kk3.meanlogr, kk.meanlogr)
np.testing.assert_allclose(kk3.xi, kk.xi)
fits_name = 'output/kk_fits.fits'
kk.write(fits_name)
kk4 = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
kk4.read(fits_name)
np.testing.assert_allclose(kk4.npairs, kk.npairs)
np.testing.assert_allclose(kk4.weight, kk.weight)
np.testing.assert_allclose(kk4.meanr, kk.meanr)
np.testing.assert_allclose(kk4.meanlogr, kk.meanlogr)
np.testing.assert_allclose(kk4.xi, kk.xi)
with assert_raises(TypeError):
kk2 += config
kk4 = treecorr.KKCorrelation(min_sep=min_sep/2, max_sep=max_sep, nbins=nbins)
with assert_raises(ValueError):
kk2 += kk4
kk5 = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep*2, nbins=nbins)
with assert_raises(ValueError):
kk2 += kk5
kk6 = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins*2)
with assert_raises(ValueError):
kk2 += kk6
def test_pairwise():
# Test the pairwise option.
ngal = 1000
s = 10.
rng = np.random.RandomState(8675309)
x1 = rng.normal(0,s, (ngal,) )
y1 = rng.normal(0,s, (ngal,) )
w1 = rng.random_sample(ngal)
k1 = rng.normal(10,1, (ngal,) )
x2 = rng.normal(0,s, (ngal,) )
y2 = rng.normal(0,s, (ngal,) )
w2 = rng.random_sample(ngal)
k2 = rng.normal(0,3, (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, k=k2)
min_sep = 5.
max_sep = 50.
nbins = 10
bin_size = np.log(max_sep/min_sep) / nbins
kk = treecorr.KKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
with assert_warns(FutureWarning):
kk.process_pairwise(cat1, cat2)
kk.finalize(cat1.vark, cat2.vark)
true_npairs = np.zeros(nbins, dtype=int)
true_weight = np.zeros(nbins, dtype=float)
true_xi = np.zeros(nbins, dtype=float)
rsq = (x1-x2)**2 + (y1-y2)**2
r = np.sqrt(rsq)
logr = np.log(r)
expmialpha = ((x1-x2) - 1j*(y1-y2)) / r
ww = w1 * w2
xi = ww * k1 * k2
index = np.floor(np.log(r/min_sep) / bin_size).astype(int)
mask = (index >= 0) & (index < nbins)
np.add.at(true_npairs, index[mask], 1)
np.add.at(true_weight, index[mask], ww[mask])
np.add.at(true_xi, index[mask], xi[mask])
true_xi /= true_weight
np.testing.assert_array_equal(kk.npairs, true_npairs)
np.testing.assert_allclose(kk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kk.xi, true_xi, rtol=1.e-4, atol=1.e-8)
# If cats have names, then the logger will mention them.
# Also, test running with optional args.
cat1.name = "first"
cat2.name = "second"
with CaptureLog() as cl:
kk.logger = cl.logger
with assert_warns(FutureWarning):
kk.process_pairwise(cat1, cat2, metric='Euclidean', num_threads=2)
assert "for cats first, second" in cl.output
def compute_2pcf(skymaps=None, mask=None, nmaps=50, binning='log',nlinbin=100, pix_units='deg',\
npatch=10, var_method='shot', return_cov=True, prefilter=False, verbose=False,\
plot=False, thetamin_fac=1.0, mask_list=None,\
thetamax=2., thetabin_size=0.3, pixel_size=7., imdim=1024, plot_title='realization', bin_slop=0.1):
''' Inputs:
thetamax (deg): Maximum theta desired. If using radian pixel units this will get converted to radians in the code.
nlinbin (int, default: 100): number of linear bins to use if doing linear binning
'''
fs = []
wthetas, varwthetas, covs = [], [], []
deg_per_pix = (pixel_size / 3600.) # arcseconds to degrees
thetamin = thetamin_fac * deg_per_pix # minimum separation is twice the pixel size
print('thetamin is ', thetamin)
print('in arcseconds this is ', thetamin * 3600.)
if skymaps is not None:
imdim = skymaps[0].shape[0]
if nmaps > len(skymaps):
nmaps = len(skymaps)
deg_per_rad = np.pi / 180.
xind = np.indices((imdim, imdim))[0].ravel() * deg_per_pix
yind = np.indices((imdim, imdim))[1].ravel() * deg_per_pix
# I checked and having the pix_units in radians or degrees does not change the result so long as
# they are consistently used within the Catalog object.
if pix_units == 'rad':
xind *= deg_per_rad
yind *= deg_per_rad
thetamax *= deg_per_rad
thetamin *= deg_per_rad
if verbose:
print('xind:', xind)
if npatch > 1 and var_method == 'shot':
print('because npatch=' + str(npatch) +
' and var_method is shot, changing var_method to jackknife')
var_method = 'jackknife'
elif npatch == 1 and var_method == 'jackknife':
print(
'because npatch=1 and var_method is set to jackknife, changing var_method to shot'
)
var_method = 'shot'
for i in range(nmaps):
if verbose:
print('i = ', i)
load_srcmap = skymaps[i]
skymap = load_srcmap.copy()
if mask_list is not None:
xind_cat, yind_cat, ks, weights = get_treecorr_catalog_vals(
xind, yind, skymap, mask=mask_list[i], prefilter=prefilter)
else:
xind_cat, yind_cat, ks, weights = get_treecorr_catalog_vals(
xind, yind, skymap, mask=mask, prefilter=prefilter)
print('sum of weights is ', np.sum(weights))
cat = treecorr.Catalog(ra=xind_cat,
dec=yind_cat,
k=ks,
w=weights,
ra_units=pix_units,
dec_units=pix_units,
npatch=npatch)
if binning == 'linear':
kk = treecorr.KKCorrelation(min_sep=thetamin,
max_sep=thetamax,
bin_type='Linear',
nbins=nlinbin,
sep_units=pix_units,
bin_slop=bin_slop)
elif binning == 'log':
kk = treecorr.KKCorrelation(min_sep=thetamin,
max_sep=thetamax,
bin_size=thetabin_size,
sep_units=pix_units,
var_method=var_method,
bin_slop=bin_slop)
kk.process(cat, metric='Euclidean')
if verbose:
print('var method is ', kk.var_method)
print('number of pairs used is ', kk.npairs)
print('kk.r_nom is ', kk.rnom * 3600)
print('kk.meanr is ', kk.meanr * 3600)
npair_mask = (kk.npairs > 0.)
if plot:
if binning == 'linear':
plot_xscale = None
else:
plot_xscale = 'log'
f = plot_wtheta(kk.rnom[npair_mask],
kk.xi[npair_mask],
kk.varxi[npair_mask],
skymap,
pix_units=pix_units,
mask=mask,
return_fig=True,
plot_xscale=plot_xscale,
title=plot_title)
fs.append(f)
wthetas.append(kk.xi[npair_mask])
varwthetas.append(kk.varxi[npair_mask])
if return_cov:
covs.append(kk.cov[npair_mask])
if plot:
if return_cov:
return wthetas, varwthetas, kk.rnom[npair_mask], covs, fs
else:
return wthetas, varwthetas, kk.rnom[npair_mask], fs
else:
if return_cov:
return wthetas, varwthetas, kk.rnom[npair_mask], covs
else:
return wthetas, varwthetas, kk.rnom[npair_mask]
density_fluctuations_RA = density_fluctuations[0, :]
density_fluctuations_dec = density_fluctuations[1, :]
density_fluctuations_del = density_fluctuations[2, :]
mean_del = np.mean(density_fluctuations_del)
var_del = np.var(density_fluctuations)
cat = treecorr.Catalog(ra=density_fluctuations_RA,
dec=density_fluctuations_dec,
ra_units='rad',
dec_units='rad',
k=density_fluctuations_del)
kk = treecorr.KKCorrelation(min_sep=1,
max_sep=150,
nbins=15,
sep_units='arcmin')
kk.process(cat)
theta_tc = kk.rnom # tc stands for treecorr
w_tc = kk.xi
mean_del_vec = np.ones(theta_tc.size) * mean_del
var_del_vec = np.ones(theta_tc.size) * var_del
i_zeta = w_tc * mean_del # integrated 3-pt function
dat = np.array(
[theta_tc, w_tc, mean_del_vec, var_del_vec, i_zeta]
) # the mean_del_vec, var_del_vec is the same value for all angles (just for output purposes putting into column)
dat = dat.T
np.savetxt(filepath_map +
'B_treecorr_patches_correlated/correlation_mice_patch_' +
str(i + 1) + '.txt',
def measure_rho(data,
max_sep,
tag=None,
use_xy=False,
alt_tt=False,
prefix='piff'):
"""Compute the rho statistics
"""
import treecorr
e1 = data['obs_e1']
e2 = data['obs_e2']
T = data['obs_T']
p_e1 = data[prefix + '_e1']
p_e2 = data[prefix + '_e2']
p_T = data[prefix + '_T']
de1 = e1 - p_e1
de2 = e2 - p_e2
dt = (T - p_T) / T
print('mean e = ', np.mean(e1), np.mean(e2))
print('mean T = ', np.mean(T))
print('mean de = ', np.mean(de1), np.mean(de2))
print('mean dT = ', np.mean(T - p_T))
print('mean dT/T = ', np.mean(dt))
if use_xy:
x = data['fov_x']
y = data['fov_y']
print('x = ', x)
print('y = ', y)
ecat = treecorr.Catalog(x=x,
y=y,
x_units='arcsec',
y_units='arcsec',
g1=e1,
g2=e2)
decat = treecorr.Catalog(x=x,
y=y,
x_units='arcsec',
y_units='arcsec',
g1=de1,
g2=de2)
dtcat = treecorr.Catalog(x=x,
y=y,
x_units='arcsec',
y_units='arcsec',
k=dt,
g1=dt * e1,
g2=dt * e2)
else:
ra = data['ra']
dec = data['dec']
print('ra = ', ra)
print('dec = ', dec)
ecat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='deg',
dec_units='deg',
g1=e1,
g2=e2)
decat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='deg',
dec_units='deg',
g1=de1,
g2=de2)
dtcat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='deg',
dec_units='deg',
k=dt,
g1=dt * e1,
g2=dt * e2)
ecat.name = 'ecat'
decat.name = 'decat'
dtcat.name = 'dtcat'
if tag is not None:
for cat in [ecat, decat, dtcat]:
cat.name = tag + ":" + cat.name
bin_config = dict(
sep_units='arcmin',
bin_slop=0.1,
min_sep=0.5,
max_sep=max_sep,
bin_size=0.2,
#min_sep = 2.5,
#max_sep = 250,
#nbins = 20,
)
results = []
for (cat1, cat2) in [(decat, decat), (ecat, decat), (dtcat, dtcat),
(decat, dtcat), (ecat, dtcat)]:
print('Doing correlation of %s vs %s' % (cat1.name, cat2.name))
rho = treecorr.GGCorrelation(bin_config, verbose=2)
if cat1 is cat2:
rho.process(cat1)
else:
rho.process(cat1, cat2)
print('mean xi+ = ', rho.xip.mean())
print('mean xi- = ', rho.xim.mean())
results.append(rho)
if alt_tt:
print('Doing alt correlation of %s vs %s' % (dtcat.name, dtcat.name))
rho = treecorr.KKCorrelation(bin_config, verbose=2)
rho.process(dtcat)
results.append(rho)
return results
def test_kk():
# cf. http://adsabs.harvard.edu/abs/2002A%26A...389..729S for the basic formulae I use here.
#
# Use kappa(r) = A exp(-r^2/2s^2)
#
# The Fourier transform is: kappa~(k) = 2 pi A s^2 exp(-s^2 k^2/2) / L^2
# P(k) = (1/2pi) <|kappa~(k)|^2> = 2 pi A^2 (s/L)^4 exp(-s^2 k^2)
# xi(r) = (1/2pi) int( dk k P(k) J0(kr) )
# = pi A^2 (s/L)^2 exp(-r^2/2s^2/4)
# Note: I'm not sure I handled the L factors correctly, but the units at the end need
# to be kappa^2, so it needs to be (s/L)^2.
s = 10.
if __name__ == '__main__':
ngal = 1000000
L = 30. * s # Not infinity, so this introduces some error. Our integrals were to infinity.
tol_factor = 1
else:
ngal = 100000
L = 30. * s
tol_factor = 2
A = 0.05
rng = np.random.RandomState(8675309)
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.)
cat = treecorr.Catalog(x=x, y=y, k=kappa, x_units='arcmin', y_units='arcmin')
kk = treecorr.KKCorrelation(bin_size=0.1, min_sep=1., max_sep=20., sep_units='arcmin',
verbose=1)
kk.process(cat)
# log(<R>) != <logR>, but it should be close:
print('meanlogr - log(meanr) = ',kk.meanlogr - np.log(kk.meanr))
np.testing.assert_allclose(kk.meanlogr, np.log(kk.meanr), atol=1.e-3)
r = kk.meanr
true_xi = np.pi * A**2 * (s/L)**2 * np.exp(-0.25*r**2/s**2)
print('kk.xi = ',kk.xi)
print('true_xi = ',true_xi)
print('ratio = ',kk.xi / true_xi)
print('diff = ',kk.xi - true_xi)
print('max diff = ',max(abs(kk.xi - true_xi)))
print('max rel diff = ',max(abs((kk.xi - true_xi)/true_xi)))
np.testing.assert_allclose(kk.xi, true_xi, rtol=0.1*tol_factor)
# It should also work as a cross-correlation of this cat with itself
kk.process(cat,cat)
np.testing.assert_allclose(kk.meanlogr, np.log(kk.meanr), atol=1.e-3)
np.testing.assert_allclose(kk.xi, true_xi, rtol=0.1*tol_factor)
# Check that we get the same result using the corr2 function
cat.write(os.path.join('data','kk.dat'))
config = treecorr.read_config('configs/kk.yaml')
config['verbose'] = 0
config['precision'] = 8
treecorr.corr2(config)
corr2_output = np.genfromtxt(os.path.join('output','kk.out'), names=True, skip_header=1)
print('kk.xi = ',kk.xi)
print('from corr2 output = ',corr2_output['xi'])
print('ratio = ',corr2_output['xi']/kk.xi)
print('diff = ',corr2_output['xi']-kk.xi)
np.testing.assert_allclose(corr2_output['xi'], kk.xi, rtol=1.e-3)
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
# Check the fits write option
out_file_name = os.path.join('output','kk_out.fits')
kk.write(out_file_name)
data = fitsio.read(out_file_name)
np.testing.assert_almost_equal(data['r_nom'], np.exp(kk.logr))
np.testing.assert_almost_equal(data['meanr'], kk.meanr)
np.testing.assert_almost_equal(data['meanlogr'], kk.meanlogr)
np.testing.assert_almost_equal(data['xi'], kk.xi)
np.testing.assert_almost_equal(data['sigma_xi'], np.sqrt(kk.varxi))
np.testing.assert_almost_equal(data['weight'], kk.weight)
np.testing.assert_almost_equal(data['npairs'], kk.npairs)
# Check the read function
kk2 = treecorr.KKCorrelation(bin_size=0.1, min_sep=1., max_sep=100., sep_units='arcmin')
kk2.read(out_file_name)
np.testing.assert_almost_equal(kk2.logr, kk.logr)
np.testing.assert_almost_equal(kk2.meanr, kk.meanr)
np.testing.assert_almost_equal(kk2.meanlogr, kk.meanlogr)
np.testing.assert_almost_equal(kk2.xi, kk.xi)
np.testing.assert_almost_equal(kk2.varxi, kk.varxi)
np.testing.assert_almost_equal(kk2.weight, kk.weight)
np.testing.assert_almost_equal(kk2.npairs, kk.npairs)
assert kk2.coords == kk.coords
assert kk2.metric == kk.metric
assert kk2.sep_units == kk.sep_units
assert kk2.bin_type == kk.bin_type
def test_kk():
# cf. http://adsabs.harvard.edu/abs/2002A%26A...389..729S for the basic formulae I use here.
#
# Use kappa(r) = A exp(-r^2/2s^2)
#
# The Fourier transform is: kappa~(k) = 2 pi A s^2 exp(-s^2 k^2/2) / L^2
# P(k) = (1/2pi) <|kappa~(k)|^2> = 2 pi A^2 (s/L)^4 exp(-s^2 k^2)
# xi(r) = (1/2pi) int( dk k P(k) J0(kr) )
# = pi A^2 (s/L)^2 exp(-r^2/2s^2/4)
# Note: I'm not sure I handled the L factors correctly, but the units at the end need
# to be kappa^2, so it needs to be (s/L)^2.
ngal = 1000000
A = 0.05
s = 10.
L = 30. * s # Not infinity, so this introduces some error. Our integrals were to infinity.
numpy.random.seed(8675309)
x = (numpy.random.random_sample(ngal) - 0.5) * L
y = (numpy.random.random_sample(ngal) - 0.5) * L
r2 = (x**2 + y**2) / s**2
kappa = A * numpy.exp(-r2 / 2.)
cat = treecorr.Catalog(x=x,
y=y,
k=kappa,
x_units='arcmin',
y_units='arcmin')
kk = treecorr.KKCorrelation(bin_size=0.1,
min_sep=1.,
max_sep=50.,
sep_units='arcmin',
verbose=1)
kk.process(cat)
# log(<R>) != <logR>, but it should be close:
print('meanlogr - log(meanr) = ', kk.meanlogr - numpy.log(kk.meanr))
numpy.testing.assert_almost_equal(kk.meanlogr,
numpy.log(kk.meanr),
decimal=3)
r = kk.meanr
true_xi = numpy.pi * A**2 * (s / L)**2 * numpy.exp(-0.25 * r**2 / s**2)
print('kk.xi = ', kk.xi)
print('true_xi = ', true_xi)
print('ratio = ', kk.xi / true_xi)
print('diff = ', kk.xi - true_xi)
print('max diff = ', max(abs(kk.xi - true_xi)))
assert max(abs(kk.xi - true_xi)) < 5.e-7
# It should also work as a cross-correlation of this cat with itself
kk.process(cat, cat)
numpy.testing.assert_almost_equal(kk.meanlogr,
numpy.log(kk.meanr),
decimal=3)
assert max(abs(kk.xi - true_xi)) < 5.e-7
# Check that we get the same result using the corr2 executable:
if __name__ == '__main__':
cat.write(os.path.join('data', 'kk.dat'))
import subprocess
corr2_exe = get_script_name('corr2')
p = subprocess.Popen([corr2_exe, "kk.params"])
p.communicate()
corr2_output = numpy.genfromtxt(os.path.join('output', 'kk.out'),
names=True)
print('kk.xi = ', kk.xi)
print('from corr2 output = ', corr2_output['xi'])
print('ratio = ', corr2_output['xi'] / kk.xi)
print('diff = ', corr2_output['xi'] - kk.xi)
numpy.testing.assert_almost_equal(corr2_output['xi'] / kk.xi,
1.,
decimal=3)
# Check the fits write option
out_file_name = os.path.join('output', 'kk_out.fits')
kk.write(out_file_name)
try:
import fitsio
data = fitsio.read(out_file_name)
numpy.testing.assert_almost_equal(data['R_nom'], numpy.exp(kk.logr))
numpy.testing.assert_almost_equal(data['meanR'], kk.meanr)
numpy.testing.assert_almost_equal(data['meanlogR'], kk.meanlogr)
numpy.testing.assert_almost_equal(data['xi'], kk.xi)
numpy.testing.assert_almost_equal(data['sigma_xi'],
numpy.sqrt(kk.varxi))
numpy.testing.assert_almost_equal(data['weight'], kk.weight)
numpy.testing.assert_almost_equal(data['npairs'], kk.npairs)
except ImportError:
print('Unable to import fitsio. Skipping fits tests.')
# Check the read function
kk2 = treecorr.KKCorrelation(bin_size=0.1,
min_sep=1.,
max_sep=100.,
sep_units='arcmin')
kk2.read(out_file_name)
numpy.testing.assert_almost_equal(kk2.logr, kk.logr)
numpy.testing.assert_almost_equal(kk2.meanr, kk.meanr)
numpy.testing.assert_almost_equal(kk2.meanlogr, kk.meanlogr)
numpy.testing.assert_almost_equal(kk2.xi, kk.xi)
numpy.testing.assert_almost_equal(kk2.varxi, kk.varxi)
numpy.testing.assert_almost_equal(kk2.weight, kk.weight)
numpy.testing.assert_almost_equal(kk2.npairs, kk.npairs)
def test_large_scale():
# Test very large scales, comparing Arc, Euclidean (spherical), and Euclidean (3d)
# Distribute points uniformly in all directions.
if __name__ == '__main__':
ngal = 100000
tol = 1
nbins = 100
half = 50
else:
# Use fewer galaxies when running nosetests, so this is faster
ngal = 10000
tol = 3 # A factor by which we scale the tolerances to work with the smaller ngal.
nbins = 50
half = 25
s = 1.
rng = np.random.RandomState(8675309)
x = rng.normal(0, s, (ngal,) )
y = rng.normal(0, s, (ngal,) )
z = rng.normal(0, s, (ngal,) )
r = np.sqrt( x*x + y*y + z*z )
dec = np.arcsin(z/r)
ra = np.arctan2(y,x)
r = np.ones_like(x)
# Use x for "kappa" so there's a strong real correlation function
cat1 = treecorr.Catalog(ra=ra, dec=dec, k=x, ra_units='rad', dec_units='rad')
cat2 = treecorr.Catalog(ra=ra, dec=dec, k=x, r=r, ra_units='rad', dec_units='rad')
config = {
'min_sep' : 0.01,
'max_sep' : 1.57,
'nbins' : nbins,
'bin_slop' : 0.2,
'verbose' : 1
}
kk_sphere = treecorr.KKCorrelation(config)
kk_chord = treecorr.KKCorrelation(config)
kk_euclid = treecorr.KKCorrelation(config)
kk_euclid.process(cat1, metric='Euclidean')
kk_sphere.process(cat1, metric='Arc')
kk_chord.process(cat2, metric='Euclidean')
for tag in [ 'rnom', 'logr', 'meanr', 'meanlogr', 'npairs', 'xi' ]:
for name, dd in [ ('Euclid', kk_euclid), ('Sphere', kk_sphere), ('Chord', kk_chord) ]:
print(name, tag, '=', getattr(dd,tag))
# rnom and logr should be identical
np.testing.assert_array_equal(kk_sphere.rnom, kk_euclid.rnom)
np.testing.assert_array_equal(kk_chord.rnom, kk_euclid.rnom)
np.testing.assert_array_equal(kk_sphere.logr, kk_euclid.logr)
np.testing.assert_array_equal(kk_chord.logr, kk_euclid.logr)
# meanr should be similar for sphere and chord, but euclid is larger, since the chord
# distances have been scaled up to the real great circle distances
np.testing.assert_allclose(kk_sphere.meanr, kk_chord.meanr, rtol=1.e-3*tol)
np.testing.assert_allclose(kk_chord.meanr[:half], kk_euclid.meanr[:half], rtol=1.e-3*tol)
np.testing.assert_array_less(kk_chord.meanr[half:], kk_euclid.meanr[half:])
np.testing.assert_allclose(kk_sphere.meanlogr, kk_chord.meanlogr, atol=2.e-2*tol)
np.testing.assert_allclose(kk_chord.meanlogr[:half], kk_euclid.meanlogr[:half], atol=2.e-2*tol)
np.testing.assert_array_less(kk_chord.meanlogr[half:], kk_euclid.meanlogr[half:])
# npairs is basically the same for chord and euclid since the only difference there comes from
# differences in where they cut off the tree traversal, so the number of pairs is almost equal,
# even though the separations in each bin are given a different nominal distance.
# Sphere is smaller than both at all scales, since it is measuring the correlation
# function on larger real scales at each position.
print('diff (c-e)/e = ',(kk_chord.npairs-kk_euclid.npairs)/kk_euclid.npairs)
print('max = ',np.max(np.abs((kk_chord.npairs-kk_euclid.npairs)/kk_euclid.npairs)))
np.testing.assert_allclose(kk_chord.npairs, kk_euclid.npairs, rtol=1.e-3*tol)
print('diff (s-e)/e = ',(kk_sphere.npairs-kk_euclid.npairs)/kk_euclid.npairs)
np.testing.assert_allclose(kk_sphere.npairs[:half], kk_euclid.npairs[:half], rtol=3.e-3*tol)
np.testing.assert_array_less(kk_sphere.npairs[half:], kk_euclid.npairs[half:])
# Renormalize by the actual spacing in log(r)
renorm_euclid = kk_euclid.npairs / np.gradient(kk_euclid.meanlogr)
renorm_sphere = kk_sphere.npairs / np.gradient(kk_sphere.meanlogr)
# Then interpolate the euclid results to the values of the sphere distances
interp_euclid = np.interp(kk_sphere.meanlogr, kk_euclid.meanlogr, renorm_euclid)
# Matches at 3e-3 over whole range now.
print('interp_euclid = ',interp_euclid)
print('renorm_sphere = ',renorm_sphere)
print('new diff = ',(renorm_sphere-interp_euclid)/renorm_sphere)
print('max = ',np.max(np.abs((renorm_sphere-interp_euclid)/renorm_sphere)))
np.testing.assert_allclose(renorm_sphere, interp_euclid, rtol=3.e-3*tol)
# And almost the full range at the same precision.
np.testing.assert_allclose(renorm_sphere[:-4], interp_euclid[:-4], rtol=2.e-3*tol)
np.testing.assert_allclose(renorm_sphere, interp_euclid, rtol=1.e-2*tol)
# The xi values are similar. The euclid and chord values start out basically identical,
# but the distances are different. The euclid and the sphere are actually the same function
# so they match when rescaled to have the same distance values.
print('diff euclid, chord = ',(kk_chord.xi-kk_euclid.xi)/kk_euclid.xi)
print('max = ',np.max(np.abs((kk_chord.xi-kk_euclid.xi)/kk_euclid.xi)))
np.testing.assert_allclose(kk_chord.xi[:-8], kk_euclid.xi[:-8], rtol=1.e-3*tol)
np.testing.assert_allclose(kk_chord.xi, kk_euclid.xi, rtol=3.e-3*tol)
interp_euclid = np.interp(kk_sphere.meanlogr, kk_euclid.meanlogr, kk_euclid.xi)
print('interp_euclid = ',interp_euclid)
print('sphere.xi = ',kk_sphere.xi)
print('diff interp euclid, sphere = ',(kk_sphere.xi-interp_euclid))
print('max = ',np.max(np.abs((kk_sphere.xi-interp_euclid))))
np.testing.assert_allclose(kk_sphere.xi, interp_euclid, atol=1.e-3*tol)
def save_corrf(jx, jy, jz, nsp, verbose=False):
import pickle as pk
import os
import numpy as np
# from mpi4py import MPI
import time
import treecorr
import pickle as pk
import gc
import time
ti = time.time()
do_gg = True
minrad = 1.5
maxrad = 70.0
nrad = 15
nthreads = 64
bin_slop = 0.2
save_dir = '/global/project/projectdirs/m1727/shivamp_lsst/data_set/dc2_v1.0/measurements/mesh_grid/nsp' + str(
nsp) + '/'
do_mm, do_gm, do_gg = True, True, True
filename = save_dir + 'corrf_DC2_meshgrid_subsample_nsp' + str(
nsp) + '_' + str(jx) + '_' + str(jy) + '_' + str(jz) + '.pk'
if os.path.isfile(filename):
save_data = pk.load(open(filename, 'rb'))
all_keys = list(save_data.keys())
if ('xi_mm_full' in all_keys):
do_mm = False
xi_mm_full = save_data['xi_mm_full']
if ('xi_gm_full' in all_keys):
do_gm = False
xi_gm_full = save_data['xi_gm_full']
if ('xi_gg_full' in all_keys):
do_gg = False
xi_gg_full = save_data['xi_gg_full']
if do_mm + do_gm + do_gg == 0:
return 0
else:
save_data = {}
if verbose:
print('opening DM cat')
ldir = '/global/project/projectdirs/m1727/shivamp_lsst/data_set/dm_mesh/'
df = pk.load(
open(
ldir + 'mesh_dm_real_Ng3000_fullbox_nsp' + str(nsp) + '_' +
str(jx) + '_' + str(jy) + '_' + str(jz) + '.pk', 'rb'))
dm_density = np.ravel(df) - 1.
del df
gc.collect()
if verbose:
print('getting x,y,z meshgrid')
x1d = np.linspace(300 * jx, 300 * (jx + 1), 300)
y1d = np.linspace(300 * jy, 300 * (jy + 1), 300)
z1d = np.linspace(300 * jz, 300 * (jz + 1), 300)
xm, ym, zm = np.meshgrid(x1d, y1d, z1d)
xarr = np.ravel(xm)
del xm
gc.collect()
yarr = np.ravel(ym)
del ym
gc.collect()
zarr = np.ravel(zm)
del zm
gc.collect()
if do_gg or do_gm:
matter_cat = treecorr.Catalog(x=xarr, y=yarr, z=zarr, k=dm_density)
del dm_density
gc.collect()
if do_mm:
t1 = time.time()
m_m = treecorr.KKCorrelation(nbins=nrad,
min_sep=minrad,
max_sep=maxrad,
verbose=0,
num_threads=nthreads,
bin_slop=bin_slop)
if verbose:
print('doing matterxmatter calculation')
m_m.process(matter_cat, matter_cat)
if verbose:
print('this took ' + str(np.around(time.time() - t1, 2)) + 's')
xi_mm_full = m_m.xi
r_mm = np.exp(m_m.meanlogr)
print(r_mm)
print(xi_mm_full)
save_data = {'xi_mm_full': xi_mm_full, 'r_mm': r_mm}
import pickle as pk
pk.dump(save_data, open(filename, "wb"), protocol=2)
if verbose:
print('opening gal cat')
ldir = '/global/project/projectdirs/m1727/shivamp_lsst/data_set/gal_mesh_24p5/'
df = pk.load(
open(
ldir + 'mesh_gal_real_Ng3000_fullbox_nsp' + str(nsp) + '_' +
str(jx) + '_' + str(jy) + '_' + str(jz) + '.pk', 'rb'))
gal_density = np.ravel(df) - 1.
del df
gc.collect()
if do_gg or do_gm:
gal_cat = treecorr.Catalog(x=xarr, y=yarr, z=zarr, k=gal_density)
del gal_density, xarr, yarr, zarr
gc.collect()
if do_gm:
if verbose:
print('doing galaxyxmatter calculation')
t1 = time.time()
g_m = treecorr.KKCorrelation(nbins=nrad,
min_sep=minrad,
max_sep=maxrad,
verbose=0,
num_threads=nthreads,
bin_slop=bin_slop)
g_m.process(gal_cat, matter_cat)
print('this took ' + str(np.around(time.time() - t1, 2)) + 's')
del matter_cat
gc.collect()
xi_gm_full = g_m.xi
r_gm = np.exp(g_m.meanlogr)
if verbose:
print(r_gm)
print(xi_gm_full)
save_data['xi_gm_full'] = xi_gm_full
save_data['r_gm'] = r_gm
import pickle as pk
pk.dump(save_data, open(filename, "wb"), protocol=2)
if do_gg:
if verbose:
print('doing galaxyxgalaxy calculation')
t1 = time.time()
g_g = treecorr.KKCorrelation(nbins=nrad,
min_sep=minrad,
max_sep=maxrad,
verbose=0,
num_threads=nthreads,
bin_slop=bin_slop)
g_g.process(gal_cat, gal_cat)
if verbose:
print('this took ' + str(np.around(time.time() - t1, 2)) + 's')
del gal_cat
gc.collect()
xi_gg_full = g_g.xi
r_gg = np.exp(g_g.meanlogr)
if verbose:
print(r_gg)
print(xi_gg_full)
save_data['xi_gg_full'] = xi_gg_full
save_data['r_gg'] = r_gg
import pickle as pk
pk.dump(save_data, open(filename, "wb"), protocol=2)
if verbose:
print('getting the ratios gg/mm')
xi_gg_mm = (xi_gg_full) / (xi_mm_full)
save_data['xi_gg_mm'] = xi_gg_mm
if verbose:
print(xi_gg_mm)
import pickle as pk
pk.dump(save_data, open(filename, "wb"), protocol=2)
if verbose:
print('getting the ratios gm/mm')
xi_gm_mm = (xi_gm_full) / (xi_mm_full)
if verbose:
print(xi_gm_mm)
save_data['xi_gm_mm'] = xi_gm_mm
import pickle as pk
pk.dump(save_data, open(filename, "wb"), protocol=2)
if verbose:
tf = time.time()
print('Total pipeline took ' + str(np.around((tf - ti), 2)) + 's')
return 0
def compute_2pcf_2d(skymaps=None, mask=None, nmaps=20, nbins=64, imdim=1024, pixel_size=7., thetamax=2., \
with_noise=True, plot=True, datapath='/Users/richardfeder/Documents/ciber2/ciber/data/', tail_name='mmin=18.4_mmax=25', brute=False, \
bin_slop=0.05):
# for some reason, TreeCorr only works with radians as the units, so this always assumes radians unlike the 1d version
if skymaps is not None:
imdim = skymaps[0].shape[0]
if nmaps > len(skymaps):
nmaps = len(skymaps)
fs, wthetas_2d, varwthetas_2d, covs = [[] for x in range(4)]
deg_per_pix = (pixel_size / 3600.) # arcseconds to degrees
deg_per_rad = np.pi / 180.
xind = np.indices((imdim, imdim))[0].ravel() * deg_per_pix * deg_per_rad
yind = np.indices((imdim, imdim))[1].ravel() * deg_per_pix * deg_per_rad
print('min/max xind:', np.min(xind), np.max(xind))
print('thetamax:', imdim * deg_per_pix * deg_per_rad)
for i in range(nmaps):
if skymaps is None:
load_srcmap = load_ciber_srcmap(i,
with_noise=with_noise,
datapath=datapath,
tail_name=tail_name)
else:
print('skymaps is not None')
load_srcmap = skymaps[i]
skymap = load_srcmap - np.mean(load_srcmap)
xind_cat, yind_cat, ks, weights = get_treecorr_catalog_vals(
xind, yind, skymap, mask=mask, prefilter=False)
cat = treecorr.Catalog(x=xind_cat,
y=yind_cat,
k=ks,
w=weights,
npatch=1)
if brute:
kk = treecorr.KKCorrelation(nbins=nbins,
bin_type='TwoD',
max_sep=imdim * deg_per_pix *
deg_per_rad,
sep_units='rad',
brute=True)
else:
kk = treecorr.KKCorrelation(nbins=nbins,
bin_type='TwoD',
max_sep=imdim * deg_per_pix *
deg_per_rad,
sep_units='rad',
bin_slop=bin_slop)
kk.process(cat, metric='Euclidean')
print('bin slop is ', kk.bin_slop)
if plot:
plt.figure()
plt.title('npairs, 256x256 image, 256^2 bins')
plt.imshow(kk.npairs, vmin=0, vmax=10)
plt.colorbar()
plt.show()
plt.figure()
plt.hist(kk.npairs.ravel(), bins=np.linspace(-1, 10, 10))
plt.show()
plt.figure()
plt.imshow(kk.xi,
vmin=np.percentile(kk.xi, 1),
vmax=np.percentile(kk.xi, 99))
plt.colorbar()
plt.show()
wthetas_2d.append(kk.xi)
wthetas_2d = np.array(wthetas_2d)
if plot:
f = plot_w_thetax_thetay(wthetas_2d)
return wthetas_2d, f
return wthetas_2d
def corr2(config, logger=None):
"""Run the full two-point correlation function code based on the parameters in the
given config dict.
The function print_corr2_params() will output information about the valid parameters
that are expected to be in the config dict.
Optionally a logger parameter maybe given, in which case it is used for logging.
If not given, the logging will be based on the verbose and log_file parameters.
:param config: The configuration dict which defines what to do.
:param logger: If desired, a logger object for logging. (default: None, in which case
one will be built according to the config dict's verbose level.)
"""
# Setup logger based on config verbose value
if logger is None:
logger = treecorr.config.setup_logger(
treecorr.config.get(config,'verbose',int,1),
config.get('log_file',None))
# Check that config doesn't have any extra parameters.
# (Such values are probably typos.)
# Also convert the given parameters to the correct type, etc.
config = treecorr.config.check_config(config, corr2_valid_params, corr2_aliases, logger)
import pprint
logger.debug('Using configuration dict:\n%s',pprint.pformat(config))
if ( 'output_dots' not in config
and config.get('log_file',None) is None
and config['verbose'] >= 2 ):
config['output_dots'] = True
# Set the number of threads
num_threads = config.get('num_threads',None)
logger.debug('From config dict, num_threads = %s',num_threads)
treecorr.set_omp_threads(num_threads, logger)
# Read in the input files. Each of these is a list.
cat1 = treecorr.read_catalogs(config, 'file_name', 'file_list', 0, logger)
if len(cat1) == 0:
raise AttributeError("Either file_name or file_list is required")
cat2 = treecorr.read_catalogs(config, 'file_name2', 'file_list2', 1, logger)
rand1 = treecorr.read_catalogs(config, 'rand_file_name', 'rand_file_list', 0, logger)
rand2 = treecorr.read_catalogs(config, 'rand_file_name2', 'rand_file_list2', 1, logger)
if len(cat2) == 0 and len(rand2) > 0:
raise AttributeError("rand_file_name2 is invalid without file_name2")
logger.info("Done reading input catalogs")
# Do GG correlation function if necessary
if 'gg_file_name' in config or 'm2_file_name' in config:
logger.warning("Performing GG calculations...")
gg = treecorr.GGCorrelation(config,logger)
gg.process(cat1,cat2)
logger.info("Done GG calculations.")
if 'gg_file_name' in config:
gg.write(config['gg_file_name'])
logger.warning("Wrote GG correlation to %s",config['gg_file_name'])
if 'm2_file_name' in config:
gg.writeMapSq(config['m2_file_name'], m2_uform=config['m2_uform'])
logger.warning("Wrote Mapsq values to %s",config['m2_file_name'])
# Do NG correlation function if necessary
if 'ng_file_name' in config or 'nm_file_name' in config or 'norm_file_name' in config:
if len(cat2) == 0:
raise AttributeError("file_name2 is required for ng correlation")
logger.warning("Performing NG calculations...")
ng = treecorr.NGCorrelation(config,logger)
ng.process(cat1,cat2)
logger.info("Done NG calculation.")
# The default ng_statistic is compensated _iff_ rand files are given.
rg = None
if len(rand1) == 0:
if config.get('ng_statistic',None) == 'compensated':
raise AttributeError("rand_files is required for ng_statistic = compensated")
elif config.get('ng_statistic','compensated') == 'compensated':
rg = treecorr.NGCorrelation(config,logger)
rg.process(rand1,cat2)
logger.info("Done RG calculation.")
if 'ng_file_name' in config:
ng.write(config['ng_file_name'], rg)
logger.warning("Wrote NG correlation to %s",config['ng_file_name'])
if 'nm_file_name' in config:
ng.writeNMap(config['nm_file_name'], rg, m2_uform=config['m2_uform'])
logger.warning("Wrote NMap values to %s",config['nm_file_name'])
if 'norm_file_name' in config:
gg = treecorr.GGCorrelation(config,logger)
gg.process(cat2)
logger.info("Done GG calculation for norm")
dd = treecorr.NNCorrelation(config,logger)
dd.process(cat1)
logger.info("Done DD calculation for norm")
rr = treecorr.NNCorrelation(config,logger)
rr.process(rand1)
logger.info("Done RR calculation for norm")
dr = None
if config['nn_statistic'] == 'compensated':
dr = treecorr.NNCorrelation(config,logger)
dr.process(cat1,rand1)
logger.info("Done DR calculation for norm")
ng.writeNorm(config['norm_file_name'],gg,dd,rr,dr,rg,m2_uform=config['m2_uform'])
logger.warning("Wrote Norm values to %s",config['norm_file_name'])
# Do NN correlation function if necessary
if 'nn_file_name' in config:
if len(rand1) == 0:
raise AttributeError("rand_file_name is required for NN correlation")
if len(cat2) > 0 and len(rand2) == 0:
raise AttributeError("rand_file_name2 is required for NN cross-correlation")
logger.warning("Performing DD calculations...")
dd = treecorr.NNCorrelation(config,logger)
dd.process(cat1,cat2)
logger.info("Done DD calculations.")
dr = None
rd = None
if len(cat2) == 0:
logger.warning("Performing RR calculations...")
rr = treecorr.NNCorrelation(config,logger)
rr.process(rand1)
logger.info("Done RR calculations.")
if config['nn_statistic'] == 'compensated':
logger.warning("Performing DR calculations...")
dr = treecorr.NNCorrelation(config,logger)
dr.process(cat1,rand1)
logger.info("Done DR calculations.")
else:
logger.warning("Performing RR calculations...")
rr = treecorr.NNCorrelation(config,logger)
rr.process(rand1,rand2)
logger.info("Done RR calculations.")
if config['nn_statistic'] == 'compensated':
logger.warning("Performing DR calculations...")
dr = treecorr.NNCorrelation(config,logger)
dr.process(cat1,rand2)
logger.info("Done DR calculations.")
rd = treecorr.NNCorrelation(config,logger)
rd.process(rand1,cat2)
logger.info("Done RD calculations.")
dd.write(config['nn_file_name'],rr,dr,rd)
logger.warning("Wrote NN correlation to %s",config['nn_file_name'])
# Do KK correlation function if necessary
if 'kk_file_name' in config:
logger.warning("Performing KK calculations...")
kk = treecorr.KKCorrelation(config,logger)
kk.process(cat1,cat2)
logger.info("Done KK calculations.")
kk.write(config['kk_file_name'])
logger.warning("Wrote KK correlation to %s",config['kk_file_name'])
# Do NG correlation function if necessary
if 'nk_file_name' in config:
if len(cat2) == 0:
raise AttributeError("file_name2 is required for nk correlation")
logger.warning("Performing NK calculations...")
nk = treecorr.NKCorrelation(config,logger)
nk.process(cat1,cat2)
logger.info("Done NK calculation.")
rk = None
if len(rand1) == 0:
if config.get('nk_statistic',None) == 'compensated':
raise AttributeError("rand_files is required for nk_statistic = compensated")
elif config.get('nk_statistic','compensated') == 'compensated':
rk = treecorr.NKCorrelation(config,logger)
rk.process(rand1,cat2)
logger.info("Done RK calculation.")
nk.write(config['nk_file_name'], rk)
logger.warning("Wrote NK correlation to %s",config['nk_file_name'])
# Do KG correlation function if necessary
if 'kg_file_name' in config:
if len(cat2) == 0:
raise AttributeError("file_name2 is required for kg correlation")
logger.warning("Performing KG calculations...")
kg = treecorr.KGCorrelation(config,logger)
kg.process(cat1,cat2)
logger.info("Done KG calculation.")
kg.write(config['kg_file_name'])
logger.warning("Wrote KG correlation to %s",config['kg_file_name'])
def test_twod():
try:
from scipy.spatial.distance import pdist, squareform
except ImportError:
print('Skipping test_twod, since uses scipy, and scipy is not installed.')
return
# N random points in 2 dimensions
rng = np.random.RandomState(8675309)
N = 200
x = rng.uniform(-20, 20, N)
y = rng.uniform(-20, 20, N)
# Give the points a multivariate Gaussian random field for kappa and gamma
L1 = [[0.33, 0.09], [-0.01, 0.26]] # Some arbitrary correlation matrix
invL1 = np.linalg.inv(L1)
dists = pdist(np.array([x,y]).T, metric='mahalanobis', VI=invL1)
K = np.exp(-0.5 * dists**2)
K = squareform(K)
np.fill_diagonal(K, 1.)
A = 2.3
kappa = rng.multivariate_normal(np.zeros(N), K*(A**2))
# Add some noise
sigma = A/10.
kappa += rng.normal(scale=sigma, size=N)
kappa_err = np.ones_like(kappa) * sigma
# Make gamma too
gamma1 = rng.multivariate_normal(np.zeros(N), K*(A**2))
gamma1 += rng.normal(scale=sigma, size=N)
gamma2 = rng.multivariate_normal(np.zeros(N), K*(A**2))
gamma2 += rng.normal(scale=sigma, size=N)
gamma = gamma1 + 1j * gamma2
gamma_err = kappa_err
# Calculate the 2D correlation using brute force
max_sep = 21.
nbins = 21
xi_brut = corr2d(x, y, kappa, kappa, w=None, rmax=max_sep, bins=nbins)
cat1 = treecorr.Catalog(x=x, y=y, k=kappa, g1=gamma1, g2=gamma2)
kk = treecorr.KKCorrelation(min_sep=0., max_sep=max_sep, nbins=nbins, bin_type='TwoD',
brute=True)
# First the simplest case to get right: cross correlation of the catalog with itself.
kk.process(cat1, cat1)
print('max abs diff = ',np.max(np.abs(kk.xi - xi_brut)))
print('max rel diff = ',np.max(np.abs(kk.xi - xi_brut)/np.abs(kk.xi)))
np.testing.assert_allclose(kk.xi, xi_brut, atol=1.e-7)
# Auto-correlation should do the same thing.
kk.process(cat1)
print('max abs diff = ',np.max(np.abs(kk.xi - xi_brut)))
print('max rel diff = ',np.max(np.abs(kk.xi - xi_brut)/np.abs(kk.xi)))
np.testing.assert_allclose(kk.xi, xi_brut, atol=1.e-7)
# Repeat with weights.
xi_brut = corr2d(x, y, kappa, kappa, w=1./kappa_err**2, rmax=max_sep, bins=nbins)
cat2 = treecorr.Catalog(x=x, y=y, k=kappa, g1=gamma1, g2=gamma2, w=1./kappa_err**2)
# NB. Testing that min_sep = 0 is default
kk = treecorr.KKCorrelation(max_sep=max_sep, nbins=nbins, bin_type='TwoD', brute=True)
kk.process(cat2, cat2)
print('max abs diff = ',np.max(np.abs(kk.xi - xi_brut)))
print('max rel diff = ',np.max(np.abs(kk.xi - xi_brut)/np.abs(kk.xi)))
np.testing.assert_allclose(kk.xi, xi_brut, atol=1.e-7)
kk.process(cat2)
np.testing.assert_allclose(kk.xi, xi_brut, atol=1.e-7)
# Check GG
xi_brut = corr2d(x, y, gamma, np.conj(gamma), rmax=max_sep, bins=nbins)
# Equivalent bin_size = 2. Check omitting nbins
gg = treecorr.GGCorrelation(max_sep=max_sep, bin_size=2., bin_type='TwoD', brute=True)
gg.process(cat1)
print('max abs diff = ',np.max(np.abs(gg.xip - xi_brut)))
print('max rel diff = ',np.max(np.abs(gg.xip - xi_brut)/np.abs(gg.xip)))
np.testing.assert_allclose(gg.xip, xi_brut, atol=2.e-7)
xi_brut = corr2d(x, y, gamma, np.conj(gamma), w=1./kappa_err**2, rmax=max_sep, bins=nbins)
# Check omitting max_sep
gg = treecorr.GGCorrelation(bin_size=2, nbins=nbins, bin_type='TwoD', brute=True)
gg.process(cat2)
print('max abs diff = ',np.max(np.abs(gg.xip - xi_brut)))
print('max rel diff = ',np.max(np.abs(gg.xip - xi_brut)/np.abs(gg.xip)))
np.testing.assert_allclose(gg.xip, xi_brut, atol=2.e-7)
# Check NK
xi_brut = corr2d(x, y, np.ones_like(kappa), kappa, rmax=max_sep, bins=nbins)
# Check slightly larger bin_size gets rounded down
nk = treecorr.NKCorrelation(max_sep=max_sep, bin_size=2.05, bin_type='TwoD', brute=True)
nk.process(cat1, cat1)
print('max abs diff = ',np.max(np.abs(nk.xi - xi_brut)))
print('max rel diff = ',np.max(np.abs(nk.xi - xi_brut)/np.abs(nk.xi)))
np.testing.assert_allclose(nk.xi, xi_brut, atol=1.e-7)
xi_brut = corr2d(x, y, np.ones_like(kappa), kappa, w=1./kappa_err**2, rmax=max_sep, bins=nbins)
# Check very small, but non-zeo min_sep
nk = treecorr.NKCorrelation(min_sep=1.e-6, max_sep=max_sep, nbins=nbins, bin_type='TwoD', brute=True)
nk.process(cat2, cat2)
print('max abs diff = ',np.max(np.abs(nk.xi - xi_brut)))
print('max rel diff = ',np.max(np.abs(nk.xi - xi_brut)/np.abs(nk.xi)))
np.testing.assert_allclose(nk.xi, xi_brut, atol=1.e-7)
# Check NN
xi_brut, counts = corr2d(x, y, np.ones_like(kappa), np.ones_like(kappa),
rmax=max_sep, bins=nbins, return_counts=True)
nn = treecorr.NNCorrelation(max_sep=max_sep, nbins=nbins, bin_type='TwoD', brute=True)
nn.process(cat1)
print('max abs diff = ',np.max(np.abs(nn.npairs - counts)))
print('max rel diff = ',np.max(np.abs(nn.npairs - counts)/np.abs(nn.npairs)))
np.testing.assert_allclose(nn.npairs, counts, atol=1.e-7)
nn.process(cat1, cat1)
print('max abs diff = ',np.max(np.abs(nn.npairs - counts)))
print('max rel diff = ',np.max(np.abs(nn.npairs - counts)/np.abs(nn.npairs)))
np.testing.assert_allclose(nn.npairs, counts, atol=1.e-7)
xi_brut, counts = corr2d(x, y, np.ones_like(kappa), np.ones_like(kappa),
w=1./kappa_err**2, rmax=max_sep, bins=nbins, return_counts=True)
nn = treecorr.NNCorrelation(max_sep=max_sep, nbins=nbins, bin_type='TwoD', brute=True)
nn.process(cat2)
print('max abs diff = ',np.max(np.abs(nn.weight - counts)))
print('max rel diff = ',np.max(np.abs(nn.weight - counts)/np.abs(nn.weight)))
np.testing.assert_allclose(nn.weight, counts, atol=1.e-7)
nn.process(cat2, cat2)
print('max abs diff = ',np.max(np.abs(nn.weight - counts)))
print('max rel diff = ',np.max(np.abs(nn.weight - counts)/np.abs(nn.weight)))
np.testing.assert_allclose(nn.weight, counts, atol=1.e-7)
# The other two, NG and KG can't really be checked with the brute force
# calculator we have here, so we're counting on the above being a sufficient
# test of all aspects of the twod binning. I think that it is sufficient, but I
# admit I would prefer if we had a real test of these other two pairs, along with
# xi- for GG.
# Check invalid constructors
assert_raises(TypeError, treecorr.NNCorrelation, max_sep=max_sep, nbins=nbins, bin_size=2,
bin_type='TwoD')
assert_raises(TypeError, treecorr.NNCorrelation, nbins=nbins, bin_type='TwoD')
assert_raises(TypeError, treecorr.NNCorrelation, bin_size=2, bin_type='TwoD')
assert_raises(TypeError, treecorr.NNCorrelation, max_sep=max_sep, bin_type='TwoD')
