warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/core/example_marching_cubes.py

@@ -0,0 +1,194 @@
+# 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 Marching Cubes
+#
+# Shows how use the built-in marching cubes functionality to extract
+# the iso-surface from a density field.
+#
+# Note: requires a CUDA-capable device
+###########################################################################
+
+import warp as wp
+import warp.render
+
+
+@wp.func
+def sdf_create_box(pos: wp.vec3, size: wp.vec3):
+    """Creates a SDF box primitive."""
+    # https://iquilezles.org/articles/distfunctions
+    q = wp.vec3(
+        wp.abs(pos[0]) - size[0],
+        wp.abs(pos[1]) - size[1],
+        wp.abs(pos[2]) - size[2],
+    )
+    qp = wp.vec3(wp.max(q[0], 0.0), wp.max(q[1], 0.0), wp.max(q[2], 0.0))
+    return wp.length(qp) + wp.min(wp.max(q[0], wp.max(q[1], q[2])), 0.0)
+
+
+@wp.func
+def sdf_create_torus(pos: wp.vec3, major_radius: float, minor_radius: float):
+    """Creates a SDF torus primitive."""
+    # https://iquilezles.org/articles/distfunctions
+    q = wp.vec2(wp.length(wp.vec2(pos[0], pos[2])) - major_radius, pos[1])
+    return wp.length(q) - minor_radius
+
+
+@wp.func
+def sdf_translate(pos: wp.vec3, offset: wp.vec3):
+    """Translates a SDF position vector with an offset."""
+    return pos - offset
+
+
+@wp.func
+def sdf_rotate(pos: wp.vec3, angles: wp.vec3):
+    """Rotates a SDF position vector using Euler angles."""
+    rot = wp.quat_rpy(
+        wp.radians(angles[0]),
+        wp.radians(angles[1]),
+        wp.radians(angles[2]),
+    )
+    return wp.quat_rotate_inv(rot, pos)
+
+
+@wp.func
+def sdf_smooth_min(a: float, b: float, radius: float):
+    """Creates a SDF torus primitive."""
+    # https://iquilezles.org/articles/smin
+    h = wp.max(radius - wp.abs(a - b), 0.0) / radius
+    return wp.min(a, b) - h * h * h * radius * (1.0 / 6.0)
+
+
+@wp.kernel(enable_backward=False)
+def make_field(
+    torus_altitude: float,
+    torus_major_radius: float,
+    torus_minor_radius: float,
+    smooth_min_radius: float,
+    dim: int,
+    time: float,
+    out_data: wp.array3d(dtype=float),
+):
+    """Kernel to generate a SDF volume based on primitives."""
+    i, j, k = wp.tid()
+
+    # Retrieve the position of the current cell in a normalized [-1, 1] range
+    # for each dimension.
+    pos = wp.vec3(
+        2.0 * ((float(i) + 0.5) / float(dim)) - 1.0,
+        2.0 * ((float(j) + 0.5) / float(dim)) - 1.0,
+        2.0 * ((float(k) + 0.5) / float(dim)) - 1.0,
+    )
+
+    box = sdf_create_box(
+        sdf_translate(pos, wp.vec3(0.0, -0.7, 0.0)),
+        wp.vec3(0.9, 0.3, 0.9),
+    )
+    torus = sdf_create_torus(
+        sdf_rotate(
+            sdf_translate(pos, wp.vec3(0.0, torus_altitude, 0.0)),
+            wp.vec3(wp.sin(time) * 90.0, wp.cos(time) * 45.0, 0.0),
+        ),
+        torus_major_radius,
+        torus_minor_radius,
+    )
+    out_data[i, j, k] = sdf_smooth_min(box, torus, smooth_min_radius)
+
+
+class Example:
+    def __init__(self, stage_path="example_marching_cubes.usd", verbose=False):
+        self.verbose = verbose
+
+        self.dim = 64
+        self.max_verts = int(1e6)
+        self.max_tris = int(1e6)
+
+        self.torus_altitude = -0.5
+        self.torus_major_radius = 0.5
+        self.torus_minor_radius = 0.1
+        self.smooth_min_radius = 0.5
+
+        self.fps = 60
+        self.frame = 0
+
+        self.field = wp.zeros((self.dim, self.dim, self.dim), dtype=float)
+        self.mc = wp.MarchingCubes(self.dim, self.dim, self.dim, self.max_verts, self.max_tris)
+
+        self.renderer = None
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path)
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            with wp.ScopedTimer("Update Field", active=self.verbose):
+                wp.launch(
+                    make_field,
+                    dim=self.field.shape,
+                    inputs=(
+                        self.torus_altitude,
+                        self.torus_major_radius,
+                        self.torus_minor_radius,
+                        self.smooth_min_radius,
+                        self.dim,
+                        self.frame / self.fps,
+                    ),
+                    outputs=(self.field,),
+                )
+
+            with wp.ScopedTimer("Surface Extraction", active=self.verbose):
+                self.mc.surface(self.field, 0.0)
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("Render"):
+            self.renderer.begin_frame(self.frame / self.fps)
+            self.renderer.render_mesh(
+                "surface",
+                self.mc.verts.numpy(),
+                self.mc.indices.numpy(),
+                colors=(0.35, 0.55, 0.9),
+                update_topology=True,
+            )
+            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_marching_cubes.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=240, help="Total number of frames.")
+    parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, verbose=args.verbose)
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+            example.frame += 1
+
+        if example.renderer is not None:
+            example.renderer.save()

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()