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

examples/fem/example_deformed_geometry.py

@@ -0,0 +1,146 @@
+"""
+This example solves a 2d diffusion problem:
+ nu Div u = 1
+with Dirichlet boundary conditions on vertical edges and homogeneous Neumann on horizontal edges.
+"""
+
+import argparse
+
+import warp as wp
+import warp.fem as fem
+
+from warp.sparse import bsr_axpy
+from warp.fem.utils import array_axpy
+
+
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .example_diffusion import linear_form, diffusion_form
+    from .bsr_utils import bsr_cg
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+    from .plot_utils import Plot
+except ImportError:
+    from example_diffusion import linear_form, diffusion_form 
+    from bsr_utils import bsr_cg
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from plot_utils import Plot
+
+@fem.integrand
+def deformation_field_expr(
+    s: fem.Sample,
+    domain: fem.Domain,
+):
+    """
+    Deformation field mapping the unique square to a circular band
+    """
+    x = domain(s)
+
+    r = x[1] + 0.5
+    t = 0.5 * 3.1416 * x[0]
+
+    return r * wp.vec2(wp.sin(t), wp.cos(t)) - x
+
+
+@fem.integrand
+def boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    """
+    Bilinear boundary condition projector form, non-zero on radial edges
+    """
+    nor = fem.normal(domain, s)
+    active = wp.select(nor[0] < -0.9999 or nor[1] < -0.9999, 0.0, 1.0)
+    return active * u(s) * v(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("--serendipity", action="store_true", default=False)
+    parser.add_argument("--viscosity", type=float, default=2.0)
+    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
+
+        # Grid or triangle mesh geometry
+        if args.mesh == "tri":
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
+            base_geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif args.mesh == "quad":
+            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
+            base_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            base_geo = fem.Grid2D(res=wp.vec2i(args.resolution))
+
+        # Construct deformation field on base geometry
+        deformation_space = fem.make_polynomial_space(base_geo, degree=args.degree, dtype=wp.vec2)
+        deformation_field = deformation_space.make_field()
+        fem.interpolate(deformation_field_expr, dest=deformation_field)
+
+        self._geo = deformation_field.make_deformed_geometry()
+
+        # Scalar function space on deformed geometry
+        element_basis = fem.ElementBasis.SERENDIPITY if args.serendipity else None
+        self._scalar_space = fem.make_polynomial_space(self._geo, degree=args.degree, element_basis=element_basis)
+
+        # Scalar field over our function space
+        self._scalar_field = self._scalar_space.make_field()
+
+        self.renderer = Plot(stage)
+
+    def update(self):
+        args = self._args
+        geo = self._geo
+
+        domain = fem.Cells(geometry=geo)
+
+        # Right-hand-side (forcing term)
+        test = fem.make_test(space=self._scalar_space, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Diffusion form
+        trial = fem.make_trial(space=self._scalar_space, domain=domain)
+        matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": args.viscosity})
+
+        # Weakly-imposed boundary conditions on all sides
+        boundary = fem.BoundarySides(geo)
+        bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
+        bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
+
+        bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+
+        fem.project_linear_system(matrix, rhs, bd_matrix)
+
+        # Solve linear system using Conjugate Gradient
+        x = wp.zeros_like(rhs)
+        bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet, tol=1.0e-6)
+
+        # Assign system result to our discrete field
+        self._scalar_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_surface("solution", self._scalar_field)
+
+
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    example.update()
+    example.render()
+
+    example.renderer.plot()

examples/fem/example_diffusion.py

@@ -7,57 +7,57 @@ with Dirichlet boundary conditions on vertical edges and homogeneous Neumann on
 import argparse
 
 import warp as wp
-
-from warp.fem.types import *
-from warp.fem.geometry import Grid2D, Trimesh2D
-from warp.fem.space import make_polynomial_space
-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 grad, normal, integrand
-from warp.fem.dirichlet import project_linear_system
+import warp.fem as fem
 
 from warp.sparse import bsr_axpy
 from warp.fem.utils import array_axpy
 
-from plot_utils import plot_surface
-from bsr_utils import bsr_cg
-from mesh_utils import gen_trimesh
 
-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 .mesh_utils import gen_trimesh, gen_quadmesh
+    from .plot_utils import Plot
+except ImportError:
+    from bsr_utils import bsr_cg
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from plot_utils import Plot
+
+wp.set_module_options({"enable_backward": False})
 
 
-@integrand
+@fem.integrand
 def linear_form(
-    s: Sample,
-    v: Field,
+    s: fem.Sample,
+    v: fem.Field,
 ):
     """Linear form with constant slope 1 -- forcing term of our problem"""
     return v(s)
 
 
-@integrand
-def diffusion_form(s: Sample, u: Field, v: Field, nu: float):
+@fem.integrand
+def diffusion_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
     """Diffusion bilinear form with constant coefficient ``nu``"""
     return nu * wp.dot(
-        grad(u, s),
-        grad(v, s),
+        fem.grad(u, s),
+        fem.grad(v, s),
     )
 
 
-@integrand
-def y_boundary_value_form(s: Sample, domain: Domain, v: Field, val: float):
+@fem.integrand
+def y_boundary_value_form(s: fem.Sample, domain: fem.Domain, v: fem.Field, val: float):
     """Linear form with coefficient val on vertical edges, zero elsewhere"""
-    nor = normal(domain, s)
+    nor = fem.normal(domain, s)
     return val * v(s) * wp.abs(nor[0])
 
 
-@integrand
+@fem.integrand
 def y_boundary_projector_form(
-    s: Sample,
-    domain: Domain,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
 ):
     """
     Bilinear boundary condition projector form, non-zero on vertical edges only.
@@ -66,65 +66,96 @@ def y_boundary_projector_form(
     return y_boundary_value_form(s, domain, v, 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=50)
     parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--serendipity", action="store_true", default=False)
     parser.add_argument("--viscosity", type=float, default=2.0)
     parser.add_argument("--boundary_value", type=float, default=5.0)
     parser.add_argument("--boundary_compliance", type=float, default=0, help="Dirichlet boundary condition compliance")
-    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh")
-    args = parser.parse_args()
-
-    # Grid or triangle mesh geometry
-    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 and function spaces
-    domain = Cells(geometry=geo)
-    scalar_space = make_polynomial_space(geo, degree=args.degree)
-
-    # Right-hand-side (forcing term)
-    test = make_test(space=scalar_space, domain=domain)
-    rhs = integrate(linear_form, fields={"v": test})
-
-    # Diffusion form
-    trial = make_trial(space=scalar_space, domain=domain)
-    matrix = integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": args.viscosity})
-
-    # Boundary conditions on Y sides
-    # Use nodal integration so that boundary conditions are specified on each node independently
-    boundary = BoundarySides(geo)
-    bd_test = make_test(space=scalar_space, domain=boundary)
-    bd_trial = make_trial(space=scalar_space, domain=boundary)
-
-    bd_matrix = integrate(y_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
-    bd_rhs = integrate(y_boundary_value_form, fields={"v": bd_test}, values={"val": args.boundary_value}, nodal=True)
-
-    # Assemble linear system
-    if args.boundary_compliance == 0.0:
-        # Hard BC: project linear system
-        project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
-    else:
-        # Weak BC: add toegether diffusion and boundary condition matrices
-        boundary_strength = 1.0 / args.boundary_compliance
-        bsr_axpy(x=bd_matrix, y=matrix, alpha=boundary_strength, beta=1)
-        array_axpy(x=bd_rhs, y=rhs, alpha=boundary_strength, beta=1)
-
-    # Solve linear system using Conjugate Gradient
-    x = wp.zeros_like(rhs)
-    bsr_cg(matrix, b=rhs, x=x)
-
-    # Assign system result to a discrete field,
-    scalar_field = scalar_space.make_field()
-    scalar_field.dof_values = x
-
-    # Visualize it with matplotlib
-    plot_surface(scalar_field)
-    plt.show()
+    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
+
+        # 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))
+
+        # Scalar function space
+        element_basis = fem.ElementBasis.SERENDIPITY if args.serendipity else None
+        self._scalar_space = fem.make_polynomial_space(self._geo, degree=args.degree, element_basis=element_basis)
+
+        # Scalar field over our function space
+        self._scalar_field = self._scalar_space.make_field()
+
+        self.renderer = Plot(stage)
+
+    def update(self):
+        args = self._args
+        geo = self._geo
+
+        domain = fem.Cells(geometry=geo)
+
+        # Right-hand-side (forcing term)
+        test = fem.make_test(space=self._scalar_space, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Diffusion form
+        trial = fem.make_trial(space=self._scalar_space, domain=domain)
+        matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": args.viscosity})
+
+        # Boundary conditions on Y sides
+        # Use nodal integration so that boundary conditions are specified on each node independently
+        boundary = fem.BoundarySides(geo)
+        bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
+        bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
+
+        bd_matrix = fem.integrate(y_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+        bd_rhs = fem.integrate(
+            y_boundary_value_form, fields={"v": bd_test}, values={"val": args.boundary_value}, nodal=True
+        )
+
+        # Assemble linear system
+        if args.boundary_compliance == 0.0:
+            # Hard BC: project linear system
+            fem.project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
+        else:
+            # Weak BC: add toegether diffusion and boundary condition matrices
+            boundary_strength = 1.0 / args.boundary_compliance
+            bsr_axpy(x=bd_matrix, y=matrix, alpha=boundary_strength, beta=1)
+            array_axpy(x=bd_rhs, y=rhs, alpha=boundary_strength, beta=1)
+
+        # Solve linear system using Conjugate Gradient
+        x = wp.zeros_like(rhs)
+        bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
+
+        # Assign system result to our discrete field
+        self._scalar_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_surface("solution", self._scalar_field)
+
+
+if __name__ == "__main__":
+    wp.init()
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    example.update()
+    example.render()
+
+    example.renderer.plot()

examples/fem/example_diffusion_3d.py

@@ -7,103 +7,133 @@ with homogeneous Neumann conditions on horizontal sides and homogeneous Dirichle
 import argparse
 
 import warp as wp
-import numpy as np
-
-from warp.fem.types import *
-from warp.fem.geometry import Grid3D, Tetmesh
-from warp.fem.space import make_polynomial_space
-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 normal, integrand
-from warp.fem.dirichlet import project_linear_system
+import warp.fem as fem
 
 from warp.sparse import bsr_axpy
 
-from plot_utils import plot_3d_scatter
-from bsr_utils import bsr_cg
-from mesh_utils import gen_tetmesh
-
-from example_diffusion import diffusion_form, linear_form
-
-import matplotlib.pyplot as plt
-
-
-@integrand
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .example_diffusion import diffusion_form, linear_form
+    from .bsr_utils import bsr_cg
+    from .mesh_utils import gen_tetmesh
+    from .plot_utils import Plot
+except ImportError:
+    from example_diffusion import diffusion_form, linear_form
+    from bsr_utils import bsr_cg
+    from mesh_utils import gen_tetmesh, gen_hexmesh
+    from plot_utils import Plot
+
+
+@fem.integrand
 def vert_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 mass on vertical sides only
-    w = 1.0 - wp.abs(normal(domain, s)[1])
+    w = 1.0 - wp.abs(fem.normal(domain, s)[1])
     return w * u(s) * v(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=10)
     parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--serendipity", action="store_true", default=False)
     parser.add_argument("--viscosity", type=float, default=2.0)
     parser.add_argument("--boundary_compliance", type=float, default=0, help="Dirichlet boundary condition compliance")
-    parser.add_argument("--tet_mesh", action="store_true", help="Use a tetrahedral mesh")
-    args = parser.parse_args()
-
-    res = vec3i(args.resolution, args.resolution // 2, args.resolution * 2)
-
-    if args.tet_mesh:
-        pos, tet_vtx_indices = gen_tetmesh(
-            res=res,
-            bounds_lo=wp.vec3(0.0, 0.0, 0.0),
-            bounds_hi=wp.vec3(1.0, 0.5, 2.0),
-        )
-        geo = Tetmesh(tet_vtx_indices, pos)
-    else:
-        geo = Grid3D(
-            res=res,
-            bounds_lo=wp.vec3(0.0, 0.0, 0.0),
-            bounds_hi=wp.vec3(1.0, 0.5, 2.0),
-        )
-
-    # Domain and function spaces
-    domain = Cells(geometry=geo)
-    scalar_space = make_polynomial_space(geo, degree=args.degree)
-
-    # Right-hand-side
-    test = make_test(space=scalar_space, domain=domain)
-    rhs = integrate(linear_form, fields={"v": test})
-
-    # Weakly-imposed boundary conditions on Y sides
-    with wp.ScopedTimer("Integrate"):
-        boundary = BoundarySides(geo)
-
-        bd_test = make_test(space=scalar_space, domain=boundary)
-        bd_trial = make_trial(space=scalar_space, domain=boundary)
-        bd_matrix = integrate(vert_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
-
-        # Diffusion form
-        trial = make_trial(space=scalar_space, domain=domain)
-        matrix = integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": args.viscosity})
-
-    if args.boundary_compliance == 0.0:
-        # Hard BC: project linear system
-        bd_rhs = wp.zeros_like(rhs)
-        project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
-    else:
-        # Weak BC: add toegether diffusion and boundary condition matrices
-        boundary_strength = 1.0 / args.boundary_compliance
-        bsr_axpy(x=bd_matrix, y=matrix, alpha=100.0, beta=1)
-
-    with wp.ScopedTimer("CG solve"):
-        x = wp.zeros_like(rhs)
-        bsr_cg(matrix, b=rhs, x=x)
-
-    scalar_field = scalar_space.make_field()
-    scalar_field.dof_values = x
-    plot_3d_scatter(scalar_field)
-
-    plt.show()
+    parser.add_argument("--mesh", choices=("grid", "tet", "hex"), 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 = wp.vec3i(args.resolution, args.resolution // 2, args.resolution * 2)
+
+        if args.mesh == "tet":
+            pos, tet_vtx_indices = gen_tetmesh(
+                res=res,
+                bounds_lo=wp.vec3(0.0, 0.0, 0.0),
+                bounds_hi=wp.vec3(1.0, 0.5, 2.0),
+            )
+            self._geo = fem.Tetmesh(tet_vtx_indices, pos)
+        elif args.mesh == "hex":
+            pos, hex_vtx_indices = gen_hexmesh(
+                res=res,
+                bounds_lo=wp.vec3(0.0, 0.0, 0.0),
+                bounds_hi=wp.vec3(1.0, 0.5, 2.0),
+            )
+            self._geo = fem.Hexmesh(hex_vtx_indices, pos)
+        else:
+            self._geo = fem.Grid3D(
+                res=res,
+                bounds_lo=wp.vec3(0.0, 0.0, 0.0),
+                bounds_hi=wp.vec3(1.0, 0.5, 2.0),
+            )
+
+        # Domain and function spaces
+        element_basis = fem.ElementBasis.SERENDIPITY if args.serendipity else None
+        self._scalar_space = fem.make_polynomial_space(self._geo, degree=args.degree, element_basis=element_basis)
+
+        # Scalar field over our function space
+        self._scalar_field: fem.DiscreteField = self._scalar_space.make_field()
+
+        self.renderer = Plot(stage)
+
+    def update(self):
+        args = self._args
+        geo = self._geo
+
+        domain = fem.Cells(geometry=geo)
+
+        # Right-hand-side
+        test = fem.make_test(space=self._scalar_space, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Weakly-imposed boundary conditions on Y sides
+        with wp.ScopedTimer("Integrate"):
+            boundary = fem.BoundarySides(geo)
+
+            bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
+            bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
+            bd_matrix = fem.integrate(vert_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+
+            # Diffusion form
+            trial = fem.make_trial(space=self._scalar_space, domain=domain)
+            matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": args.viscosity})
+
+        if args.boundary_compliance == 0.0:
+            # Hard BC: project linear system
+            bd_rhs = wp.zeros_like(rhs)
+            fem.project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
+        else:
+            # Weak BC: add toegether diffusion and boundary condition matrices
+            boundary_strength = 1.0 / args.boundary_compliance
+            bsr_axpy(x=bd_matrix, y=matrix, alpha=boundary_strength, beta=1)
+
+        with wp.ScopedTimer("CG solve"):
+            x = wp.zeros_like(rhs)
+            bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
+            self._scalar_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_volume("solution", self._scalar_field)
+
+
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    example.update()
+    example.render()
+
+    example.renderer.plot()