jaxmat.tensors.generic_tensors module

jaxmat.tensors.generic_tensors module#

jaxmat/tensors/generic_tensors.py

Second- and fourth-rank tensor classes for solid mechanics computations in JAX.

Design principles#

  • The Kelvin-Mandel array is the unique JAX leaf stored per object. All other representations are derived properties.

  • All index maps and weight arrays are module-level constants computed once at import time and shared across all instances.

  • tensor properties use gather operations (no scatter, no intermediate zero allocation) for rank-2 classes.

  • Each concrete class exposes exactly one equinox field: _array.

  • Operator @ semantics:

    • Tensor2 @ Tensor2 → dense composition → Tensor2.

    • SymmetricTensor4 @ SymmetricTensor2 → Kelvin matvec → SymmetricTensor2.

    • SymmetricTensor4 @ SymmetricTensor4 → Kelvin matmul → SymmetricTensor4.

  • Named methods for algebraically unambiguous operations: double_contract(), matvec(), fourth_contract().

  • _AbstractTensor4 centralizes all code shared between Tensor4 and SymmetricTensor4; subclasses supply only __init__, tensor, __matmul__, base_array_shape, and identity.