Merge pull request #28 from PALEOtoolkit/split_dae

SteadyState.steadystate_ptc_splitdae solver with inner Newton solve for algebraic constraints

Project.toml

@@ -1,7 +1,7 @@
 name = "PALEOmodel"
 uuid = "bf7b4fbe-ccb1-42c5-83c2-e6e9378b660c"
 authors = ["Stuart Daines <[email protected]>"]
-version = "0.15.10"
+version = "0.15.11"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"

docs/src/PALEOmodelSolvers.md

@@ -39,6 +39,32 @@ print_sol_stats
 calc_output_sol!
 ```
 
+## Steady-state solvers (Julia [`NLsolve`](https://github.com/JuliaNLSolvers/NLsolve.jl) based)
+```@meta
+CurrentModule = PALEOmodel.SteadyState
+```
+```@docs
+steadystate
+steadystate_ptc
+steadystate_ptc_splitdae
+```
+
+## Steady-state solvers (Sundials Kinsol based):
+```@meta
+CurrentModule = PALEOmodel.SteadyStateKinsol
+```
+```@docs
+steadystate_ptc
+```
+```@meta
+CurrentModule = PALEOmodel
+```
+```@docs
+Kinsol
+Kinsol.kin_create
+Kinsol.kin_solve
+```
+
 ## Fixed timestep solvers
 ```@meta
 CurrentModule = PALEOmodel.ODEfixed
@@ -83,31 +109,6 @@ ThreadBarrierAtomic
 ThreadBarrierCond
 ```
 
-## Steady-state solvers (Julia [`NLsolve`](https://github.com/JuliaNLSolvers/NLsolve.jl) based)
-```@meta
-CurrentModule = PALEOmodel.SteadyState
-```
-```@docs
-steadystate
-steadystate_ptc
-```
-
-## Steady-state solvers (Sundials Kinsol based):
-```@meta
-CurrentModule = PALEOmodel.SteadyStateKinsol
-```
-```@docs
-steadystate_ptc
-```
-```@meta
-CurrentModule = PALEOmodel
-```
-```@docs
-Kinsol
-Kinsol.kin_create
-Kinsol.kin_solve
-```
-
 ## Variable aggregation
 ```@meta
 CurrentModule = PALEOmodel

src/JacobianAD.jl

@@ -319,7 +319,7 @@ end
 ####################################################################
 
 """
-    jac_transfer_variables(model, modeldata_ad, modeldata) -> (transfer_data_arrays_ad, transfer_data_arrays)
+    jac_transfer_variables(model, modeldata_ad, modeldata; extra_vars=[]) -> (transfer_data_arrays_ad, transfer_data_arrays)
 
 Build Vectors of data arrays that need to be copied from `modeldata` to `modeldata_ad` before calculating Jacobian.
 (looks for Variables with :transfer_jacobian attribute set).  Only needed if the Jacobian calculation is optimised
@@ -333,11 +333,12 @@ The copy needed is then:
         d_ad .= d
     end
 """
-function jac_transfer_variables(model, modeldata_ad, modeldata)
+function jac_transfer_variables(model, modeldata_ad, modeldata; extra_vars=[])
 
-    # get list of Variables with attribute :transfer_jacobian = true
     transfer_vars = []
+    # get list of Variables with attribute :transfer_jacobian = true    
     for dom in model.domains
+        append!(transfer_vars, PB.get_variables(dom, v -> v.name in extra_vars))
         append!(transfer_vars, PB.get_variables(dom, v -> PB.get_attribute(v, :transfer_jacobian, false)))
     end
 

src/NonLinearNewton.jl

@@ -6,47 +6,62 @@ import LinearAlgebra
     solve(
         func, jac, u0::AbstractVector;
         reltol=1e-5,
+        miniters=0,
         maxiters=100,
         verbose=0,
+        jac_constant::Bool=false,
+        u_min=-Inf,
     ) -> (u, Lnorm_2, Lnorm_inf, niter)
 
 Minimal naive Newton solver for nonlinear function `func(u)` with jacobian `jac(u)`, starting at `u0`.
 
+
 Stopping criteria is `norm(func(u), Inf) < reltol`.
 
 Non-allocating if `u0` and hence `u` and `func(u)` are `StaticArrays.SVector`s.
 
-Set `verbose` to 1, 2, 3 for increasing levels of output.
+Set `verbose` to 1, 2, 3, 4 for increasing levels of output.
 """
 function solve(
     func,
     jac, 
     u0::AbstractVector;
     reltol=1e-5,
+    miniters::Integer=0,
     maxiters::Integer=100,
     verbose::Integer=0,
+    jac_constant::Bool=false,
+    u_min=-Inf,
 )
 
     u = copy(u0)
     residual = func(u0)
     Lnorm_2 = LinearAlgebra.norm(residual, 2)
     Lnorm_inf = LinearAlgebra.norm(residual, Inf)
     iters = 0
-    while Lnorm_inf > reltol && iters < maxiters
-        jacobian = jac(u)
+    verbose >= 1 && @info "iters $iters Lnorm_2 $Lnorm_2 Lnorm_inf $Lnorm_inf u $u residual $residual"
+
+    while (Lnorm_inf > reltol || iters < miniters) && iters < maxiters
+        if !jac_constant || iters == 0
+            global jacobian = jac(u)
+        end
+        verbose >= 4 && @info "iters $iters jac:" jacobian
         u = u - jacobian \ residual
+        if u_min != -Inf
+            u = max.(u, u_min)
+        end
         iters += 1
         residual = func(u)
         Lnorm_2 = LinearAlgebra.norm(residual, 2)
         Lnorm_inf = LinearAlgebra.norm(residual, Inf)
 
         verbose >= 2 && @info "iters $iters Lnorm_2 $Lnorm_2 Lnorm_inf $Lnorm_inf"
-        verbose >= 3 && @info "u $u residual $residual"
+        verbose >= 3 && @info "    u $u residual $residual"
     end
 
-    iters < maxiters || @warn "maxiters $maxiters reached"
+    verbose >= 1 && @info "iters $iters Lnorm_2 $Lnorm_2 Lnorm_inf $Lnorm_inf u $u residual $residual"
 
-    verbose >= 1 && @info "iters $iters Lnorm_2 $Lnorm_2 Lnorm_inf $Lnorm_inf"
+    iters < maxiters || throw(ErrorException("NonLinearNewton.solve: maxiters $maxiters reached"))
 
     return (u, Lnorm_2, Lnorm_inf, iters)
 end

src/ODE.jl

@@ -497,6 +497,21 @@ function calc_output_sol!(outputwriter, model::PB.Model, tsoln, soln,  modeldata
     return nothing
 end
 
+function calc_output_sol!(outputwriter, model::PB.Model, tsoln, soln, modelode, modeldata)
+
+    PALEOmodel.OutputWriters.initialize!(outputwriter, model, modeldata, length(tsoln))
+
+    # call model to (re)calculate
+    du = similar(first(soln)) # not used
+    for i in eachindex(tsoln)
+        tmodel = tsoln[i]     
+        modelode(du, soln[i], nothing, tmodel)
+        PALEOmodel.OutputWriters.add_record!(outputwriter, model, modeldata, tmodel)
+    end
+
+    return nothing
+end
+
 """
     print_sol_stats(sol::SciMLBase.ODESolution)
     print_sol_stats(sol::SciMLBase.DAESolution)

src/PALEOmodel.jl

@@ -35,6 +35,8 @@ include("ODELocalIMEX.jl")
 
 include("Kinsol.jl")
 
+include("SplitDAE.jl")
+
 include("SteadyState.jl")
 
 include("SteadyStateKinsol.jl")