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

examples/fem/example_convection_diffusion_dg.py

@@ -9,202 +9,195 @@ with upwind transport and Symmetric Interior Penalty
 import argparse
 
 import warp as wp
-
-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, Sides
-from warp.fem.integrate import integrate, interpolate
-from warp.fem.polynomial import Polynomial
-from warp.fem.operator import (
-    grad,
-    integrand,
-    jump,
-    average,
-    normal,
-    grad_average,
-    measure_ratio,
-    degree,
-)
+import warp.fem as fem
 
 from warp.sparse import bsr_axpy
 
-from bsr_utils import bsr_to_scipy
-from plot_utils import plot_surface
-from mesh_utils import gen_trimesh
-
-from example_convection_diffusion import (
-    initial_condition,
-    velocity,
-    inertia_form,
-    diffusion_form,
-)
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .bsr_utils import bsr_to_scipy
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+    from .plot_utils import Plot
+    from .example_convection_diffusion import (
+        initial_condition,
+        velocity,
+        inertia_form,
+        diffusion_form,
+    )
+except ImportError:
+    from bsr_utils import bsr_to_scipy
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from plot_utils import Plot
+    from example_convection_diffusion import (
+        initial_condition,
+        velocity,
+        inertia_form,
+        diffusion_form,
+    )
 
+# Non-SPD system, solve using scipy
 from scipy.sparse.linalg import factorized
 
-import matplotlib.pyplot as plt
-import matplotlib.animation as animation
-
 
 # Standard transport term, on cells' interior
-@integrand
-def transport_form(s: Sample, domain: Domain, phi: Field, psi: Field, ang_vel: float):
+@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(grad(phi, s), vel)
+    return psi(s) * wp.dot(fem.grad(phi, s), vel)
 
 
 # Upwind flux, on cell sides
-@integrand
-def upwind_transport_form(s: Sample, domain: Domain, phi: Field, psi: Field, ang_vel: float):
+@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, normal(domain, s))
+    vel_n = wp.dot(vel, fem.normal(domain, s))
 
-    return jump(phi, s) * (-average(psi, s) * vel_n + 0.5 * jump(psi, s) * wp.abs(vel_n))
+    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)
-@integrand
+@fem.integrand
 def sip_diffusion_form(
-    s: Sample,
-    domain: Domain,
-    psi: Field,
-    phi: Field,
+    s: fem.Sample,
+    domain: fem.Domain,
+    psi: fem.Field,
+    phi: fem.Field,
 ):
-    nor = normal(domain, s)
-    penalty = measure_ratio(domain, s) * float(degree(psi) * degree(phi))
+    nor = fem.normal(domain, s)
+    penalty = fem.measure_ratio(domain, s) * float(fem.degree(psi) * fem.degree(phi))
 
-    return penalty * jump(phi, s) * jump(psi, s) - (
-        wp.dot(grad_average(phi, s), nor) * jump(psi, s) + wp.dot(grad_average(psi, s), nor) * jump(phi, s)
+    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)
     )
 
 
-@integrand
-def identity(s: Sample, phi: Field):
-    return phi(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=50)
     parser.add_argument("--degree", type=int, default=2)
-    parser.add_argument("--n_frames", type=int, default=100)
+    parser.add_argument("--num_frames", type=int, default=100)
     parser.add_argument("--viscosity", type=float, default=0.001)
     parser.add_argument("--ang_vel", type=float, default=1.0)
-    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh")
-    args = parser.parse_args()
-
-    res = args.resolution
-    dt = 1.0 / (args.ang_vel * res)
-
-    if args.tri_mesh:
-        positions, tri_vidx = gen_trimesh(res=vec2i(res))
-        geo = Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
-    else:
-        geo = Grid2D(res=vec2i(res))
-
-    domain = Cells(geometry=geo)
-    sides = Sides(geo)
-    quadrature = RegularQuadrature(domain=domain, order=2 * args.degree)
-    scalar_space = make_polynomial_space(
-        geo,
-        discontinuous=True,
-        degree=args.degree,
-        family=Polynomial.GAUSS_LEGENDRE,
-    )
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+
+    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.ang_vel * res)
+        self.current_frame = 0
+
+        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))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+        scalar_space = fem.make_polynomial_space(
+            geo,
+            discontinuous=True,
+            degree=args.degree,
+            family=fem.Polynomial.GAUSS_LEGENDRE,
+        )
 
-    # Right-hand-side
-    phi0 = scalar_space.make_field()
-    interpolate(initial_condition, dest=phi0)
+        # Assemble transport, diffusion and inertia matrices
 
-    test = make_test(space=scalar_space, domain=domain)
-    trial = make_trial(space=scalar_space, domain=domain)
+        self._test = fem.make_test(space=scalar_space, domain=domain)
+        trial = fem.make_trial(space=scalar_space, domain=domain)
 
-    side_test = make_test(space=scalar_space, domain=sides)
-    side_trial = make_trial(space=scalar_space, domain=sides)
+        matrix_inertia = fem.integrate(
+            inertia_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
 
-    matrix_inertia = integrate(
-        inertia_form,
-        quadrature=quadrature,
-        fields={"phi": trial, "psi": test},
-        values={"dt": dt},
-    )
+        matrix_transport = fem.integrate(
+            transport_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"ang_vel": args.ang_vel},
+        )
 
-    matrix_transport = integrate(
-        transport_form,
-        fields={"phi": trial, "psi": test},
-        values={"ang_vel": args.ang_vel},
-    )
+        side_test = fem.make_test(space=scalar_space, domain=sides)
+        side_trial = fem.make_trial(space=scalar_space, domain=sides)
 
-    bsr_axpy(
-        integrate(
-            upwind_transport_form,
-            fields={"phi": side_trial, "psi": side_test},
-            values={"ang_vel": args.ang_vel},
-        ),
-        y=matrix_transport,
-    )
+        bsr_axpy(
+            fem.integrate(
+                upwind_transport_form,
+                fields={"phi": side_trial, "psi": side_test},
+                values={"ang_vel": args.ang_vel},
+            ),
+            y=matrix_transport,
+        )
 
-    matrix_diffusion = integrate(
-        diffusion_form,
-        fields={"u": trial, "v": test},
-    )
-    bsr_axpy(
-        integrate(
-            sip_diffusion_form,
-            fields={"phi": side_trial, "psi": side_test},
-        ),
-        y=matrix_diffusion,
-    )
+        matrix_diffusion = fem.integrate(
+            diffusion_form,
+            fields={"u": trial, "v": self._test},
+        )
+        bsr_axpy(
+            fem.integrate(
+                sip_diffusion_form,
+                fields={"phi": side_trial, "psi": side_test},
+            ),
+            y=matrix_diffusion,
+        )
+
+        self._matrix = matrix_inertia
+        bsr_axpy(x=matrix_transport, y=self._matrix)
+        bsr_axpy(x=matrix_diffusion, y=self._matrix, alpha=args.viscosity)
+        
+        # Compute LU factorization of system matrix
+        self._solve_lu = factorized(bsr_to_scipy(self._matrix))
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface("phi", self._phi_field)
 
-    matrix = matrix_inertia
-    bsr_axpy(x=matrix_transport, y=matrix)
-    bsr_axpy(x=matrix_diffusion, y=matrix, alpha=args.viscosity)
 
-    matrix_solve = factorized(bsr_to_scipy(matrix))
+    def update(self):
+        self.current_frame += 1
 
-    results = [phi0.dof_values.numpy()]
-    phik = phi0
-    for k in range(args.n_frames):
-        rhs = integrate(
+        rhs = fem.integrate(
             inertia_form,
-            quadrature=quadrature,
-            fields={"phi": phik, "psi": test},
-            values={"dt": dt},
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"dt": self.sim_dt},
         )
 
-        # Solve using Scipy
-        x = matrix_solve(rhs.numpy().flatten())
+        self._phi_field.dof_values = self._solve_lu(rhs.numpy())
 
-        phik.dof_values = x
-        results.append(x)
+    def render(self):
+        self.renderer.begin_frame(time = self.current_frame * self.sim_dt)
+        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer.end_frame()
 
-    colormesh = plot_surface(phi0)
-    ax = colormesh.axes
 
-    # Convert to continuous for visualization
-    viz_space = make_polynomial_space(geo, degree=args.degree)
-    phi_viz = viz_space.make_field()
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
 
-    def animate(i):
-        ax.clear()
-        phik.dof_values = results[i]
-        interpolate(identity, fields={"phi": phik}, dest=phi_viz)
+    args = Example.parser.parse_args()
 
-        return plot_surface(phi_viz, axes=ax)
+    example = Example(args=args)
+    for k in range(args.num_frames):
+        print(f"Frame {k}:")
+        example.update()
+        example.render()
 
-    anim = animation.FuncAnimation(
-        ax.figure,
-        animate,
-        interval=30,
-        blit=False,
-        frames=len(results),
-    )
-    plt.show()
+    example.renderer.plot()

examples/fem/example_convection_diffusion_dg0.py

@@ -9,174 +9,179 @@ This example simulates a convection-diffusion PDE using FVM with upwind transpor
 import argparse
 
 import warp as wp
-
-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, Sides
-from warp.fem.integrate import integrate, interpolate
-from warp.fem.operator import integrand, jump, average, normal
+import warp.fem as fem
 
 from warp.sparse import bsr_mm, bsr_axpy, bsr_transposed
 
-from bsr_utils import bsr_to_scipy, invert_diagonal_bsr_mass_matrix
-from plot_utils import plot_surface
-from mesh_utils import gen_trimesh
 
-from example_convection_diffusion import initial_condition, velocity, inertia_form
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .bsr_utils import bsr_to_scipy, invert_diagonal_bsr_mass_matrix
+    from .plot_utils import Plot
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+    from .example_convection_diffusion import initial_condition, velocity, inertia_form
+except ImportError:
+    from bsr_utils import bsr_to_scipy, invert_diagonal_bsr_mass_matrix
+    from plot_utils import Plot
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from example_convection_diffusion import initial_condition, velocity, inertia_form
 
 from scipy.sparse.linalg import factorized
 
-import matplotlib.pyplot as plt
-import matplotlib.animation as animation
-
 
-@integrand
+@fem.integrand
 def vel_mass_form(
-    s: Sample,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
 ):
     return wp.dot(v(s), u(s))
 
 
-@integrand
+@fem.integrand
 def half_diffusion_form(
-    s: Sample,
-    domain: Domain,
-    psi: Field,
-    u: Field,
+    s: fem.Sample,
+    domain: fem.Domain,
+    psi: fem.Field,
+    u: fem.Field,
 ):
-    return jump(psi, s) * wp.dot(average(u, s), normal(domain, s))
+    return fem.jump(psi, s) * wp.dot(fem.average(u, s), fem.normal(domain, s))
 
 
-@integrand
-def upwind_transport_form(s: Sample, domain: Domain, phi: Field, psi: Field, ang_vel: float):
+@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, normal(domain, s))
+    vel_n = wp.dot(vel, fem.normal(domain, s))
 
-    return jump(psi, s) * (average(phi, s) * vel_n + 0.5 * jump(phi, s) * wp.abs(vel_n))
+    return fem.jump(psi, s) * (fem.average(phi, s) * vel_n + 0.5 * fem.jump(phi, s) * wp.abs(vel_n))
 
 
-if __name__ == "__main__":
-    wp.init()
-    wp.set_module_options({"enable_backward": False})
-
+class Example:
     parser = argparse.ArgumentParser()
     parser.add_argument("--resolution", type=int, default=50)
-    parser.add_argument("--n_frames", type=int, default=250)
+    parser.add_argument("--num_frames", type=int, default=250)
     parser.add_argument("--viscosity", type=float, default=0.001)
     parser.add_argument("--ang_vel", type=float, default=1.0)
-    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh")
-
-    args = parser.parse_args()
-
-    res = args.resolution
-    dt = 1.0 / (args.ang_vel * res)
-
-    if args.tri_mesh:
-        positions, tri_vidx = gen_trimesh(res=vec2i(res))
-        geo = Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
-    else:
-        geo = Grid2D(res=vec2i(res))
-
-    domain = Cells(geometry=geo)
-    sides = Sides(geo)
-    quadrature = RegularQuadrature(domain=domain, order=2)
-    scalar_space = make_polynomial_space(geo, degree=0)
-
-    # Initial condition
-    phi0 = scalar_space.make_field()
-    interpolate(initial_condition, dest=phi0)
-
-    # Inertia matrix
-    test = make_test(space=scalar_space, domain=domain)
-    trial = make_trial(space=scalar_space, domain=domain)
-    matrix_inertia = integrate(
-        inertia_form,
-        quadrature=quadrature,
-        fields={"phi": trial, "psi": test},
-        values={"dt": dt},
-    )
-
-    # Upwind transport term
-    side_test = make_test(space=scalar_space, domain=sides)
-    side_trial = make_trial(space=scalar_space, domain=sides)
-    matrix_transport = integrate(
-        upwind_transport_form,
-        fields={"phi": side_trial, "psi": side_test},
-        values={"ang_vel": args.ang_vel},
-    )
-
-    # Diffusion bilinear form
-    # Since we have piecewise constant element, we cannot use the classical diffusion form
-    # Instead we assemble the matrix B M^-1 B^T, with B associated to the form psi div(u)
-    # and the diagonal matrix M to the velocity mass form u.v
-
-    velocity_space = make_polynomial_space(geo, degree=0, dtype=wp.vec2)
-    side_trial_vel = make_trial(space=velocity_space, domain=sides)
-    matrix_half_diffusion = integrate(
-        half_diffusion_form,
-        fields={"psi": side_test, "u": side_trial_vel},
-    )
-
-    # Diagonal velocity mass matrix
-    test_vel = make_test(space=velocity_space, domain=domain)
-    trial_vel = make_trial(space=velocity_space, domain=domain)
-    inv_vel_mass_matrix = integrate(vel_mass_form, domain=domain, fields={"u": trial_vel, "v": test_vel}, nodal=True)
-    invert_diagonal_bsr_mass_matrix(inv_vel_mass_matrix)
-
-    # Assemble system matrix
-
-    matrix = matrix_inertia
-    # matrix += matrix_transport
-    bsr_axpy(x=matrix_transport, y=matrix)
-    # matrix += nu * B M^-1 B^T
-    bsr_mm(
-        x=bsr_mm(matrix_half_diffusion, inv_vel_mass_matrix),
-        y=bsr_transposed(matrix_half_diffusion),
-        z=matrix,
-        alpha=args.viscosity,
-        beta=1.0,
-    )
-
-    matrix_solve = factorized(bsr_to_scipy(matrix))
-
-    results = [phi0.dof_values.numpy()]
-    phik = phi0
-    for k in range(args.n_frames):
-        # right-hand-side -- standard inertia
-        rhs = integrate(
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+
+    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.ang_vel * res)
+        self.current_frame = 0
+
+        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))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+        scalar_space = fem.make_polynomial_space(geo, degree=0)
+
+        # Inertia matrix
+        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},
+        )
+
+        # Upwind transport term
+        side_test = fem.make_test(space=scalar_space, domain=sides)
+        side_trial = fem.make_trial(space=scalar_space, domain=sides)
+        matrix_transport = fem.integrate(
+            upwind_transport_form,
+            fields={"phi": side_trial, "psi": side_test},
+            values={"ang_vel": args.ang_vel},
+        )
+
+        # Diffusion bilinear form
+        # Since we have piecewise constant element, we cannot use the classical diffusion form
+        # Instead we assemble the matrix B M^-1 B^T, with B associated to the form psi div(u)
+        # and the diagonal matrix M to the velocity mass form u.v
+
+        velocity_space = fem.make_polynomial_space(geo, degree=0, dtype=wp.vec2)
+        side_trial_vel = fem.make_trial(space=velocity_space, domain=sides)
+        matrix_half_diffusion = fem.integrate(
+            half_diffusion_form,
+            fields={"psi": side_test, "u": side_trial_vel},
+        )
+
+        # Diagonal velocity mass matrix
+        test_vel = fem.make_test(space=velocity_space, domain=domain)
+        trial_vel = fem.make_trial(space=velocity_space, domain=domain)
+        inv_vel_mass_matrix = fem.integrate(
+            vel_mass_form, domain=domain, fields={"u": trial_vel, "v": test_vel}, nodal=True
+        )
+        invert_diagonal_bsr_mass_matrix(inv_vel_mass_matrix)
+
+        # Assemble system matrix
+
+        matrix = matrix_inertia
+        # matrix += matrix_transport
+        bsr_axpy(x=matrix_transport, y=matrix)
+        # matrix += nu * B M^-1 B^T
+        bsr_mm(
+            x=bsr_mm(matrix_half_diffusion, inv_vel_mass_matrix),
+            y=bsr_transposed(matrix_half_diffusion),
+            z=matrix,
+            alpha=args.viscosity,
+            beta=1.0,
+        )
+
+        # Compute LU factorization of system matrix
+        self._solve_lu = factorized(bsr_to_scipy(matrix))
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface("phi", self._phi_field)
+
+    def update(self):
+        self.current_frame += 1
+
+        rhs = fem.integrate(
             inertia_form,
-            quadrature=quadrature,
-            fields={"phi": phik, "psi": test},
-            values={"dt": dt},
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"dt": self.sim_dt},
         )
 
-        # Solve using Scipy
-        x = matrix_solve(rhs.numpy().flatten())
+        self._phi_field.dof_values = self._solve_lu(rhs.numpy())
+
+    def render(self):
+        self.renderer.begin_frame(time = self.current_frame * self.sim_dt)
+        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer.end_frame()
+
 
-        phik.dof_values = x
-        results.append(x)
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
 
-    colormesh = plot_surface(phi0)
-    ax = colormesh.axes
+    args = Example.parser.parse_args()
 
-    def animate(i):
-        ax.clear()
-        phik.dof_values = results[i]
-        return plot_surface(phik, axes=ax)
+    example = Example(args=args)
+    for k in range(args.num_frames):
+        print(f"Frame {k}:")
+        example.update()
+        example.render()
 
-    anim = animation.FuncAnimation(
-        ax.figure,
-        animate,
-        interval=30,
-        blit=False,
-        frames=len(results),
-    )
-    plt.show()
+    example.renderer.plot()