NvidiaWarp-GarmentCode/examples/fem/example_mixed_elasticity.py

""" This example illustrates using Mixed FEM to solve a 2D linear elasticity problem
 Div[ E: D(u) ] = 0
with Dirichlet boundary conditions on horizontal sides, and E the elasticity rank-4 tensor
"""
import argparse

import warp as wp
import numpy as np

import warp.fem as fem

from warp.sparse import bsr_transposed, bsr_mm

try:
    from .plot_utils import Plot
    from .bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
    from .mesh_utils import gen_trimesh, gen_quadmesh
except ImportError:
    from plot_utils import Plot
    from bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
    from mesh_utils import gen_trimesh, gen_quadmesh


@wp.func
def compute_stress(tau: wp.mat22, E: wp.mat33):
    """Strain to stress computation"""
    tau_sym = wp.vec3(tau[0, 0], tau[1, 1], tau[0, 1] + tau[1, 0])
    sig_sym = E * tau_sym
    return wp.mat22(sig_sym[0], 0.5 * sig_sym[2], 0.5 * sig_sym[2], sig_sym[1])


@fem.integrand
def symmetric_grad_form(
    s: fem.Sample,
    u: fem.Field,
    tau: fem.Field,
):
    """D(u) : tau"""
    return wp.ddot(tau(s), fem.D(u, s))


@fem.integrand
def stress_form(s: fem.Sample, u: fem.Field, tau: fem.Field, E: wp.mat33):
    """(E : D(u)) : tau"""
    return wp.ddot(tau(s), compute_stress(fem.D(u, s), E))


@fem.integrand
def horizontal_boundary_projector_form(
    s: fem.Sample,
    domain: fem.Domain,
    u: fem.Field,
    v: fem.Field,
):
    # non zero on horizontal boundary of domain only
    nor = fem.normal(domain, s)
    return wp.dot(u(s), v(s)) * wp.abs(nor[1])


@fem.integrand
def horizontal_displacement_form(
    s: fem.Sample,
    domain: fem.Domain,
    v: fem.Field,
    displacement: float,
):
    # opposed to normal on horizontal boundary of domain only
    nor = fem.normal(domain, s)
    return -wp.abs(nor[1]) * 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))


class Example:
    parser = argparse.ArgumentParser()
    parser.add_argument("--resolution", type=int, default=25)
    parser.add_argument("--degree", type=int, default=2)
    parser.add_argument("--displacement", type=float, default=0.1)
    parser.add_argument("--young_modulus", type=float, default=1.0)
    parser.add_argument("--poisson_ratio", type=float, default=0.5)
    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)"
    )

    def __init__(self, stage=None, quiet=False, args=None, **kwargs):
        if args is None:
            # Read args from kwargs, add default arg values from parser
            args = argparse.Namespace(**kwargs)
            args = Example.parser.parse_args(args=[], namespace=args)
        self._args = args
        self._quiet = quiet

        # Grid or triangle mesh geometry
        if args.mesh == "tri":
            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
            self._geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
        elif args.mesh == "quad":
            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
            self._geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
        else:
            self._geo = fem.Grid2D(res=wp.vec2i(args.resolution))

        # Strain-stress matrix
        young = args.young_modulus
        poisson = args.poisson_ratio
        self._elasticity_mat = wp.mat33(
            young
            / (1.0 - poisson * poisson)
            * np.array(
                [
                    [1.0, poisson, 0.0],
                    [poisson, 1.0, 0.0],
                    [0.0, 0.0, (2.0 * (1.0 + poisson)) * (1.0 - poisson * poisson)],
                ]
            )
        )

        # Function spaces -- S_k for displacement, Q_{k-1}d for stress
        self._u_space = fem.make_polynomial_space(
            self._geo, degree=args.degree, dtype=wp.vec2, element_basis=fem.ElementBasis.SERENDIPITY
        )

        # Store stress degrees of freedom as symmetric tensors (3 dof) rather than full 2x2 matrices
        tau_basis = (
            fem.ElementBasis.NONCONFORMING_POLYNOMIAL if args.nonconforming_stresses else fem.ElementBasis.LAGRANGE
        )
        self._tau_space = fem.make_polynomial_space(
            self._geo,
            degree=args.degree - 1,
            discontinuous=True,
            element_basis=tau_basis,
            dof_mapper=fem.SymmetricTensorMapper(wp.mat22),
        )

        self._u_field = self._u_space.make_field()

        self.renderer = Plot(stage)

    def update(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(
            horizontal_displacement_form,
            fields={"v": u_bd_test},
            values={"displacement": self._args.displacement},
            nodal=True,
            output_dtype=wp.vec2d,
        )
        u_bd_matrix = fem.integrate(
            horizontal_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)

        sym_grad_matrix = fem.integrate(symmetric_grad_form, fields={"u": u_trial, "tau": tau_test})
        stress_matrix = fem.integrate(
            stress_form, fields={"u": u_trial, "tau": tau_test}, values={"E": self._elasticity_mat}
        )

        # 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)
        invert_diagonal_bsr_mass_matrix(tau_inv_mass_matrix)

        # Assemble system matrix
        u_matrix = bsr_mm(bsr_transposed(sym_grad_matrix), bsr_mm(tau_inv_mass_matrix, stress_matrix))

        # Enforce boundary conditions
        u_rhs = wp.zeros_like(u_bd_rhs)
        fem.project_linear_system(u_matrix, u_rhs, u_bd_matrix, u_bd_rhs)

        x = wp.zeros_like(u_rhs)
        bsr_cg(u_matrix, b=u_rhs, x=x, tol=1.0e-16, quiet=self._quiet)

        # Extract result
        self._u_field.dof_values = x

    def render(self):
        self.renderer.add_surface_vector("solution", self._u_field)


if __name__ == "__main__":
    wp.init()
    wp.set_module_options({"enable_backward": False})

    args = Example.parser.parse_args()

    example = Example(args=args)
    example.update()
    example.render()

    example.renderer.plot()