warp-lang 0.10.1__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,12 @@ 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(
@@ -69,6 +74,11 @@ def sample_mesh_query_no_sign(
     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(
@@ -97,6 +107,12 @@ def sample_mesh_query_sign_normal(
     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(
@@ -125,6 +141,12 @@ def sample_mesh_query_sign_winding_number(
     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):
@@ -237,9 +259,11 @@ def sample_mesh_brute(
 
 # 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))
 
@@ -263,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
 
@@ -478,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
 
@@ -557,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)

warp/tests/test_mlp.py

@@ -5,9 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
+
 import numpy as np
+
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -35,17 +38,17 @@ def loss_kernel(x: wp.array2d(dtype=float), loss: wp.array(dtype=float)):
 
 
 def test_mlp(test, device):
-    np.random.seed(0)
+    rng = np.random.default_rng(123)
 
     m = 10
     n = 200
 
     batches = 20000
 
-    weights = wp.array(np.random.rand(m, n) * 0.5 - 0.5, dtype=float, device=device)
-    bias = wp.array(np.random.rand(m) * 0.5 - 0.5, dtype=float, device=device)
+    weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device)
+    bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device)
 
-    x = wp.array(np.random.rand(n, batches), dtype=float, device=device)
+    x = wp.array(rng.random(size=(n, batches)), dtype=float, device=device)
     y = wp.zeros(shape=(m, batches), device=device)
 
     with wp.ScopedTimer("warp", active=False):
@@ -86,13 +89,15 @@ def create_mlp(m, n):
 def create_golden():
     import torch
 
+    rng = np.random.default_rng(123)
+
     input_size = 32
     hidden_size = 16
     batch_size = 64
 
     network = create_mlp(input_size, hidden_size)
 
-    x = torch.Tensor(np.random.rand(batch_size, input_size))
+    x = torch.Tensor(rng.random(size=(batch_size, input_size)))
     x.requires_grad = True
 
     y = network.forward(x)
@@ -169,6 +174,8 @@ def test_mlp_grad(test, device):
 def profile_mlp_torch(device):
     import torch
 
+    rng = np.random.default_rng(123)
+
     m = 128
     n = 64
 
@@ -179,7 +186,7 @@ def profile_mlp_torch(device):
 
         network = create_mlp(m, n)
 
-        x = torch.Tensor(np.random.rand(b, m))
+        x = torch.Tensor(rng.random(size=(b, m)))
 
         with wp.ScopedTimer("torch_forward" + str(b)):
             y = network.forward(x)
@@ -190,7 +197,7 @@ def profile_mlp_torch(device):
 
         network = create_mlp(m, n)
 
-        x = torch.Tensor(np.random.rand(b, m))
+        x = torch.Tensor(rng.random(size=(b, m)))
         y = network.forward(x)
 
         loss = torch.norm(y)
@@ -204,6 +211,8 @@ def profile_mlp_torch(device):
 
 
 def profile_mlp_warp(device):
+    rng = np.random.default_rng(123)
+
     m = 128
     n = 64
 
@@ -212,10 +221,10 @@ def profile_mlp_warp(device):
     for i in range(steps):
         b = 2**i
 
-        weights = wp.array(np.random.rand(m, n) * 0.5 - 0.5, dtype=float, device=device)
-        bias = wp.array(np.random.rand(m) * 0.5 - 0.5, dtype=float, device=device)
+        weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device)
+        bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device)
 
-        x = wp.array(np.random.rand(n, b), dtype=float, device=device)
+        x = wp.array(rng.random(size=(n, b)), dtype=float, device=device)
         y = wp.zeros(shape=(m, b), device=device)
 
         with wp.ScopedTimer("warp-forward" + str(b)):
@@ -225,10 +234,10 @@ def profile_mlp_warp(device):
     for i in range(steps):
         b = 2**i
 
-        weights = wp.array(np.random.rand(m, n) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
-        bias = wp.array(np.random.rand(m) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
+        weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
+        bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
 
-        x = wp.array(np.random.rand(n, b), dtype=float, device=device, requires_grad=True)
+        x = wp.array(rng.random(size=(n, b)), dtype=float, device=device, requires_grad=True)
         y = wp.zeros(shape=(m, b), device=device, requires_grad=True)
 
         loss = wp.zeros(1, dtype=float, device=device)
@@ -251,18 +260,17 @@ def profile_mlp_warp(device):
 # profile_mlp_torch("cuda")
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestMLP(parent):
-        pass
+class TestMLP(unittest.TestCase):
+    pass
 
-    add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices)
-    add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices)
 
-    return TestMLP
+add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices)
+add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_model.py

@@ -5,103 +5,106 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import warp as wp
-from warp.tests.test_base import *
-from warp.sim import ModelBuilder
+import unittest
 
 import numpy as np
 
+import warp as wp
+from warp.sim import ModelBuilder
+from warp.tests.unittest_utils import *
+
 wp.init()
 
 
-def register(parent):
-    class TestModel(parent):
-        def test_add_triangles(self):
-            pts = np.array(
-                [
-                    [-0.00585869, 0.34189449, -1.17415233],
-                    [-1.894547, 0.1788074, 0.9251329],
-                    [-1.26141048, 0.16140787, 0.08823282],
-                    [-0.08609255, -0.82722546, 0.65995427],
-                    [0.78827592, -1.77375711, -0.55582718],
-                ]
-            )
-            tris = np.array([[0, 3, 4], [0, 2, 3], [2, 1, 3], [1, 4, 3]])
-
-            builder1 = ModelBuilder()
-            builder2 = ModelBuilder()
-            for pt in pts:
-                builder1.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-                builder2.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-
-            # test add_triangle(s) with default arguments:
-            areas = builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2])
-            for i, t in enumerate(tris):
-                area = builder1.add_triangle(t[0], t[1], t[2])
-                self.assertAlmostEqual(area, areas[i], places=6)
-
-            # test add_triangle(s) with non default arguments:
-            tri_ke = np.random.randn(pts.shape[0])
-            tri_ka = np.random.randn(pts.shape[0])
-            tri_kd = np.random.randn(pts.shape[0])
-            tri_drag = np.random.randn(pts.shape[0])
-            tri_lift = np.random.randn(pts.shape[0])
-            for i, t in enumerate(tris):
-                builder1.add_triangle(
-                    t[0],
-                    t[1],
-                    t[2],
-                    tri_ke[i],
-                    tri_ka[i],
-                    tri_kd[i],
-                    tri_drag[i],
-                    tri_lift[i],
-                )
-            builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
-
-            assert_np_equal(np.array(builder1.tri_indices), np.array(builder2.tri_indices))
-            assert_np_equal(np.array(builder1.tri_poses), np.array(builder2.tri_poses), tol=1.0e-6)
-            assert_np_equal(np.array(builder1.tri_activations), np.array(builder2.tri_activations))
-            assert_np_equal(np.array(builder1.tri_materials), np.array(builder2.tri_materials))
-
-        def test_add_edges(self):
-            pts = np.array(
-                [
-                    [-0.00585869, 0.34189449, -1.17415233],
-                    [-1.894547, 0.1788074, 0.9251329],
-                    [-1.26141048, 0.16140787, 0.08823282],
-                    [-0.08609255, -0.82722546, 0.65995427],
-                    [0.78827592, -1.77375711, -0.55582718],
-                ]
+class TestModel(unittest.TestCase):
+    def test_add_triangles(self):
+        rng = np.random.default_rng(123)
+
+        pts = np.array(
+            [
+                [-0.00585869, 0.34189449, -1.17415233],
+                [-1.894547, 0.1788074, 0.9251329],
+                [-1.26141048, 0.16140787, 0.08823282],
+                [-0.08609255, -0.82722546, 0.65995427],
+                [0.78827592, -1.77375711, -0.55582718],
+            ]
+        )
+        tris = np.array([[0, 3, 4], [0, 2, 3], [2, 1, 3], [1, 4, 3]])
+
+        builder1 = ModelBuilder()
+        builder2 = ModelBuilder()
+        for pt in pts:
+            builder1.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+            builder2.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+
+        # test add_triangle(s) with default arguments:
+        areas = builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2])
+        for i, t in enumerate(tris):
+            area = builder1.add_triangle(t[0], t[1], t[2])
+            self.assertAlmostEqual(area, areas[i], places=6)
+
+        # test add_triangle(s) with non default arguments:
+        tri_ke = rng.standard_normal(size=pts.shape[0])
+        tri_ka = rng.standard_normal(size=pts.shape[0])
+        tri_kd = rng.standard_normal(size=pts.shape[0])
+        tri_drag = rng.standard_normal(size=pts.shape[0])
+        tri_lift = rng.standard_normal(size=pts.shape[0])
+        for i, t in enumerate(tris):
+            builder1.add_triangle(
+                t[0],
+                t[1],
+                t[2],
+                tri_ke[i],
+                tri_ka[i],
+                tri_kd[i],
+                tri_drag[i],
+                tri_lift[i],
             )
-            edges = np.array([[0, 4, 3, 1], [3, 2, 4, 1]])
-
-            builder1 = ModelBuilder()
-            builder2 = ModelBuilder()
-            for pt in pts:
-                builder1.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-                builder2.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-
-            # test defaults:
-            for i in range(2):
-                builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3])
-            builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3])
-
-            # test non defaults:
-            rest = np.random.randn(2)
-            edge_ke = np.random.randn(2)
-            edge_kd = np.random.randn(2)
-            for i in range(2):
-                builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3], rest[i], edge_ke[i], edge_kd[i])
-            builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3], rest, edge_ke, edge_kd)
-
-            assert_np_equal(np.array(builder1.edge_indices), np.array(builder2.edge_indices))
-            assert_np_equal(np.array(builder1.edge_rest_angle), np.array(builder2.edge_rest_angle), tol=1.0e-4)
-            assert_np_equal(np.array(builder1.edge_bending_properties), np.array(builder2.edge_bending_properties))
-
-    return TestModel
+        builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
+
+        assert_np_equal(np.array(builder1.tri_indices), np.array(builder2.tri_indices))
+        assert_np_equal(np.array(builder1.tri_poses), np.array(builder2.tri_poses), tol=1.0e-6)
+        assert_np_equal(np.array(builder1.tri_activations), np.array(builder2.tri_activations))
+        assert_np_equal(np.array(builder1.tri_materials), np.array(builder2.tri_materials))
+
+    def test_add_edges(self):
+        rng = np.random.default_rng(123)
+
+        pts = np.array(
+            [
+                [-0.00585869, 0.34189449, -1.17415233],
+                [-1.894547, 0.1788074, 0.9251329],
+                [-1.26141048, 0.16140787, 0.08823282],
+                [-0.08609255, -0.82722546, 0.65995427],
+                [0.78827592, -1.77375711, -0.55582718],
+            ]
+        )
+        edges = np.array([[0, 4, 3, 1], [3, 2, 4, 1]])
+
+        builder1 = ModelBuilder()
+        builder2 = ModelBuilder()
+        for pt in pts:
+            builder1.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+            builder2.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+
+        # test defaults:
+        for i in range(2):
+            builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3])
+        builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3])
+
+        # test non defaults:
+        rest = rng.standard_normal(size=2)
+        edge_ke = rng.standard_normal(size=2)
+        edge_kd = rng.standard_normal(size=2)
+        for i in range(2):
+            builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3], rest[i], edge_ke[i], edge_kd[i])
+        builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3], rest, edge_ke, edge_kd)
+
+        assert_np_equal(np.array(builder1.edge_indices), np.array(builder2.edge_indices))
+        assert_np_equal(np.array(builder1.edge_rest_angle), np.array(builder2.edge_rest_angle), tol=1.0e-4)
+        assert_np_equal(np.array(builder1.edge_bending_properties), np.array(builder2.edge_bending_properties))
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)