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

examples/fem/example_navier_stokes.py

@@ -11,26 +11,24 @@ and using semi-Lagrangian advection
 import argparse
 
 import warp as wp
+import warp.fem as fem
 
 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
-from warp.fem.space import make_polynomial_space
-from warp.fem.quadrature import RegularQuadrature
-from warp.fem.domain import Cells, BoundarySides
-from warp.fem.integrate import integrate
-from warp.fem.operator import integrand, D, div, lookup
-from warp.fem.dirichlet import project_linear_system, normalize_dirichlet_projector
 from warp.fem.utils import array_axpy
 
 from warp.sparse import bsr_mm, bsr_mv, bsr_copy
 
-from bsr_utils import bsr_to_scipy
-from plot_utils import plot_grid_streamlines, plot_velocities
-from mesh_utils import gen_trimesh
+try:
+    from .bsr_utils import bsr_to_scipy
+    from .plot_utils import Plot
+    from .mesh_utils import gen_trimesh
+except ImportError:
+    from bsr_utils import bsr_to_scipy
+    from plot_utils import Plot
+    from mesh_utils import gen_trimesh
 
+# need to solve a saddle-point system, use scopy for simplicity
 from scipy.sparse import bmat
 from scipy.sparse.linalg import factorized
 
@@ -38,8 +36,8 @@ import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 
 
-@integrand
-def u_boundary_value(s: Sample, domain: Domain, v: Field, top_vel: float):
+@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))
@@ -47,194 +45,209 @@ def u_boundary_value(s: Sample, domain: Domain, v: Field, top_vel: float):
     return wp.dot(wp.vec2f(0.0, 0.0), 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
-def inertia_form(s: Sample, u: Field, v: Field, dt: float):
+@fem.integrand
+def inertia_form(s: fem.Sample, u: fem.Field, v: fem.Field, dt: float):
     return mass_form(s, u, v) / dt
 
 
-@integrand
-def viscosity_form(s: Sample, u: Field, v: Field, nu: float):
-    return 2.0 * 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 2.0 * nu * wp.ddot(fem.D(u, s), fem.D(v, s))
 
 
-@integrand
-def viscosity_and_inertia_form(s: Sample, u: Field, v: Field, dt: float, nu: float):
+@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)
 
 
-@integrand
-def transported_inertia_form(s: Sample, domain: Domain, u: Field, v: Field, dt: float):
+@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 = lookup(domain, conv_pos, s)
+    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(lookup(domain, conv_pos, conv_s))
+    conv_vel = u(fem.lookup(domain, conv_pos, conv_s))
 
     return wp.dot(conv_vel, v(s)) / dt
 
 
-@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)
 
 
-if __name__ == "__main__":
-    wp.init()
-    wp.set_module_options({"enable_backward": False})
-
+class Example:
     parser = argparse.ArgumentParser()
     parser.add_argument("--resolution", type=int, default=25)
     parser.add_argument("--degree", type=int, default=2)
-    parser.add_argument("--n_frames", type=int, default=1000)
+    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")
-    args = parser.parse_args()
-
-    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))
-
-    boundary = BoundarySides(geo)
-
-    viscosity = args.top_velocity / args.Re
-    dt = 1.0 / args.resolution
-
-    domain = Cells(geometry=geo)
-
-    # Functions spaces: Q(d)-Q(d-1)
-    u_degree = args.degree
-    u_space = make_polynomial_space(geo, degree=u_degree, dtype=wp.vec2)
-    p_space = make_polynomial_space(geo, degree=u_degree - 1)
-    quadrature = RegularQuadrature(domain=domain, order=2 * u_degree)
-
-    # Viscosity and inertia
-    u_test = make_test(space=u_space, domain=domain)
-    u_trial = make_trial(space=u_space, domain=domain)
-
-    u_matrix = integrate(
-        viscosity_and_inertia_form,
-        fields={"u": u_trial, "v": u_test},
-        values={"nu": viscosity, "dt": dt},
-    )
-
-    # Pressure-velocity coupling
-    p_test = make_test(space=p_space, domain=domain)
-    div_matrix = 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 = make_test(space=u_space, domain=boundary)
-    u_bd_trial = make_trial(space=u_space, domain=boundary)
-    u_bd_projector = integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True)
-    u_bd_value = integrate(
-        u_boundary_value,
-        fields={"v": u_bd_test},
-        values={"top_vel": args.top_velocity},
-        nodal=True,
-        output_dtype=wp.vec2d,
-    )
-
-    normalize_dirichlet_projector(u_bd_projector, u_bd_value)
-
-    u_bd_rhs = wp.zeros_like(u_bd_value)
-    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_rhs, y=div_bd_rhs)
-
-    # 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 with Scipy
-    div_matrix = bsr_to_scipy(div_matrix)
-    u_matrix = bsr_to_scipy(u_matrix)
-    div_bd_rhs = div_bd_rhs.numpy()
-
-    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],
-        ],
-    )
-
-    with wp.ScopedTimer("LU factorization"):
-        solve_saddle = factorized(saddle_system)
-
-    u_k = u_space.make_field()
-    u_rhs = wp.zeros_like(u_bd_rhs)
-
-    results = [u_k.dof_values.numpy()]
-
-    for k in range(args.n_frames):
-        print("Solving step", k)
-
-        u_inertia_rhs = integrate(
+
+    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_rhs, y=div_bd_rhs)
+
+        # 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 with Scipy
+        div_matrix = bsr_to_scipy(div_matrix)
+        u_matrix = bsr_to_scipy(u_matrix)
+        div_bd_rhs = div_bd_rhs.numpy()
+
+        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],
+            ],
+        )
+
+        with wp.ScopedTimer("LU factorization"):
+            self._solve_saddle = factorized(saddle_system)
+
+        # 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 field
+
+        self._u_field = u_space.make_field()
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface_vector("velocity", self._u_field)
+
+    def update(self):
+        self.current_frame += 1
+
+        u_rhs = fem.integrate(
             transported_inertia_form,
-            quadrature=quadrature,
-            fields={"u": u_k, "v": u_test},
-            values={"dt": dt},
+            fields={"u": self._u_field, "v": self._u_test},
+            values={"dt": self.sim_dt},
             output_dtype=wp.vec2d,
         )
-        # u_rhs = (I - P) * u_inertia_rhs + u_bd_rhs
-        bsr_mv(u_bd_projector, u_inertia_rhs, y=u_rhs, alpha=-1.0, beta=0.0)
-        array_axpy(x=u_inertia_rhs, y=u_rhs, alpha=1.0, beta=1.0)
-        array_axpy(x=u_bd_rhs, y=u_rhs, alpha=1.0, beta=1.0)
+
+        # 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)
 
         # Assemble scipy saddle system rhs
-        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())
+        u_dof_count = self._u_bd_projector.shape[0]
+        p_dof_count = self._div_bd_rhs.shape[0]
+        tot_dof_count = u_dof_count + p_dof_count + 1
+
+        u_slice = slice(0, u_dof_count)
+        p_slice = slice(u_dof_count, tot_dof_count - 1)
+
+        saddle_rhs = np.zeros(tot_dof_count)
         saddle_rhs[u_slice] = u_rhs.numpy().flatten()
-        saddle_rhs[p_slice] = div_bd_rhs
+        saddle_rhs[p_slice] = self._div_bd_rhs
 
-        x = solve_saddle(saddle_rhs)
+        x = self._solve_saddle(saddle_rhs)
 
         # Extract result
-        x_u = x[u_slice].reshape((-1, 2))
-        results.append(x_u)
-
-        u_k.dof_values = x_u
-        # p_field.dof_values = x[p_slice]
-
-    if isinstance(geo, Grid2D):
-        plot_grid_streamlines(u_k)
-
-    quiver = plot_velocities(u_k)
-    ax = quiver.axes
-
-    def animate(i):
-        ax.clear()
-        u_k.dof_values = results[i]
-        return plot_velocities(u_k, axes=ax)
-
-    anim = animation.FuncAnimation(
-        ax.figure,
-        animate,
-        interval=30,
-        blit=False,
-        frames=len(results),
-    )
-    plt.show()
+        self._u_field.dof_values = x[u_slice].reshape((-1, 2))
+
+    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.init()
+    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.update()
+        example.render()
+
+    example.renderer.add_surface_vector("velocity_final", example._u_field)
+    example.renderer.plot(streamlines=set(["velocity_final"]))