warp-lang 1.7.2rc1__py3-none-macosx_10_13_universal2.whl → 1.8.1__py3-none-macosx_10_13_universal2.whl

warp/examples/interop/example_jax_callable.py

@@ -42,7 +42,7 @@ def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtyp
 
 # The Python function to call.
 # Note the argument annotations, just like Warp kernels.
-def example_func(
+def scale_func(
     # inputs
     a: wp.array(dtype=float),
     b: wp.array(dtype=wp.vec2),
@@ -55,8 +55,23 @@ def example_func(
     wp.launch(scale_vec_kernel, dim=b.shape, inputs=[b, s], outputs=[d])
 
 
+@wp.kernel
+def accum_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
+    tid = wp.tid()
+    b[tid] += a[tid]
+
+
+def in_out_func(
+    a: wp.array(dtype=float),  # input only
+    b: wp.array(dtype=float),  # input and output
+    c: wp.array(dtype=float),  # output only
+):
+    wp.launch(scale_kernel, dim=a.size, inputs=[a, 2.0], outputs=[c])
+    wp.launch(accum_kernel, dim=a.size, inputs=[a, b])  # modifies `b`
+
+
 def example1():
-    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+    jax_func = jax_callable(scale_func, num_outputs=2)
 
     @jax.jit
     def f():
@@ -78,7 +93,7 @@ def example1():
 
 
 def example2():
-    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+    jax_func = jax_callable(scale_func, num_outputs=2)
 
     # NOTE: scalar arguments must be static compile-time constants
     @partial(jax.jit, static_argnames=["s"])
@@ -100,11 +115,26 @@ def example2():
     print(r2)
 
 
+def example3():
+    # Using input-output arguments
+
+    jax_func = jax_callable(in_out_func, num_outputs=2, in_out_argnames=["b"])
+
+    f = jax.jit(jax_func)
+
+    a = jnp.ones(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32)
+
+    b, c = f(a, b)
+    print(b)
+    print(c)
+
+
 def main():
     wp.init()
     wp.load_module(device=wp.get_device())
 
-    examples = [example1, example2]
+    examples = [example1, example2, example3]
 
     for example in examples:
         print("\n===========================================================================")

warp/examples/interop/example_jax_ffi_callback.py

@@ -45,11 +45,11 @@ def example1():
     # the Python function to call
     def print_args(inputs, outputs, attrs, ctx):
         def buffer_to_string(b):
-            return str(b.dtype) + str(list(b.shape)) + " @%x" % b.data
+            return f"{b.dtype}{list(b.shape)} @{b.data:x}"
 
         print("Inputs:     ", ", ".join([buffer_to_string(b) for b in inputs]))
         print("Outputs:    ", ", ".join([buffer_to_string(b) for b in outputs]))
-        print("Attributes: ", "".join(["\n  %s: %s" % (k, str(v)) for k, v in attrs.items()]))
+        print("Attributes: ", "".join([f"\n  {k}: {str(v)}" for k, v in attrs.items()]))  # noqa: RUF010
 
     # register callback
     register_ffi_callback("print_args", print_args)

warp/examples/interop/example_jax_kernel.py

@@ -72,6 +72,17 @@ def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtyp
     output[tid] = a[tid] * s
 
 
+@wp.kernel
+def in_out_kernel(
+    a: wp.array(dtype=float),  # input only
+    b: wp.array(dtype=float),  # input and output
+    c: wp.array(dtype=float),  # output only
+):
+    tid = wp.tid()
+    b[tid] += a[tid]
+    c[tid] = 2.0 * a[tid]
+
+
 def example1():
     # two inputs and one output
     jax_add = jax_kernel(add_kernel)
@@ -189,11 +200,26 @@ def example7():
     print(f())
 
 
+def example8():
+    # Using input-output arguments
+
+    jax_func = jax_kernel(in_out_kernel, num_outputs=2, in_out_argnames=["b"])
+
+    f = jax.jit(jax_func)
+
+    a = jnp.ones(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32)
+
+    b, c = f(a, b)
+    print(b)
+    print(c)
+
+
 def main():
     wp.init()
     wp.load_module(device=wp.get_device())
 
-    examples = [example1, example2, example3, example4, example5, example6, example7]
+    examples = [example1, example2, example3, example4, example5, example6, example7, example8]
 
     for example in examples:
         print("\n===========================================================================")

warp/examples/optim/example_drone.py

@@ -708,7 +708,7 @@ class Example:
         self.tape.zero()
 
     def step(self):
-        if self.frame % int((self.num_frames / len(self.targets))) == 0:
+        if self.frame % int(self.num_frames / len(self.targets)) == 0:
             if self.verbose:
                 print(f"Choosing new flight target: {self.target_idx + 1}")
 

warp/examples/sim/example_cloth.py

@@ -219,7 +219,7 @@ if __name__ == "__main__":
             example.render()
 
         frame_times = example.profiler["step"]
-        print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
+        print(f"\nAverage frame sim time: {sum(frame_times) / len(frame_times):.2f} ms")
 
         if example.renderer:
             example.renderer.save()

warp/examples/sim/example_cloth_self_contact.py

@@ -128,6 +128,10 @@ class Example:
         self.num_substeps = 10
         self.iterations = 4
         self.dt = self.frame_dt / self.num_substeps
+        # the BVH used by VBDIntegrator will be rebuilt every self.bvh_rebuild_frames
+        # When the simulated object deforms significantly, simply refitting the BVH can lead to deterioration of the BVH's
+        # quality, in this case we need to completely rebuild the tree to achieve better query efficiency.
+        self.bvh_rebuild_frames = 10
 
         self.num_frames = num_frames
         self.sim_time = 0.0
@@ -227,69 +231,62 @@ class Example:
         self.cuda_graph = None
         if self.use_cuda_graph:
             with wp.ScopedCapture() as capture:
-                for _ in range(self.num_substeps):
-                    wp.launch(
-                        kernel=apply_rotation,
-                        dim=self.rot_point_indices.shape[0],
-                        inputs=[
-                            self.rot_point_indices,
-                            self.rot_axes,
-                            self.roots,
-                            self.roots_to_ps,
-                            self.t,
-                            self.rot_angular_velocity,
-                            self.dt,
-                            self.rot_end_time,
-                        ],
-                        outputs=[
-                            self.state0.particle_q,
-                            self.state1.particle_q,
-                        ],
-                    )
-
-                    self.integrator.simulate(self.model, self.state0, self.state1, self.dt, None)
-                    (self.state0, self.state1) = (self.state1, self.state0)
+                self.integrate_frame_substeps()
 
             self.cuda_graph = capture.graph
 
-    def step(self):
+    def integrate_frame_substeps(self):
+        for _ in range(self.num_substeps):
+            wp.launch(
+                kernel=apply_rotation,
+                dim=self.rot_point_indices.shape[0],
+                inputs=[
+                    self.rot_point_indices,
+                    self.rot_axes,
+                    self.roots,
+                    self.roots_to_ps,
+                    self.t,
+                    self.rot_angular_velocity,
+                    self.dt,
+                    self.rot_end_time,
+                ],
+                outputs=[
+                    self.state0.particle_q,
+                    self.state1.particle_q,
+                ],
+            )
+
+            self.integrator.simulate(self.model, self.state0, self.state1, self.dt, None)
+            (self.state0, self.state1) = (self.state1, self.state0)
+
+    def advance_frame(self):
         with wp.ScopedTimer("step", print=False, dict=self.profiler):
             if self.use_cuda_graph:
                 wp.capture_launch(self.cuda_graph)
             else:
-                for _ in range(self.num_substeps):
-                    wp.launch(
-                        kernel=apply_rotation,
-                        dim=self.rot_point_indices.shape[0],
-                        inputs=[
-                            self.rot_point_indices,
-                            self.rot_axes,
-                            self.roots,
-                            self.roots_to_ps,
-                            self.t,
-                            self.rot_angular_velocity,
-                            self.dt,
-                            self.rot_end_time,
-                        ],
-                        outputs=[
-                            self.state0.particle_q,
-                            self.state1.particle_q,
-                        ],
-                    )
-                    self.integrator.simulate(self.model, self.state0, self.state1, self.dt)
-
-                    (self.state0, self.state1) = (self.state1, self.state0)
+                self.integrate_frame_substeps()
 
             self.sim_time += self.dt
 
+    def run(self):
+        for i in range(self.num_frames):
+            self.advance_frame()
+            self.render()
+            print(f"[{i:4d}/{self.num_frames}]")
+
+            if i != 0 and not i % self.bvh_rebuild_frames and self.use_cuda_graph:
+                self.integrator.rebuild_bvh(self.state0)
+                with wp.ScopedCapture() as capture:
+                    self.integrate_frame_substeps()
+                self.cuda_graph = capture.graph
+
     def render(self):
         if self.renderer is None:
             return
 
-        with wp.ScopedTimer("render", print=False):
-            self.renderer.begin_frame(self.sim_time)
-            self.renderer.render(self.state0)
-            self.renderer.end_frame()
+        self.renderer.begin_frame(self.sim_time)
+        self.renderer.render(self.state0)
+        self.renderer.end_frame()
 
 
 if __name__ == "__main__":
@@ -310,13 +307,10 @@ if __name__ == "__main__":
     with wp.ScopedDevice(args.device):
         example = Example(stage_path=args.stage_path, num_frames=args.num_frames)
 
-        for i in range(example.num_frames):
-            example.step()
-            example.render()
-            print(f"[{i:4d}/{example.num_frames}]")
+        example.run()
 
         frame_times = example.profiler["step"]
-        print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
+        print(f"\nAverage frame sim time: {sum(frame_times) / len(frame_times):.2f} ms")
 
         if example.renderer:
             example.renderer.save()