[Tutorial] Hyperelastic Material Definition with Automatic Differentation using casADi

[adtzlr]

Hi @kinnala, @gdmcbain, @bhaveshshrimali,

this is not a Question, it's a tutorial - probably in favour to #441 using JAX.

I found a nice package for automatic differentiation (AD) with easy installation, good performance and straight-forward usage: casADi (LGPL licensed). To be honest, this is the first AD-package I'm happy with. I had no luck with autograd (slow) and jax (Windows -> install ❌). At least on Windows casadi performs really well! I figured out how to apply deformation gradients with trailing axes to the AD-generated functions. Performance is about four times slower compared to hardcoded stress and elasticity tensors, which is very okay - at least for me. Let's take the Neo-Hookean material formulation as an example.

Neo-Hookean material formulation
The strain energy density function (per unit undeformed volume) of the Neo-Hookean material formulation is defined as follows:

First, let's define some random deformation gradients.

import casadi as ca
import numpy as np
from functools import partial
from skfem.helpers import identity

np.random.seed(3242345)
dudX = np.random.rand(3,3,8000,8) / 2
dxdX =identity(dudX) + dudX

I wrote a little helper class Material which provides AD-calculated functions for stress and elasticity where deformation gradient tensors with trailing axes as they are used in scikit-fem are applied to the function objects generated by casadi. If you prefer something simpler, just take the apply method of this class and use it as a simple function.

class Material:
    def __init__(self, W, *args, **kwargs):
        "Material class with Automatic Differentiation."
        
        # init deformation gradient
        F = ca.SX.sym('F', 3, 3)
        
        # gradient and hessian of strain ernergy function W
        d2WdFdF, dWdF = ca.hessian(W(F, *args, **kwargs), F)
        
        # generate casadi function objects
        f_P = ca.Function("P", [F], [dWdF])
        f_A = ca.Function("A", [F], [d2WdFdF])
        
        # functions for stress P and elasticity A
        self.P = partial(self._apply, fun=f_P, fun_shape=(3,3))
        self.A = partial(self._apply, fun=f_A, fun_shape=(3,3,3,3))
    
    def _apply(self, x, fun, fun_shape, trailing_axes=2):
        "Internal helper function for the calculation of fun(x)."
        
        # get shape of trailing axes
        ax = x.shape[-trailing_axes:]
        
        # reshape input
        y = x.reshape(3, -1, order="F")
        
        # map function to reshaped input
        out = np.array(
            fun.map(np.product(ax))(y)
        )
            
        return out.reshape(*fun_shape, *ax, order="F")

Next, we define the strain energy density function of the Neo-Hookean material model, create an instance of the Material class and calculate stress (1st Piola-Kirchhoff) and elasticity tensors.

def W_neohooke(F, mu=1.0, bulk=5.0):
    "Strain energy density function for Neo-Hookean material formulation."
    
    J = ca.det(F)
    C = F.T @ F
    I_C = ca.trace(C)
    
    return mu / 2 * (I_C * J**(-2 / 3) - 3) + bulk / 2 * (J - 1)**2

# material with automatic differention
umat = Material(W_neohooke, mu=1.0, bulk=200.0)
P_AD = umat.P(dxdX)
A_AD = umat.A(dxdX)

That's it! 😎

Below you'll find some checks and performance evaluations...

I compared the results with a reference implementation of the neo-hookean material used in my own package felupe.

import felupe as fe
# reference material (hard-coded; taken from `felupe`)
nh = fe.constitution.NeoHooke(mu=1.0, bulk=200.0)
P = nh.P(dxdX)
A = nh.A(dxdX)

#check casADi-calculated PK1-stress and elasticity
print("Check results...")
print(np.allclose(P, P_AD))
print(np.allclose(A, A_AD))

# timings
import timeit

def test_autodiff():
    return umat.A(dxdX)

def test_hardcoded():
    return nh.A(dxdX)

time_autodiff = timeit.timeit(test_autodiff, number=1)
time_hardcoded = timeit.timeit(test_hardcoded, number=1)

print("")
print("Exec.time Autom.Diff. =", time_autodiff)
print("Exec.time Hardcoded =", time_hardcoded)
print("Exec.time Autom.Diff. vs. Hardcoded =", time_autodiff / time_hardcoded)

Check results...
True
True

Exec.time Autom.Diff. = 1.3052240000000666
Exec.time Hardcoded = 0.33741989999998623
Exec.time Autom.Diff. vs. Hardcoded = 3.8682484346658863

Disclaimer: I have NOT testet this stuff yet with scikit-fem, but that shoudn't be too hard... I thought probably someone of you may find this useful!

Cheers 🏖️,
Andreas

P.S.: There is one big drawback in casadi - if-statements...

[kinnala]

Thanks for the suggestion. I'll surely try it out at some point.

[adtzlr]

I just wanted to let you know that I recently created a little module matADi which basically handles the above code snippets in a more user-friendly way.

...beside that - I'm still impressed of the development of scikit-fem. This is such an amazing project 👍🏻

[adtzlr]

And this is a simple script:

import numpy as np

import skfem as fem
from skfem.assembly import LinearForm, BilinearForm
from skfem.helpers import grad, identity

from matadi import MaterialHyperelastic
from matadi.models import neo_hooke

mat = MaterialHyperelastic(neo_hooke, C10=0.5, bulk=2.0)

right = 0.5
tol = 1e-10

m  = fem.MeshHex().refined(2)
e1 = fem.ElementHex1()
e  = fem.ElementVectorH1(e1)
ib = fem.Basis(m, e, intorder=1)

def P(w):
    H = grad(w["displacement"])
    F = H + identity(H)
    return mat.gradient([F])[0]

def A(w):
    H = grad(w["displacement"])
    F = H + identity(H)
    return mat.hessian([F])[0]

@LinearForm(nthreads=4)
def L(v, w):
    return np.einsum("ij...,ij...", P(w), grad(v))

@BilinearForm(nthreads=4)
def a(u, v, w):
    return np.einsum("ijkl...,ij...,kl...", A(w), grad(u), grad(v))

u  = ib.zeros()
du = ib.zeros()
uv = ib.interpolate(u)

dofs = {
    'left' : ib.get_dofs(lambda x: x[0] == 0.0),
    'right': ib.get_dofs(lambda x: x[0] == 1.0),
}

I = ib.complement_dofs(dofs)

f = fem.asm(L, ib, displacement=uv)
K = fem.asm(a, ib, ib, displacement=uv)

for iteration in range(8):
    
    du_D = u.copy()
    du_D[dofs["right"].nodal['u^1']] = right - du_D[dofs["right"].nodal['u^1']]
    
    system = fem.condense(K, -f, x=du_D, I=I)
    du = fem.solve(*system)
    norm_du = np.linalg.norm(du)
    
    u += du
    
    uv = ib.interpolate(u)
    f = fem.asm(L, ib, displacement=uv)
    K = fem.asm(a, ib, ib, displacement=uv)
    
    print(1 + iteration, norm_du)
    
    if norm_du < tol:
        break

savedata = {"u": u[ib.nodal_dofs].T}
m.save("result.vtk", savedata)

Output:

0 3.4810410954530506
1 0.13063064781292275
2 0.009208944826176955
3 4.989804896048235e-05
4 2.3408069100383187e-09
5 3.2247002135208875e-16

[adtzlr]

The above bultin Neo-Hookean material model can be replaced by any user-defined strain energy function in terms of the deformation gradient. Stress and Elasticity tensor are carried out by automatic differentiation at reasonable speed.

# ...

from matadi import MaterialHyperelastic
from matadi.math import det, transpose, trace

def strain_energy_function(F, C10, bulk):
    J = det(F)
    C = transpose(F) @ F
    return C10 * (J**(-2/3) * trace(C) - 3) + bulk * (J - 1)**2/2

#mat = MaterialHyperelastic(neo_hooke, C10=0.5, bulk=2.0)
mat = MaterialHyperelastic(strain_energy_function, C10=0.5, bulk=2.0)

# ...

[kinnala]

It's great, good job! I was going to try this out, already started to write some code, but you finished first. This would be good material for the eventual #712.

[adtzlr]

Hi, inspired by the linear elastic material definition of scikit-fem I tried to implement the bilinear form of a Neo-Hookean solid without calculating the fourth-order (deformation-dependent) elasticity tensor explicitely. Instead I only used variational calculus. Some first tests gave me a 2x speed-up and a really nice code representation. I'll come up with a nice example in the next days. 😎

[adtzlr]

Example and discussion --> see #839 for an implementation.