Make tensortrax optional

adtzlr · Nov 9, 2024 · 5669392 · 5669392
1 parent 29571c1
commit 5669392
Show file tree

Hide file tree

Showing 21 changed files with 190 additions and 172 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,7 @@ All notable changes to this project will be documented in this file. The format
 - Change default `np.einsum(..., order="K")` to `np.einsum(..., order="C")` in the methods of `Field`, `FieldAxisymmetric`, `FieldPlaneStrain` and `FieldContainer`.
 - Change supported Python versions to 3.9 - 3.12.
 - Change the `dtype`-argument in `Region.astype(dtype)` from an optional to a required argument.
+- Make `tensortrax` an optional dependency (again). Now FElupe does only depend on NumPy and SciPy, all other extras are optional.
 
 ### Fixed
 - Fix the number of points for non-disconnected dual meshes. This reduces the assembled (sparse) vector- and matrix-shapes, which are defined on mixed-fields.

diff --git a/docs/index.rst b/docs/index.rst
@@ -65,16 +65,16 @@ where ``[all]`` is a combination of ``[io,parallel,plot,progress,view]`` and ins
 
 * `numpy <https://github.com/numpy/numpy>`_ for array operations
 * `scipy <https://github.com/scipy/scipy>`_ for sparse matrices
-* `tensortrax <https://github.com/adtzlr/tensortrax>`_ for automatic differentiation
 
 In order to make use of all features of FElupe 💎💰💍👑💎, it is suggested to install all optional dependencies.
 
 * `einsumt <https://github.com/mrkwjc/einsumt>`_ for parallel (threaded) assembly
-* `jax <https://github.com/jax-ml/jax>`_ for JAX-based material formulations
+* `jax <https://github.com/jax-ml/jax>`_ for automatic differentiation in material formulations (JAX-based)
 * `h5py <https://github.com/h5py/h5py>`_ for writing XDMF result files
 * `matplotlib <https://github.com/matplotlib/matplotlib>`_ for plotting graphs
 * `meshio <https://github.com/nschloe/meshio>`_ for mesh-related I/O
 * `pyvista <https://github.com/pyvista/pyvista>`_ for interactive visualizations
+* `tensortrax <https://github.com/adtzlr/tensortrax>`_ for automatic differentiation in material formulations (NumPy-based)
 * `tqdm <https://github.com/tqdm/tqdm>`_ for showing progress bars during job evaluations
 
 The development version may contain not yet released bug fixes and features. Consider using the ``--user`` option to install the package into your home directory (see `pip documentation <https://pip.pypa.io/en/stable/user_guide/#user-installs>`_ for more details). To install FElupe from the latest development branch, use

diff --git a/pyproject.toml b/pyproject.toml
@@ -51,7 +51,6 @@ dynamic = ["version"]
 dependencies = [
   "numpy",
   "scipy",
-  "tensortrax>=0.25.2",
 ]
 
 [project.optional-dependencies]
@@ -64,11 +63,11 @@ docs = [
     "sphinx-gallery",
     "pypandoc",
 ]
-constitution = ["jax"]
+autodiff = ["tensortrax", "jax"]
 examples = [
     "contique",
     "imageio",
-    "matadi>=0.2.2",
+    "matadi",
     "pypardiso",
     "torch",
 ]
@@ -81,8 +80,8 @@ progress = ["tqdm"]
 plot = ["matplotlib"]
 view = ["pyvista[jupyter]"]
 
-test = ["felupe[io,constitution,plot]"]
-all = ["felupe[io,constitution,parallel,plot,progress,view]"]
+test = ["felupe[io,autodiff,plot]"]
+all = ["felupe[io,autodiff,parallel,plot,progress,view]"]
 
 [tool.setuptools.dynamic]
 version = {attr = "felupe.__about__.__version__"}

diff --git a/src/felupe/__init__.py b/src/felupe/__init__.py
@@ -19,7 +19,6 @@
     AreaChange,
     CompositeMaterial,
     ConstitutiveMaterial,
-    Hyperelastic,
     Laplace,
     LinearElastic,
     LinearElasticLargeStrain,
@@ -28,7 +27,6 @@
     LinearElasticPlasticIsotropicHardening,
     LineChange,
     Material,
-    MaterialAD,
     MaterialStrain,
     NearlyIncompressible,
     NeoHooke,
@@ -39,29 +37,8 @@
     ViewMaterialIncompressible,
     VolumeChange,
     Volumetric,
-    alexander,
-    anssari_benam_bucchi,
-    arruda_boyce,
-    constitutive_material,
-    extended_tube,
-    finite_strain_viscoelastic,
-    isochoric_volumetric_split,
     linear_elastic,
     linear_elastic_plastic_isotropic_hardening,
-    lopez_pamies,
-    miehe_goektepe_lulei,
-    mooney_rivlin,
-    morph,
-    morph_representative_directions,
-    neo_hooke,
-    ogden,
-    ogden_roxburgh,
-    saint_venant_kirchhoff,
-    third_order_deformation,
-    total_lagrange,
-    updated_lagrange,
-    van_der_waals,
-    yeoh,
 )
 from .dof import Boundary
 from .element import ArbitraryOrderLagrange as ArbitraryOrderLagrangeElement
@@ -181,35 +158,13 @@
     "NearlyIncompressible",
     "Material",
     "MaterialStrain",
-    "Hyperelastic",
-    "total_lagrange",
-    "updated_lagrange",
-    "MaterialAD",
     "ViewMaterial",
     "ViewMaterialIncompressible",
     "ConstitutiveMaterial",
     "constitutive_material",
     "VolumeChange",
     "linear_elastic",
     "linear_elastic_plastic_isotropic_hardening",
-    "alexander",
-    "anssari_benam_bucchi",
-    "arruda_boyce",
-    "extended_tube",
-    "finite_strain_viscoelastic",
-    "isochoric_volumetric_split",
-    "lopez_pamies",
-    "miehe_goektepe_lulei",
-    "mooney_rivlin",
-    "morph",
-    "morph_representative_directions",
-    "neo_hooke",
-    "ogden",
-    "ogden_roxburgh",
-    "saint_venant_kirchhoff",
-    "third_order_deformation",
-    "van_der_waals",
-    "yeoh",
     "Boundary",
     "ArbitraryOrderLagrangeElement",
     "BiQuadraticQuad",
@@ -286,3 +241,58 @@
     "ViewSolid",
     "runs_on",
 ]
+
+try:
+    from .constitution import (
+        Hyperelastic,
+        MaterialAD,
+        alexander,
+        anssari_benam_bucchi,
+        arruda_boyce,
+        constitutive_material,
+        extended_tube,
+        finite_strain_viscoelastic,
+        isochoric_volumetric_split,
+        lopez_pamies,
+        miehe_goektepe_lulei,
+        mooney_rivlin,
+        morph,
+        morph_representative_directions,
+        neo_hooke,
+        ogden,
+        ogden_roxburgh,
+        saint_venant_kirchhoff,
+        third_order_deformation,
+        total_lagrange,
+        updated_lagrange,
+        van_der_waals,
+        yeoh,
+    )
+
+    __all_tensortrax = [
+        "Hyperelastic",
+        "total_lagrange",
+        "updated_lagrange",
+        "MaterialAD",
+        "alexander",
+        "anssari_benam_bucchi",
+        "arruda_boyce",
+        "extended_tube",
+        "finite_strain_viscoelastic",
+        "isochoric_volumetric_split",
+        "lopez_pamies",
+        "miehe_goektepe_lulei",
+        "mooney_rivlin",
+        "morph",
+        "morph_representative_directions",
+        "neo_hooke",
+        "ogden",
+        "ogden_roxburgh",
+        "saint_venant_kirchhoff",
+        "third_order_deformation",
+        "van_der_waals",
+        "yeoh",
+    ]
+    __all__.extend(__all_tensortrax)
+except ImportError:
+    pass
diff --git a/src/felupe/constitution/__init__.py b/src/felupe/constitution/__init__.py
@@ -20,47 +20,8 @@
     linear_elastic,
     linear_elastic_plastic_isotropic_hardening,
 )
-from .tensortrax import Hyperelastic
-from .tensortrax import Material as MaterialAD
-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,
-    lopez_pamies,
-    miehe_goektepe_lulei,
-    mooney_rivlin,
-    neo_hooke,
-    ogden,
-    ogden_roxburgh,
-    saint_venant_kirchhoff,
-    third_order_deformation,
-    van_der_waals,
-    yeoh,
-)
-from .tensortrax.models.lagrange import morph, morph_representative_directions
 
 __all__ = [
-    "alexander",
-    "anssari_benam_bucchi",
-    "arruda_boyce",
-    "extended_tube",
-    "finite_strain_viscoelastic",
-    "isochoric_volumetric_split",
-    "lopez_pamies",
-    "miehe_goektepe_lulei",
-    "mooney_rivlin",
-    "morph",
-    "morph_representative_directions",
-    "neo_hooke",
-    "ogden",
-    "ogden_roxburgh",
-    "saint_venant_kirchhoff",
-    "third_order_deformation",
-    "van_der_waals",
-    "yeoh",
     "NeoHooke",
     "NeoHookeCompressible",
     "Laplace",
@@ -75,10 +36,6 @@
     "LinearElasticPlasticIsotropicHardening",
     "Material",
     "MaterialStrain",
-    "MaterialAD",
-    "Hyperelastic",
-    "total_lagrange",
-    "updated_lagrange",
     "AreaChange",
     "LineChange",
     "VolumeChange",
@@ -93,3 +50,54 @@
     "CompositeMaterial",
     "Volumetric",
 ]
+try:
+    from .tensortrax import Hyperelastic
+    from .tensortrax import Material as MaterialAD
+    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,
+        lopez_pamies,
+        miehe_goektepe_lulei,
+        mooney_rivlin,
+        neo_hooke,
+        ogden,
+        ogden_roxburgh,
+        saint_venant_kirchhoff,
+        third_order_deformation,
+        van_der_waals,
+        yeoh,
+    )
+    from .tensortrax.models.lagrange import morph, morph_representative_directions
+
+    __all_tensortrax = [
+        "alexander",
+        "anssari_benam_bucchi",
+        "arruda_boyce",
+        "extended_tube",
+        "finite_strain_viscoelastic",
+        "isochoric_volumetric_split",
+        "lopez_pamies",
+        "miehe_goektepe_lulei",
+        "mooney_rivlin",
+        "morph",
+        "morph_representative_directions",
+        "neo_hooke",
+        "ogden",
+        "ogden_roxburgh",
+        "saint_venant_kirchhoff",
+        "third_order_deformation",
+        "van_der_waals",
+        "yeoh",
+        "MaterialAD",
+        "Hyperelastic",
+        "total_lagrange",
+        "updated_lagrange",
+    ]
+    __all__.extend(__all_tensortrax)
+
+except ImportError:
+    pass
diff --git a/src/felupe/constitution/jax/__init__.py b/src/felupe/constitution/jax/__init__.py
@@ -1,16 +1,20 @@
-from . import models
-from ._helpers import isochoric_volumetric_split, vmap
-from ._hyperelastic import Hyperelastic
-from ._material import Material
-from ._total_lagrange import total_lagrange
-from ._updated_lagrange import updated_lagrange
+try:
+    from . import models
+    from ._helpers import isochoric_volumetric_split, vmap
+    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",
-    "Material",
-    "models",
-    "total_lagrange",
-    "updated_lagrange",
-    "vmap",
-]
+    __all__ = [
+        "Hyperelastic",
+        "isochoric_volumetric_split",
+        "Material",
+        "models",
+        "total_lagrange",
+        "updated_lagrange",
+        "vmap",
+    ]
+
+except ModuleNotFoundError:
+    __all__ = []
diff --git a/src/felupe/constitution/jax/_helpers.py b/src/felupe/constitution/jax/_helpers.py
@@ -15,10 +15,13 @@
 You should have received a copy of the GNU General Public License
 along with FElupe.  If not, see <http://www.gnu.org/licenses/>.
 """
-
 import inspect
 from functools import wraps
 
+import jax
+import jax.numpy as jnp
+from jax.numpy.linalg import det
+
 
 def vmap(fun, in_axes=0, out_axes=0, **kwargs):
     """Vectorizing map. Creates a function which maps ``fun`` over argument axes. This
@@ -30,8 +33,6 @@ def vmap(fun, in_axes=0, out_axes=0, **kwargs):
         axes.
     """
 
-    import jax
-
     @wraps(fun)
     def vmap_with_static_kwargs(*args, **keywordargs):
         # sorted list of all parameter keys, including kwargs with default values
@@ -91,7 +92,6 @@ def vmap2(fun, in_axes=0, out_axes=0, **kwargs):
 
 
 def as_total_lagrange(fun):
-    import jax.numpy as jnp
 
     @wraps(fun)
     def evaluate(F, *args, **kwargs):
@@ -106,7 +106,6 @@ def evaluate(F, *args, **kwargs):
 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):