warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

examples/example_fluid.py

@@ -18,10 +18,6 @@ import math
 import warp as wp
 import warp.render
 
-import matplotlib
-import matplotlib.pyplot as plt
-import matplotlib.animation as anim
-
 wp.init()
 
 grid_width = wp.constant(256)
@@ -168,7 +164,9 @@ def init(rho: wp.array2d(dtype=float), u: wp.array2d(dtype=wp.vec2), radius: int
 
 
 class Example:
-    def __init__(self):
+    def __init__(self, **kwargs):
+        self.device = wp.get_device()
+
         self.sim_fps = 60.0
         self.sim_substeps = 2
         self.iterations = 100
@@ -191,23 +189,15 @@ class Example:
 
         # capture pressure solve as a CUDA graph
         if self.device.is_cuda:
-            wp.capture_begin()
-            self.solve()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.pressure_iterations()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
-        pass
-
-    def solve(self):
-        for i in range(self.iterations):
-            wp.launch(pressure_solve, dim=self.p0.shape, inputs=[self.p0, self.p1, self.div])
-
-            # swap pressure fields
-            (self.p0, self.p1) = (self.p1, self.p0)
-
-    def render(self, img, i):
         with wp.ScopedTimer("update"):
-            for i in range(self.sim_substeps):
+            for _ in range(self.sim_substeps):
                 shape = (grid_width, grid_height)
                 dt = self.sim_dt
 
@@ -230,7 +220,7 @@ class Example:
                     wp.capture_launch(self.graph)
 
                 else:
-                    self.solve()
+                    self.pressure_iterations()
 
                 # velocity update
                 wp.launch(pressure_apply, dim=shape, inputs=[self.p0, self.u0])
@@ -244,20 +234,39 @@ class Example:
 
                 self.sim_time += dt
 
-        # plot
+    def render(self):
+        pass
+
+    def pressure_iterations(self):
+        for _ in range(self.iterations):
+            wp.launch(pressure_solve, dim=self.p0.shape, inputs=[self.p0, self.p1, self.div])
+
+            # swap pressure fields
+            (self.p0, self.p1) = (self.p1, self.p0)
+
+    def update_and_render_frame(self, frame_num=None, img=None):
+        self.update()
+
         with wp.ScopedTimer("render"):
-            img.set_array(self.rho0.numpy())
+            if img:
+                img.set_array(self.rho0.numpy())
 
         return (img,)
 
 
 if __name__ == "__main__":
+    import matplotlib
+    import matplotlib.animation as anim
+    import matplotlib.pyplot as plt
+
     example = Example()
 
     fig = plt.figure()
 
     img = plt.imshow(example.rho0.numpy(), origin="lower", animated=True, interpolation="antialiased")
     img.set_norm(matplotlib.colors.Normalize(0.0, 1.0))
-    seq = anim.FuncAnimation(fig, lambda i: example.render(img, i), frames=100000, blit=True, interval=8, repeat=False)
+    seq = anim.FuncAnimation(
+        fig, example.update_and_render_frame, fargs=(img,), frames=100000, blit=True, interval=8, repeat=False
+    )
 
     plt.show()

examples/example_jacobian_ik.py

@@ -16,8 +16,8 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import numpy as np
 
@@ -28,22 +28,32 @@ import warp.sim.render
 wp.init()
 
 
-class Robot:
-    frame_dt = 1.0 / 60.0
-
-    render_time = 0.0
+@wp.kernel
+def compute_endeffector_position(
+    body_q: wp.array(dtype=wp.transform),
+    num_links: int,
+    ee_link_index: int,
+    ee_link_offset: wp.vec3,
+    ee_pos: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    ee_pos[tid] = wp.transform_point(body_q[tid * num_links + ee_link_index], ee_link_offset)
 
-    # step size to use for the IK updates
-    step_size = 0.1
 
-    def __init__(self, render=True, num_envs=1, device=None):
+class Example:
+    def __init__(self, stage=None, enable_rendering=True, num_envs=1, device=None):
         builder = wp.sim.ModelBuilder()
 
-        self.render = render
-
         self.num_envs = num_envs
         self.device = device
 
+        self.frame_dt = 1.0 / 60.0
+
+        self.render_time = 0.0
+
+        # step size to use for the IK updates
+        self.step_size = 0.1
+
         articulation_builder = wp.sim.ModelBuilder()
 
         wp.sim.parse_urdf(
@@ -77,7 +87,7 @@ class Robot:
             builder.add_builder(
                 articulation_builder,
                 xform=wp.transform(
-                    np.array((i * 2.0, 4.0, 0.0)), wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)
+                    wp.vec3(i * 2.0, 4.0, 0.0), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
                 ),
             )
             self.target_origin.append((i * 2.0, 4.0, 0.0))
@@ -97,31 +107,24 @@ class Robot:
 
         self.integrator = wp.sim.SemiImplicitIntegrator()
 
-        # -----------------------
-        # set up Usd renderer
-        if self.render:
-            self.renderer = wp.sim.render.SimRenderer(
-                self.model, os.path.join(os.path.dirname(__file__), "outputs/example_jacobian_ik.usd")
-            )
+        self.enable_rendering = enable_rendering
+        self.renderer = None
+        if enable_rendering:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage)
 
         self.ee_pos = wp.zeros(self.num_envs, dtype=wp.vec3, device=device, requires_grad=True)
 
-    @wp.kernel
-    def compute_endeffector_position(
-        body_q: wp.array(dtype=wp.transform),
-        num_links: int,
-        ee_link_index: int,
-        ee_link_offset: wp.vec3,
-        ee_pos: wp.array(dtype=wp.vec3),
-    ):
-        tid = wp.tid()
-        ee_pos[tid] = wp.transform_point(body_q[tid * num_links + ee_link_index], ee_link_offset)
+        self.state = self.model.state(requires_grad=True)
+
+        self.targets = self.target_origin.copy()
+
+        self.profiler = {}
 
     def compute_ee_position(self):
         # computes the end-effector position from the current joint angles
         wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
         wp.launch(
-            self.compute_endeffector_position,
+            compute_endeffector_position,
             dim=self.num_envs,
             inputs=[self.state.body_q, self.num_links, self.ee_link_index, self.ee_link_offset],
             outputs=[self.ee_pos],
@@ -163,12 +166,38 @@ class Robot:
                 jacobians[e, :, i] = (f_plus - f_minus) / (2 * eps)
         return jacobians
 
-    def run(self, render=True):
-        profiler = {}
+    def update(self):
+        with wp.ScopedTimer("jacobian", print=False, active=True, dict=self.profiler):
+            # compute jacobian
+            jacobians = self.compute_jacobian()
+            # jacobians = self.compute_fd_jacobian()
 
-        self.state = self.model.state(requires_grad=True)
+        # compute error
+        self.ee_pos_np = self.compute_ee_position().numpy()
+        error = self.targets - self.ee_pos_np
+        self.error = error.reshape(self.num_envs, 3, 1)
+
+        # compute Jacobian transpose update
+        delta_q = np.matmul(jacobians.transpose(0, 2, 1), self.error)
+
+        self.model.joint_q = wp.array(
+            self.model.joint_q.numpy() + self.step_size * delta_q.flatten(),
+            dtype=wp.float32,
+            device=self.device,
+            requires_grad=True,
+        )
 
-        if render:
+    def render(self):
+        if self.enable_rendering:
+            self.renderer.begin_frame(self.render_time)
+            self.renderer.render(self.state)
+            self.renderer.render_points("targets", self.targets, radius=0.05)
+            self.renderer.render_points("ee_pos", self.ee_pos_np, radius=0.05)
+            self.renderer.end_frame()
+            self.render_time += self.frame_dt
+
+    def run(self):
+        if self.enable_rendering:
             print("autodiff:")
             print(self.compute_jacobian())
             print("finite diff:")
@@ -176,45 +205,18 @@ class Robot:
 
         for _ in range(5):
             # select new random target points
-            targets = self.target_origin.copy()
-            targets[:, 1:] += np.random.uniform(-0.5, 0.5, size=(self.num_envs, 2))
+            self.targets = self.target_origin.copy()
+            self.targets[:, 1:] += np.random.uniform(-0.5, 0.5, size=(self.num_envs, 2))
 
             for iter in range(50):
-                with wp.ScopedTimer("jacobian", print=False, active=True, dict=profiler):
-                    # compute jacobian
-                    jacobians = self.compute_jacobian()
-                    # jacobians = self.compute_fd_jacobian()
-
-                # compute error
-                ee_pos = self.compute_ee_position().numpy()
-                error = targets - ee_pos
-                error = error.reshape(self.num_envs, 3, 1)
-
-                # compute Jacobian transpose update
-                delta_q = np.matmul(jacobians.transpose(0, 2, 1), error)
-
-                self.model.joint_q = wp.array(
-                    self.model.joint_q.numpy() + self.step_size * delta_q.flatten(),
-                    dtype=wp.float32,
-                    device=self.device,
-                    requires_grad=True,
-                )
-
-                # render
-                if render:
-                    self.render_time += self.frame_dt
-
-                    self.renderer.begin_frame(self.render_time)
-                    self.renderer.render(self.state)
-                    self.renderer.render_points("targets", targets, radius=0.05)
-                    self.renderer.render_points("ee_pos", ee_pos, radius=0.05)
-                    self.renderer.end_frame()
-
-                    self.renderer.save()
+                self.update()
+                self.render()
+                print("iter:", iter, "error:", self.error.mean())
 
-                    print("iter:", iter, "error:", error.mean())
+        if self.enable_rendering:
+            self.renderer.save()
 
-        avg_time = np.array(profiler["jacobian"]).mean()
+        avg_time = np.array(self.profiler["jacobian"]).mean()
         avg_steps_second = 1000.0 * float(self.num_envs) / avg_time
 
         print(f"envs: {self.num_envs} steps/second {avg_steps_second} avg_time {avg_time}")
@@ -222,37 +224,39 @@ class Robot:
         return 1000.0 * float(self.num_envs) / avg_time
 
 
-profile = False
+if __name__ == "__main__":
+    profile = False
 
-if profile:
-    env_count = 2
-    env_times = []
-    env_size = []
+    if profile:
+        env_count = 2
+        env_times = []
+        env_size = []
 
-    for i in range(12):
-        robot = Robot(render=False, num_envs=env_count)
-        steps_per_second = robot.run(render=False)
+        for i in range(12):
+            example = Example(enable_rendering=False, num_envs=env_count)
+            steps_per_second = example.run()
 
-        env_size.append(env_count)
-        env_times.append(steps_per_second)
+            env_size.append(env_count)
+            env_times.append(steps_per_second)
 
-        env_count *= 2
+            env_count *= 2
 
-    # dump times
-    for i in range(len(env_times)):
-        print(f"envs: {env_size[i]} steps/second: {env_times[i]}")
+        # dump times
+        for i in range(len(env_times)):
+            print(f"envs: {env_size[i]} steps/second: {env_times[i]}")
 
-    # plot
-    import matplotlib.pyplot as plt
+        # plot
+        import matplotlib.pyplot as plt
 
-    plt.figure(1)
-    plt.plot(env_size, env_times)
-    plt.xscale("log")
-    plt.xlabel("Number of Envs")
-    plt.yscale("log")
-    plt.ylabel("Steps/Second")
-    plt.show()
+        plt.figure(1)
+        plt.plot(env_size, env_times)
+        plt.xscale("log")
+        plt.xlabel("Number of Envs")
+        plt.yscale("log")
+        plt.ylabel("Steps/Second")
+        plt.show()
 
-else:
-    robot = Robot(render=True, num_envs=10)
-    robot.run()
+    else:
+        stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_jacobian_ik.usd")
+        example = Example(stage_path, enable_rendering=True, num_envs=10)
+        example.run()

examples/example_marching_cubes.py

@@ -14,12 +14,12 @@
 ###########################################################################
 
 
+import math
+import os
+
 import warp as wp
 import warp.render
 
-import os
-import math
-
 wp.init()
 
 
@@ -41,7 +41,6 @@ def sdf_box(upper: wp.vec3, p: wp.vec3):
     return wp.length(e) + wp.min(wp.max(qx, wp.max(qy, qz)), 0.0)
 
 
-# union
 @wp.func
 def op_union(d1: float, d2: float):
     return wp.min(d1, d2)
@@ -49,11 +48,6 @@ def op_union(d1: float, d2: float):
 
 @wp.func
 def op_smooth_union(d1: float, d2: float, k: float):
-    # a = wp.pow(d1, k)
-    # b = wp.pow(d2, k)
-
-    # return wp.pow((a*b)/(a+b), 1.0/k)
-
     a = d1
     b = d2
 
@@ -61,13 +55,11 @@ def op_smooth_union(d1: float, d2: float, k: float):
     return wp.lerp(b, a, h) - k * h * (1.0 - h)
 
 
-# subtraction
 @wp.func
 def op_subtract(d1: float, d2: float):
     return wp.max(-d1, d2)
 
 
-# intersection
 @wp.func
 def op_intersect(d1: float, d2: float):
     return wp.max(d1, d2)
@@ -97,6 +89,7 @@ class Example:
         self.max_tris = 10**6
 
         self.time = 0.0
+        self.frame_dt = 1.0 / 60.0
 
         self.field = wp.zeros(shape=(self.dim, self.dim, self.dim), dtype=float)
 
@@ -107,26 +100,23 @@ class Example:
         self.renderer = wp.render.UsdRenderer(stage)
 
     def update(self):
-        pass
-
-    def render(self, is_live=False):
         with wp.ScopedTimer("Update Field"):
             wp.launch(
                 make_field,
                 dim=self.field.shape,
                 inputs=[self.field, wp.vec3(self.dim / 2, self.dim / 2, self.dim / 2), self.dim / 4, self.time],
             )
+            self.time += self.frame_dt
 
         with wp.ScopedTimer("Surface Extraction"):
             self.iso.surface(field=self.field, threshold=math.sin(self.time) * self.dim / 8)
 
+    def render(self, is_live=False):
         with wp.ScopedTimer("Render"):
             self.renderer.begin_frame(self.time)
             self.renderer.render_mesh("surface", self.iso.verts.numpy(), self.iso.indices.numpy(), update_topology=True)
             self.renderer.end_frame()
 
-        self.time += 1.0 / 60.0
-
 
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_marching_cubes.usd")

examples/example_mesh.py

@@ -15,15 +15,14 @@
 #
 ###########################################################################
 
+import os
+
 import numpy as np
+from pxr import Usd, UsdGeom
 
 import warp as wp
 import warp.render
 
-from pxr import Usd, UsdGeom
-
-import os
-
 wp.init()
 
 
@@ -46,7 +45,6 @@ def simulate(
     positions: wp.array(dtype=wp.vec3),
     velocities: wp.array(dtype=wp.vec3),
     mesh: wp.uint64,
-    restitution: float,
     margin: float,
     dt: float,
 ):
@@ -96,7 +94,6 @@ class Example:
         self.sim_time = 0.0
         self.sim_timers = {}
 
-        self.sim_restitution = 0.0
         self.sim_margin = 0.1
 
         self.renderer = wp.render.UsdRenderer(stage)
@@ -128,16 +125,11 @@ class Example:
             wp.launch(
                 kernel=simulate,
                 dim=self.num_particles,
-                inputs=[
-                    self.positions,
-                    self.velocities,
-                    self.mesh.id,
-                    self.sim_restitution,
-                    self.sim_margin,
-                    self.sim_dt,
-                ],
+                inputs=[self.positions, self.velocities, self.mesh.id, self.sim_margin, self.sim_dt],
             )
 
+            self.sim_time += self.sim_dt
+
     def render(self, is_live=False):
         with wp.ScopedTimer("render", detailed=False):
             time = 0.0 if is_live else self.sim_time
@@ -147,8 +139,6 @@ class Example:
             self.renderer.render_points(name="points", points=self.positions.numpy(), radius=self.sim_margin)
             self.renderer.end_frame()
 
-        self.sim_time += self.sim_dt
-
 
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_mesh.usd")

examples/example_mesh_intersect.py

@@ -12,15 +12,16 @@
 #
 ##############################################################################
 
-import warp as wp
-import warp.render
+import os
 
 import numpy as np
+from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.render
 
 np.random.seed(42)
 
-from pxr import Usd, UsdGeom
-import os
 
 wp.init()
 
@@ -82,6 +83,9 @@ def intersect(
 
 class Example:
     def __init__(self, stage):
+        rng = np.random.default_rng()
+
+        self.device = wp.get_device()
         self.query_count = 1024
         self.has_queried = False
 
@@ -101,10 +105,10 @@ class Example:
 
         for i in range(self.query_count):
             # random offset
-            p = (np.random.rand(3) * 0.5 - 0.5) * 5.0
+            p = wp.vec3(rng.random(3) * 0.5 - 0.5) * 5.0
 
             # random orientation
-            axis = wp.normalize(np.random.rand(3) * 0.5 - 0.5)
+            axis = wp.normalize(wp.vec3(rng.random(3) * 0.5 - 0.5))
             angle = float(np.random.rand(1)[0])
 
             q = wp.quat_from_axis_angle(wp.normalize(axis), angle)
@@ -114,8 +118,8 @@ class Example:
         self.array_result = wp.zeros(self.query_count, dtype=int)
         self.array_xforms = wp.array(self.xforms, dtype=wp.transform)
 
-        # force module load (for accurate profiling)
-        wp.force_load()
+        # compile and load the module up front (for accurate profiling)
+        wp.load_module(device=self.device)
 
     def update(self):
         with wp.ScopedTimer("intersect", active=True):
@@ -124,12 +128,10 @@ class Example:
                 dim=self.query_num_faces * self.query_count,
                 inputs=[self.mesh_0.id, self.mesh_1.id, self.query_num_faces, self.array_xforms, self.array_result],
             )
-            wp.synchronize()
 
-    def render(self, is_live=False):
+    def render(self):
         # bring results back to host
         result = self.array_result.numpy()
-        print(result)
 
         with wp.ScopedTimer("render", active=True):
             self.renderer.begin_frame(0.0)
@@ -144,14 +146,14 @@ class Example:
                     os.path.join(os.path.dirname(__file__), self.path_0),
                     pos=wp.vec3(xform.p[0] + offset, xform.p[1], xform.p[2]),
                     rot=xform.q,
-                    scale=(1.0, 1.0, 1.0),
+                    scale=wp.vec3(1.0, 1.0, 1.0),
                 )
                 self.renderer.render_ref(
                     f"mesh_{i}_1",
                     os.path.join(os.path.dirname(__file__), self.path_1),
-                    pos=(offset, 0.0, 0.0),
+                    pos=wp.vec3(offset, 0.0, 0.0),
                     rot=wp.quat_identity(),
-                    scale=(1.0, 1.0, 1.0),
+                    scale=wp.vec3(1.0, 1.0, 1.0),
                 )
 
                 # if pair intersects then draw a small box above the pair