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

warp/examples/core/example_torch.py

@@ -0,0 +1,181 @@
+# 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 Torch
+#
+# Optimizes the Rosenbrock function using the PyTorch Adam optimizer
+# The Rosenbrock function is a non-convex function, and is often used
+# to test optimization algorithms. The function is defined as:
+#              f(x, y) = (a - x)^2 + b * (y - x^2)^2
+# where a = 1 and b = 100. The minimum value of the function is 0 at (1, 1).
+#
+# The example demonstrates how to set up a torch.autograd.Function to 
+# incorporate Warp kernel launches within a PyTorch graph.
+###########################################################################
+
+import warp as wp
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+import torch
+
+wp.init()
+
+
+pvec2 = wp.types.vector(length=2, dtype=wp.float32)
+
+# Define the Rosenbrock function
+@wp.func
+def rosenbrock(x: float, y: float):
+    return (1.0 - x) ** 2.0 + 100.0 * (y - x**2.0) ** 2.0
+
+@wp.kernel
+def eval_rosenbrock(
+    xs: wp.array(dtype=pvec2),
+    # outputs
+    z: wp.array(dtype=float),
+):
+    i = wp.tid()
+    x = xs[i]
+    z[i] = rosenbrock(x[0], x[1])
+
+
+class Rosenbrock(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, xy, num_particles):
+
+        ctx.xy = wp.from_torch(xy, dtype=pvec2, requires_grad=True)
+        ctx.num_particles = num_particles
+
+        # allocate output
+        ctx.z = wp.zeros(num_particles, requires_grad=True)
+
+        wp.launch(
+            kernel=eval_rosenbrock,
+            dim=ctx.num_particles,
+            inputs=[ctx.xy],
+            outputs=[ctx.z]
+        )
+
+        return wp.to_torch(ctx.z)
+
+    @staticmethod
+    def backward(ctx, adj_z):
+
+        # map incoming Torch grads to our output variables
+        ctx.z.grad = wp.from_torch(adj_z)
+
+        wp.launch(
+            kernel=eval_rosenbrock,
+            dim=ctx.num_particles,
+            inputs=[ctx.xy],
+            outputs=[ctx.z],
+            adj_inputs=[ctx.xy.grad],
+            adj_outputs=[ctx.z.grad],
+            adjoint=True
+        )
+
+        # return adjoint w.r.t. inputs
+        return (wp.to_torch(ctx.xy.grad), None)
+
+
+class Example:
+
+    def __init__(self):
+
+        self.num_particles = 1500
+
+        min_x, max_x = -2.0, 2.0
+        min_y, max_y = -2.0, 2.0
+
+        # Create a grid of points
+        x = np.linspace(min_x, max_x, 100)
+        y = np.linspace(min_y, max_y, 100)
+        X, Y = np.meshgrid(x, y)
+        xy = np.column_stack((X.flatten(), Y.flatten()))
+        N = len(xy)
+        
+        xy = wp.array(xy, dtype=pvec2)
+        Z = wp.empty(N, dtype=wp.float32)
+
+        wp.launch(eval_rosenbrock, dim=N, inputs=[xy], outputs=[Z])
+        Z = Z.numpy().reshape(X.shape)
+
+        # Plot the function as a heatmap
+        self.fig = plt.figure(figsize=(6, 6))
+        ax = plt.gca()
+        plt.imshow(
+            Z,
+            extent=[min_x, max_x, min_y, max_y],
+            origin="lower",
+            interpolation="bicubic",
+            cmap="coolwarm",
+        )
+        plt.contour(
+            X,
+            Y,
+            Z,
+            extent=[min_x, max_x, min_y, max_y],
+            levels=150,
+            colors="k",
+            alpha=0.5,
+            linewidths=0.5,
+        )
+
+        plt.title("Rosenbrock function")
+        plt.xlabel("x")
+        plt.ylabel("y")
+
+        # Create a scatter plot (initially empty)
+        self.scat = ax.scatter([], [], c="k", s=2)
+
+        # Plot optimum
+        plt.plot(1, 1, "*", color="r", markersize=10)
+
+        self.learning_rate = 5e-2
+
+        self.torch_device = wp.device_to_torch(wp.get_device())
+
+        rng = np.random.default_rng(42)
+        self.xy = torch.tensor(rng.normal(size=(self.num_particles, 2)), dtype=torch.float32, requires_grad=True, device=self.torch_device)
+        self.opt = torch.optim.Adam([self.xy], lr=self.learning_rate)
+
+    def forward(self):
+        self.z = Rosenbrock.apply(self.xy, self.num_particles)
+
+    def step(self):
+        self.opt.zero_grad()
+        self.forward()
+        self.z.backward(torch.ones_like(self.z))
+
+        self.opt.step()
+
+    def render(self):
+        # Update the scatter plot
+        xy_np = self.xy.numpy(force=True)
+        self.scat.set_offsets(np.c_[xy_np[:, 0], xy_np[:, 1]])
+
+        print(f"\rParticle mean: {np.mean(xy_np, axis=0)}     ", end="")
+
+    # Function to update the scatter plot
+    def step_and_render(self, frame):
+        for _ in range(10):
+            self.step()
+
+        self.render()
+
+
+if __name__ == "__main__":
+
+    example = Example()
+
+    # Create the animation
+    ani = FuncAnimation(example.fig, example.step_and_render, frames=10000, interval=200)
+
+    # Display the animation
+    plt.show()

warp/examples/core/example_wave.py

@@ -0,0 +1,249 @@
+# 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 Wave
+#
+# Shows how to implement a simple 2D wave-equation solver with collision
+# against a moving sphere.
+#
+###########################################################################
+
+import math
+import os
+
+import warp as wp
+import warp.render
+
+wp.init()
+
+
+@wp.func
+def sample(f: wp.array(dtype=float), x: int, y: int, width: int, height: int):
+    # clamp texture coords
+    x = wp.clamp(x, 0, width - 1)
+    y = wp.clamp(y, 0, height - 1)
+
+    s = f[y * width + x]
+    return s
+
+
+@wp.func
+def laplacian(f: wp.array(dtype=float), x: int, y: int, width: int, height: int):
+    ddx = sample(f, x + 1, y, width, height) - 2.0 * sample(f, x, y, width, height) + sample(f, x - 1, y, width, height)
+    ddy = sample(f, x, y + 1, width, height) - 2.0 * sample(f, x, y, width, height) + sample(f, x, y - 1, width, height)
+
+    return ddx + ddy
+
+
+@wp.kernel
+def wave_displace(
+    hcurrent: wp.array(dtype=float),
+    hprevious: wp.array(dtype=float),
+    width: int,
+    height: int,
+    center_x: float,
+    center_y: float,
+    r: float,
+    mag: float,
+    t: float,
+):
+    tid = wp.tid()
+
+    x = tid % width
+    y = tid // width
+
+    dx = float(x) - center_x
+    dy = float(y) - center_y
+
+    dist_sq = float(dx * dx + dy * dy)
+
+    if dist_sq < r * r:
+        h = mag * wp.sin(t)
+
+        hcurrent[tid] = h
+        hprevious[tid] = h
+
+
+@wp.kernel
+def wave_solve(
+    hprevious: wp.array(dtype=float),
+    hcurrent: wp.array(dtype=float),
+    width: int,
+    height: int,
+    inv_cell: float,
+    k_speed: float,
+    k_damp: float,
+    dt: float,
+):
+    tid = wp.tid()
+
+    x = tid % width
+    y = tid // width
+
+    l = laplacian(hcurrent, x, y, width, height) * inv_cell * inv_cell
+
+    # integrate
+    h1 = hcurrent[tid]
+    h0 = hprevious[tid]
+
+    h = 2.0 * h1 - h0 + dt * dt * (k_speed * l - k_damp * (h1 - h0))
+
+    # buffers get swapped each iteration
+    hprevious[tid] = h
+
+
+# simple kernel to apply height deltas to a vertex array
+@wp.kernel
+def grid_update(heights: wp.array(dtype=float), vertices: wp.array(dtype=wp.vec3)):
+    tid = wp.tid()
+
+    h = heights[tid]
+    v = vertices[tid]
+
+    v_new = wp.vec3(v[0], h, v[2])
+
+    vertices[tid] = v_new
+
+
+class Example:
+    def __init__(self, stage):
+        self.sim_width = 128
+        self.sim_height = 128
+
+        self.sim_fps = 60.0
+        self.sim_substeps = 16
+        self.sim_duration = 5.0
+        self.sim_frames = int(self.sim_duration * self.sim_fps)
+        self.sim_dt = (1.0 / self.sim_fps) / self.sim_substeps
+        self.sim_time = 0.0
+
+        # wave constants
+        self.k_speed = 1.0
+        self.k_damp = 0.0
+
+        # grid constants
+        self.grid_size = 0.1
+        self.grid_displace = 0.5
+
+        vertices = []
+        self.indices = []
+
+        def grid_index(x, y, stride):
+            return y * stride + x
+
+        for z in range(self.sim_height):
+            for x in range(self.sim_width):
+                pos = (
+                    float(x) * self.grid_size,
+                    0.0,
+                    float(z) * self.grid_size,
+                )
+
+                # directly modifies verts_host memory since this is a numpy alias of the same buffer
+                vertices.append(pos)
+
+                if x > 0 and z > 0:
+                    self.indices.append(grid_index(x - 1, z - 1, self.sim_width))
+                    self.indices.append(grid_index(x, z, self.sim_width))
+                    self.indices.append(grid_index(x, z - 1, self.sim_width))
+
+                    self.indices.append(grid_index(x - 1, z - 1, self.sim_width))
+                    self.indices.append(grid_index(x - 1, z, self.sim_width))
+                    self.indices.append(grid_index(x, z, self.sim_width))
+
+        # simulation grids
+        self.sim_grid0 = wp.zeros(self.sim_width * self.sim_height, dtype=float)
+        self.sim_grid1 = wp.zeros(self.sim_width * self.sim_height, dtype=float)
+        self.sim_verts = wp.array(vertices, dtype=wp.vec3)
+
+        # create surface displacement around a point
+        self.cx = self.sim_width / 2 + math.sin(self.sim_time) * self.sim_width / 3
+        self.cy = self.sim_height / 2 + math.cos(self.sim_time) * self.sim_height / 3
+
+        self.renderer = None
+        if stage:
+            self.renderer = wp.render.UsdRenderer(stage)
+    
+    def step(self):
+        with wp.ScopedTimer("step", active=True):
+            for s in range(self.sim_substeps):
+                # create surface displacement around a point
+                self.cx = self.sim_width / 2 + math.sin(self.sim_time) * self.sim_width / 3
+                self.cy = self.sim_height / 2 + math.cos(self.sim_time) * self.sim_height / 3
+
+                wp.launch(
+                    kernel=wave_displace,
+                    dim=self.sim_width * self.sim_height,
+                    inputs=[
+                        self.sim_grid0,
+                        self.sim_grid1,
+                        self.sim_width,
+                        self.sim_height,
+                        self.cx,
+                        self.cy,
+                        10.0,
+                        self.grid_displace,
+                        -math.pi * 0.5,
+                    ],
+                )
+
+                # integrate wave equation
+                wp.launch(
+                    kernel=wave_solve,
+                    dim=self.sim_width * self.sim_height,
+                    inputs=[
+                        self.sim_grid0,
+                        self.sim_grid1,
+                        self.sim_width,
+                        self.sim_height,
+                        1.0 / self.grid_size,
+                        self.k_speed,
+                        self.k_damp,
+                        self.sim_dt,
+                    ],
+                )
+
+                # swap grids
+                (self.sim_grid0, self.sim_grid1) = (self.sim_grid1, self.sim_grid0)
+
+                self.sim_time += self.sim_dt
+
+        with wp.ScopedTimer("mesh", active=False):
+            # update grid vertices from heights
+            wp.launch(kernel=grid_update, dim=self.sim_width * self.sim_height, inputs=[self.sim_grid0, self.sim_verts])
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render", active=True):
+            vertices = self.sim_verts.numpy()
+
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render_mesh("surface", vertices, self.indices, colors=((0.35, 0.55, 0.9),) * len(vertices))
+            self.renderer.render_sphere(
+                "sphere",
+                (self.cx * self.grid_size, 0.0, self.cy * self.grid_size),
+                (0.0, 0.0, 0.0, 1.0),
+                10.0 * self.grid_size,
+                color=(1.0, 1.0, 1.0),
+            )
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_wave.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.sim_frames):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()