warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/core/example_mesh.py

@@ -0,0 +1,180 @@
+# 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 Mesh
+#
+# Shows how to implement a PBD particle simulation with collision against
+# a deforming triangle mesh. The mesh collision uses wp.mesh_query_point_sign_normal()
+# to compute the closest point, and wp.Mesh.refit() to update the mesh
+# object after deformation.
+#
+###########################################################################
+
+import os
+
+import numpy as np
+from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.examples
+import warp.render
+
+
+@wp.kernel
+def deform(positions: wp.array(dtype=wp.vec3), t: float):
+    tid = wp.tid()
+
+    x = positions[tid]
+
+    offset = -wp.sin(x[0]) * 0.02
+    scale = wp.sin(t)
+
+    x = x + wp.vec3(0.0, offset * scale, 0.0)
+
+    positions[tid] = x
+
+
+@wp.kernel
+def simulate(
+    positions: wp.array(dtype=wp.vec3),
+    velocities: wp.array(dtype=wp.vec3),
+    mesh: wp.uint64,
+    margin: float,
+    dt: float,
+):
+    tid = wp.tid()
+
+    x = positions[tid]
+    v = velocities[tid]
+
+    v = v + wp.vec3(0.0, 0.0 - 9.8, 0.0) * dt - v * 0.1 * dt
+    xpred = x + v * dt
+
+    max_dist = 1.5
+
+    query = wp.mesh_query_point_sign_normal(mesh, xpred, max_dist)
+    if query.result:
+        p = wp.mesh_eval_position(mesh, query.face, query.u, query.v)
+
+        delta = xpred - p
+
+        dist = wp.length(delta) * query.sign
+        err = dist - margin
+
+        # mesh collision
+        if err < 0.0:
+            n = wp.normalize(delta) * query.sign
+            xpred = xpred - n * err
+
+    # pbd update
+    v = (xpred - x) * (1.0 / dt)
+    x = xpred
+
+    positions[tid] = x
+    velocities[tid] = v
+
+
+class Example:
+    def __init__(self, stage_path="example_mesh.usd"):
+        rng = np.random.default_rng(42)
+        self.num_particles = 1000
+
+        self.sim_dt = 1.0 / 60.0
+
+        self.sim_time = 0.0
+        self.sim_timers = {}
+
+        self.sim_margin = 0.1
+
+        usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+        usd_scale = 10.0
+
+        # create collision mesh
+        self.mesh = wp.Mesh(
+            points=wp.array(usd_geom.GetPointsAttr().Get() * usd_scale, dtype=wp.vec3),
+            indices=wp.array(usd_geom.GetFaceVertexIndicesAttr().Get(), dtype=int),
+        )
+
+        # random particles
+        init_pos = (rng.random((self.num_particles, 3)) - np.array([0.5, -1.5, 0.5])) * 10.0
+        init_vel = rng.random((self.num_particles, 3)) * 0.0
+
+        self.positions = wp.from_numpy(init_pos, dtype=wp.vec3)
+        self.velocities = wp.from_numpy(init_vel, dtype=wp.vec3)
+
+        # renderer
+        self.renderer = None
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path)
+
+    def step(self):
+        with wp.ScopedTimer("step", dict=self.sim_timers):
+            wp.launch(kernel=deform, dim=len(self.mesh.points), inputs=[self.mesh.points, self.sim_time])
+
+            # refit the mesh BVH to account for the deformation
+            self.mesh.refit()
+
+            wp.launch(
+                kernel=simulate,
+                dim=self.num_particles,
+                inputs=[self.positions, self.velocities, self.mesh.id, self.sim_margin, 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_mesh(
+                name="mesh",
+                points=self.mesh.points.numpy(),
+                indices=self.mesh.indices.numpy(),
+                colors=(0.35, 0.55, 0.9),
+            )
+            self.renderer.render_points(
+                name="points", points=self.positions.numpy(), radius=self.sim_margin, colors=(0.8, 0.3, 0.2)
+            )
+            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_mesh.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=500, 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/core/example_mesh_intersect.py

@@ -0,0 +1,211 @@
+# 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 Mesh Intersection
+#
+# Show how to use built-in BVH query to test if two triangle meshes intersect.
+#
+##############################################################################
+
+import os
+
+import numpy as np
+from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.examples
+import warp.render
+
+
+@wp.func
+def cw_min(a: wp.vec3, b: wp.vec3):
+    return wp.vec3(wp.min(a[0], b[0]), wp.min(a[1], b[1]), wp.min(a[2], b[2]))
+
+
+@wp.func
+def cw_max(a: wp.vec3, b: wp.vec3):
+    return wp.vec3(wp.max(a[0], b[0]), wp.max(a[1], b[1]), wp.max(a[2], b[2]))
+
+
+@wp.kernel
+def intersect(
+    mesh_0: wp.uint64,
+    mesh_1: wp.uint64,
+    num_faces: int,
+    xforms: wp.array(dtype=wp.transform),
+    result: wp.array(dtype=int),
+):
+    tid = wp.tid()
+
+    # mesh_0 is assumed to be the query mesh, we launch one thread
+    # for each face in mesh_0 and test it against the opposing mesh's BVH
+    face = tid % num_faces
+    batch = tid // num_faces
+
+    # transforms from mesh_0 -> mesh_1 space
+    xform = xforms[batch]
+
+    # load query triangles points and transform to mesh_1's space
+    v0 = wp.transform_point(xform, wp.mesh_eval_position(mesh_0, face, 1.0, 0.0))
+    v1 = wp.transform_point(xform, wp.mesh_eval_position(mesh_0, face, 0.0, 1.0))
+    v2 = wp.transform_point(xform, wp.mesh_eval_position(mesh_0, face, 0.0, 0.0))
+
+    # compute bounds of the query triangle
+    lower = cw_min(cw_min(v0, v1), v2)
+    upper = cw_max(cw_max(v0, v1), v2)
+
+    query = wp.mesh_query_aabb(mesh_1, lower, upper)
+
+    for f in query:
+        u0 = wp.mesh_eval_position(mesh_1, f, 1.0, 0.0)
+        u1 = wp.mesh_eval_position(mesh_1, f, 0.0, 1.0)
+        u2 = wp.mesh_eval_position(mesh_1, f, 0.0, 0.0)
+
+        # test for triangle intersection
+        i = wp.intersect_tri_tri(v0, v1, v2, u0, u1, u2)
+
+        if i > 0:
+            result[batch] = 1
+            return
+
+        # use if you want to count all intersections
+        # wp.atomic_add(result, batch, i)
+
+
+class Example:
+    def __init__(self, stage_path="example_mesh_intersect.usd"):
+        rng = np.random.default_rng(42)
+
+        self.query_count = 1024
+        self.has_queried = False
+
+        self.path_0 = os.path.join(warp.examples.get_asset_directory(), "cube.usd")
+        self.path_1 = os.path.join(warp.examples.get_asset_directory(), "sphere.usd")
+
+        self.mesh_0 = self.load_mesh(self.path_0, "/root/cube")
+        self.mesh_1 = self.load_mesh(self.path_1, "/root/sphere")
+
+        self.query_num_faces = int(len(self.mesh_0.indices) / 3)
+        self.query_num_points = len(self.mesh_0.points)
+
+        # generate random relative transforms
+        self.xforms = []
+
+        for _ in range(self.query_count):
+            # random offset
+            p = wp.vec3(rng.random(size=3) * 0.5 - 0.5) * 5.0
+
+            # random orientation
+            axis = wp.normalize(wp.vec3(rng.random(size=3) * 0.5 - 0.5))
+            angle = rng.random()
+
+            q = wp.quat_from_axis_angle(wp.normalize(axis), angle)
+
+            self.xforms.append(wp.transform(p, q))
+
+        self.array_result = wp.zeros(self.query_count, dtype=int)
+        self.array_xforms = wp.array(self.xforms, dtype=wp.transform)
+
+        # renderer
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path)
+        else:
+            self.renderer = None
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            wp.launch(
+                kernel=intersect,
+                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],
+            )
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        # bring results back to host
+        result = self.array_result.numpy()
+
+        with wp.ScopedTimer("render", active=True):
+            self.renderer.begin_frame(0.0)
+
+            for i in range(self.query_count):
+                spacing = 8.0
+                offset = i * spacing
+
+                xform = self.xforms[i]
+                self.renderer.render_ref(
+                    f"mesh_{i}_0",
+                    self.path_0,
+                    pos=wp.vec3(xform.p[0] + offset, xform.p[1], xform.p[2]),
+                    rot=xform.q,
+                    scale=wp.vec3(1.0, 1.0, 1.0),
+                )
+                self.renderer.render_ref(
+                    f"mesh_{i}_1",
+                    self.path_1,
+                    pos=wp.vec3(offset, 0.0, 0.0),
+                    rot=wp.quat_identity(),
+                    scale=wp.vec3(1.0, 1.0, 1.0),
+                )
+
+                # if pair intersects then draw a small box above the pair
+                if result[i] > 0:
+                    self.renderer.render_box(
+                        f"result_{i}",
+                        pos=wp.vec3(xform.p[0] + offset, xform.p[1] + 5.0, xform.p[2]),
+                        rot=wp.quat_identity(),
+                        extents=(0.1, 0.1, 0.1),
+                    )
+
+            self.renderer.end_frame()
+
+    # create collision meshes
+    def load_mesh(self, path, prim):
+        usd_stage = Usd.Stage.Open(path)
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath(prim))
+
+        mesh = wp.Mesh(
+            points=wp.array(usd_geom.GetPointsAttr().Get(), dtype=wp.vec3),
+            indices=wp.array(usd_geom.GetFaceVertexIndicesAttr().Get(), dtype=int),
+        )
+
+        return mesh
+
+
+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_mesh_intersect.usd",
+        help="Path to the output USD file.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        example.step()
+        example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/core/example_nvdb.py

@@ -0,0 +1,182 @@
+# 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 NanoVDB
+#
+# Shows how to implement a particle simulation with collision against
+# a NanoVDB signed-distance field. In this example the NanoVDB field
+# is created offline in Houdini. The particle kernel uses the Warp
+# wp.volume_sample_f() method to compute the SDF and normal at a point.
+#
+###########################################################################
+
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.examples
+import warp.render
+
+
+@wp.func
+def volume_grad(volume: wp.uint64, p: wp.vec3):
+    eps = 1.0
+    q = wp.volume_world_to_index(volume, p)
+
+    # compute gradient of the SDF using finite differences
+    dx = wp.volume_sample_f(volume, q + wp.vec3(eps, 0.0, 0.0), wp.Volume.LINEAR) - wp.volume_sample_f(
+        volume, q - wp.vec3(eps, 0.0, 0.0), wp.Volume.LINEAR
+    )
+    dy = wp.volume_sample_f(volume, q + wp.vec3(0.0, eps, 0.0), wp.Volume.LINEAR) - wp.volume_sample_f(
+        volume, q - wp.vec3(0.0, eps, 0.0), wp.Volume.LINEAR
+    )
+    dz = wp.volume_sample_f(volume, q + wp.vec3(0.0, 0.0, eps), wp.Volume.LINEAR) - wp.volume_sample_f(
+        volume, q - wp.vec3(0.0, 0.0, eps), wp.Volume.LINEAR
+    )
+
+    return wp.normalize(wp.vec3(dx, dy, dz))
+
+
+@wp.kernel
+def simulate(
+    positions: wp.array(dtype=wp.vec3),
+    velocities: wp.array(dtype=wp.vec3),
+    volume: wp.uint64,
+    margin: float,
+    dt: float,
+):
+    tid = wp.tid()
+
+    x = positions[tid]
+    v = velocities[tid]
+
+    v = v + wp.vec3(0.0, -9.8, 0.0) * dt - v * 0.1 * dt
+    xpred = x + v * dt
+    xpred_local = wp.volume_world_to_index(volume, xpred)
+
+    # d = wp.volume_sample_f(volume, xpred_local, wp.Volume.LINEAR)
+    n = wp.vec3()
+    d = wp.volume_sample_grad_f(volume, xpred_local, wp.Volume.LINEAR, n)
+
+    if d < margin:
+        # n = volume_grad(volume, xpred)
+        n = wp.normalize(n)
+        err = d - margin
+
+        # mesh collision
+        xpred = xpred - n * err
+
+    # ground collision
+    if xpred[1] < 0.0:
+        xpred = wp.vec3(xpred[0], 0.0, xpred[2])
+
+    # pbd update
+    v = (xpred - x) * (1.0 / dt)
+    x = xpred
+
+    positions[tid] = x
+    velocities[tid] = v
+
+
+class Example:
+    def __init__(self, stage_path="example_nvdb.usd"):
+        rng = np.random.default_rng(42)
+        self.num_particles = 10000
+
+        fps = 60
+        frame_dt = 1.0 / fps
+        self.sim_substeps = 3
+        self.sim_dt = frame_dt / self.sim_substeps
+
+        self.sim_time = 0.0
+        self.sim_timers = {}
+
+        self.sim_margin = 0.15
+
+        init_pos = 10.0 * (rng.random((self.num_particles, 3)) * 2.0 - 1.0) + np.array((0.0, 30.0, 0.0))
+        init_vel = rng.random((self.num_particles, 3))
+
+        self.positions = wp.from_numpy(init_pos.astype(np.float32), dtype=wp.vec3)
+        self.velocities = wp.from_numpy(init_vel.astype(np.float32), dtype=wp.vec3)
+
+        # load collision volume
+        with open(os.path.join(warp.examples.get_asset_directory(), "rocks.nvdb"), "rb") as file:
+            # create Volume object
+            self.volume = wp.Volume.load_from_nvdb(file)
+
+        # renderer
+        self.renderer = None
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path)
+            self.renderer.render_ground(size=100.0)
+
+    def step(self):
+        with wp.ScopedTimer("step", dict=self.sim_timers):
+            for _ in range(self.sim_substeps):
+                wp.launch(
+                    kernel=simulate,
+                    dim=self.num_particles,
+                    inputs=[self.positions, self.velocities, self.volume.id, self.sim_margin, 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_ref(
+                name="collision",
+                path=os.path.join(warp.examples.get_asset_directory(), "rocks.usd"),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                rot=wp.quat(0.0, 0.0, 0.0, 1.0),
+                scale=wp.vec3(1.0, 1.0, 1.0),
+                color=(0.35, 0.55, 0.9),
+            )
+            self.renderer.render_points(
+                name="points", points=self.positions.numpy(), radius=self.sim_margin, colors=(0.8, 0.3, 0.2)
+            )
+
+            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_nvdb.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=1000, 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/core/example_raycast.py

@@ -0,0 +1,111 @@
+# 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 Ray Cast
+#
+# Shows how to use the built-in wp.Mesh data structure and wp.mesh_query_ray()
+# function to implement a basic ray-tracer.
+#
+##############################################################################
+
+import os
+
+import numpy as np
+from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.examples
+
+
+@wp.kernel
+def draw(mesh: wp.uint64, cam_pos: wp.vec3, width: int, height: int, pixels: wp.array(dtype=wp.vec3)):
+    tid = wp.tid()
+
+    x = tid % width
+    y = tid // width
+
+    sx = 2.0 * float(x) / float(height) - 1.0
+    sy = 2.0 * float(y) / float(height) - 1.0
+
+    # compute view ray
+    ro = cam_pos
+    rd = wp.normalize(wp.vec3(sx, sy, -1.0))
+
+    color = wp.vec3(0.0, 0.0, 0.0)
+
+    query = wp.mesh_query_ray(mesh, ro, rd, 1.0e6)
+    if query.result:
+        color = query.normal * 0.5 + wp.vec3(0.5, 0.5, 0.5)
+
+    pixels[tid] = color
+
+
+class Example:
+    def __init__(self, height=1024, width=1024):
+        self.height = height
+        self.width = width
+        self.cam_pos = (0.0, 1.0, 2.0)
+
+        asset_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
+        mesh_geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/root/bunny"))
+
+        points = np.array(mesh_geom.GetPointsAttr().Get())
+        indices = np.array(mesh_geom.GetFaceVertexIndicesAttr().Get())
+
+        self.pixels = wp.zeros(self.width * self.height, dtype=wp.vec3)
+
+        # create wp mesh
+        self.mesh = wp.Mesh(
+            points=wp.array(points, dtype=wp.vec3), velocities=None, indices=wp.array(indices, dtype=int)
+        )
+
+    def render(self):
+        with wp.ScopedTimer("render"):
+            wp.launch(
+                kernel=draw,
+                dim=self.width * self.height,
+                inputs=[self.mesh.id, self.cam_pos, self.width, self.height, self.pixels],
+            )
+
+
+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("--width", type=int, default=1024, help="Output image width in pixels.")
+    parser.add_argument("--height", type=int, default=1024, help="Output image height in pixels.")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(height=args.height, width=args.width)
+        example.render()
+
+        if not args.headless:
+            import matplotlib.pyplot as plt
+
+            plt.imshow(
+                example.pixels.numpy().reshape((example.height, example.width, 3)),
+                origin="lower",
+                interpolation="antialiased",
+            )
+            plt.show()