def reference_sim_paircount(pos1,
w1,
redges,
Nmu,
boxsize,
pos2=None,
w2=None,
los=2):
"""Reference pair counting via kdcount"""
tree1 = correlate.points(pos1, boxsize=boxsize, weights=w1)
if pos2 is None:
tree2 = tree1
else:
tree2 = correlate.points(pos2, boxsize=boxsize, weights=w2)
bins = correlate.FlatSkyBinning(
redges,
Nmu,
los=los,
mu_min=0.,
absmu=True,
)
pc = correlate.paircount(tree1,
tree2,
bins,
np=0,
compute_mean_coords=True)
return numpy.nan_to_num(pc.pair_counts), numpy.nan_to_num(
pc.mean_centers[0]), pc.sum1
def reference_2pcf_smu(sedges,muedges,position1,weight1,position2=None,weight2=None,los='midpoint'):
"""Reference pair counting via kdcount"""
tree1 = correlate.points(position1,boxsize=None,weights=weight1)
if position2 is None: tree2 = tree1
else: tree2 = correlate.points(position2,boxsize=None,weights=weight2)
if los=='midpoint':
bins = correlate.RmuBinning(np.asarray(sedges),(len(muedges)-1),observer=(0,0,0),mu_min=muedges[0],mu_max=muedges[-1],absmu=False)
else:
bins = correlate.FlatSkyBinning(np.asarray(sedges),(len(muedges)-1),los='xyz'.index(los),mu_min=muedges[0],mu_max=muedges[-1],absmu=False)
pc = correlate.paircount(tree2,tree1,bins,np=0,usefast=False,compute_mean_coords=True)
return pc.sum1
def compute_brutal_corr(datasources,
redges,
Nmu=0,
comm=None,
subsample=1,
los='z',
poles=[]):
r"""
Compute the correlation function by direct pair summation, either as a function
of separation (`R`) or as a function of separation and line-of-sight angle (`R`, `mu`)
The estimator used to compute the correlation function is:
.. math::
\xi(r, \mu) = DD(r, \mu) / RR(r, \mu) - 1.
where `DD` is the number of data-data pairs, and `RR` is the number of random-random pairs,
which is determined solely by the binning used, assuming a constant number density
Parameters
----------
datasources : list of DataSource objects
the list of data instances from which the 3D correlation will be computed
redges : array_like
the bin edges for the `R` variable
Nmu : int, optional
the number of desired `mu` bins, where `mu` is the cosine
of the angle from the line-of-sight. Default is `0`, in
which case the correlation function is binned as a function of `R` only
comm : MPI.Communicator, optional
the communicator to pass to the ``ParticleMesh`` object. If not
provided, ``MPI.COMM_WORLD`` is used
subsample : int, optional
downsample the input datasources by choosing 1 out of every `N` points.
Default is `1` (no subsampling).
los : str, {'x', 'y', 'z'}, optional
the dimension to treat as the line-of-sight; default is 'z'.
poles : list of int, optional
integers specifying the multipoles to compute from the 2D correlation function
Returns
-------
pc : :class:`kdcount.correlate.paircount`
the pair counting instance
xi : array_like
the correlation function result; if `poles` supplied, the shape is
`(len(redges)-1, len(poles))`, otherwise, the shape is either `(len(redges)-1, )`
or `(len(redges)-1, Nmu)`
RR : array_like
the number of random-random pairs (used as normalization of the data-data pairs)
"""
from pmesh.domain import GridND
from kdcount import correlate
# some setup
if los not in "xyz": raise ValueError("`los` must be `x`, `y`, or `z`")
los = "xyz".index(los)
poles = numpy.array(poles)
Rmax = redges[-1]
if comm is None: comm = MPI.COMM_WORLD
# determine processor division for domain decomposition
for Nx in range(int(comm.size**0.3333) + 1, 0, -1):
if comm.size % Nx == 0: break
else:
Nx = 1
for Ny in range(int(comm.size**0.5) + 1, 0, -1):
if (comm.size // Nx) % Ny == 0: break
else:
Ny = 1
Nz = comm.size // Nx // Ny
Nproc = [Nx, Ny, Nz]
# log some info
if comm.rank == 0:
logger.info('Nproc = %s' % str(Nproc))
logger.info('Rmax = %g' % Rmax)
# domain decomposition
grid = [
numpy.linspace(0,
datasources[0].BoxSize[i],
Nproc[i] + 1,
endpoint=True) for i in range(3)
]
domain = GridND(grid, comm=comm)
# read position for field #1
with datasources[0].open() as stream:
[[pos1]] = stream.read(['Position'], full=True)
pos1 = pos1[comm.rank * subsample // comm.size::subsample]
N1 = comm.allreduce(len(pos1))
# read position for field #2
if len(datasources) > 1:
with datasources[1].open() as stream:
[[pos2]] = stream.read(['Position'], full=True)
pos2 = pos2[comm.rank * subsample // comm.size::subsample]
N2 = comm.allreduce(len(pos2))
else:
pos2 = pos1
N2 = N1
# exchange field #1 positions
layout = domain.decompose(pos1, smoothing=0)
pos1 = layout.exchange(pos1)
if comm.rank == 0: logger.info('exchange pos1')
# exchange field #2 positions
if Rmax > datasources[0].BoxSize[0] * 0.25:
pos2 = numpy.concatenate(comm.allgather(pos2), axis=0)
else:
layout = domain.decompose(pos2, smoothing=Rmax)
pos2 = layout.exchange(pos2)
if comm.rank == 0: logger.info('exchange pos2')
# initialize the trees to hold the field points
tree1 = correlate.points(pos1, boxsize=datasources[0].BoxSize)
tree2 = correlate.points(pos2, boxsize=datasources[0].BoxSize)
# log the sizes of the trees
logger.info('rank %d correlating %d x %d' %
(comm.rank, len(tree1), len(tree2)))
if comm.rank == 0: logger.info('all correlating %d x %d' % (N1, N2))
# use multipole binning
if len(poles):
bins = correlate.FlatSkyMultipoleBinning(redges,
poles,
los,
compute_mean_coords=True)
# use (R, mu) binning
elif Nmu > 0:
bins = correlate.FlatSkyBinning(redges,
Nmu,
los,
compute_mean_coords=True)
# use R binning
else:
bins = correlate.RBinning(redges, compute_mean_coords=True)
# do the pair counting
# have to set usefast = False to get mean centers, or exception thrown
pc = correlate.paircount(tree2, tree1, bins, np=0, usefast=False)
pc.sum1[:] = comm.allreduce(pc.sum1)
# get the mean bin values, reducing from all ranks
pc.pair_counts[:] = comm.allreduce(pc.pair_counts)
with numpy.errstate(invalid='ignore'):
if bins.Ndim > 1:
for i in range(bins.Ndim):
pc.mean_centers[i][:] = comm.allreduce(
pc.mean_centers_sum[i]) / pc.pair_counts
else:
pc.mean_centers[:] = comm.allreduce(
pc.mean_centers_sum[0]) / pc.pair_counts
# compute the random pairs from the fractional volume
RR = 1. * N1 * N2 / datasources[0].BoxSize.prod()
if Nmu > 0:
dr3 = numpy.diff(pc.edges[0]**3)
dmu = numpy.diff(pc.edges[1])
RR *= 2. / 3. * numpy.pi * dr3[:, None] * dmu[None, :]
else:
RR *= 4. / 3. * numpy.pi * numpy.diff(pc.edges**3)
# return the correlation and the pair count object
xi = (1. * pc.sum1 / RR) - 1.0
if len(poles):
xi = xi.T # makes ell the second axis
xi[:, poles != 0] += 1.0 # only monopole gets the minus one
return pc, xi, RR