warp-lang 1.4.2__py3-none-macosx_10_13_universal2.whl → 1.5.1__py3-none-macosx_10_13_universal2.whl

warp/examples/fem/example_convection_diffusion_dg.py

@@ -88,7 +88,7 @@ class Example:
             geo,
             discontinuous=True,
             degree=degree,
-            family=fem.Polynomial.GAUSS_LEGENDRE,
+            family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE,
         )
 
         # Assemble transport, diffusion and inertia matrices
@@ -110,23 +110,22 @@ class Example:
 
         side_test = fem.make_test(space=scalar_space, domain=sides)
         side_trial = fem.make_trial(space=scalar_space, domain=sides)
-
-        matrix_transport += fem.integrate(
+        fem.integrate(
             upwind_transport_form,
             fields={"phi": side_trial, "psi": side_test},
             values={"ang_vel": ang_vel},
+            output=matrix_transport,
+            add=True,
         )
 
         matrix_diffusion = fem.integrate(
             diffusion_form,
             fields={"u": trial, "v": self._test},
         )
-
         matrix_diffusion += fem.integrate(
             sip_diffusion_form,
             fields={"phi": side_trial, "psi": side_test},
         )
-
         self._matrix = matrix_inertia + matrix_transport + viscosity * matrix_diffusion
 
         # Initial condition

warp/examples/fem/example_deformed_geometry.py

@@ -28,7 +28,7 @@ def deformation_field_expr(
     domain: fem.Domain,
 ):
     """
-    Deformation field mapping the unique square to a circular band
+    Deformation field mapping the unit square to a circular band
     """
     x = domain(s)
 

warp/examples/fem/example_diffusion_3d.py

@@ -16,6 +16,8 @@
 # and homogeneous Dirichlet boundary conditions other sides.
 ###########################################################################
 
+import numpy as np
+
 import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
@@ -24,7 +26,7 @@ from warp.sparse import bsr_axpy
 
 
 @fem.integrand
-def vert_boundary_projector_form(
+def vertical_boundary_projector_form(
     s: fem.Sample,
     domain: fem.Domain,
     u: fem.Field,
@@ -36,6 +38,18 @@ def vert_boundary_projector_form(
     return w * u(s) * v(s)
 
 
+@fem.integrand
+def y_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # Constrain Y edges
+    tangent = fem.deformation_gradient(domain, s)
+    return wp.abs(tangent[1]) * u(s) * v(s)
+
+
 class Example:
     def __init__(
         self,
@@ -77,6 +91,28 @@ class Example:
                 bounds_hi=bounds_hi,
             )
             self._geo = fem.Nanogrid(volume)
+        elif mesh == "tri":
+            pos, quad_vtx_indices = fem_example_utils.gen_trimesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            pos = pos.numpy()
+            pos_z = np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            pos = wp.array(pos, dtype=wp.vec3)
+            self._geo = fem.Trimesh3D(quad_vtx_indices, pos)
+        elif mesh == "quad":
+            pos, quad_vtx_indices = fem_example_utils.gen_quadmesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            pos = pos.numpy()
+            pos_z = np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            pos = wp.array(pos, dtype=wp.vec3)
+            self._geo = fem.Quadmesh3D(quad_vtx_indices, pos)
         else:
             self._geo = fem.Grid3D(
                 res=res,
@@ -108,7 +144,12 @@ class Example:
 
             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)
+
+            # Pick boundary conditions depending on whether our geometry is a 3d surface or a volume
+            boundary_projector_form = (
+                vertical_boundary_projector_form if self._geo.cell_dimension == 3 else y_boundary_projector_form
+            )
+            bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
 
             # Diffusion form
             trial = fem.make_trial(space=self._scalar_space, domain=domain)
@@ -146,7 +187,9 @@ if __name__ == "__main__":
     parser.add_argument(
         "--boundary_compliance", type=float, default=0.0, help="Dirichlet boundary condition compliance."
     )
-    parser.add_argument("--mesh", choices=("grid", "tet", "hex", "nano", "anano"), default="grid", help="Mesh type.")
+    parser.add_argument(
+        "--mesh", choices=("grid", "tet", "hex", "nano", "anano", "tri", "quad"), default="grid", help="Mesh type."
+    )
     parser.add_argument(
         "--headless",
         action="store_true",
@@ -171,4 +214,4 @@ if __name__ == "__main__":
         example.render()
 
         if not args.headless:
-            example.renderer.plot(backend="matplotlib")
+            example.renderer.plot()

warp/examples/fem/example_distortion_energy.py

@@ -0,0 +1,220 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Distortion Energy
+#
+# This example illustrates using a Newton loop to minimize distortion of a
+# 3D surface (u,v) parameterization under a Symmetric Dirichlet energy,
+#
+# E(F) = 1/2 |F|^2 + |F^{-1}|^2
+#
+# with F := dx/du
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+
+@fem.integrand
+def distortion_gradient_form(s: fem.Sample, u_cur: fem.Field, v: fem.Field):
+    # Symmetric Dirichlet energy gradient (linear form)
+    # E = 1/2 (F:F + F^-T:F^-T)
+
+    F = fem.grad(u_cur, s)
+
+    F_inv_sq = wp.inverse(F * wp.transpose(F))
+    F_inv = F_inv_sq * F
+
+    dE_dF = F - F_inv_sq * F_inv
+
+    return wp.ddot(fem.grad(v, s), dE_dF)
+
+
+@fem.integrand
+def distortion_hessian_form(s: fem.Sample, u_cur: fem.Field, u: fem.Field, v: fem.Field):
+    # Symmetric Dirichlet energy approximate hessian (bilinear form)
+
+    # F:F term
+    H = wp.ddot(fem.grad(v, s), fem.grad(u, s))
+
+    # F^-T:F^-T term
+    F = fem.grad(u_cur, s)
+    F_inv_sq = wp.inverse(F * wp.transpose(F))
+
+    # Gauss--Newton (ignore F^-2 derivative)
+    H += wp.ddot(F_inv_sq * fem.grad(v, s), F_inv_sq * F_inv_sq * fem.grad(u, s))
+
+    return H
+
+
+@fem.integrand
+def initial_guess(
+    s: fem.Sample,
+    domain: fem.Domain,
+):
+    # initialization for UV parameter
+    x = domain(s)
+    return wp.vec2(x[0], x[1])
+
+
+@fem.integrand
+def boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # Fix a single point
+    # (underconstrained, solution up to a rotation in UV space)
+    w = wp.select(s.qp_index == 0, 0.0, 1.0)
+    return w * wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def checkerboard(s: fem.Sample, domain: fem.Domain, u: fem.Field):
+    # checkerboard pattern for parameter visualization
+    u_s = u(s)
+    return wp.sign(wp.cos(16.0 * u_s[0]) * wp.sin(16.0 * u_s[1]))
+
+
+class Example:
+    def __init__(
+        self,
+        quiet=False,
+        degree=2,
+        resolution=25,
+        mesh="grid",
+        nonconforming_stresses=False,
+    ):
+        self._quiet = quiet
+
+        def deform_along_z(positions, z_scale=1.0):
+            pos = positions.numpy()
+            pos_z = z_scale * np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            return wp.array(pos, dtype=wp.vec3)
+
+        # Grid or mesh geometry
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
+            self._uv_geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=wp.zeros_like(positions))
+
+            positions = deform_along_z(positions)
+            self._geo = fem.Trimesh3D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            self._uv_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=wp.zeros_like(positions))
+
+            positions = deform_along_z(positions)
+            self._geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            self._uv_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=wp.zeros_like(positions))
+
+            undef_positions = deform_along_z(positions, z_scale=0.0)
+            flat_geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=undef_positions)
+
+            deformation_field = fem.make_discrete_field(fem.make_polynomial_space(flat_geo, dtype=wp.vec3))
+            deformation_field.dof_values = deform_along_z(positions)
+
+            self._geo = deformation_field.make_deformed_geometry(relative=False)
+
+        # parameter space
+        self._u_space = fem.make_polynomial_space(self._geo, degree=degree, dtype=wp.vec2)
+        self._u_field = self._u_space.make_field()
+        self._du_field = self._u_space.make_field()
+        fem.interpolate(initial_guess, dest=self._u_field)
+
+        # scalar parameter visualization function
+        viz_space = fem.make_polynomial_space(self._geo, degree=3, dtype=float)
+        self.viz_field = viz_space.make_field()
+        # For visualization of uv in 2D space
+        uv_space = fem.make_polynomial_space(self._uv_geo, degree=degree, dtype=wp.vec2)
+        self._uv_field = uv_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        boundary = fem.BoundarySides(self._geo)
+        domain = fem.Cells(geometry=self._geo)
+
+        # Parameter 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_matrix = fem.integrate(
+            boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True, output_dtype=float
+        )
+        fem.normalize_dirichlet_projector(u_bd_matrix)
+
+        u_test = fem.make_test(space=self._u_space, domain=domain)
+        u_trial = fem.make_trial(space=self._u_space, domain=domain)
+
+        # Newton iterations (without line-search for simplicity)
+        for _newton_iteration in range(10):
+            u_matrix = fem.integrate(
+                distortion_hessian_form, fields={"u_cur": self._u_field, "u": u_trial, "v": u_test}, output_dtype=float
+            )
+
+            u_rhs = fem.integrate(
+                distortion_gradient_form, fields={"u_cur": self._u_field, "v": u_test}, output_dtype=wp.vec2
+            )
+
+            fem.project_linear_system(u_matrix, u_rhs, u_bd_matrix, normalize_projector=False)
+
+            # Solve for uv increment
+            du = self._du_field.dof_values
+            du.zero_()
+            fem_example_utils.bsr_cg(u_matrix, b=u_rhs, x=du, quiet=self._quiet)
+
+            # Accumulate to UV field
+            fem.utils.array_axpy(x=du, y=self._u_field.dof_values, alpha=-1.0, beta=1.0)
+
+    def render(self):
+        # Visualization
+        fem.interpolate(checkerboard, fields={"u": self._u_field}, dest=self.viz_field)
+
+        self._uv_field.dof_values = wp.clone(self._u_field.dof_values)
+
+        self.renderer.add_field("pattern", self.viz_field)
+        self.renderer.add_field("uv", self._uv_field)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    wp.set_module_options({"enable_backward": False})
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--resolution", type=int, default=25, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=1, help="Polynomial degree of shape functions.")
+    parser.add_argument("--mesh", choices=("tri", "quad", "deformed"), default="tri", help="Mesh type")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+    parser.add_argument("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            degree=args.degree,
+            resolution=args.resolution,
+            mesh=args.mesh,
+        )
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot(options={"uv": {"displacement": {}}})

warp/examples/fem/example_magnetostatics.py

@@ -8,17 +8,18 @@
 ###########################################################################
 # Example Magnetostatics
 #
-# This example demonstrates solving an in-plane magnetostatics problem
-# using a curl-curl formulation
+# This example demonstrates solving a 3d magnetostatics problem
+# (a copper coil with radial current around a cylindrical iron core)
+# using a curl-curl formulation and H(curl)-conforming function space
 #
 # 1/mu Curl B + j = 0
 # Div. B = 0
 #
-# solved over field A such that B = Curl A, A = (0, 0, a_z),
-# and a_z = 0 on the domain boundary
+# solved over field A such that B = Curl A,
+# and Direchlet homogeneous essential boundary conditions
 #
-# This example also illustrates using an ImplictField to warp a square mesh
-# to a circular domain
+# This example also illustrates using an ImplictField to warp a grid mesh
+# to a cylindrical domain
 ###########################################################################
 
 import numpy as np
@@ -27,138 +28,179 @@ import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
 
-# Vacuum and copper magnetic permeabilities
-MU_0 = wp.constant(np.pi * 4.0e-7)
-MU_C = wp.constant(1.25e-6)
+# Physics constants
+MU_0 = wp.constant(np.pi * 4.0e-7)  # Vacuum magnetic permeability
+MU_c = wp.constant(1.25e-6)  # Copper magnetic permeability
+MU_i = wp.constant(6.0e-3)  # Iron magnetic permeability
 
 
 @wp.func
-def square_to_disk(x: wp.vec2):
+def cube_to_cylinder(x: wp.vec3):
     # mapping from unit square to unit disk
-    return wp.normalize(x) * wp.max(wp.abs(x))
+    pos_xz = wp.vec3(x[0], 0.0, x[2])
+    return wp.max(wp.abs(pos_xz)) * wp.normalize(pos_xz) + wp.vec3(0.0, x[1], 0.0)
 
 
 @wp.func
-def square_to_disk_grad(x: wp.vec2):
+def cube_to_cylinder_grad(x: wp.vec3):
     # gradient of mapping from unit square to unit disk
-    if x == wp.vec2(0.0):
-        return wp.mat22(0.0)
+    pos_xz = wp.vec3(x[0], 0.0, x[2])
+    if pos_xz == wp.vec3(0.0):
+        grad = wp.mat33(0.0)
+    else:
+        dir_xz = wp.normalize(pos_xz)
+        dir_grad = (wp.identity(n=3, dtype=float) - wp.outer(dir_xz, dir_xz)) / wp.length(pos_xz)
+
+        abs_xz = wp.abs(pos_xz)
+        xinf_grad = wp.select(
+            abs_xz[0] > abs_xz[2], wp.vec3(0.0, 0.0, wp.sign(pos_xz[2])), wp.vec(wp.sign(pos_xz[0]), 0.0, 0.0)
+        )
+        grad = dir_grad * wp.max(abs_xz) + wp.outer(dir_xz, xinf_grad)
 
-    d = wp.normalize(x)
-    d_grad = (wp.identity(n=2, dtype=float) - wp.outer(d, d)) / wp.length(x)
+    grad[1, 1] = 1.0
+    return grad
 
-    ax = wp.abs(x)
-    xinf = wp.max(ax)
 
-    xinf_grad = wp.select(ax[0] > ax[1], wp.vec2(0.0, wp.sign(x[1])), wp.vec(wp.sign(x[0]), 0.0))
+@wp.func
+def permeability_field(
+    pos: wp.vec3,
+    core_radius: float,
+    core_height: float,
+    coil_internal_radius: float,
+    coil_external_radius: float,
+    coil_height: float,
+):
+    x = wp.abs(pos[0])
+    y = wp.abs(pos[1])
+    z = wp.abs(pos[2])
 
-    return d_grad * xinf + wp.outer(d, xinf_grad)
+    r = wp.sqrt(x * x + z * z)
 
+    if r <= core_radius:
+        return wp.select(y < core_height, MU_0, MU_i)
 
-@wp.func
-def permeability_field(pos: wp.vec2, coil_height: float, coil_internal_radius: float, coil_external_radius: float):
-    # space-varying permeability
+    if r >= coil_internal_radius and r <= coil_external_radius:
+        return wp.select(y < coil_height, MU_0, MU_c)
 
-    x = wp.abs(pos[0])
-    y = wp.abs(pos[1])
-    return wp.select(
-        y < coil_height and x > coil_internal_radius and x < coil_external_radius,
-        MU_0,
-        MU_C,
-    )
+    return MU_0
 
 
 @wp.func
 def current_field(
-    pos: wp.vec2, coil_height: float, coil_internal_radius: float, coil_external_radius: float, current: float
+    pos: wp.vec3,
+    current: float,
+    coil_internal_radius: float,
+    coil_external_radius: float,
+    coil_height: float,
 ):
-    # space-varying current direction along z axis (0, +1 or -1)
-    x = wp.abs(pos[0])
+    x = pos[0]
     y = wp.abs(pos[1])
-    return (
-        wp.select(
-            y < coil_height and x > coil_internal_radius and x < coil_external_radius,
-            0.0,
-            wp.sign(pos[0]),
-        )
-        * current
-    )
+    z = pos[2]
 
+    r = wp.sqrt(x * x + z * z)
 
-@fem.integrand
-def curl_z(u: fem.Field, s: fem.Sample):
-    # projection of curl((0, 0, u)) over z axis
-    du = fem.grad(u, s)
-    return wp.vec2(du[1], -du[0])
+    return wp.select(
+        y < coil_height and r >= coil_internal_radius and r <= coil_external_radius,
+        wp.vec3(0.0),
+        wp.vec3(z, 0.0, -x) * current / r,
+    )
 
 
 @fem.integrand
 def curl_curl_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, mu: fem.Field):
-    return wp.dot(curl_z(u, s), curl_z(v, s)) / mu(s)
+    return wp.dot(fem.curl(u, s), fem.curl(v, s)) / mu(s)
 
 
 @fem.integrand
 def mass_form(s: fem.Sample, domain: fem.Domain, v: fem.Field, u: fem.Field):
-    return u(s) * v(s)
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def curl_expr(s: fem.Sample, u: fem.Field):
+    return fem.curl(u, s)
 
 
 class Example:
-    def __init__(self, quiet=False, degree=2, resolution=32, domain_radius=2.0, current=1.0e6):
+    def __init__(self, quiet=False, mesh: str = "grid", resolution=32, domain_radius=2.0, current=1.0e6):
         # We mesh the unit disk by first meshing the unit square, then building a deformed geometry
         # from an implicit mapping field
-        square_geo = fem.Grid2D(
-            bounds_lo=wp.vec2(-domain_radius, -domain_radius),
-            bounds_hi=wp.vec2(domain_radius, domain_radius),
-            res=wp.vec2i(resolution, resolution),
-        )
 
-        def_field = fem.ImplicitField(domain=fem.Cells(square_geo), func=square_to_disk, grad_func=square_to_disk_grad)
-        disk_geo = def_field.make_deformed_geometry(relative=False)
+        if mesh == "hex":
+            positions, hex_vidx = fem_example_utils.gen_hexmesh(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+            cube_geo = fem.Hexmesh(hex_vertex_indices=hex_vidx, positions=positions)
+        elif mesh == "tet":
+            positions, tet_vidx = fem_example_utils.gen_tetmesh(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+            cube_geo = fem.Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)
+        elif mesh == "nano":
+            vol = fem_example_utils.gen_volume(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+            cube_geo = fem.Nanogrid(grid=vol)
+        else:
+            cube_geo = fem.Grid3D(
+                bounds_lo=wp.vec3(-domain_radius, -domain_radius, -domain_radius),
+                bounds_hi=wp.vec3(domain_radius, domain_radius, domain_radius),
+                res=wp.vec3i(resolution, resolution, resolution),
+            )
+
+        def_field = fem.ImplicitField(
+            domain=fem.Cells(cube_geo), func=cube_to_cylinder, grad_func=cube_to_cylinder_grad
+        )
+        sim_geo = def_field.make_deformed_geometry(relative=False)
 
-        coil_config = {"coil_height": 1.0, "coil_internal_radius": 0.1, "coil_external_radius": 0.3}
+        coil_config = {"coil_height": 0.25, "coil_internal_radius": 0.3, "coil_external_radius": 0.4}
+        core_config = {"core_height": 1.0, "core_radius": 0.2}
 
-        domain = fem.Cells(disk_geo)
-        self._permeability_field = fem.ImplicitField(domain, func=permeability_field, values=coil_config)
+        domain = fem.Cells(sim_geo)
+        self._permeability_field = fem.ImplicitField(
+            domain, func=permeability_field, values=dict(**coil_config, **core_config)
+        )
         self._current_field = fem.ImplicitField(domain, func=current_field, values=dict(current=current, **coil_config))
 
-        z_space = fem.make_polynomial_space(disk_geo, degree=degree, element_basis=fem.ElementBasis.LAGRANGE)
-        xy_space = fem.make_polynomial_space(
-            disk_geo, degree=degree, element_basis=fem.ElementBasis.LAGRANGE, dtype=wp.vec2
+        A_space = fem.make_polynomial_space(
+            sim_geo, degree=1, element_basis=fem.ElementBasis.NEDELEC_FIRST_KIND, dtype=wp.vec3
         )
+        self.A_field = A_space.make_field()
 
-        self.A_field = z_space.make_field()
-        self.B_field = xy_space.make_field()
+        B_space = fem.make_polynomial_space(sim_geo, degree=1, element_basis=fem.ElementBasis.LAGRANGE, dtype=wp.vec3)
+        self.B_field = B_space.make_field()
 
         self.renderer = fem_example_utils.Plot()
 
     def step(self):
-        z_space = self.A_field.space
-        disk_geo = z_space.geometry
+        A_space = self.A_field.space
+        sim_geo = A_space.geometry
 
-        u = fem.make_trial(space=z_space)
-        v = fem.make_test(space=z_space)
-        lhs = fem.integrate(curl_curl_form, fields={"u": u, "v": v, "mu": self._permeability_field})
-        rhs = fem.integrate(mass_form, fields={"v": v, "u": self._current_field})
+        u = fem.make_trial(space=A_space)
+        v = fem.make_test(space=A_space)
+        lhs = fem.integrate(curl_curl_form, fields={"u": u, "v": v, "mu": self._permeability_field}, output_dtype=float)
+        rhs = fem.integrate(mass_form, fields={"v": v, "u": self._current_field}, output_dtype=float)
 
         # Dirichlet BC
-        boundary = fem.BoundarySides(disk_geo)
-        u_bd = fem.make_trial(space=z_space, domain=boundary)
-        v_bd = fem.make_test(space=z_space, domain=boundary)
-        dirichlet_bd_proj = fem.integrate(mass_form, fields={"u": u_bd, "v": v_bd}, nodal=True)
+        boundary = fem.BoundarySides(sim_geo)
+        u_bd = fem.make_trial(space=A_space, domain=boundary)
+        v_bd = fem.make_test(space=A_space, domain=boundary)
+        dirichlet_bd_proj = fem.integrate(mass_form, fields={"u": u_bd, "v": v_bd}, nodal=True, output_dtype=float)
         fem.project_linear_system(lhs, rhs, dirichlet_bd_proj)
 
-        x = wp.zeros_like(rhs)
-        fem_example_utils.bsr_cg(lhs, b=rhs, x=x, tol=1.0e-8, quiet=False)
-
-        # make sure result is exactly zero outside of circle
-        wp.sparse.bsr_mv(dirichlet_bd_proj, x=x, y=x, alpha=-1.0, beta=1.0)
-        wp.utils.array_cast(in_array=x, out_array=self.A_field.dof_values)
+        # solve using Conjugate Residual (numerically rhs may not be in image of lhs)
+        fem_example_utils.bsr_cg(lhs, b=rhs, x=self.A_field.dof_values, method="cr", max_iters=250, quiet=False)
 
         # compute B as curl(A)
-        fem.interpolate(curl_z, dest=self.B_field, fields={"u": self.A_field})
+        fem.interpolate(curl_expr, dest=self.B_field, fields={"u": self.A_field})
 
     def render(self):
-        self.renderer.add_field("A", self.A_field)
         self.renderer.add_field("B", self.B_field)
 
 
@@ -170,7 +212,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
     parser.add_argument("--resolution", type=int, default=32, help="Grid resolution.")
-    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--mesh", type=str, default="grid", choices=["tet", "hex", "grid", "nano"], help="Mesh type.")
     parser.add_argument("--radius", type=float, default=2.0, help="Radius of simulation domain.")
     parser.add_argument(
         "--headless",
@@ -182,9 +224,9 @@ if __name__ == "__main__":
     args = parser.parse_known_args()[0]
 
     with wp.ScopedDevice(args.device):
-        example = Example(quiet=args.quiet, degree=args.degree, resolution=args.resolution, domain_radius=args.radius)
+        example = Example(quiet=args.quiet, mesh=args.mesh, resolution=args.resolution, domain_radius=args.radius)
         example.step()
         example.render()
 
         if not args.headless:
-            example.renderer.plot({"A": {"contours": {"levels": 30}}, "B": {"streamlines": {"density": 1.0}}})
+            example.renderer.plot({"B": {"streamlines": {"density": 1.0}}})

warp/examples/fem/example_nonconforming_contact.py

@@ -210,11 +210,11 @@ class Example:
 
         # Add boundary stress from other solid field
         other_stress_field = fem.field.field.NonconformingField(boundary, other_stress_field)
-        fem.utils.array_axpy(
-            y=u_rhs,
-            x=fem.integrate(
-                boundary_stress_form, fields={"u": u_bd_test, "tau": other_stress_field}, output_dtype=wp.vec2d
-            ),
+        fem.integrate(
+            boundary_stress_form,
+            fields={"u": u_bd_test, "tau": other_stress_field},
+            output=u_rhs,
+            add=True,
         )
 
         # Enforce boundary conditions