Make the benchmark more complex by passing it through AD

doc/demo/demo_plasticity_mohr_coulomb.py

@@ -73,6 +73,9 @@
 # ### Preamble
 
 # %%
+import os, sys 
+sys.setdlopenflags(os.RTLD_NOW | os.RTLD_GLOBAL)
+
 from mpi4py import MPI
 from petsc4py import PETSc
 
@@ -581,11 +584,12 @@ def C_tang_impl(deps):
 
     unique_iters, counts = jnp.unique(niter, return_counts=True)
 
-    print("\tInner Newton summary:")
-    print(f"\t\tUnique number of iterations: {unique_iters}")
-    print(f"\t\tCounts of unique number of iterations: {counts}")
-    print(f"\t\tMaximum f: {jnp.max(yielding)}")
-    print(f"\t\tMaximum residual: {jnp.max(norm_res)}")
+    if MPI.COMM_WORLD.rank == 0:
+        print("\tInner Newton summary:")
+        print(f"\t\tUnique number of iterations: {unique_iters}")
+        print(f"\t\tCounts of unique number of iterations: {counts}")
+        print(f"\t\tMaximum f: {jnp.max(yielding)}")
+        print(f"\t\tMaximum residual: {jnp.max(norm_res)}")
 
     return C_tang_global.reshape(-1), sigma_global.reshape(-1)
 
@@ -715,7 +719,8 @@ def F_ext(v):
     if len(points_on_process) > 0:
         results[i + 1, :] = (-u.eval(points_on_process, cells)[0], load)
 
-print(f"Slope stability factor: {-q.value[-1]*H/c}")
+if MPI.COMM_WORLD.rank == 0:
+    print(f"Slope stability factor: {-q.value[-1]*H/c}")
 
 # %% [markdown]
 # ## Verification

jax-gpu/constitutive_model.py

@@ -241,4 +241,5 @@ def constitutive_response(sigma_local, sigma_n_local):
     deps_elas = S_elas @ sigma_local
     sigma_corrected, state = return_mapping(deps_elas, sigma_n_local)
     yielding = state[2]
-    return sigma_corrected, yielding
+    return sigma_corrected, (sigma_corrected, yielding)
+

pyproject.toml

@@ -12,8 +12,8 @@ authors = [
 license = { file = "LICENSE" }
 readme = "README.md"
 dependencies = [
-    "fenics-dolfinx>=0.10.0.dev0,<=0.11.0",
-    "fenics-ufl>=2024.3.0.dev0,<=2024.4.0",
+    "fenics-dolfinx>=0.9.0,<=0.11.0",
+    "fenics-ufl>=2024.2.0,<=2024.4.0",
 ]
 
 [project.optional-dependencies]