warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/fem/example_burgers.py

@@ -0,0 +1,261 @@
+# 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 Burgers
+#
+# This example simulates an inviscid non-conservative Burgers PDE using
+# Discontinuous Galerkin with minmod slope limiter
+#
+# d u /dt + (u . grad) u = 0
+#
+###########################################################################
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+
+@fem.integrand
+def vel_mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(v(s), u(s))
+
+
+@fem.integrand
+def upwind_transport_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, w: fem.Field):
+    # Upwinding transport with discontinuous convection velocity,
+    # using jump(w v) = jump(w) avg(v) + avg(w) jump(w)
+
+    nor = fem.normal(domain, s)
+    w_avg_n = wp.dot(fem.average(w, s), nor)
+    w_jump_n = wp.dot(fem.jump(w, s), nor)
+
+    x = domain(s)
+    v_avg = fem.average(v, s)
+    if x[0] <= 0.0 or x[0] >= 1.0:  # out
+        # if x[0] >= 1.0:  # out
+        v_jump = v(s)
+    else:
+        v_jump = fem.jump(v, s)
+
+    u_avg = fem.average(u, s)
+    u_jump = fem.jump(u, s)
+
+    return wp.dot(u_avg, v_jump * w_avg_n + v_avg * w_jump_n) + 0.5 * wp.dot(v_jump, u_jump) * (
+        wp.abs(w_avg_n) + 0.5 * wp.abs(w_jump_n)
+    )
+
+
+@fem.integrand
+def cell_transport_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, w: fem.Field):
+    # ((w . grad) u) . v = v^T (grad u) w =  grad(u) : (v w^T)
+    # with integration by parts
+    # u . Div (w v^T) = u^T grad(v) w  + u^T v div (w)
+
+    return -wp.dot(fem.div(w, s) * v(s) + fem.grad(v, s) * w(s), u(s))
+
+
+@fem.integrand
+def initial_condition(s: fem.Sample, domain: fem.Domain):
+    x = domain(s)[0] * 2.0
+    wave = wp.sin(x * wp.pi)
+    return wp.vec2(wp.where(x <= 1.0, wave, 0.0), 0.0)
+
+
+@fem.integrand
+def velocity_norm(s: fem.Sample, u: fem.Field):
+    return wp.length(u(s))
+
+
+@wp.func
+def minmod(a: float, b: float):
+    sa = wp.sign(a)
+    sb = wp.sign(b)
+    return wp.where(sa == sb, sa * wp.min(wp.abs(a), wp.abs(b)), 0.0)
+
+
+@fem.integrand
+def slope_limiter(domain: fem.Domain, s: fem.Sample, u: fem.Field, dx: wp.vec2):
+    # Minmod slope limiter against P0 discretization (evaluation at cell centers)
+    # Assumes regular grid topology
+
+    center_coords = fem.Coords(0.5, 0.5, 0.0)
+    cell_center = fem.make_free_sample(s.element_index, center_coords)
+    center_pos = domain(cell_center)
+
+    u_center = u(cell_center)
+
+    delta_coords = s.element_coords - center_coords
+
+    neighbour_xp = fem.lookup(domain, center_pos + wp.vec2(dx[0], 0.0))
+    neighbour_yp = fem.lookup(domain, center_pos + wp.vec2(0.0, dx[1]))
+    neighbour_xm = fem.lookup(domain, center_pos - wp.vec2(dx[0], 0.0))
+    neighbour_ym = fem.lookup(domain, center_pos - wp.vec2(0.0, dx[1]))
+
+    u_nxp = u(neighbour_xp)
+    u_nyp = u(neighbour_yp)
+    u_nxm = u(neighbour_xm)
+    u_nym = u(neighbour_ym)
+
+    gx = minmod(u_nxp[0] - u_center[0], u_center[0] - u_nxm[0]) * delta_coords[0]
+    gy = minmod(u_nyp[1] - u_center[1], u_center[1] - u_nym[1]) * delta_coords[1]
+
+    delta_u = u(s) - u_center
+    return u_center + wp.vec2(minmod(gx, delta_u[0]), minmod(gy, delta_u[1]))
+
+
+class Example:
+    def __init__(self, quiet=False, resolution=50, degree=1):
+        self._quiet = quiet
+
+        res = resolution
+        self.sim_dt = 1.0 / res
+        self.current_frame = 0
+
+        geo = fem.Grid2D(res=wp.vec2i(resolution))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+
+        basis_space = fem.make_polynomial_basis_space(geo, degree=degree, discontinuous=True)
+        vector_space = fem.make_collocated_function_space(basis_space, dtype=wp.vec2)
+        scalar_space = fem.make_collocated_function_space(basis_space, dtype=float)
+
+        # Test function for ou vector space
+        self._test = fem.make_test(space=vector_space, domain=domain)
+        # Test function for integration on sides
+        self._side_test = fem.make_test(space=vector_space, domain=sides)
+
+        # Inertia matrix
+        # For simplicity, use nodal integration so that inertia matrix is diagonal
+        trial = fem.make_trial(space=vector_space, domain=domain)
+        matrix_inertia = fem.integrate(
+            vel_mass_form, fields={"u": trial, "v": self._test}, output_dtype=wp.float32, assembly="nodal"
+        )
+        self._inv_mass_matrix = wp.sparse.bsr_copy(matrix_inertia)
+        fem_example_utils.invert_diagonal_bsr_matrix(self._inv_mass_matrix)
+
+        # Initial condition
+        self.velocity_field = vector_space.make_field()
+        fem.interpolate(initial_condition, dest=self.velocity_field)
+
+        # Velocity norm field -- for visualization purposes
+        self.velocity_norm_field = scalar_space.make_field()
+        fem.interpolate(velocity_norm, dest=self.velocity_norm_field, fields={"u": self.velocity_field})
+
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("u_norm", self.velocity_norm_field)
+
+    def _velocity_delta(self, trial_velocity):
+        # Integration on sides
+        rhs = fem.integrate(
+            upwind_transport_form,
+            fields={"u": trial_velocity.trace(), "v": self._side_test, "w": trial_velocity.trace()},
+            output_dtype=wp.vec2,
+        )
+
+        if self.velocity_field.space.degree > 0:
+            # Integration on cells (if not piecewise-constant)
+            fem.linalg.array_axpy(
+                x=fem.integrate(
+                    cell_transport_form,
+                    fields={"u": trial_velocity, "v": self._test, "w": trial_velocity},
+                    output_dtype=wp.vec2,
+                    quadrature=fem.RegularQuadrature(
+                        order=3, domain=self._test.domain, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE
+                    ),
+                ),
+                y=rhs,
+                alpha=1.0,
+                beta=1.0,
+            )
+        return self._inv_mass_matrix @ rhs
+
+    def step(self):
+        self.current_frame += 1
+
+        # Third-order Strong Stability Preserving Runge-Kutta (SSPRK3)
+
+        k1 = self._velocity_delta(self.velocity_field)
+
+        # tmp = v0 - dt * k1
+        tmp = self.velocity_field.space.make_field()
+        fem.linalg.array_axpy(y=tmp.dof_values, x=self.velocity_field.dof_values, alpha=1.0, beta=0.0)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=k1, alpha=-self.sim_dt, beta=1.0)
+        k2 = self._velocity_delta(tmp)
+
+        # tmp = v0 - dt * (0.25 * k1 + 0.25 * k2)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=k1, alpha=0.75 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=k2, alpha=-0.25 * self.sim_dt, beta=1.0)
+        k3 = self._velocity_delta(tmp)
+
+        # v = v0 - dt * (1/6 * k1 + 1/6 * k2 + 2/3 * k3)
+        fem.linalg.array_axpy(y=self.velocity_field.dof_values, x=k1, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=self.velocity_field.dof_values, x=k2, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=self.velocity_field.dof_values, x=k3, alpha=-2.0 / 3.0 * self.sim_dt, beta=1.0)
+
+        # Apply slope limiter
+        if self.velocity_field.space.degree > 0:
+            res = self.velocity_field.space.geometry.res
+            dx = wp.vec2(1.0 / res[0], 1.0 / res[1])
+            fem.interpolate(slope_limiter, dest=tmp, fields={"u": self.velocity_field}, values={"dx": dx})
+            wp.copy(src=tmp.dof_values, dest=self.velocity_field.dof_values)
+
+        # Update velocity norm (for visualization)
+        fem.interpolate(velocity_norm, dest=self.velocity_norm_field, fields={"u": self.velocity_field})
+
+    def render(self):
+        self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
+        self.renderer.add_field("u_norm", self.velocity_norm_field)
+        self.renderer.end_frame()
+
+
+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("--num_frames", type=int, default=250, help="Total number of frames.")
+    parser.add_argument("--degree", choices=(0, 1), type=int, default=1, help="Discretization order.")
+    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")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            resolution=args.resolution,
+            degree=args.degree,
+        )
+
+        for k in range(args.num_frames):
+            print(f"Frame {k}:")
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot()

warp/examples/fem/example_convection_diffusion.py

@@ -0,0 +1,181 @@
+# 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 Convection Diffusion
+#
+# This example simulates a convection-diffusion PDE using
+# semi-Lagrangian advection
+#
+# D phi / dt - nu d2 phi / dx^2 = 0
+###########################################################################
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+
+@fem.integrand
+def initial_condition(domain: fem.Domain, s: fem.Sample):
+    """Initial condition: 1.0 in ]0.6, 0.4[ x ]0.2, 0.8[, 0.0 elsewhere"""
+    pos = domain(s)
+    if pos[0] > 0.4 and pos[0] < 0.6 and pos[1] > 0.2 and pos[1] < 0.8:
+        return 1.0
+    return 0.0
+
+
+@wp.func
+def velocity(pos: wp.vec2, ang_vel: float):
+    center = wp.vec2(0.5, 0.5)
+    offset = pos - center
+    return wp.vec2(offset[1], -offset[0]) * ang_vel
+
+
+@fem.integrand
+def inertia_form(s: fem.Sample, phi: fem.Field, psi: fem.Field, dt: float):
+    return phi(s) * psi(s) / dt
+
+
+@fem.integrand
+def transported_inertia_form(
+    s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float, dt: float
+):
+    pos = domain(s)
+    vel = velocity(pos, ang_vel)
+
+    # semi-Lagrangian advection; evaluate phi upstream
+    conv_pos = pos - vel * dt
+    # lookup operator constructs a Sample from a world position.
+    # the optional last argument provides a initial guess for the lookup
+    conv_phi = phi(fem.lookup(domain, conv_pos, s))
+
+    return conv_phi * psi(s) / dt
+
+
+@fem.integrand
+def diffusion_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(
+        fem.grad(u, s),
+        fem.grad(v, s),
+    )
+
+
+@fem.integrand
+def diffusion_and_inertia_form(s: fem.Sample, phi: fem.Field, psi: fem.Field, dt: float, nu: float):
+    return inertia_form(s, phi, psi, dt) + nu * diffusion_form(s, phi, psi)
+
+
+class Example:
+    def __init__(self, quiet=False, degree=2, resolution=50, mesh: str = "grid", viscosity=0.001, ang_vel=1.0):
+        self._quiet = quiet
+
+        self._ang_vel = ang_vel
+
+        res = resolution
+        self.sim_dt = 1.0 / (ang_vel * res)
+        self.current_frame = 0
+
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions, build_bvh=True)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(resolution))
+
+        domain = fem.Cells(geometry=geo)
+        scalar_space = fem.make_polynomial_space(geo, degree=degree)
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        # Assemble diffusion and inertia matrix
+        self._test = fem.make_test(space=scalar_space, domain=domain)
+        self._trial = fem.make_trial(space=scalar_space, domain=domain)
+        self._matrix = fem.integrate(
+            diffusion_and_inertia_form,
+            fields={"phi": self._trial, "psi": self._test},
+            values={"nu": viscosity, "dt": self.sim_dt},
+            output_dtype=float,
+        )
+
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("phi", self._phi_field)
+
+    def step(self):
+        self.current_frame += 1
+
+        # right-hand-side -- advected inertia
+        rhs = fem.integrate(
+            transported_inertia_form,
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"ang_vel": self._ang_vel, "dt": self.sim_dt},
+            output_dtype=float,
+        )
+
+        # Solve linear system
+        fem_example_utils.bsr_cg(self._matrix, x=self._phi_field.dof_values, b=rhs, quiet=self._quiet, tol=1.0e-12)
+
+    def render(self):
+        self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
+        self.renderer.add_field("phi", self._phi_field)
+        self.renderer.end_frame()
+
+
+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("--num_frames", type=int, default=250, help="Total number of frames.")
+    parser.add_argument("--viscosity", type=float, default=0.001, help="Fluid viscosity parameter.")
+    parser.add_argument("--ang_vel", type=float, default=1.0, help="Angular velocity.")
+    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,
+            viscosity=args.viscosity,
+            ang_vel=args.ang_vel,
+        )
+
+        for k in range(args.num_frames):
+            print(f"Frame {k}:")
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot()

warp/examples/fem/example_convection_diffusion_dg.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 Convection Diffusion DG
+#
+# This example simulates a convection-diffusion PDE using Discontinuous
+# Galerkin with upwind transport and Symmetric Interior Penalty
+#
+# D phi / dt - nu d2 phi / dx^2 = 0
+###########################################################################
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.examples.fem.example_convection_diffusion import (
+    diffusion_form,
+    inertia_form,
+    initial_condition,
+    velocity,
+)
+
+
+# Standard transport term, on cells' interior
+@fem.integrand
+def transport_form(s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float):
+    pos = domain(s)
+    vel = velocity(pos, ang_vel)
+
+    return psi(s) * wp.dot(fem.grad(phi, s), vel)
+
+
+# Upwind flux, on cell sides
+@fem.integrand
+def upwind_transport_form(s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float):
+    pos = domain(s)
+    vel = velocity(pos, ang_vel)
+    vel_n = wp.dot(vel, fem.normal(domain, s))
+
+    if wp.min(pos) <= 0.0 or wp.max(pos) >= 1.0:  # boundary side
+        return phi(s) * (-psi(s) * vel_n + 0.5 * psi(s) * wp.abs(vel_n))
+
+    # interior side
+    return fem.jump(phi, s) * (-fem.average(psi, s) * vel_n + 0.5 * fem.jump(psi, s) * wp.abs(vel_n))
+
+
+# Symmetric-Interior-Penalty diffusion term (See Pietro Ern 2012)
+@fem.integrand
+def sip_diffusion_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    psi: fem.Field,
+    phi: fem.Field,
+):
+    nor = fem.normal(domain, s)
+    penalty = fem.measure_ratio(domain, s) * float(fem.degree(psi) * fem.degree(phi))
+
+    return penalty * fem.jump(phi, s) * fem.jump(psi, s) - (
+        wp.dot(fem.grad_average(phi, s), nor) * fem.jump(psi, s)
+        + wp.dot(fem.grad_average(psi, s), nor) * fem.jump(phi, s)
+    )
+
+
+class Example:
+    def __init__(self, quiet=False, degree=2, resolution=50, mesh="grid", viscosity=0.0001, ang_vel=1.0):
+        self._quiet = quiet
+
+        res = resolution
+        self.sim_dt = 1.0 / (ang_vel * res)
+        self.current_frame = 0
+
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
+            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))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(resolution))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+        scalar_space = fem.make_polynomial_space(
+            geo,
+            discontinuous=True,
+            degree=degree,
+            family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE,
+        )
+
+        # Assemble transport, diffusion and inertia matrices
+
+        self._test = fem.make_test(space=scalar_space, domain=domain)
+        trial = fem.make_trial(space=scalar_space, domain=domain)
+
+        matrix_inertia = fem.integrate(
+            inertia_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
+
+        matrix_transport = fem.integrate(
+            transport_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"ang_vel": ang_vel},
+        )
+
+        self._side_test = fem.make_test(space=scalar_space, domain=sides)
+        side_trial = fem.make_trial(space=scalar_space, domain=sides)
+        fem.integrate(
+            upwind_transport_form,
+            fields={"phi": side_trial, "psi": self._side_test},
+            values={"ang_vel": ang_vel},
+            output=matrix_transport,
+            add=True,
+        )
+
+        matrix_diffusion = fem.integrate(
+            diffusion_form,
+            fields={"u": trial, "v": self._test},
+        )
+        matrix_diffusion += fem.integrate(
+            sip_diffusion_form,
+            fields={"phi": side_trial, "psi": self._side_test},
+        )
+        self._matrix = matrix_inertia + matrix_transport + viscosity * matrix_diffusion
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        self._phi_curvature_field = scalar_space.make_field()
+        self._compute_phi_curvature()
+
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("phi", self._phi_field)
+        self.renderer.add_field("phi_curvature", self._phi_curvature_field)
+
+    def step(self):
+        self.current_frame += 1
+
+        rhs = fem.integrate(
+            inertia_form,
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
+
+        phi = wp.zeros_like(rhs)
+        fem_example_utils.bsr_cg(self._matrix, b=rhs, x=phi, method="bicgstab", quiet=self._quiet)
+        wp.utils.array_cast(in_array=phi, out_array=self._phi_field.dof_values)
+
+        # for visualization purposes only
+        self._compute_phi_curvature()
+
+    def _compute_phi_curvature(self):
+        fem.integrate(
+            diffusion_form,
+            fields={"u": self._phi_field, "v": self._test},
+            output=self._phi_curvature_field.dof_values,
+        )
+        fem.integrate(
+            sip_diffusion_form,
+            fields={"phi": self._phi_field.trace(), "psi": self._side_test},
+            output=self._phi_curvature_field.dof_values,
+            add=True,
+        )
+
+    def render(self):
+        self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
+        self.renderer.add_field("phi", self._phi_field)
+        self.renderer.add_field("phi_curvature", self._phi_curvature_field)
+
+        self.renderer.end_frame()
+
+
+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("--num_frames", type=int, default=100, help="Total number of frames.")
+    parser.add_argument("--viscosity", type=float, default=0.001, help="Fluid viscosity parameter.")
+    parser.add_argument("--ang_vel", type=float, default=1.0, help="Angular velocity.")
+    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")
+
+    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,
+            viscosity=args.viscosity,
+            ang_vel=args.ang_vel,
+        )
+
+        for k in range(args.num_frames):
+            print(f"Frame {k}:")
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot()