warp-lang 1.5.0__py3-none-win_amd64.whl → 1.6.0__py3-none-win_amd64.whl

warp/tests/test_mesh.py

@@ -88,21 +88,27 @@ def read_indices_kernel(
 
 
 def test_mesh_read_properties(test, device):
-    points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
-    indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-    mesh = wp.Mesh(points=points, indices=indices)
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
 
-    assert mesh.points.size == POINT_COUNT
-    assert mesh.indices.size == VERTEX_COUNT
-    assert int(mesh.indices.size / 3) == FACE_COUNT
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
 
-    out_points = wp.empty(POINT_COUNT, dtype=wp.vec3, device=device)
-    wp.launch(read_points_kernel, dim=POINT_COUNT, inputs=[mesh.id], outputs=[out_points], device=device)
-    assert_np_equal(out_points.numpy(), np.array(POINT_POSITIONS))
+        assert mesh.points.size == POINT_COUNT
+        assert mesh.indices.size == VERTEX_COUNT
+        assert int(mesh.indices.size / 3) == FACE_COUNT
 
-    out_indices = wp.empty(VERTEX_COUNT, dtype=int, device=device)
-    wp.launch(read_indices_kernel, dim=FACE_COUNT, inputs=[mesh.id], outputs=[out_indices], device=device)
-    assert_np_equal(out_indices.numpy(), np.array(RIGHT_HANDED_FACE_VERTEX_INDICES))
+        out_points = wp.empty(POINT_COUNT, dtype=wp.vec3, device=device)
+        wp.launch(read_points_kernel, dim=POINT_COUNT, inputs=[mesh.id], outputs=[out_points], device=device)
+        assert_np_equal(out_points.numpy(), np.array(POINT_POSITIONS))
+
+        out_indices = wp.empty(VERTEX_COUNT, dtype=int, device=device)
+        wp.launch(read_indices_kernel, dim=FACE_COUNT, inputs=[mesh.id], outputs=[out_indices], device=device)
+        assert_np_equal(out_indices.numpy(), np.array(RIGHT_HANDED_FACE_VERTEX_INDICES))
 
 
 @wp.kernel(enable_backward=False)
@@ -127,17 +133,23 @@ def query_point_kernel(
 
 
 def test_mesh_query_point(test, device):
-    points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
 
-    indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-    mesh = wp.Mesh(points=points, indices=indices)
-    expected_sign = -1.0
-    wp.launch(query_point_kernel, dim=1, inputs=[mesh.id, expected_sign], device=device)
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
 
-    indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-    mesh = wp.Mesh(points=points, indices=indices)
-    expected_sign = 1.0
-    wp.launch(query_point_kernel, dim=1, inputs=[mesh.id, expected_sign], device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+        expected_sign = -1.0
+        wp.launch(query_point_kernel, dim=1, inputs=[mesh.id, expected_sign], device=device)
+
+        indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices)
+        expected_sign = 1.0
+        wp.launch(query_point_kernel, dim=1, inputs=[mesh.id, expected_sign], device=device)
 
 
 @wp.kernel(enable_backward=False)
@@ -179,53 +191,65 @@ def query_ray_kernel(
 
 
 def test_mesh_query_ray(test, device):
-    points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
-
-    indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-    mesh = wp.Mesh(points=points, indices=indices)
-    expected_sign = -1.0
-    wp.launch(
-        query_ray_kernel,
-        dim=1,
-        inputs=[
-            mesh.id,
-            expected_sign,
-        ],
-        device=device,
-    )
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
 
-    indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-    mesh = wp.Mesh(points=points, indices=indices)
-    expected_sign = 1.0
-    wp.launch(
-        query_ray_kernel,
-        dim=1,
-        inputs=[
-            mesh.id,
-            expected_sign,
-        ],
-        device=device,
-    )
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+        expected_sign = -1.0
+        wp.launch(
+            query_ray_kernel,
+            dim=1,
+            inputs=[
+                mesh.id,
+                expected_sign,
+            ],
+            device=device,
+        )
+
+        indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices)
+        expected_sign = 1.0
+        wp.launch(
+            query_ray_kernel,
+            dim=1,
+            inputs=[
+                mesh.id,
+                expected_sign,
+            ],
+            device=device,
+        )
 
 
 def test_mesh_refit_graph(test, device):
-    points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
 
-    indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
-    mesh = wp.Mesh(points=points, indices=indices)
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
 
-    wp.capture_begin(device, force_module_load=False)
-    try:
-        mesh.refit()
-    finally:
-        graph = wp.capture_end(device)
+        wp.capture_begin(device, force_module_load=False)
+        try:
+            mesh.refit()
+        finally:
+            graph = wp.capture_end(device)
 
-    # replay
-    num_iters = 10
-    for _ in range(num_iters):
-        wp.capture_launch(graph)
+        # replay
+        num_iters = 10
+        for _ in range(num_iters):
+            wp.capture_launch(graph)
 
-    wp.synchronize_device(device)
+        wp.synchronize_device(device)
 
 
 def test_mesh_exceptions(test, device):

warp/tests/test_mesh_query_aabb.py

@@ -5,11 +5,13 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import os
 import unittest
 
 import numpy as np
 
 import warp as wp
+import warp.examples
 from warp.tests.unittest_utils import *
 
 
@@ -192,6 +194,162 @@ def test_mesh_query_aabb_count_nonoverlap(test, device):
         test.assertTrue(c == 1)
 
 
+@wp.kernel
+def compute_num_contact_with_checksums(
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+    mesh_id: wp.uint64,
+    counts: wp.array(dtype=int),
+    check_sums: wp.array(dtype=int),
+):
+    face_index = int(0)
+    face_u = float(0.0)
+    face_v = float(0.0)
+    sign = float(0.0)
+
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    query = wp.mesh_query_aabb(mesh_id, lower, upper)
+    count = int(0)
+
+    check_sum = int(0)
+    for _index in query:
+        check_sum = check_sum ^ _index
+        count = count + 1
+
+    counts[tid] = count
+    check_sums[tid] = check_sum
+
+
+@wp.func
+def intersect_aabb_aabb(a_lower: wp.vec3, a_upper: wp.vec3, b_lower: wp.vec3, b_upper: wp.vec3):
+    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 False
+    else:
+        return True
+
+
+@wp.kernel
+def compute_num_contact_with_checksums_brutal(
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+    mesh_points: wp.array(dtype=wp.vec3),
+    mesh_indices: wp.array(dtype=int),
+    counts: wp.array(dtype=int),
+    check_sums: wp.array(dtype=int),
+):
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    check_sum = int(0)
+    count = int(0)
+    num_faces = mesh_indices.shape[0] / 3
+
+    for face in range(num_faces):
+        i = mesh_indices[face * 3 + 0]
+        j = mesh_indices[face * 3 + 1]
+        k = mesh_indices[face * 3 + 2]
+
+        x0 = mesh_points[i]  # point zero
+        x1 = mesh_points[j]  # point one
+        x2 = mesh_points[k]  # point two
+
+        tri_lower = min_vec3(min_vec3(x0, x1), x2)
+        tri_upper = max_vec3(max_vec3(x0, x1), x2)
+
+        if intersect_aabb_aabb(lower, upper, tri_lower, tri_upper):
+            check_sum = check_sum ^ face
+            count = count + 1
+
+    counts[tid] = count
+    check_sums[tid] = check_sum
+
+
+def load_mesh():
+    from pxr import Usd, UsdGeom
+
+    usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+    usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+    vertices = np.array(usd_geom.GetPointsAttr().Get())
+    faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+    return vertices, faces
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_mesh_query_aabb_count_overlap_with_checksum(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    points, indices = load_mesh()
+    points_wp = wp.array(points, dtype=wp.vec3, device=device)
+    indices_wp = wp.array(indices, dtype=int, device=device)
+
+    for constructor in constructors:
+        m = wp.Mesh(points=points_wp, indices=indices_wp, bvh_constructor=constructor)
+
+        num_test_bounds = 10000
+        test_bound_relative_size = 0.01
+
+        world_min = np.min(points, axis=0)
+        world_max = np.max(points, axis=0)
+
+        world_center = 0.5 * (world_min + world_max)
+        world_size = world_max - world_min
+
+        rng = np.random.default_rng(123)
+
+        centers = (
+            rng.uniform(-0.5, 0.5, size=num_test_bounds * 3).reshape(num_test_bounds, 3) * world_size + world_center
+        )
+        diffs = (
+            0.5 * test_bound_relative_size * rng.random(num_test_bounds * 3).reshape(num_test_bounds, 3) * world_size
+        )
+
+        lowers = wp.array(centers - diffs, dtype=wp.vec3, device=device)
+        uppers = wp.array(centers + diffs, dtype=wp.vec3, device=device)
+
+        counts = wp.empty(n=num_test_bounds, dtype=int, device=device)
+        checksums = wp.empty(n=num_test_bounds, dtype=int, device=device)
+
+        wp.launch(
+            kernel=compute_num_contact_with_checksums,
+            dim=num_test_bounds,
+            inputs=[lowers, uppers, m.id],
+            outputs=[counts, checksums],
+            device=device,
+        )
+
+        counts_brutal = wp.empty(n=num_test_bounds, dtype=int, device=device)
+        checksums_brutal = wp.empty(n=num_test_bounds, dtype=int, device=device)
+
+        wp.launch(
+            kernel=compute_num_contact_with_checksums_brutal,
+            dim=num_test_bounds,
+            inputs=[lowers, uppers, points_wp, indices_wp],
+            outputs=[counts_brutal, checksums_brutal],
+            device=device,
+        )
+
+        assert_array_equal(counts, counts_brutal)
+        assert_array_equal(checksums, checksums_brutal)
+
+
 devices = get_test_devices()
 
 
@@ -220,6 +378,13 @@ add_function_test(
     test_mesh_query_aabb_count_nonoverlap,
     devices=devices,
 )
+add_function_test(
+    TestMeshQueryAABBMethods,
+    "test_mesh_query_aabb_count_overlap_with_checksum",
+    test_mesh_query_aabb_count_overlap_with_checksum,
+    devices=devices,
+)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()