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

warp/examples/core/example_raymarch.py

@@ -0,0 +1,178 @@
+# 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 March
+#
+# Shows how to implement an SDF ray marching based renderer. Please see
+# https://iquilezles.org/www/articles/distfunctions/distfunctions.htm
+# for reference on different distance functions.
+#
+###########################################################################
+
+
+import warp as wp
+
+wp.init()
+
+
+@wp.func
+def sdf_sphere(p: wp.vec3, r: float):
+    return wp.length(p) - r
+
+
+@wp.func
+def sdf_box(upper: wp.vec3, p: wp.vec3):
+    qx = wp.abs(p[0]) - upper[0]
+    qy = wp.abs(p[1]) - upper[1]
+    qz = wp.abs(p[2]) - upper[2]
+
+    e = wp.vec3(wp.max(qx, 0.0), wp.max(qy, 0.0), wp.max(qz, 0.0))
+
+    return wp.length(e) + wp.min(wp.max(qx, wp.max(qy, qz)), 0.0)
+
+
+@wp.func
+def sdf_plane(p: wp.vec3, plane: wp.vec4):
+    return plane[0] * p[0] + plane[1] * p[1] + plane[2] * p[2] + plane[3]
+
+
+@wp.func
+def op_union(d1: float, d2: float):
+    return wp.min(d1, d2)
+
+
+@wp.func
+def op_subtract(d1: float, d2: float):
+    return wp.max(-d1, d2)
+
+
+@wp.func
+def op_intersect(d1: float, d2: float):
+    return wp.max(d1, d2)
+
+
+# simple scene
+@wp.func
+def sdf(p: wp.vec3):
+    sphere_1 = wp.vec3(0.0, 0.75, 0.0)
+
+    d = op_subtract(sdf_sphere(p - sphere_1, 0.75), sdf_box(wp.vec3(1.0, 0.5, 0.5), p))
+
+    # ground plane
+    d = op_union(d, sdf_plane(p, wp.vec4(0.0, 1.0, 0.0, 1.0)))
+
+    return d
+
+
+@wp.func
+def normal(p: wp.vec3):
+    eps = 1.0e-5
+
+    # compute gradient of the SDF using finite differences
+    dx = sdf(p + wp.vec3(eps, 0.0, 0.0)) - sdf(p - wp.vec3(eps, 0.0, 0.0))
+    dy = sdf(p + wp.vec3(0.0, eps, 0.0)) - sdf(p - wp.vec3(0.0, eps, 0.0))
+    dz = sdf(p + wp.vec3(0.0, 0.0, eps)) - sdf(p - wp.vec3(0.0, 0.0, eps))
+
+    return wp.normalize(wp.vec3(dx, dy, dz))
+
+
+@wp.func
+def shadow(ro: wp.vec3, rd: wp.vec3):
+    t = float(0.0)
+    s = float(1.0)
+
+    for _ in range(64):
+        d = sdf(ro + t * rd)
+        t = t + wp.clamp(d, 0.0001, 2.0)
+
+        h = wp.clamp(4.0 * d / t, 0.0, 1.0)
+        s = wp.min(s, h * h * (3.0 - 2.0 * h))
+
+        if t > 8.0:
+            return 1.0
+
+    return s
+
+
+@wp.kernel
+def draw(cam_pos: wp.vec3, cam_rot: wp.quat, width: int, height: int, pixels: wp.array(dtype=wp.vec3)):
+    tid = wp.tid()
+
+    x = tid % width
+    y = tid // width
+
+    # compute pixel coordinates
+    sx = (2.0 * float(x) - float(width)) / float(height)
+    sy = (2.0 * float(y) - float(height)) / float(height)
+
+    # compute view ray
+    ro = cam_pos
+    rd = wp.quat_rotate(cam_rot, wp.normalize(wp.vec3(sx, sy, -2.0)))
+
+    t = float(0.0)
+
+    # ray march
+    for _ in range(128):
+        d = sdf(ro + rd * t)
+        t = t + d
+
+    if d < 0.01:
+        p = ro + rd * t
+        n = normal(p)
+        l = wp.normalize(wp.vec3(0.6, 0.4, 0.5))
+
+        # half-vector
+        h = wp.normalize(l - rd)
+
+        diffuse = wp.dot(n, l)
+        specular = wp.clamp(wp.dot(n, h), 0.0, 1.0) ** 80.0
+        fresnel = 0.04 + 0.96 * wp.clamp(1.0 - wp.dot(h, l), 0.0, 1.0) ** 5.0
+
+        intensity = 2.0
+        result = (
+            wp.vec3(0.85, 0.9, 0.95) * (diffuse * (1.0 - fresnel) + specular * fresnel * 10.0) * shadow(p, l) * intensity
+        )
+
+        # gamma
+        pixels[tid] = wp.vec3(wp.clamp(result[0] ** 2.2, 0.0, 1.0), wp.clamp(result[1] ** 2.2, 0.0, 1.0), wp.clamp(result[2] ** 2.2, 0.0, 1.0))
+
+    else:
+        pixels[tid] = wp.vec3(0.4, 0.45, 0.5) * 1.5
+
+
+class Example:
+    def __init__(self, **kwargs):
+        self.width = 2048
+        self.height = 1024
+        self.cam_pos = (-1.25, 1.0, 2.0)
+        self.cam_rot = wp.quat_rpy(-0.5, -0.5, 0.0)
+
+        self.pixels = wp.zeros(self.width * self.height, dtype=wp.vec3)
+
+    def step(self):
+        pass
+
+    def render(self):
+        with wp.ScopedTimer("render"):
+            wp.launch(
+                kernel=draw,
+                dim=self.width * self.height,
+                inputs=[self.cam_pos, self.cam_rot, 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()

warp/examples/core/example_render_opengl.py

@@ -0,0 +1,141 @@
+# 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.
+
+###########################################################################
+# OpenGL renderer example
+#
+# Demonstrates how to set up tiled rendering and retrieves the pixels from
+# OpenGLRenderer as a Warp array while keeping all memory on the GPU.
+#
+###########################################################################
+
+import warp as wp
+import warp.render
+import numpy as np
+
+wp.init()
+
+# number of viewports to render in a single frame
+num_tiles = 4
+# whether to split tiles into subplots
+split_up_tiles = True
+# whether to apply custom tile arrangement
+custom_tile_arrangement = False
+# whether to display the pixels in a matplotlib figure
+show_plot = True
+# whether to render depth image to a Warp array
+render_mode = "depth"
+
+renderer = wp.render.OpenGLRenderer(vsync=False)
+instance_ids = []
+
+if custom_tile_arrangement:
+    positions = []
+    sizes = []
+else:
+    positions = None
+    sizes = None
+
+if num_tiles > 1:
+    # set up instances to hide one of the capsules in each tile
+    for i in range(num_tiles):
+        instances = [j for j in np.arange(13) if j != i + 2]
+        instance_ids.append(instances)
+        if custom_tile_arrangement:
+            angle = np.pi * 2.0 / num_tiles * i
+            positions.append((int(np.cos(angle) * 150 + 250), int(np.sin(angle) * 150 + 250)))
+            sizes.append((150, 150))
+    renderer.setup_tiled_rendering(instance_ids, tile_positions=positions, tile_sizes=sizes)
+
+renderer.render_ground()
+
+channels = 1 if render_mode == "depth" else 3
+if show_plot:
+    import matplotlib.pyplot as plt
+
+    if split_up_tiles:
+        pixels = wp.zeros((num_tiles, renderer.tile_height, renderer.tile_width, channels), dtype=wp.float32)
+        ncols = int(np.ceil(np.sqrt(num_tiles)))
+        nrows = int(np.ceil(num_tiles / float(ncols)))
+        img_plots = []
+        aspect_ratio = renderer.tile_height / renderer.tile_width
+        fig, axes = plt.subplots(
+            ncols=ncols,
+            nrows=nrows,
+            constrained_layout=True,
+            figsize=(ncols * 3.5, nrows * 3.5 * aspect_ratio),
+            squeeze=False,
+            sharex=True,
+            sharey=True,
+            num=1,
+        )
+        tile_temp = np.zeros((renderer.tile_height, renderer.tile_width, channels), dtype=np.float32)
+        for dim in range(ncols * nrows):
+            ax = axes[dim // ncols, dim % ncols]
+            if dim >= num_tiles:
+                ax.axis("off")
+                continue
+            if render_mode == "depth":
+                img_plots.append(ax.imshow(tile_temp, vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane))
+            else:
+                img_plots.append(ax.imshow(tile_temp))
+    else:
+        fig = plt.figure(1)
+        pixels = wp.zeros((renderer.screen_height, renderer.screen_width, channels), dtype=wp.float32)
+        if render_mode == "depth":
+            img_plot = plt.imshow(pixels.numpy(), vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane)
+        else:
+            img_plot = plt.imshow(pixels.numpy())
+
+    plt.ion()
+    plt.show()
+
+while renderer.is_running():
+    time = renderer.clock_time
+    renderer.begin_frame(time)
+    for i in range(10):
+        renderer.render_capsule(
+            f"capsule_{i}", [i - 5.0, np.sin(time + i * 0.2), -3.0], [0.0, 0.0, 0.0, 1.0], radius=0.5, half_height=0.8
+        )
+    renderer.render_cylinder(
+        "cylinder",
+        [3.2, 1.0, np.sin(time + 0.5)],
+        np.array(wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.sin(time + 0.5))),
+        radius=0.5,
+        half_height=0.8,
+    )
+    renderer.render_cone(
+        "cone",
+        [-1.2, 1.0, 0.0],
+        np.array(wp.quat_from_axis_angle(wp.vec3(0.707, 0.707, 0.0), time)),
+        radius=0.5,
+        half_height=0.8,
+    )
+    renderer.end_frame()
+
+    if show_plot and plt.fignum_exists(1):
+        if split_up_tiles:
+            pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, channels)
+        else:
+            pixel_shape = (renderer.screen_height, renderer.screen_width, channels)
+
+        if pixel_shape != pixels.shape:
+            # make sure we resize the pixels array to the right dimensions if the user resizes the window
+            pixels = wp.zeros(pixel_shape, dtype=wp.float32)
+
+        renderer.get_pixels(pixels, split_up_tiles=split_up_tiles, mode=render_mode)
+
+        if split_up_tiles:
+            pixels_np = pixels.numpy()
+            for i, img_plot in enumerate(img_plots):
+                img_plot.set_data(pixels_np[i])
+        else:
+            img_plot.set_data(pixels.numpy())
+        fig.canvas.draw()
+        fig.canvas.flush_events()
+
+renderer.clear()

warp/examples/core/example_sph.py

@@ -0,0 +1,389 @@
+# 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 Smoothed Particle Hydrodynamics
+#
+# Shows how to implement a SPH fluid simulation.
+#
+# Neighbors are found using the wp.HashGrid class, and
+# wp.hash_grid_query(), wp.hash_grid_query_next() kernel methods.
+#
+# Reference Publication
+# Matthias Müller, David Charypar, and Markus H. Gross.
+# "Particle-based fluid simulation for interactive applications."
+# Symposium on Computer animation. Vol. 2. 2003.
+#
+###########################################################################
+
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.render
+
+wp.init()
+
+
+@wp.func
+def square(x: float):
+    return x * x
+
+
+@wp.func
+def cube(x: float):
+    return x * x * x
+
+
+@wp.func
+def fifth(x: float):
+    return x * x * x * x * x
+
+
+@wp.func
+def density_kernel(xyz: wp.vec3, smoothing_length: float):
+    # calculate distance
+    distance = wp.dot(xyz, xyz)
+
+    return wp.max(cube(square(smoothing_length) - distance), 0.0)
+
+
+@wp.func
+def diff_pressure_kernel(
+    xyz: wp.vec3, pressure: float, neighbor_pressure: float, neighbor_rho: float, smoothing_length: float
+):
+    # calculate distance
+    distance = wp.sqrt(wp.dot(xyz, xyz))
+
+    if distance < smoothing_length:
+        # calculate terms of kernel
+        term_1 = -xyz / distance
+        term_2 = (neighbor_pressure + pressure) / (2.0 * neighbor_rho)
+        term_3 = square(smoothing_length - distance)
+        return term_1 * term_2 * term_3
+    else:
+        return wp.vec3()
+
+
+@wp.func
+def diff_viscous_kernel(xyz: wp.vec3, v: wp.vec3, neighbor_v: wp.vec3, neighbor_rho: float, smoothing_length: float):
+    # calculate distance
+    distance = wp.sqrt(wp.dot(xyz, xyz))
+
+    # calculate terms of kernel
+    if distance < smoothing_length:
+        term_1 = (neighbor_v - v) / neighbor_rho
+        term_2 = smoothing_length - distance
+        return term_1 * term_2
+    else:
+        return wp.vec3()
+
+
+@wp.kernel
+def compute_density(
+    grid: wp.uint64,
+    particle_x: wp.array(dtype=wp.vec3),
+    particle_rho: wp.array(dtype=float),
+    density_normalization: float,
+    smoothing_length: float,
+):
+    tid = wp.tid()
+
+    # order threads by cell
+    i = wp.hash_grid_point_id(grid, tid)
+
+    # get local particle variables
+    x = particle_x[i]
+
+    # store density
+    rho = float(0.0)
+
+    # particle contact
+    neighbors = wp.hash_grid_query(grid, x, smoothing_length)
+
+    # loop through neighbors to compute density
+    for index in neighbors:
+        # compute distance
+        distance = x - particle_x[index]
+
+        # compute kernel derivative
+        rho += density_kernel(distance, smoothing_length)
+
+    # add external potential
+    particle_rho[i] = density_normalization * rho
+
+
+@wp.kernel
+def get_acceleration(
+    grid: wp.uint64,
+    particle_x: wp.array(dtype=wp.vec3),
+    particle_v: wp.array(dtype=wp.vec3),
+    particle_rho: wp.array(dtype=float),
+    particle_a: wp.array(dtype=wp.vec3),
+    isotropic_exp: float,
+    base_density: float,
+    gravity: float,
+    pressure_normalization: float,
+    viscous_normalization: float,
+    smoothing_length: float,
+):
+    tid = wp.tid()
+
+    # order threads by cell
+    i = wp.hash_grid_point_id(grid, tid)
+
+    # get local particle variables
+    x = particle_x[i]
+    v = particle_v[i]
+    rho = particle_rho[i]
+    pressure = isotropic_exp * (rho - base_density)
+
+    # store forces
+    pressure_force = wp.vec3()
+    viscous_force = wp.vec3()
+
+    # particle contact
+    neighbors = wp.hash_grid_query(grid, x, smoothing_length)
+
+    # loop through neighbors to compute acceleration
+    for index in neighbors:
+        if index != i:
+            # get neighbor velocity
+            neighbor_v = particle_v[index]
+
+            # get neighbor density and pressures
+            neighbor_rho = particle_rho[index]
+            neighbor_pressure = isotropic_exp * (neighbor_rho - base_density)
+
+            # compute relative position
+            relative_position = particle_x[index] - x
+
+            # calculate pressure force
+            pressure_force += diff_pressure_kernel(
+                relative_position, pressure, neighbor_pressure, neighbor_rho, smoothing_length
+            )
+
+            # compute kernel derivative
+            viscous_force += diff_viscous_kernel(relative_position, v, neighbor_v, neighbor_rho, smoothing_length)
+
+    # sum all forces
+    force = pressure_normalization * pressure_force + viscous_normalization * viscous_force
+
+    # add external potential
+    particle_a[i] = force / rho + wp.vec3(0.0, gravity, 0.0)
+
+
+@wp.kernel
+def apply_bounds(
+    particle_x: wp.array(dtype=wp.vec3),
+    particle_v: wp.array(dtype=wp.vec3),
+    damping_coef: float,
+    width: float,
+    height: float,
+    length: float,
+):
+    tid = wp.tid()
+
+    # get pos and velocity
+    x = particle_x[tid]
+    v = particle_v[tid]
+
+    # clamp x left
+    if x[0] < 0.0:
+        x = wp.vec3(0.0, x[1], x[2])
+        v = wp.vec3(v[0] * damping_coef, v[1], v[2])
+
+    # clamp x right
+    if x[0] > width:
+        x = wp.vec3(width, x[1], x[2])
+        v = wp.vec3(v[0] * damping_coef, v[1], v[2])
+
+    # clamp y bot
+    if x[1] < 0.0:
+        x = wp.vec3(x[0], 0.0, x[2])
+        v = wp.vec3(v[0], v[1] * damping_coef, v[2])
+
+    # clamp z left
+    if x[2] < 0.0:
+        x = wp.vec3(x[0], x[1], 0.0)
+        v = wp.vec3(v[0], v[1], v[2] * damping_coef)
+
+    # clamp z right
+    if x[2] > length:
+        x = wp.vec3(x[0], x[1], length)
+        v = wp.vec3(v[0], v[1], v[2] * damping_coef)
+
+    # apply clamps
+    particle_x[tid] = x
+    particle_v[tid] = v
+
+
+@wp.kernel
+def kick(particle_v: wp.array(dtype=wp.vec3), particle_a: wp.array(dtype=wp.vec3), dt: float):
+    tid = wp.tid()
+    v = particle_v[tid]
+    particle_v[tid] = v + particle_a[tid] * dt
+
+
+@wp.kernel
+def drift(particle_x: wp.array(dtype=wp.vec3), particle_v: wp.array(dtype=wp.vec3), dt: float):
+    tid = wp.tid()
+    x = particle_x[tid]
+    particle_x[tid] = x + particle_v[tid] * dt
+
+
+@wp.kernel
+def initialize_particles(
+    particle_x: wp.array(dtype=wp.vec3), smoothing_length: float, width: float, height: float, length: float
+):
+    tid = wp.tid()
+
+    # grid size
+    nr_x = wp.int32(width / 4.0 / smoothing_length)
+    nr_y = wp.int32(height / smoothing_length)
+    nr_z = wp.int32(length / 4.0 / smoothing_length)
+
+    # calculate particle position
+    z = wp.float(tid % nr_z)
+    y = wp.float((tid // nr_z) % nr_y)
+    x = wp.float((tid // (nr_z * nr_y)) % nr_x)
+    pos = smoothing_length * wp.vec3(x, y, z)
+
+    # add small jitter
+    state = wp.rand_init(123, tid)
+    pos = pos + 0.001 * smoothing_length * wp.vec3(wp.randn(state), wp.randn(state), wp.randn(state))
+
+    # set position
+    particle_x[tid] = pos
+
+
+class Example:
+    def __init__(self, stage):
+        # render params
+        self.frame_dt = 1.0 / 60.0
+        self.frame_count = 600
+        self.sim_time = 0.0
+
+        # simulation params
+        self.smoothing_length = 0.8  # NOTE change this to adjust number of particles
+        self.width = 80.0  # x
+        self.height = 80.0  # y
+        self.length = 80.0  # z
+        self.isotropic_exp = 20
+        self.base_density = 1.0
+        self.particle_mass = 0.01 * self.smoothing_length**3  # reduce according to smoothing length
+        self.dt = 0.01 * self.smoothing_length  # decrease sim dt by smoothing length
+        self.dynamic_visc = 0.025
+        self.damping_coef = -0.95
+        self.gravity = -0.1
+        self.n = int(
+            self.height * (self.width / 4.0) * (self.height / 4.0) / (self.smoothing_length**3)
+        )  # number particles (small box in corner)
+        self.sim_step_to_frame_ratio = int(32 / self.smoothing_length)
+
+        # constants
+        self.density_normalization = (315.0 * self.particle_mass) / (
+            64.0 * np.pi * self.smoothing_length**9
+        )  # integrate density kernel
+        self.pressure_normalization = -(45.0 * self.particle_mass) / (np.pi * self.smoothing_length**6)
+        self.viscous_normalization = (45.0 * self.dynamic_visc * self.particle_mass) / (
+            np.pi * self.smoothing_length**6
+        )
+
+        # allocate arrays
+        self.x = wp.empty(tuple([self.n]), dtype=wp.vec3)
+        self.v = wp.zeros(tuple([self.n]), dtype=wp.vec3)
+        self.rho = wp.zeros(tuple([self.n]), dtype=float)
+        self.a = wp.zeros(tuple([self.n]), dtype=wp.vec3)
+
+        # set random positions
+        wp.launch(
+            kernel=initialize_particles,
+            dim=self.n,
+            inputs=[self.x, self.smoothing_length, self.width, self.height, self.length],
+        )  # initialize in small area
+
+        # create hash array
+        grid_size = int(self.height / (4.0 * self.smoothing_length))
+        self.grid = wp.HashGrid(grid_size, grid_size, grid_size)
+
+        # renderer
+        self.renderer = None
+        if stage:
+            self.renderer = wp.render.UsdRenderer(stage)
+
+    def step(self):
+        with wp.ScopedTimer("step", active=True):
+            for _ in range(self.sim_step_to_frame_ratio):
+                with wp.ScopedTimer("grid build", active=False):
+                    # build grid
+                    self.grid.build(self.x, self.smoothing_length)
+
+                with wp.ScopedTimer("forces", active=False):
+                    # compute density of points
+                    wp.launch(
+                        kernel=compute_density,
+                        dim=self.n,
+                        inputs=[self.grid.id, self.x, self.rho, self.density_normalization, self.smoothing_length],
+                    )
+
+                    # get new acceleration
+                    wp.launch(
+                        kernel=get_acceleration,
+                        dim=self.n,
+                        inputs=[
+                            self.grid.id,
+                            self.x,
+                            self.v,
+                            self.rho,
+                            self.a,
+                            self.isotropic_exp,
+                            self.base_density,
+                            self.gravity,
+                            self.pressure_normalization,
+                            self.viscous_normalization,
+                            self.smoothing_length,
+                        ],
+                    )
+
+                    # apply bounds
+                    wp.launch(
+                        kernel=apply_bounds,
+                        dim=self.n,
+                        inputs=[self.x, self.v, self.damping_coef, self.width, self.height, self.length],
+                    )
+
+                    # kick
+                    wp.launch(kernel=kick, dim=self.n, inputs=[self.v, self.a, self.dt])
+
+                    # drift
+                    wp.launch(kernel=drift, dim=self.n, inputs=[self.x, self.v, self.dt])
+
+            self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render", active=True):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render_points(points=self.x.numpy(), radius=self.smoothing_length, name="points", colors=((0.8, 0.3, 0.2),) * len(self.x))
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_sph.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.frame_count):
+        example.render()
+        example.step()
+
+    if example.renderer:
+        example.renderer.save()