def test_count_symmetric():
pos1 = numpy.random.uniform(size=(1000000, 3)).astype('f4')
pos2 = numpy.array([[0.3, 0.5, 0.1]], dtype='f4')
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
assert_equal(tree2.count(tree1, (0, 0.1, 1.0)),
tree1.count(tree2, (0, 0.1, 1.0)))
def test_force2_slow():
from kdcount import force_kernels
from time import time
numpy.random.seed(13)
pos = numpy.random.uniform(size=(32 * 32 * 32 * 8, 3))
tt = KDTree(pos[:1], thresh=1)
tree = KDTree(pos, thresh=1)
mass = KDAttr(tree, numpy.ones(pos.shape[0]))
xmass = KDAttr(tree, pos * mass.input)
print(tt.root, tt.root.min, tt.root.max)
t0 = time()
force2 = tree.root.force2(force_kernels['count'],
tt.root,
mass,
xmass,
1.0 / len(pos)**0.3333 * 8,
eta=1.)
t2 = time() - t0
print('time', t2)
t0 = time()
force1 = tree.root.force(force_kernels['count'],
pos[:1],
mass,
xmass,
1.0 / len(pos)**0.3333 * 8,
eta=1.)
t1 = time() - t0
print('time', t1)
# FIXME: add more tests
assert_array_equal(force1, force2)
def test_fof():
numpy.random.seed(1000)
pos = numpy.linspace(0, 1, 10000, endpoint=False).reshape(-1, 1)
tree = KDTree(pos).root
label = tree.fof(1.1 / len(pos))
assert_equal(numpy.unique(label).size, 1)
label = tree.fof(0.8 / len(pos))
assert_equal(numpy.unique(label).size, len(pos))
def test_fof():
numpy.random.seed(1000)
pos = numpy.linspace(0, 1, 10000, endpoint=False).reshape(-1, 1)
tree = KDTree(pos).root
label = tree.fof(1.1/ len(pos))
assert_equal(numpy.unique(label).size, 1)
label = tree.fof(0.8/ len(pos))
assert_equal(numpy.unique(label).size, len(pos))
def test_enum_count_agree():
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f4')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f4')
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
N = [0]
def process1(r, i, j):
N[0] += len(r)
tree1.enum(tree2, rmax=1.0, process=process1)
c = tree1.count(tree2, r=1.0)
assert_equal(N[0], c)
def test_build():
numpy.random.seed(1000)
pos = numpy.random.uniform(size=(1000, 3))
tree = KDTree(pos).root
N = [0, 0]
def process1(r, i, j):
N[0] += len(r)
def process2(r, i, j):
N[1] += len(r)
tree.enum(tree, rmax=0.1, process=process1)
def test_count():
numpy.random.seed(1234)
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f8')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f8')
dist = ((pos1[:, None, :] - pos2[None, :, :])**2).sum(axis=-1)
truth = (dist < 1).sum()
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
c = tree1.count(tree2, r=1.0)
assert_equal(c, truth)
def test_count():
numpy.random.seed(1234)
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f8')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f8')
dist = ((pos1[:, None, :] - pos2[None, :, :]) ** 2).sum(axis=-1)
truth = (dist < 1).sum()
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
c = tree1.count(tree2, r=1.0)
assert_equal(c, truth)
def test_build():
numpy.random.seed(1000)
pos = numpy.random.uniform(size=(1000, 3))
tree = KDTree(pos).root
N = [0, 0]
def process1(r, i, j):
N[0] += len(r)
def process2(r, i, j):
N[1] += len(r)
tree.enum(tree, rmax=0.1, process=process1)
def test_enum():
numpy.random.seed(1234)
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f8')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f8')
dist = ((pos1[:, None, :] - pos2[None, :, :]) ** 2).sum(axis=-1)
truth = (dist < 1).sum()
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
N = [0]
def process1(r, i, j):
N[0] += (r < 1.0).sum()
tree1.enum(tree2, rmax=1.0, process=process1)
assert_equal(N[0], truth)
def test_fof_ind(method):
numpy.random.seed(1000)
pos = numpy.arange(100000).reshape(-1, 1).astype('f4')
ind = numpy.arange(len(pos))[::-2]
tree = KDTree(pos, ind=ind).root
label = tree.fof(2.1, method=method)
assert len(label) == len(pos)
correct_label = numpy.arange(len(pos))
correct_label[::-2] = label[1]
assert_equal(label, correct_label)
label = tree.fof(0.8, method=method)
correct_label = numpy.arange(len(pos)) # one group per active particle.
assert_equal(label, correct_label)
def test_enum():
numpy.random.seed(1234)
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f8')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f8')
dist = ((pos1[:, None, :] - pos2[None, :, :])**2).sum(axis=-1)
truth = (dist < 1).sum()
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
N = [0]
def process1(r, i, j):
N[0] += (r < 1.0).sum()
tree1.enum(tree2, rmax=1.0, process=process1)
assert_equal(N[0], truth)
def Force(self, state, ai, ac, af):
nbar = 1.0 * state.csize / self.pm.Nmesh.prod()
support = max([self.solver.r_cut, self.pm.resampler.support * 0.5])
layout, X1, rho = self.prepare_force(state, smoothing=support)
state.RHO[...] = layout.gather(rho.readout(X1))
delta_k = rho.r2c(out=Ellipsis)
state.F[...] = layout.gather(
self.solver.compute_longrange(X1,
delta_k,
factor=1.5 * self.cosmology.Om0))
tree = KDTree(X1, boxsize=self.pm.BoxSize)
Fs = layout.gather(self.solver.compute_shortrange(tree,
tree,
factor=state.GM0 /
state.H0**2),
mode='local')
state.F[...] += Fs
state.a['F'] = af
def get_tree(self, state, bin):
if self.Trees[bin] is None:
#print("Tree bin ==", bin)
self.Trees[bin] = KDTree(state.X,
ind=self.select(bin),
boxsize=state.pm.BoxSize)
return self.Trees[bin]
def match(coord, center, R):
"""
Returns a veto mask for coord. any coordinate within R of center
is vet.
Parameters
----------
coord : (RA, DEC)
center : (RA, DEC)
R : degrees
Returns
-------
Vetomask : True for veto, False for keep.
"""
from kdcount import KDTree
pos = radec2pos(center[0], center[1])
tree = KDTree(pos)
if numpy.isscalar(R):
#print('This is the value of R =%g'%(R))
R = center[0] * 0 + R
R = 2 * numpy.sin(numpy.radians(R) * 0.5)
pos = radec2pos(coord[0], coord[1])
other = KDTree(pos)
#vetoflag = numpy.zeros(len(pos), dtype='?')
test_j = []
test_i = []
Rmax = R.max()
def process(r, i, j):
# i is tycho, j is objects
rcut = R[i]
jcut = (j[r < rcut], i)
j1 = jcut[0]
j2 = jcut[1]
test_j.append(j1)
test_i.append(j2)
tree.root.enum(other.root, Rmax, process)
return numpy.array(test_j[0]), numpy.array(test_i[0])
def test_empty():
pos = numpy.arange(1000).astype('f4').reshape(-1, 1)
shape = ()
data = numpy.ones((len(pos)), dtype=('f4', shape))
tree = KDTree(pos, ind=[])
attr = KDAttr(tree, data)
assert_equal(attr[tree.root], 0)
def test_dtype():
numpy.random.seed(1000)
pos = numpy.arange(10).reshape(-1, 1)
try:
tree = KDTree(pos).root
assert True
except TypeError:
pass
def test_integrate2d():
N = 1000
pos = numpy.arange(N).astype('f4').reshape(-1, 1)
pos = numpy.concatenate([pos, pos], axis=-1)
tree = KDTree(pos)
root = tree.root
assert_equal(root.integrate(-numpy.inf, numpy.inf), N)
assert_equal(root.integrate(0, pos), numpy.arange(N))
def test_integrate1d():
pos = numpy.arange(1000).astype('f4').reshape(-1, 1)
tree = KDTree(pos)
root = tree.root
assert_equal(root.integrate(-numpy.inf, numpy.inf), 1000)
assert_equal(root.integrate(0, 1), 1)
assert_equal(root.integrate(0, 0), 0)
assert_equal(root.integrate(999, 999), 0)
assert_equal(root.integrate(0, pos), pos[:, 0])
def test_enum_count_weighted():
numpy.random.seed(1234)
pos1 = numpy.random.uniform(size=(1000, 3)).astype('f4')
pos2 = numpy.random.uniform(size=(1000, 3)).astype('f4')
w1 = numpy.ones(len(pos1))
w2 = numpy.ones(len(pos2))
tree1 = KDTree(pos1)
tree2 = KDTree(pos2)
a1 = KDAttr(tree1, w1)
a2 = KDAttr(tree2, w2)
N = [0]
def process1(r, i, j):
N[0] += len(r)
tree1.root.enum(tree2.root, rmax=1.0, process=process1)
c, w = tree1.root.count(tree2.root, r=1.0, attrs=(a1, a2))
assert_equal(N[0], c)
assert_equal(N[0], w)
def veto(coord, center, R):
"""
Returns a veto mask for coord. any coordinate within R of center
is vet.
Parameters
----------
coord : (RA, DEC)
center : (RA, DEC)
R : degrees
Returns
-------
Vetomask : True for veto, False for keep.
"""
from kdcount import KDTree
pos = radec2pos(center[0], center[1])
tree = KDTree(pos)
if numpy.isscalar(R):
#print('This is the value of R =%g'%(R))
R = center[0] * 0 + R
R = 2 * numpy.sin(numpy.radians(R) * 0.5)
pos = radec2pos(coord[0], coord[1])
other = KDTree(pos)
vetoflag = numpy.zeros(len(pos), dtype='?')
Rmax = R.max()
def process(r, i, j):
# i is tycho, j is objects
rcut = R[i]
jcut = j[(r < rcut) & (r != 0)]
vetoflag[jcut] |= True
tree.root.enum(other.root, Rmax, process)
return vetoflag
def test_count_symmetric():
pos1 = numpy.random.uniform(size=(1000000, 3)).astype('f4')
pos2 = numpy.array([[0.3, 0.5, 0.1]], dtype='f4')
tree1 = KDTree(pos1).root
tree2 = KDTree(pos2).root
assert_equal(tree2.count(tree1, (0, 0.1, 1.0)),
tree1.count(tree2, (0, 0.1, 1.0)))
def CompletenessEstimator(fluxes, noises, confidence):
"""
Create a completeness estimator for an object type,
based on the object type, intrinsic fluxes and intrinsic noise levels.
Parameters
----------
fluxes : array_like (N, Nbands)
intrinsic fluxes from DECAM; usually calculated by imglss-mpi-select-objects.py.
in nano-maggies.
noises : array_like (N, Nbands)
intrinsic 1-sigma noise level from DECAM; usually calculated by imglss-mpi-select-objects.py,
in nano-maggies.
confidence : array_like (Nbands)
confidence (in sigma) for each band.
"""
# we use the integrator in kdcount for the estimator
from kdcount import KDTree
# the targetselection type knows which bands are useful
# we only put limits on those bands.
#
# FIXME: @ekitanidis what about adding a link to your talk slides
# explaining this?
# This will be the 100% completeness limit for the given confidence
lim = fluxes.min(axis=0)
noises = confidence[None, :] * noises
mask = (noises < lim).all(axis=-1)
model = fluxes[mask]
tree = KDTree(model)
root = tree.root
def fcmodelfunc(query_noises):
query_noises = confidence[None, :] * query_noises
seen = root.integrate(query_noises, np.inf)
mask = (query_noises < lim).all(axis=-1)
# Watchout:
# Only 100% complete area has fcomp == 1.0
# otherwise we give a completeness slightly less than 1.0
fcomp = 1.0 * seen / (len(model) + 1.0)
fcomp[mask] = 1.0
return fcomp
return fcmodelfunc
def test_constattr():
pos = numpy.arange(100).astype('f4').reshape(-1, 1)
shapes = [(), (1, ), (2, )]
for shape in shapes:
data = constant_array((len(pos)), dtype=('f4', shape))
data.value[...] = 1.0
tree = KDTree(pos)
attr = KDAttr(tree, data)
assert_equal(tree.root.index, 0)
assert_equal(tree.root.less.index, 1)
assert_equal(tree.root.greater.index, 2)
assert_equal(attr.buffer.shape[0], tree.size)
assert_equal(attr.buffer.shape[1:], shape)
assert_equal(attr[tree.root], data.sum(axis=0))
def test_force():
from kdcount import force_kernels
pos = numpy.array([[0., 0., 0.], [1., 1., 1.], [2., 2., 2.], [3., 3., 3.]])
#pos = numpy.arange(0, 4, dtype='f8').reshape(-1, 1)
tree = KDTree(pos, thresh=1)
mass = KDAttr(tree, numpy.ones(pos.shape[0]))
xmass = KDAttr(tree, pos * mass.input)
force = tree.root.force(force_kernels['plummer'](1),
pos,
mass,
xmass,
2.,
eta=0.1)
# FIXME: add more tests
print(force)
def test_force_slow():
from kdcount import force_kernels
numpy.random.seed(13)
pos = numpy.random.uniform(size=(32 * 32 * 32, 3))
tree = KDTree(pos, thresh=1)
mass = KDAttr(tree, numpy.ones(pos.shape[0]))
xmass = KDAttr(tree, pos * mass.input)
#force = tree.root.force(force_kernels['plummer'](1), pos, mass, xmass, 1.0 / len(pos) ** 0.3333 * 4, eta=0.1)
force = tree.root.force(force_kernels['count'],
pos,
mass,
xmass,
1.0 / len(pos)**0.3333 * 4,
eta=0.1)
# FIXME: add more tests
print(force)
def veto_ellip(coord, center, a, b, l):
"""
Returns a veto mask for coord. any coordinate within R of center
is vet.
Parameters
----------
coord : (RA, DEC)
center : (RA, DEC)
R : degrees
Returns
-------
Vetomask : True for veto, False for keep.
"""
from kdcount import KDTree
pos = radec2pos(center[0], center[1])
tree = KDTree(pos)
#R = 2 * np.sin(np.radians(a) * 0.5)
R = a
pos = radec2pos(coord[0], coord[1])
other = KDTree(pos)
vetoflag = numpy.zeros(len(pos), dtype='?')
#j_near = numpy.zeros(len(pos), dtype=numpy.float64)
Rmax = R.max()
#print (Rmax)
x = numpy.array(coord[0])
y = numpy.array(coord[1])
h = numpy.array(center[0])
k = numpy.array(center[1])
a1 = numpy.array(a)
b1 = numpy.array(b)
l1 = numpy.array(l)
#j2 = numpy.array()
def process(r, i, j):
# i is 2mass, j is objects
#rcut = R[i]
X = x[j]
Y = y[j]
H = h[i]
K = k[i]
A = a1[i]
B = b1[i]
L = numpy.radians(l1[i])
#JEXT = j2[i]
c1 = (1 / A**2) * (numpy.cos(L))**2 + (1 / B**2) * (numpy.sin(L))**2
c2 = 2 * numpy.cos(L) * numpy.sin(L) * (1 / A**2 - 1 / B**2)
c3 = (1 / A**2) * (numpy.sin(L))**2 + (1 / B**2) * (numpy.cos(L))**2
rell = c1 * (X - H)**2 + c2 * (X - H) * (Y - K) + c3 * (Y - K)**2
jcut = j[rell < 1]
vetoflag[jcut] |= True
#j_near[jcut] = JEXT
tree.root.enum(other.root, Rmax, process)
return vetoflag
def test_ind():
numpy.random.seed(1000)
pos = numpy.arange(10).reshape(-1, 1).astype('f4')
tree = KDTree(pos, ind=[0, 1, 2, 3])
assert tree.root.size == 4
