warp-lang 1.0.1__py3-none-macosx_10_13_universal2.whl → 1.1.0__py3-none-macosx_10_13_universal2.whl

warp/examples/core/example_nvdb.py

@@ -1,169 +1,176 @@
-# 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))
-
-            self.renderer.end_frame()
-
-
-if __name__ == "__main__":
-    stage_path = os.path.join(wp.examples.get_output_directory(), "example_nvdb.usd")
-
-    example = Example(stage_path)
-
-    for i in range(example.sim_steps):
-        example.step()
-        example.render()
-
-    if example.renderer:
-        example.renderer.save()
+# 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 os
+
+import numpy as np
+
+import warp as wp
+import warp.examples
+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, -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

@@ -1,89 +1,105 @@
-# 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()
+# 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
+import warp.examples
+
+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, 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()