def test_direct():
nobj = 5000
rng = np.random.RandomState(8675309)
x = rng.random_sample(nobj)
y = rng.random_sample(nobj)
ra = rng.random_sample(nobj)
dec = rng.random_sample(nobj)
w = rng.random_sample(nobj)
g1 = rng.random_sample(nobj)
g2 = rng.random_sample(nobj)
k = rng.random_sample(nobj)
cat1 = treecorr.Catalog(x=x, y=y, w=w, g1=g1, g2=g2, k=k)
np.testing.assert_almost_equal(cat1.x, x)
np.testing.assert_almost_equal(cat1.y, y)
np.testing.assert_almost_equal(cat1.w, w)
np.testing.assert_almost_equal(cat1.g1, g1)
np.testing.assert_almost_equal(cat1.g2, g2)
np.testing.assert_almost_equal(cat1.k, k)
cat2 = treecorr.Catalog(ra=ra, dec=dec, w=w, g1=g1, g2=g2, k=k,
ra_units='hours', dec_units='degrees')
np.testing.assert_almost_equal(cat2.ra, ra * coord.hours / coord.radians)
np.testing.assert_almost_equal(cat2.dec, dec * coord.degrees / coord.radians)
np.testing.assert_almost_equal(cat2.w, w)
np.testing.assert_almost_equal(cat2.g1, g1)
np.testing.assert_almost_equal(cat2.g2, g2)
np.testing.assert_almost_equal(cat2.k, k)
do_pickle(cat1)
do_pickle(cat2)
assert_raises(TypeError, treecorr.Catalog, x=x)
assert_raises(TypeError, treecorr.Catalog, y=y)
assert_raises(TypeError, treecorr.Catalog, z=x)
assert_raises(TypeError, treecorr.Catalog, r=x)
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, r=x)
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, ra=ra, dec=dec)
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, ra=ra, dec=dec,
ra_units='hours', dec_units='degrees')
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, ra_units='hours')
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, dec_units='degrees')
assert_raises(TypeError, treecorr.Catalog, ra=ra, ra_units='hours')
assert_raises(TypeError, treecorr.Catalog, dec=dec, dec_units='degrees')
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, g1=g1)
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, g2=g2)
assert_raises(TypeError, treecorr.Catalog, ra=ra, dec=dec,
ra_units='hours', dec_units='degrees', x_units='arcmin')
assert_raises(TypeError, treecorr.Catalog, ra=ra, dec=dec,
ra_units='hours', dec_units='degrees', y_units='arcmin')
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, x_units='arcmin')
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, y_units='arcmin')
assert_raises(TypeError, treecorr.Catalog, x=x, y=y, z=x, z_units='arcmin')
assert_raises(ValueError, treecorr.Catalog, x=x, y=y[4:])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, z=x[4:])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=w[4:])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=w, wpos=w[4:])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=w, g1=g1[4:], g2=g2[4:])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=w, g1=g1[4:], g2=g2)
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=w, g1=g1, g2=g2[4:])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=w, k=k[4:])
assert_raises(ValueError, treecorr.Catalog, ra=ra, dec=dec[4:], w=w, g1=g1, g2=g2, k=k,
ra_units='hours', dec_units='degrees')
assert_raises(ValueError, treecorr.Catalog, ra=ra[4:], dec=dec, w=w, g1=g1, g2=g2, k=k,
ra_units='hours', dec_units='degrees')
assert_raises(ValueError, treecorr.Catalog, ra=ra, dec=dec, r=x[4:], w=w, g1=g1, g2=g2, k=k,
ra_units='hours', dec_units='degrees')
assert_raises(ValueError, treecorr.Catalog, x=[], y=[])
assert_raises(ValueError, treecorr.Catalog, x=x, y=y, w=np.zeros_like(x))
def test_ascii():
nobj = 5000
rng = np.random.RandomState(8675309)
x = rng.random_sample(nobj)
y = rng.random_sample(nobj)
z = rng.random_sample(nobj)
ra = rng.random_sample(nobj)
dec = rng.random_sample(nobj)
r = rng.random_sample(nobj)
wpos = rng.random_sample(nobj)
g1 = rng.random_sample(nobj)
g2 = rng.random_sample(nobj)
k = rng.random_sample(nobj)
# Some elements have both w and wpos = 0.
w = wpos.copy()
use = rng.randint(30, size=nobj).astype(float)
w[use == 0] = 0
wpos[use == 0] = 0
# Others just have w = 0
use = rng.randint(30, size=nobj).astype(float)
w[use == 0] = 0
flags = np.zeros(nobj).astype(int)
for flag in [ 1, 2, 4, 8, 16 ]:
sub = rng.random_sample(nobj) < 0.1
flags[sub] = np.bitwise_or(flags[sub], flag)
file_name = os.path.join('data','test.dat')
with open(file_name, 'w') as fid:
# These are intentionally in a different order from the order we parse them.
fid.write('# ra,dec,x,y,k,g1,g2,w,z,r,wpos,flag\n')
for i in range(nobj):
fid.write((('%.8f '*11)+'%d\n')%(
ra[i],dec[i],x[i],y[i],k[i],g1[i],g2[i],w[i],z[i],r[i],wpos[i],flags[i]))
# Check basic input
config = {
'x_col' : 3,
'y_col' : 4,
'z_col' : 9,
'x_units' : 'rad',
'y_units' : 'rad',
'w_col' : 8,
'wpos_col' : 11,
'k_col' : 5,
'g1_col' : 6,
'g2_col' : 7,
'kk_file_name' : 'kk.out', # These make sure k and g are required.
'gg_file_name' : 'gg.out',
}
cat1 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat1.x, x)
np.testing.assert_almost_equal(cat1.y, y)
np.testing.assert_almost_equal(cat1.z, z)
np.testing.assert_almost_equal(cat1.w, w)
np.testing.assert_almost_equal(cat1.g1, g1)
np.testing.assert_almost_equal(cat1.g2, g2)
np.testing.assert_almost_equal(cat1.k, k)
np.testing.assert_almost_equal(cat1.wpos, wpos)
assert_raises(TypeError, treecorr.Catalog, file_name)
assert_raises(TypeError, treecorr.Catalog, file_name, config, x=x)
assert_raises(TypeError, treecorr.Catalog, file_name, config, y=y)
assert_raises(TypeError, treecorr.Catalog, file_name, config, z=z)
assert_raises(TypeError, treecorr.Catalog, file_name, config, ra=ra)
assert_raises(TypeError, treecorr.Catalog, file_name, config, dec=dec)
assert_raises(TypeError, treecorr.Catalog, file_name, config, r=r)
assert_raises(TypeError, treecorr.Catalog, file_name, config, g2=g2)
assert_raises(TypeError, treecorr.Catalog, file_name, config, k=k)
assert_raises(TypeError, treecorr.Catalog, file_name, config, w=w)
assert_raises(TypeError, treecorr.Catalog, file_name, config, wpos=wpos)
assert_raises(TypeError, treecorr.Catalog, file_name, config, flag=flag)
assert_raises(ValueError, treecorr.Catalog, file_name, config, file_type='Invalid')
assert_raises(TypeError, treecorr.Catalog, file_name, config, x_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, x_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, y_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, y_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, z_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, z_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, w_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, w_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, wpos_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, wpos_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, k_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, k_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, g1_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, g1_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, g2_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, g2_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, flag_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, flag_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, ra_col=4)
assert_raises(TypeError, treecorr.Catalog, file_name, config, dec_col=4)
assert_raises(TypeError, treecorr.Catalog, file_name, config, r_col=4)
assert_raises(TypeError, treecorr.Catalog, file_name, config, x_col=0)
assert_raises(TypeError, treecorr.Catalog, file_name, config, y_col=0)
assert_raises(TypeError, treecorr.Catalog, file_name, config, x_col=0, y_col=0, z_col=0)
# Check flags
config['flag_col'] = 12
print('config = ',config)
cat2 = treecorr.Catalog(file_name, config, file_type='ASCII')
np.testing.assert_almost_equal(cat2.w[flags==0], w[flags==0])
np.testing.assert_almost_equal(cat2.w[flags!=0], 0.)
# Check ok_flag
config['ok_flag'] = 4
cat3 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat3.w[np.logical_or(flags==0, flags==4)],
w[np.logical_or(flags==0, flags==4)])
np.testing.assert_almost_equal(cat3.w[np.logical_and(flags!=0, flags!=4)], 0.)
# Check ignore_flag
del config['ok_flag']
config['ignore_flag'] = 16
cat4 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat4.w[flags < 16], w[flags < 16])
np.testing.assert_almost_equal(cat4.w[flags >= 16], 0.)
# If weight is missing, automatically make it when there are flags
del config['w_col']
del config['wpos_col']
cat4 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat4.w[flags < 16], 1.)
np.testing.assert_almost_equal(cat4.w[flags >= 16], 0.)
config['w_col'] = 8 # Put them back for later.
config['wpos_col'] = 11
# Check different units for x,y
config['x_units'] = 'arcsec'
config['y_units'] = 'arcsec'
del config['z_col']
cat5 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat5.x, x * (pi/180./3600.))
np.testing.assert_almost_equal(cat5.y, y * (pi/180./3600.))
config['x_units'] = 'arcmin'
config['y_units'] = 'arcmin'
cat5 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat5.x, x * (pi/180./60.))
np.testing.assert_almost_equal(cat5.y, y * (pi/180./60.))
config['x_units'] = 'deg'
config['y_units'] = 'deg'
cat5 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat5.x, x * (pi/180.))
np.testing.assert_almost_equal(cat5.y, y * (pi/180.))
del config['x_units'] # Default is radians
del config['y_units']
cat5 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat5.x, x)
np.testing.assert_almost_equal(cat5.y, y)
# Check ra,dec
del config['x_col']
del config['y_col']
config['ra_col'] = 1
config['dec_col'] = 2
config['r_col'] = 10
config['ra_units'] = 'rad'
config['dec_units'] = 'rad'
cat6 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat6.ra, ra)
np.testing.assert_almost_equal(cat6.dec, dec)
config['ra_units'] = 'deg'
config['dec_units'] = 'deg'
cat6 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat6.ra, ra * (pi/180.))
np.testing.assert_almost_equal(cat6.dec, dec * (pi/180.))
config['ra_units'] = 'hour'
config['dec_units'] = 'deg'
cat6 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat6.ra, ra * (pi/12.))
np.testing.assert_almost_equal(cat6.dec, dec * (pi/180.))
assert_raises(TypeError, treecorr.Catalog, file_name, config, ra_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, ra_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, dec_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, dec_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, r_col=-1)
assert_raises(TypeError, treecorr.Catalog, file_name, config, r_col=100)
assert_raises(TypeError, treecorr.Catalog, file_name, config, x_col=4)
assert_raises(TypeError, treecorr.Catalog, file_name, config, y_col=4)
assert_raises(TypeError, treecorr.Catalog, file_name, config, z_col=4)
# Check using a different delimiter, comment marker
csv_file_name = os.path.join('data','test.csv')
with open(csv_file_name, 'w') as fid:
# These are intentionally in a different order from the order we parse them.
fid.write('% This file uses commas for its delimiter')
fid.write('% And more than one header line.')
fid.write('% Plus some extra comment lines every so often.')
fid.write('% And we use a weird comment marker to boot.')
fid.write('% ra,dec,x,y,k,g1,g2,w,flag\n')
for i in range(nobj):
fid.write((('%.8f,'*11)+'%d\n')%(
ra[i],dec[i],x[i],y[i],k[i],g1[i],g2[i],w[i],z[i],r[i],wpos[i],flags[i]))
if i%100 == 0:
fid.write('%%%% Line %d\n'%i)
config['delimiter'] = ','
config['comment_marker'] = '%'
cat7 = treecorr.Catalog(csv_file_name, config)
np.testing.assert_almost_equal(cat7.ra, ra * (pi/12.))
np.testing.assert_almost_equal(cat7.dec, dec * (pi/180.))
np.testing.assert_almost_equal(cat7.r, r)
np.testing.assert_almost_equal(cat7.g1, g1)
np.testing.assert_almost_equal(cat7.g2, g2)
np.testing.assert_almost_equal(cat7.w[flags < 16], w[flags < 16])
np.testing.assert_almost_equal(cat7.w[flags >= 16], 0.)
# Check flip_g1, flip_g2
del config['delimiter']
del config['comment_marker']
config['flip_g1'] = True
cat8 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat8.g1, -g1)
np.testing.assert_almost_equal(cat8.g2, g2)
config['flip_g2'] = 'true'
cat8 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat8.g1, -g1)
np.testing.assert_almost_equal(cat8.g2, -g2)
config['flip_g1'] = 'n'
config['flip_g2'] = 'yes'
cat8 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat8.g1, g1)
np.testing.assert_almost_equal(cat8.g2, -g2)
# Check overriding values with kwargs
cat8 = treecorr.Catalog(file_name, config, flip_g1=True, flip_g2=False)
np.testing.assert_almost_equal(cat8.g1, -g1)
np.testing.assert_almost_equal(cat8.g2, g2)
# Check copy command
cat9 = cat8.copy()
np.testing.assert_almost_equal(cat9.ra, cat8.ra)
np.testing.assert_almost_equal(cat9.dec, cat8.dec)
np.testing.assert_almost_equal(cat9.r, cat8.r)
np.testing.assert_almost_equal(cat9.g1, cat8.g1)
np.testing.assert_almost_equal(cat9.g2, cat8.g2)
np.testing.assert_almost_equal(cat9.w, cat8.w)
# Swapping w and wpos leads to zeros being copied from wpos to w
cat10 = treecorr.Catalog(file_name, config, w_col=11, wpos_col=8, flag_col=0)
np.testing.assert_almost_equal(cat10.wpos, w)
np.testing.assert_almost_equal(cat10.w, w)
# And if there is wpos, but no w, copy over the zeros, but not the other values
with CaptureLog() as cl:
cat10 = treecorr.Catalog(file_name, config, w_col=0, wpos_col=11, flag_col=0,
logger=cl.logger)
np.testing.assert_almost_equal(cat10.wpos, wpos)
np.testing.assert_almost_equal(cat10.w[wpos==0], 0)
np.testing.assert_almost_equal(cat10.w[wpos!=0], 1)
assert 'Some wpos values are zero, setting w=0 for these points' in cl.output
do_pickle(cat1)
do_pickle(cat2)
do_pickle(cat3)
do_pickle(cat4)
do_pickle(cat5)
do_pickle(cat6)
do_pickle(cat7)
do_pickle(cat8)
do_pickle(cat9)
do_pickle(cat10)
def test_fits():
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
get_from_wiki('Aardvark.fit')
file_name = os.path.join('data','Aardvark.fit')
config = treecorr.read_config('Aardvark.yaml')
config['verbose'] = 1
config['kk_file_name'] = 'kk.fits'
config['gg_file_name'] = 'gg.fits'
# Just test a few random particular values
cat1 = treecorr.Catalog(file_name, config)
np.testing.assert_equal(len(cat1.ra), 390935)
np.testing.assert_equal(cat1.nobj, 390935)
np.testing.assert_almost_equal(cat1.ra[0], 56.4195 * (pi/180.))
np.testing.assert_almost_equal(cat1.ra[390934], 78.4782 * (pi/180.))
np.testing.assert_almost_equal(cat1.dec[290333], 83.1579 * (pi/180.))
np.testing.assert_almost_equal(cat1.g1[46392], 0.0005066675)
np.testing.assert_almost_equal(cat1.g2[46392], -0.0001006742)
np.testing.assert_almost_equal(cat1.k[46392], -0.0008628797)
assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, dec_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, r_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, w_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, wpos_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, flag_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g1_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g2_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, k_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col='0')
assert_raises(ValueError, treecorr.Catalog, file_name, config, dec_col='0')
assert_raises(ValueError, treecorr.Catalog, file_name, config, x_col='x')
assert_raises(ValueError, treecorr.Catalog, file_name, config, y_col='y')
assert_raises(ValueError, treecorr.Catalog, file_name, config, z_col='z')
assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col='0', dec_col='0')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g1_col='0')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g2_col='0')
assert_raises(ValueError, treecorr.Catalog, file_name, config, k_col='0')
assert_raises(TypeError, treecorr.Catalog, file_name, config, x_units='arcmin')
assert_raises(TypeError, treecorr.Catalog, file_name, config, y_units='arcmin')
del config['ra_units']
assert_raises(TypeError, treecorr.Catalog, file_name, config)
del config['dec_units']
assert_raises(TypeError, treecorr.Catalog, file_name, config, ra_units='deg')
# The catalog doesn't have x, y, or w, but test that functionality as well.
del config['ra_col']
del config['dec_col']
config['x_col'] = 'RA'
config['y_col'] = 'DEC'
config['w_col'] = 'MU'
config['flag_col'] = 'INDEX'
config['ignore_flag'] = 64
cat2 = treecorr.Catalog(file_name, config)
np.testing.assert_almost_equal(cat2.x[390934], 78.4782, decimal=4)
np.testing.assert_almost_equal(cat2.y[290333], 83.1579, decimal=4)
np.testing.assert_almost_equal(cat2.w[46392], 0.) # index = 1200379
np.testing.assert_almost_equal(cat2.w[46393], 0.9995946) # index = 1200386
assert_raises(ValueError, treecorr.Catalog, file_name, config, x_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, y_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, z_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, dec_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, r_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, w_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, wpos_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, flag_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g1_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g2_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, k_col='invalid')
# Test using a limited set of rows
config['first_row'] = 101
config['last_row'] = 50000
cat3 = treecorr.Catalog(file_name, config)
np.testing.assert_equal(len(cat3.x), 49900)
np.testing.assert_equal(cat3.ntot, 49900)
np.testing.assert_equal(cat3.nobj, sum(cat3.w != 0))
np.testing.assert_equal(cat3.sumw, sum(cat3.w))
np.testing.assert_equal(cat3.sumw, sum(cat2.w[100:50000]))
np.testing.assert_almost_equal(cat3.g1[46292], 0.0005066675)
np.testing.assert_almost_equal(cat3.g2[46292], -0.0001006742)
np.testing.assert_almost_equal(cat3.k[46292], -0.0008628797)
cat4 = treecorr.read_catalogs(config, key='file_name', is_rand=True)[0]
np.testing.assert_equal(len(cat4.x), 49900)
np.testing.assert_equal(cat4.ntot, 49900)
np.testing.assert_equal(cat4.nobj, sum(cat4.w != 0))
np.testing.assert_equal(cat4.sumw, sum(cat4.w))
np.testing.assert_equal(cat4.sumw, sum(cat2.w[100:50000]))
assert cat4.g1 is None
assert cat4.g2 is None
assert cat4.k is None
do_pickle(cat1)
do_pickle(cat2)
do_pickle(cat3)
do_pickle(cat4)
assert_raises(ValueError, treecorr.Catalog, file_name, config, first_row=-10)
assert_raises(ValueError, treecorr.Catalog, file_name, config, first_row=0)
assert_raises(ValueError, treecorr.Catalog, file_name, config, first_row=60000)
assert_raises(ValueError, treecorr.Catalog, file_name, config, first_row=50001)
assert_raises(TypeError, treecorr.read_catalogs, config)
assert_raises(TypeError, treecorr.read_catalogs, config, key='file_name', list_key='file_name')
# If gg output not given, it is still invalid to only have one or the other of g1,g2.
del config['gg_file_name']
assert_raises(ValueError, treecorr.Catalog, file_name, config, g1_col='0')
assert_raises(ValueError, treecorr.Catalog, file_name, config, g2_col='0')
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=True.
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(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,) )
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 = 1.
max_sep = 50.
nbins = 50
bin_size = np.log(max_sep/min_sep) / nbins
kg = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, brute=True)
kg.process(cat1, cat2)
true_npairs = np.zeros(nbins, dtype=int)
true_weight = np.zeros(nbins, dtype=float)
true_xi = np.zeros(nbins, dtype=complex)
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] * (g12 + 1j*g22) * expmialpha**2
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 = ',kg.npairs - true_npairs)
np.testing.assert_array_equal(kg.npairs, true_npairs)
print('true_weight = ',true_weight)
print('diff = ',kg.weight - true_weight)
np.testing.assert_allclose(kg.weight, true_weight, rtol=1.e-5, atol=1.e-8)
print('true_xi = ',true_xi)
print('kg.xi = ',kg.xi)
print('kg.xi_im = ',kg.xi_im)
np.testing.assert_allclose(kg.xi, true_xi.real, rtol=1.e-4, atol=1.e-8)
np.testing.assert_allclose(kg.xi_im, true_xi.imag, 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/kg_direct.yaml')
cat1.write(config['file_name'])
cat2.write(config['file_name2'])
treecorr.corr2(config)
data = fitsio.read(config['kg_file_name'])
np.testing.assert_allclose(data['r_nom'], kg.rnom)
np.testing.assert_allclose(data['npairs'], kg.npairs)
np.testing.assert_allclose(data['weight'], kg.weight)
np.testing.assert_allclose(data['kgamT'], kg.xi, rtol=1.e-3)
np.testing.assert_allclose(data['kgamX'], kg.xi_im, rtol=1.e-3)
# Invalid with only one file_name
del config['file_name2']
with assert_raises(TypeError):
treecorr.corr2(config)
# Repeat with binslop = 0, since code is different for bin_slop=0 and brute=True.
# And don't do any top-level recursion so we actually test not going to the leaves.
kg = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, bin_slop=0,
max_top=0)
kg.process(cat1, cat2)
np.testing.assert_array_equal(kg.npairs, true_npairs)
np.testing.assert_allclose(kg.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kg.xi, true_xi.real, rtol=1.e-3, atol=1.e-3)
np.testing.assert_allclose(kg.xi_im, true_xi.imag, rtol=1.e-3, atol=1.e-3)
# Check a few basic operations with a KGCorrelation object.
do_pickle(kg)
kg2 = kg.copy()
kg2 += kg
np.testing.assert_allclose(kg2.npairs, 2*kg.npairs)
np.testing.assert_allclose(kg2.weight, 2*kg.weight)
np.testing.assert_allclose(kg2.meanr, 2*kg.meanr)
np.testing.assert_allclose(kg2.meanlogr, 2*kg.meanlogr)
np.testing.assert_allclose(kg2.xi, 2*kg.xi)
np.testing.assert_allclose(kg2.xi_im, 2*kg.xi_im)
kg2.clear()
kg2 += kg
np.testing.assert_allclose(kg2.npairs, kg.npairs)
np.testing.assert_allclose(kg2.weight, kg.weight)
np.testing.assert_allclose(kg2.meanr, kg.meanr)
np.testing.assert_allclose(kg2.meanlogr, kg.meanlogr)
np.testing.assert_allclose(kg2.xi, kg.xi)
np.testing.assert_allclose(kg2.xi_im, kg.xi_im)
ascii_name = 'output/kg_ascii.txt'
kg.write(ascii_name, precision=16)
kg3 = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
kg3.read(ascii_name)
np.testing.assert_allclose(kg3.npairs, kg.npairs)
np.testing.assert_allclose(kg3.weight, kg.weight)
np.testing.assert_allclose(kg3.meanr, kg.meanr)
np.testing.assert_allclose(kg3.meanlogr, kg.meanlogr)
np.testing.assert_allclose(kg3.xi, kg.xi)
np.testing.assert_allclose(kg3.xi_im, kg.xi_im)
fits_name = 'output/kg_fits.fits'
kg.write(fits_name)
kg4 = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
kg4.read(fits_name)
np.testing.assert_allclose(kg4.npairs, kg.npairs)
np.testing.assert_allclose(kg4.weight, kg.weight)
np.testing.assert_allclose(kg4.meanr, kg.meanr)
np.testing.assert_allclose(kg4.meanlogr, kg.meanlogr)
np.testing.assert_allclose(kg4.xi, kg.xi)
np.testing.assert_allclose(kg4.xi_im, kg.xi_im)
with assert_raises(TypeError):
kg2 += config
kg4 = treecorr.KGCorrelation(min_sep=min_sep/2, max_sep=max_sep, nbins=nbins)
with assert_raises(ValueError):
kg2 += kg4
kg5 = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep*2, nbins=nbins)
with assert_raises(ValueError):
kg2 += kg5
kg6 = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins*2)
with assert_raises(ValueError):
kg2 += kg6
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=True.
ngal = 100
s = 10.
rng = np.random.RandomState(8675309)
x = rng.normal(0,s, (ngal,) )
y = rng.normal(0,s, (ngal,) )
w = rng.random_sample(ngal)
g1 = rng.normal(0,0.2, (ngal,) )
g2 = rng.normal(0,0.2, (ngal,) )
cat = treecorr.Catalog(x=x, y=y, w=w, g1=g1, g2=g2)
min_sep = 1.
bin_size = 0.2
nrbins = 10
nubins = 5
nvbins = 5
max_sep = min_sep * np.exp(nrbins * bin_size)
ggg = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins, brute=True)
ggg.process(cat, num_threads=2)
true_ntri = np.zeros((nrbins, nubins, 2*nvbins), dtype=int)
true_weight = np.zeros((nrbins, nubins, 2*nvbins), dtype=float)
true_gam0 = np.zeros((nrbins, nubins, 2*nvbins), dtype=complex)
true_gam1 = np.zeros((nrbins, nubins, 2*nvbins), dtype=complex)
true_gam2 = np.zeros((nrbins, nubins, 2*nvbins), dtype=complex)
true_gam3 = np.zeros((nrbins, nubins, 2*nvbins), dtype=complex)
for i in range(ngal):
for j in range(i+1,ngal):
for k in range(j+1,ngal):
d12 = np.sqrt((x[i]-x[j])**2 + (y[i]-y[j])**2)
d23 = np.sqrt((x[j]-x[k])**2 + (y[j]-y[k])**2)
d31 = np.sqrt((x[k]-x[i])**2 + (y[k]-y[i])**2)
d3, d2, d1 = sorted([d12, d23, d31])
rindex = np.floor(np.log(d2/min_sep) / bin_size).astype(int)
if rindex < 0 or rindex >= nrbins: continue
if [d1, d2, d3] == [d23, d31, d12]: ii,jj,kk = i,j,k
elif [d1, d2, d3] == [d23, d12, d31]: ii,jj,kk = i,k,j
elif [d1, d2, d3] == [d31, d12, d23]: ii,jj,kk = j,k,i
elif [d1, d2, d3] == [d31, d23, d12]: ii,jj,kk = j,i,k
elif [d1, d2, d3] == [d12, d23, d31]: ii,jj,kk = k,i,j
elif [d1, d2, d3] == [d12, d31, d23]: ii,jj,kk = k,j,i
else: assert False
# Now use ii, jj, kk rather than i,j,k, to get the indices
# that correspond to the points in the right order.
u = d3/d2
v = (d1-d2)/d3
if (x[jj]-x[ii])*(y[kk]-y[ii]) < (x[kk]-x[ii])*(y[jj]-y[ii]):
v = -v
uindex = np.floor(u / bin_size).astype(int)
assert 0 <= uindex < nubins
vindex = np.floor((v+1) / bin_size).astype(int)
assert 0 <= vindex < 2*nvbins
# Rotate shears to coordinates where line connecting to center is horizontal.
cenx = (x[i] + x[j] + x[k])/3.
ceny = (y[i] + y[j] + y[k])/3.
expmialpha1 = (x[ii]-cenx) - 1j*(y[ii]-ceny)
expmialpha1 /= abs(expmialpha1)
expmialpha2 = (x[jj]-cenx) - 1j*(y[jj]-ceny)
expmialpha2 /= abs(expmialpha2)
expmialpha3 = (x[kk]-cenx) - 1j*(y[kk]-ceny)
expmialpha3 /= abs(expmialpha3)
www = w[i] * w[j] * w[k]
g1p = (g1[ii] + 1j*g2[ii]) * expmialpha1**2
g2p = (g1[jj] + 1j*g2[jj]) * expmialpha2**2
g3p = (g1[kk] + 1j*g2[kk]) * expmialpha3**2
gam0 = www * g1p * g2p * g3p
gam1 = www * np.conjugate(g1p) * g2p * g3p
gam2 = www * g1p * np.conjugate(g2p) * g3p
gam3 = www * g1p * g2p * np.conjugate(g3p)
true_ntri[rindex,uindex,vindex] += 1
true_weight[rindex,uindex,vindex] += www
true_gam0[rindex,uindex,vindex] += gam0
true_gam1[rindex,uindex,vindex] += gam1
true_gam2[rindex,uindex,vindex] += gam2
true_gam3[rindex,uindex,vindex] += gam3
pos = true_weight > 0
true_gam0[pos] /= true_weight[pos]
true_gam1[pos] /= true_weight[pos]
true_gam2[pos] /= true_weight[pos]
true_gam3[pos] /= true_weight[pos]
np.testing.assert_array_equal(ggg.ntri, true_ntri)
np.testing.assert_allclose(ggg.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0r, true_gam0.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0i, true_gam0.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam1r, true_gam1.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam1i, true_gam1.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam2r, true_gam2.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam2i, true_gam2.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam3r, true_gam3.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam3i, true_gam3.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0, true_gam0, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam1, true_gam1, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam2, true_gam2, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam3, true_gam3, rtol=1.e-5, atol=1.e-8)
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
# Check that running via the corr3 script works correctly.
config = treecorr.config.read_config('configs/ggg_direct.yaml')
cat.write(config['file_name'])
treecorr.corr3(config)
data = fitsio.read(config['ggg_file_name'])
np.testing.assert_allclose(data['r_nom'], ggg.rnom.flatten())
np.testing.assert_allclose(data['u_nom'], ggg.u.flatten())
np.testing.assert_allclose(data['v_nom'], ggg.v.flatten())
np.testing.assert_allclose(data['ntri'], ggg.ntri.flatten())
np.testing.assert_allclose(data['weight'], ggg.weight.flatten())
np.testing.assert_allclose(data['gam0r'], ggg.gam0r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam0i'], ggg.gam0i.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam1r'], ggg.gam1r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam1i'], ggg.gam1i.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam2r'], ggg.gam2r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam2i'], ggg.gam2i.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam3r'], ggg.gam3r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam3i'], ggg.gam3i.flatten(), rtol=1.e-3)
# Also check the "cross" calculation. (Real cross doesn't work, but this should.)
ggg = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins, brute=True)
ggg.process(cat, cat, cat, num_threads=2)
np.testing.assert_array_equal(ggg.ntri, true_ntri)
np.testing.assert_allclose(ggg.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0r, true_gam0.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0i, true_gam0.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam1r, true_gam1.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam1i, true_gam1.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam2r, true_gam2.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam2i, true_gam2.imag, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam3r, true_gam3.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam3i, true_gam3.imag, rtol=1.e-5, atol=1.e-8)
config['file_name2'] = config['file_name']
config['file_name3'] = config['file_name']
treecorr.corr3(config)
data = fitsio.read(config['ggg_file_name'])
np.testing.assert_allclose(data['r_nom'], ggg.rnom.flatten())
np.testing.assert_allclose(data['u_nom'], ggg.u.flatten())
np.testing.assert_allclose(data['v_nom'], ggg.v.flatten())
np.testing.assert_allclose(data['ntri'], ggg.ntri.flatten())
np.testing.assert_allclose(data['weight'], ggg.weight.flatten())
np.testing.assert_allclose(data['gam0r'], ggg.gam0r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam0i'], ggg.gam0i.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam1r'], ggg.gam1r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam1i'], ggg.gam1i.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam2r'], ggg.gam2r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam2i'], ggg.gam2i.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam3r'], ggg.gam3r.flatten(), rtol=1.e-3)
np.testing.assert_allclose(data['gam3i'], ggg.gam3i.flatten(), rtol=1.e-3)
# Repeat with binslop = 0, since the code flow is different from brute=True.
# And don't do any top-level recursion so we actually test not going to the leaves.
ggg = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
bin_slop=0, max_top=0)
ggg.process(cat)
np.testing.assert_array_equal(ggg.ntri, true_ntri)
np.testing.assert_allclose(ggg.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0r, true_gam0.real, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(ggg.gam0i, true_gam0.imag, rtol=1.e-5, atol=1.e-4)
np.testing.assert_allclose(ggg.gam1r, true_gam1.real, rtol=1.e-3, atol=1.e-4)
np.testing.assert_allclose(ggg.gam1i, true_gam1.imag, rtol=1.e-3, atol=1.e-4)
np.testing.assert_allclose(ggg.gam2r, true_gam2.real, rtol=1.e-3, atol=1.e-4)
np.testing.assert_allclose(ggg.gam2i, true_gam2.imag, rtol=1.e-3, atol=1.e-4)
np.testing.assert_allclose(ggg.gam3r, true_gam3.real, rtol=1.e-3, atol=1.e-4)
np.testing.assert_allclose(ggg.gam3i, true_gam3.imag, rtol=1.e-3, atol=1.e-4)
# Check a few basic operations with a GGCorrelation object.
do_pickle(ggg)
ggg2 = ggg.copy()
ggg2 += ggg
np.testing.assert_allclose(ggg2.ntri, 2*ggg.ntri)
np.testing.assert_allclose(ggg2.weight, 2*ggg.weight)
np.testing.assert_allclose(ggg2.meand1, 2*ggg.meand1)
np.testing.assert_allclose(ggg2.meand2, 2*ggg.meand2)
np.testing.assert_allclose(ggg2.meand3, 2*ggg.meand3)
np.testing.assert_allclose(ggg2.meanlogd1, 2*ggg.meanlogd1)
np.testing.assert_allclose(ggg2.meanlogd2, 2*ggg.meanlogd2)
np.testing.assert_allclose(ggg2.meanlogd3, 2*ggg.meanlogd3)
np.testing.assert_allclose(ggg2.meanu, 2*ggg.meanu)
np.testing.assert_allclose(ggg2.meanv, 2*ggg.meanv)
np.testing.assert_allclose(ggg2.gam0r, 2*ggg.gam0r)
np.testing.assert_allclose(ggg2.gam0i, 2*ggg.gam0i)
np.testing.assert_allclose(ggg2.gam1r, 2*ggg.gam1r)
np.testing.assert_allclose(ggg2.gam1i, 2*ggg.gam1i)
np.testing.assert_allclose(ggg2.gam2r, 2*ggg.gam2r)
np.testing.assert_allclose(ggg2.gam2i, 2*ggg.gam2i)
np.testing.assert_allclose(ggg2.gam3r, 2*ggg.gam3r)
np.testing.assert_allclose(ggg2.gam3i, 2*ggg.gam3i)
ggg2.clear()
ggg2 += ggg
np.testing.assert_allclose(ggg2.ntri, ggg.ntri)
np.testing.assert_allclose(ggg2.weight, ggg.weight)
np.testing.assert_allclose(ggg2.meand1, ggg.meand1)
np.testing.assert_allclose(ggg2.meand2, ggg.meand2)
np.testing.assert_allclose(ggg2.meand3, ggg.meand3)
np.testing.assert_allclose(ggg2.meanlogd1, ggg.meanlogd1)
np.testing.assert_allclose(ggg2.meanlogd2, ggg.meanlogd2)
np.testing.assert_allclose(ggg2.meanlogd3, ggg.meanlogd3)
np.testing.assert_allclose(ggg2.meanu, ggg.meanu)
np.testing.assert_allclose(ggg2.meanv, ggg.meanv)
np.testing.assert_allclose(ggg2.gam0r, ggg.gam0r)
np.testing.assert_allclose(ggg2.gam0i, ggg.gam0i)
np.testing.assert_allclose(ggg2.gam1r, ggg.gam1r)
np.testing.assert_allclose(ggg2.gam1i, ggg.gam1i)
np.testing.assert_allclose(ggg2.gam2r, ggg.gam2r)
np.testing.assert_allclose(ggg2.gam2i, ggg.gam2i)
np.testing.assert_allclose(ggg2.gam3r, ggg.gam3r)
np.testing.assert_allclose(ggg2.gam3i, ggg.gam3i)
ascii_name = 'output/ggg_ascii.txt'
ggg.write(ascii_name, precision=16)
ggg3 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins)
ggg3.read(ascii_name)
np.testing.assert_allclose(ggg3.ntri, ggg.ntri)
np.testing.assert_allclose(ggg3.weight, ggg.weight)
np.testing.assert_allclose(ggg3.meand1, ggg.meand1)
np.testing.assert_allclose(ggg3.meand2, ggg.meand2)
np.testing.assert_allclose(ggg3.meand3, ggg.meand3)
np.testing.assert_allclose(ggg3.meanlogd1, ggg.meanlogd1)
np.testing.assert_allclose(ggg3.meanlogd2, ggg.meanlogd2)
np.testing.assert_allclose(ggg3.meanlogd3, ggg.meanlogd3)
np.testing.assert_allclose(ggg3.meanu, ggg.meanu)
np.testing.assert_allclose(ggg3.meanv, ggg.meanv)
np.testing.assert_allclose(ggg3.gam0r, ggg.gam0r)
np.testing.assert_allclose(ggg3.gam0i, ggg.gam0i)
np.testing.assert_allclose(ggg3.gam1r, ggg.gam1r)
np.testing.assert_allclose(ggg3.gam1i, ggg.gam1i)
np.testing.assert_allclose(ggg3.gam2r, ggg.gam2r)
np.testing.assert_allclose(ggg3.gam2i, ggg.gam2i)
np.testing.assert_allclose(ggg3.gam3r, ggg.gam3r)
np.testing.assert_allclose(ggg3.gam3i, ggg.gam3i)
fits_name = 'output/ggg_fits.fits'
ggg.write(fits_name)
ggg4 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins)
ggg4.read(fits_name)
np.testing.assert_allclose(ggg4.ntri, ggg.ntri)
np.testing.assert_allclose(ggg4.weight, ggg.weight)
np.testing.assert_allclose(ggg4.meand1, ggg.meand1)
np.testing.assert_allclose(ggg4.meand2, ggg.meand2)
np.testing.assert_allclose(ggg4.meand3, ggg.meand3)
np.testing.assert_allclose(ggg4.meanlogd1, ggg.meanlogd1)
np.testing.assert_allclose(ggg4.meanlogd2, ggg.meanlogd2)
np.testing.assert_allclose(ggg4.meanlogd3, ggg.meanlogd3)
np.testing.assert_allclose(ggg4.meanu, ggg.meanu)
np.testing.assert_allclose(ggg4.meanv, ggg.meanv)
np.testing.assert_allclose(ggg4.gam0r, ggg.gam0r)
np.testing.assert_allclose(ggg4.gam0i, ggg.gam0i)
np.testing.assert_allclose(ggg4.gam1r, ggg.gam1r)
np.testing.assert_allclose(ggg4.gam1i, ggg.gam1i)
np.testing.assert_allclose(ggg4.gam2r, ggg.gam2r)
np.testing.assert_allclose(ggg4.gam2i, ggg.gam2i)
np.testing.assert_allclose(ggg4.gam3r, ggg.gam3r)
np.testing.assert_allclose(ggg4.gam3i, ggg.gam3i)
with assert_raises(TypeError):
ggg2 += config
ggg5 = treecorr.GGGCorrelation(min_sep=min_sep/2, bin_size=bin_size, nbins=nrbins)
with assert_raises(ValueError):
ggg2 += ggg5
ggg6 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size/2, nbins=nrbins)
with assert_raises(ValueError):
ggg2 += ggg6
ggg7 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins*2)
with assert_raises(ValueError):
ggg2 += ggg7
ggg8 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
min_u=0.1)
with assert_raises(ValueError):
ggg2 += ggg8
ggg0 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
max_u=0.1)
with assert_raises(ValueError):
ggg2 += ggg0
ggg10 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
nubins=nrbins*2)
with assert_raises(ValueError):
ggg2 += ggg10
ggg11 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
min_v=0.1)
with assert_raises(ValueError):
ggg2 += ggg11
ggg12 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
max_v=0.1)
with assert_raises(ValueError):
ggg2 += ggg12
ggg13 = treecorr.GGGCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
nvbins=nrbins*2)
with assert_raises(ValueError):
ggg2 += ggg13
# Currently not implemented to only have cat2 or cat3
with assert_raises(NotImplementedError):
ggg.process(cat, cat2=cat)
with assert_raises(NotImplementedError):
ggg.process(cat, cat3=cat)
with assert_raises(NotImplementedError):
ggg.process_cross21(cat, cat)
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)
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)
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
nk = treecorr.NKCorrelation(min_sep=min_sep,
max_sep=max_sep,
nbins=nbins,
brute=True)
nk.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)
ww = w1[i] * w2
xi = ww * 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 = ', nk.npairs - true_npairs)
np.testing.assert_array_equal(nk.npairs, true_npairs)
print('true_weight = ', true_weight)
print('diff = ', nk.weight - true_weight)
np.testing.assert_allclose(nk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
print('true_xi = ', true_xi)
print('nk.xi = ', nk.xi)
np.testing.assert_allclose(nk.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/nk_direct.yaml')
cat1.write(config['file_name'])
cat2.write(config['file_name2'])
treecorr.corr2(config)
data = fitsio.read(config['nk_file_name'])
np.testing.assert_allclose(data['r_nom'], nk.rnom)
np.testing.assert_allclose(data['npairs'], nk.npairs)
np.testing.assert_allclose(data['weight'], nk.weight)
np.testing.assert_allclose(data['kappa'], nk.xi, rtol=1.e-3)
# Invalid with only one file_name
del config['file_name2']
with assert_raises(TypeError):
treecorr.corr2(config)
config['file_name2'] = 'data/nk_direct_cat2.fits'
# Invalid to request compoensated if no rand_file
config['nk_statistic'] = 'compensated'
with assert_raises(TypeError):
treecorr.corr2(config)
# Repeat with binslop = 0, since the code flow is different from brute=True
# And don't do any top-level recursion so we actually test not going to the leaves.
nk = treecorr.NKCorrelation(min_sep=min_sep,
max_sep=max_sep,
nbins=nbins,
bin_slop=0,
max_top=0)
nk.process(cat1, cat2)
np.testing.assert_array_equal(nk.npairs, true_npairs)
np.testing.assert_allclose(nk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(nk.xi, true_xi, rtol=1.e-4, atol=1.e-8)
# Check a few basic operations with a NKCorrelation object.
do_pickle(nk)
nk2 = nk.copy()
nk2 += nk
np.testing.assert_allclose(nk2.npairs, 2 * nk.npairs)
np.testing.assert_allclose(nk2.weight, 2 * nk.weight)
np.testing.assert_allclose(nk2.meanr, 2 * nk.meanr)
np.testing.assert_allclose(nk2.meanlogr, 2 * nk.meanlogr)
np.testing.assert_allclose(nk2.xi, 2 * nk.xi)
nk2.clear()
nk2 += nk
np.testing.assert_allclose(nk2.npairs, nk.npairs)
np.testing.assert_allclose(nk2.weight, nk.weight)
np.testing.assert_allclose(nk2.meanr, nk.meanr)
np.testing.assert_allclose(nk2.meanlogr, nk.meanlogr)
np.testing.assert_allclose(nk2.xi, nk.xi)
ascii_name = 'output/nk_ascii.txt'
nk.write(ascii_name, precision=16)
nk3 = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
nk3.read(ascii_name)
np.testing.assert_allclose(nk3.npairs, nk.npairs)
np.testing.assert_allclose(nk3.weight, nk.weight)
np.testing.assert_allclose(nk3.meanr, nk.meanr)
np.testing.assert_allclose(nk3.meanlogr, nk.meanlogr)
np.testing.assert_allclose(nk3.xi, nk.xi)
with assert_raises(TypeError):
nk2 += config
nk4 = treecorr.NKCorrelation(min_sep=min_sep / 2,
max_sep=max_sep,
nbins=nbins)
with assert_raises(ValueError):
nk2 += nk4
nk5 = treecorr.NKCorrelation(min_sep=min_sep,
max_sep=max_sep * 2,
nbins=nbins)
with assert_raises(ValueError):
nk2 += nk5
nk6 = treecorr.NKCorrelation(min_sep=min_sep,
max_sep=max_sep,
nbins=nbins * 2)
with assert_raises(ValueError):
nk2 += nk6
fits_name = 'output/nk_fits.fits'
nk.write(fits_name)
nk4 = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
nk4.read(fits_name)
np.testing.assert_allclose(nk4.npairs, nk.npairs)
np.testing.assert_allclose(nk4.weight, nk.weight)
np.testing.assert_allclose(nk4.meanr, nk.meanr)
np.testing.assert_allclose(nk4.meanlogr, nk.meanlogr)
np.testing.assert_allclose(nk4.xi, nk.xi)
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=True.
ngal = 100
s = 10.
rng = np.random.RandomState(8675309)
x = rng.normal(0, s, (ngal, ))
y = rng.normal(0, s, (ngal, ))
w = rng.random_sample(ngal)
kap = rng.normal(0, 3, (ngal, ))
cat = treecorr.Catalog(x=x, y=y, w=w, k=kap)
min_sep = 1.
bin_size = 0.2
nrbins = 10
nubins = 5
nvbins = 5
max_sep = min_sep * np.exp(nrbins * bin_size)
kkk = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
brute=True)
kkk.process(cat, num_threads=2)
true_ntri = np.zeros((nrbins, nubins, 2 * nvbins), dtype=int)
true_weight = np.zeros((nrbins, nubins, 2 * nvbins), dtype=float)
true_zeta = np.zeros((nrbins, nubins, 2 * nvbins), dtype=float)
for i in range(ngal):
for j in range(i + 1, ngal):
for k in range(j + 1, ngal):
d12 = np.sqrt((x[i] - x[j])**2 + (y[i] - y[j])**2)
d23 = np.sqrt((x[j] - x[k])**2 + (y[j] - y[k])**2)
d31 = np.sqrt((x[k] - x[i])**2 + (y[k] - y[i])**2)
d3, d2, d1 = sorted([d12, d23, d31])
rindex = np.floor(np.log(d2 / min_sep) / bin_size).astype(int)
if rindex < 0 or rindex >= nrbins: continue
if [d1, d2, d3] == [d23, d31, d12]: ii, jj, kk = i, j, k
elif [d1, d2, d3] == [d23, d12, d31]: ii, jj, kk = i, k, j
elif [d1, d2, d3] == [d31, d12, d23]: ii, jj, kk = j, k, i
elif [d1, d2, d3] == [d31, d23, d12]: ii, jj, kk = j, i, k
elif [d1, d2, d3] == [d12, d23, d31]: ii, jj, kk = k, i, j
elif [d1, d2, d3] == [d12, d31, d23]: ii, jj, kk = k, j, i
else: assert False
# Now use ii, jj, kk rather than i,j,k, to get the indices
# that correspond to the points in the right order.
u = d3 / d2
v = (d1 - d2) / d3
if (x[jj] - x[ii]) * (y[kk] - y[ii]) < (x[kk] - x[ii]) * (
y[jj] - y[ii]):
v = -v
uindex = np.floor(u / bin_size).astype(int)
assert 0 <= uindex < nubins
vindex = np.floor((v + 1) / bin_size).astype(int)
assert 0 <= vindex < 2 * nvbins
www = w[i] * w[j] * w[k]
zeta = www * kap[i] * kap[j] * kap[k]
true_ntri[rindex, uindex, vindex] += 1
true_weight[rindex, uindex, vindex] += www
true_zeta[rindex, uindex, vindex] += zeta
pos = true_weight > 0
true_zeta[pos] /= true_weight[pos]
np.testing.assert_array_equal(kkk.ntri, true_ntri)
np.testing.assert_allclose(kkk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kkk.zeta, true_zeta, rtol=1.e-5, atol=1.e-8)
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
# Check that running via the corr3 script works correctly.
config = treecorr.config.read_config('configs/kkk_direct.yaml')
cat.write(config['file_name'])
treecorr.corr3(config)
data = fitsio.read(config['kkk_file_name'])
np.testing.assert_allclose(data['r_nom'], kkk.rnom.flatten())
np.testing.assert_allclose(data['u_nom'], kkk.u.flatten())
np.testing.assert_allclose(data['v_nom'], kkk.v.flatten())
np.testing.assert_allclose(data['ntri'], kkk.ntri.flatten())
np.testing.assert_allclose(data['weight'], kkk.weight.flatten())
np.testing.assert_allclose(data['zeta'], kkk.zeta.flatten(), rtol=1.e-3)
# Also check the "cross" calculation. (Real cross doesn't work, but this should.)
kkk = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
brute=True)
kkk.process(cat, cat, cat, num_threads=2)
np.testing.assert_array_equal(kkk.ntri, true_ntri)
np.testing.assert_allclose(kkk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kkk.zeta, true_zeta, rtol=1.e-5, atol=1.e-8)
config['file_name2'] = config['file_name']
config['file_name3'] = config['file_name']
treecorr.corr3(config)
data = fitsio.read(config['kkk_file_name'])
np.testing.assert_allclose(data['r_nom'], kkk.rnom.flatten())
np.testing.assert_allclose(data['u_nom'], kkk.u.flatten())
np.testing.assert_allclose(data['v_nom'], kkk.v.flatten())
np.testing.assert_allclose(data['ntri'], kkk.ntri.flatten())
np.testing.assert_allclose(data['weight'], kkk.weight.flatten())
np.testing.assert_allclose(data['zeta'], kkk.zeta.flatten(), 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.
kkk = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
bin_slop=0,
max_top=0)
kkk.process(cat)
np.testing.assert_array_equal(kkk.ntri, true_ntri)
np.testing.assert_allclose(kkk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kkk.zeta, true_zeta, rtol=1.e-5, atol=1.e-8)
# Check a few basic operations with a GGCorrelation object.
do_pickle(kkk)
kkk2 = kkk.copy()
kkk2 += kkk
np.testing.assert_allclose(kkk2.ntri, 2 * kkk.ntri)
np.testing.assert_allclose(kkk2.weight, 2 * kkk.weight)
np.testing.assert_allclose(kkk2.meand1, 2 * kkk.meand1)
np.testing.assert_allclose(kkk2.meand2, 2 * kkk.meand2)
np.testing.assert_allclose(kkk2.meand3, 2 * kkk.meand3)
np.testing.assert_allclose(kkk2.meanlogd1, 2 * kkk.meanlogd1)
np.testing.assert_allclose(kkk2.meanlogd2, 2 * kkk.meanlogd2)
np.testing.assert_allclose(kkk2.meanlogd3, 2 * kkk.meanlogd3)
np.testing.assert_allclose(kkk2.meanu, 2 * kkk.meanu)
np.testing.assert_allclose(kkk2.meanv, 2 * kkk.meanv)
np.testing.assert_allclose(kkk2.zeta, 2 * kkk.zeta)
kkk2.clear()
kkk2 += kkk
np.testing.assert_allclose(kkk2.ntri, kkk.ntri)
np.testing.assert_allclose(kkk2.weight, kkk.weight)
np.testing.assert_allclose(kkk2.meand1, kkk.meand1)
np.testing.assert_allclose(kkk2.meand2, kkk.meand2)
np.testing.assert_allclose(kkk2.meand3, kkk.meand3)
np.testing.assert_allclose(kkk2.meanlogd1, kkk.meanlogd1)
np.testing.assert_allclose(kkk2.meanlogd2, kkk.meanlogd2)
np.testing.assert_allclose(kkk2.meanlogd3, kkk.meanlogd3)
np.testing.assert_allclose(kkk2.meanu, kkk.meanu)
np.testing.assert_allclose(kkk2.meanv, kkk.meanv)
np.testing.assert_allclose(kkk2.zeta, kkk.zeta)
ascii_name = 'output/kkk_ascii.txt'
kkk.write(ascii_name, precision=16)
kkk3 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins)
kkk3.read(ascii_name)
np.testing.assert_allclose(kkk3.ntri, kkk.ntri)
np.testing.assert_allclose(kkk3.weight, kkk.weight)
np.testing.assert_allclose(kkk3.meand1, kkk.meand1)
np.testing.assert_allclose(kkk3.meand2, kkk.meand2)
np.testing.assert_allclose(kkk3.meand3, kkk.meand3)
np.testing.assert_allclose(kkk3.meanlogd1, kkk.meanlogd1)
np.testing.assert_allclose(kkk3.meanlogd2, kkk.meanlogd2)
np.testing.assert_allclose(kkk3.meanlogd3, kkk.meanlogd3)
np.testing.assert_allclose(kkk3.meanu, kkk.meanu)
np.testing.assert_allclose(kkk3.meanv, kkk.meanv)
np.testing.assert_allclose(kkk3.zeta, kkk.zeta)
fits_name = 'output/kkk_fits.fits'
kkk.write(fits_name)
kkk4 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins)
kkk4.read(fits_name)
np.testing.assert_allclose(kkk4.ntri, kkk.ntri)
np.testing.assert_allclose(kkk4.weight, kkk.weight)
np.testing.assert_allclose(kkk4.meand1, kkk.meand1)
np.testing.assert_allclose(kkk4.meand2, kkk.meand2)
np.testing.assert_allclose(kkk4.meand3, kkk.meand3)
np.testing.assert_allclose(kkk4.meanlogd1, kkk.meanlogd1)
np.testing.assert_allclose(kkk4.meanlogd2, kkk.meanlogd2)
np.testing.assert_allclose(kkk4.meanlogd3, kkk.meanlogd3)
np.testing.assert_allclose(kkk4.meanu, kkk.meanu)
np.testing.assert_allclose(kkk4.meanv, kkk.meanv)
np.testing.assert_allclose(kkk4.zeta, kkk.zeta)
with assert_raises(TypeError):
kkk2 += config
kkk5 = treecorr.KKKCorrelation(min_sep=min_sep / 2,
bin_size=bin_size,
nbins=nrbins)
with assert_raises(ValueError):
kkk2 += kkk5
kkk6 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size / 2,
nbins=nrbins)
with assert_raises(ValueError):
kkk2 += kkk6
kkk7 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins * 2)
with assert_raises(ValueError):
kkk2 += kkk7
kkk8 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
min_u=0.1)
with assert_raises(ValueError):
kkk2 += kkk8
kkk0 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
max_u=0.1)
with assert_raises(ValueError):
kkk2 += kkk0
kkk10 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
nubins=nrbins * 2)
with assert_raises(ValueError):
kkk2 += kkk10
kkk11 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
min_v=0.1)
with assert_raises(ValueError):
kkk2 += kkk11
kkk12 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
max_v=0.1)
with assert_raises(ValueError):
kkk2 += kkk12
kkk13 = treecorr.KKKCorrelation(min_sep=min_sep,
bin_size=bin_size,
nbins=nrbins,
nvbins=nrbins * 2)
with assert_raises(ValueError):
kkk2 += kkk13
# Currently not implemented to only have cat2 or cat3
with assert_raises(NotImplementedError):
kkk.process(cat, cat2=cat)
with assert_raises(NotImplementedError):
kkk.process(cat, cat3=cat)
with assert_raises(NotImplementedError):
kkk.process_cross21(cat, cat)
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=True.
ngal = 100
s = 10.
rng = np.random.RandomState(8675309)
x = rng.normal(0,s, (ngal,) )
y = rng.normal(0,s, (ngal,) )
w = rng.random_sample(ngal)
kap = rng.normal(0,3, (ngal,) )
cat = treecorr.Catalog(x=x, y=y, w=w, k=kap)
min_sep = 1.
bin_size = 0.2
nrbins = 10
nubins = 5
nvbins = 5
max_sep = min_sep * np.exp(nrbins * bin_size)
kkk = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins, brute=True)
kkk.process(cat, num_threads=2)
true_ntri = np.zeros((nrbins, nubins, 2*nvbins), dtype=int)
true_weight = np.zeros((nrbins, nubins, 2*nvbins), dtype=float)
true_zeta = np.zeros((nrbins, nubins, 2*nvbins), dtype=float)
for i in range(ngal):
for j in range(i+1,ngal):
for k in range(j+1,ngal):
d12 = np.sqrt((x[i]-x[j])**2 + (y[i]-y[j])**2)
d23 = np.sqrt((x[j]-x[k])**2 + (y[j]-y[k])**2)
d31 = np.sqrt((x[k]-x[i])**2 + (y[k]-y[i])**2)
d3, d2, d1 = sorted([d12, d23, d31])
rindex = np.floor(np.log(d2/min_sep) / bin_size).astype(int)
if rindex < 0 or rindex >= nrbins: continue
if [d1, d2, d3] == [d23, d31, d12]: ii,jj,kk = i,j,k
elif [d1, d2, d3] == [d23, d12, d31]: ii,jj,kk = i,k,j
elif [d1, d2, d3] == [d31, d12, d23]: ii,jj,kk = j,k,i
elif [d1, d2, d3] == [d31, d23, d12]: ii,jj,kk = j,i,k
elif [d1, d2, d3] == [d12, d23, d31]: ii,jj,kk = k,i,j
elif [d1, d2, d3] == [d12, d31, d23]: ii,jj,kk = k,j,i
else: assert False
# Now use ii, jj, kk rather than i,j,k, to get the indices
# that correspond to the points in the right order.
u = d3/d2
v = (d1-d2)/d3
if (x[jj]-x[ii])*(y[kk]-y[ii]) < (x[kk]-x[ii])*(y[jj]-y[ii]):
v = -v
uindex = np.floor(u / bin_size).astype(int)
assert 0 <= uindex < nubins
vindex = np.floor((v+1) / bin_size).astype(int)
assert 0 <= vindex < 2*nvbins
www = w[i] * w[j] * w[k]
zeta = www * kap[i] * kap[j] * kap[k]
true_ntri[rindex,uindex,vindex] += 1
true_weight[rindex,uindex,vindex] += www
true_zeta[rindex,uindex,vindex] += zeta
pos = true_weight > 0
true_zeta[pos] /= true_weight[pos]
np.testing.assert_array_equal(kkk.ntri, true_ntri)
np.testing.assert_allclose(kkk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kkk.zeta, true_zeta, rtol=1.e-5, atol=1.e-8)
try:
import fitsio
except ImportError:
print('Skipping FITS tests, since fitsio is not installed')
return
# Check that running via the corr3 script works correctly.
config = treecorr.config.read_config('configs/kkk_direct.yaml')
cat.write(config['file_name'])
treecorr.corr3(config)
data = fitsio.read(config['kkk_file_name'])
np.testing.assert_allclose(data['r_nom'], kkk.rnom.flatten())
np.testing.assert_allclose(data['u_nom'], kkk.u.flatten())
np.testing.assert_allclose(data['v_nom'], kkk.v.flatten())
np.testing.assert_allclose(data['ntri'], kkk.ntri.flatten())
np.testing.assert_allclose(data['weight'], kkk.weight.flatten())
np.testing.assert_allclose(data['zeta'], kkk.zeta.flatten(), rtol=1.e-3)
# Also check the "cross" calculation. (Real cross doesn't work, but this should.)
kkk = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins, brute=True)
kkk.process(cat, cat, cat, num_threads=2)
np.testing.assert_array_equal(kkk.ntri, true_ntri)
np.testing.assert_allclose(kkk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kkk.zeta, true_zeta, rtol=1.e-5, atol=1.e-8)
config['file_name2'] = config['file_name']
config['file_name3'] = config['file_name']
treecorr.corr3(config)
data = fitsio.read(config['kkk_file_name'])
np.testing.assert_allclose(data['r_nom'], kkk.rnom.flatten())
np.testing.assert_allclose(data['u_nom'], kkk.u.flatten())
np.testing.assert_allclose(data['v_nom'], kkk.v.flatten())
np.testing.assert_allclose(data['ntri'], kkk.ntri.flatten())
np.testing.assert_allclose(data['weight'], kkk.weight.flatten())
np.testing.assert_allclose(data['zeta'], kkk.zeta.flatten(), 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.
kkk = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
bin_slop=0, max_top=0)
kkk.process(cat)
np.testing.assert_array_equal(kkk.ntri, true_ntri)
np.testing.assert_allclose(kkk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(kkk.zeta, true_zeta, rtol=1.e-5, atol=1.e-8)
# Check a few basic operations with a GGCorrelation object.
do_pickle(kkk)
kkk2 = kkk.copy()
kkk2 += kkk
np.testing.assert_allclose(kkk2.ntri, 2*kkk.ntri)
np.testing.assert_allclose(kkk2.weight, 2*kkk.weight)
np.testing.assert_allclose(kkk2.meand1, 2*kkk.meand1)
np.testing.assert_allclose(kkk2.meand2, 2*kkk.meand2)
np.testing.assert_allclose(kkk2.meand3, 2*kkk.meand3)
np.testing.assert_allclose(kkk2.meanlogd1, 2*kkk.meanlogd1)
np.testing.assert_allclose(kkk2.meanlogd2, 2*kkk.meanlogd2)
np.testing.assert_allclose(kkk2.meanlogd3, 2*kkk.meanlogd3)
np.testing.assert_allclose(kkk2.meanu, 2*kkk.meanu)
np.testing.assert_allclose(kkk2.meanv, 2*kkk.meanv)
np.testing.assert_allclose(kkk2.zeta, 2*kkk.zeta)
kkk2.clear()
kkk2 += kkk
np.testing.assert_allclose(kkk2.ntri, kkk.ntri)
np.testing.assert_allclose(kkk2.weight, kkk.weight)
np.testing.assert_allclose(kkk2.meand1, kkk.meand1)
np.testing.assert_allclose(kkk2.meand2, kkk.meand2)
np.testing.assert_allclose(kkk2.meand3, kkk.meand3)
np.testing.assert_allclose(kkk2.meanlogd1, kkk.meanlogd1)
np.testing.assert_allclose(kkk2.meanlogd2, kkk.meanlogd2)
np.testing.assert_allclose(kkk2.meanlogd3, kkk.meanlogd3)
np.testing.assert_allclose(kkk2.meanu, kkk.meanu)
np.testing.assert_allclose(kkk2.meanv, kkk.meanv)
np.testing.assert_allclose(kkk2.zeta, kkk.zeta)
ascii_name = 'output/kkk_ascii.txt'
kkk.write(ascii_name, precision=16)
kkk3 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins)
kkk3.read(ascii_name)
np.testing.assert_allclose(kkk3.ntri, kkk.ntri)
np.testing.assert_allclose(kkk3.weight, kkk.weight)
np.testing.assert_allclose(kkk3.meand1, kkk.meand1)
np.testing.assert_allclose(kkk3.meand2, kkk.meand2)
np.testing.assert_allclose(kkk3.meand3, kkk.meand3)
np.testing.assert_allclose(kkk3.meanlogd1, kkk.meanlogd1)
np.testing.assert_allclose(kkk3.meanlogd2, kkk.meanlogd2)
np.testing.assert_allclose(kkk3.meanlogd3, kkk.meanlogd3)
np.testing.assert_allclose(kkk3.meanu, kkk.meanu)
np.testing.assert_allclose(kkk3.meanv, kkk.meanv)
np.testing.assert_allclose(kkk3.zeta, kkk.zeta)
fits_name = 'output/kkk_fits.fits'
kkk.write(fits_name)
kkk4 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins)
kkk4.read(fits_name)
np.testing.assert_allclose(kkk4.ntri, kkk.ntri)
np.testing.assert_allclose(kkk4.weight, kkk.weight)
np.testing.assert_allclose(kkk4.meand1, kkk.meand1)
np.testing.assert_allclose(kkk4.meand2, kkk.meand2)
np.testing.assert_allclose(kkk4.meand3, kkk.meand3)
np.testing.assert_allclose(kkk4.meanlogd1, kkk.meanlogd1)
np.testing.assert_allclose(kkk4.meanlogd2, kkk.meanlogd2)
np.testing.assert_allclose(kkk4.meanlogd3, kkk.meanlogd3)
np.testing.assert_allclose(kkk4.meanu, kkk.meanu)
np.testing.assert_allclose(kkk4.meanv, kkk.meanv)
np.testing.assert_allclose(kkk4.zeta, kkk.zeta)
with assert_raises(TypeError):
kkk2 += config
kkk5 = treecorr.KKKCorrelation(min_sep=min_sep/2, bin_size=bin_size, nbins=nrbins)
with assert_raises(ValueError):
kkk2 += kkk5
kkk6 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size/2, nbins=nrbins)
with assert_raises(ValueError):
kkk2 += kkk6
kkk7 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins*2)
with assert_raises(ValueError):
kkk2 += kkk7
kkk8 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
min_u=0.1)
with assert_raises(ValueError):
kkk2 += kkk8
kkk0 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
max_u=0.1)
with assert_raises(ValueError):
kkk2 += kkk0
kkk10 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
nubins=nrbins*2)
with assert_raises(ValueError):
kkk2 += kkk10
kkk11 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
min_v=0.1)
with assert_raises(ValueError):
kkk2 += kkk11
kkk12 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
max_v=0.1)
with assert_raises(ValueError):
kkk2 += kkk12
kkk13 = treecorr.KKKCorrelation(min_sep=min_sep, bin_size=bin_size, nbins=nrbins,
nvbins=nrbins*2)
with assert_raises(ValueError):
kkk2 += kkk13
# Currently not implemented to only have cat2 or cat3
with assert_raises(NotImplementedError):
kkk.process(cat, cat2=cat)
with assert_raises(NotImplementedError):
kkk.process(cat, cat3=cat)
with assert_raises(NotImplementedError):
kkk.process_cross21(cat, cat)
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)
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)
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
nk = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, brute=True)
nk.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)
ww = w1[i] * w2
xi = ww * 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 = ',nk.npairs - true_npairs)
np.testing.assert_array_equal(nk.npairs, true_npairs)
print('true_weight = ',true_weight)
print('diff = ',nk.weight - true_weight)
np.testing.assert_allclose(nk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
print('true_xi = ',true_xi)
print('nk.xi = ',nk.xi)
np.testing.assert_allclose(nk.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/nk_direct.yaml')
cat1.write(config['file_name'])
cat2.write(config['file_name2'])
treecorr.corr2(config)
data = fitsio.read(config['nk_file_name'])
np.testing.assert_allclose(data['r_nom'], nk.rnom)
np.testing.assert_allclose(data['npairs'], nk.npairs)
np.testing.assert_allclose(data['weight'], nk.weight)
np.testing.assert_allclose(data['kappa'], nk.xi, rtol=1.e-3)
# Invalid with only one file_name
del config['file_name2']
with assert_raises(TypeError):
treecorr.corr2(config)
config['file_name2'] = 'data/nk_direct_cat2.fits'
# Invalid to request compoensated if no rand_file
config['nk_statistic'] = 'compensated'
with assert_raises(TypeError):
treecorr.corr2(config)
# Repeat with binslop = 0, since the code flow is different from brute=True
# And don't do any top-level recursion so we actually test not going to the leaves.
nk = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins, bin_slop=0,
max_top=0)
nk.process(cat1, cat2)
np.testing.assert_array_equal(nk.npairs, true_npairs)
np.testing.assert_allclose(nk.weight, true_weight, rtol=1.e-5, atol=1.e-8)
np.testing.assert_allclose(nk.xi, true_xi, rtol=1.e-4, atol=1.e-8)
# Check a few basic operations with a NKCorrelation object.
do_pickle(nk)
nk2 = nk.copy()
nk2 += nk
np.testing.assert_allclose(nk2.npairs, 2*nk.npairs)
np.testing.assert_allclose(nk2.weight, 2*nk.weight)
np.testing.assert_allclose(nk2.meanr, 2*nk.meanr)
np.testing.assert_allclose(nk2.meanlogr, 2*nk.meanlogr)
np.testing.assert_allclose(nk2.xi, 2*nk.xi)
nk2.clear()
nk2 += nk
np.testing.assert_allclose(nk2.npairs, nk.npairs)
np.testing.assert_allclose(nk2.weight, nk.weight)
np.testing.assert_allclose(nk2.meanr, nk.meanr)
np.testing.assert_allclose(nk2.meanlogr, nk.meanlogr)
np.testing.assert_allclose(nk2.xi, nk.xi)
ascii_name = 'output/nk_ascii.txt'
nk.write(ascii_name, precision=16)
nk3 = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
nk3.read(ascii_name)
np.testing.assert_allclose(nk3.npairs, nk.npairs)
np.testing.assert_allclose(nk3.weight, nk.weight)
np.testing.assert_allclose(nk3.meanr, nk.meanr)
np.testing.assert_allclose(nk3.meanlogr, nk.meanlogr)
np.testing.assert_allclose(nk3.xi, nk.xi)
with assert_raises(TypeError):
nk2 += config
nk4 = treecorr.NKCorrelation(min_sep=min_sep/2, max_sep=max_sep, nbins=nbins)
with assert_raises(ValueError):
nk2 += nk4
nk5 = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep*2, nbins=nbins)
with assert_raises(ValueError):
nk2 += nk5
nk6 = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins*2)
with assert_raises(ValueError):
nk2 += nk6
fits_name = 'output/nk_fits.fits'
nk.write(fits_name)
nk4 = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins)
nk4.read(fits_name)
np.testing.assert_allclose(nk4.npairs, nk.npairs)
np.testing.assert_allclose(nk4.weight, nk.weight)
np.testing.assert_allclose(nk4.meanr, nk.meanr)
np.testing.assert_allclose(nk4.meanlogr, nk.meanlogr)
np.testing.assert_allclose(nk4.xi, nk.xi)
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 bin_slop=0.
ngal = 200
s = 10.
np.random.seed(8675309)
x1 = np.random.normal(0, s, (ngal, ))
y1 = np.random.normal(0, s, (ngal, ))
w1 = np.random.random(ngal)
k1 = np.random.normal(10, 1, (ngal, ))
x2 = np.random.normal(0, s, (ngal, ))
y2 = np.random.normal(0, s, (ngal, ))
w2 = np.random.random(ngal)
k2 = np.random.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,
bin_slop=0.)
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 not precisely 0, since the code flow is different for bin_slop == 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=1.e-16,
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)