A Differentiable Framework for Gradient Enhanced Damage with Physics-Augmented Neural Networks in JAX-FEM

Abstract: Soft materials such as rubbers, hydrogels, and biological tissues undergo damage in the form of stiffness degradation without apparent changes in their stress-free geometry. Accurate simulation of this behavior is critical in applications ranging from soft robotics to the design of medical devices, yet two persistent challenges are the difficulty of constructing flexible, thermodynamically consistent constitutive models, and the mesh dependence of finite element solutions caused by strain softening. Here we address both challenges simultaneously by combining physics-augmented neural network constitutive models with a gradient-enhanced damage formulation implemented within the differentiable finite element framework JAX-FEM. The elastic strain energy and the damage yield function are each parameterized by input-convex neural networks (ICNNs), which enforce polyconvexity and satisfaction of the Clausius--Duhem inequality by design. The gradient-enhanced formulation introduces a non-local damage field governed by an additional partial differential equation, regularizing the spatial distribution of damage and eliminating mesh dependence. The implementation is validated through local stress-strain fits, single-element parametric studies, a mesh and solution strategy study for a uniform deformation case, and a notched plate simulation. The results demonstrate that the proposed framework enables flexible, data-driven, mesh-independent damage simulation for a broad class of soft materials. We anticipate that the open-source implementation will lower the barrier to adopting physics-augmented neural network constitutive models.

• The paper introduces a differentiable framework that fuses gradient-enhanced continuum damage mechanics with physics-augmented neural networks to simulate damage in soft materials.
• It demonstrates that a monolithic solution scheme with automatic differentiation in JAX-FEM yields superior numerical stability and mesh-independent results.
• The framework flexibly integrates both analytical and data-driven constitutive models, validated by uniaxial tests and notched plate simulations, setting a new benchmark in computational mechanics.

Differentiable Gradient-Enhanced Damage Modeling with Physics-Augmented Neural Networks in JAX-FEM

Formulation and Theoretical Foundations

The paper develops a differentiable framework for modeling damage in soft materials by integrating gradient-enhanced continuum damage mechanics, physics-augmented neural network constitutive models, and automatic differentiation using JAX-FEM (2604.03411). The underlying physical phenomena—manifested in rubbers, hydrogels, and biological tissues—are characterized by hysteretic stress-strain responses and strain softening that drive mesh-dependent failure localization in conventional FE implementations.

The formulation introduces both local ($\kappa$) and non-local ($\phi$) damage variables via an additively split free energy, where non-local terms regularize damage spatially and enforce mesh-objectivity. Damage evolution adheres to thermodynamic consistency: polyconvexity is ensured through ICNN parameterization of the strain energy and monotonic dissipation yield functions, guaranteeing the Clausius–Duhem inequality. The variational framework leads to coupled PDEs for displacement and non-local damage fields, with local damage governed via associative flow rules and KKT conditions.

Implementation in JAX-FEM

The differentiable programming paradigm simplifies constitutive model integration into JAX-FEM. The residual formulation leverages automatic differentiation to compute stress and non-local driving forces directly from specified energy potentials, allowing for seamless substitution between closed-form and neural network-based models. Two solution schemes are discussed: monolithic (simultaneous global update of variables) and staggered (separate evolution of local history variables), with the monolithic scheme yielding superior numerical stability and physical accuracy at high damage states.

The computational treatment of $\kappa$ employs arc-length continuation and Newton-Raphson procedures, while mesh discretization and global assembly are managed by JAX-FEM with direct compatibility for GPU acceleration and scalable design.

Constitutive Model Flexibility

The framework supports both analytical and data-driven models for constitutive behavior. The closed-form neo-Hookean model provides benchmarks, whereas physics-augmented neural networks enable flexible, thermodynamically consistent modeling across materials with markedly different damage characteristics. The ICNNs enforce polyconvexity by architectural design, and the monotonicity of the damage yield function is guaranteed via increasing neural network mappings.

Numerical Validation and Mesh Independence

Robust validation is presented through multiple numerical demonstrations:

• Local Element Response: Uniaxial tension tests on single elements reproduce the expected softening and damage initiation behaviors, showing parameter sensitivity in $\kappa_d$ and $\eta_d$, with damage evolution accurately tracked using arc-length solvers.

  Figure 1: Axial stress $\sigma_{11}$, local damage $\kappa$, and damage degradation $d$ vs. stretch under uniaxial loading.

• Data-Driven Fitting: The ICNN-based model fit to experimental data captures diverse Mullins-type degradation across three material types, showing precise agreement between predicted and measured stress-stretch behavior and confirming the flexibility of the architecture.

  Figure 2: Physics-augmented neural networks fitted to three distinct continuum damage datasets for soft materials.

• Mesh and Solution Strategy Study: Uniform plate tension simulations demonstrate that gradient enhancement eliminates mesh-dependence, even for coarse and non-uniform meshes. Comparison between staggered and monolithic schemes reveals that only the monolithic, gradient-regularized approach achieves damage evolution free from spurious localization.

  Figure 3: Comparison of solution schemes and mesh refinement; gradient-enhanced monolithic formulation achieves mesh-independent damage evolution.

Applied Example: Notched Plate

A notched plate simulation utilizing the data-driven model validates the practical applicability of the framework. The global damage field $\phi$ diffuses over the domain, coupled to $\kappa$ via penalty and gradient terms, preventing damage localization at the notch. The procedure exhibits convergence and stability across mesh refinements and load-step increments, further attesting to the robustness of the differentiable gradient-enhanced framework.

Figure 4: Notched plate simulation convergence with respect to increment steps and mesh refinement; mesh-independent damage observed.

Implications and Future Directions

The integration of differentiable programming, physics-augmented neural networks, and gradient-enhanced continuum damage mechanics offers a modular and efficient paradigm for modeling complex material degradation. The framework is broadly applicable to soft matter and biomedical device design, enabling gradient-based optimization and inverse problems directly within the FE solver by virtue of end-to-end automatic differentiation.

Practical implications include:

• Flexible and accurate simulation of a wide range of soft materials without mesh-dependent artifacts.
• Open-source accessibility and extensibility for data-driven constitutive modeling in scientific machine learning environments.
• Potential for scalable, gradient-based design and parameter identification workflows for damaged materials.

Theoretical implications center on the validation of physics-augmented architectures for enforcing polyconvexity and dissipation constraints, setting a precedent for neural constitutive modeling with intrinsic thermodynamic consistency.

Future developments are anticipated in:

• Extending the framework to anisotropic and tensorial damage models using suitable neural network parameterizations.
• Integration with multi-scale modeling and multi-physics couplings, particularly for biomechanics and advanced manufacturing.
• Incorporation of experimental datasets for automated inverse model discovery and optimization under physically realistic damage scenarios.

Conclusion

This work establishes an authoritative framework for differentiable, mesh-objective, and data-driven damage simulation in soft materials, implemented via JAX-FEM with physically constrained neural constitutive models. The proposed approach yields accurate, flexible, and stable FE simulations devoid of mesh dependence, and facilitates future advancements in both computational mechanics and scientific machine learning for damaged materials.