-
Notifications
You must be signed in to change notification settings - Fork 493
Expand file tree
/
Copy pathexample_mixed_elasticity.py
More file actions
307 lines (245 loc) · 11.1 KB
/
example_mixed_elasticity.py
File metadata and controls
307 lines (245 loc) · 11.1 KB
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
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
###########################################################################
# 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)
@fem.integrand
def stress_norm_form(s: fem.Sample, stress: fem.Field):
"""Frobenius norm of the stress tensor"""
P = stress(s)
return wp.sqrt(wp.ddot(P, P))
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.NONCONFORMING_POLYNOMIAL if nonconforming_stresses else 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._stress_field = self._tau_space.make_field()
self._stress_norm_space = fem.make_polynomial_space(
self._geo,
degree=degree,
dtype=float,
)
self._stress_norm_field = self._stress_norm_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.vec2,
)
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.types.vector(length=stress_matrix.block_shape[0], dtype=float),
)
# 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)
# Compute stress field from last Newton iteration: sigma = M_tau^{-1} * stress_rhs
stress_dofs_f64 = wp.zeros_like(stress_rhs)
wp.sparse.bsr_mv(tau_inv_mass_matrix, x=stress_rhs, y=stress_dofs_f64)
wp.utils.array_cast(in_array=stress_dofs_f64, out_array=self._stress_field.dof_values)
# Compute scalar stress norm for visualization
fem.interpolate(stress_norm_form, dest=self._stress_norm_field, fields={"stress": self._stress_field})
# 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)
self.renderer.add_field("stress", self._stress_norm_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": {}}})