warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/tests/test_examples.py

@@ -28,7 +28,7 @@ def add_example_test(cls, name, devices=None, options={}):
         test_options = options.copy()
 
         try:
-            module = importlib.import_module(f"examples.{name}")
+            module = importlib.import_module(f"warp.examples.{name}")
 
             torch_cuda_required = test_options.setdefault("torch_cuda_required", False)
             test_options.pop("torch_cuda_required", None)
@@ -45,9 +45,11 @@ def add_example_test(cls, name, devices=None, options={}):
         # create default USD stage output path which many examples expect
         test_options.setdefault(
             "stage",
-            os.path.join(os.path.dirname(__file__), f"outputs/{name}_{sanitize_identifier(device)}.usd")
-            if USD_AVAILABLE
-            else None,
+            (
+                os.path.join(os.path.dirname(__file__), f"outputs/{name}_{sanitize_identifier(device)}.usd")
+                if USD_AVAILABLE
+                else None
+            ),
         )
 
         if test_options["stage"]:
@@ -60,33 +62,34 @@ def add_example_test(cls, name, devices=None, options={}):
         test_options.pop("num_frames", None)
 
         # Don't want to force load all modules by default for serial test runner
-        wp.config.graph_capture_module_load_default = False
+        wp.config.enable_graph_capture_module_load_by_default = False
 
-        try:
-            enable_backward = test_options.get("enable_backward", True)
-            wp.set_module_options({"enable_backward": enable_backward}, module)
-            test_options.pop("enable_backward", None)
-
-            with wp.ScopedDevice(device):
-                wp.load_module(module, device=wp.get_device())
-                extra_load_modules = test_options.get("load_modules", [])
-                for module_name in extra_load_modules:
-                    wp.load_module(module_name, device=wp.get_device())
-                test_options.pop("load_modules", None)
-
-                e = module.Example(**test_options)
-
-                # disable scoped timer to avoid log spam from time steps
-                wp.ScopedTimer.enabled = False
-
-                for _ in range(num_frames):
-                    e.update()
-                    e.render()
-        except Exception as e:
-            test.fail(f"{e}")
-        finally:
-            wp.ScopedTimer.enabled = True
-            wp.config.graph_capture_module_load_default = True
+        with wp.ScopedTimer(f"{name}_{sanitize_identifier(device)}"):
+            try:
+                enable_backward = test_options.get("enable_backward", True)
+                wp.set_module_options({"enable_backward": enable_backward}, module)
+                test_options.pop("enable_backward", None)
+
+                with wp.ScopedDevice(device):
+                    wp.load_module(module, device=wp.get_device())
+                    extra_load_modules = test_options.get("load_modules", [])
+                    for module_name in extra_load_modules:
+                        wp.load_module(module_name, device=wp.get_device())
+                    test_options.pop("load_modules", None)
+
+                    e = module.Example(**test_options)
+
+                    # disable scoped timer to avoid log spam from time steps
+                    wp.ScopedTimer.enabled = False
+
+                    for _ in range(num_frames):
+                        e.step()
+                        e.render()
+            except Exception as e:
+                test.fail(f"{e}")
+            finally:
+                wp.ScopedTimer.enabled = True
+                wp.config.enable_graph_capture_module_load_by_default = True
 
     from warp.tests.unittest_utils import add_function_test
 
@@ -109,87 +112,176 @@ class TestExamples(unittest.TestCase):
 # Exclude unless we can run headless somehow
 # add_example_test(TestExamples, name="example_render_opengl", options={})
 
-add_example_test(TestExamples, name="example_dem", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_diffray", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_fluid", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_jacobian_ik", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_marching_cubes", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_mesh", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_mesh_intersect", devices=cuda_test_devices, options={"num_frames": 1})
-add_example_test(TestExamples, name="example_nvdb", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_raycast", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_raymarch", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_sph", devices=cuda_test_devices)
-add_example_test(TestExamples, name="example_wave", devices=cuda_test_devices, options={"resx": 256, "resy": 256})
+add_example_test(TestExamples, name="core.example_dem", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_fluid", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_marching_cubes", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_mesh", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_mesh_intersect", devices=cuda_test_devices, options={"num_frames": 1})
+add_example_test(TestExamples, name="core.example_nvdb", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_raycast", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_raymarch", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_sph", devices=cuda_test_devices)
+add_example_test(TestExamples, name="core.example_wave", devices=cuda_test_devices)
 
 
 class TestSimExamples(unittest.TestCase):
     pass
 
 
+warp_sim_modules = [
+    "warp.sim.integrator",
+    "warp.sim.integrator_euler",
+    "warp.sim.particles",
+    "warp.sim.collide",
+    "warp.sim.articulation",
+    "warp.sim.integrator_xpbd",
+    "warp.sim.integrator_featherstone",
+    "warp.sim.integrator_euler",
+    "warp.sim.integrator",
+    "warp.sim.utils",
+]
+
+add_example_test(
+    TestSimExamples,
+    name="optim.example_bounce",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_cartpole",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_cloth",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
 add_example_test(
     TestSimExamples,
-    name="example_sim_cartpole",
+    name="optim.example_cloth_throw",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.collide", "warp.sim.integrator_euler", "warp.sim.articulation"]},
+    options={"load_modules": warp_sim_modules},
 )
+add_example_test(TestSimExamples, name="optim.example_diffray", devices=cuda_test_devices)
 add_example_test(
     TestSimExamples,
-    name="example_sim_cloth",
+    name="optim.example_drone",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.collide", "warp.sim.integrator_euler", "warp.sim.particles"]},
+    options={
+        "load_modules": warp_sim_modules + ["warp.optim.sgd"],
+        "num_frames": 1,
+        "drone_path": os.path.join(os.path.dirname(__file__), "..", "examples", "assets", "crazyflie.usd"),
+    },
 )
-add_example_test(TestSimExamples, name="example_sim_fk_grad", devices=cuda_test_devices)
 add_example_test(
-    TestSimExamples, name="example_sim_fk_grad_torch", devices=cuda_test_devices, options={"torch_cuda_required": True}
+    TestSimExamples,
+    name="sim.example_granular",
+    devices=cuda_test_devices,
+    options={
+        "load_modules": warp_sim_modules,
+    },
 )
 add_example_test(
     TestSimExamples,
-    name="example_sim_grad_bounce",
+    name="sim.example_granular_collision_sdf",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.integrator_euler", "warp.sim.particles"]},
+    options={
+        "load_modules": warp_sim_modules,
+    },
 )
 add_example_test(
     TestSimExamples,
-    name="example_sim_grad_cloth",
+    name="optim.example_inverse_kinematics",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.integrator_euler", "warp.sim.particles"]},
+    options={
+        "load_modules": warp_sim_modules,
+    },
 )
-add_example_test(TestSimExamples, name="example_sim_granular", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_granular_collision_sdf", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_neo_hookean", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_particle_chain", devices=cuda_test_devices)
 add_example_test(
     TestSimExamples,
-    name="example_sim_quadruped",
+    name="optim.example_inverse_kinematics_torch",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.integrator_xpbd", "warp.sim.integrator_euler"]},
+    options={
+        "torch_cuda_required": True,
+        "load_modules": warp_sim_modules,
+    },
 )
+add_example_test(TestExamples, name="sim.example_jacobian_ik", devices=cuda_test_devices)
 add_example_test(
     TestSimExamples,
-    name="example_sim_rigid_chain",
+    name="sim.example_particle_chain",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.integrator_xpbd", "warp.sim.integrator_euler"]},
+    options={"load_modules": warp_sim_modules},
 )
 add_example_test(
     TestSimExamples,
-    name="example_sim_rigid_contact",
+    name="sim.example_quadruped",
     devices=cuda_test_devices,
-    options={"load_modules": ["warp.sim.integrator_euler"]},
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_rigid_chain",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_rigid_contact",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_rigid_soft_contact",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_rigid_force",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_rigid_gyroscopic",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="sim.example_soft_body",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="optim.example_spring_cage",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
+)
+add_example_test(
+    TestSimExamples,
+    name="optim.example_trajectory",
+    devices=cuda_test_devices,
+    options={"load_modules": warp_sim_modules},
 )
-add_example_test(TestSimExamples, name="example_sim_rigid_fem", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_rigid_force", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_rigid_gyroscopic", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_rigid_kinematics", devices=cuda_test_devices)
-add_example_test(TestSimExamples, name="example_sim_trajopt", devices=cuda_test_devices)
 
 
 class TestFemExamples(unittest.TestCase):
     pass
 
 
+class TestFemDiffusionExamples(unittest.TestCase):
+    pass
+
+
 add_example_test(
-    TestFemExamples,
+    TestFemDiffusionExamples,
     name="fem.example_diffusion_mgpu",
     devices=cuda_test_devices,
     options={"quiet": True, "num_frames": 1, "enable_backward": False},
@@ -203,13 +295,13 @@ add_example_test(
     options={"quiet": True, "res": [16, 16, 16], "enable_backward": False},
 )
 add_example_test(
-    TestFemExamples,
+    TestFemDiffusionExamples,
     name="fem.example_diffusion",
     devices=test_devices,
     options={"quiet": True, "resolution": 10, "mesh": "tri", "num_frames": 1, "enable_backward": False},
 )
 add_example_test(
-    TestFemExamples,
+    TestFemDiffusionExamples,
     name="fem.example_diffusion_3d",
     devices=test_devices,
     options={"quiet": True, "resolution": 10, "num_frames": 1, "enable_backward": False},
@@ -227,13 +319,13 @@ add_example_test(
     options={"quiet": True, "resolution": 20, "enable_backward": False},
 )
 add_example_test(
-    TestExamples,
+    TestFemExamples,
     name="fem.example_convection_diffusion_dg0",
     devices=test_devices,
     options={"quiet": True, "resolution": 20, "num_frames": 25, "mesh": "quad", "enable_backward": False},
 )
 add_example_test(
-    TestExamples,
+    TestFemExamples,
     name="fem.example_convection_diffusion_dg",
     devices=test_devices,
     options={"quiet": True, "resolution": 20, "num_frames": 25, "mesh": "tri", "enable_backward": False},
@@ -251,7 +343,7 @@ add_example_test(
     options={"quiet": True, "num_frames": 1, "enable_backward": False},
 )
 add_example_test(
-    TestExamples,
+    TestFemExamples,
     name="fem.example_stokes",
     devices=test_devices,
     options={
@@ -263,7 +355,7 @@ add_example_test(
     },
 )
 add_example_test(
-    TestExamples,
+    TestFemExamples,
     name="fem.example_navier_stokes",
     devices=test_devices,
     options={"quiet": True, "num_frames": 100, "resolution": 10, "tri_mesh": True, "enable_backward": False},

warp/tests/test_fem.py

@@ -101,9 +101,15 @@ def test_interpolate_gradient(test_case, device):
         # Interpolate gradient at center point
         vector_field = vector_space.make_field()
         vector_field.dof_values.requires_grad = True
+        vector_field_restriction = fem.make_restriction(vector_field)
         tape = wp.Tape()
         with tape:
-            fem.interpolate(grad_field, dest=vector_field, fields={"p": scalar_field})
+            fem.interpolate(
+                grad_field,
+                dest=vector_field_restriction,
+                fields={"p": scalar_field},
+                kernel_options={"enable_backward": True},
+            )
 
         assert_np_equal(vector_field.dof_values.numpy(), np.array([[1.0, 1.0]]))
 
@@ -146,6 +152,7 @@ def test_interpolate_gradient(test_case, device):
         )
         assert_np_equal(point_coords.grad.numpy(), np.array([[[2.0, 0.0, 0.0]]]))
 
+
 @integrand
 def vector_divergence_form(s: Sample, u: Field, q: Field):
     return div(u, s) * q(s)

warp/tests/test_generics.py

@@ -164,6 +164,7 @@ def generic_accumulator_kernel(a: wp.array(dtype=wp.float64), value: Any):
 wp.overload(generic_accumulator_kernel, {"value": int})
 wp.overload(generic_accumulator_kernel, {"value": float})
 wp.overload(generic_accumulator_kernel, {"value": wp.float64})
+wp.overload(generic_accumulator_kernel, {"value": wp.bool})
 
 
 def test_generic_accumulator_kernel(test, device):
@@ -174,8 +175,9 @@ def test_generic_accumulator_kernel(test, device):
         wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, 25])
         wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, 17.0])
         wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, wp.float64(8.0)])
+        wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, wp.bool(True)])
 
-        assert_np_equal(a.numpy(), np.full((n,), 50.0, dtype=np.float64))
+        assert_np_equal(a.numpy(), np.full((n,), 51.0, dtype=np.float64))
 
 
 # generic kernel used to automatically generate overloads from launch args
@@ -191,14 +193,17 @@ def test_generic_fill(test, device):
         ai = wp.zeros(n, dtype=int)
         af = wp.zeros(n, dtype=float)
         a3 = wp.zeros(n, dtype=wp.vec3)
+        ab = wp.zeros(n, dtype=wp.bool)
 
         wp.launch(generic_fill, dim=ai.size, inputs=[ai, 42])
         wp.launch(generic_fill, dim=af.size, inputs=[af, 17.0])
         wp.launch(generic_fill, dim=a3.size, inputs=[a3, wp.vec3(5.0, 5.0, 5.0)])
+        wp.launch(generic_fill, dim=ab.size, inputs=[ab, wp.bool(True)])
 
         assert_np_equal(ai.numpy(), np.full((n,), 42, dtype=np.int32))
         assert_np_equal(af.numpy(), np.full((n,), 17.0, dtype=np.float32))
         assert_np_equal(a3.numpy(), np.full((n, 3), 5.0, dtype=np.float32))
+        assert_np_equal(ab.numpy(), np.full((n,), True, dtype=np.bool_))
 
 
 # generic kernel used to create and launch explicit overloads
@@ -208,10 +213,12 @@ def generic_fill_v2(a: wp.array(dtype=Any), value: Any):
     a[tid] = value
 
 
+vec3b_type = wp.vec(3, wp.bool)
 # create explicit overloads to be launched directly
 fill_int = wp.overload(generic_fill_v2, [wp.array(dtype=int), int])
 fill_float = wp.overload(generic_fill_v2, [wp.array(dtype=float), float])
 fill_vec3 = wp.overload(generic_fill_v2, [wp.array(dtype=wp.vec3), wp.vec3])
+fill_vec3b = wp.overload(generic_fill_v2, [wp.array(dtype=vec3b_type), vec3b_type])
 
 
 def test_generic_fill_overloads(test, device):
@@ -220,19 +227,22 @@ def test_generic_fill_overloads(test, device):
         ai = wp.zeros(n, dtype=int)
         af = wp.zeros(n, dtype=float)
         a3 = wp.zeros(n, dtype=wp.vec3)
+        a3b = wp.zeros(n, dtype=vec3b_type)
 
         wp.launch(fill_int, dim=ai.size, inputs=[ai, 42])
         wp.launch(fill_float, dim=af.size, inputs=[af, 17.0])
         wp.launch(fill_vec3, dim=a3.size, inputs=[a3, wp.vec3(5.0, 5.0, 5.0)])
+        wp.launch(fill_vec3b, dim=a3b.size, inputs=[a3b, vec3b_type([True, True, True])])
 
         assert_np_equal(ai.numpy(), np.full((n,), 42, dtype=np.int32))
         assert_np_equal(af.numpy(), np.full((n,), 17.0, dtype=np.float32))
         assert_np_equal(a3.numpy(), np.full((n, 3), 5.0, dtype=np.float32))
+        assert_np_equal(a3b.numpy(), np.full((n, 3), True, dtype=np.bool_))
 
 
 # custom vector/matrix types
-my_vec5 = wp.types.vector(length=5, dtype=wp.float32)
-my_mat55 = wp.types.matrix(shape=(5, 5), dtype=wp.float32)
+my_vec5 = wp.vec(length=5, dtype=wp.float32)
+my_mat55 = wp.mat(shape=(5, 5), dtype=wp.float32)
 
 
 @wp.kernel
@@ -458,7 +468,7 @@ def test_generic_type_as_argument(test, device):
         wp.synchronize()
 
 
-def test_type_operator_mispell(test, device):
+def test_type_operator_misspell(test, device):
     @wp.kernel
     def kernel():
         i = wp.tid()
@@ -553,7 +563,7 @@ add_kernel_test(
     inputs=[bar],
     devices=devices,
 )
-add_function_test(TestGenerics, "test_type_operator_mispell", test_type_operator_mispell, devices=devices)
+add_function_test(TestGenerics, "test_type_operator_misspell", test_type_operator_misspell, devices=devices)
 add_function_test(TestGenerics, "test_type_attribute_error", test_type_attribute_error, devices=devices)
 
 if __name__ == "__main__":

warp/tests/test_grad.py

@@ -428,7 +428,7 @@ def test_multi_valued_function_grad(test, device):
 
     @wp.func
     def multi_valued(x: float, y: float, z: float):
-        return wp.sin(x), wp.cos(y) * z, wp.sqrt(z) / wp.abs(x)
+        return wp.sin(x), wp.cos(y) * z, wp.sqrt(wp.abs(z)) / wp.abs(x)
 
     # test multi-valued functions
     def check_multi_valued(vs: wp.array(dtype=wp.vec3), out: wp.array(dtype=float)):