warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_bvh.py

@@ -5,13 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
-
-np.random.seed(42)
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -70,70 +69,103 @@ def intersect_ray_aabb(start, dir, lower, upper):
         return 0
 
 
-def test_bvh_query_aabb(test, device):
+def test_bvh(test, type, device):
+    rng = np.random.default_rng(123)
+
     num_bounds = 100
-    lowers = np.random.rand(num_bounds, 3) * 5.0
-    uppers = lowers + np.random.rand(num_bounds, 3) * 5.0
+    lowers = rng.random(size=(num_bounds, 3)) * 5.0
+    uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
 
-    bvh = wp.Bvh(
-        lowers=wp.array(lowers, dtype=wp.vec3, device=device), uppers=wp.array(uppers, dtype=wp.vec3, device=device)
-    )
+    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)
 
-    wp.launch(
-        kernel=bvh_query_aabb, dim=1, inputs=[bvh.id, query_lower, query_upper, bounds_intersected], device=device
-    )
+    query_start = wp.vec3(0.0, 0.0, 0.0)
+    query_dir = wp.normalize(wp.vec3(1.0, 1.0, 1.0))
 
-    device_intersected = bounds_intersected.numpy()
+    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, 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_()
 
-    for i in range(num_bounds):
-        lower = lowers[i]
-        upper = uppers[i]
-        host_intersected = aabb_overlap(lower, upper, query_lower, query_upper)
-        test.assertEqual(host_intersected, device_intersected[i])
+
+def test_bvh_query_aabb(test, device):
+    test_bvh(test, "AABB", device)
 
 
 def test_bvh_query_ray(test, device):
-    num_bounds = 100
-    lowers = np.random.rand(num_bounds, 3) * 5.0
-    uppers = lowers + np.random.rand(num_bounds, 3) * 5.0
+    test_bvh(test, "ray", device)
 
-    bvh = wp.Bvh(
-        lowers=wp.array(lowers, dtype=wp.vec3, device=device), uppers=wp.array(uppers, dtype=wp.vec3, device=device)
-    )
 
-    bounds_intersected = wp.zeros(shape=(num_bounds), dtype=int, device=device)
+def test_bvh_codegen_adjoints_with_select(test, device):
+    def kernel_fn(
+        bvh: wp.uint64,
+    ):
+        v = wp.vec3(0.0, 0.0, 0.0)
+        bounds_nr = int(0)
 
-    query_start = wp.vec3(0.0, 0.0, 0.0)
-    query_dir = wp.normalize(wp.vec3(1.0, 1.0, 1.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.launch(kernel=bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device)
+            wp.bvh_query_next(query_1, bounds_nr)
+            wp.bvh_query_next(query_2, bounds_nr)
 
-    device_intersected = bounds_intersected.numpy()
+    wp.Kernel(func=kernel_fn)
 
-    for i in range(num_bounds):
-        lower = lowers[i]
-        upper = uppers[i]
-        host_intersected = intersect_ray_aabb(query_start, query_dir, lower, upper)
-        test.assertEqual(host_intersected, device_intersected[i])
 
+devices = get_test_devices()
 
-def register(parent):
-    devices = get_test_devices()
 
-    class TestBvh(parent):
-        pass
+class TestBvh(unittest.TestCase):
+    pass
 
-    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)
 
-    return TestBvh
+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_bvh_codegen_adjoints_with_select", test_bvh_codegen_adjoints_with_select, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_closest_point_edge_edge.py

@@ -5,12 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
-
 import unittest
 
+import numpy as np
+
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 epsilon = 0.00001
@@ -171,61 +171,58 @@ def test_edge_edge_perpendicular_s0_t1(test, device):
     test.assertAlmostEqual(st0[1], 1.0)  # t value
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestClosestPointEdgeEdgeMethods(parent):
-        pass
-
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods,
-        "test_edge_edge_middle_crossing",
-        test_edge_edge_middle_crossing,
-        devices=devices,
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods, "test_edge_edge_parallel_s1_t0", test_edge_edge_parallel_s1_t0, devices=devices
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods, "test_edge_edge_parallel_s0_t1", test_edge_edge_parallel_s0_t1, devices=devices
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods,
-        "test_edge_edge_both_degenerate_case",
-        test_edge_edge_both_degenerate_case,
-        devices=devices,
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods,
-        "test_edge_edge_degenerate_first_edge",
-        test_edge_edge_degenerate_first_edge,
-        devices=devices,
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods,
-        "test_edge_edge_degenerate_second_edge",
-        test_edge_edge_degenerate_second_edge,
-        devices=devices,
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods, "test_edge_edge_parallel", test_edge_edge_parallel, devices=devices
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods,
-        "test_edge_edge_perpendicular_s1_t0",
-        test_edge_edge_perpendicular_s1_t0,
-        devices=devices,
-    )
-    add_function_test(
-        TestClosestPointEdgeEdgeMethods,
-        "test_edge_edge_perpendicular_s0_t1",
-        test_edge_edge_perpendicular_s0_t1,
-        devices=devices,
-    )
-
-    return TestClosestPointEdgeEdgeMethods
+devices = get_test_devices()
+
+
+class TestClosestPointEdgeEdgeMethods(unittest.TestCase):
+    pass
+
+
+add_function_test(
+    TestClosestPointEdgeEdgeMethods,
+    "test_edge_edge_middle_crossing",
+    test_edge_edge_middle_crossing,
+    devices=devices,
+)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods, "test_edge_edge_parallel_s1_t0", test_edge_edge_parallel_s1_t0, devices=devices
+)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods, "test_edge_edge_parallel_s0_t1", test_edge_edge_parallel_s0_t1, devices=devices
+)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods,
+    "test_edge_edge_both_degenerate_case",
+    test_edge_edge_both_degenerate_case,
+    devices=devices,
+)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods,
+    "test_edge_edge_degenerate_first_edge",
+    test_edge_edge_degenerate_first_edge,
+    devices=devices,
+)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods,
+    "test_edge_edge_degenerate_second_edge",
+    test_edge_edge_degenerate_second_edge,
+    devices=devices,
+)
+add_function_test(TestClosestPointEdgeEdgeMethods, "test_edge_edge_parallel", test_edge_edge_parallel, devices=devices)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods,
+    "test_edge_edge_perpendicular_s1_t0",
+    test_edge_edge_perpendicular_s1_t0,
+    devices=devices,
+)
+add_function_test(
+    TestClosestPointEdgeEdgeMethods,
+    "test_edge_edge_perpendicular_s0_t1",
+    test_edge_edge_perpendicular_s0_t1,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)