2D Elastoplasticity FEM (#18)

* transfer changes from original branch

* rewrite all f-strings safely

* Modular FEM + 2D and 3D elastoplasticity with unit test

* Saint_Venant_Material_Model

* Improved_constitutive_material_model

* Updates_on_comments

* fix SaintVenan import issue

---------

Co-authored-by: RezaNajian <r.najian@hotmail.com>
Co-authored-by: Rishabhgit config --global user.email rishab.arora236@gmail.co.comgit config --global user.name Rishabh <rishab.arora236@gmail.com>
Co-authored-by: RezaNajian <62375973+RezaNajian@users.noreply.github.com>

Author: 3 people

examples/elastoplasticity/mechanical_2d_elastoplasticity.py

@@ -0,0 +1,124 @@
+import sys
+import os
+import shutil
+import jax
+import numpy as np
+
+from fol.loss_functions.mechanical_elastoplasticity import ElastoplasticityLoss3DTetra
+from fol.controls.fourier_control import FourierControl
+from fol.mesh_input_output.mesh import Mesh
+from fol.controls.voronoi_control2D import VoronoiControl2D
+from fol.solvers.fe_nonlinear_residual_based_solver_with_history_update import FiniteElementNonLinearResidualBasedSolverWithStateUpdate
+from fol.tools.usefull_functions import *
+from fol.tools.logging_functions import Logger
+import matplotlib.pyplot as plt
+import pickle
+
+
+def main(solve_FE=True, clean_dir=False):
+    # directory & save handling
+    working_directory_name = "box_3D_tetra"
+    case_dir = os.path.join('.', working_directory_name)
+    create_clean_directory(working_directory_name)
+    sys.stdout = Logger(os.path.join(case_dir,working_directory_name+".log"))
+
+    # create mesh_io
+    fe_mesh = Mesh("fol_io","box_3D_coarse.med",'../meshes/')
+
+    # create fe-based loss function
+    bc_dict = {"Ux":{"left":0.0},
+                "Uy":{"left":0.0,"right":-0.06},
+                "Uz":{"left":0.0}}
+    
+    # fourier control
+    fourier_control_settings = {"x_freqs":np.array([2,4,6]),"y_freqs":np.array([2,4,6]),"z_freqs":np.array([2,4,6]),
+                                "beta":20,"min":1,"max":1}
+    fourier_control = FourierControl("fourier_control",fourier_control_settings,fe_mesh)
+
+
+    material_dict = {"young_modulus": 3.0, "poisson_ratio": 0.3, "iso_hardening_parameter_1": 0.4, "iso_hardening_param_2" :10.0, "yield_limit" :0.2}
+    mechanical_loss_3d = ElastoplasticityLoss3DTetra("mechanical_loss_3d",loss_settings={"dirichlet_bc_dict":bc_dict,
+                                                                                   "material_dict":material_dict},
+                                                                                   fe_mesh=fe_mesh)
+
+    fe_mesh.Initialize()
+    mechanical_loss_3d.Initialize()
+    fourier_control.Initialize()
+    create_random_coefficients = True
+    if create_random_coefficients:
+        number_of_random_samples = 200
+        coeffs_matrix,K_matrix = create_random_fourier_samples(fourier_control,number_of_random_samples)
+        export_dict = {}
+        export_dict["coeffs_matrix"] = coeffs_matrix
+        export_dict["x_freqs"] = fourier_control.x_freqs
+        export_dict["y_freqs"] = fourier_control.y_freqs
+        export_dict["z_freqs"] = fourier_control.z_freqs
+        with open(f'fourier_control_dict.pkl', 'wb') as f:
+            pickle.dump(export_dict,f)
+    else:
+        with open(f'fourier_control_dict.pkl', 'rb') as f:
+            loaded_dict = pickle.load(f)
+        
+        coeffs_matrix = loaded_dict["coeffs_matrix"]
+
+    K_matrix = fourier_control.ComputeBatchControlledVariables(coeffs_matrix)
+
+    # now save K matrix 
+    export_Ks = False
+    if export_Ks:
+        for i in range(K_matrix.shape[0]):
+            fe_mesh[f'K_{i}'] = np.array(K_matrix[i,:])
+        fe_mesh.Finalize(export_dir=case_dir)
+
+    eval_id = 69
+    fe_mesh['K'] = np.array(K_matrix[eval_id,:])
+
+
+    # choose which sample to evaluate
+
+    # classical FE solve (no ML)
+    if solve_FE:
+        fe_setting = {
+            "linear_solver_settings": {
+                "solver": "JAX-direct",
+                "tol": 1e-6,
+                "atol": 1e-6,
+                "maxiter": 1000,
+                "pre-conditioner": "ilu"
+            },
+            "nonlinear_solver_settings": {
+                "rel_tol": 1e-5,
+                "abs_tol": 1e-5,
+                "maxiter": 100,
+                "load_incr": 10
+            }
+        }
+
+        nonlinear_fe_solver = FiniteElementNonLinearResidualBasedSolverWithStateUpdate(
+            "nonlinear_fe_solver",
+            mechanical_loss_3d,
+            fe_setting,
+            history_plot_settings={"plot":True,"save_directory":case_dir}
+        )
+        nonlinear_fe_solver.Initialize()
+
+        # Solve for the chosen K-field and zero initial guess
+        load_steps_solutions, load_steps_states, solution_history_dict = nonlinear_fe_solver.Solve(
+                        K_matrix[eval_id], np.zeros(3 * fe_mesh.GetNumberOfNodes()),return_all_steps=True)
+        
+        n_incr = fe_setting["nonlinear_solver_settings"]["load_incr"]
+        FE_UVW=load_steps_solutions[n_incr-1,:]
+        fe_mesh['U_FE'] = FE_UVW.reshape((fe_mesh.GetNumberOfNodes(), 3))
+        
+    # finalize and export mesh data
+    fe_mesh.Finalize(export_dir=case_dir)
+
+    if clean_dir:
+        shutil.rmtree(case_dir)
+
+if __name__ == "__main__":
+    # Defaults
+    solve_FE = True
+    clean_dir = False
+
+    main(solve_FE, clean_dir)

examples/elastoplasticity/mechanical_3d_elastoplasticity.py

@@ -0,0 +1,27 @@
+"""
+ Authors: Rishabh Arora, https://github.com/rishabharora236-cell
+ Date: Dec, 2025
+ License: FOL/LICENSE
+"""
+from abc import ABC, abstractmethod
+from typing import Optional, Any, Tuple
+import jax.numpy as jnp
+from jax import Array
+
+class BaseConstitutiveModel(ABC):
+    """
+    Abstract base class for constitutive models.
+    Minimal interface - subclasses define their own evaluate signature.
+    """
+    
+    @abstractmethod
+    def evaluate(self, *args, **kwargs) -> Tuple:
+        """
+        Evaluate constitutive relation.
+        
+        Returns vary by material type:
+        - Hyperelastic: (energy, stress, tangent)
+        - Small-strain plastic: (stress, new_state)
+        - Large-strain plastic: (stress, new_state)
+        """
+        pass

examples/elastoplasticity/voroni_control_dict.pkl

@@ -0,0 +1,201 @@
+import jax
+import jax.numpy as jnp
+from .base import BaseConstitutiveModel
+from .utils import TensorOperations as TO
+from .utils import TensorVoigtArray as TVA
+
+
+
+# -----------------------------------------
+class NeoHookianModel2D(BaseConstitutiveModel):
+    """
+    Material model.
+    """
+    def evaluate(self, F, k, mu):
+        """
+        Evaluate the stress and tangent operator at given local coordinates.
+        This method should be overridden by subclasses.
+
+        Parameters:
+        F (ndarray): Deformation gradient.
+        args (float): Optional material constants
+
+        Returns:
+        jnp.ndarray: Values of stress and tangent operator at given local coordinates.
+        """
+        # Supporting functions:
+
+        C = jnp.dot(F.T,F)
+        invC = jnp.linalg.inv(C)
+        J = jnp.linalg.det(F)
+        p = 0.5*k*(J-(1/J))
+        dp_dJ = 0.5*k*(1 + J**(-2))
+
+        # Strain Energy
+        xsie_vol = (k/4)*(J**2 - 2*jnp.log(J) -1)
+        I1_bar = (J**(-2/2))*jnp.trace(C)
+        xsie_iso = 0.5*mu*(I1_bar - 2)
+        xsie = xsie_vol + xsie_iso
+
+        # Stress Tensor
+        S_vol = J*p*invC
+        I_fourth = TO.fourth_order_identity_tensor(C.shape[0])
+        P = I_fourth - (1/2)*jnp.einsum('ij,kl->ijkl', invC, C)
+        S_bar = mu*jnp.eye(C.shape[0])
+        S_iso = (J**(-2/2))*jnp.einsum('ijkl,kl->ij',P,S_bar)
+        Se = S_vol + S_iso
+
+        C_ = jnp.einsum('ij,kl->ijkl',jnp.zeros(C.shape),jnp.zeros(C.shape))
+        P_double_C = jnp.einsum('ijkl,klpq->ijpq',P,C_)
+        P_bar = TO.diad_special(invC,invC,invC.shape[0]) - (1/2)*jnp.einsum('ij,kl->ijkl',invC,invC)
+        C_vol = (J*p + dp_dJ*J**2)*jnp.einsum('ij,kl->ijkl',invC,invC) - 2*J*p*TO.diad_special(invC,invC,invC.shape[0])
+        C_iso = jnp.einsum('ijkl,pqkl->ijpq',P_double_C,P) + \
+                (2/2)*(J**(-2/2))*jnp.vdot(S_bar,C)*P_bar - \
+                (2/2)*(jnp.einsum('ij,kl->ijkl',invC,S_iso) + jnp.einsum('ij,kl->ijkl',S_iso,invC))
+        C_tangent_fourth = C_vol + C_iso
+        Se_voigt = TVA.TensorToVoigt(Se)
+        C_tangent = TVA.FourthTensorToVoigt(C_tangent_fourth)
+        return xsie, Se_voigt, C_tangent
+    
+class NeoHookianModel(BaseConstitutiveModel):
+    """
+    Material model.
+    """
+    def evaluate(self, F, k, mu):
+        """
+        Evaluate the stress and tangent operator at given local coordinates.
+        This method should be overridden by subclasses.
+
+        Parameters:
+        F (ndarray): Deformation gradient.
+        args (float): Optional material constants
+
+        Returns:
+        jnp.ndarray: Values of stress and tangent operator at given local coordinates.
+        """
+        # Supporting functions:
+
+        C = jnp.dot(F.T,F)
+        invC = jnp.linalg.inv(C)
+        J = jnp.linalg.det(F)
+        ph = 0.5*k*(J-(1/J))
+        #dp_dJ = (k/4)*(2 + 2*J**(-2))
+        dp_dJ = 0.5*k*(1 + J**(-2))
+
+        # Strain Energy
+        xsie_vol = (k/4)*(J**2 - 2*jnp.log(J) -1)
+        I1_bar = (J**(-2/3))*jnp.trace(C)
+        xsie_iso = 0.5*mu*(I1_bar - 3)
+        xsie = xsie_vol + xsie_iso
+
+        # Stress Tensor
+        S_vol = J*ph*invC
+        I_fourth = TO.fourth_order_identity_tensor(C.shape[0])
+        P = I_fourth - (1/3)*jnp.einsum('ij,kl->ijkl', invC, C)
+        S_bar = mu*jnp.eye(C.shape[0])
+        S_iso = (J**(-2/3))*jnp.einsum('ijkl,kl->ij',P,S_bar)
+        Se = S_vol + S_iso
+
+        C_ = jnp.einsum('ij,kl->ijkl',jnp.zeros(C.shape),jnp.zeros(C.shape))
+        P_double_C = jnp.einsum('ijkl,klpq->ijpq',P,C_)
+        P_bar = TO.diad_special(invC,invC,invC.shape[0]) - (1/3)*jnp.einsum('ij,kl->ijkl',invC,invC)
+        C_vol = (J*ph + dp_dJ*J**2)*jnp.einsum('ij,kl->ijkl',invC,invC) - 2*J*ph*TO.diad_special(invC,invC,invC.shape[0])
+        C_iso = jnp.einsum('ijkl,pqkl->ijpq',P_double_C,P) + \
+                (2/3)*(J**(-2/3))*jnp.vdot(S_bar,C)*P_bar - \
+                (2/3)*(jnp.einsum('ij,kl->ijkl',invC,S_iso) + jnp.einsum('ij,kl->ijkl',S_iso,invC))
+        C_tangent_fourth = C_vol + C_iso
+        Se_voigt = TVA.TensorToVoigt(Se)
+        C_tangent = TVA.FourthTensorToVoigt(C_tangent_fourth)
+        return xsie, Se_voigt, C_tangent
+        
+    
+class NeoHookianModelAD(BaseConstitutiveModel):
+    """
+    Material model.
+    """
+    def evaluate(self, C_mat, k, mu, lambda_, *args, **keyargs):
+        """
+        Evaluate the stress and tangent operator at given local coordinates.
+        This method should be overridden by subclasses.
+
+        Parameters:
+        F (ndarray): Deformation gradient.
+        args (float): Optional material constants
+
+        Returns:
+        jnp.ndarray: Values of stress and tangent operator at given local coordinates.
+        """
+
+        def strain_energy(C_voigt):
+            C = TVA.VoigtToTensor(C_voigt)
+            J = jnp.sqrt(jnp.linalg.det(C))
+            xsie_vol = (k/4)*(J**2 - 2*jnp.log(J) -1)
+            I1_bar = (J**(-2/3))*jnp.trace(C)
+            xsie_iso = 0.5*mu*(I1_bar - 3)
+            return 0.5*mu*(I1_bar - 3) - mu*jnp.log(J) + (lambda_/2)*(jnp.log(J))**2
+        
+        def strain_energy_paper(C_voigt):
+            C = TVA.VoigtToTensor(C_voigt)
+            J = jnp.sqrt(jnp.linalg.det(C))
+            xsie_vol = (k/4)*(J**2 - 2*jnp.log(J) -1)
+            I1_bar = (J**(-2/3))*jnp.trace(C)
+            xsie_iso = 0.5*mu*(I1_bar - 3)
+            return xsie_vol + xsie_iso
+        
+        def second_piola(C_voigt):
+            return 2*jax.grad(strain_energy)(C_voigt)
+        
+        def tangent(C_voigt):
+            return 2*jax.jacfwd(second_piola)(C_voigt)
+        
+        C_voigt = TVA.TensorToVoigt(C_mat)
+
+        xsie = strain_energy(C_voigt)
+        Se_voigt = second_piola(C_voigt)
+        C_tangent = tangent(C_voigt)
+
+        return xsie, Se_voigt, C_tangent.squeeze()
+    
+class NeoHookianModel2DAD(BaseConstitutiveModel):
+    """
+    Material model.
+    """
+    def evaluate(self, C_mat, k, mu, lambda_, *args, **keyargs):
+        """
+        Evaluate the stress and tangent operator at given local coordinates.
+        This method should be overridden by subclasses.
+
+        Parameters:
+        F (ndarray): Deformation gradient.
+        args (float): Optional material constants
+
+        Returns:
+        jnp.ndarray: Values of stress and tangent operator at given local coordinates.
+        """
+        # Supporting functions:
+        # Strain Energy
+
+        def strain_energy(C_voigt):
+            C = TVA.VoigtToTensor(C_voigt)
+            J = jnp.sqrt(jnp.linalg.det(C))
+            return 0.5*mu*(jnp.linalg.trace(C) - 2) - mu*jnp.log(J) + 0.5*lambda_*(jnp.log(J)**2)
+
+        
+        def strain_energy_paper(C_voigt):
+            C = TVA.VoigtToTensor(C_voigt)
+            J = jnp.sqrt(jnp.linalg.det(C))
+            return (k/4)*(J**2 - 2*jnp.log(J) -1) + 0.5*mu*((J**(-2/2))*jnp.trace(C) - 2)
+        
+        def second_piola(C_voigt):
+            return 2*jax.grad(strain_energy)(C_voigt)
+        
+        def tangent(C_voigt):
+            return 2*jax.jacfwd(second_piola)(C_voigt)
+        
+        # C_mat = jnp.dot(F.T,F)
+        C_voigt = TVA.TensorToVoigt(C_mat)
+
+        xsie = strain_energy(C_voigt)
+        Se_voigt = second_piola(C_voigt)
+        C_tangent = tangent(C_voigt)
+        return xsie, Se_voigt, C_tangent.squeeze()