-
Notifications
You must be signed in to change notification settings - Fork 493
Expand file tree
/
Copy pathexample_nonconforming_contact.py
More file actions
314 lines (261 loc) · 10.9 KB
/
example_nonconforming_contact.py
File metadata and controls
314 lines (261 loc) · 10.9 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
308
309
310
311
312
313
314
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
###########################################################################
# Example Nonconforming Contact
#
# This example demonstrates using nonconforming fields (warp.fem.NonconformingField)
# to solve a no-slip contact problem between two elastic bodies discretized separately.
#
# Div[ E: D(u) ] = g over each body
# u_top = u_bottom along the contact surface (top body bottom boundary)
# u_bottom = 0 along the bottom boundary of the bottom body
#
# with E the rank-4 elasticity tensor
#
# Below we use a simple staggered scheme for solving bodies iteratively,
# but more robust methods could be considered (e.g. Augmented Lagrangian)
###########################################################################
from typing import Any
import numpy as np
import warp as wp
import warp.examples.fem.utils as fem_example_utils
import warp.fem as fem
@wp.func
def compute_stress(tau: Any, E: Any):
"""Strain to stress computation (using Voigt notation to drop tensor order)"""
tau_sym = wp.vector(tau[0, 0], tau[1, 1], tau[0, 1] + tau[1, 0], dtype=tau.dtype)
sig_sym = E * tau_sym
h = tau.dtype(0.5)
return type(tau)(sig_sym[0], h * sig_sym[2], h * sig_sym[2], sig_sym[1])
@fem.integrand
def stress_form(s: fem.Sample, u: fem.Field, tau: fem.Field, E: wp.mat33):
"""Stress inside body: (E : D(u)) : tau"""
return wp.ddot(tau(s), compute_stress(fem.D(u, s), E))
@fem.integrand
def stress_form_fp64(s: fem.Sample, u: fem.Field, tau: fem.Field, E: wp.mat33d):
"""Stress inside body (fp64): (E : D(u)) : tau"""
return wp.ddot(tau(s), compute_stress(fem.D(u, s), E))
@fem.integrand
def boundary_stress_form(
s: fem.Sample,
domain: fem.Domain,
u: fem.Field,
tau: fem.Field,
):
"""Stress on boundary: u' tau n"""
return wp.dot(tau(s) * fem.normal(domain, s), u(s))
@fem.integrand
def symmetric_grad_form(
s: fem.Sample,
u: fem.Field,
tau: fem.Field,
):
"""Symmetric part of gradient of displacement: D(u) : tau"""
return wp.ddot(tau(s), fem.D(u, s))
@fem.integrand
def gravity_form(
s: fem.Sample,
domain: fem.Domain,
v: fem.Field,
gravity: float,
):
return -fem.scalar_type(domain)(gravity) * v(s)[1]
@fem.integrand
def bottom_boundary_projector_form(
s: fem.Sample,
domain: fem.Domain,
u: fem.Field,
v: fem.Field,
):
# non zero on bottom boundary only
nor = fem.normal(domain, s)
return wp.dot(u(s), v(s)) * wp.max(fem.scalar_type(domain)(0.0), -nor[1])
@fem.integrand
def tensor_mass_form(
s: fem.Sample,
sig: fem.Field,
tau: fem.Field,
):
return wp.ddot(tau(s), sig(s))
class Example:
def __init__(
self,
degree=2,
resolution=16,
young_modulus=1.0,
poisson_ratio=0.5,
nonconforming_stresses=False,
fp64=False,
):
scalar_type = wp.float64 if fp64 else wp.float32
vec2_type = wp.vec2d if fp64 else wp.vec2
mat22_type = wp.mat22d if fp64 else wp.mat22
mat33_type = wp.mat33d if fp64 else wp.mat33
self._geo1 = fem.Grid2D(bounds_hi=vec2_type(1.0, 0.5), res=wp.vec2i(resolution), scalar_type=scalar_type)
self._geo2 = fem.Grid2D(
bounds_lo=vec2_type(0.33, 0.5),
bounds_hi=vec2_type(0.67, 0.5 + 0.33),
res=wp.vec2i(resolution),
scalar_type=scalar_type,
)
# Strain-stress matrix
young = young_modulus
poisson = poisson_ratio
self._elasticity_mat = mat33_type(
young
/ (1.0 - poisson * poisson)
* np.array(
[
[1.0, poisson, 0.0],
[poisson, 1.0, 0.0],
[0.0, 0.0, (1.0 - poisson * poisson) / (2.0 * (1.0 + poisson))],
]
)
)
# Displacement spaces and fields -- S_k
self._u1_space = fem.make_polynomial_space(
self._geo1, degree=degree, dtype=vec2_type, element_basis=fem.ElementBasis.SERENDIPITY
)
self._u2_space = fem.make_polynomial_space(
self._geo2, degree=degree, dtype=vec2_type, element_basis=fem.ElementBasis.SERENDIPITY
)
self._u1_field = self._u1_space.make_field()
self._u2_field = self._u2_space.make_field()
# Stress spaces and fields -- Q_{k-1}d
# Store stress degrees of freedom as symmetric tensors (3 dof) rather than full 2x2 matrices
self._tau1_space = fem.make_polynomial_space(
self._geo1,
degree=degree,
discontinuous=True,
element_basis=fem.ElementBasis.LAGRANGE,
dof_mapper=fem.SymmetricTensorMapper(mat22_type),
)
self._tau2_space = fem.make_polynomial_space(
self._geo2,
degree=degree,
discontinuous=True,
element_basis=fem.ElementBasis.LAGRANGE,
dof_mapper=fem.SymmetricTensorMapper(mat22_type),
)
self._sig1_field = self._tau1_space.make_field()
self._sig2_field = self._tau2_space.make_field()
self._sig2_field_new = self._tau2_space.make_field()
self.renderer = fem_example_utils.Plot()
def step(self):
# Solve for the two bodies separately
# Body (top) is 25x more dense and 5x stiffer than top body
# Body 1 affects body 2 through bottom displacement dirichlet BC
# Body 2 affects body 1 through applied strain on top
self.solve_solid(self._u1_field, self._sig1_field, self._u2_field, self._sig2_field, gravity=1.0, stiffness=1.0)
self.solve_solid(
self._u2_field, self._sig2_field_new, self._u1_field, self._sig1_field, gravity=25.0, stiffness=5.0
)
# Damped update of coupling stress (for stability)
alpha = 0.1
fem.linalg.array_axpy(
x=self._sig2_field_new.dof_values, y=self._sig2_field.dof_values, alpha=alpha, beta=1.0 - alpha
)
def solve_solid(
self,
u_field,
stress_field,
other_u_field,
other_stress_field,
gravity: float,
stiffness: float,
):
u_space = u_field.space
stress_space = stress_field.space
geo = u_field.space.geometry
vec2_type = wp.vec2d if geo.scalar_type == wp.float64 else wp.vec2
domain = fem.Cells(geometry=geo)
boundary = fem.BoundarySides(geometry=geo)
u_test = fem.make_test(space=u_space, domain=domain)
u_trial = fem.make_trial(space=u_space, domain=domain)
tau_test = fem.make_test(space=stress_space, domain=domain)
tau_trial = fem.make_trial(space=stress_space, domain=domain)
u_bd_test = fem.make_test(space=u_space, domain=boundary)
u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
# Assemble stiffness matrix
# (Note: this is constant per body, this could be precomputed)
sym_grad_matrix = fem.integrate(symmetric_grad_form, fields={"u": u_trial, "tau": tau_test})
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)
_stress_form = stress_form_fp64 if geo.scalar_type == wp.float64 else stress_form
E = self._elasticity_mat * geo.scalar_type(stiffness)
stress_matrix = tau_inv_mass_matrix @ fem.integrate(
_stress_form, fields={"u": u_trial, "tau": tau_test}, values={"E": E}
)
stiffness_matrix = sym_grad_matrix.transpose() @ stress_matrix
# Right-hand-side
u_rhs = fem.integrate(gravity_form, fields={"v": u_test}, values={"gravity": gravity}, output_dtype=vec2_type)
# Add boundary stress from other solid field
other_stress_field = fem.NonconformingField(boundary, other_stress_field)
fem.integrate(
boundary_stress_form,
fields={"u": u_bd_test, "tau": other_stress_field},
output=u_rhs,
add=True,
)
# Enforce boundary conditions
u_bd_matrix = fem.integrate(
bottom_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal"
)
# displacement from other body defines bottom boundary Dirichlet BC
other_u_field = fem.NonconformingField(boundary, other_u_field)
u_bd_rhs = fem.integrate(
bottom_boundary_projector_form, fields={"u": other_u_field, "v": u_bd_test}, assembly="nodal"
)
fem.project_linear_system(stiffness_matrix, u_rhs, u_bd_matrix, u_bd_rhs)
# solve
x = wp.zeros_like(u_rhs)
wp.utils.array_cast(in_array=u_field.dof_values, out_array=x)
fem_example_utils.bsr_cg(stiffness_matrix, b=u_rhs, x=x, tol=1.0e-6, quiet=not wp.config.verbose)
# Extract result
stress = stress_matrix @ x
wp.utils.array_cast(in_array=x, out_array=u_field.dof_values)
wp.utils.array_cast(in_array=stress, out_array=stress_field.dof_values)
def render(self):
self.renderer.add_field("u1", self._u1_field)
self.renderer.add_field("u2", self._u2_field)
self.renderer.add_field("sig1", self._sig1_field)
self.renderer.add_field("sig2", self._sig2_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("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
parser.add_argument("--young-modulus", type=float, default=10.0)
parser.add_argument("--poisson-ratio", type=float, default=0.9)
parser.add_argument("--num-steps", type=int, default=50)
parser.add_argument("--fp64", action="store_true", help="Use fp64 precision.")
parser.add_argument(
"--headless",
action="store_true",
help="Run in headless mode, suppressing the opening of any graphical windows.",
)
args = parser.parse_known_args()[0]
with wp.ScopedDevice(args.device):
example = Example(
degree=args.degree,
resolution=args.resolution,
young_modulus=args.young_modulus,
poisson_ratio=args.poisson_ratio,
fp64=args.fp64,
)
for i in range(args.num_steps):
print("Step", i)
example.step()
example.render()
if not args.headless:
example.renderer.plot(
{
"rows": 2,
"u1": {"displacement": {}},
"u2": {"displacement": {}},
},
)