Skip to content

Commit b5a8036

Browse files
Q-MinhQ-Minh
committed
Attempt non-linear elasticity energy
1 parent 6fe9fad commit b5a8036

File tree

1 file changed

+129
-59
lines changed

1 file changed

+129
-59
lines changed

‎python/examples/multigrid/linear_restriction_prolongation.py‎

Lines changed: 129 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,16 @@ def min_max_eigs(A, n=10):
1616
return lmin, lmax
1717

1818

19+
def rest_pose_hessian(mesh, Y, nu):
20+
x = mesh.X.reshape(math.prod(mesh.X.shape), order='f')
21+
energy = pbat.fem.HyperElasticEnergy.StableNeoHookean
22+
hep, detJeU, GNeU = pbat.fem.hyper_elastic_potential(
23+
mesh, Y, nu, energy=energy)
24+
hep.compute_element_elasticity(x)
25+
HU = hep.hessian()
26+
return HU
27+
28+
1929
def laplacian_energy(mesh, k=1, dims=1):
2030
L = -pbat.fem.laplacian(mesh, dims=dims)[0]
2131
M = None if k == 1 else pbat.fem.mass_matrix(mesh, dims=1, lump=True)
@@ -25,6 +35,40 @@ def laplacian_energy(mesh, k=1, dims=1):
2535
return U
2636

2737

38+
def display_quadrature(mesh, Xg):
39+
V = mesh.X
40+
C = mesh.E
41+
bvh = pbat.geometry.bvh(V, C, cell=pbat.geometry.Cell.Tetrahedron)
42+
e, d = bvh.nearest_primitives_to_points(Xg, parallelize=True)
43+
nelems = mesh.E.shape[1]
44+
efree = np.setdiff1d(list(range(nelems)), np.unique(e))
45+
if len(efree) > 0:
46+
ps.register_volume_mesh(f"Invalid mesh nelems={
47+
nelems}", V.T, C[:, efree].T)
48+
49+
50+
def gradient_operator(mesh, Xg, dims=1):
51+
nevalpts = Xg.shape[1]
52+
nnodes = mesh.X.shape[1]
53+
V = mesh.X
54+
C = mesh.E
55+
bvh = pbat.geometry.bvh(V, C, cell=pbat.geometry.Cell.Tetrahedron)
56+
e, d = bvh.nearest_primitives_to_points(Xg, parallelize=True)
57+
Xi = pbat.fem.reference_positions(mesh, e, Xg)
58+
dphi = pbat.fem.shape_function_gradients_at(mesh, e, Xi)
59+
data = np.hstack([dphi[:, d::mesh.dims].flatten(order="F")
60+
for d in range(mesh.dims)])
61+
rows = np.hstack([np.repeat(list(range(nevalpts)), dphi.shape[0]) + d*nevalpts
62+
for d in range(mesh.dims)])
63+
cols = np.hstack([mesh.E[:, e].flatten(order="F")
64+
for d in range(mesh.dims)])
65+
G = sp.sparse.coo_matrix((data, (rows, cols)),
66+
shape=(mesh.dims*nevalpts, nnodes))
67+
if dims > 1:
68+
G = sp.sparse.kron(G, sp.sparse.eye(dims))
69+
return G.asformat('csc')
70+
71+
2872
def shape_function_matrix(mesh, Xg, dims=1):
2973
nevalpts = Xg.shape[1]
3074
# Unfortunately, we have to perform the copy here...
@@ -59,12 +103,12 @@ def __init__(self, MD: pbat.fem.Mesh, MS: pbat.fem.Mesh, MT: pbat.fem.Mesh, H):
59103
MS (pbat.fem.Mesh): Source mesh
60104
MT (pbat.fem.Mesh): Target mesh
61105
"""
62-
nelems = MD.E.shape[1]
63106
quadrature_order = 2*max(MS.order, MT.order)
64107
Xg = MD.quadrature_points(quadrature_order)
65-
wg = np.tile(MD.quadrature_weights(quadrature_order), nelems)
66-
Ig = sp.sparse.diags(wg)
67-
Ig = sp.sparse.kron(Ig, sp.sparse.eye(MT.dims))
108+
wg = np.tile(MD.quadrature_weights(quadrature_order), MD.E.shape[1])
109+
from scipy.sparse import kron, eye, diags
110+
Ig = diags(wg)
111+
Ig = kron(Ig, eye(MT.dims))
68112
NS = shape_function_matrix(MS, Xg, dims=MT.dims)
69113
NT = shape_function_matrix(MT, Xg, dims=MT.dims)
70114
A = (NT.T @ Ig @ NT).asformat('csc')
@@ -75,60 +119,79 @@ def __init__(self, MD: pbat.fem.Mesh, MS: pbat.fem.Mesh, MT: pbat.fem.Mesh, H):
75119
self.A = A
76120
self.P = P
77121
self.U = laplacian_energy(MT, dims=MT.dims)
122+
self.GS = gradient_operator(MS, Xg, dims=MT.dims)
123+
self.GT = gradient_operator(MT, Xg, dims=MT.dims)
124+
self.IG = kron(kron(eye(MT.dims), diags(wg)), eye(MT.dims))
78125
self.H = H
79126

80127
def __matmul__(self, X):
81128
pass
82129

83130

84131
class CholFemFunctionTransferOperator(BaseFemFunctionTransferOperator):
85-
def __init__(self, MD: pbat.fem.Mesh, MS: pbat.fem.Mesh, MT: pbat.fem.Mesh, H, lreg=5, hreg=1):
132+
def __init__(self, MD: pbat.fem.Mesh, MS: pbat.fem.Mesh, MT: pbat.fem.Mesh, H, lreg=5, hreg=1, greg=1, Y=1e6, nu=0.45, hxreg=1e-4, rho=1e3):
86133
super().__init__(MD, MS, MT, H)
87134
n = self.A.shape[0]
88-
A = self.A + lreg*self.U + hreg*self.H
89-
lmin, lmax = min_max_eigs(A, n=1)
90-
tau = 0.
91-
if lmin[0] <= 0:
92-
# Regularize A (due to positive semi-definiteness)
93-
tau = abs(lmin[0]) + 1e-10
94-
tau = sp.sparse.diags(np.full(n, tau))
95-
AR = A + tau
96-
solver = pbat.math.linalg.SolverBackend.Eigen
97-
self.Ainv = pbat.math.linalg.chol(AR, solver=solver)
98-
self.Ainv.compute(AR)
99-
100-
def __matmul__(self, B):
101-
B = self.P @ B
102-
X = self.Ainv.solve(B)
103-
return X
135+
energy = pbat.fem.HyperElasticEnergy.StableNeoHookean
136+
self.hep, self.detJeU, self.GNeU = pbat.fem.hyper_elastic_potential(
137+
MT, Y, nu, energy=energy)
138+
self.MT = MT
139+
self.rho = rho
140+
self.M = self.NT.T @ (rho * self.Ig) @ self.NT # + lreg*self.U + hreg*self.H + greg * self.GT.T @ self.IG @ self.GT
141+
self.P = self.NT.T @ (rho * self.Ig) @ self.NS
142+
self.hxreg = hxreg
143+
# lmin, lmax = min_max_eigs(A, n=1)
144+
# tau = 0.
145+
# if lmin[0] <= 0:
146+
# # Regularize A (due to positive semi-definiteness)
147+
# print("WARNING: A is semi-definite")
148+
# tau = abs(lmin[0]) + 1e-10
149+
# tau = sp.sparse.diags(np.full(n, tau))
150+
# AR = A + tau
151+
# solver = pbat.math.linalg.SolverBackend.Eigen
152+
# self.Ainv = pbat.math.linalg.chol(AR, solver=solver)
153+
# self.Ainv.compute(AR)
154+
self.uk = np.zeros(math.prod(self.MT.X.shape))
155+
156+
def __matmul__(self, u):
157+
self.uk = np.zeros(math.prod(self.MT.X.shape))
158+
xr = self.MT.X.reshape(math.prod(self.MT.X.shape), order='f')
159+
for k in range(1):
160+
x = xr + self.uk
161+
self.hep.compute_element_elasticity(x)
162+
duku = self.NT @ self.uk - self.NS @ u
163+
E = 0.5 * duku.T @ (self.rho * self.Ig) @ duku + self.hxreg * self.hep.eval()
164+
H = self.M + self.hxreg * self.hep.hessian()
165+
solver = pbat.math.linalg.SolverBackend.Eigen
166+
Hinv = pbat.math.linalg.ldlt(H, solver=solver)
167+
Hinv.compute(H)
168+
# B = self.P @ u + self.GT.T @ self.IG @ self.GS @ u
169+
ng = self.P @ u - self.hxreg * self.hep.gradient()
170+
r = ng.T @ ng
171+
if r < 1e-3:
172+
break
173+
# X = self.Ainv.solve(B)
174+
du = Hinv.solve(ng).squeeze()
175+
self.uk += du
176+
return self.uk
104177

105178

106179
class RankKApproximateFemFunctionTransferOperator(BaseFemFunctionTransferOperator):
107-
def __init__(self, MD: pbat.fem.Mesh, MS: pbat.fem.Mesh, MT: pbat.fem.Mesh, H, lreg=5, hreg=1, modes=30):
180+
def __init__(self, MD: pbat.fem.Mesh, MS: pbat.fem.Mesh, MT: pbat.fem.Mesh, H, lreg=5, hreg=1, greg=1, modes=30):
108181
super().__init__(MD, MS, MT, H)
109-
A = self.A + lreg*self.U + hreg*self.H
182+
A = self.A + lreg*self.U + hreg*self.H + greg * self.GT.T @ self.IG @ self.GT
110183
l, V = sp.sparse.linalg.eigsh(A, k=modes, sigma=1e-5, which='LM')
111184
keep = np.nonzero(l > 1e-5)[0]
112185
self.l, self.V = l[keep], V[:, keep]
113186

114187
def __matmul__(self, B):
115-
B = self.P @ B
188+
B = self.P @ B + self.GT.T @ self.IG @ self.GS @ B
116189
B = self.V.T @ B
117190
B = B @ sp.sparse.diags(1 / self.l)
118191
X = self.V @ B
119192
return X
120193

121194

122-
def rest_pose_hessian(mesh, Y, nu):
123-
x = mesh.X.reshape(math.prod(mesh.X.shape), order='f')
124-
energy = pbat.fem.HyperElasticEnergy.StableNeoHookean
125-
hep, detJeU, GNeU = pbat.fem.hyper_elastic_potential(
126-
mesh, Y, nu, energy=energy)
127-
hep.compute_element_elasticity(x)
128-
HU = hep.hessian()
129-
return HU
130-
131-
132195
def linear_elastic_deformation_modes(mesh, rho, Y, nu, modes=30, sigma=-1e-5):
133196
M, detJeM = pbat.fem.mass_matrix(mesh, rho=rho)
134197
HU = rest_pose_hessian(mesh, Y, nu)
@@ -185,52 +248,59 @@ def signal(w: float, v: np.ndarray, t: float, c: float, k: float):
185248
w, L = linear_elastic_deformation_modes(
186249
mesh, args.rho, args.Y, args.nu, args.modes)
187250
HC = rest_pose_hessian(cmesh, args.Y, args.nu)
188-
lreg, hreg = 0, 1e-2
251+
lreg, hreg, greg, hxreg = 0, 0, 0, 1e-4
189252
Fldl = CholFemFunctionTransferOperator(
190-
mesh, mesh, cmesh, HC, lreg=lreg, hreg=hreg)
191-
Krestrict = 30
192-
Frank = RankKApproximateFemFunctionTransferOperator(
193-
mesh, mesh, cmesh, HC, lreg=lreg, hreg=hreg, modes=Krestrict)
253+
mesh, mesh, cmesh, HC, lreg=lreg, hreg=hreg, greg=greg, hxreg=hxreg, rho=args.rho)
254+
# Krestrict = 30
255+
# Frank = RankKApproximateFemFunctionTransferOperator(
256+
# mesh, mesh, cmesh, HC, lreg=lreg, hreg=hreg, greg=greg, modes=Krestrict)
194257

195258
ps.set_up_dir("z_up")
196259
ps.set_front_dir("neg_y_front")
197260
ps.set_ground_plane_mode("shadow_only")
198261
ps.init()
199262
vm = ps.register_surface_mesh("model", V, F)
200263
ldlvm = ps.register_surface_mesh("LDL cage", CV, CF)
201-
rankvm = ps.register_surface_mesh("Rank K cage", CV, CF)
264+
# rankvm = ps.register_surface_mesh("Rank K cage", CV, CF)
202265
mode = 6
203266
t0 = time.time()
204267
t = 0
205268
c = 0.15
206269
k = 0.05
207270
theta = 0
208271
dtheta = np.pi/120
272+
animate = False
273+
step = False
209274

210275
def callback():
211276
global mode, c, k, theta, dtheta
277+
global animate, step
278+
212279
changed, mode = imgui.InputInt("Mode", mode)
213280
changed, c = imgui.InputFloat("Wave amplitude", c)
214281
changed, k = imgui.InputFloat("Wave frequency", k)
215-
216-
t = time.time() - t0
217-
218-
R = sp.spatial.transform.Rotation.from_quat(
219-
[0, np.sin(theta/2), 0, np.cos(theta/4)]).as_matrix()
220-
X = (V - V.mean(axis=0)) @ R.T + V.mean(axis=0)
221-
uf = signal(w[mode], L[:, mode], t, c, k)
222-
ur = (X - V).flatten(order="C")
223-
ut = 1
224-
u = uf + ur + ut
225-
XCldl = CV + (Fldl @ u).reshape(CV.shape)
226-
XCrank = CV + (Frank @ u).reshape(CV.shape)
227-
vm.update_vertex_positions(V + (uf + ur).reshape(V.shape))
228-
ldlvm.update_vertex_positions(XCldl)
229-
rankvm.update_vertex_positions(XCrank)
230-
231-
theta += dtheta
232-
if theta > 2*np.pi:
233-
theta = 0
282+
changed, animate = imgui.Checkbox("animate", animate)
283+
step = imgui.Button("step")
284+
285+
if animate or step:
286+
t = time.time() - t0
287+
288+
R = sp.spatial.transform.Rotation.from_quat(
289+
[0, np.sin(theta/2), 0, np.cos(theta/4)]).as_matrix()
290+
X = (V - V.mean(axis=0)) @ R.T + V.mean(axis=0)
291+
uf = signal(w[mode], L[:, mode], t, c, k)
292+
ur = (X - V).flatten(order="C")
293+
ut = 1
294+
u = uf + ur + ut
295+
XCldl = CV + (Fldl @ u).reshape(CV.shape)
296+
# XCrank = CV + (Frank @ u).reshape(CV.shape)
297+
vm.update_vertex_positions(V + (uf + ur).reshape(V.shape))
298+
ldlvm.update_vertex_positions(XCldl)
299+
# rankvm.update_vertex_positions(XCrank)
300+
301+
theta += dtheta
302+
if theta > 2*np.pi:
303+
theta = 0
234304

235305
ps.set_user_callback(callback)
236306
ps.show()

0 commit comments

Comments
 (0)