warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/sim/example_rigid_force.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Sim Rigid Force
+#
+# Shows how to apply an external force (torque) to a rigid body causing
+# it to roll.
+#
+###########################################################################
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+class Example:
+    def __init__(self, stage_path="example_rigid_force.usd", use_opengl=False):
+        fps = 60
+        self.frame_dt = 1.0 / fps
+        self.sim_substeps = 5
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+
+        builder = wp.sim.ModelBuilder()
+
+        b = builder.add_body(origin=wp.transform((0.0, 10.0, 0.0), wp.quat_identity()))
+        builder.add_shape_box(body=b, hx=1.0, hy=1.0, hz=1.0, density=100.0)
+
+        self.model = builder.finalize()
+        self.model.ground = True
+
+        self.integrator = wp.sim.XPBDIntegrator()
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        if use_opengl:
+            self.renderer = wp.sim.render.SimRendererOpenGL(self.model, stage_path)
+        elif stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path)
+        else:
+            self.renderer = None
+
+        # simulate() allocates memory via a clone, so we can't use graph capture if the device does not support mempools
+        self.use_cuda_graph = wp.get_device().is_cuda and wp.is_mempool_enabled(wp.get_device())
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            wp.sim.collide(self.model, self.state_0)
+
+            self.state_0.clear_forces()
+            self.state_1.clear_forces()
+
+            self.state_0.body_f.assign(
+                [
+                    [0.0, 0.0, -7000.0, 0.0, 0.0, 0.0],
+                ]
+            )
+
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+
+            # swap states
+            (self.state_0, self.state_1) = (self.state_1, self.state_0)
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+        self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_rigid_force.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
+    parser.add_argument(
+        "--opengl",
+        action="store_true",
+        help="Open an interactive window to play back animations in real time. Ignores --num_frames if used.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, use_opengl=args.opengl)
+
+        if args.opengl:
+            while example.renderer.is_running():
+                example.step()
+                example.render()
+        else:
+            for _ in range(args.num_frames):
+                example.step()
+                example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/sim/example_rigid_gyroscopic.py

@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Sim Rigid Gyroscopic
+#
+# Demonstrates the Dzhanibekov effect where rigid bodies will tumble in
+# free space due to unstable axes of rotation.
+#
+###########################################################################
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+class Example:
+    def __init__(self, stage_path="example_rigid_gyroscopic.usd"):
+        fps = 120
+        self.sim_dt = 1.0 / fps
+        self.sim_time = 0.0
+
+        self.scale = 0.5
+
+        builder = wp.sim.ModelBuilder()
+
+        b = builder.add_body()
+
+        # axis shape
+        builder.add_shape_box(
+            pos=wp.vec3(0.3 * self.scale, 0.0, 0.0),
+            hx=0.25 * self.scale,
+            hy=0.1 * self.scale,
+            hz=0.1 * self.scale,
+            density=100.0,
+            body=b,
+        )
+
+        # tip shape
+        builder.add_shape_box(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            hx=0.05 * self.scale,
+            hy=0.2 * self.scale,
+            hz=1.0 * self.scale,
+            density=100.0,
+            body=b,
+        )
+
+        # initial spin
+        builder.body_qd[0] = (25.0, 0.01, 0.01, 0.0, 0.0, 0.0)
+
+        builder.gravity = 0.0
+        self.model = builder.finalize()
+        self.model.ground = False
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+        self.state = self.model.state()
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=100.0)
+        else:
+            self.renderer = None
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            self.state.clear_forces()
+            self.state = self.integrator.simulate(self.model, self.state, self.state, self.sim_dt)
+            self.sim_time += self.sim_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_rigid_gyroscopic.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=2000, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/sim/example_rigid_soft_contact.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Sim Rigid FEM
+#
+# Shows how to set up a rigid sphere colliding with an FEM beam
+# using wp.sim.ModelBuilder().
+#
+###########################################################################
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+class Example:
+    def __init__(self, stage_path="example_rigid_soft_contact.usd"):
+        self.sim_width = 8
+        self.sim_height = 8
+
+        fps = 60
+        self.frame_dt = 1.0 / fps
+        self.sim_substeps = 32
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+        self.sim_iterations = 1
+        self.sim_relaxation = 1.0
+        self.profiler = {}
+
+        builder = wp.sim.ModelBuilder()
+        builder.default_particle_radius = 0.01
+
+        builder.add_soft_grid(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_identity(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=20,
+            dim_y=10,
+            dim_z=10,
+            cell_x=0.1,
+            cell_y=0.1,
+            cell_z=0.1,
+            density=100.0,
+            k_mu=50000.0,
+            k_lambda=20000.0,
+            k_damp=0.0,
+        )
+
+        b = builder.add_body(origin=wp.transform((0.5, 2.5, 0.5), wp.quat_identity()))
+        builder.add_shape_sphere(body=b, radius=0.75, density=100.0)
+
+        self.model = builder.finalize()
+        self.model.ground = True
+        self.model.soft_contact_ke = 1.0e3
+        self.model.soft_contact_kd = 0.0
+        self.model.soft_contact_kf = 1.0e3
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1.0)
+        else:
+            self.renderer = None
+
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _s in range(self.sim_substeps):
+            wp.sim.collide(self.model, self.state_0)
+
+            self.state_0.clear_forces()
+            self.state_1.clear_forces()
+
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+
+            # swap states
+            (self.state_0, self.state_1) = (self.state_1, self.state_0)
+
+    def step(self):
+        with wp.ScopedTimer("step", dict=self.profiler):
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+            self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_rigid_soft_contact.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/sim/example_soft_body.py

@@ -0,0 +1,196 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Sim Neo-Hookean
+#
+# Shows a simulation of an Neo-Hookean FEM beam being twisted through a
+# 180 degree rotation.
+#
+###########################################################################
+import math
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+@wp.kernel
+def twist_points(
+    rest: wp.array(dtype=wp.vec3), points: wp.array(dtype=wp.vec3), mass: wp.array(dtype=float), xform: wp.transform
+):
+    tid = wp.tid()
+
+    r = rest[tid]
+    p = points[tid]
+    m = mass[tid]
+
+    # twist the top layer of particles in the beam
+    if m == 0 and p[1] != 0.0:
+        points[tid] = wp.transform_point(xform, r)
+
+
+@wp.kernel
+def compute_volume(points: wp.array(dtype=wp.vec3), indices: wp.array2d(dtype=int), volume: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    i = indices[tid, 0]
+    j = indices[tid, 1]
+    k = indices[tid, 2]
+    l = indices[tid, 3]
+
+    x0 = points[i]
+    x1 = points[j]
+    x2 = points[k]
+    x3 = points[l]
+
+    x10 = x1 - x0
+    x20 = x2 - x0
+    x30 = x3 - x0
+
+    v = wp.dot(x10, wp.cross(x20, x30)) / 6.0
+
+    wp.atomic_add(volume, 0, v)
+
+
+class Example:
+    def __init__(self, stage_path="example_soft_body.usd", num_frames=300):
+        self.sim_substeps = 64
+        self.num_frames = num_frames
+        fps = 60
+        sim_duration = self.num_frames / fps
+        self.frame_dt = 1.0 / fps
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+        self.lift_speed = 2.5 / sim_duration * 2.0  # from Smith et al.
+        self.rot_speed = math.pi / sim_duration
+
+        builder = wp.sim.ModelBuilder()
+
+        cell_dim = 15
+        cell_size = 2.0 / cell_dim
+
+        center = cell_size * cell_dim * 0.5
+
+        builder.add_soft_grid(
+            pos=wp.vec3(-center, 0.0, -center),
+            rot=wp.quat_identity(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=cell_dim,
+            dim_y=cell_dim,
+            dim_z=cell_dim,
+            cell_x=cell_size,
+            cell_y=cell_size,
+            cell_z=cell_size,
+            density=100.0,
+            fix_bottom=True,
+            fix_top=True,
+            k_mu=1000.0,
+            k_lambda=5000.0,
+            k_damp=0.0,
+        )
+
+        self.model = builder.finalize()
+        self.model.ground = False
+        self.model.gravity[1] = 0.0
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.rest = self.model.state()
+        self.rest_vol = (cell_size * cell_dim) ** 3
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        self.volume = wp.zeros(1, dtype=wp.float32)
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=20.0)
+        else:
+            self.renderer = None
+
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            self.state_0.clear_forces()
+            self.state_1.clear_forces()
+
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+
+            # swap states
+            (self.state_0, self.state_1) = (self.state_1, self.state_0)
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            xform = wp.transform(
+                (0.0, self.lift_speed * self.sim_time, 0.0),
+                wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
+            )
+            wp.launch(
+                kernel=twist_points,
+                dim=len(self.state_0.particle_q),
+                inputs=[self.rest.particle_q, self.state_0.particle_q, self.model.particle_mass, xform],
+            )
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+            self.volume.zero_()
+            wp.launch(
+                kernel=compute_volume,
+                dim=self.model.tet_count,
+                inputs=[self.state_0.particle_q, self.model.tet_indices, self.volume],
+            )
+        self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_soft_body.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, num_frames=args.num_frames)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/tile/example_tile_cholesky.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Tile Cholesky
+#
+# Shows how to write a simple kernel computing a Cholesky factorize and
+# triangular solve using Warp Cholesky Tile APIs.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.init()
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 128
+TILE = 32
+
+# Both should work
+np_type, wp_type = np.float64, wp.float64
+# np_type, wp_type = np.float32, wp.float32
+
+
+@wp.kernel
+def cholesky(
+    A: wp.array2d(dtype=wp_type),
+    L: wp.array2d(dtype=wp_type),
+    X: wp.array1d(dtype=wp_type),
+    Y: wp.array1d(dtype=wp_type),
+):
+    i, j, _ = wp.tid()
+
+    a = wp.tile_load(A, shape=(TILE, TILE))
+    l = wp.tile_cholesky(a)
+    wp.tile_store(L, l)
+
+    x = wp.tile_load(X, shape=TILE)
+    y = wp.tile_cholesky_solve(l, x)
+    wp.tile_store(Y, y)
+
+
+if __name__ == "__main__":
+    wp.set_device("cuda:0")
+
+    A_h = np.ones((TILE, TILE), dtype=np_type) + 5 * np.diag(np.ones(TILE), 0)
+    L_h = np.zeros_like(A_h)
+
+    A_wp = wp.array2d(A_h, dtype=wp_type)
+    L_wp = wp.array2d(L_h, dtype=wp_type)
+
+    X_h = np.arange(TILE, dtype=np_type)
+    Y_h = np.zeros_like(X_h)
+
+    X_wp = wp.array1d(X_h, dtype=wp_type)
+    Y_wp = wp.array1d(Y_h, dtype=wp_type)
+
+    wp.launch_tiled(cholesky, dim=[1, 1], inputs=[A_wp, L_wp, X_wp, Y_wp], block_dim=BLOCK_DIM)
+
+    L_np = np.linalg.cholesky(A_h)
+    Y_np = np.linalg.solve(A_h, X_h)
+
+    print("A:\n", A_h)
+    print("L (Warp):\n", L_wp)
+    print("L (Numpy):\n", L_np)
+
+    print("x:\n", X_h)
+    print("A\\n (Warp):\n", Y_wp.numpy())
+    print("A\\x (Numpy):\n", Y_np)
+
+    assert np.allclose(Y_wp.numpy(), Y_np) and np.allclose(L_wp.numpy(), L_np)
+
+    print("Example Tile Cholesky passed")