warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/tests/geometry/test_bvh.py

@@ -0,0 +1,200 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def bvh_query_aabb(bvh_id: wp.uint64, lower: wp.vec3, upper: wp.vec3, bounds_intersected: wp.array(dtype=int)):
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
+    bounds_nr = int(0)
+
+    while wp.bvh_query_next(query, bounds_nr):
+        bounds_intersected[bounds_nr] = 1
+
+
+@wp.kernel
+def bvh_query_ray(bvh_id: wp.uint64, start: wp.vec3, dir: wp.vec3, bounds_intersected: wp.array(dtype=int)):
+    query = wp.bvh_query_ray(bvh_id, start, dir)
+    bounds_nr = int(0)
+
+    while wp.bvh_query_next(query, bounds_nr):
+        bounds_intersected[bounds_nr] = 1
+
+
+def aabb_overlap(a_lower, a_upper, b_lower, b_upper):
+    if (
+        a_lower[0] > b_upper[0]
+        or a_lower[1] > b_upper[1]
+        or a_lower[2] > b_upper[2]
+        or a_upper[0] < b_lower[0]
+        or a_upper[1] < b_lower[1]
+        or a_upper[2] < b_lower[2]
+    ):
+        return 0
+    else:
+        return 1
+
+
+def intersect_ray_aabb(start, rcp_dir, lower, upper):
+    l1 = (lower[0] - start[0]) * rcp_dir[0]
+    l2 = (upper[0] - start[0]) * rcp_dir[0]
+    lmin = min(l1, l2)
+    lmax = max(l1, l2)
+
+    l1 = (lower[1] - start[1]) * rcp_dir[1]
+    l2 = (upper[1] - start[1]) * rcp_dir[1]
+    lmin = max(min(l1, l2), lmin)
+    lmax = min(max(l1, l2), lmax)
+
+    l1 = (lower[2] - start[2]) * rcp_dir[2]
+    l2 = (upper[2] - start[2]) * rcp_dir[2]
+    lmin = max(min(l1, l2), lmin)
+    lmax = min(max(l1, l2), lmax)
+
+    if lmax >= 0.0 and lmax >= lmin:
+        return 1
+    else:
+        return 0
+
+
+def test_bvh(test, type, device):
+    rng = np.random.default_rng(123)
+
+    num_bounds = 100
+    lowers = rng.random(size=(num_bounds, 3)) * 5.0
+    uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
+
+    device_lowers = wp.array(lowers, dtype=wp.vec3, device=device)
+    device_uppers = wp.array(uppers, dtype=wp.vec3, device=device)
+
+    bvh = wp.Bvh(device_lowers, device_uppers)
+
+    bounds_intersected = wp.zeros(shape=(num_bounds), dtype=int, device=device)
+
+    query_lower = wp.vec3(2.0, 2.0, 2.0)
+    query_upper = wp.vec3(8.0, 8.0, 8.0)
+
+    query_start = wp.vec3(0.0, 0.0, 0.0)
+    query_dir = wp.normalize(wp.vec3(1.0, 1.0, 1.0))
+
+    for test_case in range(2):
+        if type == "AABB":
+            wp.launch(
+                kernel=bvh_query_aabb,
+                dim=1,
+                inputs=[bvh.id, query_lower, query_upper, bounds_intersected],
+                device=device,
+            )
+        else:
+            wp.launch(
+                kernel=bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device
+            )
+
+        device_intersected = bounds_intersected.numpy()
+
+        for i in range(num_bounds):
+            lower = lowers[i]
+            upper = uppers[i]
+            if type == "AABB":
+                host_intersected = aabb_overlap(lower, upper, query_lower, query_upper)
+            else:
+                host_intersected = intersect_ray_aabb(query_start, 1.0 / query_dir, lower, upper)
+
+            test.assertEqual(host_intersected, device_intersected[i])
+
+        if test_case == 0:
+            lowers = rng.random(size=(num_bounds, 3)) * 5.0
+            uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
+            wp.copy(device_lowers, wp.array(lowers, dtype=wp.vec3, device=device))
+            wp.copy(device_uppers, wp.array(uppers, dtype=wp.vec3, device=device))
+            bvh.refit()
+            bounds_intersected.zero_()
+
+
+def test_bvh_query_aabb(test, device):
+    test_bvh(test, "AABB", device)
+
+
+def test_bvh_query_ray(test, device):
+    test_bvh(test, "ray", device)
+
+
+def test_gh_288(test, device):
+    num_bounds = 1
+    lowers = ((0.5, -1.0, -1.0),) * num_bounds
+    uppers = ((1.0, 1.0, 1.0),) * num_bounds
+
+    device_lowers = wp.array(lowers, dtype=wp.vec3f, device=device)
+    device_uppers = wp.array(uppers, dtype=wp.vec3f, device=device)
+
+    bvh = wp.Bvh(device_lowers, device_uppers)
+
+    bounds_intersected = wp.zeros(shape=num_bounds, dtype=int, device=device)
+
+    for x in (0.0, 0.75):
+        query_start = wp.vec3(x, 0.0, 0.0)
+        query_dir = wp.vec3(1.0, 0.0, 0.0)
+
+        wp.launch(
+            kernel=bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device
+        )
+
+        device_intersected = bounds_intersected.numpy()
+        test.assertEqual(device_intersected.sum(), num_bounds)
+
+
+devices = get_test_devices()
+
+
+class TestBvh(unittest.TestCase):
+    def test_bvh_codegen_adjoints_with_select(self):
+        def kernel_fn(bvh: wp.uint64):
+            v = wp.vec3(0.0, 0.0, 0.0)
+            bounds_nr = int(0)
+
+            if True:
+                query_1 = wp.bvh_query_aabb(bvh, v, v)
+                query_2 = wp.bvh_query_ray(bvh, v, v)
+
+                wp.bvh_query_next(query_1, bounds_nr)
+                wp.bvh_query_next(query_2, bounds_nr)
+            else:
+                query_1 = wp.bvh_query_aabb(bvh, v, v)
+                query_2 = wp.bvh_query_ray(bvh, v, v)
+
+                wp.bvh_query_next(query_1, bounds_nr)
+                wp.bvh_query_next(query_2, bounds_nr)
+
+        wp.Kernel(func=kernel_fn)
+
+    def test_bvh_new_del(self):
+        # test the scenario in which a bvh is created but not initialized before gc
+        instance = wp.Bvh.__new__(wp.Bvh)
+        instance.__del__()
+
+
+add_function_test(TestBvh, "test_bvh_aabb", test_bvh_query_aabb, devices=devices)
+add_function_test(TestBvh, "test_bvh_ray", test_bvh_query_ray, devices=devices)
+add_function_test(TestBvh, "test_gh_288", test_gh_288, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/geometry/test_hash_grid.py

@@ -0,0 +1,221 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+num_points = 4096
+dim_x = 128
+dim_y = 128
+dim_z = 128
+
+scale = 150.0
+
+cell_radius = 8.0
+query_radius = 8.0
+
+num_runs = 4
+
+print_enabled = False
+
+
+@wp.kernel
+def count_neighbors(grid: wp.uint64, radius: float, points: wp.array(dtype=wp.vec3), counts: wp.array(dtype=int)):
+    tid = wp.tid()
+
+    # order threads by cell
+    i = wp.hash_grid_point_id(grid, tid)
+
+    # query point
+    p = points[i]
+    count = int(0)
+
+    # construct query around point p
+    for index in wp.hash_grid_query(grid, p, radius):
+        # compute distance to point
+        d = wp.length(p - points[index])
+
+        if d <= radius:
+            count += 1
+
+    counts[i] = count
+
+
+@wp.kernel
+def count_neighbors_reference(
+    radius: float, points: wp.array(dtype=wp.vec3), counts: wp.array(dtype=int), num_points: int
+):
+    tid = wp.tid()
+
+    i = tid % num_points
+    j = tid // num_points
+
+    # query point
+    p = points[i]
+    q = points[j]
+
+    # compute distance to point
+    d = wp.length(p - q)
+
+    if d <= radius:
+        wp.atomic_add(counts, i, 1)
+
+
+def particle_grid(dim_x, dim_y, dim_z, lower, radius, jitter):
+    rng = np.random.default_rng(123)
+    points = np.meshgrid(np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z))
+    points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
+    points_t = points_t + rng.random(size=points_t.shape) * radius * jitter
+
+    return points_t.reshape((-1, 3))
+
+
+def test_hashgrid_query(test, device):
+    wp.load_module(device=device)
+
+    grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+
+    for i in range(num_runs):
+        if print_enabled:
+            print(f"Run: {i + 1}")
+            print("---------")
+
+        points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)
+
+        points_arr = wp.array(points, dtype=wp.vec3, device=device)
+        counts_arr = wp.zeros(len(points), dtype=int, device=device)
+        counts_arr_ref = wp.zeros(len(points), dtype=int, device=device)
+
+        profiler = {}
+
+        with wp.ScopedTimer("grid operations", print=print_enabled, dict=profiler, synchronize=True):
+            with wp.ScopedTimer("brute", print=print_enabled, dict=profiler, synchronize=True):
+                wp.launch(
+                    kernel=count_neighbors_reference,
+                    dim=len(points) * len(points),
+                    inputs=[query_radius, points_arr, counts_arr_ref, len(points)],
+                    device=device,
+                )
+                wp.synchronize_device(device)
+
+            with wp.ScopedTimer("grid build", print=print_enabled, dict=profiler, synchronize=True):
+                grid.build(points_arr, cell_radius)
+
+            with wp.ScopedTimer("grid query", print=print_enabled, dict=profiler, synchronize=True):
+                wp.launch(
+                    kernel=count_neighbors,
+                    dim=len(points),
+                    inputs=[grid.id, query_radius, points_arr, counts_arr],
+                    device=device,
+                )
+
+        counts = counts_arr.numpy()
+        counts_ref = counts_arr_ref.numpy()
+
+        if print_enabled:
+            print(f"Grid min: {np.min(counts)} max: {np.max(counts)} avg: {np.mean(counts)}")
+            print(f"Ref min: {np.min(counts_ref)} max: {np.max(counts_ref)} avg: {np.mean(counts_ref)}")
+
+            print(f"Passed: {np.array_equal(counts, counts_ref)}")
+
+        assert_np_equal(counts, counts_ref)
+
+
+def test_hashgrid_inputs(test, device):
+    points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)
+    points_ref = wp.array(points, dtype=wp.vec3, device=device)
+    counts_ref = wp.zeros(len(points), dtype=int, device=device)
+
+    grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+    grid.build(points_ref, cell_radius)
+
+    # get reference counts
+    wp.launch(
+        kernel=count_neighbors, dim=len(points), inputs=[grid.id, query_radius, points_ref, counts_ref], device=device
+    )
+
+    # test with strided 1d input arrays
+    for stride in [2, 3]:
+        with test.subTest(msg=f"stride_{stride}"):
+            points_buffer = wp.zeros(len(points) * stride, dtype=wp.vec3, device=device)
+            points_strided = points_buffer[::stride]
+            wp.copy(points_strided, points_ref)
+            counts_strided = wp.zeros(len(points), dtype=int, device=device)
+
+            grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+            grid.build(points_strided, cell_radius)
+
+            wp.launch(
+                kernel=count_neighbors,
+                dim=len(points),
+                inputs=[grid.id, query_radius, points_ref, counts_strided],
+                device=device,
+            )
+
+            assert_array_equal(counts_strided, counts_ref)
+
+    # test with multidimensional input arrays
+    for ndim in [2, 3, 4]:
+        with test.subTest(msg=f"ndim_{ndim}"):
+            shape = (len(points) // (2 ** (ndim - 1)), *((ndim - 1) * (2,)))
+            points_ndim = wp.zeros(shape, dtype=wp.vec3, device=device)
+            wp.copy(points_ndim, points_ref)
+            counts_ndim = wp.zeros(len(points), dtype=int, device=device)
+
+            grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+            grid.build(points_ndim, cell_radius)
+
+            wp.launch(
+                kernel=count_neighbors,
+                dim=len(points),
+                inputs=[grid.id, query_radius, points_ref, counts_ndim],
+                device=device,
+            )
+
+            assert_array_equal(counts_ndim, counts_ref)
+
+
+devices = get_test_devices()
+
+
+class TestHashGrid(unittest.TestCase):
+    def test_hashgrid_codegen_adjoints_with_select(self):
+        def kernel_fn(grid: wp.uint64):
+            v = wp.vec3(0.0, 0.0, 0.0)
+
+            if True:
+                query = wp.hash_grid_query(grid, v, 0.0)
+            else:
+                query = wp.hash_grid_query(grid, v, 0.0)
+
+        wp.Kernel(func=kernel_fn)
+
+    def test_hashgrid_new_del(self):
+        # test the scenario in which a hashgrid is created but not initialized before gc
+        instance = wp.HashGrid.__new__(wp.HashGrid)
+        instance.__del__()
+
+
+add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
+add_function_test(TestHashGrid, "test_hashgrid_inputs", test_hashgrid_inputs, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)

warp/tests/geometry/test_marching_cubes.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def make_field(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
+    i, j, k = wp.tid()
+
+    p = wp.vec3(float(i), float(j), float(k))
+
+    d = wp.length(p - center) - radius
+
+    field[i, j, k] = d
+
+
+def test_marching_cubes(test, device):
+    dim = 64
+    max_verts = 10**6
+    max_tris = 10**6
+
+    field = wp.zeros(shape=(dim, dim, dim), dtype=float, device=device)
+
+    iso = wp.MarchingCubes(nx=dim, ny=dim, nz=dim, max_verts=max_verts, max_tris=max_tris, device=device)
+
+    radius = dim / 4.0
+
+    wp.launch(make_field, dim=field.shape, inputs=[field, wp.vec3(dim / 2, dim / 2, dim / 2), radius], device=device)
+
+    iso.surface(field=field, threshold=0.0)
+
+    # check that all returned vertices lie on the surface of the sphere
+    length = np.linalg.norm(iso.verts.numpy() - np.array([dim / 2, dim / 2, dim / 2]), axis=1)
+    error = np.abs(length - radius)
+
+    test.assertTrue(np.max(error) < 1.0)
+
+    iso.resize(nx=dim * 2, ny=dim * 2, nz=dim * 2, max_verts=max_verts, max_tris=max_tris)
+
+
+devices = get_selected_cuda_test_devices()
+
+
+class TestMarchingCubes(unittest.TestCase):
+    def test_marching_cubes_new_del(self):
+        # test the scenario in which a MarchingCubes instance is created but not initialized before gc
+        instance = wp.MarchingCubes.__new__(wp.MarchingCubes)
+        instance.__del__()
+
+
+add_function_test(TestMarchingCubes, "test_marching_cubes", test_marching_cubes, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)