Move constitution.jax.isochoric_volumetric_split()

src/felupe/constitution/__init__.py

@@ -22,14 +22,13 @@
 )
 from .tensortrax import Hyperelastic
 from .tensortrax import Material as MaterialAD
-from .tensortrax import total_lagrange, updated_lagrange
+from .tensortrax import isochoric_volumetric_split, total_lagrange, updated_lagrange
 from .tensortrax.models.hyperelastic import (
     alexander,
     anssari_benam_bucchi,
     arruda_boyce,
     extended_tube,
     finite_strain_viscoelastic,
-    isochoric_volumetric_split,
     lopez_pamies,
     miehe_goektepe_lulei,
     mooney_rivlin,

src/felupe/constitution/jax/__init__.py

@@ -1,12 +1,13 @@
 from . import models
+from ._helpers import isochoric_volumetric_split, vmap
 from ._hyperelastic import Hyperelastic
 from ._material import Material
-from ._tools import vmap
 from ._total_lagrange import total_lagrange
 from ._updated_lagrange import updated_lagrange
 
 __all__ = [
     "Hyperelastic",
+    "isochoric_volumetric_split",
     "Material",
     "models",
     "total_lagrange",

src/felupe/constitution/jax/_tools.py → src/felupe/constitution/jax/_helpers.py

@@ -90,7 +90,7 @@ def vmap2(fun, in_axes=0, out_axes=0, **kwargs):
     )
 
 
-def total_lagrange(fun):
+def as_total_lagrange(fun):
     import jax.numpy as jnp
 
     @wraps(fun)
@@ -101,3 +101,15 @@ def evaluate(F, *args, **kwargs):
         return fun(C, *args, **kwargs)
 
     return evaluate
+
+
+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

src/felupe/constitution/jax/_hyperelastic.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 from .._material import Material
-from ._tools import total_lagrange, vmap2
+from ._helpers import as_total_lagrange, vmap2
 
 
 class Hyperelastic(Material):
@@ -147,7 +147,7 @@ def __init__(self, fun, nstatevars=0, jit=True, parallel=False, **kwargs):
         import jax
 
         has_aux = nstatevars > 0
-        self.fun = total_lagrange(fun)
+        self.fun = as_total_lagrange(fun)
 
         if parallel:
             warnings.warn("Parallel execution is not implemented.")

src/felupe/constitution/jax/_material.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 from .._material import Material as MaterialDefault
-from ._tools import vmap2
+from ._helpers import vmap2
 
 
 class Material(MaterialDefault):

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

@@ -1,10 +1,8 @@
-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/tensortrax/__init__.py

@@ -1,11 +1,13 @@
 from . import models
+from ._helpers import isochoric_volumetric_split
 from ._hyperelastic import Hyperelastic
 from ._material import Material
 from ._total_lagrange import total_lagrange
 from ._updated_lagrange import updated_lagrange
 
 __all__ = [
     "Hyperelastic",
+    "isochoric_volumetric_split",
     "models",
     "total_lagrange",
     "updated_lagrange",

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

@@ -14,7 +14,6 @@
 from ._arruda_boyce import arruda_boyce
 from ._extended_tube import extended_tube
 from ._finite_strain_viscoelastic import finite_strain_viscoelastic
-from ._helpers import isochoric_volumetric_split
 from ._lopez_pamies import lopez_pamies
 from ._miehe_goektepe_lulei import miehe_goektepe_lulei
 from ._mooney_rivlin import mooney_rivlin
@@ -33,7 +32,6 @@
     "arruda_boyce",
     "extended_tube",
     "finite_strain_viscoelastic",
-    "isochoric_volumetric_split",
     "lopez_pamies",
     "miehe_goektepe_lulei",
     "mooney_rivlin",

tests/test_constitution_newton.py

@@ -42,7 +42,7 @@ def test_visco_newton():
     field = fem.FieldContainer([fem.FieldPlaneStrain(region, dim=2)])
     boundaries, loadcase = fem.dof.uniaxial(field, clamped=True)
 
-    @mat.models.hyperelastic.isochoric_volumetric_split
+    @mat.isochoric_volumetric_split
     def finite_strain_viscoelastic_newton(C, Cin, mu, eta, dtime):
         "Finite Strain Viscoelastic material formulation."
 
@@ -94,7 +94,7 @@ def test_visco_newton_jax():
     field = fem.FieldContainer([fem.FieldPlaneStrain(region, dim=2)])
     boundaries, loadcase = fem.dof.uniaxial(field, clamped=True)
 
-    @mat.models.hyperelastic.isochoric_volumetric_split
+    @mat.isochoric_volumetric_split
     def finite_strain_viscoelastic_newton(C, statevars, mu, eta, dtime):
         "Finite Strain Viscoelastic material formulation."