warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/fem/example_diffusion.py

@@ -0,0 +1,196 @@
+# 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 Diffusion
+#
+# This example solves a 2d diffusion problem:
+#
+# nu Div u = 1
+#
+# with Dirichlet boundary conditions on vertical edges and
+# homogeneous Neumann on horizontal edges.
+###########################################################################
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.fem.linalg import array_axpy
+
+
+@fem.integrand
+def linear_form(
+    s: fem.Sample,
+    v: fem.Field,
+):
+    """Linear form with constant slope 1 -- forcing term of our problem"""
+    return v(s)
+
+
+@fem.integrand
+def diffusion_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
+    """Diffusion bilinear form with constant coefficient ``nu``"""
+    return nu * wp.dot(
+        fem.grad(u, s),
+        fem.grad(v, s),
+    )
+
+
+@fem.integrand
+def y_boundary_value_form(s: fem.Sample, domain: fem.Domain, v: fem.Field, val: float):
+    """Linear form with coefficient val on vertical edges, zero elsewhere"""
+    nor = fem.normal(domain, s)
+    return val * v(s) * wp.abs(nor[0])
+
+
+@fem.integrand
+def y_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    """
+    Bilinear boundary condition projector form, non-zero on vertical edges only.
+    """
+    # Reuse the above linear form implementation by evaluating one of the participating field and passing it as a normal scalar argument.
+    return y_boundary_value_form(s, domain, v, u(s))
+
+
+class Example:
+    def __init__(
+        self,
+        quiet=False,
+        degree=2,
+        resolution=50,
+        mesh="grid",
+        serendipity=False,
+        viscosity=2.0,
+        boundary_value=5.0,
+        boundary_compliance=0.0,
+    ):
+        self._quiet = quiet
+
+        self._viscosity = viscosity
+        self._boundary_value = boundary_value
+        self._boundary_compliance = boundary_compliance
+
+        # Grid or triangle mesh geometry
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
+            self._geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            self._geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            self._geo = fem.Grid2D(res=wp.vec2i(resolution))
+
+        # Scalar function space
+        element_basis = fem.ElementBasis.SERENDIPITY if serendipity else None
+        self._scalar_space = fem.make_polynomial_space(self._geo, degree=degree, element_basis=element_basis)
+
+        # Scalar field over our function space
+        self._scalar_field = self._scalar_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        geo = self._geo
+
+        domain = fem.Cells(geometry=geo)
+
+        # Right-hand-side (forcing term)
+        test = fem.make_test(space=self._scalar_space, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Diffusion form
+        trial = fem.make_trial(space=self._scalar_space, domain=domain)
+        matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": self._viscosity})
+
+        # Boundary conditions on Y sides
+        # Use nodal integration so that boundary conditions are specified on each node independently
+        boundary = fem.BoundarySides(geo)
+        bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
+        bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
+
+        bd_matrix = fem.integrate(y_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, assembly="nodal")
+        bd_rhs = fem.integrate(
+            y_boundary_value_form, fields={"v": bd_test}, values={"val": self._boundary_value}, assembly="nodal"
+        )
+
+        # Assemble linear system
+        if self._boundary_compliance == 0.0:
+            # Hard BC: project linear system
+            fem.project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
+        else:
+            # Weak BC: add together diffusion and boundary condition matrices
+            boundary_strength = 1.0 / self._boundary_compliance
+            matrix += bd_matrix * boundary_strength
+            array_axpy(x=bd_rhs, y=rhs, alpha=boundary_strength, beta=1)
+
+        # Solve linear system using Conjugate Gradient
+        x = wp.zeros_like(rhs)
+        fem_example_utils.bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
+
+        # Assign system result to our discrete field
+        self._scalar_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_field("solution", self._scalar_field)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    wp.set_module_options({"enable_backward": False})
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--resolution", type=int, default=50, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--serendipity", action="store_true", default=False, help="Use Serendipity basis functions.")
+    parser.add_argument("--viscosity", type=float, default=2.0, help="Fluid viscosity parameter.")
+    parser.add_argument(
+        "--boundary_value", type=float, default=5.0, help="Value of Dirichlet boundary condition on vertical edges."
+    )
+    parser.add_argument(
+        "--boundary_compliance", type=float, default=0.0, help="Dirichlet boundary condition compliance."
+    )
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type.")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+    parser.add_argument("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            degree=args.degree,
+            resolution=args.resolution,
+            mesh=args.mesh,
+            serendipity=args.serendipity,
+            viscosity=args.viscosity,
+            boundary_value=args.boundary_value,
+            boundary_compliance=args.boundary_compliance,
+        )
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot()

warp/examples/fem/example_diffusion_3d.py

@@ -0,0 +1,225 @@
+# 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 Diffusion 3D
+#
+# This example solves a 3d diffusion problem:
+#
+# nu Div u = 1
+#
+# with homogeneous Neumann conditions on horizontal sides
+# and homogeneous Dirichlet boundary conditions other sides.
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.examples.fem.example_diffusion import diffusion_form, linear_form
+from warp.sparse import bsr_axpy
+
+
+@fem.integrand
+def vertical_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # Constrain XY and YZ faces
+    nor = fem.normal(domain, s)
+    w = 1.0 - wp.abs(nor[1])
+    return w * u(s) * v(s)
+
+
+@fem.integrand
+def y_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # Constrain Y edges
+    tangent = fem.deformation_gradient(domain, s)
+    return wp.abs(tangent[1]) * u(s) * v(s)
+
+
+class Example:
+    def __init__(
+        self,
+        quiet=False,
+        degree=2,
+        resolution=10,
+        mesh="grid",
+        serendipity=False,
+        viscosity=2.0,
+        boundary_compliance=0.0,
+    ):
+        self._quiet = quiet
+
+        self._viscosity = viscosity
+        self._boundary_compliance = boundary_compliance
+
+        res = wp.vec3i(resolution, max(1, resolution // 2), resolution * 2)
+        bounds_lo = wp.vec3(0.0, 0.0, 0.0)
+        bounds_hi = wp.vec3(1.0, 0.5, 2.0)
+
+        if mesh == "tet":
+            pos, tet_vtx_indices = fem_example_utils.gen_tetmesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            self._geo = fem.Tetmesh(tet_vtx_indices, pos)
+        elif mesh == "hex":
+            pos, hex_vtx_indices = fem_example_utils.gen_hexmesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            self._geo = fem.Hexmesh(hex_vtx_indices, pos)
+        elif mesh == "nano":
+            volume = fem_example_utils.gen_volume(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            self._geo = fem.Nanogrid(volume)
+        elif mesh == "tri":
+            pos, quad_vtx_indices = fem_example_utils.gen_trimesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            pos = pos.numpy()
+            pos_z = np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            pos = wp.array(pos, dtype=wp.vec3)
+            self._geo = fem.Trimesh3D(quad_vtx_indices, pos)
+        elif mesh == "quad":
+            pos, quad_vtx_indices = fem_example_utils.gen_quadmesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            pos = pos.numpy()
+            pos_z = np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            pos = wp.array(pos, dtype=wp.vec3)
+            self._geo = fem.Quadmesh3D(quad_vtx_indices, pos)
+        else:
+            self._geo = fem.Grid3D(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+
+        # Domain and function spaces
+        element_basis = fem.ElementBasis.SERENDIPITY if serendipity else None
+        self._scalar_space = fem.make_polynomial_space(self._geo, degree=degree, element_basis=element_basis)
+
+        # Scalar field over our function space
+        self._scalar_field: fem.DiscreteField = self._scalar_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        geo = self._geo
+
+        domain = fem.Cells(geometry=geo)
+
+        # Right-hand-side
+        test = fem.make_test(space=self._scalar_space, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Weakly-imposed boundary conditions on Y sides
+        with wp.ScopedTimer("Integrate"):
+            boundary = fem.BoundarySides(geo)
+
+            bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
+            bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
+
+            # Pick boundary conditions depending on whether our geometry is a 3d surface or a volume
+            boundary_projector_form = (
+                vertical_boundary_projector_form if self._geo.cell_dimension == 3 else y_boundary_projector_form
+            )
+            bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, assembly="nodal")
+
+            # Diffusion form
+            trial = fem.make_trial(space=self._scalar_space, domain=domain)
+            matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": self._viscosity})
+
+        if self._boundary_compliance == 0.0:
+            # Hard BC: project linear system
+            bd_rhs = wp.zeros_like(rhs)
+            fem.project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
+        else:
+            # Weak BC: add together diffusion and boundary condition matrices
+            boundary_strength = 1.0 / self._boundary_compliance
+            bsr_axpy(x=bd_matrix, y=matrix, alpha=boundary_strength, beta=1)
+
+        with wp.ScopedTimer("CG solve"):
+            x = wp.zeros_like(rhs)
+            fem_example_utils.bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
+            self._scalar_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_field("solution", self._scalar_field)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    wp.set_module_options({"enable_backward": False})
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--resolution", type=int, default=10, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--serendipity", action="store_true", default=False, help="Use Serendipity basis functions.")
+    parser.add_argument("--viscosity", type=float, default=2.0, help="Fluid viscosity parameter.")
+    parser.add_argument(
+        "--boundary_compliance", type=float, default=0.0, help="Dirichlet boundary condition compliance."
+    )
+    parser.add_argument(
+        "--mesh", choices=("grid", "tet", "hex", "nano", "anano", "tri", "quad"), default="grid", help="Mesh type."
+    )
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+    parser.add_argument("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            degree=args.degree,
+            resolution=args.resolution,
+            mesh=args.mesh,
+            serendipity=args.serendipity,
+            viscosity=args.viscosity,
+            boundary_compliance=args.boundary_compliance,
+        )
+
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot()

warp/examples/fem/example_diffusion_mgpu.py

@@ -0,0 +1,225 @@
+# 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 Diffusion MGPU
+#
+# This example illustrates using domain decomposition to
+# solve a diffusion PDE over multiple devices
+###########################################################################
+
+from typing import Tuple
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.examples.fem.example_diffusion import diffusion_form, linear_form
+from warp.sparse import bsr_axpy, bsr_mv
+from warp.utils import array_cast
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return u(s) * v(s)
+
+
+@wp.kernel
+def scal_kernel(a: wp.array(dtype=wp.float64), res: wp.array(dtype=wp.float64), alpha: wp.float64):
+    res[wp.tid()] = a[wp.tid()] * alpha
+
+
+@wp.kernel
+def sum_kernel(a: wp.indexedarray(dtype=wp.float64), b: wp.array(dtype=wp.float64)):
+    a[wp.tid()] = a[wp.tid()] + b[wp.tid()]
+
+
+def sum_vecs(vecs, indices, sum: wp.array, tmp: wp.array):
+    for v, idx in zip(vecs, indices):
+        wp.copy(dest=tmp, src=v)
+        idx_sum = wp.indexedarray(sum, idx)
+        wp.launch(kernel=sum_kernel, dim=idx.shape, device=sum.device, inputs=[idx_sum, tmp])
+
+    return sum
+
+
+class DistributedSystem:
+    device = None
+    scalar_type: type
+    tmp_buf: wp.array
+
+    nrow: int
+    shape = Tuple[int, int]
+    rank_data = None
+
+    def mv_routine(self, x: wp.array, y: wp.array, z: wp.array, alpha=1.0, beta=0.0):
+        """Distributed matrix-vector multiplication routine, for example purposes"""
+
+        tmp = self.tmp_buf
+
+        wp.launch(kernel=scal_kernel, dim=y.shape, device=y.device, inputs=[y, z, wp.float64(beta)])
+
+        stream = wp.get_stream()
+
+        for mat_i, x_i, y_i, idx in zip(*self.rank_data):
+            tmp_i = tmp[: idx.size]
+
+            # Compress rhs on rank 0
+            x_idx = wp.indexedarray(x, idx)
+            wp.copy(dest=tmp_i, src=x_idx, count=idx.size, stream=stream)
+
+            # Send to rank i
+            wp.copy(dest=x_i, src=tmp_i, count=idx.size, stream=stream)
+
+            with wp.ScopedDevice(x_i.device):
+                wp.wait_stream(stream)
+                bsr_mv(A=mat_i, x=x_i, y=y_i, alpha=alpha, beta=0.0)
+
+            wp.wait_stream(wp.get_stream(x_i.device))
+
+            # Back to rank 0 for sum
+            wp.copy(dest=tmp_i, src=y_i, count=idx.size, stream=stream)
+            z_idx = wp.indexedarray(z, idx)
+            wp.launch(kernel=sum_kernel, dim=idx.shape, device=z_idx.device, inputs=[z_idx, tmp_i], stream=stream)
+
+        wp.wait_stream(stream)
+
+
+class Example:
+    def __init__(self, quiet=False, device=None):
+        self._bd_weight = 100.0
+        self._quiet = quiet
+
+        self._geo = fem.Grid2D(res=wp.vec2i(25))
+
+        self._main_device = wp.get_device(device)
+
+        with wp.ScopedDevice(self._main_device):
+            self._scalar_space = fem.make_polynomial_space(self._geo, degree=3)
+            self._scalar_field = self._scalar_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        devices = wp.get_cuda_devices()
+        main_device = self._main_device
+
+        rhs_vecs = []
+        res_vecs = []
+        matrices = []
+        indices = []
+
+        # Build local system for each device
+        for k, device in enumerate(devices):
+            with wp.ScopedDevice(device):
+                # Construct the partition corresponding to the k'th device
+                geo_partition = fem.LinearGeometryPartition(self._geo, k, len(devices))
+                matrix, rhs, partition_node_indices = self._assemble_local_system(geo_partition)
+
+                rhs_vecs.append(rhs)
+                res_vecs.append(wp.empty_like(rhs))
+                matrices.append(matrix)
+                indices.append(partition_node_indices.to(main_device))
+
+        # Global rhs as sum of all local rhs
+        glob_rhs = wp.zeros(n=self._scalar_space.node_count(), dtype=wp.float64, device=main_device)
+
+        # This temporary buffer will be used for peer-to-peer copying during graph capture,
+        # so we allocate it using the default CUDA allocator.  This ensures that the copying
+        # will succeed without enabling mempool access between devices, which is not supported
+        # on all systems.
+        with wp.ScopedMempool(main_device, False):
+            tmp = wp.empty_like(glob_rhs)
+
+        sum_vecs(rhs_vecs, indices, glob_rhs, tmp)
+
+        # Distributed CG
+        global_res = wp.zeros_like(glob_rhs)
+        A = DistributedSystem()
+        A.device = main_device
+        A.dtype = glob_rhs.dtype
+        A.nrow = self._scalar_space.node_count()
+        A.shape = (A.nrow, A.nrow)
+        A.tmp_buf = tmp
+        A.rank_data = (matrices, rhs_vecs, res_vecs, indices)
+
+        with wp.ScopedDevice(main_device):
+            fem_example_utils.bsr_cg(
+                A,
+                x=global_res,
+                b=glob_rhs,
+                use_diag_precond=False,
+                quiet=self._quiet,
+                mv_routine=A.mv_routine,
+                mv_routine_uses_multiple_cuda_contexts=True,
+            )
+
+        array_cast(in_array=global_res, out_array=self._scalar_field.dof_values)
+
+    def render(self):
+        self.renderer.add_field("solution", self._scalar_field)
+
+    def _assemble_local_system(self, geo_partition: fem.GeometryPartition):
+        scalar_space = self._scalar_space
+        space_partition = fem.make_space_partition(scalar_space, geo_partition)
+
+        domain = fem.Cells(geometry=geo_partition)
+
+        # Right-hand-side
+        test = fem.make_test(space=scalar_space, space_partition=space_partition, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Weakly-imposed boundary conditions on all sides
+        boundary = fem.BoundarySides(geometry=geo_partition)
+        bd_test = fem.make_test(space=scalar_space, space_partition=space_partition, domain=boundary)
+        bd_trial = fem.make_trial(space=scalar_space, space_partition=space_partition, domain=boundary)
+        bd_matrix = fem.integrate(mass_form, fields={"u": bd_trial, "v": bd_test})
+
+        # Diffusion form
+        trial = fem.make_trial(space=scalar_space, space_partition=space_partition, domain=domain)
+        matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": 1.0})
+
+        bsr_axpy(y=matrix, x=bd_matrix, alpha=self._bd_weight)
+
+        return matrix, rhs, space_partition.space_node_indices()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    wp.set_module_options({"enable_backward": False})
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+    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.ScopedTimer(__file__):
+        example = Example(quiet=args.quiet, device=args.device)
+
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot()