def test_reverse_query_ball(self):
radius = 0.1
r = np.random.RandomState(123)
in_coords = r.uniform(size=(na, 3))
out_coords = r.uniform(size=(nb, 3))
in_tree = tree_utils.cKDTree(in_coords)
out_tree = tree_utils.cKDTree(out_coords)
arr = tree_utils.query_ball_tree(in_tree, out_tree, radius)
rel_coords = np.repeat(out_coords, arr.row_lengths, axis=0) - \
in_coords[arr.flat_values]
rel_dists = np.linalg.norm(rel_coords, axis=-1)
rev_arr, rev_indices = tree_utils.reverse_query_ball(arr, na)
rel_coords_inv = rel_coords[rev_indices]
arr_rev, rel_dists_inv = tree_utils.reverse_query_ball(
arr, na, rel_dists)
np.testing.assert_allclose(np.linalg.norm(rel_coords_inv, axis=-1),
rel_dists_inv)
naive_arr_rev = tree_utils.query_ball_tree(out_tree, in_tree, radius)
np.testing.assert_equal(naive_arr_rev.flat_values, rev_arr.flat_values)
np.testing.assert_equal(naive_arr_rev.row_splits, rev_arr.row_splits)
naive_rel_coords_inv = (
np.repeat(in_coords, naive_arr_rev.row_lengths, axis=0) -
out_coords[naive_arr_rev.flat_values])
naive_rel_dists_inv = np.linalg.norm(naive_rel_coords_inv, axis=-1)
np.testing.assert_allclose(rel_coords_inv, -naive_rel_coords_inv)
np.testing.assert_allclose(rel_dists_inv, naive_rel_dists_inv)
def test_query_ball_tree(self):
radius = 1
pa = [
[0, 0],
[10, 0],
[0, 0.5],
]
pb = [[0, 0.1], [-0.9, 0]]
in_tree = tree_utils.cKDTree(pa)
out_tree = tree_utils.cKDTree(pb)
arr = tree_utils.query_ball_tree(in_tree, out_tree, radius)
np.testing.assert_equal(arr.flat_values, [0, 2, 0])
np.testing.assert_equal(arr.row_lengths, [2, 1])
def test_query_knn(self):
coords = np.random.uniform(size=(100, 3))
k = 8
tree = tree_utils.cKDTree(coords)
expected = tree_utils.query_knn(tree, k)
actual = self.evaluate(q.query_knn(coords, k))
self.assertAllEqual(actual[0], expected[0])
self.assertAllClose(actual[1], expected[1])
def test_query_pairs(self):
coords = np.random.uniform(size=(100, 3))
radius = 0.1
tree = tree_utils.cKDTree(coords)
expected = tree_utils.query_pairs(tree, radius)
actual = self.evaluate(q.query_pairs(coords, radius))
self.assertAllEqual(actual.values, expected.values)
self.assertAllEqual(actual.row_splits, expected.row_splits)
def test_query_gather(self):
# logic test: should be the same
# query_pairs(in_tree, radius)[mask]
# query_ball_tree(in_tree, cKDTree(in_tree.data[mask]), radius)
radius = 0.1
r = np.random.RandomState(123)
in_coords = r.uniform(size=(na, 3))
in_tree = tree_utils.cKDTree(in_coords)
neighbors = tree_utils.query_pairs(in_tree, radius)
indices = sample.inverse_density_sample(
in_coords.shape[0] // 2, neighborhood_size=neighbors.row_lengths)
out_coords = in_coords[indices]
out_tree = tree_utils.cKDTree(out_coords)
out_neigh = tree_utils.query_ball_tree(in_tree, out_tree, radius)
out_neigh2 = neighbors.gather(indices)
np.testing.assert_equal(out_neigh.values, out_neigh2.values)
np.testing.assert_equal(out_neigh.row_splits, out_neigh2.row_splits)
def test_query_mask(self):
# logic test: should be the same
# query_pairs(in_tree, radius)[mask]
# query_ball_tree(in_tree, cKDTree(in_tree.data[mask]), radius)
radius = 0.1
r = np.random.RandomState(123)
in_coords = r.uniform(size=(na, 3))
in_tree = tree_utils.cKDTree(in_coords)
neighbors = tree_utils.query_pairs(in_tree, radius)
mask = sample.inverse_density_mask(
neighborhood_size=neighbors.row_lengths, mean_keep_rate=0.5)
out_coords = in_coords[mask]
out_tree = tree_utils.cKDTree(out_coords)
out_neigh = tree_utils.query_ball_tree(in_tree, out_tree, radius)
out_neigh2 = neighbors.mask(mask)
np.testing.assert_equal(out_neigh.values, out_neigh2.values)
np.testing.assert_equal(out_neigh.row_splits, out_neigh2.row_splits)
def test_query_pairs(self):
radius = 0.1
r = np.random.RandomState(123)
in_coords = r.uniform(size=(na, 3))
tree = tree_utils.cKDTree(in_coords)
base = ragged_array.RaggedArray.from_ragged_lists(tree.query_ball_tree(
tree, radius),
dtype=np.int64)
efficient = tree_utils.query_pairs(tree, radius)
np.testing.assert_equal(base.flat_values, efficient.flat_values)
np.testing.assert_equal(base.row_splits, efficient.row_splits)
clamped = tree_utils.query_pairs(tree, radius, max_neighbors=5)
self.assertLessEqual(np.max(clamped.row_lengths), 5)
def load(self, fp, path=None):
import trimesh
from scipy.sparse.linalg import svds
vertices, faces = mn.load_off_mesh(fp)
# quick recentering to avoid numerical instabilities in sampling
center = (np.max(vertices, axis=0) + np.min(vertices, axis=0)) / 2
vertices -= center
scale = np.max(vertices)
vertices /= scale
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
positions, face_indices = trimesh.sample.sample_surface(
mesh, self._num_points_start)
normals = mesh.face_normals[face_indices]
tree = tree_utils.cKDTree(positions)
indices, neighborhood_size = tree_utils.query_ball_tree(
tree, tree, self._radius)
mask = sample.inverse_density_mask(
neighborhood_size,
prob_scale_factor=self._prob_scale_factor,
indices=indices,
coords=positions,
scale_factor=self._sample_scale_factor,
weight_fn=self._weight_fn)
positions = positions[mask]
if positions.shape[0] < 2:
logging.info('Skipping %s' % path)
self._bad_paths.append(path)
return None
normals = normals[mask]
r, c = mn.naive_bounding_sphere(positions)
positions -= c
positions /= r
center += c * scale
scale *= r
_, _, v = svds(positions, k=1)
v = v[0]
theta = np.arctan2(v[1], v[0])
positions, normals = augment.rotate_np(positions, normals, -theta)
return dict(
positions=positions.astype(np.float32),
normals=normals.astype(np.float32),
original_center=center,
original_radius=scale,
)
def fn(coords, k_tensor=None):
kv = k if k_tensor is None else k_tensor.numpy()
tree = tree_utils.cKDTree(coords.numpy())
values, dists = tree_utils.query_knn(tree, kv)
return values, dists
def fn(coords, radius_tensor=None):
r = radius if radius_tensor is None else radius_tensor.numpy()
tree = tree_utils.cKDTree(coords.numpy())
pairs = tree_utils.query_pairs(tree, r, max_neighbors)
return pairs.values, pairs.row_splits
