warp-lang 1.6.2__py3-none-win_amd64.whl → 1.7.1__py3-none-win_amd64.whl

warp/tests/{test_mesh_query_point.py → geometry/test_mesh_query_point.py}

@@ -298,7 +298,7 @@ def triangulate(face_counts, face_indices):
 def test_mesh_query_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 = 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()
@@ -526,7 +526,7 @@ def mesh_query_point_loss(
 def test_adj_mesh_query_point(test, device):
     from pxr import Usd, UsdGeom
 
-    mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/torus.usda")))
+    mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "torus.usda")))
     mesh_geom = UsdGeom.Mesh(mesh.GetPrimAtPath("/World/Torus"))
 
     mesh_counts = mesh_geom.GetFaceVertexCountsAttr().Get()
@@ -663,7 +663,7 @@ def sample_furthest_points_brute(
 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 = 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()
@@ -750,11 +750,11 @@ def triangle_closest_point_for_test(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.ve
     return a + v * ab + w * ac, bary
 
 
-def load_mesh():
+def load_mesh(model_name="bunny"):
     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"))
+    usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), model_name + ".usd"))
+    usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/" + model_name))
 
     vertices = np.array(usd_geom.GetPointsAttr().Get())
     faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
@@ -820,76 +820,79 @@ def test_set_mesh_points(test, device):
 
     rng = np.random.default_rng(123)
 
-    vs, fs = load_mesh()
-    vertices1 = wp.array(vs, dtype=wp.vec3, device=device)
-    velocities1_np = rng.standard_normal(size=(vertices1.shape[0], 3))
-    velocities1 = wp.array(velocities1_np, dtype=wp.vec3, device=device)
+    models = ["bunny", "nonuniform"]
 
-    faces = wp.array(fs, dtype=wp.int32, device=device)
+    for model in models:
+        vs, fs = load_mesh(model)
+        vertices1 = wp.array(vs, dtype=wp.vec3, device=device)
+        velocities1_np = rng.standard_normal(size=(vertices1.shape[0], 3))
+        velocities1 = wp.array(velocities1_np, dtype=wp.vec3, device=device)
 
-    n = 1000
-    query_radius = 0.2
-    pts1 = wp.array(rng.standard_normal(size=(n, 3)), dtype=wp.vec3, device=device)
+        faces = wp.array(fs, dtype=wp.int32, device=device)
 
-    query_results_num_cols1 = wp.zeros(n, dtype=wp.int32, device=device)
-    query_results_min_dist1 = wp.zeros(n, dtype=float, device=device)
-    query_results_closest_point_velocity1 = wp.zeros(n, dtype=wp.vec3, device=device)
+        n = 1000
+        query_radius = 0.2
+        pts1 = wp.array(rng.standard_normal(size=(n, 3)), dtype=wp.vec3, device=device)
 
-    for constructor in constructors:
-        mesh = wp.Mesh(vertices1, faces, velocities=velocities1, bvh_constructor=constructor)
-        fs_2D = faces.reshape((-1, 3))
+        query_results_num_cols1 = wp.zeros(n, dtype=wp.int32, device=device)
+        query_results_min_dist1 = wp.zeros(n, dtype=float, device=device)
+        query_results_closest_point_velocity1 = wp.zeros(n, dtype=wp.vec3, device=device)
 
-        wp.launch(
-            kernel=point_query_aabb_and_closest,
-            inputs=[
-                query_radius,
-                mesh.id,
-                pts1,
-                vertices1,
-                fs_2D,
-                query_results_num_cols1,
-                query_results_min_dist1,
-                query_results_closest_point_velocity1,
-            ],
-            dim=n,
-            device=device,
-        )
+        for constructor in constructors:
+            mesh = wp.Mesh(vertices1, faces, velocities=velocities1, bvh_constructor=constructor)
+            fs_2D = faces.reshape((-1, 3))
+
+            wp.launch(
+                kernel=point_query_aabb_and_closest,
+                inputs=[
+                    query_radius,
+                    mesh.id,
+                    pts1,
+                    vertices1,
+                    fs_2D,
+                    query_results_num_cols1,
+                    query_results_min_dist1,
+                    query_results_closest_point_velocity1,
+                ],
+                dim=n,
+                device=device,
+            )
 
-        shift = rng.standard_normal(size=3)
+            shift = rng.standard_normal(size=3)
 
-        vs_higher = vs + shift
-        vertices2 = wp.array(vs_higher, dtype=wp.vec3, device=device)
+            vs_higher = vs + shift
+            vertices2 = wp.array(vs_higher, dtype=wp.vec3, device=device)
 
-        velocities2_np = velocities1_np + shift[None, ...]
-        velocities2 = wp.array(velocities2_np, dtype=wp.vec3, device=device)
+            velocities2_np = velocities1_np + shift[None, ...]
+            velocities2 = wp.array(velocities2_np, dtype=wp.vec3, device=device)
 
-        pts2 = wp.array(pts1.numpy() + shift, dtype=wp.vec3, device=device)
+            pts2 = wp.array(pts1.numpy() + shift, dtype=wp.vec3, device=device)
 
-        mesh.points = vertices2
-        mesh.velocities = velocities2
+            mesh.points = vertices2
+            mesh.velocities = velocities2
 
-        query_results_num_cols2 = wp.zeros(n, dtype=wp.int32, device=device)
-        query_results_min_dist2 = wp.zeros(n, dtype=float, device=device)
-        query_results_closest_point_velocity2 = wp.array([shift for i in range(n)], dtype=wp.vec3, device=device)
+            query_results_num_cols2 = wp.zeros(n, dtype=wp.int32, device=device)
+            query_results_min_dist2 = wp.zeros(n, dtype=float, device=device)
+            query_results_closest_point_velocity2 = wp.array([shift for i in range(n)], dtype=wp.vec3, device=device)
 
-        wp.launch(
-            kernel=point_query_aabb_and_closest,
-            inputs=[
-                query_radius,
-                mesh.id,
-                pts2,
-                vertices2,
-                fs_2D,
-                query_results_num_cols2,
-                query_results_min_dist2,
-                query_results_closest_point_velocity2,
-            ],
-            dim=n,
-            device=device,
-        )
+            wp.launch(
+                kernel=point_query_aabb_and_closest,
+                inputs=[
+                    query_radius,
+                    mesh.id,
+                    pts2,
+                    vertices2,
+                    fs_2D,
+                    query_results_num_cols2,
+                    query_results_min_dist2,
+                    query_results_closest_point_velocity2,
+                ],
+                dim=n,
+                device=device,
+            )
 
-        test.assertTrue((query_results_num_cols1.numpy() == query_results_num_cols2.numpy()).all())
-        test.assertTrue(((query_results_min_dist1.numpy() - query_results_min_dist2.numpy()) < 1e-5).all())
+            test.assertTrue((query_results_num_cols1.numpy() == query_results_num_cols2.numpy()).all())
+            test.assertTrue(((query_results_min_dist1.numpy() - query_results_min_dist2.numpy()) < 1e-5).all())
 
 
 devices = get_test_devices()

warp/tests/{test_mesh_query_ray.py → geometry/test_mesh_query_ray.py}

@@ -88,7 +88,7 @@ def test_mesh_query_ray_grad(test, device):
     # mesh_points = wp.array(np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0]]), dtype=wp.vec3, device=device)
     # mesh_indices = wp.array(np.array([0,1,2]), dtype=int, device=device)
 
-    mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/torus.usda")))
+    mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "torus.usda")))
     mesh_geom = UsdGeom.Mesh(mesh.GetPrimAtPath("/World/Torus"))
 
     mesh_counts = mesh_geom.GetFaceVertexCountsAttr().Get()

warp/tests/{test_volume.py → geometry/test_volume.py}

@@ -370,6 +370,22 @@ def test_volume_store_i(volume: wp.uint64, points: wp.array(dtype=wp.vec3), valu
     values[tid] = wp.volume_lookup_i(volume, i, j, k)
 
 
+@wp.kernel
+def test_volume_store_v4(volume: wp.uint64, points: wp.array(dtype=wp.vec3), values: wp.array(dtype=wp.vec4)):
+    tid = wp.tid()
+
+    p = points[tid]
+    i = int(p[0])
+    j = int(p[1])
+    k = int(p[2])
+
+    v = wp.vec4(p[0], p[1], p[2], float(i + 100 * j + 10000 * k))
+
+    wp.volume_store(volume, i, j, k, v)
+
+    values[tid] = wp.volume_lookup(volume, i, j, k, dtype=wp.vec4)
+
+
 devices = get_test_devices()
 rng = np.random.default_rng(101215)
 
@@ -393,12 +409,12 @@ rng = np.random.default_rng(101215)
 #   (-90 degrees rotation along X)
 #   voxel size: 0.1
 volume_paths = {
-    "float": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_grid.nvdb")),
-    "int32": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_int32_grid.nvdb")),
-    "vec3f": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_vec_grid.nvdb")),
-    "index": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_index_grid.nvdb")),
-    "torus": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/torus.nvdb")),
-    "float_write": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_grid.nvdb")),
+    "float": os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "test_grid.nvdb")),
+    "int32": os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "test_int32_grid.nvdb")),
+    "vec3f": os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "test_vec_grid.nvdb")),
+    "index": os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "test_index_grid.nvdb")),
+    "torus": os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "torus.nvdb")),
+    "float_write": os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "test_grid.nvdb")),
 }
 
 test_volume_tiles = (
@@ -635,6 +651,22 @@ def test_volume_allocation_i(test, device):
     np.testing.assert_equal(values_res, values_ref)
 
 
+def test_volume_allocation_v4(test, device):
+    bg_value = (-1, 2.0, -3, 5)
+    points_np = np.append(point_grid, [[8096, 8096, 8096]], axis=0)
+
+    w_ref = np.array([x + 100 * y + 10000 * z for x, y, z in point_grid])[:, np.newaxis]
+    values_ref = np.append(np.hstack((point_grid, w_ref)), [bg_value], axis=0)
+
+    volume = wp.Volume.allocate(min=[-11, -11, -11], max=[11, 11, 11], voxel_size=0.1, bg_value=bg_value, device=device)
+    points = wp.array(points_np, dtype=wp.vec3, device=device)
+    values = wp.empty(len(points_np), dtype=wp.vec4, device=device)
+    wp.launch(test_volume_store_v4, dim=len(points_np), inputs=[volume.id, points, values], device=device)
+
+    values_res = values.numpy()
+    np.testing.assert_equal(values_res, values_ref)
+
+
 def test_volume_introspection(test, device):
     for volume_names in ("float", "vec3f"):
         with test.subTest(volume_names=volume_names):
@@ -967,6 +999,9 @@ add_function_test(
 add_function_test(
     TestVolume, "test_volume_allocation_i", test_volume_allocation_i, devices=get_selected_cuda_test_devices()
 )
+add_function_test(
+    TestVolume, "test_volume_allocation_v4", test_volume_allocation_v4, devices=get_selected_cuda_test_devices()
+)
 add_function_test(TestVolume, "test_volume_introspection", test_volume_introspection, devices=devices)
 add_function_test(
     TestVolume, "test_volume_from_numpy", test_volume_from_numpy, devices=get_selected_cuda_test_devices()

warp/tests/{test_dlpack.py → interop/test_dlpack.py}

@@ -389,11 +389,23 @@ def test_dlpack_paddle_to_warp(test, device):
 def test_dlpack_warp_to_jax(test, device):
     import jax
     import jax.dlpack
+    import jax.numpy as jnp
 
-    a = wp.array(data=np.arange(N, dtype=np.float32), device=device)
+    cpu_device = jax.devices("cpu")[0]
+
+    # Create a numpy array from a JAX array to respect XLA alignment needs
+    with jax.default_device(cpu_device):
+        x_jax = jnp.arange(N, dtype=jnp.float32)
+        x_numpy = np.asarray(x_jax)
+        test.assertEqual(x_jax.unsafe_buffer_pointer(), np.lib.array_utils.byte_bounds(x_numpy)[0])
+
+    a = wp.array(x_numpy, device=device, dtype=wp.float32, copy=False)
+
+    if device.is_cpu:
+        test.assertEqual(a.ptr, np.lib.array_utils.byte_bounds(x_numpy)[0])
 
     # use generic dlpack conversion
-    j1 = jax.dlpack.from_dlpack(wp.to_dlpack(a))
+    j1 = jax.dlpack.from_dlpack(a, copy=False)
 
     # use jax wrapper
     j2 = wp.to_jax(a)
@@ -423,14 +435,25 @@ def test_dlpack_warp_to_jax(test, device):
 @unittest.skipUnless(_jax_version() >= (0, 4, 15), "Jax version too old")
 def test_dlpack_warp_to_jax_v2(test, device):
     # same as original test, but uses newer __dlpack__() method
-
     import jax
     import jax.dlpack
+    import jax.numpy as jnp
 
-    a = wp.array(data=np.arange(N, dtype=np.float32), device=device)
+    cpu_device = jax.devices("cpu")[0]
+
+    # Create a numpy array from a JAX array to respect XLA alignment needs
+    with jax.default_device(cpu_device):
+        x_jax = jnp.arange(N, dtype=jnp.float32)
+        x_numpy = np.asarray(x_jax)
+        test.assertEqual(x_jax.unsafe_buffer_pointer(), np.lib.array_utils.byte_bounds(x_numpy)[0])
+
+    a = wp.array(x_numpy, device=device, dtype=wp.float32, copy=False)
+
+    if device.is_cpu:
+        test.assertEqual(a.ptr, np.lib.array_utils.byte_bounds(x_numpy)[0])
 
     # pass warp array directly
-    j1 = jax.dlpack.from_dlpack(a)
+    j1 = jax.dlpack.from_dlpack(a, copy=False)
 
     # use jax wrapper
     j2 = wp.to_jax(a)

warp/tests/{disabled_kinematics.py → sim/disabled_kinematics.py}

@@ -17,6 +17,7 @@ import math
 import unittest
 
 import warp as wp
+import warp.examples
 import warp.sim
 from warp.tests.unittest_utils import *
 
@@ -24,11 +25,7 @@ from warp.tests.unittest_utils import *
 def build_ant(num_envs):
     builder = wp.sim.ModelBuilder()
     for i in range(num_envs):
-        wp.sim.parse_mjcf(
-            os.path.join(os.path.dirname(__file__), "../../examples/assets/nv_ant.xml"),
-            builder,
-            up_axis="y",
-        )
+        wp.sim.parse_mjcf(os.path.join(warp.examples.get_asset_directory(), "nv_ant.xml"), builder, up_axis="y")
 
         coord_count = 15
         dof_count = 14
@@ -37,8 +34,10 @@ def build_ant(num_envs):
         dof_start = i * dof_count
 
         # base
-        builder.joint_q[coord_start : coord_start + 3] = [i * 2.0, 0.70, 0.0]
-        builder.joint_q[coord_start + 3 : coord_start + 7] = wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)
+        p = [i * 2.0, 0.70, 0.0]
+        q = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
+        builder.joint_q[coord_start : coord_start + 3] = p
+        builder.joint_q[coord_start + 3 : coord_start + 7] = q
 
         # joints
         builder.joint_q[coord_start + 7 : coord_start + coord_count] = [0.0, 1.0, 0.0, -1.0, 0.0, -1.0, 0.0, 1.0]
@@ -56,9 +55,9 @@ def build_complex_joint_mechanism(chain_length):
     ax1 = wp.normalize(wp.vec3(4.0, -1.0, 2.0))
     ax2 = wp.normalize(wp.vec3(-3.0, 4.0, -1.0))
     # declare some transforms with nonzero translation and orientation
-    tf0 = wp.transform(wp.vec3(1.0, 2.0, 3.0), wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi * 0.25))
-    tf1 = wp.transform(wp.vec3(4.0, 5.0, 6.0), wp.quat_from_axis_angle((0.0, 1.0, 0.0), math.pi * 0.5))
-    tf2 = wp.transform(wp.vec3(7.0, 8.0, 9.0), wp.quat_from_axis_angle((0.0, 0.0, 1.0), math.pi * 0.75))
+    tf0 = wp.transform(wp.vec3(1.0, 2.0, 3.0), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.25))
+    tf1 = wp.transform(wp.vec3(4.0, 5.0, 6.0), wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.5))
+    tf2 = wp.transform(wp.vec3(7.0, 8.0, 9.0), wp.quat_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), math.pi * 0.75))
 
     parent = -1
     for _i in range(chain_length):