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

warp/examples/fem/example_mixed_elasticity.py

@@ -0,0 +1,222 @@
+# 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 Mixed Elasticity
+#
+# This example illustrates using Mixed FEM to solve a
+# 2D linear elasticity problem:
+#
+# Div[ E: D(u) ] = 0
+# 
+# with Dirichlet boundary conditions on horizontal sides,
+# and E the elasticity rank-4 tensor
+###########################################################################
+
+import argparse
+
+import warp as wp
+import numpy as np
+
+import warp.fem as fem
+
+from warp.sparse import bsr_transposed, bsr_mm
+
+try:
+    from .plot_utils import Plot
+    from .bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+except ImportError:
+    from plot_utils import Plot
+    from bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
+    from mesh_utils import gen_trimesh, gen_quadmesh
+
+wp.init()
+
+
+@wp.func
+def compute_stress(tau: wp.mat22, E: wp.mat33):
+    """Strain to stress computation"""
+    tau_sym = wp.vec3(tau[0, 0], tau[1, 1], tau[0, 1] + tau[1, 0])
+    sig_sym = E * tau_sym
+    return wp.mat22(sig_sym[0], 0.5 * sig_sym[2], 0.5 * sig_sym[2], sig_sym[1])
+
+
+@fem.integrand
+def symmetric_grad_form(
+    s: fem.Sample,
+    u: fem.Field,
+    tau: fem.Field,
+):
+    """D(u) : tau"""
+    return wp.ddot(tau(s), fem.D(u, s))
+
+
+@fem.integrand
+def stress_form(s: fem.Sample, u: fem.Field, tau: fem.Field, E: wp.mat33):
+    """(E : D(u)) : tau"""
+    return wp.ddot(tau(s), compute_stress(fem.D(u, s), E))
+
+
+@fem.integrand
+def horizontal_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # non zero on horizontal boundary of domain only
+    nor = fem.normal(domain, s)
+    return wp.dot(u(s), v(s)) * wp.abs(nor[1])
+
+
+@fem.integrand
+def horizontal_displacement_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    v: fem.Field,
+    displacement: float,
+):
+    # opposed to normal on horizontal boundary of domain only
+    nor = fem.normal(domain, s)
+    return -wp.abs(nor[1]) * displacement * wp.dot(nor, v(s))
+
+
+@fem.integrand
+def tensor_mass_form(
+    s: fem.Sample,
+    sig: fem.Field,
+    tau: fem.Field,
+):
+    return wp.ddot(tau(s), sig(s))
+
+
+class Example:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resolution", type=int, default=25)
+    parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--displacement", type=float, default=0.1)
+    parser.add_argument("--young_modulus", type=float, default=1.0)
+    parser.add_argument("--poisson_ratio", type=float, default=0.5)
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+    parser.add_argument(
+        "--nonconforming_stresses", action="store_true", help="For grid, use non-conforming stresses (Q_d/P_d)"
+    )
+
+    def __init__(self, stage=None, quiet=False, args=None, **kwargs):
+        if args is None:
+            # Read args from kwargs, add default arg values from parser
+            args = argparse.Namespace(**kwargs)
+            args = Example.parser.parse_args(args=[], namespace=args)
+        self._args = args
+        self._quiet = quiet
+
+        # Grid or triangle mesh geometry
+        if args.mesh == "tri":
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
+            self._geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif args.mesh == "quad":
+            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
+            self._geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            self._geo = fem.Grid2D(res=wp.vec2i(args.resolution))
+
+        # Strain-stress matrix
+        young = args.young_modulus
+        poisson = args.poisson_ratio
+        self._elasticity_mat = wp.mat33(
+            young
+            / (1.0 - poisson * poisson)
+            * np.array(
+                [
+                    [1.0, poisson, 0.0],
+                    [poisson, 1.0, 0.0],
+                    [0.0, 0.0, (2.0 * (1.0 + poisson)) * (1.0 - poisson * poisson)],
+                ]
+            )
+        )
+
+        # Function spaces -- S_k for displacement, Q_{k-1}d for stress
+        self._u_space = fem.make_polynomial_space(
+            self._geo, degree=args.degree, dtype=wp.vec2, element_basis=fem.ElementBasis.SERENDIPITY
+        )
+
+        # Store stress degrees of freedom as symmetric tensors (3 dof) rather than full 2x2 matrices
+        tau_basis = (
+            fem.ElementBasis.NONCONFORMING_POLYNOMIAL if args.nonconforming_stresses else fem.ElementBasis.LAGRANGE
+        )
+        self._tau_space = fem.make_polynomial_space(
+            self._geo,
+            degree=args.degree - 1,
+            discontinuous=True,
+            element_basis=tau_basis,
+            dof_mapper=fem.SymmetricTensorMapper(wp.mat22),
+        )
+
+        self._u_field = self._u_space.make_field()
+
+        self.renderer = Plot(stage)
+
+    def step(self):
+        boundary = fem.BoundarySides(self._geo)
+        domain = fem.Cells(geometry=self._geo)
+
+        # Displacement boundary conditions
+        u_bd_test = fem.make_test(space=self._u_space, domain=boundary)
+        u_bd_trial = fem.make_trial(space=self._u_space, domain=boundary)
+        u_bd_rhs = fem.integrate(
+            horizontal_displacement_form,
+            fields={"v": u_bd_test},
+            values={"displacement": self._args.displacement},
+            nodal=True,
+            output_dtype=wp.vec2d,
+        )
+        u_bd_matrix = fem.integrate(
+            horizontal_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True
+        )
+
+        # Stress/velocity coupling
+        u_trial = fem.make_trial(space=self._u_space, domain=domain)
+        tau_test = fem.make_test(space=self._tau_space, domain=domain)
+        tau_trial = fem.make_trial(space=self._tau_space, domain=domain)
+
+        sym_grad_matrix = fem.integrate(symmetric_grad_form, fields={"u": u_trial, "tau": tau_test})
+        stress_matrix = fem.integrate(
+            stress_form, fields={"u": u_trial, "tau": tau_test}, values={"E": self._elasticity_mat}
+        )
+
+        # Compute inverse of the (block-diagonal) tau mass matrix
+        tau_inv_mass_matrix = fem.integrate(tensor_mass_form, fields={"sig": tau_trial, "tau": tau_test}, nodal=True)
+        invert_diagonal_bsr_mass_matrix(tau_inv_mass_matrix)
+
+        # Assemble system matrix
+        u_matrix = bsr_mm(bsr_transposed(sym_grad_matrix), bsr_mm(tau_inv_mass_matrix, stress_matrix))
+
+        # Enforce boundary conditions
+        u_rhs = wp.zeros_like(u_bd_rhs)
+        fem.project_linear_system(u_matrix, u_rhs, u_bd_matrix, u_bd_rhs)
+
+        x = wp.zeros_like(u_rhs)
+        bsr_cg(u_matrix, b=u_rhs, x=x, tol=1.0e-16, quiet=self._quiet)
+
+        # Extract result
+        self._u_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_surface_vector("solution", self._u_field)
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    example.step()
+    example.render()
+
+    example.renderer.plot()

warp/examples/fem/example_navier_stokes.py

@@ -0,0 +1,243 @@
+# 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 Navier Stokes
+#
+# This example solves a 2D Navier-Stokes flow problem:
+#
+# Du/dt -nu D(u) + grad p = 0
+# Div u = 0
+#
+# with (hard) velocity-Dirichlet boundary conditions
+# and using semi-Lagrangian advection
+###########################################################################
+
+import argparse
+
+import warp as wp
+import warp.fem as fem
+
+from warp.fem.utils import array_axpy
+
+from warp.sparse import bsr_mm, bsr_mv, bsr_copy 
+
+try:
+    from .bsr_utils import bsr_solve_saddle, SaddleSystem
+    from .plot_utils import Plot
+    from .mesh_utils import gen_trimesh
+except ImportError:
+    from bsr_utils import bsr_solve_saddle, SaddleSystem
+    from plot_utils import Plot
+    from mesh_utils import gen_trimesh
+
+wp.init()
+
+
+@fem.integrand
+def u_boundary_value(s: fem.Sample, domain: fem.Domain, v: fem.Field, top_vel: float):
+    # Horizontal velocity on top of domain, zero elsewhere
+    if domain(s)[1] == 1.0:
+        return wp.dot(wp.vec2f(top_vel, 0.0), v(s))
+
+    return wp.dot(wp.vec2f(0.0, 0.0), v(s))
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def inertia_form(s: fem.Sample, u: fem.Field, v: fem.Field, dt: float):
+    return mass_form(s, u, v) / dt
+
+
+@fem.integrand
+def viscosity_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
+    return 2.0 * nu * wp.ddot(fem.D(u, s), fem.D(v, s))
+
+
+@fem.integrand
+def viscosity_and_inertia_form(s: fem.Sample, u: fem.Field, v: fem.Field, dt: float, nu: float):
+    return inertia_form(s, u, v, dt) + viscosity_form(s, u, v, nu)
+
+
+@fem.integrand
+def transported_inertia_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, dt: float):
+    pos = domain(s)
+    vel = u(s)
+
+    conv_pos = pos - 0.5 * vel * dt
+    conv_s = fem.lookup(domain, conv_pos, s)
+    conv_vel = u(conv_s)
+
+    conv_pos = conv_pos - 0.5 * conv_vel * dt
+    conv_vel = u(fem.lookup(domain, conv_pos, conv_s))
+
+    return wp.dot(conv_vel, v(s)) / dt
+
+
+@fem.integrand
+def div_form(
+    s: fem.Sample,
+    u: fem.Field,
+    q: fem.Field,
+):
+    return -q(s) * fem.div(u, s)
+
+
+class Example:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resolution", type=int, default=25)
+    parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--num_frames", type=int, default=1000)
+    parser.add_argument("--top_velocity", type=float, default=1.0)
+    parser.add_argument("--Re", type=float, default=1000.0)
+    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh")
+
+    def __init__(self, stage=None, quiet=False, args=None, **kwargs):
+        if args is None:
+            # Read args from kwargs, add default arg values from parser
+            args = argparse.Namespace(**kwargs)
+            args = Example.parser.parse_args(args=[], namespace=args)
+        self._args = args
+        self._quiet = quiet
+
+        res = args.resolution
+        self.sim_dt = 1.0 / args.resolution
+        self.current_frame = 0
+
+        viscosity = args.top_velocity / args.Re
+
+        if args.tri_mesh:
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(res))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(res))
+
+        domain = fem.Cells(geometry=geo)
+        boundary = fem.BoundarySides(geo)
+
+        # Functions spaces: Q(d)-Q(d-1)
+        u_degree = args.degree
+        u_space = fem.make_polynomial_space(geo, degree=u_degree, dtype=wp.vec2)
+        p_space = fem.make_polynomial_space(geo, degree=u_degree - 1)
+
+        # Viscosity and inertia
+        u_test = fem.make_test(space=u_space, domain=domain)
+        u_trial = fem.make_trial(space=u_space, domain=domain)
+
+        u_matrix = fem.integrate(
+            viscosity_and_inertia_form,
+            fields={"u": u_trial, "v": u_test},
+            values={"nu": viscosity, "dt": self.sim_dt},
+        )
+
+        # Pressure-velocity coupling
+        p_test = fem.make_test(space=p_space, domain=domain)
+        div_matrix = fem.integrate(div_form, fields={"u": u_trial, "q": p_test})
+
+        # Enforcing the Dirichlet boundary condition the hard way;
+        # build projector for velocity left- and right-hand-sides
+        u_bd_test = fem.make_test(space=u_space, domain=boundary)
+        u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
+        u_bd_projector = fem.integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True)
+        u_bd_value = fem.integrate(
+            u_boundary_value,
+            fields={"v": u_bd_test},
+            values={"top_vel": args.top_velocity},
+            nodal=True,
+            output_dtype=wp.vec2d,
+        )
+
+        fem.normalize_dirichlet_projector(u_bd_projector, u_bd_value)
+
+        u_bd_rhs = wp.zeros_like(u_bd_value)
+        fem.project_linear_system(u_matrix, u_bd_rhs, u_bd_projector, u_bd_value, normalize_projector=False)
+
+        # div_bd_rhs = div_matrix * u_bd_rhs
+        div_bd_rhs = wp.zeros(shape=(div_matrix.nrow,), dtype=div_matrix.scalar_type)
+        bsr_mv(div_matrix, u_bd_value, y=div_bd_rhs, alpha=-1.0)
+
+        # div_matrix = div_matrix - div_matrix * bd_projector
+        bsr_mm(x=bsr_copy(div_matrix), y=u_bd_projector, z=div_matrix, alpha=-1.0, beta=1.0)
+
+        # Assemble saddle system
+        self._saddle_system = SaddleSystem(u_matrix, div_matrix)
+
+        # Save data for computing time steps rhs
+        self._u_bd_projector = u_bd_projector
+        self._u_bd_rhs = u_bd_rhs
+        self._u_test = u_test
+        self._div_bd_rhs = div_bd_rhs
+
+        # Velocitiy and pressure fields
+        self._u_field = u_space.make_field()
+        self._p_field = p_space.make_field()
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface_vector("velocity", self._u_field)
+
+    def step(self):
+        self.current_frame += 1
+
+        u_rhs = fem.integrate(
+            transported_inertia_form,
+            fields={"u": self._u_field, "v": self._u_test},
+            values={"dt": self.sim_dt},
+            output_dtype=wp.vec2d,
+        )
+
+        # Apply boundary conditions
+        # u_rhs = (I - P) * u_rhs + u_bd_rhs
+        bsr_mv(self._u_bd_projector, x=u_rhs, y=u_rhs, alpha=-1.0, beta=1.0)
+        array_axpy(x=self._u_bd_rhs, y=u_rhs, alpha=1.0, beta=1.0)
+
+        p_rhs = self._div_bd_rhs
+
+        x_u = wp.empty_like(u_rhs)
+        x_p = wp.empty_like(p_rhs)
+        wp.utils.array_cast(out_array=x_u, in_array=self._u_field.dof_values)
+        wp.utils.array_cast(out_array=x_p, in_array=self._p_field.dof_values)
+
+        bsr_solve_saddle(
+            saddle_system=self._saddle_system,
+            tol=1.0e-6,
+            x_u=x_u,
+            x_p=x_p,
+            b_u=u_rhs,
+            b_p=p_rhs,
+            quiet=self._quiet,
+        )
+
+        wp.utils.array_cast(in_array=x_u, out_array=self._u_field.dof_values)
+        wp.utils.array_cast(in_array=x_p, out_array=self._p_field.dof_values)
+
+    def render(self):
+        self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
+        self.renderer.add_surface_vector("velocity", self._u_field)
+        self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    for k in range(args.num_frames):
+        print(f"Frame {k}:")
+        example.step()
+        example.render()
+
+    example.renderer.add_surface_vector("velocity_final", example._u_field)
+    example.renderer.plot(streamlines=set(["velocity_final"]))

warp/examples/fem/example_stokes.py

@@ -0,0 +1,192 @@
+# 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 Stokes
+#
+# This example solves a 2D Stokes flow problem
+#
+# -nu D(u) + grad p = 0
+# Div u = 0
+# 
+# with (soft) velocity-Dirichlet boundary conditions
+###########################################################################
+
+import argparse
+
+import warp as wp
+
+import warp.fem as fem
+from warp.fem.utils import array_axpy
+
+import warp.sparse as sparse
+
+try:
+    from .plot_utils import Plot
+    from .bsr_utils import bsr_solve_saddle, SaddleSystem
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+except ImportError:
+    from plot_utils import Plot
+    from bsr_utils import bsr_solve_saddle, SaddleSystem
+    from mesh_utils import gen_trimesh, gen_quadmesh
+
+wp.init()
+
+
+@fem.integrand
+def constant_form(val: wp.vec2):
+    return val
+
+
+@fem.integrand
+def viscosity_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
+    return nu * wp.ddot(fem.D(u, s), fem.D(v, s))
+
+
+@fem.integrand
+def top_mass_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # non zero on top boundary of domain only
+    nor = fem.normal(domain, s)
+    return wp.dot(u(s), v(s)) * wp.max(0.0, nor[1])
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def div_form(
+    s: fem.Sample,
+    u: fem.Field,
+    q: fem.Field,
+):
+    return q(s) * fem.div(u, s)
+
+
+class Example:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resolution", type=int, default=50)
+    parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--top_velocity", type=float, default=1.0)
+    parser.add_argument("--viscosity", type=float, default=1.0)
+    parser.add_argument("--boundary_strength", type=float, default=100.0)
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+    parser.add_argument(
+        "--nonconforming_pressures", action="store_true", help="For grid, use non-conforming pressure (Q_d/P_{d-1})"
+    )
+
+    def __init__(self, stage=None, quiet=False, args=None, **kwargs):
+        if args is None:
+            # Read args from kwargs, add default arg values from parser
+            args = argparse.Namespace(**kwargs)
+            args = Example.parser.parse_args(args=[], namespace=args)
+        self._args = args
+        self._quiet = quiet
+
+        # Grid or triangle mesh geometry
+        if args.mesh == "tri":
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif args.mesh == "quad":
+            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(args.resolution))
+
+        # Function spaces -- Q_d for vel, P_{d-1} for pressure
+        u_space = fem.make_polynomial_space(geo, degree=args.degree, dtype=wp.vec2)
+        if args.mesh != "tri" and args.nonconforming_pressures:
+            p_space = fem.make_polynomial_space(
+                geo, degree=args.degree - 1, element_basis=fem.ElementBasis.NONCONFORMING_POLYNOMIAL
+            )
+        else:
+            p_space = fem.make_polynomial_space(geo, degree=args.degree - 1)
+
+        # Vector and scalar fields
+        self._u_field = u_space.make_field()
+        self._p_field = p_space.make_field()
+
+        # Interpolate initial condition on boundary (for example purposes)
+        self._bd_field = u_space.make_field()
+        f_boundary = fem.make_restriction(self._bd_field, domain=fem.BoundarySides(geo))
+        top_velocity = wp.vec2(args.top_velocity, 0.0)
+        fem.interpolate(constant_form, dest=f_boundary, values={"val": top_velocity})
+
+        self.renderer = Plot(stage)
+
+    def step(self):
+        args = self._args
+        u_space = self._u_field.space
+        p_space = self._p_field.space
+        geo = u_space.geometry
+
+        domain = fem.Cells(geometry=geo)
+        boundary = fem.BoundarySides(geo)
+
+        # Viscosity
+        u_test = fem.make_test(space=u_space, domain=domain)
+        u_trial = fem.make_trial(space=u_space, domain=domain)
+
+        u_visc_matrix = fem.integrate(
+            viscosity_form,
+            fields={"u": u_trial, "v": u_test},
+            values={"nu": args.viscosity},
+        )
+
+        # Weak velocity boundary conditions
+        u_bd_test = fem.make_test(space=u_space, domain=boundary)
+        u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
+        u_rhs = fem.integrate(
+            top_mass_form, fields={"u": self._bd_field.trace(), "v": u_bd_test}, output_dtype=wp.vec2d
+        )
+        u_bd_matrix = fem.integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test})
+
+        # Pressure-velocity coupling
+        p_test = fem.make_test(space=p_space, domain=domain)
+        div_matrix = fem.integrate(div_form, fields={"u": u_trial, "q": p_test})
+
+        # Define and solve the saddle-point system
+        u_matrix = u_visc_matrix
+        sparse.bsr_axpy(x=u_bd_matrix, y=u_matrix, alpha=args.boundary_strength, beta=1.0)
+        array_axpy(x=u_rhs, y=u_rhs, alpha=0.0, beta=args.boundary_strength)
+
+        p_rhs = wp.zeros(p_space.node_count(), dtype=wp.float64)
+        x_u = wp.zeros_like(u_rhs)
+        x_p = wp.zeros_like(p_rhs)
+
+        bsr_solve_saddle(
+            SaddleSystem(A=u_matrix, B=div_matrix), x_u=x_u, x_p=x_p, b_u=u_rhs, b_p=p_rhs, quiet=self._quiet
+        )
+
+        wp.utils.array_cast(in_array=x_u, out_array=self._u_field.dof_values)
+        wp.utils.array_cast(in_array=x_p, out_array=self._p_field.dof_values)
+
+    def render(self):
+        self.renderer.add_surface("pressure", self._p_field)
+        self.renderer.add_surface_vector("velocity", self._u_field)
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    example.step()
+    example.render()
+
+    example.renderer.plot()