def test_point_correctness():
import itertools
stencil = [-1, 0, 1]
ndim = 3
n = 2000
stencil = itertools.product(*[stencil]*ndim)
stencil = np.array(list(stencil)).astype(np.int32)
points = (np.random.rand(n, ndim) * [1, 2, 3]).astype(np.float32)
scale = 0.1
spec = GridSpec(points, float(scale))
offsets = spec.stencil(stencil).astype(np.int32)
grid = PointGrid(spec, points, offsets)
pairs = grid.pairs()
from scipy.spatial import cKDTree
tree = cKDTree(points)
tree_pairs = tree.query_pairs(scale, output_type='ndarray')
print(tree_pairs)
print(pairs)
assert np.alltrue(npi.unique(tree_pairs) == npi.unique(np.sort(pairs, axis=1)))
python类cKDTree()的实例源码
def hp_kd_tree(nside=set_default_nside(), leafsize=100):
"""
Generate a KD-tree of healpixel locations
Parameters
----------
nside : int
A valid healpix nside
leafsize : int (100)
Leafsize of the kdtree
Returns
-------
tree : scipy kdtree
"""
hpid = np.arange(hp.nside2npix(nside))
ra, dec = _hpid2RaDec(nside, hpid)
x, y, z = treexyz(ra, dec)
tree = kdtree(list(zip(x, y, z)), leafsize=leafsize, balanced_tree=False, compact_nodes=False)
return tree
def _set_altaz_fields(self, leafsize=10):
"""
Have a fixed grid of alt,az pointings to use
"""
tmp = read_fields()
names = ['alt', 'az']
types = [float, float]
self.fields = np.zeros(tmp.size, dtype=list(zip(names, types)))
self.fields['alt'] = tmp['dec']
self.fields['az'] = tmp['RA']
x, y, z = treexyz(self.fields['az'], self.fields['alt'])
self.field_tree = kdtree(list(zip(x, y, z)), leafsize=leafsize,
balanced_tree=False, compact_nodes=False)
hpids = np.arange(hp.nside2npix(self.nside))
self.reward_ra, self.reward_dec = _hpid2RaDec(self.nside, hpids)
def _spin_fields(self, lon=None, lat=None):
"""Spin the field tesselation
"""
if lon is None:
lon = np.random.rand()*np.pi*2
if lat is None:
lat = np.random.rand()*np.pi*2
# rotate longitude
ra = (self.fields['RA'] + lon) % (2.*np.pi)
dec = self.fields['dec'] + 0
# Now to rotate ra and dec about the x-axis
x, y, z = thetaphi2xyz(self.fields['RA'], self.fields['dec']+np.pi/2.)
xp, yp, zp = rotx(lat, x, y, z)
theta, phi = xyz2thetaphi(xp, yp, zp)
dec = phi - np.pi/2
ra = theta + np.pi
self.fields['RA'] = ra
self.fields['dec'] = dec
# Rebuild the kdtree with the new positions
# XXX-may be doing some ra,dec to conversions xyz more than needed.
self._hp2fieldsetup(ra, dec)
def fit(self, X, w):
if len(X) == 0:
raise NotEnoughParticles("Fitting not possible.")
self.X_arr = X.as_matrix()
ctree = cKDTree(X)
_, indices = ctree.query(X, k=min(self.k + 1, X.shape[0]))
covs, inv_covs, dets = list(zip(*[self._cov_and_inv(n, indices)
for n in range(X.shape[0])]))
self.covs = sp.array(covs)
self.inv_covs = sp.array(inv_covs)
self.determinants = sp.array(dets)
self.normalization = sp.sqrt(
(2 * sp.pi) ** self.X_arr.shape[1] * self.determinants)
if not sp.isreal(self.normalization).all():
raise Exception("Normalization not real")
self.normalization = sp.real(self.normalization)
def __init__(self, init_num, size, stp):
self.num = init_num
self.size = size
self.stp = stp
self.one = 1.0/size
self.rad = 0.04
edge = 0.2
self.xy = edge + random((init_num, 2))*(1.0-2*edge)
self.v = zeros((init_num, 2), 'float')
self.dx = zeros((init_num, 2), 'float')
self.random_velocity(self.stp)
self.tree = kdtree(self.xy)
self.slowdown = 0.99999
def get_cKDTree():
try:
from scipy.spatial import cKDTree
except ImportError:
cKDTree = None
return cKDTree
# getheader gives a 87a header and a color palette (two elements in a list).
# getdata()[0] gives the Image Descriptor up to (including) "LZW min code size".
# getdatas()[1:] is the image data itself in chuncks of 256 bytes (well
# technically the first byte says how many bytes follow, after which that
# amount (max 255) follows).
def quantize_with_scipy(self, image):
w,h = image.size
px = np.asarray(image).copy()
px2 = px[:,:,:3].reshape((w*h,3))
cKDTree = get_cKDTree()
kdtree = cKDTree(self.colormap[:,:3],leafsize=10)
result = kdtree.query(px2)
colorindex = result[1]
print("Distance: %1.2f" % (result[0].sum()/(w*h)) )
px2[:] = self.colormap[colorindex,:3]
return Image.fromarray(px).convert("RGB").quantize(palette=self.paletteImage())
def compute_index(codebook):
return cKDTree(codebook)
def get_cKDTree():
try:
from scipy.spatial import cKDTree
except ImportError:
cKDTree = None
return cKDTree
# getheader gives a 87a header and a color palette (two elements in a list).
# getdata()[0] gives the Image Descriptor up to (including) "LZW min code size".
# getdatas()[1:] is the image data itself in chuncks of 256 bytes (well
# technically the first byte says how many bytes follow, after which that
# amount (max 255) follows).
def quantize_with_scipy(self, image):
w, h = image.size
px = np.asarray(image).copy()
px2 = px[:, :, :3].reshape((w * h, 3))
cKDTree = get_cKDTree()
kdtree = cKDTree(self.colormap[:, :3], leafsize=10)
result = kdtree.query(px2)
colorindex = result[1]
print("Distance: %1.2f" % (result[0].sum() / (w * h)))
px2[:] = self.colormap[colorindex, :3]
return Image.fromarray(px).convert("RGB").quantize(palette=self.paletteImage())
def invalidate_distances(self):
self._tree = cKDTree(self.params.X)
self._distances, self._neighbors = None, None
def where_close(pos, separation, intensity=None):
""" Returns indices of features that are closer than separation from other
features. When intensity is given, the one with the lowest intensity is
returned: else the most topleft is returned (to avoid randomness)
To be implemented in trackpy v0.4"""
if len(pos) == 0:
return []
separation = validate_tuple(separation, pos.shape[1])
if any([s == 0 for s in separation]):
return []
# Rescale positions, so that pairs are identified below a distance
# of 1.
pos_rescaled = pos / separation
duplicates = cKDTree(pos_rescaled, 30).query_pairs(1 - 1e-7)
if len(duplicates) == 0:
return []
index_0 = np.fromiter((x[0] for x in duplicates), dtype=int)
index_1 = np.fromiter((x[1] for x in duplicates), dtype=int)
if intensity is None:
to_drop = np.where(np.sum(pos_rescaled[index_0], 1) >
np.sum(pos_rescaled[index_1], 1),
index_1, index_0)
else:
intensity_0 = intensity[index_0]
intensity_1 = intensity[index_1]
to_drop = np.where(intensity_0 > intensity_1, index_1, index_0)
edge_cases = intensity_0 == intensity_1
if np.any(edge_cases):
index_0 = index_0[edge_cases]
index_1 = index_1[edge_cases]
to_drop[edge_cases] = np.where(np.sum(pos_rescaled[index_0], 1) >
np.sum(pos_rescaled[index_1], 1),
index_1, index_0)
return np.unique(to_drop)
def sort_positions(actual, expected):
assert_equal(len(actual), len(expected))
tree = cKDTree(actual)
devs, argsort = tree.query([expected])
return devs, actual[argsort][0]
def __init__(self, X=None, z=None, leafsize=10):
if not X is None:
self.tree = cKDTree(X, leafsize=leafsize )
if not z is None:
self.z = z
def _make_kdtree(self, x):
self.kdtree = cKDTree(mpiops.random_full_points(x, Napprox=self.nodes))
if not np.isfinite(self.kdtree.query(x, k=self.k)[0]).all():
log.warning('Kdtree computation encountered problem. '
'Not enough neighbors available to compute '
'kdtree. Printing kdtree for debugging purpose')
raise ValueError('Computed kdtree is not fully populated.'
'Not enough valid neighbours available.')
def __init__(self, train, weights=None, truncation=3.0, nmin=4, factor=1.0):
"""
Parameters
----------
train : np.ndarray
The training data set. Should be a 1D array of samples or a 2D array of shape (n_samples, n_dim).
weights : np.ndarray, optional
An array of weights. If not specified, equal weights are assumed.
truncation : float, optional
The maximum deviation (in sigma) to use points in the KDE
nmin : int, optional
The minimum number of points required to estimate the density
factor : float, optional
Send bandwidth to this factor of the data estimate
"""
self.truncation = truncation
self.nmin = nmin
self.train = train
if len(train.shape) == 1:
train = np.atleast_2d(train).T
self.num_points, self.num_dim = train.shape
if weights is None:
weights = np.ones(self.num_points)
self.weights = weights
self.mean = np.average(train, weights=weights, axis=0)
dx = train - self.mean
cov = np.atleast_2d(np.cov(dx.T, aweights=weights))
self.A = np.linalg.cholesky(np.linalg.inv(cov)) # The sphere-ifying transform
self.d = np.dot(dx, self.A) # Sphere-ified data
self.tree = spatial.cKDTree(self.d) # kD tree of data
self.sigma = 2.0 * factor * np.power(self.num_points, -1. / (4 + self.num_dim)) # Starting sigma (bw) of Gauss
self.sigma_fact = -0.5 / (self.sigma * self.sigma)
# Cant get normed probs to work atm, turning off for now as I don't need normed pdfs for contours
# self.norm = np.product(np.diagonal(self.A)) * (2 * np.pi) ** (-0.5 * self.num_dim) # prob norm
# self.scaling = np.power(self.norm * self.sigma, -self.num_dim)
def entropy(x, k=3, base=2):
""" The classic K-L k-nearest neighbor continuous entropy estimator
x should be a list of vectors, e.g. x = [[1.3], [3.7], [5.1], [2.4]]
if x is a one-dimensional scalar and we have four samples
"""
assert k <= len(x) - 1, "Set k smaller than num. samples - 1"
d = len(x[0])
N = len(x)
intens = 1e-10 # small noise to break degeneracy, see doc.
x = [list(p + intens * nr.rand(len(x[0]))) for p in x]
tree = ss.cKDTree(x)
nn = [tree.query(point, k + 1, p=float('inf'))[0][k] for point in x]
const = digamma(N) - digamma(k) + d * log(2)
return (const + d * np.mean(map(log, nn))) / log(base)
def avgdigamma(points, dvec):
# This part finds number of neighbors in some radius in the marginal space
# returns expectation value of <psi(nx)>
N = len(points)
tree = ss.cKDTree(points)
avg = 0.
for i in range(N):
dist = dvec[i]
# subtlety, we don't include the boundary point,
# but we are implicitly adding 1 to kraskov def bc center point is included
num_points = len(tree.query_ball_point(points[i], dist - 1e-15, p=float('inf')))
avg += digamma(num_points) / N
return avg
def maskMeshByDistance(x,y,d,grid):
"""Filters all (x,y) coordinates that are more than d
in meshgrid given some actual coordinates (x,y).
Args:
x (numpy.ndarray): x-coordinates.
y (numpy.ndarray): y-coordinates.
d (float): Maximum distance.
grid (numpy.ndarray): Numpy meshgrid.
Returns:
idxs (list): List of booleans.
"""
#Convert x/y into useful array
xy=np.vstack((x,y))
#Compute distances to nearest neighbors
tree = cKDTree(xy.T)
xi = _ndim_coords_from_arrays(tuple(grid))
dists, indexes = tree.query(xi)
#Get indices of distances that are further apart than d
idxs = (dists > d)
return idxs,
