def _set_num_threads(self, num_threads):
if num_threads is None:
num_threads = self.config.get('num_threads',None)
if num_threads is None:
self.logger.debug('Set num_threads automatically from ncpu')
else:
self.logger.debug('Set num_threads = %d',num_threads)
treecorr.set_omp_threads(num_threads, self.logger)
def _set_num_threads(self, num_threads):
if num_threads is None:
num_threads = self.config.get('num_threads', None)
if num_threads is None:
self.logger.debug('Set num_threads automatically from ncpu')
else:
self.logger.debug('Set num_threads = %d', num_threads)
treecorr.set_omp_threads(num_threads, self.logger)
def test_zero_weight():
# Based on test_ra_dec, but where many galaxies have w=0.
# There used to be a bug where w=0 objects were not assigned to any patch.
ngal = 10000
s = 10.
rng = np.random.RandomState(8675309)
x = rng.normal(0, s, (ngal, ))
y = rng.normal(
0, s,
(ngal, )) + 100 # Put everything at large y, so smallish angle on sky
z = rng.normal(0, s, (ngal, ))
w = np.zeros(ngal)
w[np.random.choice(range(ngal), ngal // 10, replace=False)] = 1.0
ra, dec = coord.CelestialCoord.xyz_to_radec(x, y, z)
print('minra = ', np.min(ra) * coord.radians / coord.degrees)
print('maxra = ', np.max(ra) * coord.radians / coord.degrees)
print('mindec = ', np.min(dec) * coord.radians / coord.degrees)
print('maxdec = ', np.max(dec) * coord.radians / coord.degrees)
cat = treecorr.Catalog(ra=ra,
dec=dec,
ra_units='rad',
dec_units='rad',
w=w,
keep_zero_weight=True)
treecorr.set_omp_threads(1)
npatch = 16
field = cat.getNField()
t0 = time.time()
p, c = field.run_kmeans(npatch)
t1 = time.time()
print('patches = ', np.unique(p))
assert len(p) == cat.ntot
assert min(p) == 0
assert max(p) == npatch - 1
print('w>0 patches = ', np.unique(p[w > 0]))
print('w==0 patches = ', np.unique(p[w == 0]))
assert set(p[w > 0]) == set(p[w == 0])
def corr3(config, logger=None):
"""Run the full three-point correlation function code based on the parameters in the
given config dict.
The function print_corr3_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, corr3_valid_params, corr3_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 = %d',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', 'rand_file_list2', 1, logger)
cat3 = treecorr.read_catalogs(config, 'file_name3', 'rand_file_list3', 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)
rand3 = treecorr.read_catalogs(config, 'rand_file_name3', 'rand_file_list3', 1, logger)
if len(cat2) == 0 and len(rand2) > 0:
raise AttributeError("rand_file_name2 is invalid without file_name2")
if len(cat3) == 0 and len(rand3) > 0:
raise AttributeError("rand_file_name3 is invalid without file_name3")
logger.info("Done reading input catalogs")
# Do GGG correlation function if necessary
if 'ggg_file_name' in config: #or 'm3_file_name' in config:
logger.info("Start GGG calculations...")
ggg = treecorr.GGGCorrelation(config,logger)
ggg.process(cat1,cat2,cat3)
logger.info("Done GGG calculations.")
if 'ggg_file_name' in config:
ggg.write(config['ggg_file_name'])
if 'm3_file_name' in config:
ggg.writeMapSq(config['m3_file_name'])
# Do NNN correlation function if necessary
if 'nnn_file_name' in config:
if len(rand1) == 0:
raise AttributeError("rand_file_name is required for NNN correlation")
if len(cat2) > 0 and len(rand2) == 0:
raise AttributeError("rand_file_name2 is required for NNN cross-correlation")
if len(cat3) > 0 and len(rand3) == 0:
raise AttributeError("rand_file_name3 is required for NNN cross-correlation")
if (len(cat2) > 0) != (len(cat3) > 0):
raise NotImplementedError(
"Cannot yet handle 3-point corrleations with only two catalogs. "+
"Need both cat2 and cat3.")
logger.info("Start DDD calculations...")
ddd = treecorr.NNNCorrelation(config,logger)
ddd.process(cat1,cat2,cat3)
logger.info("Done DDD calculations.")
if len(cat2) == 0:
rrr = treecorr.NNNCorrelation(config,logger)
rrr.process(rand1)
logger.info("Done RRR calculations.")
if config['nnn_statistic'] == 'compensated':
drr = treecorr.NNNCorrelation(config,logger)
drr.process(cat1,rand2,rand3)
logger.info("Done DRR calculations.")
ddr = treecorr.NNNCorrelation(config,logger)
ddr.process(cat1,cat2,rand3)
logger.info("Done DDR calculations.")
ddd.write(config['nnn_file_name'],rrr,drr,ddr)
else:
ddd.write(config['nnn_file_name'],rrr)
else:
rrr = treecorr.NNNCorrelation(config,logger)
rrr.process(rand1,rand2,rand3)
logger.info("Done RRR calculations.")
if config['nnn_statistic'] == 'compensated':
drr = treecorr.NNNCorrelation(config,logger)
drr.process(cat1,rand2,rand3)
logger.info("Done DRR calculations.")
ddr = treecorr.NNNCorrelation(config,logger)
ddr.process(cat1,cat2,rand3)
logger.info("Done DDR calculations.")
rdr = treecorr.NNNCorrelation(config,logger)
rdr.process(rand1,cat2,rand3)
logger.info("Done RDR calculations.")
rrd = treecorr.NNNCorrelation(config,logger)
rrd.process(rand1,rand2,cat3)
logger.info("Done RRD calculations.")
drd = treecorr.NNNCorrelation(config,logger)
drd.process(cat1,rand2,cat3)
logger.info("Done DRD calculations.")
rdd = treecorr.NNNCorrelation(config,logger)
rdd.process(rand1,cat2,cat3)
logger.info("Done RDD calculations.")
ddd.write(config['nnn_file_name'],rrr,drr,ddr,rdr,rrd,drd,rdd)
else:
ddd.write(config['nnn_file_name'],rrr)
# Do KKK correlation function if necessary
if 'kkk_file_name' in config:
logger.info("Start KKK calculations...")
kkk = treecorr.KKKCorrelation(config,logger)
kkk.process(cat1,cat2,cat3)
logger.info("Done KKK calculations.")
kkk.write(config['kkk_file_name'])
# Do NNG correlation function if necessary
if False:
#if 'nng_file_name' in config or 'nnm_file_name' in config:
if len(cat3) == 0:
raise AttributeError("file_name3 is required for nng correlation")
logger.info("Start NNG calculations...")
nng = treecorr.NNGCorrelation(config,logger)
nng.process(cat1,cat2,cat3)
logger.info("Done NNG calculation.")
# The default ng_statistic is compensated _iff_ rand files are given.
rrg = None
if len(rand1) == 0:
if config.get('nng_statistic',None) == 'compensated':
raise AttributeError("rand_files is required for nng_statistic = compensated")
elif config.get('nng_statistic','compensated') == 'compensated':
rrg = treecorr.NNGCorrelation(config,logger)
rrg.process(rand1,rand1,cat2)
logger.info("Done RRG calculation.")
if 'nng_file_name' in config:
nng.write(config['nng_file_name'], rrg)
if 'nnm_file_name' in config:
nng.writeNNMap(config['nnm_file_name'], rrg)
# Do NNK correlation function if necessary
if False:
#if 'nnk_file_name' in config:
if len(cat3) == 0:
raise AttributeError("file_name3 is required for nnk correlation")
logger.info("Start NNK calculations...")
nnk = treecorr.NNKCorrelation(config,logger)
nnk.process(cat1,cat2,cat3)
logger.info("Done NNK calculation.")
rrk = None
if len(rand1) == 0:
if config.get('nnk_statistic',None) == 'compensated':
raise AttributeError("rand_files is required for nnk_statistic = compensated")
elif config.get('nnk_statistic','compensated') == 'compensated':
rrk = treecorr.NNKCorrelation(config,logger)
rrk.process(rand1,rand1,cat2)
logger.info("Done RRK calculation.")
nnk.write(config['nnk_file_name'], rrk)
# Do KKG correlation function if necessary
if False:
#if 'kkg_file_name' in config:
if len(cat3) == 0:
raise AttributeError("file_name3 is required for kkg correlation")
logger.info("Start KKG calculations...")
kkg = treecorr.KKGCorrelation(config,logger)
kkg.process(cat1,cat2,cat3)
logger.info("Done KKG calculation.")
kkg.write(config['kkg_file_name'])
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_omp():
"""Test setting the number of omp threads.
"""
import multiprocessing
# If num_threads <= 0 or None, get num from cpu_count
cpus = multiprocessing.cpu_count()
assert treecorr.set_omp_threads(0) > 0
assert treecorr.set_omp_threads(0) <= cpus
assert treecorr.set_omp_threads(None) > 0
assert treecorr.set_omp_threads(None) <= cpus
# If num_threads == 1, it should always set to 1
assert treecorr.set_omp_threads(1) == 1
# If num_threads > 1, it could be 1 or up to the input num_threads
assert treecorr.set_omp_threads(2) >= 1
assert treecorr.set_omp_threads(2) <= 2
assert treecorr.set_omp_threads(20) >= 1
assert treecorr.set_omp_threads(20) <= 20
# Repeat and check that appropriate messages are emitted
with CaptureLog() as cl:
num_threads = treecorr.set_omp_threads(0, logger=cl.logger)
assert "multiprocessing.cpu_count() = " in cl.output
assert "Telling OpenMP to use %s threads" % cpus in cl.output
with CaptureLog() as cl:
treecorr.set_omp_threads(None, logger=cl.logger)
assert "multiprocessing.cpu_count() = " in cl.output
assert "Telling OpenMP to use %s threads" % cpus in cl.output
with CaptureLog() as cl:
treecorr.set_omp_threads(1, logger=cl.logger)
assert "multiprocessing.cpu_count() = " not in cl.output
assert "Telling OpenMP to use 1 threads" in cl.output
assert "Using %s threads" % num_threads not in cl.output
assert "Unable to use multiple threads" not in cl.output
with CaptureLog() as cl:
treecorr.set_omp_threads(2, logger=cl.logger)
assert "multiprocessing.cpu_count() = " not in cl.output
assert "Telling OpenMP to use 2 threads" in cl.output
# It's hard to tell what happens in the next step, since we can't control what
# treecorr._lib.SetOMPThreads does. It depends on whether OpenMP is enabled and
# how many cores are available. So let's mock it up.
if sys.version_info < (3, ): return # mock only available on python 3
from unittest import mock
with mock.patch('treecorr._lib') as _lib:
# First mock with OpenMP enables and able to use lots of threads
_lib.SetOMPThreads = lambda x: x
assert treecorr.set_omp_threads(20) == 20
with CaptureLog() as cl:
treecorr.set_omp_threads(20, logger=cl.logger)
assert "OpenMP reports that it will use 20 threads" in cl.output
assert "Using 20 threads" in cl.output
# Next only 4 threads available
_lib.SetOMPThreads = lambda x: 4 if x > 4 else x
assert treecorr.set_omp_threads(20) == 4
with CaptureLog() as cl:
treecorr.set_omp_threads(20, logger=cl.logger)
assert "OpenMP reports that it will use 4 threads" in cl.output
assert "Using 4 threads" in cl.output
assert treecorr.set_omp_threads(2) == 2
with CaptureLog() as cl:
treecorr.set_omp_threads(2, logger=cl.logger)
assert "OpenMP reports that it will use 2 threads" in cl.output
# Finally, no OpenMP
_lib.SetOMPThreads = lambda x: 1
assert treecorr.set_omp_threads(20) == 1
with CaptureLog() as cl:
treecorr.set_omp_threads(20, logger=cl.logger)
assert "OpenMP reports that it will use 1 threads" in cl.output
assert "Unable to use multiple threads" in cl.output
def corr3(config, logger=None):
"""Run the full three-point correlation function code based on the parameters in the
given config dict.
The function print_corr3_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, corr3_valid_params,
corr3_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', 'rand_file_list2', 1,
logger)
cat3 = treecorr.read_catalogs(config, 'file_name3', 'rand_file_list3', 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)
rand3 = treecorr.read_catalogs(config, 'rand_file_name3',
'rand_file_list3', 1, logger)
if len(cat2) == 0 and len(rand2) > 0:
raise AttributeError("rand_file_name2 is invalid without file_name2")
if len(cat3) == 0 and len(rand3) > 0:
raise AttributeError("rand_file_name3 is invalid without file_name3")
logger.info("Done reading input catalogs")
# Do GGG correlation function if necessary
if 'ggg_file_name' in config: #or 'm3_file_name' in config:
logger.info("Start GGG calculations...")
ggg = treecorr.GGGCorrelation(config, logger)
ggg.process(cat1, cat2, cat3)
logger.info("Done GGG calculations.")
if 'ggg_file_name' in config:
ggg.write(config['ggg_file_name'])
if 'm3_file_name' in config:
ggg.writeMapSq(config['m3_file_name'])
# Do NNN correlation function if necessary
if 'nnn_file_name' in config:
if len(rand1) == 0:
raise AttributeError(
"rand_file_name is required for NNN correlation")
if len(cat2) > 0 and len(rand2) == 0:
raise AttributeError(
"rand_file_name2 is required for NNN cross-correlation")
if len(cat3) > 0 and len(rand3) == 0:
raise AttributeError(
"rand_file_name3 is required for NNN cross-correlation")
if (len(cat2) > 0) != (len(cat3) > 0):
raise NotImplementedError(
"Cannot yet handle 3-point corrleations with only two catalogs. "
+ "Need both cat2 and cat3.")
logger.info("Start DDD calculations...")
ddd = treecorr.NNNCorrelation(config, logger)
ddd.process(cat1, cat2, cat3)
logger.info("Done DDD calculations.")
if len(cat2) == 0:
rrr = treecorr.NNNCorrelation(config, logger)
rrr.process(rand1)
logger.info("Done RRR calculations.")
# For the next step, just make cat2 = cat3 = cat1 and rand2 = rand3 = rand1.
cat2 = cat3 = cat1
rand2 = rand3 = rand1
else:
rrr = treecorr.NNNCorrelation(config, logger)
rrr.process(rand1, rand2, rand3)
logger.info("Done RRR calculations.")
if config['nnn_statistic'] == 'compensated':
drr = treecorr.NNNCorrelation(config, logger)
drr.process(cat1, rand2, rand3)
logger.info("Done DRR calculations.")
ddr = treecorr.NNNCorrelation(config, logger)
ddr.process(cat1, cat2, rand3)
logger.info("Done DDR calculations.")
rdr = treecorr.NNNCorrelation(config, logger)
rdr.process(rand1, cat2, rand3)
logger.info("Done RDR calculations.")
rrd = treecorr.NNNCorrelation(config, logger)
rrd.process(rand1, rand2, cat3)
logger.info("Done RRD calculations.")
drd = treecorr.NNNCorrelation(config, logger)
drd.process(cat1, rand2, cat3)
logger.info("Done DRD calculations.")
rdd = treecorr.NNNCorrelation(config, logger)
rdd.process(rand1, cat2, cat3)
logger.info("Done RDD calculations.")
ddd.write(config['nnn_file_name'], rrr, drr, rdr, rrd, ddr, drd,
rdd)
else:
ddd.write(config['nnn_file_name'], rrr)
# Do KKK correlation function if necessary
if 'kkk_file_name' in config:
logger.info("Start KKK calculations...")
kkk = treecorr.KKKCorrelation(config, logger)
kkk.process(cat1, cat2, cat3)
logger.info("Done KKK calculations.")
kkk.write(config['kkk_file_name'])
# Do NNG correlation function if necessary
if False:
#if 'nng_file_name' in config or 'nnm_file_name' in config:
if len(cat3) == 0:
raise AttributeError("file_name3 is required for nng correlation")
logger.info("Start NNG calculations...")
nng = treecorr.NNGCorrelation(config, logger)
nng.process(cat1, cat2, cat3)
logger.info("Done NNG calculation.")
# The default ng_statistic is compensated _iff_ rand files are given.
rrg = None
if len(rand1) == 0:
if config.get('nng_statistic', None) == 'compensated':
raise AttributeError(
"rand_files is required for nng_statistic = compensated")
elif config.get('nng_statistic', 'compensated') == 'compensated':
rrg = treecorr.NNGCorrelation(config, logger)
rrg.process(rand1, rand1, cat2)
logger.info("Done RRG calculation.")
if 'nng_file_name' in config:
nng.write(config['nng_file_name'], rrg)
if 'nnm_file_name' in config:
nng.writeNNMap(config['nnm_file_name'], rrg)
# Do NNK correlation function if necessary
if False:
#if 'nnk_file_name' in config:
if len(cat3) == 0:
raise AttributeError("file_name3 is required for nnk correlation")
logger.info("Start NNK calculations...")
nnk = treecorr.NNKCorrelation(config, logger)
nnk.process(cat1, cat2, cat3)
logger.info("Done NNK calculation.")
rrk = None
if len(rand1) == 0:
if config.get('nnk_statistic', None) == 'compensated':
raise AttributeError(
"rand_files is required for nnk_statistic = compensated")
elif config.get('nnk_statistic', 'compensated') == 'compensated':
rrk = treecorr.NNKCorrelation(config, logger)
rrk.process(rand1, rand1, cat2)
logger.info("Done RRK calculation.")
nnk.write(config['nnk_file_name'], rrk)
# Do KKG correlation function if necessary
if False:
#if 'kkg_file_name' in config:
if len(cat3) == 0:
raise AttributeError("file_name3 is required for kkg correlation")
logger.info("Start KKG calculations...")
kkg = treecorr.KKGCorrelation(config, logger)
kkg.process(cat1, cat2, cat3)
logger.info("Done KKG calculation.")
kkg.write(config['kkg_file_name'])
def test_omp():
"""Test setting the number of omp threads.
"""
import multiprocessing
# If num_threads <= 0 or None, get num from cpu_count
cpus = multiprocessing.cpu_count()
assert treecorr.set_omp_threads(0) > 0
assert treecorr.set_omp_threads(0) <= cpus
assert treecorr.set_omp_threads(None) > 0
assert treecorr.set_omp_threads(None) <= cpus
# If num_threads == 1, it should always set to 1
assert treecorr.set_omp_threads(1) == 1
# If num_threads > 1, it could be 1 or up to the input num_threads
assert treecorr.set_omp_threads(2) >= 1
assert treecorr.set_omp_threads(2) <= 2
assert treecorr.set_omp_threads(20) >= 1
assert treecorr.set_omp_threads(20) <= 20
# Repeat and check that appropriate messages are emitted
with CaptureLog() as cl:
num_threads = treecorr.set_omp_threads(0, logger=cl.logger)
assert "multiprocessing.cpu_count() = " in cl.output
assert "Telling OpenMP to use %s threads"%cpus in cl.output
with CaptureLog() as cl:
treecorr.set_omp_threads(None, logger=cl.logger)
assert "multiprocessing.cpu_count() = " in cl.output
assert "Telling OpenMP to use %s threads"%cpus in cl.output
with CaptureLog() as cl:
treecorr.set_omp_threads(1, logger=cl.logger)
assert "multiprocessing.cpu_count() = " not in cl.output
assert "Telling OpenMP to use 1 threads" in cl.output
assert "Using %s threads"%num_threads not in cl.output
assert "Unable to use multiple threads" not in cl.output
with CaptureLog() as cl:
treecorr.set_omp_threads(2, logger=cl.logger)
assert "multiprocessing.cpu_count() = " not in cl.output
assert "Telling OpenMP to use 2 threads" in cl.output
# It's hard to tell what happens in the next step, since we can't control what
# treecorr._lib.SetOMPThreads does. It depends on whether OpenMP is enabled and
# how many cores are available. So let's mock it up.
if sys.version_info < (3,): return # mock only available on python 3
from unittest import mock
with mock.patch('treecorr.config._lib') as _lib:
# First mock with OpenMP enables and able to use lots of threads
_lib.SetOMPThreads = lambda x: x
assert treecorr.set_omp_threads(20) == 20
with CaptureLog() as cl:
treecorr.set_omp_threads(20, logger=cl.logger)
assert "OpenMP reports that it will use 20 threads" in cl.output
assert "Using 20 threads" in cl.output
# Next only 4 threads available
_lib.SetOMPThreads = lambda x: 4 if x > 4 else x
assert treecorr.set_omp_threads(20) == 4
with CaptureLog() as cl:
treecorr.set_omp_threads(20, logger=cl.logger)
assert "OpenMP reports that it will use 4 threads" in cl.output
assert "Using 4 threads" in cl.output
assert treecorr.set_omp_threads(2) == 2
with CaptureLog() as cl:
treecorr.set_omp_threads(2, logger=cl.logger)
assert "OpenMP reports that it will use 2 threads" in cl.output
# Finally, no OpenMP
_lib.SetOMPThreads = lambda x: 1
assert treecorr.set_omp_threads(20) == 1
with CaptureLog() as cl:
treecorr.set_omp_threads(20, logger=cl.logger)
assert "OpenMP reports that it will use 1 threads" in cl.output
assert "Unable to use multiple threads" in cl.output
def corr3(config, logger=None):
"""Run the full three-point correlation function code based on the parameters in the
given config dict.
The function `print_corr3_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, corr3_valid_params,
corr3_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)
cat2 = treecorr.read_catalogs(config, 'file_name2', 'rand_file_list2', 1,
logger)
cat3 = treecorr.read_catalogs(config, 'file_name3', 'rand_file_list3', 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)
rand3 = treecorr.read_catalogs(config, 'rand_file_name3',
'rand_file_list3', 1, logger)
if len(cat1) == 0:
raise TypeError("Either file_name or file_list is required")
if len(cat2) == 0: cat2 = None
if len(cat3) == 0: cat3 = None
if len(rand1) == 0: rand1 = None
if len(rand2) == 0: rand2 = None
if len(rand3) == 0: rand3 = None
if cat2 is None and rand2 is not None:
raise TypeError("rand_file_name2 is invalid without file_name2")
if cat3 is None and rand3 is not None:
raise TypeError("rand_file_name3 is invalid without file_name3")
if (cat2 is None) != (cat3 is None):
raise NotImplementedError(
"Cannot yet handle 3-point corrleations with only two catalogs. " +
"Need both cat2 and cat3.")
logger.info("Done reading input catalogs")
# Do GGG correlation function if necessary
if 'ggg_file_name' in config or 'm3_file_name' in config:
logger.warning("Performing GGG calculations...")
ggg = treecorr.GGGCorrelation(config, logger)
ggg.process(cat1, cat2, cat3)
logger.info("Done GGG calculations.")
if 'ggg_file_name' in config:
ggg.write(config['ggg_file_name'])
logger.warning("Wrote GGG correlation to %s",
config['ggg_file_name'])
if 'm3_file_name' in config:
ggg.writeMap3(config['m3_file_name'])
logger.warning("Wrote Map3 values to %s", config['m3_file_name'])
# Do NNN correlation function if necessary
if 'nnn_file_name' in config:
logger.warning("Performing DDD calculations...")
ddd = treecorr.NNNCorrelation(config, logger)
ddd.process(cat1, cat2, cat3)
logger.info("Done DDD calculations.")
drr = None
rdr = None
rrd = None
ddr = None
drd = None
rdd = None
if rand1 is None:
if rand2 is not None or rand3 is not None:
raise TypeError(
"rand_file_name is required if rand2 or rand3 is given")
logger.warning(
"No random catalogs given. Only doing ntri calculation.")
rrr = None
elif cat2 is None:
logger.warning("Performing RRR calculations...")
rrr = treecorr.NNNCorrelation(config, logger)
rrr.process(rand1)
logger.info("Done RRR calculations.")
# For the next step, just make cat2 = cat3 = cat1 and rand2 = rand3 = rand1.
cat2 = cat3 = cat1
rand2 = rand3 = rand1
else:
if rand2 is None:
raise TypeError(
"rand_file_name2 is required when file_name2 is given")
if cat3 is not None and rand3 is None:
raise TypeError(
"rand_file_name3 is required when file_name3 is given")
logger.warning("Performing RRR calculations...")
rrr = treecorr.NNNCorrelation(config, logger)
rrr.process(rand1, rand2, rand3)
logger.info("Done RRR calculations.")
if rrr is not None and config['nnn_statistic'] == 'compensated':
logger.warning("Performing DRR calculations...")
drr = treecorr.NNNCorrelation(config, logger)
drr.process(cat1, rand2, rand3)
logger.info("Done DRR calculations.")
logger.warning("Performing DDR calculations...")
ddr = treecorr.NNNCorrelation(config, logger)
ddr.process(cat1, cat2, rand3)
logger.info("Done DDR calculations.")
logger.warning("Performing RDR calculations...")
rdr = treecorr.NNNCorrelation(config, logger)
rdr.process(rand1, cat2, rand3)
logger.info("Done RDR calculations.")
logger.warning("Performing RRD calculations...")
rrd = treecorr.NNNCorrelation(config, logger)
rrd.process(rand1, rand2, cat3)
logger.info("Done RRD calculations.")
logger.warning("Performing DRD calculations...")
drd = treecorr.NNNCorrelation(config, logger)
drd.process(cat1, rand2, cat3)
logger.info("Done DRD calculations.")
logger.warning("Performing RDD calculations...")
rdd = treecorr.NNNCorrelation(config, logger)
rdd.process(rand1, cat2, cat3)
logger.info("Done RDD calculations.")
ddd.write(config['nnn_file_name'], rrr, drr, rdr, rrd, ddr, drd, rdd)
logger.warning("Wrote NNN correlation to %s", config['nnn_file_name'])
# Do KKK correlation function if necessary
if 'kkk_file_name' in config:
logger.warning("Performing KKK calculations...")
kkk = treecorr.KKKCorrelation(config, logger)
kkk.process(cat1, cat2, cat3)
logger.info("Done KKK calculations.")
kkk.write(config['kkk_file_name'])
logger.warning("Wrote KKK correlation to %s", config['kkk_file_name'])
def corr3(config, logger=None):
"""Run the full three-point correlation function code based on the parameters in the
given config dict.
The function `print_corr3_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, corr3_valid_params, corr3_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)
cat2 = treecorr.read_catalogs(config, 'file_name2', 'rand_file_list2', 1, logger)
cat3 = treecorr.read_catalogs(config, 'file_name3', 'rand_file_list3', 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)
rand3 = treecorr.read_catalogs(config, 'rand_file_name3', 'rand_file_list3', 1, logger)
if len(cat1) == 0:
raise TypeError("Either file_name or file_list is required")
if len(cat2) == 0: cat2 = None
if len(cat3) == 0: cat3 = None
if len(rand1) == 0: rand1 = None
if len(rand2) == 0: rand2 = None
if len(rand3) == 0: rand3 = None
if cat2 is None and rand2 is not None:
raise TypeError("rand_file_name2 is invalid without file_name2")
if cat3 is None and rand3 is not None:
raise TypeError("rand_file_name3 is invalid without file_name3")
if (cat2 is None) != (cat3 is None):
raise NotImplementedError(
"Cannot yet handle 3-point corrleations with only two catalogs. "+
"Need both cat2 and cat3.")
logger.info("Done reading input catalogs")
# Do GGG correlation function if necessary
if 'ggg_file_name' in config or 'm3_file_name' in config:
logger.warning("Performing GGG calculations...")
ggg = treecorr.GGGCorrelation(config,logger)
ggg.process(cat1,cat2,cat3)
logger.info("Done GGG calculations.")
if 'ggg_file_name' in config:
ggg.write(config['ggg_file_name'])
logger.warning("Wrote GGG correlation to %s",config['ggg_file_name'])
if 'm3_file_name' in config:
ggg.writeMap3(config['m3_file_name'])
logger.warning("Wrote Map3 values to %s",config['m3_file_name'])
# Do NNN correlation function if necessary
if 'nnn_file_name' in config:
logger.warning("Performing DDD calculations...")
ddd = treecorr.NNNCorrelation(config,logger)
ddd.process(cat1,cat2,cat3)
logger.info("Done DDD calculations.")
drr = None
rdr = None
rrd = None
ddr = None
drd = None
rdd = None
if rand1 is None:
if rand2 is not None or rand3 is not None:
raise TypeError("rand_file_name is required if rand2 or rand3 is given")
logger.warning("No random catalogs given. Only doing ntri calculation.")
rrr = None
elif cat2 is None:
logger.warning("Performing RRR calculations...")
rrr = treecorr.NNNCorrelation(config,logger)
rrr.process(rand1)
logger.info("Done RRR calculations.")
# For the next step, just make cat2 = cat3 = cat1 and rand2 = rand3 = rand1.
cat2 = cat3 = cat1
rand2 = rand3 = rand1
else:
if rand2 is None:
raise TypeError("rand_file_name2 is required when file_name2 is given")
if cat3 is not None and rand3 is None:
raise TypeError("rand_file_name3 is required when file_name3 is given")
logger.warning("Performing RRR calculations...")
rrr = treecorr.NNNCorrelation(config,logger)
rrr.process(rand1,rand2,rand3)
logger.info("Done RRR calculations.")
if rrr is not None and config['nnn_statistic'] == 'compensated':
logger.warning("Performing DRR calculations...")
drr = treecorr.NNNCorrelation(config,logger)
drr.process(cat1,rand2,rand3)
logger.info("Done DRR calculations.")
logger.warning("Performing DDR calculations...")
ddr = treecorr.NNNCorrelation(config,logger)
ddr.process(cat1,cat2,rand3)
logger.info("Done DDR calculations.")
logger.warning("Performing RDR calculations...")
rdr = treecorr.NNNCorrelation(config,logger)
rdr.process(rand1,cat2,rand3)
logger.info("Done RDR calculations.")
logger.warning("Performing RRD calculations...")
rrd = treecorr.NNNCorrelation(config,logger)
rrd.process(rand1,rand2,cat3)
logger.info("Done RRD calculations.")
logger.warning("Performing DRD calculations...")
drd = treecorr.NNNCorrelation(config,logger)
drd.process(cat1,rand2,cat3)
logger.info("Done DRD calculations.")
logger.warning("Performing RDD calculations...")
rdd = treecorr.NNNCorrelation(config,logger)
rdd.process(rand1,cat2,cat3)
logger.info("Done RDD calculations.")
ddd.write(config['nnn_file_name'],rrr,drr,rdr,rrd,ddr,drd,rdd)
logger.warning("Wrote NNN correlation to %s",config['nnn_file_name'])
# Do KKK correlation function if necessary
if 'kkk_file_name' in config:
logger.warning("Performing KKK calculations...")
kkk = treecorr.KKKCorrelation(config,logger)
kkk.process(cat1,cat2,cat3)
logger.info("Done KKK calculations.")
kkk.write(config['kkk_file_name'])
logger.warning("Wrote KKK correlation to %s",config['kkk_file_name'])