Add isochoric_volumetric_split() for JAX-hyperelastic material

CHANGELOG.md

@@ -13,7 +13,7 @@ All notable changes to this project will be documented in this file. The format
 - Add `constitution.jax.Hyperelastic` as a feature-equivalent alternative to `Hyperelastic` with `jax` as backend.
 - Add `constitution.jax.Material(..., jacobian=None)` with JAX as backend. A custom jacobian-callable may be passed to switch between forward- and backward-mode automatic differentiation.
 - Add material models for JAX-based materials: `felupe.constitution.jax.models.hyperelastic.mooney_rivlin()`, `felupe.constitution.jax.models.hyperelastic.yeoh()`, `felupe.constitution.jax.models.hyperelastic.third_order_deformation()`, `felupe.constitution.jax.models.lagrange.morph()`, `felupe.constitution.jax.models.lagrange.morph_representative_directions()`.
-- Add `felupe.constitution.jax.total_lagrange()` and `felupe.constitution.jax.updated_lagrange()` function decorators for JAX materials.
+- Add `felupe.constitution.jax.total_lagrange()`, `felupe.constitution.jax.updated_lagrange()` and `felupe.constitution.jax.models.hyperelastic.isochoric_volumetric_split()` function decorators for the JAX hyperelastic material class.
 
 ### Changed
 - Change default `np.einsum(..., order="K")` to `np.einsum(..., order="C")` in the methods of `Field`, `FieldAxisymmetric`, `FieldPlaneStrain` and `FieldContainer`.

src/felupe/constitution/jax/models/hyperelastic/__init__.py

@@ -1,8 +1,10 @@
+from ._helpers import isochoric_volumetric_split
 from ._mooney_rivlin import mooney_rivlin
 from ._third_order_deformation import third_order_deformation
 from ._yeoh import yeoh
 
 __all__ = [
+    "isochoric_volumetric_split",
     "mooney_rivlin",
     "third_order_deformation",
     "yeoh",

src/felupe/constitution/jax/models/hyperelastic/_helpers.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+"""
+This file is part of FElupe.
+
+FElupe is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+FElupe is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with FElupe.  If not, see <http://www.gnu.org/licenses/>.
+"""
+from functools import wraps
+
+
+def isochoric_volumetric_split(fun):
+    """Apply the material formulation only on the isochoric part of the
+    multiplicative split of the deformation gradient."""
+    from jax.numpy.linalg import det
+
+    @wraps(fun)
+    def apply_iso(C, *args, **kwargs):
+        return fun(det(C) ** (-1 / 3) * C, *args, **kwargs)
+
+    return apply_iso

tests/test_constitution_newton.py

@@ -25,21 +25,24 @@
 
 """
 import numpy as np
-import tensortrax as tr
-from tensortrax.math import trace
-from tensortrax.math.linalg import det, inv
-from tensortrax.math.special import dev, from_triu_1d, triu_1d
 
 import felupe as fem
 
 
 def test_visco_newton():
+    import tensortrax as tr
+    from tensortrax.math import trace
+    from tensortrax.math.linalg import det, inv
+    from tensortrax.math.special import from_triu_1d, triu_1d
+
+    import felupe.constitution.tensortrax as mat
+
     mesh = fem.Rectangle(n=3)
     region = fem.RegionQuad(mesh)
     field = fem.FieldContainer([fem.FieldPlaneStrain(region, dim=2)])
     boundaries, loadcase = fem.dof.uniaxial(field, clamped=True)
 
-    @fem.constitution.isochoric_volumetric_split
+    @mat.models.hyperelastic.isochoric_volumetric_split
     def finite_strain_viscoelastic_newton(C, Cin, mu, eta, dtime):
         "Finite Strain Viscoelastic material formulation."
 
@@ -65,11 +68,10 @@ def evolution(Ci, Cin, C, mu, eta, dtime):
         # strain energy function and state variable
         return mu / 2 * (I1 - 3), triu_1d(Ci)
 
-    umat = fem.Hyperelastic(
+    umat = mat.Hyperelastic(
         finite_strain_viscoelastic_newton, mu=1.0, eta=1.0, dtime=0.1, nstatevars=6
     )
     solid = fem.SolidBodyNearlyIncompressible(umat, field, bulk=5000)
-    # solid.results.statevars[[0, 3, 5]] += 1
 
     move = fem.math.linsteps([0, 0.3], num=3)
     step = fem.Step(
@@ -79,5 +81,66 @@ def evolution(Ci, Cin, C, mu, eta, dtime):
     assert np.all([norm[-1] < 1e-7 for norm in job.fnorms])
 
 
+def test_visco_newton_jax():
+    import jax
+    from jax.numpy import trace
+    from jax.numpy.linalg import det, inv
+    from jax.scipy.optimize import minimize
+
+    import felupe.constitution.jax as mat
+
+    mesh = fem.Rectangle(n=3)
+    region = fem.RegionQuad(mesh)
+    field = fem.FieldContainer([fem.FieldPlaneStrain(region, dim=2)])
+    boundaries, loadcase = fem.dof.uniaxial(field, clamped=True)
+
+    @mat.models.hyperelastic.isochoric_volumetric_split
+    def finite_strain_viscoelastic_newton(C, statevars, mu, eta, dtime):
+        "Finite Strain Viscoelastic material formulation."
+
+        def evolution(Ci, Cin, C, mu, eta, dtime):
+            "Viscoelastic evolution equation."
+            Ci = Ci.reshape(3, 3)
+            residuals = mu / eta * C - (Ci - Cin) / dtime
+            return residuals.ravel() @ residuals.ravel()
+
+        # update of state variables by evolution equation
+        Cin = statevars[:9].reshape(3, 3)
+        Ci = minimize(
+            evolution,
+            x0=Cin.ravel(),
+            args=(Cin, jax.lax.stop_gradient(C), mu, eta, dtime),
+            method="BFGS",
+        ).x.reshape(3, 3)
+        Ci *= det(Ci) ** (-1 / 3)
+
+        # first invariant of elastic part of right Cauchy-Green deformation tensor
+        I1 = trace(C @ inv(Ci))
+
+        # strain energy function and state variable
+        statevars_new = Ci.ravel()
+        return mu / 2 * (I1 - 3), statevars_new
+
+    umat = mat.Hyperelastic(
+        finite_strain_viscoelastic_newton,
+        mu=1.0,
+        eta=1.0,
+        dtime=0.1,
+        nstatevars=9,
+        jit=True,
+    )
+    solid = fem.SolidBodyNearlyIncompressible(umat, field, bulk=5000)
+
+    move = fem.math.linsteps([0, 0.3], num=3)
+    step = fem.Step(
+        items=[solid], ramp={boundaries["move"]: move}, boundaries=boundaries
+    )
+    job = fem.CharacteristicCurve(steps=[step], boundary=boundaries["move"])
+    job.evaluate(tol=1e-3)
+
+    assert np.all([fnorm[-1] < 1e-3 for fnorm in job.fnorms])
+
+
 if __name__ == "__main__":
     test_visco_newton()
+    test_visco_newton_jax()