FluxFEM Finte Element Method with JAX

What is FluxFEM?

A weak-form-centric differentiable finite element framework in JAX, where variational forms are treated as first-class, differentiable programs.

https://github.com/kevin-tofu/fluxfem

Coding nonlinear structural analysis with FluxFEM

Neohookean with Weakform

import fluxfem as ff
import fluxfem.helpers_wf as h_wf

def neo_hookean_residual_wf(v, u, params):
    mu = params["mu"]
    lam = params["lam"]

    F = h_wf.I(3) + h_wf.grad(u)
    C = h_wf.matmul(h_wf.transpose(F), F)
    C_inv = h_wf.inv(C)
    J = h_wf.det(F)

    S = mu * (h_wf.I(3) - C_inv) + lam * h_wf.log(J) * C_inv
    dE = 0.5 * (
        h_wf.matmul(h_wf.grad(v), F)
        + h_wf.transpose(h_wf.matmul(h_wf.grad(v), F))
    )
    return h_wf.ddot(S, dE) * h_wf.dOmega()

residual_form = ff.ResidualForm.volume(neo_hookean_residual_wf)

Newton method interface

import numpy as np
import jax.numpy as jnp
import fluxfem as ff
import fluxfem.helpers_ts as h_ts

dtype = jnp.float64
mesh, facets, facet_tags = ff.load_gmsh_tet_mesh("tension_bar.msh")
space = ff.make_tet_space(mesh, dim=3, intorder=2)
U = ff.NamedSpace("U", space)
V = ff.NamedSpace("V", space)

analysis = ff.NonlinearAnalysis(
    space=space,
    residual_form=ff.neo_hookean_residual_form,
    params=params,
    base_external_vector=F_ext,
    dirichlet=ff.DirichletBC(dir_dofs, None),
    jacobian_pattern=ff.make_sparsity_pattern(space, with_idx=False),
    dtype=dtype,
)

newton_cfg = ff.NewtonLoopConfig(
    tol=1e-6,
    atol=0.0,
    maxiter=10,
    linear_solver="petsc_shell",
    petsc_ksp_type="preonly",
    petsc_pc_type="lu",
    petsc_use_pmat=True,
    n_steps=20,
)

runner = ff.NewtonSolveRunner(analysis, newton_cfg)
u, history = runner.run(u0=jnp.zeros(space.n_dofs, dtype=dtype))

Result

It seems completely fine.