warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

examples/fem/example_stokes.py

@@ -13,152 +13,184 @@ import warp as wp
 
 import numpy as np
 
-from warp.fem.types import *
-from warp.fem.geometry import Grid2D, Trimesh2D
-from warp.fem.field import make_test, make_trial, make_restriction
-from warp.fem.space import make_polynomial_space
-from warp.fem.domain import Cells, BoundarySides
-from warp.fem.integrate import integrate, interpolate
-from warp.fem.operator import normal, integrand, D, div
-
-from plot_utils import plot_velocities, plot_surface
-from bsr_utils import bsr_to_scipy
-from mesh_utils import gen_trimesh
-
+import warp.fem as fem
+
+try:
+    from .plot_utils import Plot
+    from .bsr_utils import bsr_to_scipy
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+except ImportError:
+    from plot_utils import Plot
+    from bsr_utils import bsr_to_scipy
+    from mesh_utils import gen_trimesh, gen_quadmesh
+
+# Need to solve a saddle-point system, use scipy for simplicity
 from scipy.sparse import bmat
 from scipy.sparse.linalg import spsolve
 
-import matplotlib.pyplot as plt
-
 
-@integrand
+@fem.integrand
 def constant_form(val: wp.vec2):
     return val
 
 
-@integrand
-def viscosity_form(s: Sample, u: Field, v: Field, nu: float):
-    return nu * wp.ddot(D(u, s), D(v, s))
+@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))
 
 
-@integrand
+@fem.integrand
 def top_mass_form(
-    s: Sample,
-    domain: Domain,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
 ):
     # non zero on top boundary of domain only
-    nor = normal(domain, s)
+    nor = fem.normal(domain, s)
     return wp.dot(u(s), v(s)) * wp.max(0.0, nor[1])
 
 
-@integrand
+@fem.integrand
 def mass_form(
-    s: Sample,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
 ):
     return wp.dot(u(s), v(s))
 
 
-@integrand
+@fem.integrand
 def div_form(
-    s: Sample,
-    u: Field,
-    q: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    q: fem.Field,
 ):
-    return q(s) * div(u, s)
-
-
-if __name__ == "__main__":
-    wp.init()
-    wp.set_module_options({"enable_backward": False})
+    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("--tri_mesh", action="store_true", help="Use a triangular mesh")
-    args = parser.parse_args()
-    
-    top_velocity = wp.vec2(args.top_velocity, 0.0)
-
-    if args.tri_mesh:
-        positions, tri_vidx = gen_trimesh(res=vec2i(args.resolution))
-        geo = Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
-    else:
-        geo = Grid2D(res=vec2i(args.resolution))
-    
-    domain = Cells(geometry=geo)
-    boundary = BoundarySides(geo)
-
-    # Function spaces -- Q_d for vel, Q_{d-1} for pressure
-    u_space = make_polynomial_space(geo, degree=args.degree, dtype=wp.vec2)
-    p_space = make_polynomial_space(geo, degree=args.degree-1)
-
-    # Interpolate initial condition on boundary (mostly for testing)
-    f = u_space.make_field()
-    f_boundary = make_restriction(f, domain=boundary)
-    interpolate(constant_form, dest=f_boundary, values={"val": top_velocity})
-
-    # Viscosity
-    u_test = make_test(space=u_space, domain=domain)
-    u_trial = make_trial(space=u_space, domain=domain)
-
-    u_visc_matrix = integrate(
-        viscosity_form,
-        fields={"u": u_trial, "v": u_test},
-        values={"nu": args.viscosity},
+    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})"
     )
 
-    # Weak velocity boundary conditions
-    u_bd_test = make_test(space=u_space, domain=boundary)
-    u_bd_trial = make_trial(space=u_space, domain=boundary)
-    u_rhs = integrate(top_mass_form, fields={"u": f.trace(), "v": u_bd_test})
-    u_bd_matrix = integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test})
-
-    # Pressure-velocity coupling
-    p_test = make_test(space=p_space, domain=domain)
-    div_matrix = integrate(div_form, fields={"u": u_trial, "q": p_test})
-
-    # Solve with scipy
-    # Assemble saddle-point system with velocity, pressure, and zero-average-pressure constraint
-    u_rhs = u_rhs.numpy() * args.boundary_strength
-    u_matrix = bsr_to_scipy(u_visc_matrix) + args.boundary_strength * bsr_to_scipy(u_bd_matrix)
-
-    div_matrix = bsr_to_scipy(div_matrix)
-
-    ones = np.ones(shape=(p_space.node_count(), 1), dtype=float)
-    saddle_system = bmat(
-        [
-            [u_matrix, div_matrix.transpose(), None],
-            [div_matrix, None, ones],
-            [None, ones.transpose(), None],
-        ],
-        format="csr",
-    )
-
-    saddle_rhs = np.zeros(saddle_system.shape[0])
-    u_slice = slice(0, 2 * u_space.node_count())
-    p_slice = slice(
-        2 * u_space.node_count(), 2 * u_space.node_count() + p_space.node_count()
-    )
-    saddle_rhs[u_slice] = u_rhs.flatten()
+    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 update(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})
+        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})
+
+        # Solve with scipy
+        # Assemble saddle-point system with velocity, pressure, and zero-average-pressure constraint
+        u_rhs = u_rhs.numpy() * args.boundary_strength
+        u_matrix = bsr_to_scipy(u_visc_matrix) + args.boundary_strength * bsr_to_scipy(u_bd_matrix)
+
+        div_matrix = bsr_to_scipy(div_matrix)
+
+        ones = np.ones(shape=(p_space.node_count(), 1), dtype=float)
+        saddle_system = bmat(
+            [
+                [u_matrix, div_matrix.transpose(), None],
+                [div_matrix, None, ones],
+                [None, ones.transpose(), None],
+            ],
+            format="csr",
+        )
+
+        saddle_rhs = np.zeros(saddle_system.shape[0])
+        u_slice = slice(0, 2 * u_space.node_count())
+        p_slice = slice(2 * u_space.node_count(), 2 * u_space.node_count() + p_space.node_count())
+        saddle_rhs[u_slice] = u_rhs.flatten()
+
+        x = spsolve(saddle_system, saddle_rhs)
+
+        # Extract result
+
+        self._u_field.dof_values = x[u_slice].reshape((-1, 2))
+        self._p_field.dof_values = x[p_slice]
+
+    def render(self):
+        self.renderer.add_surface("pressure", self._p_field)
+        self.renderer.add_surface_vector("velocity", self._u_field)
 
-    x = spsolve(saddle_system, saddle_rhs)
 
-    # Extract result
-    u_field = u_space.make_field()
-    p_field = p_space.make_field()
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
 
-    u_field.dof_values = x[u_slice].reshape((-1, 2))
-    p_field.dof_values = x[p_slice]
+    args = Example.parser.parse_args()
 
-    plot_surface(p_field)
-    plot_velocities(u_field)
+    example = Example(args=args)
+    example.update()
+    example.render()
 
-    plt.show()
+    example.renderer.plot()

examples/fem/example_stokes_transfer.py

@@ -9,65 +9,62 @@ import math
 import warp as wp
 import numpy as np
 
-from warp.fem.types import *
-from warp.fem.geometry import Grid2D, ExplicitGeometryPartition
-from warp.fem.field import make_test, make_trial
-from warp.fem.space import make_polynomial_space, make_space_partition
-from warp.fem.domain import Cells
-from warp.fem.integrate import integrate, interpolate
-from warp.fem.operator import integrand, D, div
-from warp.fem.quadrature import PicQuadrature
-from warp.fem.utils import array_axpy
+import warp.fem as fem
 
 from warp.utils import array_cast
-from warp.sparse import bsr_transposed, bsr_mm, bsr_axpy
-
-from plot_utils import plot_grid_streamlines, plot_grid_surface
-from bsr_utils import bsr_cg, bsr_mv
+from warp.fem.utils import array_axpy
+from warp.sparse import bsr_transposed, bsr_mm, bsr_axpy, bsr_mv
 
-import matplotlib.pyplot as plt
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .bsr_utils import bsr_cg
+    from .plot_utils import Plot
+except ImportError:
+    from bsr_utils import bsr_cg
+    from plot_utils import Plot
 
 
-@integrand
-def vel_from_particles_form(s: Sample, particle_vel: wp.array(dtype=wp.vec2), v: Field):
+@fem.integrand
+def vel_from_particles_form(s: fem.Sample, particle_vel: wp.array(dtype=wp.vec2), v: fem.Field):
     vel = particle_vel[s.qp_index]
     return wp.dot(vel, v(s))
 
 
-@integrand
-def viscosity_form(s: Sample, u: Field, v: Field, nu: float):
-    return nu * wp.ddot(D(u, s), D(v, s))
+@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))
 
 
-@integrand
+@fem.integrand
 def mass_form(
-    s: Sample,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
 ):
     return wp.dot(u(s), v(s))
 
 
-@integrand
+@fem.integrand
 def scalar_mass_form(
-    s: Sample,
-    p: Field,
-    q: Field,
+    s: fem.Sample,
+    p: fem.Field,
+    q: fem.Field,
 ):
     return p(s) * q(s)
 
 
-@integrand
+@fem.integrand
 def div_form(
-    s: Sample,
-    u: Field,
-    q: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    q: fem.Field,
 ):
-    return q(s) * div(u, s)
+    return q(s) * fem.div(u, s)
 
 
-@integrand
-def cell_activity(s: Sample, domain: Domain, c1: wp.vec2, c2: wp.vec2, radius: float):
+@fem.integrand
+def cell_activity(s: fem.Sample, domain: fem.Domain, c1: wp.vec2, c2: wp.vec2, radius: float):
     pos = domain(s)
     if wp.length(pos - c1) < radius:
         return 0.0
@@ -81,130 +78,162 @@ def inverse_array_kernel(m: wp.array(dtype=wp.float64)):
     m[wp.tid()] = wp.float64(1.0) / m[wp.tid()]
 
 
+class Example:
+    def __init__(self, stage=None, quiet=False):
+        self._quiet = quiet
+
+        self.res = 50
+        self.cell_size = 1.0 / self.res
+
+        self.vel = 1.0
+        self.viscosity = 100.0
+        self.compliance = 0.01
+        self.bd_strength = 100000.0
+
+        geo = fem.Grid2D(res=wp.vec2i(self.res))
+
+        # Displacement boundary conditions are defined by two circles going in opposite directions
+        # Sample particles along those
+        circle_radius = 0.15
+        c1_center = wp.vec2(0.25, 0.5)
+        c2_center = wp.vec2(0.75, 0.5)
+        particles, particle_areas, particle_velocities = self._gen_particles(circle_radius, c1_center, c2_center)
+
+        # Disable cells that are interior to the circles
+        cell_space = fem.make_polynomial_space(geo, degree=0)
+        activity = cell_space.make_field()
+        fem.interpolate(
+            cell_activity,
+            dest=activity,
+            values={"c1": c1_center, "c2": c2_center, "radius": circle_radius - self.cell_size},
+        )
+
+        # Explicitly define the active geometry partition from those cells
+        self._active_partition = fem.ExplicitGeometryPartition(geo, wp.array(activity.dof_values.numpy(), dtype=int))
+        if not self._quiet:
+            print("Active cells:", self._active_partition.cell_count())
+
+        # Function spaces -- Q1 for vel, Q0 for pressure
+        u_space = fem.make_polynomial_space(geo, degree=1, dtype=wp.vec2)
+        p_space = fem.make_polynomial_space(geo, degree=0)
+
+        self._active_space_partition = fem.make_space_partition(
+            space=u_space, geometry_partition=self._active_partition
+        )
+        self._active_p_space_partition = fem.make_space_partition(
+            space=p_space, geometry_partition=self._active_partition
+        )
+
+        self._u_field = u_space.make_field()
+        self._p_field = p_space.make_field()
+
+        # Particle-based quadrature rule over active cells
+        domain = fem.Cells(geometry=self._active_partition)
+        self._pic_quadrature = fem.PicQuadrature(domain, particles, particle_areas)
+        self._particle_velocities = particle_velocities
+
+        self.renderer = Plot(stage)
+
+    def update(self):
+        u_space = self._u_field.space
+        p_space = self._p_field.space
+
+        # Weakly-enforced boundary condition on particles
+        u_test = fem.make_test(space=u_space, space_partition=self._active_space_partition)
+        u_trial = fem.make_trial(space=u_space, space_partition=self._active_space_partition)
+
+        u_rhs = fem.integrate(
+            vel_from_particles_form,
+            quadrature=self._pic_quadrature,
+            fields={"v": u_test},
+            values={"particle_vel": self._particle_velocities},
+            output_dtype=wp.vec2d,
+        )
+        u_bd_matrix = fem.integrate(mass_form, quadrature=self._pic_quadrature, fields={"u": u_trial, "v": u_test})
+
+        # Viscosity
+        u_visc_matrix = fem.integrate(
+            viscosity_form,
+            fields={"u": u_trial, "v": u_test},
+            values={"nu": self.viscosity},
+        )
+
+        # Pressure-velocity coupling
+        p_test = fem.make_test(space=p_space, space_partition=self._active_p_space_partition)
+        p_trial = fem.make_trial(space=p_space, space_partition=self._active_p_space_partition)
+
+        div_matrix = fem.integrate(div_form, fields={"u": u_trial, "q": p_test})
+        inv_p_mass_matrix = fem.integrate(scalar_mass_form, fields={"p": p_trial, "q": p_test})
+        wp.launch(
+            kernel=inverse_array_kernel,
+            dim=inv_p_mass_matrix.values.shape,
+            device=inv_p_mass_matrix.values.device,
+            inputs=[inv_p_mass_matrix.values],
+        )
+
+        # Assemble linear system
+        u_matrix = u_visc_matrix
+        bsr_axpy(u_bd_matrix, u_matrix, alpha=self.bd_strength)
+
+        div_matrix_t = bsr_transposed(div_matrix)
+        gradient_matrix = bsr_mm(div_matrix_t, inv_p_mass_matrix)
+        bsr_mm(gradient_matrix, div_matrix, u_matrix, alpha=1.0 / self.compliance, beta=1.0)
+
+        array_axpy(u_rhs, u_rhs, alpha=0.0, beta=self.bd_strength)
+
+        # Solve for displacement
+        u_res = wp.zeros_like(u_rhs)
+        bsr_cg(u_matrix, x=u_res, b=u_rhs, quiet=self._quiet)
+
+        # Compute pressure from displacement
+        div_u = bsr_mv(A=div_matrix, x=u_res)
+        p_res = bsr_mv(A=inv_p_mass_matrix, x=div_u, alpha=-1)
+
+        # Copy to fields
+        u_nodes = wp.indexedarray(self._u_field.dof_values, indices=self._active_space_partition.space_node_indices())
+        p_nodes = wp.indexedarray(self._p_field.dof_values, indices=self._active_p_space_partition.space_node_indices())
+
+        array_cast(in_array=u_res, out_array=u_nodes)
+        array_cast(in_array=p_res, out_array=p_nodes)
+
+    def render(self):
+        self.renderer.add_surface("pressure", self._p_field)
+        self.renderer.add_surface_vector("velocity", self._u_field)
+
+    def _gen_particles(self, circle_radius, c1_center, c2_center):
+        """Generate some particles along two circles defining velocity boundary conditions"""
+
+        # Generate particles defining the transfer displacement
+        particles_per_circle = int(2.0 * math.pi * circle_radius * self.res)
+
+        angles = np.linspace(0, 2.0 * math.pi, particles_per_circle, endpoint=False)
+
+        n_particles = 2 * particles_per_circle
+        particles = np.empty((n_particles, 2), dtype=float)
+
+        particles[:particles_per_circle, 0] = c1_center[0] + circle_radius * np.cos(angles)
+        particles[:particles_per_circle, 1] = c1_center[1] + circle_radius * np.sin(angles)
+        particles[particles_per_circle:, 0] = c2_center[0] + circle_radius * np.cos(angles)
+        particles[particles_per_circle:, 1] = c2_center[1] + circle_radius * np.sin(angles)
+
+        particle_areas = np.ones(n_particles) * self.cell_size**2
+        particle_velocities = np.zeros_like(particles)
+        particle_velocities[:particles_per_circle, 0] = self.vel
+        particle_velocities[particles_per_circle:, 0] = -self.vel
+
+        particles = wp.array(particles, dtype=wp.vec2)
+        particle_areas = wp.array(particle_areas, dtype=float)
+        particle_velocities = wp.array(particle_velocities, dtype=wp.vec2)
+
+        return particles, particle_areas, particle_velocities
+
+
 if __name__ == "__main__":
     wp.init()
     wp.set_module_options({"enable_backward": False})
 
-    res = 50
-    cell_size = 1.0 / res
-    geo = Grid2D(res=vec2i(res))
-
-    vel = 1.0
-    viscosity = 100.0
-    compliance = 0.01
-    bd_strength = 100000.0
-
-    # Generate particles defining the transfer displacement
-    circle_radius = 0.15
-    c1_center = wp.vec2(0.25, 0.5)
-    c2_center = wp.vec2(0.75, 0.5)
-    particles_per_circle = int(2.0 * math.pi * circle_radius * res)
-
-    angles = np.linspace(0, 2.0 * math.pi, particles_per_circle, endpoint=False)
-
-    n_particles = 2 * particles_per_circle
-    particles = np.empty((n_particles, 2), dtype=float)
-
-    particles[:particles_per_circle, 0] = c1_center[0] + circle_radius * np.cos(angles)
-    particles[:particles_per_circle, 1] = c1_center[1] + circle_radius * np.sin(angles)
-    particles[particles_per_circle:, 0] = c2_center[0] + circle_radius * np.cos(angles)
-    particles[particles_per_circle:, 1] = c2_center[1] + circle_radius * np.sin(angles)
-
-    particle_areas = np.ones(n_particles) * cell_size**2
-    particle_velocities = np.zeros_like(particles)
-    particle_velocities[:particles_per_circle, 0] = vel
-    particle_velocities[particles_per_circle:, 0] = -vel
-
-    particles = wp.array(particles, dtype=wp.vec2)
-    particle_areas = wp.array(particle_areas, dtype=float)
-    particle_velocities = wp.array(particle_velocities, dtype=wp.vec2)
-
-    # Disable cells that are interior to the circles
-    cell_space = make_polynomial_space(geo, degree=0)
-    activity = cell_space.make_field()
-    interpolate(
-        cell_activity,
-        dest=activity,
-        values={"c1": c1_center, "c2": c2_center, "radius": circle_radius - cell_size},
-    )
-
-    active_partition = ExplicitGeometryPartition(geo, wp.array(activity.dof_values.numpy(), dtype=int))
-    print("Active cells:", active_partition.cell_count())
-
-    # Function spaces -- Q1 for vel, Q0 for pressure
-    u_space = make_polynomial_space(geo, degree=1, dtype=wp.vec2)
-    p_space = make_polynomial_space(geo, degree=0)
-    active_space_partition = make_space_partition(space=u_space, geometry_partition=active_partition)
-    active_p_space_partition = make_space_partition(space=p_space, geometry_partition=active_partition)
-
-    domain = Cells(geometry=active_partition)
-    pic_quadrature = PicQuadrature(domain, particles, particle_areas)
-
-    # Boundary condition on particles
-    u_test = make_test(space=u_space, space_partition=active_space_partition, domain=domain)
-    u_trial = make_trial(space=u_space, space_partition=active_space_partition, domain=domain)
-
-    u_rhs = integrate(
-        vel_from_particles_form,
-        quadrature=pic_quadrature,
-        fields={"v": u_test},
-        values={"particle_vel": particle_velocities},
-        output_dtype=wp.vec2d
-    )
-
-    u_bd_matrix = integrate(mass_form, quadrature=pic_quadrature, fields={"u": u_trial, "v": u_test})
-
-    # Viscosity
-    u_visc_matrix = integrate(
-        viscosity_form,
-        fields={"u": u_trial, "v": u_test},
-        values={"nu": viscosity},
-    )
-
-    # Pressure-velocity coupling
-    p_test = make_test(space=p_space, space_partition=active_p_space_partition, domain=domain)
-    p_trial = make_trial(space=p_space, space_partition=active_p_space_partition, domain=domain)
-
-    div_matrix = integrate(div_form, fields={"u": u_trial, "q": p_test})
-    inv_p_mass_matrix = integrate(scalar_mass_form, fields={"p": p_trial, "q": p_test})
-    wp.launch(
-        kernel=inverse_array_kernel,
-        dim=inv_p_mass_matrix.values.shape,
-        device=inv_p_mass_matrix.values.device,
-        inputs=[inv_p_mass_matrix.values],
-    )
-
-    # Assemble linear system
-    u_matrix = u_visc_matrix
-    bsr_axpy(u_bd_matrix, u_matrix, alpha=bd_strength)
-
-    div_matrix_t = bsr_transposed(div_matrix)
-    gradient_matrix = bsr_mm(div_matrix_t, inv_p_mass_matrix)
-    bsr_mm(gradient_matrix, div_matrix, u_matrix, alpha=1.0 / compliance, beta=1.0)
-
-    array_axpy(u_rhs, u_rhs, alpha=0.0, beta=bd_strength)
-
-    # Solve for displacement
-    u_res = wp.zeros_like(u_rhs)
-    bsr_cg(u_matrix, x=u_res, b=u_rhs)
-
-    # Recompute pressure
-    p_res = wp.zeros(n=active_p_space_partition.node_count(), dtype=wp.float64)
-    p_tmp = wp.empty_like(p_res)
-    bsr_mv(A=div_matrix, x=u_res, y=p_tmp)
-    bsr_mv(A=inv_p_mass_matrix, x=p_tmp, y=p_res, alpha=-1)
-
-    # Display result
-    u_field = u_space.make_field()
-    p_field = p_space.make_field()
-
-    u_nodes = wp.indexedarray(u_field.dof_values, indices=active_space_partition.space_node_indices())
-    p_nodes = wp.indexedarray(p_field.dof_values, indices=active_p_space_partition.space_node_indices())
-
-    array_cast(in_array=u_res, out_array=u_nodes)
-    array_cast(in_array=p_res, out_array=p_nodes)
-
-    plot_grid_surface(p_field)
-    plot_grid_streamlines(u_field)
-
-    plt.show()
+    example = Example()
+    example.update()
+    example.render()
+
+    example.renderer.plot(streamlines=["velocity"])