warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/examples/fem/example_mixed_elasticity.py

@@ -0,0 +1,291 @@
+# 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 isinstance(self._geo.reference_cell(), fem.geometry.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},
+            nodal=True,
+            output_dtype=wp.vec2d,
+        )
+        u_bd_matrix = fem.integrate(
+            vertical_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True
+        )
+
+        # Stress/velocity coupling
+        u_trial = fem.make_trial(space=self._u_space, domain=domain)
+        tau_test = fem.make_test(space=self._tau_space, domain=domain)
+        tau_trial = fem.make_trial(space=self._tau_space, domain=domain)
+
+        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}, nodal=True)
+        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.utils.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": {}}})

warp/examples/fem/example_navier_stokes.py

@@ -0,0 +1,261 @@
+# 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 Navier Stokes
+#
+# This example solves a 2D Navier-Stokes flow problem:
+#
+# Du/dt -nu D(u) + grad p = 0
+# Div u = 0
+#
+# with (hard) velocity-Dirichlet boundary conditions
+# and using semi-Lagrangian advection
+###########################################################################
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.fem.utils import array_axpy
+
+
+@wp.func
+def u_boundary_value(x: wp.vec2, box_height: float, top_velocity: float):
+    # Horizontal velocity on top of domain, zero elsewhere
+    if x[1] >= box_height:
+        return wp.vec2(top_velocity, 0.0)
+
+    return wp.vec2(0.0, 0.0)
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def inertia_form(s: fem.Sample, u: fem.Field, v: fem.Field, dt: float):
+    return mass_form(s, u, v) / dt
+
+
+@fem.integrand
+def viscosity_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
+    return 2.0 * nu * wp.ddot(fem.D(u, s), fem.D(v, s))
+
+
+@fem.integrand
+def viscosity_and_inertia_form(s: fem.Sample, u: fem.Field, v: fem.Field, dt: float, nu: float):
+    return inertia_form(s, u, v, dt) + viscosity_form(s, u, v, nu)
+
+
+@fem.integrand
+def transported_inertia_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, dt: float):
+    pos = domain(s)
+    vel = u(s)
+
+    conv_pos = pos - 0.5 * vel * dt
+    conv_s = fem.lookup(domain, conv_pos, s)
+    conv_vel = u(conv_s)
+
+    conv_pos = conv_pos - 0.5 * conv_vel * dt
+    conv_vel = u(fem.lookup(domain, conv_pos, conv_s))
+
+    return wp.dot(conv_vel, v(s)) / dt
+
+
+@fem.integrand
+def div_form(
+    s: fem.Sample,
+    u: fem.Field,
+    q: fem.Field,
+):
+    return -q(s) * fem.div(u, s)
+
+
+class Example:
+    def __init__(self, quiet=False, degree=2, resolution=25, Re=1000.0, top_velocity=1.0, tri_mesh=False):
+        self._quiet = quiet
+
+        res = resolution
+        self.sim_dt = 1.0 / resolution
+        self.current_frame = 0
+
+        viscosity = top_velocity / Re
+
+        if tri_mesh:
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(res))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions, build_bvh=True)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(res))
+
+        domain = fem.Cells(geometry=geo)
+        boundary = fem.BoundarySides(geo)
+
+        # Functions spaces: Q(d)-Q(d-1)
+        u_degree = degree
+        u_space = fem.make_polynomial_space(geo, degree=u_degree, dtype=wp.vec2)
+        p_space = fem.make_polynomial_space(geo, degree=u_degree - 1)
+
+        # Viscosity and inertia
+        u_test = fem.make_test(space=u_space, domain=domain)
+        u_trial = fem.make_trial(space=u_space, domain=domain)
+
+        u_matrix = fem.integrate(
+            viscosity_and_inertia_form,
+            fields={"u": u_trial, "v": u_test},
+            values={"nu": viscosity, "dt": self.sim_dt},
+        )
+
+        # Pressure-velocity coupling
+        p_test = fem.make_test(space=p_space, domain=domain)
+        div_matrix = fem.integrate(div_form, fields={"u": u_trial, "q": p_test})
+
+        # Enforcing the Dirichlet boundary condition the hard way;
+        # build projector for velocity left- and right-hand-sides
+        u_bd_test = fem.make_test(space=u_space, domain=boundary)
+        u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
+        u_bd_projector = fem.integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True)
+
+        # Define an implicit field for our boundary condition value and integrate
+        u_bd_field = fem.ImplicitField(
+            domain=boundary, func=u_boundary_value, values={"top_velocity": top_velocity, "box_height": 1.0}
+        )
+        u_bd_value = fem.integrate(
+            mass_form,
+            fields={"u": u_bd_field, "v": u_bd_test},
+            nodal=True,
+            output_dtype=wp.vec2d,
+        )
+
+        fem.normalize_dirichlet_projector(u_bd_projector, u_bd_value)
+
+        u_bd_rhs = wp.zeros_like(u_bd_value)
+        fem.project_linear_system(u_matrix, u_bd_rhs, u_bd_projector, u_bd_value, normalize_projector=False)
+
+        div_bd_rhs = -div_matrix @ u_bd_value
+        div_matrix -= div_matrix @ u_bd_projector
+
+        # Assemble saddle system
+        self._saddle_system = fem_example_utils.SaddleSystem(u_matrix, div_matrix)
+
+        # Save data for computing time steps rhs
+        self._u_bd_projector = u_bd_projector
+        self._u_bd_rhs = u_bd_rhs
+        self._u_test = u_test
+        self._div_bd_rhs = div_bd_rhs
+
+        # Velocitiy and pressure fields
+        self._u_field = u_space.make_field()
+        self._p_field = p_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+        self.renderer.add_field("velocity", self._u_field)
+
+    def step(self):
+        self.current_frame += 1
+
+        u_rhs = fem.integrate(
+            transported_inertia_form,
+            fields={"u": self._u_field, "v": self._u_test},
+            values={"dt": self.sim_dt},
+            output_dtype=wp.vec2d,
+        )
+
+        # Apply boundary conditions
+        # u_rhs = (I - P) * u_rhs + u_bd_rhs
+        wp.sparse.bsr_mv(self._u_bd_projector, x=u_rhs, y=u_rhs, alpha=-1.0, beta=1.0)
+        array_axpy(x=self._u_bd_rhs, y=u_rhs, alpha=1.0, beta=1.0)
+
+        p_rhs = self._div_bd_rhs
+
+        x_u = wp.empty_like(u_rhs)
+        x_p = wp.empty_like(p_rhs)
+        wp.utils.array_cast(out_array=x_u, in_array=self._u_field.dof_values)
+        wp.utils.array_cast(out_array=x_p, in_array=self._p_field.dof_values)
+
+        fem_example_utils.bsr_solve_saddle(
+            saddle_system=self._saddle_system,
+            tol=1.0e-6,
+            x_u=x_u,
+            x_p=x_p,
+            b_u=u_rhs,
+            b_p=p_rhs,
+            quiet=self._quiet,
+        )
+
+        wp.utils.array_cast(in_array=x_u, out_array=self._u_field.dof_values)
+        wp.utils.array_cast(in_array=x_p, out_array=self._p_field.dof_values)
+
+    def render(self):
+        self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
+        self.renderer.add_field("velocity", self._u_field)
+        self.renderer.end_frame()
+
+
+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("--num_frames", type=int, default=1000, help="Total number of frames.")
+    parser.add_argument(
+        "--top_velocity",
+        type=float,
+        default=1.0,
+        help="Horizontal velocity boundary condition at the top of the domain.",
+    )
+    parser.add_argument("--Re", type=float, default=1000.0, help="Reynolds number.")
+    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh.")
+    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,
+            Re=args.Re,
+            top_velocity=args.top_velocity,
+            tri_mesh=args.tri_mesh,
+        )
+
+        for k in range(args.num_frames):
+            print(f"Frame {k}:")
+            example.step()
+            example.render()
+
+        example.renderer.add_field("velocity_final", example._u_field)
+
+        if not args.headless:
+            example.renderer.plot(
+                options={
+                    "velocity": {"arrows": {"glyph_scale": 0.25}},
+                    "velocity_final": {"streamlines": {"density": 2.0}},
+                }
+            )