warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/core/example_nvdb.py

@@ -0,0 +1,182 @@
+# 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 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.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

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

warp/examples/core/example_raymarch.py

@@ -0,0 +1,205 @@
+# 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 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.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, height=1024, width=2048):
+        self.width = width
+        self.height = height
+        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 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 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(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+    parser.add_argument("--width", type=int, default=2048, help="Output image width in pixels.")
+    parser.add_argument("--height", type=int, default=1024, help="Output image height in pixels.")
+
+    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()

warp/examples/core/example_render_opengl.py

@@ -0,0 +1,193 @@
+# 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.
+
+###########################################################################
+# 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 numpy as np
+
+import warp as wp
+import warp.render
+
+
+class Example:
+    def __init__(self, num_tiles=4, custom_tile_arrangement=False):
+        if num_tiles < 1:
+            raise ValueError("num_tiles must be greater than or equal to 1.")
+
+        self.renderer = wp.render.OpenGLRenderer(vsync=False)
+        instance_ids = []
+
+        if custom_tile_arrangement:
+            positions = []
+            sizes = []
+        else:
+            positions = None
+            sizes = None
+
+        # 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))
+        self.renderer.setup_tiled_rendering(instance_ids, tile_positions=positions, tile_sizes=sizes)
+
+        self.renderer.render_ground()
+
+    def render(self):
+        time = self.renderer.clock_time
+        self.renderer.begin_frame(time)
+        for i in range(10):
+            self.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,
+            )
+        self.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,
+        )
+        self.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,
+        )
+        self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+    import distutils.util
+
+    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_tiles", type=int, default=4, help="Number of viewports to render in a single frame.")
+    parser.add_argument(
+        "--show_plot",
+        type=lambda x: bool(distutils.util.strtobool(x.strip())),
+        default=True,
+        help="Display the pixels in an additional matplotlib figure.",
+    )
+    parser.add_argument("--render_mode", type=str, choices=("depth", "rgb"), default="depth", help="")
+    parser.add_argument(
+        "--split_up_tiles",
+        type=lambda x: bool(distutils.util.strtobool(x.strip())),
+        default=True,
+        help="Whether to split tiles into subplots when --show_plot is True.",
+    )
+    parser.add_argument("--custom_tile_arrangement", action="store_true", help="Apply custom tile arrangement.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(num_tiles=args.num_tiles, custom_tile_arrangement=args.custom_tile_arrangement)
+
+        channels = 1 if args.render_mode == "depth" else 3
+
+        if args.show_plot:
+            import matplotlib.pyplot as plt
+
+            if args.split_up_tiles:
+                pixels = wp.zeros(
+                    (args.num_tiles, example.renderer.tile_height, example.renderer.tile_width, channels),
+                    dtype=wp.float32,
+                )
+                ncols = int(np.ceil(np.sqrt(args.num_tiles)))
+                nrows = int(np.ceil(args.num_tiles / float(ncols)))
+                img_plots = []
+                aspect_ratio = example.renderer.tile_height / example.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(
+                    (example.renderer.tile_height, example.renderer.tile_width, channels), dtype=np.float32
+                )
+                for dim in range(ncols * nrows):
+                    ax = axes[dim // ncols, dim % ncols]
+                    if dim >= args.num_tiles:
+                        ax.axis("off")
+                        continue
+                    if args.render_mode == "depth":
+                        img_plots.append(
+                            ax.imshow(
+                                tile_temp,
+                                vmin=example.renderer.camera_near_plane,
+                                vmax=example.renderer.camera_far_plane,
+                            )
+                        )
+                    else:
+                        img_plots.append(ax.imshow(tile_temp))
+            else:
+                fig = plt.figure(1)
+                pixels = wp.zeros(
+                    (example.renderer.screen_height, example.renderer.screen_width, channels), dtype=wp.float32
+                )
+                if args.render_mode == "depth":
+                    img_plot = plt.imshow(
+                        pixels.numpy(), vmin=example.renderer.camera_near_plane, vmax=example.renderer.camera_far_plane
+                    )
+                else:
+                    img_plot = plt.imshow(pixels.numpy())
+
+            plt.ion()
+            plt.show()
+
+        while example.renderer.is_running():
+            example.render()
+
+            if args.show_plot and plt.fignum_exists(1):
+                if args.split_up_tiles:
+                    pixel_shape = (args.num_tiles, example.renderer.tile_height, example.renderer.tile_width, channels)
+                else:
+                    pixel_shape = (example.renderer.screen_height, example.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)
+
+                example.renderer.get_pixels(pixels, split_up_tiles=args.split_up_tiles, mode=args.render_mode)
+
+                if args.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()
+
+        example.renderer.clear()