#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on April 10th 07:43:08 2024
@author: nchevaug
"""
# pylint: disable=invalid-name
from warnings import warn
import numpy as np
import numpy.typing as npt
import scipy as sp
from numba import njit, prange
import gmsh
from pyxfem.core.config import npfloat, FloatArray
import pyxfem.core.function_spaces as fs
import pyxfem.core.quadrature as quad
import pyxfem.core.mesh_simplex_one_level as sm
import pyxfem.postpro.gmsh_post_pro_interface as gp
from pyxfem.core.nonlinear_solver import newton
import pyxfem.solidmechanics.hyperelasticity as hy
import pyxfem.core.sparse_tools as st
import pyxfem.solidmechanics.boundary_conditions as bc
from pyxfem.solidmechanics.largestrain import PKI_to_Cauchy
from pyxfem.postpro.ioh5 import ParaViewXDMFWriter
# def computeNeumann(space, physName2Neumann):
# ''' return the vector F such as F.T.dot(U) = int_neumann ( t.u_h) dGamma_h '''
# F = np.zeros((space.scalarSize(), space.vecDim))
# edgequad = quad.Edge_gauss(0)
# for physName, dir2Value in physName2Neumann.items():
# eids, ___ = space.m.getEdgeGroupe(physName)
# N = space.evalOpOnEdges(eids, edgequad.s)
# xy = space.m.getEdgesCog(eids)
# l = space.m.getEdgesLength(eids)
# t = np.zeros((eids.size, 2))
# for directionName, value in dir2Value.items():
# direction = dirName2dirVal[directionName]
# if direction is None:
# funName = computeNeumann.__name__
# raise ValueError('directionName '+directionName+' unknown in '+funName)
# t[:, direction] = l*setValues(xy, value)
# F += (edgequad.w*N.T.dot(t.flatten())).reshape((-1, 2))
# return F.flatten()
# def computeVolumeLoad(space, source):
# """ compute the vector F such as F.U = int_Omega f.u_h dOmega_h """
# assert isinstance(source, np.ndarray)
# quadSource = quad.T3_gauss(0)
# N = space.operator_dof2val_tri(quadSource.uv)
# J = space.get_tris_mapping().J
# wSource = quadSource.w[0]*source[np.newaxis, :]*J[:, np.newaxis]
# Fint = N.T.dot(wSource.flatten())
# return Fint
[docs]
class HyperElasticitySolver():
""" A class to set up and compute a 2d linear Hyper Elastic problem with FEM method """
def __init__(self, mesh:sm.SMesh, matlaw:hy.HyperElasticMaterial,
physname2dirichlet:dict|None= None,
physname2neumann:dict|None= None,
source=np.array([0., 0.]),
space_constructor=fs.FEMSpaceP1):
self.mesh = mesh
self.law = matlaw
self.quad_bilinform = quad.T3_gauss(0)
self.quad_corners = quad.T3_nodes()
self.space = space_constructor("DisplacementSpace", self.mesh, vecdim=2)
self.ndofs = self.space.size()
self.B:sp.sparse.spmatrix = self.space.operator_dof2grad_tri(
self.quad_bilinform.uv).tocsc()
self.B_csc = self.B
self.B_csc_T = self.B_csc.T
self.Bb = self.space.operator_dof2grad_tri(
self.quad_bilinform.uv, raw_data = True).squeeze()
self.BbT = self.Bb.swapaxes(-1,-2)
self.NatVertices = self.space.operator_dof2val_tri(
self.quad_corners.uv)
self.physname2dirichlet = {} if physname2dirichlet is None else physname2dirichlet
self.physname2neumann = {} if physname2neumann is None else physname2neumann
self.source = source
gname2dofid, gname2dofxy = bc.get_dofs_and_interpolation_points(
mesh, self.physname2dirichlet.keys(), self.space)
self.free_dofs, self.fixed_dofs, self.fixed_values = bc.compute_dirichlet(
self.space, self.physname2dirichlet, gname2dofid, gname2dofxy)
# Next is a tentative to imporve the speed of assembly, see
# _compute_internal_stiffness_blocs
def Ksparse_indexes():
scal_t2dofs = self.space.a2dofs
vecdim = self.space.vecdim
# Interleave DOFs: [ux0, uy0, ux1, uy1, ux2, uy2]
vec_t2dofs = (vecdim * scal_t2dofs[..., np.newaxis] + np.arange(vecdim)).reshape(scal_t2dofs.shape[0], -1)
row = np.repeat(vec_t2dofs, vec_t2dofs.shape[1], axis=1).reshape(-1)
col = np.repeat(vec_t2dofs, vec_t2dofs.shape[1], axis=0).reshape(-1)
return (row, col)
self.Kindexes = Ksparse_indexes()
# --- PRECOMPUTE CSR SPARSITY MAPPING ---
row, col = self.Kindexes
# Create a dummy array to let Scipy compute the compressed indices once
dummy_csr = sp.sparse.csr_array((np.ones_like(row), (row, col)), shape=(self.ndofs, self.ndofs))
self.K_indptr = dummy_csr.indptr
self.K_indices = dummy_csr.indices
# Create unique 1D flat indices for both the original blocks and the compressed CSR format
flat_rc = row.astype(np.int64) * self.ndofs + col
csr_rows = np.repeat(np.arange(self.ndofs), np.diff(self.K_indptr))
csr_flat_rc = csr_rows.astype(np.int64) * self.ndofs + self.K_indices
# Map each element in K_blocs to its final index in the compressed data array
self.K_mapping = np.searchsorted(csr_flat_rc, flat_rc)
# this need to be implemented latter
# self.Fneumann = computeNeumann(self.space, self.physname2neumann)
# self.Fint = computeVolumeLoad(self.space, source)
[docs]
def compute_energy(self, U:FloatArray)-> float:
""" return the stored energy associated to nodal displacement e"""
F = self.compute_F(U)
J0 = self.space.get_tris_mapping().J
e = self.law.potentialF(F)
return np.sum(J0*self.quad_bilinform.w[0]*e)
[docs]
def compute_internal_forces(self, U:FloatArray)->FloatArray:
""" compute the internal force vecor F_int, such as F_int.delta_u = delta w_int """
F = self.compute_F(U)
J0 = self.space.get_tris_mapping().J
w = J0*self.quad_bilinform.w[0]
wPKI = w[..., np.newaxis, np.newaxis]*self.law.PKI(F)
internal_forces = self.B.T@wPKI.flatten()
return internal_forces
[docs]
def compute_external_forces(self, U: FloatArray, lamb:npfloat):
warn("External Forces not implemented", category = UserWarning, stacklevel =2)
return np.zeros(self.ndofs)
[docs]
def compute_internal_stiffness(self, U: FloatArray) -> sp.sparse.csr_array:
F = self.compute_F(U)# --- PRECOMPUTE CSR SPARSITY MAPPING ---
row, col = self.Kindexes
# Create a dummy array to let Scipy compute the compressed indices once
dummy_csr = sp.sparse.csr_array((np.ones_like(row), (row, col)), shape=(self.ndofs, self.ndofs))
self.K_indptr = dummy_csr.indptr
self.K_indices = dummy_csr.indices
# Create unique 1D flat indices for both the original blocks and the compressed CSR format
flat_rc = row.astype(np.int64) * self.ndofs + col
csr_rows = np.repeat(np.arange(self.ndofs), np.diff(self.K_indptr))
csr_flat_rc = csr_rows.astype(np.int64) * self.ndofs + self.K_indices
# Map each element in K_blocs to its final index in the compressed data array
self.K_mapping = np.searchsorted(csr_flat_rc, flat_rc)
J0 = self.space.get_tris_mapping().J
w = J0 * self.quad_bilinform.w[0]
# Calculate the 4x4 blocks: shape is (91514, 4, 4)
blocks = w[..., np.newaxis, np.newaxis] * (self.law.dPKIdF(F).reshape(-1, 4, 4))
# Build A directly as a Block Sparse Row (BSR) matrix
n_blocks = blocks.shape[0]
indptr = np.arange(n_blocks + 1, dtype=np.int32)
indices = np.arange(n_blocks, dtype=np.int32)
A_bsr = sp.sparse.bsr_matrix((blocks, indices, indptr), shape=(n_blocks * 4, n_blocks * 4))
# Perform the multiplication with the cached matrices
# A_bsr @ B_csc -> convert to CSC for the final left-multiplication
C = (A_bsr @ self.B_csc).tocsc()
return self.B_csc_T @ C
@staticmethod
@njit(parallel=True, fastmath=True, cache=True)
def _fast_k_blocs(w, dPKIdF, Bb, K_blocs):
N = dPKIdF.shape[0]
# Bb is expected to be (N, 4, 6), dPKIdF is (N, 4, 4)
for n in prange(N):
wn = w[n]
for i in range(6):
for j in range(6):
val = 0.0
# Standard matrix multiplication unrolled
for k in range(4):
for l in range(4):
val += Bb[n, k, i] * dPKIdF[n, k, l] * Bb[n, l, j]
# Write directly to the flattened array structure
K_blocs[n * 36 + i * 6 + j] = val * wn
[docs]
def compute_internal_stiffness_blocs_v0(self, U:FloatArray)->sp.sparse.csr_array:
F = self.compute_F(U)
J0 = self.space.get_tris_mapping().J
w = J0*self.quad_bilinform.w[0]
wdPKIdF = w[..., np.newaxis, np.newaxis]*(self.law.dPKIdF(F).reshape(-1, 4, 4))
Bb = self.Bb
BbT = self.BbT
K_blocs = BbT@wdPKIdF@Bb
K_blocs.shape = K_blocs.size
# --- FAST ASSEMBLY ---
# Accumulate overlapping block values instantly using the precomputed mapping
K_data = np.bincount(self.K_mapping, weights=K_blocs, minlength=len(self.K_indices))
# Construct the CSR array directly from pre-sorted arrays (bypasses _coo_to_compressed entirely)
K_csr = sp.sparse.csr_array((K_data, self.K_indices, self.K_indptr), shape=(self.ndofs, self.ndofs), copy=False)
return K_csr
[docs]
def compute_internal_stiffness_blocs(self, U:FloatArray)->sp.sparse.csr_array:
F = self.compute_F(U)
J0 = self.space.get_tris_mapping().J
w = J0 * self.quad_bilinform.w[0]
# Get raw material tangent and reshape strictly to (N, 4, 4)
dPKIdF = self.law.dPKIdF(F).reshape(-1, 4, 4)
# 1. Pre-allocate the exact output size for the flattened 6x6 blocks
K_blocs = np.empty(w.shape[0] * 36, dtype=F.dtype)
# 2. Execute parallel C-loop (Bypasses all NumPy batched @ operations)
self._fast_k_blocs(w, dPKIdF, self.Bb, K_blocs)
# 3. Apply mask to only accumulate the upper triangle
#K_blocs_upper = K_blocs[self.K_upper_mask]
# Fast Assembly
K_data = np.bincount(self.K_mapping, weights=K_blocs, minlength=len(self.K_indices))
# Construct the CSR array directly
K_csr = sp.sparse.csr_array((K_data, self.K_indices, self.K_indptr), shape=(self.ndofs, self.ndofs), copy=False)
return K_csr
def _compute_Hsparse(self, U:npt.NDArray)->npt.NDArray:
""" compute H, the gradient of the displacment field at each gauss point """
Ushape = U.shape
nelem = self.mesh.ntris()
ndofs = self.space.size()
U = U.reshape(ndofs, copy = False)
vecdim = self.space.vecdim
H = self.B@U
H.shape = (nelem, vecdim, vecdim)
U = U.reshape(Ushape, copy = False)
return H
def _compute_Hblocs(self, U:npt.NDArray)->npt.NDArray:
""" compute H, the gradient of the displacment field at each gauss point
This is a tentative version to speed up the computation compared to compute Hsparse
Not effective as of now
"""
nelem = self.mesh.ntris()
ndofs = self.space.size()
vecdim = self.space.vecdim
U = U.reshape((ndofs,), copy = False)
Bb = self.Bb
Ushape = U.shape
nelem = Bb.shape[0]
U = U.reshape((ndofs//vecdim, vecdim), copy =False)
Utris = U[self.mesh.t2v()]
Utris = Utris.reshape((nelem, 3*vecdim, 1), copy = False)
H = Bb@Utris
H = H.reshape(nelem, 2, 2)
U = U.reshape(Ushape)
return H
[docs]
def compute_F(self, U:npt.NDArray)->npt.NDArray:
""" compute F, the gradient of the position field at each gauss point """
return self._compute_Hsparse(U) + np.eye(2)[np.newaxis]
[docs]
def predictor_corrector(self, U0: FloatArray, loadfac0:npfloat, loadinc:npfloat, **newton_kwargs)-> FloatArray:
""" solve the non linear system """
# Default options for newton, with itmax mapped to nitermax
newton_opts = {
"linsolver": sp.sparse.linalg.spsolve,
"rtol": 1.0e-6,
"atol": 1.0e-6,
"nitermax": 10,
"fail_on_increasing_res": False,
"verbose": True,
}
# Allow callers to override defaults or pass other parameters (norm, isub, etc.)
newton_opts.update(newton_kwargs)
#atol = newton_opts["atol"]
free_dofs = self.free_dofs
diri_dofs = self.fixed_dofs
#Fext = self.compute_external_forces(U0, loadfac0)
#Fint = self.compute_internal_forces(U0)
#check residual at U0
# R = Fint - Fext
# Rf = R[free_dofs]
# nRf = np.linalg.norm(Rf)
# print(f"Starting res {nRf:.3e}")
# assert np.linalg.norm(Rf) < atol
# predictor
K = self.compute_internal_stiffness_blocs(U0)
Kfa = K[free_dofs, :]
Kff = Kfa[:,free_dofs]
U0d = self.fixed_values
# U0d_check = U0[diri_dofs]
# assert(np.linalg.norm(U0d - U0d_check) < atol)
FUD_f = Kfa[:,diri_dofs]@U0d
Flamb1_f = self.compute_external_forces(U0, npfloat(1.))[free_dofs]
# assert(np.linalg.norm(Flamb1_f) == 0.)
linsolve = newton_opts["linsolver"]
DeltaUpred_f = linsolve(Kff, loadinc*(-FUD_f+Flamb1_f ))
Upred = U0
Upred[diri_dofs] = (loadfac0+loadinc)*self.fixed_values
Upred[free_dofs] += DeltaUpred_f
U = Upred.copy()
def fun(x):
U[self.free_dofs] = x
R = self.compute_internal_forces(U)
return R[self.free_dofs]
def jac(x):
U[self.free_dofs] = x
K = self.compute_internal_stiffness_blocs(U)
K_free = K[:,self.free_dofs][self.free_dofs, :]
return K_free
res = newton(fun, x0 = Upred[free_dofs], jac = jac, **newton_opts)
U[self.free_dofs] = res.x
return U
[docs]
def solve_loadstep(mesh:sm.SMesh, dirichlet_bc, loadsteps, \
savedir:None|str = None, popup_gmsh = False, paraviewout = True, **newton_kwargs):
""" example main program """
law = hy.CompressibleNeoHookeen(E=1., nu=0.3, hyp=hy.PLANESTRESS)
pb = HyperElasticitySolver(mesh= mesh, matlaw = law, physname2dirichlet=dirichlet_bc)
U= None
inc = 0
if paraviewout:
writer = ParaViewXDMFWriter("output/hypersolveres", mesh.t2v(), fields_config={
"Cauchy_Stress": {"type": "Tensor", "center": "Cell"}
},
overwrite=True)
gmsh.initialize()
gmsh.option.setNumber('PostProcessing.AnimationCycle', 1)
if savedir is not None:
gmsh.write(savedir+"post_options.opt")
Xtri= mesh.t2v_coord()
X = mesh.vcoord()
loadsteps = np.array(loadsteps)
loadincs = loadsteps[1:] - loadsteps[:-1]
U = np.zeros(pb.ndofs)
def post(U, basename):
u = U.reshape(-1,2)
utri = u[mesh.t2v()]
x = X.copy()
x[mesh.used_v()] +=u
xtri = Xtri + utri
F = pb.compute_F(U).reshape(-1, 2, 2) # pylint: disable=C0103
P = law.PKI(F) # pylint: disable=C0103
Cauchy = PKI_to_Cauchy(PKI = P, F = F) # pylint: disable=C0103
# if paraviewout:
# writer.write_step(i, loadfac0+loadinc, x, Cauchy_Stress = Cauchy)
# if law.hyp == hy.PLANESTRESS:
# lambda3 = law.l3_planestress(np.linalg.det(F))
# vL3 = gp.listPlotFieldTri(x, lambda3, #pylint: disable=C0103
# f"lambda3_{inc:04d}", P0 = True)
# vF = gp.listPlotFieldTri(x, F, f"F_{inc:04d}", P0 = True)
vP = gp.listPlotFieldTri(xtri, P, f"{basename}_{inc:04d}", P0 = True)
if savedir is not None :
#gmsh.view.write(vF, savedir+f"F_{inc:04d}"+".pos")
gmsh.view.write(vP, savedir+f"P_{inc:04d}"+".pos")
# if law.hyp == hy.PLANESTRESS:
# gmsh.view.write(vL3, savedir+f"L3_{inc:04d}"+".pos")
for i, (loadfac0, loadinc) in enumerate( zip(loadsteps[:-1], loadincs)):
print(f"# increment {inc:5d}, loading factor 0 {loadfac0:.3e}, loadinc {loadinc:.3e}")
U = pb.predictor_corrector(U0= U, loadfac0= loadfac0, loadinc= loadinc, **newton_kwargs)
post(U, "Sol")
inc += 1
if paraviewout:
writer.finalize()
if popup_gmsh:
gmsh.fltk.run()
gmsh.finalize()