warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/fem/example_diffusion_mgpu.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Diffusion MGPU
+#
+# This example illustrates using domain decomposition to
+# solve a diffusion PDE over multiple devices
+###########################################################################
+
+from typing import Tuple
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.examples.fem.example_diffusion import diffusion_form, linear_form
+from warp.sparse import bsr_axpy, bsr_mv
+from warp.utils import array_cast
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return u(s) * v(s)
+
+
+@wp.kernel
+def scal_kernel(a: wp.array(dtype=wp.float64), res: wp.array(dtype=wp.float64), alpha: wp.float64):
+    res[wp.tid()] = a[wp.tid()] * alpha
+
+
+@wp.kernel
+def sum_kernel(a: wp.indexedarray(dtype=wp.float64), b: wp.array(dtype=wp.float64)):
+    a[wp.tid()] = a[wp.tid()] + b[wp.tid()]
+
+
+def sum_vecs(vecs, indices, sum: wp.array, tmp: wp.array):
+    for v, idx in zip(vecs, indices):
+        wp.copy(dest=tmp, src=v)
+        idx_sum = wp.indexedarray(sum, idx)
+        wp.launch(kernel=sum_kernel, dim=idx.shape, device=sum.device, inputs=[idx_sum, tmp])
+
+    return sum
+
+
+class DistributedSystem:
+    device = None
+    scalar_type: type
+    tmp_buf: wp.array
+
+    nrow: int
+    shape = Tuple[int, int]
+    rank_data = None
+
+    def mv_routine(self, x: wp.array, y: wp.array, z: wp.array, alpha=1.0, beta=0.0):
+        """Distributed matrix-vector multiplication routine, for example purposes"""
+
+        tmp = self.tmp_buf
+
+        wp.launch(kernel=scal_kernel, dim=y.shape, device=y.device, inputs=[y, z, wp.float64(beta)])
+
+        stream = wp.get_stream()
+
+        for mat_i, x_i, y_i, idx in zip(*self.rank_data):
+            # WAR copy with indexed array requiring matching shape
+            tmp_i = wp.array(ptr=tmp.ptr, device=tmp.device, capacity=tmp.capacity, dtype=tmp.dtype, shape=idx.shape)
+
+            # Compress rhs on rank 0
+            x_idx = wp.indexedarray(x, idx)
+            wp.copy(dest=tmp_i, src=x_idx, count=idx.size, stream=stream)
+
+            # Send to rank i
+            wp.copy(dest=x_i, src=tmp_i, count=idx.size, stream=stream)
+
+            with wp.ScopedDevice(x_i.device):
+                wp.wait_stream(stream)
+                bsr_mv(A=mat_i, x=x_i, y=y_i, alpha=alpha, beta=0.0)
+
+            wp.wait_stream(wp.get_stream(x_i.device))
+
+            # Back to rank 0 for sum
+            wp.copy(dest=tmp_i, src=y_i, count=idx.size, stream=stream)
+            z_idx = wp.indexedarray(z, idx)
+            wp.launch(kernel=sum_kernel, dim=idx.shape, device=z_idx.device, inputs=[z_idx, tmp_i], stream=stream)
+
+        wp.wait_stream(stream)
+
+
+class Example:
+    def __init__(self, quiet=False, device=None):
+        self._bd_weight = 100.0
+        self._quiet = quiet
+
+        self._geo = fem.Grid2D(res=wp.vec2i(25))
+
+        self._main_device = wp.get_device(device)
+
+        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 = fem_example_utils.Plot()
+
+    def step(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)
+
+        # This temporary buffer will be used for peer-to-peer copying during graph capture,
+        # so we allocate it using the default CUDA allocator.  This ensures that the copying
+        # will succeed without enabling mempool access between devices, which is not supported
+        # on all systems.
+        with wp.ScopedMempool(main_device, False):
+            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 = main_device
+        A.dtype = 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)
+
+        with wp.ScopedDevice(main_device):
+            fem_example_utils.bsr_cg(
+                A, x=global_res, b=glob_rhs, use_diag_precond=False, quiet=self._quiet, mv_routine=A.mv_routine
+            )
+
+        array_cast(in_array=global_res, out_array=self._scalar_field.dof_values)
+
+    def render(self):
+        self.renderer.add_field("solution", self._scalar_field)
+
+    def _assemble_local_system(self, geo_partition: fem.GeometryPartition):
+        scalar_space = self._scalar_space
+        space_partition = fem.make_space_partition(scalar_space, geo_partition)
+
+        domain = fem.Cells(geometry=geo_partition)
+
+        # 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})
+
+        # 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})
+
+        bsr_axpy(y=matrix, x=bd_matrix, alpha=self._bd_weight)
+
+        return matrix, rhs, space_partition.space_node_indices()
+
+
+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("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedTimer(__file__):
+        example = Example(quiet=args.quiet, device=args.device)
+
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot()

warp/examples/fem/example_distortion_energy.py

@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# 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.where(s.qp_index == 0, 1.0, 0.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

@@ -0,0 +1,240 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Magnetostatics
+#
+# 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,
+# and Direchlet homogeneous essential boundary conditions
+#
+# This example also illustrates using an ImplictField to warp a grid mesh
+# to a cylindrical domain
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+# 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 cube_to_cylinder(x: wp.vec3):
+    # mapping from unit square to unit disk
+    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 cube_to_cylinder_grad(x: wp.vec3):
+    # gradient of mapping from unit square to unit disk
+    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.where(
+            abs_xz[0] > abs_xz[2], wp.vec(wp.sign(pos_xz[0]), 0.0, 0.0), wp.vec3(0.0, 0.0, wp.sign(pos_xz[2]))
+        )
+        grad = dir_grad * wp.max(abs_xz) + wp.outer(dir_xz, xinf_grad)
+
+    grad[1, 1] = 1.0
+    return grad
+
+
+@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])
+
+    r = wp.sqrt(x * x + z * z)
+
+    if r <= core_radius:
+        return wp.where(y < core_height, MU_i, MU_0)
+
+    if r >= coil_internal_radius and r <= coil_external_radius:
+        return wp.where(y < coil_height, MU_c, MU_0)
+
+    return MU_0
+
+
+@wp.func
+def current_field(
+    pos: wp.vec3,
+    current: float,
+    coil_internal_radius: float,
+    coil_external_radius: float,
+    coil_height: float,
+):
+    x = pos[0]
+    y = wp.abs(pos[1])
+    z = pos[2]
+
+    r = wp.sqrt(x * x + z * z)
+
+    return wp.where(
+        y < coil_height and r >= coil_internal_radius and r <= coil_external_radius,
+        wp.vec3(z, 0.0, -x) * current / r,
+        wp.vec3(0.0),
+    )
+
+
+@fem.integrand
+def curl_curl_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, mu: fem.Field):
+    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 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, 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
+
+        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": 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(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))
+
+        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()
+
+        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):
+        A_space = self.A_field.space
+        sim_geo = A_space.geometry
+
+        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(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)
+
+        # 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_expr, dest=self.B_field, fields={"u": self.A_field})
+
+    def render(self):
+        self.renderer.add_field("B", self.B_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=32, help="Grid resolution.")
+    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",
+        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, mesh=args.mesh, resolution=args.resolution, domain_radius=args.radius)
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot({"B": {"streamlines": {"density": 1.0}}})