warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/core/example_dem.py

@@ -0,0 +1,241 @@
+# 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 DEM
+#
+# Shows how to implement a DEM particle simulation with cohesion between
+# particles. Neighbors are found using the wp.HashGrid class, and
+# wp.hash_grid_query(), wp.hash_grid_query_next() kernel methods.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.render
+
+
+@wp.func
+def contact_force(n: wp.vec3, v: wp.vec3, c: float, k_n: float, k_d: float, k_f: float, k_mu: float):
+    vn = wp.dot(n, v)
+    jn = c * k_n
+    jd = min(vn, 0.0) * k_d
+
+    # contact force
+    fn = jn + jd
+
+    # friction force
+    vt = v - n * vn
+    vs = wp.length(vt)
+
+    if vs > 0.0:
+        vt = vt / vs
+
+    # Coulomb condition
+    ft = wp.min(vs * k_f, k_mu * wp.abs(fn))
+
+    # total force
+    return -n * fn - vt * ft
+
+
+@wp.kernel
+def apply_forces(
+    grid: wp.uint64,
+    particle_x: wp.array(dtype=wp.vec3),
+    particle_v: wp.array(dtype=wp.vec3),
+    particle_f: wp.array(dtype=wp.vec3),
+    radius: float,
+    k_contact: float,
+    k_damp: float,
+    k_friction: float,
+    k_mu: float,
+):
+    tid = wp.tid()
+
+    # order threads by cell
+    i = wp.hash_grid_point_id(grid, tid)
+
+    x = particle_x[i]
+    v = particle_v[i]
+
+    f = wp.vec3()
+
+    # ground contact
+    n = wp.vec3(0.0, 1.0, 0.0)
+    c = wp.dot(n, x)
+
+    cohesion_ground = 0.02
+    cohesion_particle = 0.0075
+
+    if c < cohesion_ground:
+        f = f + contact_force(n, v, c, k_contact, k_damp, 100.0, 0.5)
+
+    # particle contact
+    neighbors = wp.hash_grid_query(grid, x, radius * 5.0)
+
+    for index in neighbors:
+        if index != i:
+            # compute distance to point
+            n = x - particle_x[index]
+            d = wp.length(n)
+            err = d - radius * 2.0
+
+            if err <= cohesion_particle:
+                n = n / d
+                vrel = v - particle_v[index]
+
+                f = f + contact_force(n, vrel, err, k_contact, k_damp, k_friction, k_mu)
+
+    particle_f[i] = f
+
+
+@wp.kernel
+def integrate(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    f: wp.array(dtype=wp.vec3),
+    gravity: wp.vec3,
+    dt: float,
+    inv_mass: float,
+):
+    tid = wp.tid()
+
+    v_new = v[tid] + f[tid] * inv_mass * dt + gravity * dt
+    x_new = x[tid] + v_new * dt
+
+    v[tid] = v_new
+    x[tid] = x_new
+
+
+class Example:
+    def __init__(self, stage_path="example_dem.usd"):
+        fps = 60
+        self.frame_dt = 1.0 / fps
+
+        self.sim_substeps = 64
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+
+        self.point_radius = 0.1
+
+        self.k_contact = 8000.0
+        self.k_damp = 2.0
+        self.k_friction = 1.0
+        self.k_mu = 100000.0  # for cohesive materials
+
+        self.inv_mass = 64.0
+
+        self.grid = wp.HashGrid(128, 128, 128)
+        self.grid_cell_size = self.point_radius * 5.0
+
+        self.points = self.particle_grid(32, 64, 32, (0.0, 0.5, 0.0), self.point_radius, 0.1)
+
+        self.x = wp.array(self.points, dtype=wp.vec3)
+        self.v = wp.array(np.ones([len(self.x), 3]) * np.array([0.0, 0.0, 15.0]), dtype=wp.vec3)
+        self.f = wp.zeros_like(self.v)
+
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path)
+            self.renderer.render_ground()
+        else:
+            self.renderer = None
+
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            wp.launch(
+                kernel=apply_forces,
+                dim=len(self.x),
+                inputs=[
+                    self.grid.id,
+                    self.x,
+                    self.v,
+                    self.f,
+                    self.point_radius,
+                    self.k_contact,
+                    self.k_damp,
+                    self.k_friction,
+                    self.k_mu,
+                ],
+            )
+            wp.launch(
+                kernel=integrate,
+                dim=len(self.x),
+                inputs=[self.x, self.v, self.f, (0.0, -9.8, 0.0), self.sim_dt, self.inv_mass],
+            )
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            with wp.ScopedTimer("grid build", active=False):
+                self.grid.build(self.x, self.grid_cell_size)
+
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+
+            self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render_points(
+                points=self.x.numpy(), radius=self.point_radius, name="points", colors=(0.8, 0.3, 0.2)
+            )
+            self.renderer.end_frame()
+
+    # creates a grid of particles
+    def particle_grid(self, dim_x, dim_y, dim_z, lower, radius, jitter):
+        rng = np.random.default_rng(42)
+        points = np.meshgrid(np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z))
+        points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
+        points_t = points_t + rng.random(size=points_t.shape) * radius * jitter
+
+        return points_t.reshape((-1, 3))
+
+
+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_dem.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=200, 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_fluid.py

@@ -0,0 +1,299 @@
+# 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 Fluid
+#
+# Shows how to implement a simple 2D Stable Fluids solver using
+# multidimensional arrays and launches.
+#
+###########################################################################
+
+import math
+
+import warp as wp
+import warp.render
+
+grid_width = wp.constant(256)
+grid_height = wp.constant(128)
+
+
+@wp.func
+def lookup_float(f: wp.array2d(dtype=float), x: int, y: int):
+    x = wp.clamp(x, 0, grid_width - 1)
+    y = wp.clamp(y, 0, grid_height - 1)
+
+    return f[x, y]
+
+
+@wp.func
+def sample_float(f: wp.array2d(dtype=float), x: float, y: float):
+    lx = int(wp.floor(x))
+    ly = int(wp.floor(y))
+
+    tx = x - float(lx)
+    ty = y - float(ly)
+
+    s0 = wp.lerp(lookup_float(f, lx, ly), lookup_float(f, lx + 1, ly), tx)
+    s1 = wp.lerp(lookup_float(f, lx, ly + 1), lookup_float(f, lx + 1, ly + 1), tx)
+
+    s = wp.lerp(s0, s1, ty)
+    return s
+
+
+@wp.func
+def lookup_vel(f: wp.array2d(dtype=wp.vec2), x: int, y: int):
+    if x < 0 or x >= grid_width:
+        return wp.vec2()
+    if y < 0 or y >= grid_height:
+        return wp.vec2()
+
+    return f[x, y]
+
+
+@wp.func
+def sample_vel(f: wp.array2d(dtype=wp.vec2), x: float, y: float):
+    lx = int(wp.floor(x))
+    ly = int(wp.floor(y))
+
+    tx = x - float(lx)
+    ty = y - float(ly)
+
+    s0 = wp.lerp(lookup_vel(f, lx, ly), lookup_vel(f, lx + 1, ly), tx)
+    s1 = wp.lerp(lookup_vel(f, lx, ly + 1), lookup_vel(f, lx + 1, ly + 1), tx)
+
+    s = wp.lerp(s0, s1, ty)
+    return s
+
+
+@wp.kernel
+def advect(
+    u0: wp.array2d(dtype=wp.vec2),
+    u1: wp.array2d(dtype=wp.vec2),
+    rho0: wp.array2d(dtype=float),
+    rho1: wp.array2d(dtype=float),
+    dt: float,
+):
+    i, j = wp.tid()
+
+    u = u0[i, j]
+
+    # trace backward
+    p = wp.vec2(float(i), float(j))
+    p = p - u * dt
+
+    # advect
+    u1[i, j] = sample_vel(u0, p[0], p[1])
+    rho1[i, j] = sample_float(rho0, p[0], p[1])
+
+
+@wp.kernel
+def divergence(u: wp.array2d(dtype=wp.vec2), div: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+
+    if i == grid_width - 1:
+        return
+    if j == grid_height - 1:
+        return
+
+    dx = (u[i + 1, j][0] - u[i, j][0]) * 0.5
+    dy = (u[i, j + 1][1] - u[i, j][1]) * 0.5
+
+    div[i, j] = dx + dy
+
+
+@wp.kernel
+def pressure_solve(p0: wp.array2d(dtype=float), p1: wp.array2d(dtype=float), div: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+
+    s1 = lookup_float(p0, i - 1, j)
+    s2 = lookup_float(p0, i + 1, j)
+    s3 = lookup_float(p0, i, j - 1)
+    s4 = lookup_float(p0, i, j + 1)
+
+    # Jacobi update
+    err = s1 + s2 + s3 + s4 - div[i, j]
+
+    p1[i, j] = err * 0.25
+
+
+@wp.kernel
+def pressure_apply(p: wp.array2d(dtype=float), u: wp.array2d(dtype=wp.vec2)):
+    i, j = wp.tid()
+
+    if i == 0 or i == grid_width - 1:
+        return
+    if j == 0 or j == grid_height - 1:
+        return
+
+    # pressure gradient
+    f_p = wp.vec2(p[i + 1, j] - p[i - 1, j], p[i, j + 1] - p[i, j - 1]) * 0.5
+
+    u[i, j] = u[i, j] - f_p
+
+
+@wp.kernel
+def integrate(u: wp.array2d(dtype=wp.vec2), rho: wp.array2d(dtype=float), dt: float):
+    i, j = wp.tid()
+
+    # gravity
+    f_g = wp.vec2(-90.8, 0.0) * rho[i, j]
+
+    # integrate
+    u[i, j] = u[i, j] + dt * f_g
+
+    # fade
+    rho[i, j] = rho[i, j] * (1.0 - 0.1 * dt)
+
+
+@wp.kernel
+def init(rho: wp.array2d(dtype=float), u: wp.array2d(dtype=wp.vec2), radius: int, dir: wp.vec2):
+    i, j = wp.tid()
+
+    d = wp.length(wp.vec2(float(i - grid_width / 2), float(j - grid_height / 2)))
+
+    if d < radius:
+        rho[i, j] = 1.0
+        u[i, j] = dir
+
+
+class Example:
+    def __init__(self):
+        fps = 60
+        self.frame_dt = 1.0 / fps
+        self.sim_substeps = 2
+        self.iterations = 100  # Number of pressure iterations
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+
+        shape = (grid_width, grid_height)
+
+        self.u0 = wp.zeros(shape, dtype=wp.vec2)
+        self.u1 = wp.zeros(shape, dtype=wp.vec2)
+
+        self.rho0 = wp.zeros(shape, dtype=float)
+        self.rho1 = wp.zeros(shape, dtype=float)
+
+        self.p0 = wp.zeros(shape, dtype=float)
+        self.p1 = wp.zeros(shape, dtype=float)
+        self.div = wp.zeros(shape, dtype=float)
+
+        # capture pressure solve as a CUDA graph
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.pressure_iterations()
+            self.graph = capture.graph
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            for _ in range(self.sim_substeps):
+                shape = (grid_width, grid_height)
+                dt = self.sim_dt
+
+                speed = 400.0
+                angle = math.sin(self.sim_time * 4.0) * 1.5
+                vel = wp.vec2(math.cos(angle) * speed, math.sin(angle) * speed)
+
+                # update emitters
+                wp.launch(init, dim=shape, inputs=[self.rho0, self.u0, 5, vel])
+
+                # force integrate
+                wp.launch(integrate, dim=shape, inputs=[self.u0, self.rho0, dt])
+                wp.launch(divergence, dim=shape, inputs=[self.u0, self.div])
+
+                # pressure solve
+                self.p0.zero_()
+                self.p1.zero_()
+
+                if self.use_cuda_graph:
+                    wp.capture_launch(self.graph)
+                else:
+                    self.pressure_iterations()
+
+                # velocity update
+                wp.launch(pressure_apply, dim=shape, inputs=[self.p0, self.u0])
+
+                # semi-Lagrangian advection
+                wp.launch(advect, dim=shape, inputs=[self.u0, self.u1, self.rho0, self.rho1, dt])
+
+                # swap buffers
+                (self.u0, self.u1) = (self.u1, self.u0)
+                (self.rho0, self.rho1) = (self.rho1, self.rho0)
+
+                self.sim_time += dt
+
+    def pressure_iterations(self):
+        for _ in range(self.iterations):
+            wp.launch(pressure_solve, dim=self.p0.shape, inputs=[self.p0, self.p1, self.div])
+
+            # swap pressure fields
+            (self.p0, self.p1) = (self.p1, self.p0)
+
+    def step_and_render_frame(self, frame_num=None, img=None):
+        self.step()
+
+        with wp.ScopedTimer("render"):
+            if img:
+                img.set_array(self.rho0.numpy())
+
+        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=100000, 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 args.headless:
+            for _ in range(args.num_frames):
+                example.step()
+        else:
+            import matplotlib
+            import matplotlib.animation as anim
+            import matplotlib.pyplot as plt
+
+            fig = plt.figure()
+
+            img = plt.imshow(
+                example.rho0.numpy(),
+                origin="lower",
+                animated=True,
+                interpolation="antialiased",
+            )
+            img.set_norm(matplotlib.colors.Normalize(0.0, 1.0))
+            seq = anim.FuncAnimation(
+                fig,
+                example.step_and_render_frame,
+                fargs=(img,),
+                frames=args.num_frames,
+                blit=True,
+                interval=8,
+                repeat=False,
+            )
+
+            plt.show()

warp/examples/core/example_graph_capture.py

@@ -0,0 +1,150 @@
+# 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 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()