warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/fem/example_magnetostatics.py

@@ -0,0 +1,242 @@
+# 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 Dirichlet homogeneous essential boundary conditions
+#
+# This example also illustrates using an ImplicitField 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}, assembly="nodal", 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}}})

warp/examples/fem/example_mixed_elasticity.py

@@ -0,0 +1,293 @@
+# 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 Mixed Elasticity
+#
+# This example illustrates using Mixed FEM to solve a nonlinear static elasticity equilibrium problem:
+#
+# Div[ d/dF Psi(F(u)) ] = 0
+#
+# with Dirichlet boundary conditions on vertical sides and Psi an elastic potential function of the deformation gradient.
+# Here we choose Psi Neo-Hookean, as per Sec 3.2 of "Stable Neo-Hookean Flesh Simulation" (Smith et al. 2018),
+# Psi(F) = mu ||F||^2 + lambda (det J - 1 - mu/lambda)^2
+#
+# which we write as a sequence of Newton iterations:
+# int {sigma : grad v}  = 0   for all displacement test functions v
+# int {sigma : tau} = int{dPsi/dF : tau} + int{grad du : d2 Psi/dF2  : tau} for all stress test functions tau
+###########################################################################
+
+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 displacement_gradient_form(
+    s: fem.Sample,
+    u: fem.Field,
+    tau: fem.Field,
+):
+    """grad(u) : tau"""
+    return wp.ddot(tau(s), fem.grad(u, s))
+
+
+@wp.func
+def nh_parameters_from_lame(lame: wp.vec2):
+    """Parameters such that for small strains model behaves according to Hooke's law"""
+    mu_nh = lame[1]
+    lambda_nh = lame[0] + lame[1]
+
+    return mu_nh, lambda_nh
+
+
+@fem.integrand
+def nh_stress_form(s: fem.Sample, tau: fem.Field, u_cur: fem.Field, lame: wp.vec2):
+    """d Psi/dF : tau"""
+
+    # Deformation gradient
+    F = wp.identity(n=2, dtype=float) + fem.grad(u_cur, s)
+
+    # Area term and its derivative w.r.t F
+    J = wp.determinant(F)
+    dJ_dF = wp.mat22(F[1, 1], -F[1, 0], -F[0, 1], F[0, 0])
+
+    mu_nh, lambda_nh = nh_parameters_from_lame(lame)
+    nh_stress = mu_nh * F + (lambda_nh * (J - 1.0) - mu_nh) * dJ_dF
+
+    return wp.ddot(tau(s), nh_stress)
+
+
+@fem.integrand
+def nh_stress_delta_form(s: fem.Sample, tau: fem.Field, u: fem.Field, u_cur: fem.Field, lame: wp.vec2):
+    """grad(u) : d2 Psi/dF2 : tau"""
+
+    tau_s = tau(s)
+    sigma_s = fem.grad(u, s)
+
+    F = wp.identity(n=2, dtype=float) + fem.grad(u_cur, s)
+    dJ_dF = wp.mat22(F[1, 1], -F[1, 0], -F[0, 1], F[0, 0])
+
+    # Gauss--Newton approximation; ignore d2J/dF2 term
+    mu_nh, lambda_nh = nh_parameters_from_lame(lame)
+    return mu_nh * wp.ddot(tau_s, sigma_s) + lambda_nh * wp.ddot(dJ_dF, tau_s) * wp.ddot(dJ_dF, sigma_s)
+
+
+@fem.integrand
+def vertical_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # non zero on vertical boundary of domain only
+    nor = fem.normal(domain, s)
+    return wp.dot(u(s), v(s)) * wp.abs(nor[0])
+
+
+@fem.integrand
+def vertical_displacement_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    v: fem.Field,
+    displacement: float,
+):
+    # opposed to normal on vertical boundary of domain only
+    nor = fem.normal(domain, s)
+    return -wp.abs(nor[0]) * displacement * wp.dot(nor, v(s))
+
+
+@fem.integrand
+def tensor_mass_form(
+    s: fem.Sample,
+    sig: fem.Field,
+    tau: fem.Field,
+):
+    return wp.ddot(tau(s), sig(s))
+
+
+@fem.integrand
+def area_form(s: fem.Sample, u_cur: fem.Field):
+    F = wp.identity(n=2, dtype=float) + fem.grad(u_cur, s)
+    return wp.determinant(F)
+
+
+class Example:
+    def __init__(
+        self,
+        quiet=False,
+        degree=2,
+        resolution=25,
+        mesh="grid",
+        displacement=0.1,
+        poisson_ratio=0.5,
+        nonconforming_stresses=False,
+    ):
+        self._quiet = quiet
+
+        self._displacement = displacement
+
+        # Grid or mesh geometry
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
+            self._geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            self._geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            self._geo = fem.Grid2D(res=wp.vec2i(resolution))
+
+        # Lame coefficients from Young modulus and Poisson ratio
+        self._lame = wp.vec2(1.0 / (1.0 + poisson_ratio) * np.array([poisson_ratio / (1.0 - poisson_ratio), 0.5]))
+
+        # Function spaces -- S_k for displacement, Q_k or P_{k-1}d for stress
+        self._u_space = fem.make_polynomial_space(
+            self._geo, degree=degree, dtype=wp.vec2, element_basis=fem.ElementBasis.SERENDIPITY
+        )
+
+        if self._geo.reference_cell() == fem.Element.TRIANGLE:
+            # triangle elements
+            tau_basis = fem.ElementBasis.NONCONFORMING_POLYNOMIAL
+            tau_degree = degree - 1
+        else:
+            # square elements
+            tau_basis = fem.ElementBasis.LAGRANGE
+            tau_degree = degree
+
+        self._tau_space = fem.make_polynomial_space(
+            self._geo,
+            degree=tau_degree,
+            discontinuous=True,
+            element_basis=tau_basis,
+            family=fem.Polynomial.GAUSS_LEGENDRE,
+            dtype=wp.mat22,
+        )
+
+        self._u_field = self._u_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        boundary = fem.BoundarySides(self._geo)
+        domain = fem.Cells(geometry=self._geo)
+
+        # 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(
+            vertical_displacement_form,
+            fields={"v": u_bd_test},
+            values={"displacement": self._displacement},
+            assembly="nodal",
+            output_dtype=wp.vec2d,
+        )
+        u_bd_matrix = fem.integrate(
+            vertical_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal"
+        )
+
+        # 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)
+
+        gradient_matrix = fem.integrate(displacement_gradient_form, fields={"u": u_trial, "tau": tau_test}).transpose()
+
+        # 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}, assembly="nodal"
+        )
+        fem_example_utils.invert_diagonal_bsr_matrix(tau_inv_mass_matrix)
+
+        # Newton iterations (without line-search for simplicity)
+        for newton_iteration in range(5):
+            stress_matrix = fem.integrate(
+                nh_stress_delta_form,
+                fields={"u_cur": self._u_field, "u": u_trial, "tau": tau_test},
+                values={"lame": self._lame},
+            )
+
+            stress_rhs = fem.integrate(
+                nh_stress_form,
+                fields={"u_cur": self._u_field, "tau": tau_test},
+                values={"lame": self._lame},
+                output_dtype=wp.vec(length=stress_matrix.block_shape[0], dtype=wp.float64),
+            )
+
+            # Assemble system matrix
+            u_matrix = gradient_matrix @ tau_inv_mass_matrix @ stress_matrix
+
+            # Enforce boundary conditions (apply displacement only at first iteration)
+            u_rhs = -gradient_matrix @ (tau_inv_mass_matrix @ stress_rhs)
+            fem.project_linear_system(u_matrix, u_rhs, u_bd_matrix, u_bd_rhs if newton_iteration == 0 else None)
+
+            x = wp.zeros_like(u_rhs)
+            fem_example_utils.bsr_cg(u_matrix, b=u_rhs, x=x, quiet=self._quiet)
+
+            # Extract result -- cast to float32 and accumulate to displacement field
+            delta_u = wp.empty_like(self._u_field.dof_values)
+            wp.utils.array_cast(in_array=x, out_array=delta_u)
+            fem.linalg.array_axpy(x=delta_u, y=self._u_field.dof_values)
+
+        # Evaluate area conservation, should converge to 1.0 as Poisson ratio approaches 1.0
+        final_area = fem.integrate(
+            area_form, quadrature=fem.RegularQuadrature(domain, order=4), fields={"u_cur": self._u_field}
+        )
+        print(f"Area gain: {final_area}  (using Poisson ratio={self._lame[0] / (self._lame[0] + 2.0 * self._lame[1])})")
+
+    def render(self):
+        self.renderer.add_field("solution", self._u_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=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--displacement", type=float, default=-0.5)
+    parser.add_argument("--poisson_ratio", type=float, default=0.99)
+    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)"
+    )
+    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,
+            displacement=args.displacement,
+            poisson_ratio=args.poisson_ratio,
+            nonconforming_stresses=args.nonconforming_stresses,
+        )
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot(options={"solution": {"displacement": {}}})