Preallocate some buffers for local solve

Author: termi-official

src/modeling/functions.jl

@@ -92,6 +92,10 @@ struct QuasiStaticFunction{I <: NonlinearIntegrator, DH <: Ferrite.AbstractDofHa
 end
 
 solution_size(f::QuasiStaticFunction) = ndofs(f.dh)+ndofs(f.lvh)
+function local_function_size(f::QuasiStaticFunction)
+    length(f.lvh.dh.subdofhandlers) == 0 && return 0
+    return sum(Ferrite.n_components.(f.lvh.dh.subdofhandlers[1].field_interpolations); init=0)
+end
 function default_initial_condition!(u::AbstractVector, f::QuasiStaticFunction)
     fill!(u, 0.0)
     ndofs(f.lvh) == 0 && return # no internal variable

src/modeling/solid/materials.jl

@@ -223,7 +223,7 @@ setup_internal_cache(material_model::Union{<:ElastodynamicsModel{<:ActiveStressM
 
 # TODO this actually belongs to the multi-level newton file :)
 # Dual (global cache and element-level cache) use for now to make it non-allocating.
-struct GenericFirstOrderRateIndependentCondensationMaterialStateCache{LocalModelType, LocalModelCacheType, QType, QType2, T, LVH} <: RateIndependentCondensationMaterialStateCache
+struct GenericFirstOrderRateIndependentCondensationMaterialStateCache{LocalModelType, LocalModelCacheType, LocalSolverType, QType, QType2, T, LVH} <: RateIndependentCondensationMaterialStateCache
     # The actual model
     model::LocalModelType
     model_cache::LocalModelCacheType
@@ -233,7 +233,7 @@ struct GenericFirstOrderRateIndependentCondensationMaterialStateCache{LocalModel
     Qprev::QType
     # t - tprev
     Δt::T
-    # local_solver::...?
+    local_solver_cache::LocalSolverType
     lvh::LVH
     # These are used locally
     localQ::QType2
@@ -274,24 +274,23 @@ function solve_local_constraint(F::Tensor{2,dim}, coefficients, material_model::
         function local_residual!(R, Q, λ, dλdt)
             dQ = zeros(eltype(Q), length(Q)) # TODO preallocate during setup
             sarcomere_rhs!(dQ, Q, λ, dλdt, Ca, time, material_model.contraction_model)
-            R .= (Q .- Qprev) ./ state_cache.Δt .- dQ
+            @.. R = (Q - Qprev) / state_cache.Δt - dQ
             return nothing
         end
 
         function local_residual_jac_wrap!(R, Q)
             return local_residual!(R, Q, λ, dλdt)
         end
 
-        # TODO preallocate during setup
-        R  = zeros(length(Q))
-        J  = zeros(length(Q),length(Q))
-        for newton_iter in 1:10
+        R  = state_cache.local_solver_cache.residual
+        J  = state_cache.local_solver_cache.J
+        for newton_iter in 1:state_cache.local_solver_cache.params.max_iters
             ForwardDiff.jacobian!(J, local_residual_jac_wrap!, R, Q)
             local_residual!(R, Q, λ, dλdt)
             ΔQ = J \ R
             Q .-= ΔQ
             # @info qp.i, norm(R), norm(ΔQ)
-            if norm(R) < 1e-16
+            if norm(R) < state_cache.local_solver_cache.params.tol
                 break
             elseif newton_iter == 10
                 error("Local Newton did not converge")
@@ -307,23 +306,16 @@ function solve_local_constraint(F::Tensor{2,dim}, coefficients, material_model::
     _store_local_state!(state_cache, geometry_cache, qp)
 
     # Solve corrector problem
-    # function local_residual_rhs_wrap!(R, λ)
-    #     dQ = zeros(eltype(λ), 20) # TODO preallocate during setup
-    #     sarcomere_rhs!(dQ, Q, λ, dλdt, Ca, time, material_model.contraction_model)
-    #     # R .= (Q .- Qprev) .- dQ
-    #     R .= (Q .- Qprevflat) .- dQ
-    #     return nothing
-    # end
-    function local_residual_rhs_wrap(λ)
+    function local_residual_rhs_wrap!(R, λ)
         dQ = zeros(eltype(λ), length(Q)) # TODO preallocate during setup
         sarcomere_rhs!(dQ, Q, λ, dλdt, Ca, time, material_model.contraction_model)
-        # R .= (Q .- Qprev) .- dQ
-        return (Q .- Qprevflat) ./ state_cache.Δt .- dQ
+        @.. R = (Q - Qprevflat) / state_cache.Δt - dQ
         return nothing
     end
-    ∂fₗ∂λ = zeros(length(Q))
-    ForwardDiff.derivative!(∂fₗ∂λ, local_residual_rhs_wrap, λ)
-    dQdλ = - J \ ∂fₗ∂λ
+    R     = state_cache.local_solver_cache.residual
+    ∂fₗ∂λ = state_cache.local_solver_cache.rhs_corrector
+    ForwardDiff.derivative!(∂fₗ∂λ, local_residual_rhs_wrap!, R, λ)
+    dQdλ = J \ -∂fₗ∂λ
 
     return Q, _solve_local_sarcomere_dQdF(dQdλ, dλdF, λ, F, coefficients, material_model.active_stress_model, material_model.contraction_model)
 end

src/solver/nonlinear/multilevel_newton_raphson.jl

@@ -1,11 +1,13 @@
-struct GenericLocalNonlinearSolver <: AbstractNonlinearSolver
+Base.@kwdef struct GenericLocalNonlinearSolver <: AbstractNonlinearSolver
+    max_iters::Int = 10
+    tol::Float64 = 1e-16
 end
 
-struct GenericLocalNonlinearSolverCache{JacobianType, ResidualType, CorrectorType}
+struct GenericLocalNonlinearSolverCache{JacobianType, ResidualType, CorrectorRhsType}
+    params::GenericLocalNonlinearSolver
     J::JacobianType
     residual::ResidualType
-    corrector_rhs::CorrectorType
-    corrector::CorrectorType
+    rhs_corrector::CorrectorRhsType
 end
 
 """

src/solver/time/euler.jl

@@ -172,7 +172,7 @@ mutable struct BackwardEulerStageFunctionWrapper{F,U,T,S, LVH} <: AbstractTimeDi
     const u::U
     const uprev::U
     Δt::T
-    const inner_solver::S
+    const local_solver_cache::S
     const lvh::LVH
 end
 
@@ -193,6 +193,14 @@ end
     @unpack volume_model, face_model = integrator
     @unpack local_solver, newton = solver
 
+    singleQsize = local_function_size(f)
+    local_solver_cache = GenericLocalNonlinearSolverCache(
+        solver.local_solver,
+        zeros(singleQsize, singleQsize),
+        zeros(singleQsize),
+        zeros(singleQsize),
+    )
+
     # Extract condensable parts
     Q     = @view wrapper.u[(ndofs(dh)+1):end]
     Qprev = @view wrapper.uprev[(ndofs(dh)+1):end]
@@ -201,7 +209,7 @@ end
         volume_model,
         Q, Qprev,
         0.0,
-        nothing, # inner_solver_cache.local_solver_cache,
+        local_solver_cache,
         f.lvh,
     )
     face_wrapper = BackwardEulerStageFunctionWrapper(
@@ -231,9 +239,8 @@ end
     newton_cache = NewtonRaphsonSolverCache(op, residual, newton, inner_cache, T[], 0)
 
     return MultiLevelNewtonRaphsonSolverCache(
-        # FIXME global_f and local_f :)
         newton_cache, # setup_solver_cache(G, solver.newton),
-        nothing, #setup_solver_cache(L, solver.local_newton), # FIXME pass
+        local_solver_cache, #setup_solver_cache(L, solver.local_newton), # FIXME pass
     )
 end
 
@@ -315,6 +322,7 @@ function setup_internal_cache(wrapper::BackwardEulerStageFunctionWrapper{<:Quasi
         wrapper.u,
         wrapper.uprev,
         wrapper.Δt,
+        wrapper.local_solver_cache,
         wrapper.lvh,
         zeros(n_ivs_per_qp),
         zeros(n_ivs_per_qp),