Spaces:
Sleeping
Sleeping
File size: 6,971 Bytes
66c9c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | """ 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()
|