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

warp/tests/test_mesh_query_point.py

@@ -5,14 +5,13 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
-import numpy as np
 import math
+import unittest
 
-import warp as wp
-from warp.tests.test_base import *
+import numpy as np
 
-np.random.seed(42)
+import warp as wp
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -44,6 +43,110 @@ def sample_mesh_query(
     query_faces[tid] = face_index
     query_dist[tid] = wp.length(cp - p)
 
+    query = wp.mesh_query_point(mesh, p, max_dist)
+    wp.expect_eq(query.sign, sign)
+    wp.expect_eq(query.face, face_index)
+    wp.expect_eq(query.u, face_u)
+    wp.expect_eq(query.v, face_v)
+
+
+@wp.kernel
+def sample_mesh_query_no_sign(
+    mesh: wp.uint64,
+    query_points: wp.array(dtype=wp.vec3),
+    query_faces: wp.array(dtype=int),
+    query_dist: wp.array(dtype=float),
+):
+    tid = wp.tid()
+
+    face_index = int(0)
+    face_u = float(0.0)
+    face_v = float(0.0)
+
+    max_dist = 10012.0
+
+    p = query_points[tid]
+
+    wp.mesh_query_point_no_sign(mesh, p, max_dist, face_index, face_u, face_v)
+
+    cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
+
+    query_faces[tid] = face_index
+    query_dist[tid] = wp.length(cp - p)
+
+    query = wp.mesh_query_point_no_sign(mesh, p, max_dist)
+    wp.expect_eq(query.face, face_index)
+    wp.expect_eq(query.u, face_u)
+    wp.expect_eq(query.v, face_v)
+
+
+@wp.kernel
+def sample_mesh_query_sign_normal(
+    mesh: wp.uint64,
+    query_points: wp.array(dtype=wp.vec3),
+    query_faces: wp.array(dtype=int),
+    query_signs: wp.array(dtype=float),
+    query_dist: wp.array(dtype=float),
+):
+    tid = wp.tid()
+
+    face_index = int(0)
+    face_u = float(0.0)
+    face_v = float(0.0)
+    sign = float(0.0)
+
+    max_dist = 10012.0
+
+    p = query_points[tid]
+
+    wp.mesh_query_point_sign_normal(mesh, p, max_dist, sign, face_index, face_u, face_v)
+
+    cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
+
+    query_signs[tid] = sign
+    query_faces[tid] = face_index
+    query_dist[tid] = wp.length(cp - p)
+
+    query = wp.mesh_query_point_sign_normal(mesh, p, max_dist)
+    wp.expect_eq(query.sign, sign)
+    wp.expect_eq(query.face, face_index)
+    wp.expect_eq(query.u, face_u)
+    wp.expect_eq(query.v, face_v)
+
+
+@wp.kernel
+def sample_mesh_query_sign_winding_number(
+    mesh: wp.uint64,
+    query_points: wp.array(dtype=wp.vec3),
+    query_faces: wp.array(dtype=int),
+    query_signs: wp.array(dtype=float),
+    query_dist: wp.array(dtype=float),
+):
+    tid = wp.tid()
+
+    face_index = int(0)
+    face_u = float(0.0)
+    face_v = float(0.0)
+    sign = float(0.0)
+
+    max_dist = 10012.0
+
+    p = query_points[tid]
+
+    wp.mesh_query_point_sign_winding_number(mesh, p, max_dist, sign, face_index, face_u, face_v)
+
+    cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
+
+    query_signs[tid] = sign
+    query_faces[tid] = face_index
+    query_dist[tid] = wp.length(cp - p)
+
+    query = wp.mesh_query_point_sign_winding_number(mesh, p, max_dist)
+    wp.expect_eq(query.sign, sign)
+    wp.expect_eq(query.face, face_index)
+    wp.expect_eq(query.u, face_u)
+    wp.expect_eq(query.v, face_v)
+
 
 @wp.func
 def triangle_closest_point(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
@@ -107,7 +210,7 @@ def solid_angle(v0: wp.vec3, v1: wp.vec3, v2: wp.vec3, p: wp.vec3):
     det = wp.dot(a, wp.cross(b, c))
     den = a_len * b_len * c_len + wp.dot(a, b) * c_len + wp.dot(b, c) * a_len + wp.dot(c, a) * b_len
 
-    return wp.atan2(det, den)
+    return 2.0 * wp.atan2(det, den)
 
 
 @wp.kernel
@@ -125,7 +228,7 @@ def sample_mesh_brute(
     min_face = int(0)
     min_dist = float(1.0e6)
 
-    winding_angle = float(0.0)
+    sum_solid_angle = float(0.0)
 
     p = query_points[tid]
 
@@ -134,7 +237,7 @@ def sample_mesh_brute(
         b = tri_points[tri_indices[i * 3 + 1]]
         c = tri_points[tri_indices[i * 3 + 2]]
 
-        winding_angle += solid_angle(a, b, c, p)
+        sum_solid_angle += solid_angle(a, b, c, p)
 
         bary = triangle_closest_point(a, b, c, p)
         u = bary[0]
@@ -147,16 +250,20 @@ def sample_mesh_brute(
             min_dist = cp_dist
             min_face = i
 
+    # for an inside point, the sum of the solid angle should be 4PI
+    # for an outside point, the sum should be 0
     query_faces[tid] = min_face
-    query_signs[tid] = winding_angle
+    query_signs[tid] = sum_solid_angle
     query_dist[tid] = min_dist
 
 
 # constructs a grid of evenly spaced particles
 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 + np.random.rand(*points_t.shape) * radius * jitter
+    points_t = points_t + rng.random(points_t.shape) * radius * jitter
 
     return points_t.reshape((-1, 3))
 
@@ -180,6 +287,7 @@ def triangulate(face_counts, face_indices):
     return tri_indices
 
 
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
 def test_mesh_query_point(test, device):
     from pxr import Usd, UsdGeom
 
@@ -195,7 +303,7 @@ def test_mesh_query_point(test, device):
     mesh_indices = wp.array(np.array(tri_indices), dtype=int, device=device)
 
     # create mesh
-    mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices)
+    mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices, support_winding_number=True)
 
     p = particle_grid(32, 32, 32, np.array([-1.1, -1.1, -1.1]), 0.05, 0.0)
 
@@ -206,6 +314,17 @@ def test_mesh_query_point(test, device):
     faces_query = wp.zeros(query_count, dtype=int, device=device)
     dist_query = wp.zeros(query_count, dtype=float, device=device)
 
+    faces_query_no_sign = wp.zeros(query_count, dtype=int, device=device)
+    dist_query_no_sign = wp.zeros(query_count, dtype=float, device=device)
+
+    signs_query_normal = wp.zeros(query_count, dtype=float, device=device)
+    faces_query_normal = wp.zeros(query_count, dtype=int, device=device)
+    dist_query_normal = wp.zeros(query_count, dtype=float, device=device)
+
+    signs_query_winding_number = wp.zeros(query_count, dtype=float, device=device)
+    faces_query_winding_number = wp.zeros(query_count, dtype=int, device=device)
+    dist_query_winding_number = wp.zeros(query_count, dtype=float, device=device)
+
     signs_brute = wp.zeros(query_count, dtype=float, device=device)
     faces_brute = wp.zeros(query_count, dtype=int, device=device)
     dist_brute = wp.zeros(query_count, dtype=float, device=device)
@@ -216,6 +335,34 @@ def test_mesh_query_point(test, device):
         inputs=[mesh.id, query_points, faces_query, signs_query, dist_query],
         device=device,
     )
+
+    wp.launch(
+        kernel=sample_mesh_query_no_sign,
+        dim=query_count,
+        inputs=[mesh.id, query_points, faces_query_no_sign, dist_query_no_sign],
+        device=device,
+    )
+
+    wp.launch(
+        kernel=sample_mesh_query_sign_normal,
+        dim=query_count,
+        inputs=[mesh.id, query_points, faces_query_normal, signs_query_normal, dist_query_normal],
+        device=device,
+    )
+
+    wp.launch(
+        kernel=sample_mesh_query_sign_winding_number,
+        dim=query_count,
+        inputs=[
+            mesh.id,
+            query_points,
+            faces_query_winding_number,
+            signs_query_winding_number,
+            dist_query_winding_number,
+        ],
+        device=device,
+    )
+
     wp.launch(
         kernel=sample_mesh_brute,
         dim=query_count,
@@ -235,23 +382,44 @@ def test_mesh_query_point(test, device):
     faces_query = faces_query.numpy()
     dist_query = dist_query.numpy()
 
+    faces_query_no_sign = faces_query_no_sign.numpy()
+    dist_query_no_sign = dist_query_no_sign.numpy()
+
+    signs_query_normal = signs_query_normal.numpy()
+    faces_query_normal = faces_query_normal.numpy()
+    dist_query_normal = dist_query_normal.numpy()
+
+    signs_query_winding_number = signs_query_winding_number.numpy()
+    faces_query_winding_number = faces_query_winding_number.numpy()
+    dist_query_winding_number = dist_query_winding_number.numpy()
+
     signs_brute = signs_brute.numpy()
     faces_brute = faces_brute.numpy()
     dist_brute = dist_brute.numpy()
 
     query_points = query_points.numpy()
 
-    inside_query = []
-    inside_brute = []
+    inside_query = [[0.0, 0.0, 0.0]]
+    inside_query_normal = [[0.0, 0.0, 0.0]]
+    inside_query_winding_number = [[0.0, 0.0, 0.0]]
+    inside_brute = [[0.0, 0.0, 0.0]]
 
     for i in range(query_count):
         if signs_query[i] < 0.0:
             inside_query.append(query_points[i].tolist())
 
-        if signs_brute[i] > 6.0:
+        if signs_query_normal[i] < 0.0:
+            inside_query_normal.append(query_points[i].tolist())
+
+        if signs_query_winding_number[i] < 0.0:
+            inside_query_winding_number.append(query_points[i].tolist())
+
+        if signs_brute[i] > math.pi * 2.0:
             inside_brute.append(query_points[i].tolist())
 
     inside_query = np.array(inside_query)
+    inside_query_normal = np.array(inside_query_normal)
+    inside_query_winding_number = np.array(inside_query_winding_number)
     inside_brute = np.array(inside_brute)
 
     # import warp.render
@@ -269,13 +437,43 @@ def test_mesh_query_point(test, device):
     # stage.save()
 
     test.assertTrue(len(inside_query) == len(inside_brute))
+    test.assertTrue(len(inside_query_normal) == len(inside_brute))
+    test.assertTrue(len(inside_query_winding_number) == len(inside_brute))
 
+    tolerance = 1.5e-4
     dist_error = np.max(np.abs(dist_query - dist_brute))
     sign_error = np.max(np.abs(inside_query - inside_brute))
 
-    tolerance = 1.5e-4
-    test.assertTrue(dist_error < tolerance, f"dist_error is {dist_error} which is >= {tolerance}")
-    test.assertTrue(sign_error < tolerance, f"sign_error is {sign_error} which is >= {tolerance}")
+    test.assertTrue(dist_error < tolerance, f"mesh_query_point dist_error is {dist_error} which is >= {tolerance}")
+    test.assertTrue(sign_error < tolerance, f"mesh_query_point sign_error is {sign_error} which is >= {tolerance}")
+
+    dist_error = np.max(np.abs(dist_query_no_sign - dist_brute))
+
+    test.assertTrue(
+        dist_error < tolerance, f"mesh_query_point_no_sign dist_error is {dist_error} which is >= {tolerance}"
+    )
+
+    dist_error = np.max(np.abs(dist_query_normal - dist_brute))
+    sign_error = np.max(np.abs(inside_query_normal - inside_brute))
+
+    test.assertTrue(
+        dist_error < tolerance, f"mesh_query_point_sign_normal dist_error is {dist_error} which is >= {tolerance}"
+    )
+    test.assertTrue(
+        sign_error < tolerance, f"mesh_query_point_sign_normal sign_error is {sign_error} which is >= {tolerance}"
+    )
+
+    dist_error = np.max(np.abs(dist_query_winding_number - dist_brute))
+    sign_error = np.max(np.abs(inside_query_winding_number - inside_brute))
+
+    test.assertTrue(
+        dist_error < tolerance,
+        f"mesh_query_point_sign_winding_number dist_error is {dist_error} which is >= {tolerance}",
+    )
+    test.assertTrue(
+        sign_error < tolerance,
+        f"mesh_query_point_sign_winding_number sign_error is {sign_error} which is >= {tolerance}",
+    )
 
 
 @wp.kernel
@@ -305,6 +503,7 @@ def mesh_query_point_loss(
     loss[tid] = dist
 
 
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
 def test_adj_mesh_query_point(test, device):
     from pxr import Usd, UsdGeom
 
@@ -325,7 +524,7 @@ def test_adj_mesh_query_point(test, device):
     # mesh_indices = wp.array(np.array([0,1,2]), dtype=int, device=device)
 
     # create mesh
-    mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices)
+    mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices, support_winding_number=True)
 
     # p = particle_grid(32, 32, 32, np.array([-5.0, -5.0, -5.0]), 0.1, 0.1)*100.0
     p = wp.vec3(50.0, 50.0, 50.0)
@@ -384,28 +583,120 @@ def test_adj_mesh_query_point(test, device):
     test.assertTrue(error < tolerance, f"error is {error} which is >= {tolerance}")
 
 
-def register(parent):
-    devices = get_test_devices()
+@wp.kernel
+def sample_furthest_points(mesh: wp.uint64, query_points: wp.array(dtype=wp.vec3), query_result: wp.array(dtype=float)):
+    tid = wp.tid()
 
-    class TestMeshQuery(parent):
-        pass
+    p = query_points[tid]
 
-    # USD import failures should not count as a test failure
-    try:
-        from pxr import Usd, UsdGeom
+    face = int(0)
+    bary_u = float(0.0)
+    bary_v = float(0.0)
 
-        have_usd = True
-    except:
-        have_usd = False
+    if wp.mesh_query_furthest_point_no_sign(mesh, p, 0.0, face, bary_u, bary_v):
+        closest = wp.mesh_eval_position(mesh, face, bary_u, bary_v)
 
-    if have_usd:
-        add_function_test(TestMeshQuery, "test_mesh_query_point", test_mesh_query_point, devices=devices)
-        add_function_test(TestMeshQuery, "test_adj_mesh_query_point", test_adj_mesh_query_point, devices=devices)
+        query_result[tid] = wp.length_sq(p - closest)
 
-    return TestMeshQuery
+    query = wp.mesh_query_furthest_point_no_sign(mesh, p, 0.0)
+    wp.expect_eq(query.face, face)
+    wp.expect_eq(query.u, bary_u)
+    wp.expect_eq(query.v, bary_v)
 
 
-if __name__ == "__main__":
-    c = register(unittest.TestCase)
+@wp.kernel
+def sample_furthest_points_brute(
+    mesh_points: wp.array(dtype=wp.vec3), query_points: wp.array(dtype=wp.vec3), query_result: wp.array(dtype=float)
+):
+    tid = wp.tid()
+
+    p = query_points[tid]
+    max_dist_sq = float(0.0)
+
+    for i in range(mesh_points.shape[0]):
+        dist_sq = wp.length_sq(p - mesh_points[i])
+
+        if dist_sq > max_dist_sq:
+            max_dist_sq = dist_sq
+
+    query_result[tid] = max_dist_sq
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_mesh_query_furthest_point(test, device):
+    from pxr import Usd, UsdGeom
+
+    mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/spiky.usd")))
+    mesh_geom = UsdGeom.Mesh(mesh.GetPrimAtPath("/Cube/Cube"))
+
+    mesh_counts = mesh_geom.GetFaceVertexCountsAttr().Get()
+    mesh_indices = mesh_geom.GetFaceVertexIndicesAttr().Get()
+
+    tri_indices = triangulate(mesh_counts, mesh_indices)
+
+    mesh_points = wp.array(np.array(mesh_geom.GetPointsAttr().Get()), dtype=wp.vec3, device=device)
+    mesh_indices = wp.array(np.array(tri_indices), dtype=int, device=device)
+
+    # create mesh
+    mesh = wp.Mesh(points=mesh_points, indices=mesh_indices)
+
+    p = particle_grid(32, 32, 32, np.array([-1.1, -1.1, -1.1]), 0.05, 0.0)
 
+    query_count = len(p)
+    query_points = wp.array(p, dtype=wp.vec3, device=device)
+
+    dist_query = wp.zeros(query_count, dtype=float, device=device)
+    dist_brute = wp.zeros(query_count, dtype=float, device=device)
+
+    wp.launch(sample_furthest_points, dim=query_count, inputs=[mesh.id, query_points, dist_query], device=device)
+    wp.launch(
+        sample_furthest_points_brute, dim=query_count, inputs=[mesh_points, query_points, dist_brute], device=device
+    )
+
+    assert_np_equal(dist_query.numpy(), dist_brute.numpy(), tol=1.0e-3)
+
+
+def test_mesh_query_point_codegen_adjoints_with_select(test, device):
+    def kernel_fn(
+        mesh: wp.uint64,
+    ):
+        v = wp.vec3(0.0, 0.0, 0.0)
+        d = 1e-6
+
+        if True:
+            query_1 = wp.mesh_query_point(mesh, v, d)
+            query_2 = wp.mesh_query_point_no_sign(mesh, v, d)
+            query_3 = wp.mesh_query_furthest_point_no_sign(mesh, v, d)
+            query_4 = wp.mesh_query_point_sign_normal(mesh, v, d)
+            query_5 = wp.mesh_query_point_sign_winding_number(mesh, v, d)
+        else:
+            query_1 = wp.mesh_query_point(mesh, v, d)
+            query_2 = wp.mesh_query_point_no_sign(mesh, v, d)
+            query_3 = wp.mesh_query_furthest_point_no_sign(mesh, v, d)
+            query_4 = wp.mesh_query_point_sign_normal(mesh, v, d)
+            query_5 = wp.mesh_query_point_sign_winding_number(mesh, v, d)
+
+    wp.Kernel(func=kernel_fn)
+
+
+devices = get_test_devices()
+
+
+class TestMeshQueryPoint(unittest.TestCase):
+    pass
+
+
+add_function_test(TestMeshQueryPoint, "test_mesh_query_point", test_mesh_query_point, devices=devices)
+add_function_test(TestMeshQueryPoint, "test_mesh_query_furthest_point", test_mesh_query_furthest_point, devices=devices)
+add_function_test(TestMeshQueryPoint, "test_adj_mesh_query_point", test_adj_mesh_query_point, devices=devices)
+add_function_test(
+    TestMeshQueryPoint,
+    "test_mesh_query_point_codegen_adjoints_with_select",
+    test_mesh_query_point_codegen_adjoints_with_select,
+    devices=devices,
+)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_mesh_query_ray.py

@@ -5,14 +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 math
 
 import warp as wp
-from warp.tests.test_base import *
-
-np.random.seed(42)
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -66,9 +64,18 @@ def mesh_query_ray_loss(
     l = q[0]
     loss[tid] = l
 
+    query = wp.mesh_query_ray(mesh, p, D, max_t)
+    wp.expect_eq(query.t, t)
+    wp.expect_eq(query.u, bary_u)
+    wp.expect_eq(query.v, bary_v)
+    wp.expect_eq(query.sign, sign)
+    wp.expect_eq(query.normal, normal)
+    wp.expect_eq(query.face, face_index)
+
 
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
 def test_mesh_query_ray_grad(test, device):
-    from pxr import Usd, UsdGeom, Gf, Sdf
+    from pxr import Usd, UsdGeom
 
     # test tri
     # print("Testing Single Triangle")
@@ -210,7 +217,7 @@ def raycast_kernel(
     sign = float(0.0)  # hit face sign
     n = wp.vec3()  # hit face normal
     f = int(0)  # hit face index
-    max_dist = 1e6  # max raycast disance
+    max_dist = 1e6  # max raycast distance
 
     # ray cast against the mesh
     tid = wp.tid()
@@ -259,28 +266,38 @@ def test_mesh_query_ray_edge(test, device):
     test.assertEqual(counts.numpy()[0], n)
 
 
-def register(parent):
-    devices = get_test_devices()
+def test_mesh_query_codegen_adjoints_with_select(test, device):
+    def kernel_fn(
+        mesh: wp.uint64,
+    ):
+        v = wp.vec3(0.0, 0.0, 0.0)
+        d = 1e-6
+
+        if True:
+            query = wp.mesh_query_ray(mesh, v, v, d)
+        else:
+            query = wp.mesh_query_ray(mesh, v, v, d)
+
+    wp.Kernel(func=kernel_fn)
 
-    class TestMeshQueryRay(parent):
-        pass
 
-    add_function_test(TestMeshQueryRay, "test_mesh_query_ray_edge", test_mesh_query_ray_edge, devices=devices)
+devices = get_test_devices()
 
-    # USD import failures should not count as a test failure
-    try:
-        from pxr import Usd, UsdGeom
 
-        have_usd = True
-    except:
-        have_usd = False
+class TestMeshQueryRay(unittest.TestCase):
+    pass
 
-    if have_usd:
-        add_function_test(TestMeshQueryRay, "test_mesh_query_ray_grad", test_mesh_query_ray_grad, devices=devices)
 
-    return TestMeshQueryRay
+add_function_test(TestMeshQueryRay, "test_mesh_query_ray_edge", test_mesh_query_ray_edge, devices=devices)
+add_function_test(TestMeshQueryRay, "test_mesh_query_ray_grad", test_mesh_query_ray_grad, devices=devices)
+add_function_test(
+    TestMeshQueryRay,
+    "test_mesh_query_codegen_adjoints_with_select",
+    test_mesh_query_codegen_adjoints_with_select,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)