Speeding up the stress tensor jacobian using ForwardDiff per element (#180)

* speedup stress tensor jacobian using ForwardDiff per element

* explicitly define variables as global in stress example

* bump version

Author: mohdibntarek

Project.toml

@@ -1,7 +1,7 @@
 name = "TopOpt"
 uuid = "53a1e1a5-51bb-58a9-8a02-02056cc81109"
 authors = ["mohamed82008 <[email protected]>", "yijiangh <[email protected]>"]
-version = "0.9.1"
+version = "0.9.2"
 
 [deps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"

src/Functions/stress_tensor.jl

@@ -55,10 +55,15 @@ function (f::ElementStressTensor)(u::DisplacementResult; element_dofs=false)
     st = f.stress_tensor
     reinit!(f, f.cellidx)
     if element_dofs
-        cellu = u.u
+        return _element_stres_tensor(f, u)
     else
-        cellu = u.u[copy(st.global_dofs)]
+        return _element_stress_tensor(f, DisplacementResult(u.u[copy(st.global_dofs)]))
     end
+end
+
+function _element_stress_tensor(f::ElementStressTensor, u::DisplacementResult)
+    st = f.stress_tensor
+    cellu = u.u
     n_basefuncs = getnbasefunctions(st.cellvalues)
     n_quad = getnquadpoints(st.cellvalues)
     dim = TopOptProblems.getdim(st.problem)
@@ -76,6 +81,18 @@ function (f::ElementStressTensor)(u::DisplacementResult; element_dofs=false)
     ) ./ V
 end
 
+function ChainRulesCore.rrule(
+    ::typeof(_element_stress_tensor), f::ElementStressTensor, u::DisplacementResult
+)
+    J = ForwardDiff.jacobian(
+        vec ∘ (u -> _element_stress_tensor(f, DisplacementResult(u))), u.u
+    )
+    return _element_stress_tensor(f, u),
+    Δ -> begin
+        NoTangent(), NoTangent(), Tangent{typeof(u)}(; u=J' * vec(Δ))
+    end
+end
+
 struct ElementStressTensorKernel{T,Tc} <: AbstractFunction{T}
     E::T
     ν::T

test/examples/local_stress.jl

@@ -39,7 +39,7 @@ threshold = 3 * maximum(stress(filter(PseudoDensities(x0))))
 x = copy(x0)
 x .= 1
 for p in 1.0:1.0:3.0
-    global x
+    global x, solver, stress, filter, volfrac, obj, constr, alg, options, model, r
     penalty = TopOpt.PowerPenalty(p)
     # Define a finite element solver
     solver = FEASolver(Direct, problem; xmin=xmin, penalty=penalty)