def __setstate_pre000305__(self, state):
""" compatible version of setstate for files generated before 0.3.5 """
edges = state['edges']
self.attrs = state['attrs']
self.poles = BinnedStatistic(['k'], [edges],
state['poles'],
fields_to_sum=['modes'])
def _make_datasets(self, edges, poles, power, coords, attrs):
if self.attrs['mode'] == '1d':
power = BinnedStatistic(['k'],
edges,
power,
fields_to_sum=['modes'],
coords=coords,
**attrs)
else:
power = BinnedStatistic(['k', 'mu'],
edges,
power,
fields_to_sum=['modes'],
coords=coords,
**attrs)
if poles is not None:
poles = BinnedStatistic(['k'], [power.edges['k']],
poles,
fields_to_sum=['modes'],
coords=[power.coords['k']],
**attrs)
return power, poles
def test_2d_load(comm):
# load plaintext format
with pytest.warns(FutureWarning):
ds1 = BinnedStatistic.from_plaintext(['k', 'mu'], os.path.join(data_dir, 'dataset_2d_deprecated.dat'))
# load from JSON
ds2 = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# same data?
for name in ds1:
testing.assert_almost_equal(ds1[name], ds2[name])
def test_2d_load(comm):
# load plaintext format
with pytest.warns(FutureWarning):
ds1 = BinnedStatistic.from_plaintext(['k', 'mu'], os.path.join(data_dir, 'dataset_2d_deprecated.dat'))
# load from JSON
ds2 = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# same data?
for name in ds1:
testing.assert_almost_equal(ds1[name], ds2[name])
def __setstate__(self, state):
edges = state.pop('edges')
dims = state.pop('dims')
self.__dict__.update(state)
self.corr = BinnedStatistic(dims, edges, self.corr)
if self.wp is not None:
# NOTE: only edges[0], second dimension was summed over
self.wp = BinnedStatistic(dims[:1], edges[:1], self.wp)
for pc in ['D1D2', 'D1R2', 'D2R1', 'R1R2']:
val = getattr(self, pc)
if val is not None:
setattr(self, pc, BinnedStatistic(dims, edges, val))
def run(self):
"""
Compute the power spectrum multipoles. This function does not return
anything, but adds several attributes (see below).
Attributes
----------
edges : array_like
the edges of the wavenumber bins
poles : :class:`~nbodykit.binned_statistic.BinnedStatistic`
a BinnedStatistic object that behaves similar to a structured array, with
fancy slicing and re-indexing; it holds the measured multipole
results, as well as the number of modes (``modes``) and average
wavenumbers values in each bin (``k``)
attrs : dict
dictionary holding input parameters and several important quantites
computed during execution:
#. data.N, randoms.N :
the unweighted number of data and randoms objects
#. data.W, randoms.W :
the weighted number of data and randoms objects, using the
column specified as the completeness weights
#. alpha :
the ratio of ``data.W`` to ``randoms.W``
#. data.norm, randoms.norm :
the normalization of the power spectrum, computed from either
the "data" or "randoms" catalog (they should be similar).
See equations 13 and 14 of arxiv:1312.4611.
#. data.shotnoise, randoms.shotnoise :
the shot noise values for the "data" and "random" catalogs;
See equation 15 of arxiv:1312.4611.
#. shotnoise :
the total shot noise for the power spectrum, equal to
``data.shotnoise`` + ``randoms.shotnoise``; this should be subtracted from
the monopole.
#. BoxSize :
the size of the Cartesian box used to grid the data and
randoms objects on a Cartesian mesh.
For further details on the meta-data, see
:ref:`the documentation <fkp-meta-data>`.
"""
pm = self.first.pm
# setup the binning in k out to the minimum nyquist frequency
dk = 2 * numpy.pi / pm.BoxSize.min(
) if self.attrs['dk'] is None else self.attrs['dk']
self.edges = numpy.arange(
self.attrs['kmin'],
numpy.pi * pm.Nmesh.min() / pm.BoxSize.max() + dk / 2, dk)
# measure the binned 1D multipoles in Fourier space
poles = self._compute_multipoles()
# set all the necessary results
self.poles = BinnedStatistic(['k'], [self.edges],
poles,
fields_to_sum=['modes'],
**self.attrs)
def _load(kind, box=None):
"""
Internal function to load the QPM results.
"""
d = os.environ['THESIS_DIR']
d = os.path.join(d, 'boss_dr12_mocks', 'Results', 'QPM', 'nbodykit',
'redshift', kind)
# the pattern
box = "*" if box is None else "%04d" % box
if kind == 'power':
files = glob(
os.path.join(d,
f"pkmu_qpm_unscaled_{box}_0.6452_dk005_Nmu100.dat"))
dims = ['k', 'mu']
names = ['k', 'mu', 'power', 'modes']
else:
files = glob(
os.path.join(d,
f"poles_qpm_unscaled_{box}_0.6452_dk005_Nmu100.dat"))
dims = ['k']
names = ['k', 'power_0', 'power_2', 'power_4', 'modes']
toret = []
for f in files:
d = BinnedStatistic.from_plaintext(dims, f)
if kind == 'poles':
for i, name in enumerate(names):
d.rename_variable('col_%d' % i, name)
toret.append(d)
return toret
def test_reindex(comm):
import warnings
dataset = BinnedStatistic.from_json(
os.path.join(data_dir, 'dataset_2d.json'))
with pytest.raises(ValueError):
new, spacing = dataset.reindex('k',
0.005,
force=True,
return_spacing=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
weights = numpy.random.random(dataset.shape)
dataset['weights'] = weights
new, spacing = dataset.reindex('k',
0.02,
weights='weights',
force=True,
return_spacing=True)
diff = numpy.diff(new.coords['k'])
assert numpy.alltrue(diff > numpy.diff(dataset.coords['k'])[0])
with pytest.raises(ValueError):
new = dataset.reindex('mu', 0.4, force=False)
new = dataset.reindex('mu', 0.4, force=True)
def test_copy(comm):
dataset = BinnedStatistic.from_json(
os.path.join(data_dir, 'dataset_2d.json'))
copy = dataset.copy()
for var in dataset:
testing.assert_array_equal(dataset[var], copy[var])
def test_average(comm):
import warnings
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# unweighted
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
avg = dataset.average('mu')
for var in dataset.variables:
if var in dataset._fields_to_sum:
x = numpy.nansum(dataset[var], axis=-1)
else:
x = numpy.nanmean(dataset[var], axis=-1)
testing.assert_allclose(x, avg[var])
# weighted
weights = numpy.random.random(dataset.shape)
dataset['weights'] = weights
avg = dataset.average('mu', weights='weights')
for var in dataset:
if var in dataset._fields_to_sum:
x = numpy.nansum(dataset[var], axis=-1)
else:
x = numpy.nansum(dataset[var]*dataset['weights'], axis=-1)
x /= dataset['weights'].sum(axis=-1)
testing.assert_allclose(x, avg[var])
def test_to_json(comm):
# load from JSON
ds1 = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_1d.json'))
# to JSON
with tempfile.NamedTemporaryFile(delete=False) as ff:
ds1.to_json(ff.name)
ds2 = BinnedStatistic.from_json(ff.name)
# same data?
for name in ds1:
testing.assert_almost_equal(ds1[name], ds2[name])
# cleanup
os.remove(ff.name)
def test_average(comm):
import warnings
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# unweighted
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
avg = dataset.average('mu')
for var in dataset.variables:
if var in dataset._fields_to_sum:
x = numpy.nansum(dataset[var], axis=-1)
else:
x = numpy.nanmean(dataset[var], axis=-1)
testing.assert_allclose(x, avg[var])
# weighted
weights = numpy.random.random(dataset.shape)
dataset['weights'] = weights
avg = dataset.average('mu', weights='weights')
for var in dataset:
if var in dataset._fields_to_sum:
x = numpy.nansum(dataset[var], axis=-1)
else:
x = numpy.nansum(dataset[var]*dataset['weights'], axis=-1)
x /= dataset['weights'].sum(axis=-1)
testing.assert_allclose(x, avg[var])
def NaturalEstimator(data_paircount):
"""
Internal function to computing the correlation function using
analytic randoms and the so-called "natural" correlation function
estimator, :math:`DD/RR - 1`.
"""
# data1 x data2
D1D2 = data_paircount.pairs
attrs = data_paircount.attrs
# determine the sample sizes
ND1, ND2 = attrs['N1'], attrs['N2']
edges = D1D2.edges
mode = attrs['mode']
BoxSize = attrs['BoxSize']
# analytic randoms - randoms calculation assuming uniform distribution
_R1R2 = AnalyticUniformRandoms(mode, edges, BoxSize)(ND1, ND2)
edges = [D1D2.edges[d] for d in D1D2.dims]
R1R2 = BinnedStatistic(D1D2.dims, edges, _R1R2.view([('npairs', 'f8')]))
# and compute the correlation function as DD/RR - 1
CF = (D1D2['npairs'] * D1D2['weightavg']) / R1R2['npairs'] - 1.
# create a BinnedStatistic holding the CF
CF = _create_tpcf_result(D1D2, CF)
return R1R2, CF
def test_to_json(comm):
# load from JSON
ds1 = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_1d.json'))
# to JSON
with tempfile.NamedTemporaryFile(delete=False) as ff:
ds1.to_json(ff.name)
ds2 = BinnedStatistic.from_json(ff.name)
# same data?
for name in ds1:
testing.assert_almost_equal(ds1[name], ds2[name])
# cleanup
os.remove(ff.name)
def __setstate__(self, state):
self.__dict__.update(state)
edges = self.attrs['edges']
# reconstruct the result based on mode
kws = {'fields_to_sum': ['npairs']}
if self.attrs['mode'] == '1d':
dims, edges = ['r'], [edges]
elif self.attrs['mode'] == '2d':
muedges = numpy.linspace(0, 1., self.attrs['Nmu'] + 1)
dims, edges = ['r', 'mu'], [edges, muedges]
elif self.attrs['mode'] == 'projected':
piedges = numpy.linspace(0, self.attrs['pimax'],
self.attrs['pimax'] + 1)
dims, edges = ['rp', 'pi'], [edges, piedges]
elif self.attrs['mode'] == 'angular':
dims, edges = ['theta'], [edges]
else:
valid = ['1d', '2d', 'angular', 'projected']
args = (self.attrs['mode'], valid)
raise ValueError("mode = '%s' should be one of %s" % args)
# save the result as a BinnedStatistic
self.pairs = BinnedStatistic(dims, edges, self.pairs, **kws)
def test_sel(comm):
dataset = BinnedStatistic.from_json(
os.path.join(data_dir, 'dataset_2d.json'))
# no exact match fails
with pytest.raises(IndexError):
sliced = dataset.sel(k=0.1)
# this should be squeezed
sliced = dataset.sel(k=0.1, method='nearest')
assert len(sliced.dims) == 1
# this is not squeezed
sliced = dataset.sel(k=[0.1], method='nearest')
assert sliced.shape[0] == 1
# this return empty k with arbitary edges.
sliced = dataset.sel(k=[], method='nearest')
assert sliced.shape[0] == 0
# slice in a specific k-range
sliced = dataset.sel(k=slice(0.02, 0.15), mu=[0.5], method='nearest')
assert sliced.shape[1] == 1
assert numpy.alltrue((sliced['k'] >= 0.02) & (sliced['k'] <= 0.15))
def test_rename_variable(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
test = numpy.zeros(dataset.shape)
dataset['test'] = test
dataset.rename_variable('test', 'renamed_test')
assert 'renamed_test' in dataset
assert 'test' not in dataset
def to_pkmu(self, mu_edges, max_ell):
"""
Invert the measured multipoles :math:`P_\ell(k)` into power
spectrum wedges, :math:`P(k,\mu)`.
Parameters
----------
mu_edges : array_like
the edges of the :math:`\mu` bins
max_ell : int
the maximum multipole to use when computing the wedges;
all even multipoles with :math:`ell` less than or equal
to this number are included
Returns
-------
pkmu : BinnedStatistic
a data set holding the :math:`P(k,\mu)` wedges
"""
from scipy.special import legendre
from scipy.integrate import quad
def compute_coefficient(ell, mumin, mumax):
"""
Compute how much each multipole contributes to a given wedges.
This returns:
.. math::
\frac{1}{\mu_{max} - \mu_{max}} \int_{\mu_{min}}^{\mu^{max}} \mathcal{L}_\ell(\mu)
"""
norm = 1.0 / (mumax - mumin)
return norm * quad(lambda mu: legendre(ell)(mu), mumin, mumax)[0]
# make sure we have all the poles measured
ells = list(range(0, max_ell+1, 2))
if any('power_%d' %ell not in self.poles for ell in ells):
raise ValueError("measurements for ells=%s required if max_ell=%d" %(ells, max_ell))
# new data array
dtype = numpy.dtype([('power', 'c8'), ('k', 'f8'), ('mu', 'f8')])
data = numpy.zeros((self.poles.shape[0], len(mu_edges)-1), dtype=dtype)
# loop over each wedge
bounds = list(zip(mu_edges[:-1], mu_edges[1:]))
for imu, mulims in enumerate(bounds):
# add the contribution from each Pell
for ell in ells:
coeff = compute_coefficient(ell, *mulims)
data['power'][:,imu] += coeff * self.poles['power_%d' %ell]
data['k'][:,imu] = self.poles['k']
data['mu'][:,imu] = numpy.ones(len(data))*0.5*(mulims[1]+mulims[0])
dims = ['k', 'mu']
edges = [self.poles.edges['k'], mu_edges]
return BinnedStatistic(dims=dims, edges=edges, data=data, coords=[self.poles.coords['k'], None], **self.attrs)
def test_rename_variable(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
test = numpy.zeros(dataset.shape)
dataset['test'] = test
dataset.rename_variable('test', 'renamed_test')
assert 'renamed_test' in dataset
assert 'test' not in dataset
def _make_datasets(self):
if self.attrs['mode'] == '1d':
self.power = BinnedStatistic(['k'], [self.edges],
self.power,
fields_to_sum=['modes'],
**self.attrs)
else:
self.power = BinnedStatistic(['k', 'mu'],
self.edges,
self.power,
fields_to_sum=['modes'],
**self.attrs)
if self.poles is not None:
self.poles = BinnedStatistic(['k'], [self.power.edges['k']],
self.poles,
fields_to_sum=['modes'],
**self.attrs)
def _create_tpcf_result(D1D2, CF):
"""
Create a BinnedStatistic holding the correlation function
and average bin separation.
"""
x = D1D2.dims[0]
data = numpy.empty_like(CF, dtype=[('corr', 'f8'), (x, 'f8')])
data['corr'] = CF[:]
data[x] = D1D2[x]
edges = [D1D2.edges[d] for d in D1D2.dims]
return BinnedStatistic(D1D2.dims, edges, data)
def test_list_array_slice(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# get the first and last mu column
sliced = dataset[:,[0, -1]]
assert len(sliced.shape) == 2
assert sliced.dims == ['k', 'mu']
# make sure we grabbed the right data
for var in dataset:
testing.assert_array_equal(dataset[var][:,[0,-1]], sliced[var])
def test_list_array_slice(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# get the first and last mu column
sliced = dataset[:,[0, -1]]
assert len(sliced.shape) == 2
assert sliced.dims == ['k', 'mu']
# make sure we grabbed the right data
for var in dataset:
testing.assert_array_equal(dataset[var][:,[0,-1]], sliced[var])
def test_str(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# list all variable names
s = str(dataset)
# now just list total number of variables
dataset['test1'] = numpy.ones(dataset.shape)
dataset['test2'] = numpy.ones(dataset.shape)
s = str(dataset)
# this is the same as str
r = repr(dataset)
def test_str(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# list all variable names
s = str(dataset)
# now just list total number of variables
dataset['test1'] = numpy.ones(dataset.shape)
dataset['test2'] = numpy.ones(dataset.shape)
s = str(dataset)
# this is the same as str
r = repr(dataset)
def test_array_slice(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# get the first mu column
sliced = dataset[:,0]
assert sliced.shape[0] == dataset.shape[0]
assert len(sliced.shape) == 1
assert sliced.dims == ['k']
# get the first mu column but keep dimension
sliced = dataset[:,[0]]
assert sliced.shape[0] == dataset.shape[0]
assert sliced.shape[1] == 1
assert sliced.dims == ['k', 'mu']
def test_array_slice(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# get the first mu column
sliced = dataset[:,0]
assert sliced.shape[0] == dataset.shape[0]
assert len(sliced.shape) == 1
assert sliced.dims == ['k']
# get the first mu column but keep dimension
sliced = dataset[:,[0]]
assert sliced.shape[0] == dataset.shape[0]
assert sliced.shape[1] == 1
assert sliced.dims == ['k', 'mu']
def test_variable_set(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
modes = numpy.ones(dataset.shape)
# add new variable
dataset['TEST'] = modes
assert 'TEST' in dataset
# override existing variable
dataset['modes'] = modes
assert numpy.all(dataset['modes'] == 1.0)
# needs right shape
with pytest.raises(ValueError):
dataset['TEST'] = 10.
def test_variable_set(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
modes = numpy.ones(dataset.shape)
# add new variable
dataset['TEST'] = modes
assert 'TEST' in dataset
# override existing variable
dataset['modes'] = modes
assert numpy.all(dataset['modes'] == 1.0)
# needs right shape
with pytest.raises(ValueError):
dataset['TEST'] = 10.
def test_squeeze(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# need to specify which dimension to squeeze
with pytest.raises(ValueError):
squeezed = dataset.squeeze()
with pytest.raises(ValueError):
squeezed = dataset[[0],[0]].squeeze()
sliced = dataset[:,[2]]
with pytest.raises(ValueError):
squeezed = sliced.squeeze('k')
squeezed = sliced.squeeze('mu')
assert len(squeezed.dims) == 1
assert squeezed.shape[0] == sliced.shape[0]
def test_squeeze(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# need to specify which dimension to squeeze
with pytest.raises(ValueError):
squeezed = dataset.squeeze()
with pytest.raises(ValueError):
squeezed = dataset[[0],[0]].squeeze()
sliced = dataset[:,[2]]
with pytest.raises(ValueError):
squeezed = sliced.squeeze('k')
squeezed = sliced.squeeze('mu')
assert len(squeezed.dims) == 1
assert squeezed.shape[0] == sliced.shape[0]
def to_poles(self, poles):
r"""
Invert the measured wedges :math:`\xi(r,mu)` into correlation
multipoles, :math:`\xi_\ell(r)`.
To select a mu_range, use
.. code:: python
poles = self.sel(mu=slice(*mu_range), method='nearest').to_poles(poles)
Parameters
----------
poles: array_like
the list of multipoles to compute
Returns
-------
xi_ell : BinnedStatistic
a data set holding the :math:`\xi_\ell(r)` multipoles
"""
from scipy.special import legendre
from scipy.integrate import quad
# new data array
x = str(self.dims[0])
dtype = numpy.dtype([(x, 'f8')] + [('corr_%d' % ell, 'f8')
for ell in poles])
data = numpy.zeros((self.shape[0]), dtype=dtype)
dims = [x]
edges = [self.edges[x]]
# FIXME: use something fancier than the central point.
mu_bins = numpy.diff(self.edges['mu'])
mu_mid = (self.edges['mu'][1:] + self.edges['mu'][:-1]) / 2.
for ell in poles:
legendrePolynomial = (2. * ell + 1.) * legendre(ell)(mu_mid)
data['corr_%d' %
ell] = numpy.sum(self['corr'] * legendrePolynomial * mu_bins,
axis=-1) / numpy.sum(mu_bins)
data[x] = numpy.mean(self[x], axis=-1)
return BinnedStatistic(dims=dims, edges=edges, data=data, poles=poles)
def test_take(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
sliced = dataset.take(k=[8])
assert sliced.shape[0] == 1
assert len(sliced.dims) == 2
sliced = dataset.take(k=[])
assert sliced.shape[0] == 0
assert len(sliced.dims) == 2
dataset.take(k=dataset.coords['k'] < 0.3)
assert len(sliced.dims) == 2
dataset.take(dataset['modes'] > 0)
assert len(sliced.dims) == 2
dataset.take(dataset['k'] < 0.3)
assert len(sliced.dims) == 2
def test_take(comm):
dataset = BinnedStatistic.from_json(
os.path.join(data_dir, 'dataset_2d.json'))
sliced = dataset.take(k=[8])
assert sliced.shape[0] == 1
assert len(sliced.dims) == 2
sliced = dataset.take(k=[])
assert sliced.shape[0] == 0
assert len(sliced.dims) == 2
dataset.take(k=dataset.coords['k'] < 0.3)
assert len(sliced.dims) == 2
dataset.take(dataset['modes'] > 0)
assert len(sliced.dims) == 2
dataset.take(dataset['k'] < 0.3)
assert len(sliced.dims) == 2
def test_reindex(comm):
import warnings
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
with pytest.raises(ValueError):
new, spacing = dataset.reindex('k', 0.005, force=True, return_spacing=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=RuntimeWarning)
weights = numpy.random.random(dataset.shape)
dataset['weights'] = weights
new, spacing = dataset.reindex('k', 0.02, weights='weights', force=True, return_spacing=True)
diff = numpy.diff(new.coords['k'])
assert numpy.alltrue(diff > numpy.diff(dataset.coords['k'])[0])
with pytest.raises(ValueError):
new = dataset.reindex('mu', 0.4, force=False)
new = dataset.reindex('mu', 0.4, force=True)
def test_getitem(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# invalid key
with pytest.raises(KeyError):
bad = dataset['error']
# slice columns
sliced = dataset[['k', 'mu', 'power']]
sliced = dataset[('k', 'mu', 'power')]
# invalid slice
with pytest.raises(KeyError):
bad =dataset[['k', 'mu', 'error']]
# too many dims in slice
with pytest.raises(IndexError):
bad = dataset[0,0,0]
# cannot access single element of 2D power
with pytest.raises(IndexError):
bad = dataset[0,0]
def test_getitem(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# invalid key
with pytest.raises(KeyError):
bad = dataset['error']
# slice columns
sliced = dataset[['k', 'mu', 'power']]
sliced = dataset[('k', 'mu', 'power')]
# invalid slice
with pytest.raises(KeyError):
bad =dataset[['k', 'mu', 'error']]
# too many dims in slice
with pytest.raises(IndexError):
bad = dataset[0,0,0]
# cannot access single element of 2D power
with pytest.raises(IndexError):
bad = dataset[0,0]
def test_sel(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
# no exact match fails
with pytest.raises(IndexError):
sliced = dataset.sel(k=0.1)
# this should be squeezed
sliced = dataset.sel(k=0.1, method='nearest')
assert len(sliced.dims) == 1
# this is not squeezed
sliced = dataset.sel(k=[0.1], method='nearest')
assert sliced.shape[0] == 1
# this return empty k with arbitary edges.
sliced = dataset.sel(k=[], method='nearest')
assert sliced.shape[0] == 0
# slice in a specific k-range
sliced = dataset.sel(k=slice(0.02, 0.15), mu=[0.5], method='nearest')
assert sliced.shape[1] == 1
assert numpy.alltrue((sliced['k'] >= 0.02)&(sliced['k'] <= 0.15))
def _load(box=None, los=None):
"""
Internal function to load the N-cubic results.
"""
d = os.environ['THESIS_DIR']
d = os.path.join(d, 'boss_dr12_mocks', 'Results', 'ChallengeMocks',
'nbodykit', 'power')
if los is not None:
assert los in "xyz"
# the pattern
box = "*" if box is None else "%d" % box
los = "*" if los is None else los
files = glob(
os.path.join(
d, f"pkmu_challenge_boxN{box}_unscaled_dk005_Nmu100_{los}los.dat"))
toret = []
for f in files:
toret.append(BinnedStatistic.from_plaintext(['k', 'mu'], f))
return toret
def __setstate__(self, state):
self.__dict__.update(state)
self.power = BinnedStatistic(['k'], [self.edges], self.power)
def __setstate__(self, state):
self.attrs = state['attrs']
self.power = BinnedStatistic.from_state(state['power'])
if state['poles'] is not None:
self.poles = BinnedStatistic.from_state(state['poles'])
def _run(self, pos, w, pos_sec, w_sec, boxsize=None, bunchsize=10000):
"""
Internal function to run the 3PCF algorithm on the input data and
weights.
The input data/weights have already been domain-decomposed, and
the loads should be balanced on all ranks.
"""
# maximum radius
rmax = numpy.max(self.attrs['edges'])
# the array to hold output values
nbins = len(self.attrs['edges'])-1
Nell = len(self.attrs['poles'])
zeta = numpy.zeros((Nell,nbins,nbins), dtype='f8')
alms = {}
walms = {}
# compute the Ylm expressions we need
if self.comm.rank == 0:
self.logger.info("computing Ylm expressions...")
Ylm_cache = YlmCache(self.attrs['poles'], self.comm)
if self.comm.rank == 0:
self.logger.info("...done")
# make the KD-tree holding the secondaries
tree_sec = kdcount.KDTree(pos_sec, boxsize=boxsize).root
def callback(r, i, j, iprim=None):
# remove self pairs
valid = r > 0.
r = r[valid]; i = i[valid]
# normalized, re-centered position array (periodic)
dpos = (pos_sec[i] - pos[iprim])
# enforce periodicity in dpos
if boxsize is not None:
for axis, col in enumerate(dpos.T):
col[col > boxsize[axis]*0.5] -= boxsize[axis]
col[col <= -boxsize[axis]*0.5] += boxsize[axis]
recen_pos = dpos / r[:,numpy.newaxis]
# find the mapping of r to rbins
dig = numpy.searchsorted(self.attrs['edges'], r, side='left')
# evaluate all Ylms
Ylms = Ylm_cache(recen_pos[:,0]+1j*recen_pos[:,1], recen_pos[:,2])
# sqrt of primary weight
w0 = w[iprim]
# loop over each (l,m) pair
for (l,m) in Ylms:
# the Ylm evaluated at galaxy positions
weights = Ylms[(l,m)] * w_sec[i]
# sum over for each radial bin
alm = alms.setdefault((l, m), numpy.zeros(nbins, dtype='c16'))
walm = walms.setdefault((l, m), numpy.zeros(nbins, dtype='c16'))
r1 = numpy.bincount(dig, weights=weights.real, minlength=nbins+2)[1:-1]
alm[...] += r1
walm[...] += w0 * r1
if m != 0:
i1 = numpy.bincount(dig, weights=weights.imag, minlength=nbins+2)[1:-1]
alm[...] += 1j*i1
walm[...] += w0*1j*i1
# determine rank with largest load
loads = self.comm.allgather(len(pos))
largest_load = numpy.argmax(loads)
chunk_size = max(loads) // 10
# compute multipoles for each primary (s vector in the paper)
for iprim in range(len(pos)):
# alms must be clean for each primary particle; (s) in eq 15 and 8 of arXiv:1506.02040v2
alms.clear()
walms.clear()
tree_prim = kdcount.KDTree(numpy.atleast_2d(pos[iprim]), boxsize=boxsize).root
tree_sec.enum(tree_prim, rmax, process=callback, iprim=iprim, bunch=bunchsize)
if self.comm.rank == largest_load and iprim % chunk_size == 0:
self.logger.info("%d%% done" % (10*iprim//chunk_size))
# combine alms into zeta(s);
# this cannot be done in the callback because
# it is a nonlinear function (outer product) of alm.
for (l, m) in alms:
alm = alms[(l, m)]
walm = walms[(l, m)]
# compute alm * conjugate(alm)
alm_w_alm = numpy.outer(walm, alm.conj())
if m != 0: alm_w_alm += alm_w_alm.T # add in the -m contribution for m != 0
zeta[Ylm_cache.ell_to_iell[l], ...] += alm_w_alm.real
# sum across all ranks
zeta = self.comm.allreduce(zeta)
# normalize according to Eq. 15 of Slepian et al. 2015
# differs by factor of (4 pi)^2 / (2l+1) from the C++ code
zeta /= (4*numpy.pi)
# make a BinnedStatistic
dtype = numpy.dtype([('corr_%d' % ell, zeta.dtype) for ell in self.attrs['poles']])
data = numpy.empty(zeta.shape[-2:], dtype=dtype)
for i, ell in enumerate(self.attrs['poles']):
data['corr_%d' % ell] = zeta[i]
# save the result
edges = self.attrs['edges']
poles = BinnedStatistic(['r1', 'r2'], [edges, edges], data)
return poles
def __setstate__(self, state):
self.attrs = state['attrs']
self.power = BinnedStatistic.from_state(state['power'])
if state['poles'] is not None:
self.poles = BinnedStatistic.from_state(state['poles'])
def run(self):
"""
Run the algorithm. This attaches the following attributes to the class:
- :attr:`edges`
- :attr:`power`
Attributes
----------
edges : array_like
the edges of the wavenumber bins
power : :class:`~nbodykit.binned_statistic.BinnedStatistic`
a BinnedStatistic object that holds the projected power.
It stores the following variables:
- k :
the mean value for each ``k`` bin
- power :
complex array holding the real and imaginary components of the
projected power
- modes :
the number of Fourier modes averaged together in each bin
"""
c1 = self.first.compute(Nmesh=self.attrs['Nmesh'], mode='complex')
r1 = c1.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
# average along projected axes;
# part of product is the rfftn vs r2c (for axes)
# the rest is for the mean (Nmesh - axes)
c1 = numpy.fft.rfftn(r1) / self.attrs['Nmesh'].prod()
# compute the auto power of single supplied field
if self.first is self.second:
c2 = c1
else:
c2 = self.second.compute(Nmesh=self.attrs['Nmesh'], mode='complex')
r2 = c2.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
c2 = numpy.fft.rfftn(r2) / self.attrs['Nmesh'].prod() # average along projected axes
pk = c1 * c2.conj()
# clear the zero mode
pk.flat[0] = 0
shape = numpy.array([self.attrs['Nmesh'][i] for i in self.attrs['axes']], dtype='int')
boxsize = numpy.array([self.attrs['BoxSize'][i] for i in self.attrs['axes']])
I = numpy.eye(len(shape), dtype='int') * -2 + 1
k = [numpy.fft.fftfreq(N, 1. / (N * 2 * numpy.pi / L))[:pkshape].reshape(kshape) for N, L, kshape, pkshape in zip(shape, boxsize, I, pk.shape)]
kmag = sum(ki ** 2 for ki in k) ** 0.5
W = numpy.empty(pk.shape, dtype='f4')
W[...] = 2.0
W[..., 0] = 1.0
W[..., -1] = 1.0
dk = self.attrs['dk']
kmin = self.attrs['kmin']
axes = list(self.attrs['axes'])
kedges = numpy.arange(kmin, numpy.pi * self.attrs['Nmesh'][axes].min() / self.attrs['BoxSize'][axes].max() + dk/2, dk)
xsum = numpy.zeros(len(kedges) + 1)
Psum = numpy.zeros(len(kedges) + 1, dtype='complex128')
Nsum = numpy.zeros(len(kedges) + 1)
dig = numpy.digitize(kmag.flat, kedges)
xsum.flat += numpy.bincount(dig, weights=(W * kmag).flat, minlength=xsum.size)
Psum.real.flat += numpy.bincount(dig, weights=(W * pk.real).flat, minlength=xsum.size)
Psum.imag.flat += numpy.bincount(dig, weights=(W * pk.imag).flat, minlength=xsum.size)
Nsum.flat += numpy.bincount(dig, weights=W.flat, minlength=xsum.size)
self.power = numpy.empty(len(kedges) - 1,
dtype=[('k', 'f8'), ('power', 'c16'), ('modes', 'f8')])
with numpy.errstate(invalid='ignore', divide='ignore'):
self.power['k'] = (xsum / Nsum)[1:-1]
self.power['power'] = (Psum / Nsum)[1:-1] * boxsize.prod() # dimension is 'volume'
self.power['modes'] = Nsum[1:-1]
self.edges = kedges
self.power = BinnedStatistic(['k'], [self.edges], self.power)
def test_copy(comm):
dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))
copy = dataset.copy()
for var in dataset:
testing.assert_array_equal(dataset[var], copy[var])
