Coverage for pyxfem/solidmechanics/hyperelasticity_solver.py: 48%
215 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-24 07:17 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Created on April 10th 07:43:08 2024
6@author: nchevaug
7"""
8# pylint: disable=invalid-name
9from warnings import warn
10import numpy as np
11import numpy.typing as npt
12import scipy as sp
13from numba import njit, prange
14import gmsh
15from pyxfem.core.config import npfloat, FloatArray
16import pyxfem.core.function_spaces as fs
17import pyxfem.core.quadrature as quad
18import pyxfem.core.mesh_simplex_one_level as sm
19import pyxfem.postpro.gmsh_post_pro_interface as gp
20from pyxfem.core.nonlinear_solver import newton
21import pyxfem.solidmechanics.hyperelasticity as hy
22import pyxfem.core.sparse_tools as st
23import pyxfem.solidmechanics.boundary_conditions as bc
24from pyxfem.solidmechanics.largestrain import PKI_to_Cauchy
25from pyxfem.postpro.ioh5 import ParaViewXDMFWriter
27# def computeNeumann(space, physName2Neumann):
28# ''' return the vector F such as F.T.dot(U) = int_neumann ( t.u_h) dGamma_h '''
29# F = np.zeros((space.scalarSize(), space.vecDim))
30# edgequad = quad.Edge_gauss(0)
31# for physName, dir2Value in physName2Neumann.items():
32# eids, ___ = space.m.getEdgeGroupe(physName)
33# N = space.evalOpOnEdges(eids, edgequad.s)
34# xy = space.m.getEdgesCog(eids)
35# l = space.m.getEdgesLength(eids)
36# t = np.zeros((eids.size, 2))
37# for directionName, value in dir2Value.items():
38# direction = dirName2dirVal[directionName]
39# if direction is None:
40# funName = computeNeumann.__name__
41# raise ValueError('directionName '+directionName+' unknown in '+funName)
42# t[:, direction] = l*setValues(xy, value)
43# F += (edgequad.w*N.T.dot(t.flatten())).reshape((-1, 2))
44# return F.flatten()
46# def computeVolumeLoad(space, source):
47# """ compute the vector F such as F.U = int_Omega f.u_h dOmega_h """
48# assert isinstance(source, np.ndarray)
49# quadSource = quad.T3_gauss(0)
50# N = space.operator_dof2val_tri(quadSource.uv)
51# J = space.get_tris_mapping().J
52# wSource = quadSource.w[0]*source[np.newaxis, :]*J[:, np.newaxis]
53# Fint = N.T.dot(wSource.flatten())
54# return Fint
56class HyperElasticitySolver():
57 """ A class to set up and compute a 2d linear Hyper Elastic problem with FEM method """
58 def __init__(self, mesh:sm.SMesh, matlaw:hy.HyperElasticMaterial,
59 physname2dirichlet:dict|None= None,
60 physname2neumann:dict|None= None,
61 source=np.array([0., 0.]),
62 space_constructor=fs.FEMSpaceP1):
63 self.mesh = mesh
64 self.law = matlaw
65 self.quad_bilinform = quad.T3_gauss(0)
66 self.quad_corners = quad.T3_nodes()
67 self.space = space_constructor("DisplacementSpace", self.mesh, vecdim=2)
68 self.ndofs = self.space.size()
69 self.B:sp.sparse.spmatrix = self.space.operator_dof2grad_tri(
70 self.quad_bilinform.uv).tocsc()
71 self.B_csc = self.B
72 self.B_csc_T = self.B_csc.T
74 self.Bb = self.space.operator_dof2grad_tri(
75 self.quad_bilinform.uv, raw_data = True).squeeze()
76 self.BbT = self.Bb.swapaxes(-1,-2)
77 self.NatVertices = self.space.operator_dof2val_tri(
78 self.quad_corners.uv)
79 self.physname2dirichlet = {} if physname2dirichlet is None else physname2dirichlet
80 self.physname2neumann = {} if physname2neumann is None else physname2neumann
81 self.source = source
82 gname2dofid, gname2dofxy = bc.get_dofs_and_interpolation_points(
83 mesh, self.physname2dirichlet.keys(), self.space)
84 self.free_dofs, self.fixed_dofs, self.fixed_values = bc.compute_dirichlet(
85 self.space, self.physname2dirichlet, gname2dofid, gname2dofxy)
86 # Next is a tentative to imporve the speed of assembly, see
87 # _compute_internal_stiffness_blocs
89 def Ksparse_indexes():
90 scal_t2dofs = self.space.a2dofs
91 vecdim = self.space.vecdim
92 # Interleave DOFs: [ux0, uy0, ux1, uy1, ux2, uy2]
93 vec_t2dofs = (vecdim * scal_t2dofs[..., np.newaxis] + np.arange(vecdim)).reshape(scal_t2dofs.shape[0], -1)
94 row = np.repeat(vec_t2dofs, vec_t2dofs.shape[1], axis=1).reshape(-1)
95 col = np.repeat(vec_t2dofs, vec_t2dofs.shape[1], axis=0).reshape(-1)
96 return (row, col)
100 self.Kindexes = Ksparse_indexes()
101 # --- PRECOMPUTE CSR SPARSITY MAPPING ---
102 row, col = self.Kindexes
104 # Create a dummy array to let Scipy compute the compressed indices once
105 dummy_csr = sp.sparse.csr_array((np.ones_like(row), (row, col)), shape=(self.ndofs, self.ndofs))
106 self.K_indptr = dummy_csr.indptr
107 self.K_indices = dummy_csr.indices
109 # Create unique 1D flat indices for both the original blocks and the compressed CSR format
110 flat_rc = row.astype(np.int64) * self.ndofs + col
111 csr_rows = np.repeat(np.arange(self.ndofs), np.diff(self.K_indptr))
112 csr_flat_rc = csr_rows.astype(np.int64) * self.ndofs + self.K_indices
114 # Map each element in K_blocs to its final index in the compressed data array
115 self.K_mapping = np.searchsorted(csr_flat_rc, flat_rc)
117 # this need to be implemented latter
118 # self.Fneumann = computeNeumann(self.space, self.physname2neumann)
119 # self.Fint = computeVolumeLoad(self.space, source)
121 def compute_energy(self, U:FloatArray)-> float:
122 """ return the stored energy associated to nodal displacement e"""
124 F = self.compute_F(U)
125 J0 = self.space.get_tris_mapping().J
126 e = self.law.potentialF(F)
127 return np.sum(J0*self.quad_bilinform.w[0]*e)
129 def compute_internal_forces(self, U:FloatArray)->FloatArray:
130 """ compute the internal force vecor F_int, such as F_int.delta_u = delta w_int """
132 F = self.compute_F(U)
133 J0 = self.space.get_tris_mapping().J
134 w = J0*self.quad_bilinform.w[0]
135 wPKI = w[..., np.newaxis, np.newaxis]*self.law.PKI(F)
136 internal_forces = self.B.T@wPKI.flatten()
137 return internal_forces
139 def compute_external_forces(self, U: FloatArray, lamb:npfloat):
140 warn("External Forces not implemented", category = UserWarning, stacklevel =2)
141 return np.zeros(self.ndofs)
143 def compute_internal_stiffness(self, U: FloatArray) -> sp.sparse.csr_array:
144 F = self.compute_F(U)# --- PRECOMPUTE CSR SPARSITY MAPPING ---
145 row, col = self.Kindexes
147 # Create a dummy array to let Scipy compute the compressed indices once
148 dummy_csr = sp.sparse.csr_array((np.ones_like(row), (row, col)), shape=(self.ndofs, self.ndofs))
149 self.K_indptr = dummy_csr.indptr
150 self.K_indices = dummy_csr.indices
152 # Create unique 1D flat indices for both the original blocks and the compressed CSR format
153 flat_rc = row.astype(np.int64) * self.ndofs + col
154 csr_rows = np.repeat(np.arange(self.ndofs), np.diff(self.K_indptr))
155 csr_flat_rc = csr_rows.astype(np.int64) * self.ndofs + self.K_indices
157 # Map each element in K_blocs to its final index in the compressed data array
158 self.K_mapping = np.searchsorted(csr_flat_rc, flat_rc)
159 J0 = self.space.get_tris_mapping().J
160 w = J0 * self.quad_bilinform.w[0]
162 # Calculate the 4x4 blocks: shape is (91514, 4, 4)
163 blocks = w[..., np.newaxis, np.newaxis] * (self.law.dPKIdF(F).reshape(-1, 4, 4))
165 # Build A directly as a Block Sparse Row (BSR) matrix
166 n_blocks = blocks.shape[0]
167 indptr = np.arange(n_blocks + 1, dtype=np.int32)
168 indices = np.arange(n_blocks, dtype=np.int32)
169 A_bsr = sp.sparse.bsr_matrix((blocks, indices, indptr), shape=(n_blocks * 4, n_blocks * 4))
171 # Perform the multiplication with the cached matrices
172 # A_bsr @ B_csc -> convert to CSC for the final left-multiplication
173 C = (A_bsr @ self.B_csc).tocsc()
175 return self.B_csc_T @ C
177 @staticmethod
178 @njit(parallel=True, fastmath=True, cache=True)
179 def _fast_k_blocs(w, dPKIdF, Bb, K_blocs):
180 N = dPKIdF.shape[0]
181 # Bb is expected to be (N, 4, 6), dPKIdF is (N, 4, 4)
182 for n in prange(N):
183 wn = w[n]
184 for i in range(6):
185 for j in range(6):
186 val = 0.0
187 # Standard matrix multiplication unrolled
188 for k in range(4):
189 for l in range(4):
190 val += Bb[n, k, i] * dPKIdF[n, k, l] * Bb[n, l, j]
191 # Write directly to the flattened array structure
192 K_blocs[n * 36 + i * 6 + j] = val * wn
194 def compute_internal_stiffness_blocs_v0(self, U:FloatArray)->sp.sparse.csr_array:
195 F = self.compute_F(U)
196 J0 = self.space.get_tris_mapping().J
197 w = J0*self.quad_bilinform.w[0]
198 wdPKIdF = w[..., np.newaxis, np.newaxis]*(self.law.dPKIdF(F).reshape(-1, 4, 4))
200 Bb = self.Bb
201 BbT = self.BbT
202 K_blocs = BbT@wdPKIdF@Bb
203 K_blocs.shape = K_blocs.size
205 # --- FAST ASSEMBLY ---
206 # Accumulate overlapping block values instantly using the precomputed mapping
207 K_data = np.bincount(self.K_mapping, weights=K_blocs, minlength=len(self.K_indices))
209 # Construct the CSR array directly from pre-sorted arrays (bypasses _coo_to_compressed entirely)
210 K_csr = sp.sparse.csr_array((K_data, self.K_indices, self.K_indptr), shape=(self.ndofs, self.ndofs), copy=False)
211 return K_csr
213 def compute_internal_stiffness_blocs(self, U:FloatArray)->sp.sparse.csr_array:
214 F = self.compute_F(U)
215 J0 = self.space.get_tris_mapping().J
216 w = J0 * self.quad_bilinform.w[0]
217 # Get raw material tangent and reshape strictly to (N, 4, 4)
218 dPKIdF = self.law.dPKIdF(F).reshape(-1, 4, 4)
219 # 1. Pre-allocate the exact output size for the flattened 6x6 blocks
220 K_blocs = np.empty(w.shape[0] * 36, dtype=F.dtype)
221 # 2. Execute parallel C-loop (Bypasses all NumPy batched @ operations)
222 self._fast_k_blocs(w, dPKIdF, self.Bb, K_blocs)
224 # 3. Apply mask to only accumulate the upper triangle
225 #K_blocs_upper = K_blocs[self.K_upper_mask]
227 # Fast Assembly
228 K_data = np.bincount(self.K_mapping, weights=K_blocs, minlength=len(self.K_indices))
230 # Construct the CSR array directly
231 K_csr = sp.sparse.csr_array((K_data, self.K_indices, self.K_indptr), shape=(self.ndofs, self.ndofs), copy=False)
232 return K_csr
234 def _compute_Hsparse(self, U:npt.NDArray)->npt.NDArray:
235 """ compute H, the gradient of the displacment field at each gauss point """
237 Ushape = U.shape
238 nelem = self.mesh.ntris()
239 ndofs = self.space.size()
240 U = U.reshape(ndofs, copy = False)
241 vecdim = self.space.vecdim
242 H = self.B@U
243 H.shape = (nelem, vecdim, vecdim)
244 U = U.reshape(Ushape, copy = False)
245 return H
247 def _compute_Hblocs(self, U:npt.NDArray)->npt.NDArray:
248 """ compute H, the gradient of the displacment field at each gauss point
249 This is a tentative version to speed up the computation compared to compute Hsparse
250 Not effective as of now
251 """
254 nelem = self.mesh.ntris()
255 ndofs = self.space.size()
256 vecdim = self.space.vecdim
257 U = U.reshape((ndofs,), copy = False)
258 Bb = self.Bb
259 Ushape = U.shape
260 nelem = Bb.shape[0]
261 U = U.reshape((ndofs//vecdim, vecdim), copy =False)
263 Utris = U[self.mesh.t2v()]
264 Utris = Utris.reshape((nelem, 3*vecdim, 1), copy = False)
265 H = Bb@Utris
266 H = H.reshape(nelem, 2, 2)
267 U = U.reshape(Ushape)
268 return H
270 def compute_F(self, U:npt.NDArray)->npt.NDArray:
271 """ compute F, the gradient of the position field at each gauss point """
272 return self._compute_Hsparse(U) + np.eye(2)[np.newaxis]
274 def predictor_corrector(self, U0: FloatArray, loadfac0:npfloat, loadinc:npfloat, **newton_kwargs)-> FloatArray:
275 """ solve the non linear system """
276 # Default options for newton, with itmax mapped to nitermax
277 newton_opts = {
278 "linsolver": sp.sparse.linalg.spsolve,
279 "rtol": 1.0e-6,
280 "atol": 1.0e-6,
281 "nitermax": 10,
282 "fail_on_increasing_res": False,
283 "verbose": True,
284 }
285 # Allow callers to override defaults or pass other parameters (norm, isub, etc.)
286 newton_opts.update(newton_kwargs)
287 #atol = newton_opts["atol"]
289 free_dofs = self.free_dofs
290 diri_dofs = self.fixed_dofs
291 #Fext = self.compute_external_forces(U0, loadfac0)
292 #Fint = self.compute_internal_forces(U0)
293 #check residual at U0
294 # R = Fint - Fext
295 # Rf = R[free_dofs]
296 # nRf = np.linalg.norm(Rf)
297 # print(f"Starting res {nRf:.3e}")
298 # assert np.linalg.norm(Rf) < atol
299 # predictor
300 K = self.compute_internal_stiffness_blocs(U0)
301 Kfa = K[free_dofs, :]
302 Kff = Kfa[:,free_dofs]
303 U0d = self.fixed_values
304 # U0d_check = U0[diri_dofs]
305 # assert(np.linalg.norm(U0d - U0d_check) < atol)
306 FUD_f = Kfa[:,diri_dofs]@U0d
307 Flamb1_f = self.compute_external_forces(U0, npfloat(1.))[free_dofs]
308 # assert(np.linalg.norm(Flamb1_f) == 0.)
309 linsolve = newton_opts["linsolver"]
310 DeltaUpred_f = linsolve(Kff, loadinc*(-FUD_f+Flamb1_f ))
311 Upred = U0
312 Upred[diri_dofs] = (loadfac0+loadinc)*self.fixed_values
313 Upred[free_dofs] += DeltaUpred_f
314 U = Upred.copy()
316 def fun(x):
317 U[self.free_dofs] = x
318 R = self.compute_internal_forces(U)
319 return R[self.free_dofs]
321 def jac(x):
322 U[self.free_dofs] = x
323 K = self.compute_internal_stiffness_blocs(U)
324 K_free = K[:,self.free_dofs][self.free_dofs, :]
325 return K_free
326 res = newton(fun, x0 = Upred[free_dofs], jac = jac, **newton_opts)
327 U[self.free_dofs] = res.x
328 return U
330def solve_loadstep(mesh:sm.SMesh, dirichlet_bc, loadsteps, \
331 savedir:None|str = None, popup_gmsh = False, paraviewout = True, **newton_kwargs):
332 """ example main program """
333 law = hy.CompressibleNeoHookeen(E=1., nu=0.3, hyp=hy.PLANESTRESS)
334 pb = HyperElasticitySolver(mesh= mesh, matlaw = law, physname2dirichlet=dirichlet_bc)
335 U= None
336 inc = 0
337 if paraviewout:
338 writer = ParaViewXDMFWriter("output/hypersolveres", mesh.t2v(), fields_config={
339 "Cauchy_Stress": {"type": "Tensor", "center": "Cell"}
340 },
341 overwrite=True)
343 gmsh.initialize()
344 gmsh.option.setNumber('PostProcessing.AnimationCycle', 1)
345 if savedir is not None:
346 gmsh.write(savedir+"post_options.opt")
347 Xtri= mesh.t2v_coord()
348 X = mesh.vcoord()
349 loadsteps = np.array(loadsteps)
350 loadincs = loadsteps[1:] - loadsteps[:-1]
351 U = np.zeros(pb.ndofs)
353 def post(U, basename):
354 u = U.reshape(-1,2)
355 utri = u[mesh.t2v()]
356 x = X.copy()
357 x[mesh.used_v()] +=u
358 xtri = Xtri + utri
359 F = pb.compute_F(U).reshape(-1, 2, 2) # pylint: disable=C0103
360 P = law.PKI(F) # pylint: disable=C0103
361 Cauchy = PKI_to_Cauchy(PKI = P, F = F) # pylint: disable=C0103
362 # if paraviewout:
363 # writer.write_step(i, loadfac0+loadinc, x, Cauchy_Stress = Cauchy)
364 # if law.hyp == hy.PLANESTRESS:
365 # lambda3 = law.l3_planestress(np.linalg.det(F))
366 # vL3 = gp.listPlotFieldTri(x, lambda3, #pylint: disable=C0103
367 # f"lambda3_{inc:04d}", P0 = True)
368 # vF = gp.listPlotFieldTri(x, F, f"F_{inc:04d}", P0 = True)
369 vP = gp.listPlotFieldTri(xtri, P, f"{basename}_{inc:04d}", P0 = True)
370 if savedir is not None :
371 #gmsh.view.write(vF, savedir+f"F_{inc:04d}"+".pos")
372 gmsh.view.write(vP, savedir+f"P_{inc:04d}"+".pos")
373 # if law.hyp == hy.PLANESTRESS:
374 # gmsh.view.write(vL3, savedir+f"L3_{inc:04d}"+".pos")
378 for i, (loadfac0, loadinc) in enumerate( zip(loadsteps[:-1], loadincs)):
379 print(f"# increment {inc:5d}, loading factor 0 {loadfac0:.3e}, loadinc {loadinc:.3e}")
380 U = pb.predictor_corrector(U0= U, loadfac0= loadfac0, loadinc= loadinc, **newton_kwargs)
381 post(U, "Sol")
383 inc += 1
384 if paraviewout:
385 writer.finalize()
386 if popup_gmsh:
387 gmsh.fltk.run()
388 gmsh.finalize()