def test_invert_neighbors_list_shape_checking(ml):
num_points = 3
inp_neighbors_index = np.array([0, 1, 2, 2, 1, 2], dtype=np.int32)
inp_neighbors_row_splits = np.array([0, 3, 4, 6], dtype=np.int64)
inp_neighbors_attributes = np.array([10, 20, 30, 40, 50, 60],
dtype=np.float32)
# test the shape checking by passing arrays with wrong rank and/or size
with pytest.raises(Exception) as einfo:
_ = mltest.run_op(ml,
ml.cpu_device,
False,
ml.ops.invert_neighbors_list,
num_points=num_points,
inp_neighbors_index=inp_neighbors_index[1:],
inp_neighbors_row_splits=inp_neighbors_row_splits,
inp_neighbors_attributes=inp_neighbors_attributes)
assert 'invalid shape' in str(einfo.value)
with pytest.raises(Exception) as einfo:
_ = mltest.run_op(ml,
ml.cpu_device,
False,
ml.ops.invert_neighbors_list,
num_points=num_points,
inp_neighbors_index=inp_neighbors_index[:,
np.newaxis],
inp_neighbors_row_splits=inp_neighbors_row_splits,
inp_neighbors_attributes=inp_neighbors_attributes)
assert 'invalid shape' in str(einfo.value)
with pytest.raises(Exception) as einfo:
_ = mltest.run_op(
ml,
ml.cpu_device,
False,
ml.ops.invert_neighbors_list,
num_points=num_points,
inp_neighbors_index=inp_neighbors_index,
inp_neighbors_row_splits=inp_neighbors_row_splits[:, np.newaxis],
inp_neighbors_attributes=inp_neighbors_attributes)
assert 'invalid shape' in str(einfo.value)
with pytest.raises(Exception) as einfo:
_ = mltest.run_op(
ml,
ml.cpu_device,
False,
ml.ops.invert_neighbors_list,
num_points=num_points,
inp_neighbors_index=inp_neighbors_index,
inp_neighbors_row_splits=inp_neighbors_row_splits,
inp_neighbors_attributes=inp_neighbors_attributes[1:])
assert 'invalid shape' in str(einfo.value)
def test_roi_pool(ml):
values0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/values0.npy'
)
values1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/values1.npy'
)
values2 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/values2.npy'
)
sampled_pts_num = 512
ans0, ans1 = mltest.run_op(ml, ml.device, True, ml.ops.roi_pool, values0,
values1, values2, sampled_pts_num)
expected0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/out0.npy'
)
expected1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/out1.npy'
)
np.testing.assert_equal(ans0, expected0)
np.testing.assert_equal(ans1, expected1)
def test_reduce_subarrays_sum_random(seed, dtype, device, ml):
rng = np.random.RandomState(seed)
values_shape = [rng.randint(100, 200)]
values = rng.uniform(0, 10, size=values_shape).astype(dtype)
row_splits = [0]
for _ in range(rng.randint(1, 10)):
row_splits.append(
rng.randint(0, values_shape[0] - row_splits[-1]) + row_splits[-1])
row_splits.extend(values_shape)
expected_result = []
for start, stop in zip(row_splits, row_splits[1:]):
# np.sum correctly handles zero length arrays and returns 0
expected_result.append(np.sum(values[start:stop]))
np.array(expected_result, dtype=dtype)
row_splits = np.array(row_splits, dtype=np.int64)
ans = mltest.run_op(ml,
device,
True,
ml.ops.reduce_subarrays_sum,
values=values,
row_splits=row_splits)
if np.issubdtype(dtype, np.integer):
np.testing.assert_equal(ans, expected_result)
else: # floating point types
np.testing.assert_allclose(ans, expected_result, rtol=1e-5, atol=1e-8)
def test_voxelize_random(ml, point_dtype, ndim, max_voxels,
max_points_per_voxel):
rng = np.random.RandomState(123)
points = rng.rand(rng.randint(0, 10000), ndim).astype(point_dtype)
voxel_size = rng.uniform(0.01, 0.1, size=(ndim, )).astype(point_dtype)
point_range_min = rng.uniform(0.0, 0.3, size=(ndim, )).astype(point_dtype)
point_range_max = rng.uniform(0.7, 1.0, size=(ndim, )).astype(point_dtype)
ans = mltest.run_op(ml,
ml.device,
True,
ml.ops.voxelize,
points,
voxel_size,
point_range_min,
point_range_max,
max_points_per_voxel=max_points_per_voxel,
max_voxels=max_voxels)
voxels = convert_output_to_voxel_dict(ans)
voxels_reference = voxelize_python(
points,
voxel_size,
point_range_min,
point_range_max,
max_points_per_voxel=max_points_per_voxel,
max_voxels=max_voxels)
assert_equal_voxel_dicts(voxels, ref=voxels_reference)
def sparse_conv_transpose_filter(filters):
return mltest.run_op(ml, ml.device, True, ml.ops.sparse_conv_transpose,
filters, inp_importance, y_arr, neighbors_index,
neighbors_importance_sum, neighbors_row_splits,
inv_neighbors_index, inv_neighbors_kernel_index,
inv_neighbors_importance,
inv_neighbors_row_splits, **conv_attrs)
def test_fixed_radius_search_empty_point_sets(ml):
rng = np.random.RandomState(123)
dtype = np.float32
radius = 1
hash_table_size_factor = 1 / 64
# no query points
points = rng.random(size=(100, 3)).astype(dtype)
queries = rng.random(size=(0, 3)).astype(dtype)
layer = ml.layers.FixedRadiusSearch(return_distances=True)
ans = mltest.run_op(
ml,
ml.device,
True,
layer,
points,
queries=queries,
radius=radius,
hash_table_size_factor=hash_table_size_factor,
)
assert ans.neighbors_index.shape == (0, )
assert ans.neighbors_row_splits.shape == (1, )
assert ans.neighbors_distance.shape == (0, )
# no input points
points = rng.random(size=(0, 3)).astype(dtype)
queries = rng.random(size=(100, 3)).astype(dtype)
ans = mltest.run_op(
ml,
ml.device,
True,
layer,
points,
queries=queries,
radius=radius,
hash_table_size_factor=hash_table_size_factor,
)
assert ans.neighbors_index.shape == (0, )
assert ans.neighbors_row_splits.shape == (101, )
np.testing.assert_array_equal(np.zeros_like(ans.neighbors_row_splits),
ans.neighbors_row_splits)
assert ans.neighbors_distance.shape == (0, )
def sparse_conv_transpose_infeats(inp_features):
return mltest.run_op(ml, ml.device, True, ml.ops.sparse_conv_transpose,
filters.transpose([0, 2, 1]), inp_importance,
inp_features, neighbors_index,
neighbors_importance_sum, neighbors_row_splits,
inv_neighbors_index, inv_neighbors_kernel_index,
inv_neighbors_importance,
inv_neighbors_row_splits, **conv_attrs)
def test_knn_search_empty_point_sets(ml):
rng = np.random.RandomState(123)
device = mltest.cpu_device
dtype = np.float32
# no query points
points = rng.random(size=(100, 3)).astype(dtype)
queries = rng.random(size=(0, 3)).astype(dtype)
k = rng.randint(1, 11)
layer = ml.layers.KNNSearch(return_distances=True)
ans = mltest.run_op(
ml,
device,
layer,
True,
points,
queries=queries,
k=k,
)
assert ans.neighbors_index.shape == (0, )
assert ans.neighbors_row_splits.shape == (1, )
assert ans.neighbors_distance.shape == (0, )
# no input points
points = rng.random(size=(0, 3)).astype(dtype)
queries = rng.random(size=(100, 3)).astype(dtype)
layer = ml.layers.KNNSearch(return_distances=True)
ans = mltest.run_op(
ml,
device,
layer,
True,
points,
queries=queries,
k=k,
)
assert ans.neighbors_index.shape == (0, )
assert ans.neighbors_row_splits.shape == (101, )
np.testing.assert_array_equal(np.zeros_like(ans.neighbors_row_splits),
ans.neighbors_row_splits)
assert ans.neighbors_distance.shape == (0, )
def test_radius_search_empty_point_sets(ml):
rng = np.random.RandomState(123)
dtype = np.float32
# no query points
points = rng.random(size=(100, 3)).astype(dtype)
queries = rng.random(size=(0, 3)).astype(dtype)
radii = rng.uniform(0.1, 0.3, size=(0, )).astype(dtype)
layer = ml.layers.RadiusSearch(return_distances=True)
ans = mltest.run_op(
ml,
ml.device,
True,
layer,
points,
queries=queries,
radii=radii,
)
assert ans.neighbors_index.shape == (0, )
assert ans.neighbors_row_splits.shape == (1, )
assert ans.neighbors_distance.shape == (0, )
# no input points
points = rng.random(size=(0, 3)).astype(dtype)
queries = rng.random(size=(100, 3)).astype(dtype)
radii = rng.uniform(0.1, 0.3, size=(100, )).astype(dtype)
ans = mltest.run_op(
ml,
ml.device,
True,
layer,
points,
queries=queries,
radii=radii,
)
assert ans.neighbors_index.shape == (0, )
assert ans.neighbors_row_splits.shape == (101, )
np.testing.assert_array_equal(np.zeros_like(ans.neighbors_row_splits),
ans.neighbors_row_splits)
assert ans.neighbors_distance.shape == (0, )
def test_fixed_radius_search(dtype, device, ml, num_points_queries, radius,
hash_table_size_factor, metric, ignore_query_point,
return_distances):
# skip dtype not supported on GPU
if 'GPU' in device and not dtype in gpu_dtypes:
return
rng = np.random.RandomState(123)
num_points, num_queries = num_points_queries
points = rng.random(size=(num_points, 3)).astype(dtype)
if ignore_query_point:
queries = points
else:
queries = rng.random(size=(num_queries, 3)).astype(dtype)
# kd tree for computing the ground truth
tree = cKDTree(points, copy_data=True)
p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric]
gt_neighbors_index = tree.query_ball_point(queries, radius, p=p_norm)
layer = ml.layers.FixedRadiusSearch(metric=metric,
ignore_query_point=ignore_query_point,
return_distances=return_distances)
ans = mltest.run_op(
ml,
device,
True,
layer,
points,
queries=queries,
radius=radius,
hash_table_size_factor=hash_table_size_factor,
)
for i, q in enumerate(queries):
# check neighbors
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
q_neighbors_index = ans.neighbors_index[start:end]
gt_set = set(gt_neighbors_index[i])
if ignore_query_point:
gt_set.remove(i)
assert gt_set == set(q_neighbors_index)
# check distances
if return_distances:
q_neighbors_dist = ans.neighbors_distance[start:end]
for j, dist in zip(q_neighbors_index, q_neighbors_dist):
if metric == 'L2':
gt_dist = np.sum((q - points[j])**2)
else:
gt_dist = np.linalg.norm(q - points[j], ord=p_norm)
np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)
def test_voxelize_simple(ml, point_dtype, point_range_max, max_voxels,
max_points_per_voxel):
# global iteration
# if iteration > 0:
# return
# iteration += 1
# if 'GPU' in ml.device:
# return
# yapf: disable
points = np.array([
# 2 points in voxel
[0.5, 0.5, 0.5],
[0.7, 0.2, 0.3],
# new batch
# 1 point in two voxels
[0.7, 0.5, 0.9],
[10.7, 10.2, 10.3],
# 2 points in another voxel
[1.4, 1.5, 1.4],
[1.7, 1.2, 1.3],
], dtype=point_dtype)
row_splits = np.array([0, 2, 4, 6], dtype=np.int64)
# row_splits = np.array([0,6], dtype=np.int64)
# yapf: enable
voxel_size = np.array([1.0, 1.1, 1.2], dtype=point_dtype)
point_range_min = np.zeros((3, ), dtype=point_dtype)
point_range_max = np.array(point_range_max, dtype=point_dtype)
ans = mltest.run_op(ml,
ml.device,
True,
ml.ops.voxelize,
points,
row_splits,
voxel_size,
point_range_min,
point_range_max,
max_points_per_voxel=max_points_per_voxel,
max_voxels=max_voxels)
voxels = convert_output_to_voxel_dict(ans)
voxels_reference = voxelize_python(
points,
row_splits,
voxel_size,
point_range_min,
point_range_max,
max_points_per_voxel=max_points_per_voxel,
max_voxels=max_voxels)
# print(voxels)
# print(voxels_reference)
# print('max_voxels', max_voxels, 'max_points_per_voxel',
# max_points_per_voxel)
assert_equal_voxel_dicts(voxels, ref_b=voxels_reference)
def test_voxel_pooling_empty_point_set(ml, pos_dtype, feat_dtype, position_fn,
feature_fn):
points = np.zeros(shape=[0, 3], dtype=pos_dtype)
features = np.zeros(shape=[0, 5], dtype=feat_dtype)
voxel_size = 1
ans = mltest.run_op(ml, ml.device, True, ml.ops.voxel_pooling, points,
features, voxel_size, position_fn, feature_fn)
np.testing.assert_array_equal(points, ans.pooled_positions)
np.testing.assert_array_equal(features, ans.pooled_features)
def test_furthest_point_sampling(ml):
values = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/sampling/values.npy'
)
samples = 4096
ans = mltest.run_op(ml, ml.device, True, ml.ops.furthest_point_sampling,
values, samples)
expected = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/sampling/out.npy'
)
np.testing.assert_equal(ans, expected)
def test_nms_empty(ml):
boxes = np.zeros((0, 5), dtype=np.float32)
scores = np.array([], dtype=np.float32)
nms_overlap_thresh = 0.7
keep_indices_ref = np.array([]).astype(np.int64)
keep_indices = mltest.run_op(ml,
ml.device,
True,
ml.ops.nms,
boxes,
scores,
nms_overlap_thresh=nms_overlap_thresh)
np.testing.assert_equal(keep_indices, keep_indices_ref)
assert keep_indices.dtype == keep_indices_ref.dtype
def test_group_pts(ml):
values0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/values0.npy'
)
values1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/values1.npy'
)
ans = mltest.run_op(ml, ml.device, True, ml.ops.group_points, values0,
values1)
expected = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/out.npy'
)
np.testing.assert_equal(ans, expected)
def test_cublas_matmul(ml):
# This test checks if calling cublas functionality from open3d and the ml framework works.
rng = np.random.RandomState(123)
n = 20
arr = rng.rand(n, n).astype(np.float32)
# do matmul with open3d
A = o3d.core.Tensor.from_numpy(arr).cuda()
B = A @ A
# now use the ml framework cublas
C = mltest.run_op(ml, ml.device, True, ml.module.matmul, arr, arr)
np.testing.assert_allclose(B.cpu().numpy(), C)
def conv_filter(filters):
return mltest.run_op(ml,
ml.device,
True,
ml.ops.continuous_conv,
filters=filters,
out_positions=out_positions,
extents=extent,
offset=offset,
inp_positions=inp_positions,
inp_features=inp_features,
inp_importance=inp_importance,
neighbors_index=neighbors_index,
neighbors_importance=neighbors_importance,
neighbors_row_splits=neighbors_row_splits,
**conv_attrs)
def test_ragged_to_dense(dtype, ml):
values = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=dtype)
row_splits = np.array([0, 2, 4, 4, 5, 12, 13], dtype=np.int64)
out_col_size = 4
default_value = np.array(-1, dtype=dtype)
ans = mltest.run_op(ml, ml.device, True, ml.ops.ragged_to_dense, values,
row_splits, out_col_size, default_value)
expected = np.full((row_splits.shape[0] - 1, out_col_size), default_value)
for i in range(row_splits.shape[0] - 1):
for j, value_idx in enumerate(range(row_splits[i], row_splits[i + 1])):
if j < expected.shape[1]:
expected[i, j] = values[value_idx]
np.testing.assert_equal(ans, expected)
def test_query_pts(ml):
values0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/query_pts/values0.npy'
)
values1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/query_pts/values1.npy'
)
sample = 16
radius = 0.1
ans = mltest.run_op(ml, ml.device, True, ml.ops.ball_query, values0,
values1, radius, sample)
expected = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/query_pts/out.npy'
)
np.testing.assert_equal(ans, expected)
def test_three_interp(ml):
values0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/values0.npy'
)
values1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/values1.npy'
)
values2 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/values2.npy'
)
ans = mltest.run_op(ml, ml.device, True, ml.ops.three_interpolate, values0,
values1, values2)
expected = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/out.npy'
)
np.testing.assert_equal(ans, expected)
def test_three_nn(ml):
values0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/values0.npy'
)
values1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/values1.npy'
)
ans0, ans1 = mltest.run_op(ml, ml.device, True, ml.ops.three_nn, values0,
values1)
expected0 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/out0.npy'
)
expected1 = mltest.fetch_numpy(
'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/out1.npy'
)
np.testing.assert_equal(ans0, expected0)
np.testing.assert_equal(ans1, expected1)
def conv_transpose_infeats(inp_features):
return mltest.run_op(
ml,
ml.device,
True,
ml.ops.continuous_conv_transpose,
filters=filters.transpose([0, 1, 2, 4, 3]),
out_positions=inp_positions,
out_importance=inp_importance,
extents=extent,
offset=offset,
inp_positions=out_positions,
inp_features=inp_features,
inp_neighbors_index=neighbors_index,
inp_neighbors_importance_sum=neighbors_importance_sum,
inp_neighbors_row_splits=neighbors_row_splits,
neighbors_index=inverted_neighbors_index,
neighbors_importance=inverted_neighbors_importance,
neighbors_row_splits=inverted_neighbors_row_splits,
**conv_attrs)
def test_ragged_to_dense_more_dims(dtype, ml):
values = np.array([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6],
[7, 7], [8, 8], [9, 9], [10, 10], [11, 11], [12, 12]],
dtype=dtype)
row_splits = np.array([0, 2, 4, 4, 5, 12, 13], dtype=np.int64)
out_col_size = 4
default_value = np.array([-1, -1], dtype=dtype)
ans = mltest.run_op(ml, ml.device, True, ml.ops.ragged_to_dense, values,
row_splits, out_col_size, default_value)
expected = np.full((
row_splits.shape[0] - 1,
out_col_size,
) + default_value.shape, default_value)
for i in range(row_splits.shape[0] - 1):
for j, value_idx in enumerate(range(row_splits[i], row_splits[i + 1])):
if j < expected.shape[1]:
expected[i, j] = values[value_idx]
np.testing.assert_equal(ans, expected)
def test_nms(ml):
boxes = np.array([[15.0811, -7.9803, 15.6721, -6.8714, 0.5152],
[15.1166, -7.9261, 15.7060, -6.8137, 0.6501],
[15.1304, -7.8129, 15.7069, -6.8903, 0.7296],
[15.2050, -7.8447, 15.8311, -6.7437, 1.0506],
[15.1343, -7.8136, 15.7121, -6.8479, 1.0352],
[15.0931, -7.9552, 15.6675, -7.0056, 0.5979]],
dtype=np.float32)
scores = np.array([3, 1.1, 5, 2, 1, 0], dtype=np.float32)
nms_overlap_thresh = 0.7
keep_indices_ref = np.array([2, 3, 5]).astype(np.int64)
keep_indices = mltest.run_op(ml,
ml.device,
True,
ml.ops.nms,
boxes,
scores,
nms_overlap_thresh=nms_overlap_thresh)
np.testing.assert_equal(keep_indices, keep_indices_ref)
assert keep_indices.dtype == keep_indices_ref.dtype
def test_reduce_subarrays_sum_zero_length_values(ml):
rng = np.random.RandomState(1)
shape = [rng.randint(100, 200)]
values = np.array([], dtype=np.float32)
row_splits = [0]
for _ in range(rng.randint(1, 10)):
row_splits.append(
rng.randint(0, shape[0] - row_splits[-1]) + row_splits[-1])
row_splits.extend(shape)
row_splits = np.array(row_splits, dtype=np.int64)
ans = mltest.run_op(ml,
ml.device,
True,
ml.ops.reduce_subarrays_sum,
values=values,
row_splits=row_splits)
assert ans.shape == values.shape
assert ans.dtype == values.dtype
def test_voxelize_random(ml, point_dtype, ndim, batch_size, max_voxels,
max_points_per_voxel):
rng = np.random.RandomState(123)
points = rng.rand(rng.randint(0, 10000), ndim).astype(point_dtype)
row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64)
for i in range(batch_size):
row_splits[i + 1] = rng.randint(
points.shape[0] // batch_size) + row_splits[i]
points = points[:row_splits[batch_size]]
voxel_size = rng.uniform(0.01, 0.1, size=(ndim, )).astype(point_dtype)
point_range_min = rng.uniform(0.0, 0.3, size=(ndim, )).astype(point_dtype)
point_range_max = rng.uniform(0.7, 1.0, size=(ndim, )).astype(point_dtype)
ans = mltest.run_op(ml,
ml.device,
True,
ml.ops.voxelize,
points,
row_splits,
voxel_size,
point_range_min,
point_range_max,
max_points_per_voxel=max_points_per_voxel,
max_voxels=max_voxels)
voxels = convert_output_to_voxel_dict(ans)
voxels_reference = voxelize_python(
points,
row_splits,
voxel_size,
point_range_min,
point_range_max,
max_points_per_voxel=max_points_per_voxel,
max_voxels=max_voxels)
assert_equal_voxel_dicts(voxels, ref_b=voxels_reference)
def test_ragged_to_dense_random(dtype, ml, seed):
rng = np.random.RandomState(seed)
values = rng.random(size=(10000, )).astype(dtype)
row_splits = [0]
while row_splits[-1] < values.shape[0]:
row_splits.append(row_splits[-1] + rng.randint(0, 10))
row_splits[-1] = values.shape[0]
row_splits = np.array(row_splits, dtype=np.int64)
out_col_size = rng.randint(1, 37)
default_value = np.array(-1, dtype=dtype)
ans = mltest.run_op(ml, ml.device, True, ml.ops.ragged_to_dense, values,
row_splits, out_col_size, default_value)
expected = np.full((row_splits.shape[0] - 1, out_col_size), default_value)
for i in range(row_splits.shape[0] - 1):
for j, value_idx in enumerate(range(row_splits[i], row_splits[i + 1])):
if j < expected.shape[1]:
expected[i, j] = values[value_idx]
np.testing.assert_equal(ans, expected)
def test_invert_neighbors_list(dtype, attributes, ml):
# yapf: disable
# define connectivity for 3 query points and 3 input points
num_points = 3
edges = np.array(
[
[0, 0], [0, 1], [0, 2], # 3 neighbors
[1, 2], # 1 neighbors
[2, 1], [2, 2], # 2 neighbors
],
dtype=np.int32)
# the neighbors_index is the second column
neighbors_index = edges[:, 1]
# exclusive prefix sum of the number of neighbors
neighbors_row_splits = np.array([0, 3, 4, edges.shape[0]], dtype=np.int64)
if attributes == 'scalar':
neighbors_attributes = np.array([
10, 20, 30,
40,
50, 60,
], dtype=dtype)
elif attributes == 'none':
neighbors_attributes = np.array([], dtype=dtype)
elif attributes == 'multidim':
neighbors_attributes = np.array([
[10, 1], [20, 2], [30, 3],
[40, 4],
[50, 5], [60, 6],
], dtype=dtype)
# yapf: enable
ans = mltest.run_op(ml,
ml.device,
True,
ml.ops.invert_neighbors_list,
num_points=num_points,
inp_neighbors_index=neighbors_index,
inp_neighbors_row_splits=neighbors_row_splits,
inp_neighbors_attributes=neighbors_attributes)
expected_neighbors_row_splits = [0, 1, 3, edges.shape[0]]
np.testing.assert_equal(ans.neighbors_row_splits,
expected_neighbors_row_splits)
# checking the neighbors_index is more complicated because the order
# of the neighbors for each query point is not defined.
expected_neighbors_index = [
set([0]),
set([0, 2]),
set([0, 1, 2]),
]
for i, expected_neighbors_i in enumerate(expected_neighbors_index):
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
neighbors_i = set(ans.neighbors_index[start:end])
assert neighbors_i == expected_neighbors_i
if neighbors_attributes.shape == (0, ):
# if the input is a zero length vector then the returned attributes
# vector also must be a zero length vector
assert ans.neighbors_attributes.shape == (0, )
else:
# check if the attributes are still associated with the same edge
edge_attr_map = {
tuple(k): v
for k, v in zip(edges, neighbors_attributes)
}
for i, _ in enumerate(expected_neighbors_index):
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
# neighbors and attributes for point i
neighbors_i = ans.neighbors_index[start:end]
attributes_i = ans.neighbors_attributes[start:end]
for j, attr in zip(neighbors_i, attributes_i):
key = (j, i)
np.testing.assert_equal(attr, edge_attr_map[key])
def test_knn_search(dtype, ml, num_points_queries, metric, ignore_query_point,
return_distances):
rng = np.random.RandomState(123)
device = mltest.cpu_device
num_points, num_queries = num_points_queries
points = rng.random(size=(num_points, 3)).astype(dtype)
if ignore_query_point:
queries = points
else:
queries = rng.random(size=(num_queries, 3)).astype(dtype)
k = rng.randint(1, 11)
# kd tree for computing the ground truth
tree = cKDTree(points, copy_data=True)
p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric]
if k > num_points:
gt_neighbors_index = [tree.query(q, k, p=p_norm) for q in queries]
gt_neighbors_index = [
idxs[np.isfinite(dists)] for dists, idxs in gt_neighbors_index
]
else:
gt_neighbors_index = [tree.query(q, k, p=p_norm)[1] for q in queries]
layer = ml.layers.KNNSearch(metric=metric,
ignore_query_point=ignore_query_point,
return_distances=return_distances)
ans = mltest.run_op(
ml,
device,
layer,
True,
points,
queries=queries,
k=k,
)
for i, q in enumerate(queries):
# check neighbors
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
q_neighbors_index = ans.neighbors_index[start:end]
if k == 1:
gt_set = set([gt_neighbors_index[i]])
else:
gt_set = set(gt_neighbors_index[i])
if ignore_query_point:
gt_set.remove(i)
assert gt_set == set(q_neighbors_index)
# check distances
if return_distances:
q_neighbors_dist = ans.neighbors_distance[start:end]
for j, dist in zip(q_neighbors_index, q_neighbors_dist):
if metric == 'L2':
gt_dist = np.sum((q - points[j])**2)
else:
gt_dist = np.linalg.norm(q - points[j], ord=p_norm)
np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)
def test_knn_search_batches(ml, batch_size):
device = mltest.cpu_device
dtype = np.float32
metric = 'L2'
p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric]
ignore_query_point = False
return_distances = True
rng = np.random.RandomState(123)
# create array defining start and end of each batch
points_row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64)
queries_row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64)
for i in range(batch_size):
points_row_splits[i + 1] = rng.randint(15) + points_row_splits[i]
queries_row_splits[i + 1] = rng.randint(15) + queries_row_splits[i]
num_points = points_row_splits[-1]
num_queries = queries_row_splits[-1]
points = rng.random(size=(num_points, 3)).astype(dtype)
if ignore_query_point:
queries = points
queries_row_splits = points_row_splits
else:
queries = rng.random(size=(num_queries, 3)).astype(dtype)
k = rng.randint(1, 11)
# kd tree for computing the ground truth
gt_neighbors_index = []
for i in range(batch_size):
points_i = points[points_row_splits[i]:points_row_splits[i + 1]]
queries_i = queries[queries_row_splits[i]:queries_row_splits[i + 1]]
tree = cKDTree(points_i, copy_data=True)
if k > points_i.shape[0]:
tmp = [tree.query(q, k, p=p_norm) for q in queries_i]
tmp = [
list(idxs[np.isfinite(dists)] + points_row_splits[i])
for dists, idxs in tmp
]
else:
tmp = [
list(tree.query(q, k, p=p_norm)[1] + points_row_splits[i])
for q in queries_i
]
gt_neighbors_index.extend(tmp)
layer = ml.layers.KNNSearch(metric=metric,
ignore_query_point=ignore_query_point,
return_distances=return_distances)
ans = mltest.run_op(ml,
device,
layer,
True,
points,
queries=queries,
k=k,
points_row_splits=points_row_splits,
queries_row_splits=queries_row_splits)
for i, q in enumerate(queries):
# check neighbors
start = ans.neighbors_row_splits[i]
end = ans.neighbors_row_splits[i + 1]
q_neighbors_index = ans.neighbors_index[start:end]
if k == 1:
gt_set = set([gt_neighbors_index[i]])
else:
gt_set = set(gt_neighbors_index[i])
if ignore_query_point:
gt_set.remove(i)
assert gt_set == set(q_neighbors_index)
# check distances
if return_distances:
q_neighbors_dist = ans.neighbors_distance[start:end]
for j, dist in zip(q_neighbors_index, q_neighbors_dist):
if metric == 'L2':
gt_dist = np.sum((q - points[j])**2)
else:
gt_dist = np.linalg.norm(q - points[j], ord=p_norm)
np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)
