warp-lang 1.0.2__py3-none-win_amd64.whl → 1.2.0__py3-none-win_amd64.whl

warp/examples/core/example_graph_capture.py

@@ -1,126 +1,142 @@
-# 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 Graph Capture
-#
-# Shows how to implement CUDA graph capture using wp.ScopedCapture().
-#
-###########################################################################
-
-import warp as wp
-import numpy as np
-import matplotlib.pyplot as plt
-import matplotlib.colors
-from matplotlib.animation import FuncAnimation
-
-wp.init()
-
-
-@wp.kernel
-def fbm(
-    kernel_seed: int,
-    frequency: float,
-    amplitude: float,
-    x: wp.array(dtype=float),
-    y: wp.array(dtype=float),
-    z: wp.array2d(dtype=float),
-):
-    i, j = wp.tid()
-    state = wp.rand_init(kernel_seed)
-
-    p = frequency * wp.vec2(x[j], y[i])
-    n = amplitude * wp.noise(state, p)
-
-    z[i, j] += n
-
-
-@wp.kernel
-def slide(x: wp.array(dtype=float), shift: float):
-    tid = wp.tid()
-    x[tid] += shift
-
-
-class Example:
-    def __init__(self):
-        self.W = 128
-        self.H = 128
-        min_x, max_x = 0.0, 2.0
-        min_y, max_y = 0.0, 2.0
-
-        # create a grid of pixels
-        x = np.linspace(min_x, max_x, self.W)
-        y = np.linspace(min_y, max_y, self.H)
-
-        self.x = wp.array(x, dtype=float)
-        self.y = wp.array(y, dtype=float)
-        self.pixel_values = wp.zeros((self.W, self.H), dtype=float)
-
-        self.seed = 42
-        self.shift = 2e-2
-        self.frequency = 1.0
-        self.amplitude = 1.0
-
-        # use graph capture if launching from a CUDA-capable device
-        self.use_graph = wp.get_device().is_cuda
-        if self.use_graph:
-            # record launches
-            with wp.ScopedCapture() as capture:
-                self.fbm()
-            self.graph = capture.graph
-
-    def fbm(self):
-        for _ in range(16):
-            wp.launch(
-                kernel=fbm,
-                dim=(self.H, self.W),
-                inputs=[self.seed, self.frequency, self.amplitude, self.x, self.y],
-                outputs=[self.pixel_values],
-            )
-            self.frequency *= 2.0
-            self.amplitude *= 0.5
-
-    def step(self):
-        self.pixel_values.zero_()
-        self.frequency = 1.0
-        self.amplitude = 1.0
-
-        with wp.ScopedTimer("step", active=True):
-            wp.launch(kernel=slide, dim=self.W, inputs=[self.x, self.shift])
-
-            if self.use_graph:
-                wp.capture_launch(self.graph)
-            else:  # cpu path
-                self.fbm()
-
-    def render(self):
-        pass
-
-    def step_and_render(self, frame_num=None, img=None):
-        self.step()
-
-        with wp.ScopedTimer("render"):
-            if img:
-                pixels = self.pixel_values.numpy()
-                pixels = (pixels + 1.0) / 2.0
-                img.set_array(pixels)
-
-        return (img,)
-
-
-if __name__ == "__main__":
-    example = Example()
-
-    # Create the animation
-    fig = plt.figure()
-    img = plt.imshow(example.pixel_values.numpy(), "gray", origin="lower", animated=True)
-    img.set_norm(matplotlib.colors.Normalize(0.0, 1.0))
-
-    ani = FuncAnimation(fig, example.step_and_render, fargs=(img,), frames=1000, interval=30)
-
-    # Display the animation
-    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 Graph Capture
+#
+# Shows how to implement CUDA graph capture using wp.ScopedCapture().
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+
+@wp.kernel
+def fbm(
+    kernel_seed: int,
+    frequency: float,
+    amplitude: float,
+    x: wp.array(dtype=float),
+    y: wp.array(dtype=float),
+    z: wp.array2d(dtype=float),
+):
+    i, j = wp.tid()
+    state = wp.rand_init(kernel_seed)
+
+    p = frequency * wp.vec2(x[j], y[i])
+    n = amplitude * wp.noise(state, p)
+
+    z[i, j] += n
+
+
+@wp.kernel
+def slide(x: wp.array(dtype=float), shift: float):
+    tid = wp.tid()
+    x[tid] += shift
+
+
+class Example:
+    def __init__(self):
+        self.width = 128
+        self.height = 128
+        min_x, max_x = 0.0, 2.0
+        min_y, max_y = 0.0, 2.0
+
+        # create a grid of pixels
+        x = np.linspace(min_x, max_x, self.width)
+        y = np.linspace(min_y, max_y, self.height)
+
+        self.x = wp.array(x, dtype=float)
+        self.y = wp.array(y, dtype=float)
+        self.pixel_values = wp.zeros((self.width, self.height), dtype=float)
+
+        self.seed = 42
+        self.shift = 2e-2
+        self.frequency = 1.0
+        self.amplitude = 1.0
+
+        # use graph capture if launching from a CUDA-capable device
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            # record launches
+            with wp.ScopedCapture() as capture:
+                self.fbm()
+            self.graph = capture.graph
+
+    def fbm(self):
+        for _ in range(16):
+            wp.launch(
+                kernel=fbm,
+                dim=(self.height, self.width),
+                inputs=[self.seed, self.frequency, self.amplitude, self.x, self.y],
+                outputs=[self.pixel_values],
+            )
+            self.frequency *= 2.0
+            self.amplitude *= 0.5
+
+    def step(self):
+        self.pixel_values.zero_()
+        self.frequency = 1.0
+        self.amplitude = 1.0
+
+        with wp.ScopedTimer("step", active=True):
+            wp.launch(kernel=slide, dim=self.width, inputs=[self.x, self.shift])
+
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:  # cpu path
+                self.fbm()
+
+    def step_and_render(self, frame_num=None, img=None):
+        self.step()
+
+        with wp.ScopedTimer("render"):
+            if img:
+                pixels = self.pixel_values.numpy()
+                pixels = (pixels + 1.0) / 2.0
+                img.set_array(pixels)
+
+        return (img,)
+
+
+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("--num_frames", type=int, default=1000, help="Total number of frames.")
+    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()
+
+        if not args.headless:
+            import matplotlib.colors
+            import matplotlib.pyplot as plt
+            from matplotlib.animation import FuncAnimation
+
+            # Create the animation
+            fig = plt.figure()
+            img = plt.imshow(example.pixel_values.numpy(), "gray", origin="lower", animated=True)
+            img.set_norm(matplotlib.colors.Normalize(0.0, 1.0))
+
+            ani = FuncAnimation(fig, example.step_and_render, fargs=(img,), frames=1000, interval=30)
+
+            # Display the animation
+            plt.show()
+
+        else:
+            for _ in range(args.num_frames):
+                example.step()

warp/examples/core/example_marching_cubes.py

@@ -1,174 +1,186 @@
-# 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 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
-###########################################################################
-
-
-from typing import Optional
-
-import warp as wp
-import warp.render
-
-wp.init()
-
-
-@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: Optional[str]):
-        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.0
-        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 is not None:
-            self.renderer = wp.render.UsdRenderer(stage)
-
-    def step(self):
-        with wp.ScopedTimer("Update Field"):
-            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"):
-            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),) * len(self.mc.verts),
-                update_topology=True,
-            )
-            self.renderer.end_frame()
-
-
-if __name__ == "__main__":
-    stage_path = "example_marching_cubes.usd"
-
-    example = Example(stage_path)
-    for _ in range(240):
-        example.step()
-        example.render()
-        example.frame += 1
-
-    if example.renderer is not None:
-        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 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),) * len(self.mc.verts),
+                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()