Bugfix: AD for RepeatingAffineFEStateMaps

.gitignore

@@ -19,8 +19,6 @@ deps/src/
 # Build artifacts for creating documentation generated by the Documenter package
 docs/build/
 docs/site/
-Results/*
-results/*
 
 # File generated by Pkg, the package manager, based on a corresponding Project.toml
 # It records a fixed state of all packages used by the project. As such, it should not be

results/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

scripts/_dev/mem_leak.jl

@@ -0,0 +1,114 @@
+using Gridap, LevelSetTopOpt
+
+# FE parameters
+order = 1                                               # Finite element order
+xmax=ymax=zmax=1.0                                      # Domain size
+dom = (0,xmax,0,ymax,0,zmax)                            # Bounding domain
+el_size = (10,10,10)                                    # Mesh partition size
+prop_Γ_N = 0.4                                          # Γ_N size parameter
+prop_Γ_D = 0.2                                          # Γ_D size parameter
+f_Γ_N(x) = (x[1] ≈ xmax) &&                             # Γ_N indicator function
+  (ymax/2-ymax*prop_Γ_N/4 - eps() <= x[2] <= ymax/2+ymax*prop_Γ_N/4 + eps()) &&
+  (zmax/2-zmax*prop_Γ_N/4 - eps() <= x[3] <= zmax/2+zmax*prop_Γ_N/4 + eps())
+f_Γ_D(x) = (x[1] ≈ 0.0) &&                              # Γ_D indicator function
+  (x[2] <= ymax*prop_Γ_D + eps() || x[2] >= ymax-ymax*prop_Γ_D - eps()) &&
+  (x[3] <= zmax*prop_Γ_D + eps() || x[3] >= zmax-zmax*prop_Γ_D - eps())
+# FD parameters
+γ = 0.1                                                 # HJ equation time step coefficient
+γ_reinit = 0.5                                          # Reinit. equation time step coefficient
+max_steps = floor(Int,minimum(el_size)/3)               # Max steps for advection
+tol = 1/(2*order^2)*prod(inv,minimum(el_size))           # Advection tolerance
+# Problem parameters
+κ = 1                                                   # Diffusivity
+g = 1                                                   # Heat flow in
+vf = 0.4                                                # Volume fraction constraint
+lsf_func = initial_lsf(4,0.2)                           # Initial level set function
+iter_mod = 10                                           # Output VTK files every 10th iteration
+
+# Model
+model = CartesianDiscreteModel(dom,el_size);
+update_labels!(1,model,f_Γ_D,"Gamma_D")
+update_labels!(2,model,f_Γ_N,"Gamma_N")
+# Triangulation and measures
+Ω = Triangulation(model)
+Γ_N = BoundaryTriangulation(model,tags="Gamma_N")
+dΩ = Measure(Ω,2*order)
+dΓ_N = Measure(Γ_N,2*order)
+# Spaces
+reffe = ReferenceFE(lagrangian,Float64,order)
+V = TestFESpace(model,reffe;dirichlet_tags=["Gamma_D"])
+U = TrialFESpace(V,0.0)
+V_φ = TestFESpace(model,reffe)
+V_reg = TestFESpace(model,reffe;dirichlet_tags=["Gamma_N"])
+U_reg = TrialFESpace(V_reg,0)
+# Level set and interpolator
+φh = interpolate(lsf_func,V_φ)
+interp = SmoothErsatzMaterialInterpolation(η = 2*maximum(get_el_Δ(model)))
+I,H,DH,ρ = interp.I,interp.H,interp.DH,interp.ρ
+# Weak formulation
+a(u,v,φ,dΩ,dΓ_N) = ∫((I ∘ φ)*κ*∇(u)⋅∇(v))dΩ
+l(v,φ,dΩ,dΓ_N) = ∫(g*v)dΓ_N
+state_map = AffineFEStateMap(a,l,U,V,V_φ,U_reg,φh,dΩ,dΓ_N)
+# Objective and constraints
+J(u,φ,dΩ,dΓ_N) = ∫((I ∘ φ)*κ*∇(u)⋅∇(u))dΩ
+dJ(q,u,φ,dΩ,dΓ_N) = ∫(κ*∇(u)⋅∇(u)*q*(DH ∘ φ)*(norm ∘ ∇(φ)))dΩ
+vol_D = sum(∫(1)dΩ)
+C(u,φ,dΩ,dΓ_N) = ∫(((ρ ∘ φ) - vf)/vol_D)dΩ
+dC(q,u,φ,dΩ,dΓ_N) = ∫(-1/vol_D*q*(DH ∘ φ)*(norm ∘ ∇(φ)))dΩ
+pcfs = PDEConstrainedFunctionals(J,[C],state_map,analytic_dJ=dJ,analytic_dC=[dC])
+# Velocity extension
+α = 4*maximum(get_el_Δ(model))
+a_hilb(p,q) =∫(α^2*∇(p)⋅∇(q) + p*q)dΩ
+vel_ext = VelocityExtension(a_hilb,U_reg,V_reg)
+# Finite difference scheme
+scheme = FirstOrderStencil(length(el_size),Float64)
+stencil = AdvectionStencil(scheme,model,V_φ,tol,max_steps)
+# Optimiser
+optimiser = AugmentedLagrangian(pcfs,stencil,vel_ext,φh;γ,γ_reinit,verbose=true,constraint_names=[:Vol])
+# Solve
+for i = 1:10
+  LevelSetTopOpt.rrule(state_map,φh)
+end
+
+evaluate!(pcfs,φh);
+uh = get_state(state_map)
+
+GC.gc()
+Sys.free_memory() / 2^20
+
+for i = 1:10
+  begin
+    u  = LevelSetTopOpt.forward_solve!(state_map,φh);
+    uh = FEFunction(LevelSetTopOpt.get_trial_space(state_map),u);
+    # LevelSetTopOpt.update_adjoint_caches!(state_map,uh,φh);
+  end;
+end;
+
+Sys.free_memory() / 2^20
+
+GC.gc()
+Sys.free_memory() / 2^20
+
+for i = 1:10
+  begin
+    # u  = LevelSetTopOpt.forward_solve!(state_map,φh);
+    # uh = FEFunction(LevelSetTopOpt.get_trial_space(state_map),u);
+
+    LevelSetTopOpt.update_adjoint_caches!(state_map,uh,φh);
+  end;
+end;
+
+Sys.free_memory() / 2^20
+
+GC.gc()
+Sys.free_memory() / 2^20
+
+for i = 1:10
+  begin
+    u  = LevelSetTopOpt.forward_solve!(state_map,φh);
+    uh = FEFunction(LevelSetTopOpt.get_trial_space(state_map),u);
+    LevelSetTopOpt.update_adjoint_caches!(state_map,uh,φh);
+  end;
+end;
+
+Sys.free_memory() / 2^20

src/ChainRules.jl

@@ -707,10 +707,15 @@ struct RepeatingAffineFEStateMap{A,B,C,D,E,F,G} <: AbstractFEStateMap
     @check nblocks == length(l)
 
     spaces_0 = (U0,V0)
+    assem_U0 = assem_U
+
     biform = IntegrandWithMeasure(a,dΩ)
     liforms = map(li -> IntegrandWithMeasure(li,dΩ),l)
     U, V = repeat_spaces(nblocks,U0,V0)
     spaces = (U,V,V_φ,U_reg)
+    assem_U = SparseMatrixAssembler(
+      get_matrix_type(assem_U0),get_vector_type(assem_U0),U,V,FESpaces.get_assembly_strategy(assem_U0)
+    )
 
     ## Pullback cache
     uhd = zero(U0)
@@ -722,12 +727,12 @@ struct RepeatingAffineFEStateMap{A,B,C,D,E,F,G} <: AbstractFEStateMap
     plb_caches = (dudφ_vec,assem_deriv)
 
     ## Forward cache
-    K  = assemble_matrix((u,v) -> biform(u,v,φh),assem_U,U0,V0)
+    K  = assemble_matrix((u,v) -> biform(u,v,φh),assem_U0,U0,V0)
     b  = allocate_in_range(K); fill!(b,zero(eltype(b)))
     b0 = allocate_in_range(K); fill!(b0,zero(eltype(b0)))
     x  = mortar(map(i -> allocate_in_domain(K), 1:nblocks)); fill!(x,zero(eltype(x)))
     ns = numerical_setup(symbolic_setup(ls,K),K)
-    fwd_caches = (ns,K,b,x,uhd,assem_U,b0)
+    fwd_caches = (ns,K,b,x,uhd,assem_U0,b0,assem_U)
 
     ## Adjoint cache
     adjoint_K  = assemble_matrix((u,v)->biform(v,u,φh),assem_adjoint,V0,U0)
@@ -761,26 +766,26 @@ end
 get_state(m::RepeatingAffineFEStateMap) = FEFunction(get_trial_space(m),m.fwd_caches[4])
 get_measure(m::RepeatingAffineFEStateMap) = m.biform.dΩ
 get_spaces(m::RepeatingAffineFEStateMap) = m.spaces
-get_assemblers(m::RepeatingAffineFEStateMap) = (m.fwd_caches[6],m.plb_caches[2],m.adj_caches[4])
+get_assemblers(m::RepeatingAffineFEStateMap) = (m.fwd_caches[8],m.plb_caches[2],m.adj_caches[4])
 
 function forward_solve!(φ_to_u::RepeatingAffineFEStateMap,φh)
   biform, liforms = φ_to_u.biform, φ_to_u.liform
   U0, V0 = φ_to_u.spaces_0
-  ns, K, b, x, uhd, assem_U, b0 = φ_to_u.fwd_caches
+  ns, K, b, x, uhd, assem_U0, b0, _ = φ_to_u.fwd_caches
 
   a_fwd(u,v) = biform(u,v,φh)
-  assemble_matrix!(a_fwd,K,assem_U,U0,V0)
+  assemble_matrix!(a_fwd,K,assem_U0,U0,V0)
   numerical_setup!(ns,K)
 
   l0_fwd(v) = a_fwd(uhd,v)
-  assemble_vector!(l0_fwd,b0,assem_U,V0)
+  assemble_vector!(l0_fwd,b0,assem_U0,V0)
   rmul!(b0,-1)
 
   v = get_fe_basis(V0)
   map(blocks(x),liforms) do xi, li
     copy!(b,b0)
     vecdata = collect_cell_vector(V0,li(v,φh))
-    assemble_vector_add!(b,assem_U,vecdata)
+    assemble_vector_add!(b,assem_U0,vecdata)
     solve!(xi,ns,b)
   end
   return x
@@ -804,9 +809,9 @@ function update_adjoint_caches!(φ_to_u::RepeatingAffineFEStateMap,uh,φh)
   return φ_to_u.adj_caches
 end
 
-function adjoint_solve!(φ_to_u::RepeatingAffineFEStateMap,du::AbstractVector)
+function adjoint_solve!(φ_to_u::RepeatingAffineFEStateMap,du::AbstractBlockVector)
   adjoint_ns, _, adjoint_x, _ = φ_to_u.adj_caches
-  map(blocks(adjoint_x),du) do xi, dui
+  map(blocks(adjoint_x),blocks(du)) do xi, dui
     solve!(xi,adjoint_ns,dui)
   end
   return adjoint_x