warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/examples/core/example_mesh.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# 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.render
+
+wp.init()
+
+
+@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):
+        rng = np.random.default_rng(42)
+        self.num_particles = 1000
+
+        self.sim_steps = 500
+        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(os.path.dirname(__file__), "../assets/bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/bunny/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:
+            self.renderer = wp.render.UsdRenderer(stage)
+
+    def step(self):
+        with wp.ScopedTimer("step", detailed=False, 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", detailed=False):
+            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),) * len(self.mesh.points))
+            self.renderer.render_points(name="points", points=self.positions.numpy(), radius=self.sim_margin, colors=((0.8, 0.3, 0.2),) * len(self.mesh.points))
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_mesh.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.sim_steps):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/core/example_mesh_intersect.py

@@ -0,0 +1,193 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#############################################################################
+# 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.render
+
+
+wp.init()
+
+
+@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):
+        rng = np.random.default_rng(42)
+
+        self.query_count = 1024
+        self.has_queried = False
+
+        self.path_0 = "../assets/cube.usda"
+        self.path_1 = "../assets/sphere.usda"
+
+        self.mesh_0 = self.load_mesh(self.path_0, "/Cube/Cube_001")
+        self.mesh_1 = self.load_mesh(self.path_1, "/Sphere/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 i in range(self.query_count):
+            # random offset
+            p = wp.vec3(rng.random(3) * 0.5 - 0.5) * 5.0
+
+            # random orientation
+            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)
+
+            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
+        self.renderer = None
+        if stage is not None:
+            self.renderer = wp.render.UsdRenderer(stage)
+
+    def step(self):
+        with wp.ScopedTimer("intersect", active=True):
+            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",
+                    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=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=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_path = os.path.join(os.path.dirname(__file__), path)
+        usd_stage = Usd.Stage.Open(usd_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__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_mesh_intersect.usd")
+
+    example = Example(stage_path)
+
+    example.step()
+    example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/core/example_nvdb.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# 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 math
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.render
+
+wp.init()
+
+
+@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, 0.0, -980.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[2] < 0.0:
+        xpred = wp.vec3(xpred[0], xpred[1], 0.0)
+
+    # pbd update
+    v = (xpred - x) * (1.0 / dt)
+    x = xpred
+
+    positions[tid] = x
+    velocities[tid] = v
+
+
+class Example:
+    def __init__(self, stage):
+        rng = np.random.default_rng(42)
+        self.num_particles = 10000
+
+        self.sim_steps = 1000
+        frame_dt = 1.0 / 60.0
+        self.sim_substeps = 3
+        self.sim_dt = frame_dt / self.sim_substeps
+
+        self.sim_time = 0.0
+        self.sim_timers = {}
+
+        self.sim_margin = 15.0
+
+        init_pos = 1000.0 * (rng.random((self.num_particles, 3)) * 2.0 - 1.0) + np.array((0.0, 0.0, 3000.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
+        file = open(os.path.join(os.path.dirname(__file__), "../assets/rocks.nvdb"), "rb")
+
+        # create Volume object
+        self.volume = wp.Volume.load_from_nvdb(file)
+
+        file.close()
+
+        # renderer
+        self.renderer = None
+        if stage:
+            self.renderer = wp.render.UsdRenderer(stage, up_axis="z")
+            self.renderer.render_ground(size=10000.0)
+
+    def step(self):
+        with wp.ScopedTimer("step", detailed=False, 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", detailed=False):
+            self.renderer.begin_frame(self.sim_time)
+
+            self.renderer.render_ref(
+                name="collision",
+                path=os.path.join(os.path.dirname(__file__), "../assets/rocks.usd"),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
+                scale=wp.vec3(1.0, 1.0, 1.0),
+            )
+            self.renderer.render_points(name="points", points=self.positions.numpy(), radius=self.sim_margin, colors=((0.8, 0.3, 0.2),) * len(self.positions))
+
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_nvdb.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.sim_steps):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/core/example_raycast.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#############################################################################
+# 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
+
+wp.init()
+
+
+@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, **kwargs):
+        self.width = 1024
+        self.height = 1024
+        self.cam_pos = (0.0, 1.0, 2.0)
+
+        asset_stage = Usd.Stage.Open(os.path.join(os.path.dirname(__file__), "../assets/bunny.usd"))
+        mesh_geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/bunny/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 step(self):
+        pass
+
+    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 matplotlib.pyplot as plt
+
+    example = Example()
+    example.render()
+
+    plt.imshow(
+        example.pixels.numpy().reshape((example.height, example.width, 3)), origin="lower", interpolation="antialiased"
+    )
+    plt.show()