A JAX-based constitutive material

[adtzlr]

JAX is a popular differentiable library (there is also a FEM-package built on it, JAX-FEM). jax.vmap must be applied twice on the last two dimensions when used in FElupe. Otherwise, the usage is straight forward!

import felupe as fem

mesh = fem.Cube(n=11)
region = fem.RegionHexahedron(mesh)
field = fem.FieldContainer([fem.Field(region, dim=3)])

boundaries, loadcase = fem.dof.uniaxial(field, clamped=True)

import jax
import jax.numpy as jnp


def W(F):
    C = F.T @ F
    I1 = jnp.trace(C)
    J = jnp.linalg.det(F)
    return 1 / 2 * (J ** (-2 / 3) * I1 - 3) + 2 * (J - 1) ** 2 / 2


class JAXMaterial:
    def __init__(self, W):
        self.W = W
        self.dWdF = jax.jit(self.vmap(self.vmap(jax.grad(self.W))))
        self.d2WdFdF = jax.jit(self.vmap(self.vmap(jax.jacobian(jax.grad(self.W)))))

    def vmap(self, fun, in_axes=-1, out_axes=-1, **kwargs):
        return jax.vmap(fun, in_axes=in_axes, out_axes=out_axes, **kwargs)

    def gradient(self, x):
        F, statevars = x[0], x[-1]
        return [self.dWdF(F), statevars]

    def hessian(self, x):
        F = x[0]
        return [self.d2WdFdF(F)]


umat = JAXMaterial(W)
solid = fem.SolidBody(umat=umat, field=field)

move = fem.math.linsteps([0, 1], num=5)
ramp = {boundaries["move"]: move}
step = fem.Step(items=[solid], ramp=ramp, boundaries=boundaries)

job = fem.Job(steps=[step])
job.evaluate(verbose=2, tol=1e-6)

ax = solid.imshow("Principal Values of Cauchy Stress")

[adtzlr]

The JAX-hessian is about 10 times faster compared to the hard-coded Neo-Hookean material formulation in FElupe. However, for these simple models, the integration is much more costly than the hessian of the strain energy density function. For more lengthy strain energy density functions JAX will be really useful. The Neo-Hookean material formulation is probably too easy to see any real benefits. With integration, the evaluation of the form given below is meant, where d2ψdFdF is already pre-evaluated (e.g. by JAX).

JAX:

$$\mathbb{A}_ {iJkL(qc)} = \left( \frac{\partial^2 \psi}{\partial F_ {iJ} \partial F_ {kL}} \right)_ {(qc)}$$

FElupe:

$$\hat{K}_ {aibj(c)} = \sum_q \frac{\partial h_a}{\partial X_J}_ {(qc)} \mathbb{A}_ {iJkL(qc)} \frac{\partial h_b}{\partial X_L}_ {(qc)}\ dV_ {(qc)}$$

[adtzlr]

It's ready! 🚀 🚀

https://github.com/adtzlr/felupe/blob/main/src/felupe/constitution/jax/_hyperelastic.py
#876