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

examples/fem/example_diffusion_mgpu.py

@@ -5,30 +5,27 @@ This example illustrates using domain decomposition to solve a diffusion PDE ove
 from typing import Tuple
 
 import warp as wp
-
-from warp.fem.types import *
-from warp.fem.geometry import Grid2D, LinearGeometryPartition
-from warp.fem.space import make_polynomial_space, make_space_partition
-from warp.fem.field import make_test, make_trial
-from warp.fem.domain import Cells, BoundarySides
-from warp.fem.integrate import integrate
-from warp.fem.operator import integrand
-
+import warp.fem as fem
 from warp.sparse import bsr_axpy, bsr_mv
 from warp.utils import array_cast
 
-from example_diffusion import linear_form, diffusion_form
-from bsr_utils import bsr_cg
-from plot_utils import plot_grid_surface
-
-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 .example_diffusion import diffusion_form, linear_form
+    from .plot_utils import Plot
+except ImportError:
+    from bsr_utils import bsr_cg
+    from example_diffusion import diffusion_form, linear_form
+    from plot_utils import Plot
 
 
-@integrand
+@fem.integrand
 def mass_form(
-    s: Sample,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
 ):
     return u(s) * v(s)
 
@@ -96,77 +93,103 @@ def mv_routine(A: DistributedSystem, x: wp.array, y: wp.array, alpha=1.0, beta=0
         wp.launch(kernel=sum_kernel, dim=idx.shape, device=y_idx.device, inputs=[y_idx, tmp_i], stream=stream)
 
 
-if __name__ == "__main__":
-    wp.init()
-    wp.set_module_options({"enable_backward": False})
-
-    geo = Grid2D(res=vec2i(25))
-    bd_weight = 100.0
-
-    scalar_space = make_polynomial_space(geo, degree=3)
-
-    devices = wp.get_cuda_devices()
-
-    rhs_vecs = []
-    res_vecs = []
-    matrices = []
-    indices = []
-
-    main_device = devices[0]
-
-    # Build local system for each device
-    for k, device in enumerate(devices):
-        with wp.ScopedDevice(device):
-            # Construct the partition corresponding to the k'th device
-            geo_partition = LinearGeometryPartition(geo, k, len(devices))
-            space_partition = make_space_partition(scalar_space, geo_partition)
-
-            domain = Cells(geometry=geo_partition)
+class Example:
+    def __init__(self, stage=None, quiet=False):
+        self._bd_weight = 100.0
+        self._quiet = quiet
+
+        self._geo = fem.Grid2D(res=wp.vec2i(25))
+
+        self._main_device = wp.get_device("cuda")
+
+        with wp.ScopedDevice(self._main_device):
+            self._scalar_space = fem.make_polynomial_space(self._geo, degree=3)
+            self._scalar_field = self._scalar_space.make_field()
+
+        self.renderer = Plot(stage)
+
+    def update(self):
+        devices = wp.get_cuda_devices()
+        main_device = self._main_device
+
+        rhs_vecs = []
+        res_vecs = []
+        matrices = []
+        indices = []
+
+        # Build local system for each device
+        for k, device in enumerate(devices):
+            with wp.ScopedDevice(device):
+                # Construct the partition corresponding to the k'th device
+                geo_partition = fem.LinearGeometryPartition(self._geo, k, len(devices))
+                matrix, rhs, partition_node_indices = self._assemble_local_system(geo_partition)
+
+                rhs_vecs.append(rhs)
+                res_vecs.append(wp.empty_like(rhs))
+                matrices.append(matrix)
+                indices.append(partition_node_indices.to(main_device))
+
+        # Global rhs as sum of all local rhs
+        glob_rhs = wp.zeros(n=self._scalar_space.node_count(), dtype=wp.float64, device=main_device)
+        tmp = wp.empty_like(glob_rhs)
+        sum_vecs(rhs_vecs, indices, glob_rhs, tmp)
+
+        # Distributed CG
+        global_res = wp.zeros_like(glob_rhs)
+        A = DistributedSystem()
+        A.device = device
+        A.scalar_type = glob_rhs.dtype
+        A.nrow = self._scalar_space.node_count()
+        A.shape = (A.nrow, A.nrow)
+        A.tmp_buf = tmp
+        A.rank_data = (matrices, rhs_vecs, res_vecs, indices)
+
+        bsr_cg(
+            A,
+            x=global_res,
+            b=glob_rhs,
+            use_diag_precond=False,
+            quiet=self._quiet,
+            mv_routine=mv_routine,
+            device=main_device,
+        )
 
-            # Right-hand-side
-            test = make_test(space_partition=space_partition, domain=domain)
-            rhs = integrate(linear_form, fields={"v": test})
+        array_cast(in_array=global_res, out_array=self._scalar_field.dof_values)
 
-            # Weakly-imposed boundary conditions on all sides
-            boundary = BoundarySides(geometry=geo_partition)
-            bd_test = make_test(space_partition=space_partition, domain=boundary)
-            bd_trial = make_trial(space_partition=space_partition, domain=boundary)
-            bd_matrix = integrate(mass_form, fields={"u": bd_trial, "v": bd_test})
+    def render(self):
+        self.renderer.add_surface("solution", self._scalar_field)
 
-            # Diffusion form
-            trial = make_trial(space_partition=space_partition, domain=domain)
-            matrix = integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": 1.0})
+    def _assemble_local_system(self, geo_partition: fem.GeometryPartition):
+        scalar_space = self._scalar_space
+        space_partition = fem.make_space_partition(scalar_space, geo_partition)
 
-            bsr_axpy(y=matrix, x=bd_matrix, alpha=bd_weight)
+        domain = fem.Cells(geometry=geo_partition)
 
-            rhs_vecs.append(rhs)
-            res_vecs.append(wp.empty_like(rhs))
-            matrices.append(matrix)
-            indices.append(space_partition.space_node_indices().to(main_device))
+        # Right-hand-side
+        test = fem.make_test(space=scalar_space, space_partition=space_partition, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
 
+        # Weakly-imposed boundary conditions on all sides
+        boundary = fem.BoundarySides(geometry=geo_partition)
+        bd_test = fem.make_test(space=scalar_space, space_partition=space_partition, domain=boundary)
+        bd_trial = fem.make_trial(space=scalar_space, space_partition=space_partition, domain=boundary)
+        bd_matrix = fem.integrate(mass_form, fields={"u": bd_trial, "v": bd_test})
 
-    # Global rhs as sum of all local rhs
-    glob_rhs = wp.zeros(n=scalar_space.node_count(), dtype=wp.float64, device=main_device)
-    tmp = wp.empty_like(glob_rhs)
-    sum_vecs(rhs_vecs, indices, glob_rhs, tmp)
+        # Diffusion form
+        trial = fem.make_trial(space=scalar_space, space_partition=space_partition, domain=domain)
+        matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": 1.0})
 
-    # Distributed CG
-    glob_res = wp.zeros_like(glob_rhs)
-    A = DistributedSystem()
-    A.device = device
-    A.scalar_type = glob_rhs.dtype
-    A.nrow = scalar_space.node_count()
-    A.shape = (A.nrow, A.nrow)
-    A.tmp_buf = tmp
-    A.rank_data = (matrices, rhs_vecs, res_vecs, indices)
+        bsr_axpy(y=matrix, x=bd_matrix, alpha=self._bd_weight)
 
-    with wp.ScopedDevice(main_device):
-        bsr_cg(A, x=glob_res, b=glob_rhs, use_diag_precond=False, mv_routine=mv_routine, device=main_device)
+        return matrix, rhs, space_partition.space_node_indices()
 
-    scalar_field = scalar_space.make_field()
 
-    array_cast(in_array=glob_res, out_array=scalar_field.dof_values)
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
 
-    plot_grid_surface(scalar_field)
+    example = Example()
+    example.update()
+    example.render()
 
-    plt.show()
+    example.renderer.plot()

examples/fem/example_mixed_elasticity.py

@@ -5,25 +5,20 @@ with Dirichlet boundary conditions on horizontal sides, and E the elasticity ran
 import argparse
 
 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
-from warp.fem.space import make_polynomial_space, SymmetricTensorMapper
-from warp.fem.domain import Cells, BoundarySides
-from warp.fem.integrate import integrate
-from warp.fem.operator import normal, integrand, D
-from warp.fem.dirichlet import project_linear_system
+import warp.fem as fem
 
 from warp.sparse import bsr_transposed, bsr_mm
 
-from plot_utils import plot_velocities 
-from bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
-from mesh_utils import gen_trimesh
-
-import matplotlib.pyplot as plt
+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.func
@@ -34,134 +29,178 @@ def compute_stress(tau: wp.mat22, E: wp.mat33):
     return wp.mat22(sig_sym[0], 0.5 * sig_sym[2], 0.5 * sig_sym[2], sig_sym[1])
 
 
-@integrand
+@fem.integrand
 def symmetric_grad_form(
-    s: Sample,
-    u: Field,
-    tau: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    tau: fem.Field,
 ):
     """D(u) : tau"""
-    return wp.ddot(tau(s), D(u, s))
+    return wp.ddot(tau(s), fem.D(u, s))
 
 
-@integrand
-def stress_form(s: Sample, u: Field, tau: Field, E: wp.mat33):
+@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(D(u, s), E))
+    return wp.ddot(tau(s), compute_stress(fem.D(u, s), E))
 
 
-@integrand
+@fem.integrand
 def horizontal_boundary_projector_form(
-    s: Sample,
-    domain: Domain,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
 ):
     # non zero on horizontal boundary of domain only
-    nor = normal(domain, s)
+    nor = fem.normal(domain, s)
     return wp.dot(u(s), v(s)) * wp.abs(nor[1])
 
 
-@integrand
+@fem.integrand
 def horizontal_displacement_form(
-    s: Sample,
-    domain: Domain,
-    v: Field,
+    s: fem.Sample,
+    domain: fem.Domain,
+    v: fem.Field,
     displacement: float,
 ):
     # opposed to normal on horizontal boundary of domain only
-    nor = normal(domain, s)
+    nor = fem.normal(domain, s)
     return -wp.abs(nor[1]) * displacement * wp.dot(nor, v(s))
 
 
-@integrand
+@fem.integrand
 def tensor_mass_form(
-    s: Sample,
-    sig: Field,
-    tau: Field,
+    s: fem.Sample,
+    sig: fem.Field,
+    tau: fem.Field,
 ):
     return wp.ddot(tau(s), sig(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("--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("--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)
-
-    # Strain-stress matrix
-    young = args.young_modulus
-    poisson = args.poisson_ratio
-    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)]]
-        )
+    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)"
     )
 
-    domain = Cells(geometry=geo)
+    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 -- Q_k for displacement, Q_{k-1}d for stress
-    u_space = make_polynomial_space(geo, degree=args.degree, dtype=wp.vec2)
-    # Store stress degrees of freedom as symmetric tensors (3 dof) rather than full 2x2 matrices
-    tau_space = make_polynomial_space(
-        geo, degree=args.degree - 1, discontinuous=True, dof_mapper=SymmetricTensorMapper(wp.mat22)
-    )
+        # 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
+        )
 
-    # Displacement boundary conditions
-    u_bd_test = make_test(space=u_space, domain=boundary)
-    u_bd_trial = make_trial(space=u_space, domain=boundary)
-    u_bd_rhs = integrate(
-        horizontal_displacement_form,
-        fields={"v": u_bd_test},
-        values={"displacement": args.displacement},
-        nodal=True,
-        output_dtype=wp.vec2d,
-    )
-    u_bd_matrix = integrate(horizontal_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True)
+        # 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),
+        )
 
-    # Stress/velocity coupling
-    u_trial = make_trial(space=u_space, domain=domain)
-    tau_test = make_test(space=tau_space, domain=domain)
-    tau_trial = make_trial(space=tau_space, domain=domain)
+        self._u_field = self._u_space.make_field()
 
-    sym_grad_matrix = integrate(symmetric_grad_form, fields={"u": u_trial, "tau": tau_test})
-    stress_matrix = integrate(stress_form, fields={"u": u_trial, "tau": tau_test}, values={"E": elasticity_mat})
+        self.renderer = Plot(stage)
 
-    # Compute inverse of the (block-diagonal) tau mass matrix
-    tau_inv_mass_matrix = integrate(tensor_mass_form, fields={"sig": tau_trial, "tau": tau_test}, nodal=True)
-    invert_diagonal_bsr_mass_matrix(tau_inv_mass_matrix)
+    def update(self):
+        boundary = fem.BoundarySides(self._geo)
+        domain = fem.Cells(geometry=self._geo)
 
-    # Assemble system matrix
-    u_matrix = bsr_mm(bsr_transposed(sym_grad_matrix), bsr_mm(tau_inv_mass_matrix, stress_matrix))
+        # 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)
 
-    # Enforce boundary conditions
-    u_rhs = wp.zeros_like(u_bd_rhs)
-    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)
 
-    x = wp.zeros_like(u_rhs)
-    bsr_cg(u_matrix, b=u_rhs, x=x, tol=1.0e-16)
+        # Extract result
+        self._u_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_surface_vector("solution", self._u_field)
+
+
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
 
-    # Extract result
-    u_field = u_space.make_field()
-    u_field.dof_values = x  # .reshape((-1, 2))
+    args = Example.parser.parse_args()
 
-    plot_velocities(u_field)
+    example = Example(args=args)
+    example.update()
+    example.render()
 
-    plt.show()
+    example.renderer.plot()