Strongly Connected Component (SCC) Scheduled Nonlinear Problems

[ChrisRackauckas]

• https://link.springer.com/book/10.1007/0-387-30260-3 Chapter 7
• https://escholarship.org/content/qt80t380tv/qt80t380tv_noSplash_18a25fa49e487cb213e64ecee5b003fb.pdf?t=li54ql
• https://apps.dtic.mil/sti/tr/pdf/ADA317005.pdf

The idea is that instead of solving one giant Newton method, you can instead split it up into successive Newton solves that are then recursed based on the strongly connected components.

The algorithms of ModelingToolkit are already able to put the system in BLT form and do the tearing process to reduce the systems. I think as a general nonlinear solver, it would be good to, instead of baking this into MTK in some form, have a way to be able to express such structural systems with NonlinearSolve.jl and have MTK generate code to that interface.

I could see this happening in a few forms:

1. We could do this on split vectors.
2. We could do this by passing extra structural information for scheduling.

SCCNonlinearFunction((f1,f2,f3), jac = (J1,J2,J3)) would be needed for both. But the question is do you make u0 = [a1,a2,a3] where each a[i] is an array for the SCC[i], or do you make this segmented?

Effectively the question is as follows. If you do this in the segmented way, then f1(resid, a1, p), f2(resid, a1, a2, p), f3(resid, a1, a2, a3, p) is the target, i.e. each has more arguments as you go along. I could see this having some compilation difficulties, so you could just do this as fi(resid,u,p) where u is an array-of-arrays.

Or, at that point, you could just give the whole u array but with a known structure used by the solver that is not enforced in the definition of fi. For this interface, you would still have to give split scheduled functions SCCNonlinearFunction((f1,f2,f3), jac = (J1,J2,J3)), but with sizes of resid being chosen to match the size of the current nonlinear system. You could just create one resid with the maximum size of the strongly connected components, and then just use views of that in each such component. Thus SCCNonlinearProblem(f,u0, lengths, p) would just need to know the length of each SCC, and then schedule the algorithm into the cascading functions from there. This would probably need to impose an AbstractVector structure on u0, whereas the split form would be easier for handling more array types.

The split form could be nicely handled by keeping the cache of arrays of arrays but having a mapping method to_arrarr!(a,u0) that takes a flattened version of the vector into the SCC split structure. So I don't think it would necessarily have more allocations beyond the setup phase.

So of the ways to represent this, I think I'd rank:

1. Array of arrays fi(resid,u,p)
2. Flattened fi(resid,u,p) with sizes
3. f1(resid, a1, p), f2(resid, a1, a2, p), f3(resid, a1, a2, a3, p)

With the last one effectively eliminated for compile time concerns. I like that the first enforces that the BLT form is respected, to an extent. Someone could still modify u, but that's always a concern 😅.

[ChrisRackauckas]

A good case to test this on would be the double pendulum https://sciml.github.io/ModelingToolkitCourse/dev/lectures/lecture8/#Bonus-Demo

[YingboMa]

I like ∀i, fi(resid, u::Vector{<:StateVector}, p) where u[1] is the actual state vector that the solver needs to solve, too. Each u[i] is naturally contiguous from BLT ordering for equations, while the concatenation of u is not contiguous in general, so that rules out Flattened fi(resid, u, p) with sizes.

[ChrisRackauckas]

using NonlinearSolve, ForwardDiff, FunctionWrappersWrappers, ArrayInterface

dualgen(::Type{T}) where {T} = ForwardDiff.Dual{ForwardDiff.Tag{SciMLTag, T}, T, 2}
struct SciMLTag end

function wrapfun_iip(_ff,
        inputs::Tuple{T1, T2}) where {T1, T2, T3}
	ff = SciMLBase.Void(_ff)
    T = eltype(T1)
    dualT = dualgen(T)
    dualT1 = ArrayInterface.promote_eltype(T1, dualT)

    iip_arglists = (Tuple{T1, T1, T2},
        Tuple{dualT1, dualT1, T2})

    iip_returnlists = ntuple(x -> Nothing, 2)

    fwt = map(iip_arglists, iip_returnlists) do A, R
        FunctionWrappersWrappers.FunctionWrappers.FunctionWrapper{R, A}(SciMLBase.Void(ff))
    end
    FunctionWrappersWrappers.FunctionWrappersWrapper{typeof(fwt), false}(fwt)
end

function f(du,u,p)
	du[1] = cos(u[2]) - u[1]
	du[2] = sin(u[1] + u[2]) + u[2]
	du[3] = 2u[4] + u[3] + 1.0
	du[4] = u[5]^2 + u[4]
	du[5] = u[3]^2 + u[5]
	du[6] = u[1] + u[2] + u[3] + u[4] + u[5]    + 2.0u[6] + 2.5u[7] + 1.5u[8]
	du[7] = u[1] + u[2] + u[3] + 2.0u[4] + u[5] + 4.0u[6] - 1.5u[7] + 1.5u[8]
	du[8] = u[1] + 2.0u[2] + 3.0u[3] + 5.0u[4] + 6.0u[5] + u[6] - u[7] - u[8]
end
prob = NonlinearProblem(f, zeros(8))
sol = solve(prob)

u0 = zeros(2)
p = (zeros(0), nothing)
function f1(du,u,(cache,p))
	du[1] = cos(u[2]) - u[1]
	du[2] = sin(u[1] + u[2]) + u[2]
end
prob1 = NonlinearProblem(NonlinearFunction{true}(wrapfun_iip(f1, (u0, p))), zeros(2),(zeros(0),nothing))
sol1 = solve(prob1, NewtonRaphson(autodiff=AutoFiniteDiff()))

function f2(du,u,(cache,p))
	du[1] = 2u[2] + u[1] + 1.0
	du[2] = u[3]^2 + u[2]
	du[3] = u[1]^2 + u[3]
end
prob2 = NonlinearProblem(NonlinearFunction{true}(wrapfun_iip(f2, (u0, p))), zeros(3),(zeros(2),nothing))
sol2 = solve(prob2, NewtonRaphson(autodiff=AutoFiniteDiff()))

function f3(du,u,(uprev,p))
	du[1] = uprev[1] + uprev[2] + uprev[3] + uprev[4] + uprev[5]             + 2.0u[1] + 2.5u[2] + 1.5u[3]
	du[2] = uprev[1] + uprev[2] + uprev[3] + 2.0uprev[4] + uprev[5]          + 4.0u[1] - 1.5u[2] + 1.5u[3]
	du[3] = uprev[1] + 2.0uprev[2] + 3.0uprev[3] + 5.0uprev[4] + 6.0uprev[5] + u[1] - u[2] - u[3]
end
prob3 = NonlinearProblem(NonlinearFunction{true}(wrapfun_iip(f3, (u0, p))), zeros(3),(zeros(3),nothing))
sol3 = solve(prob3, NewtonRaphson(autodiff=AutoFiniteDiff()))

function SciMLBase.solve(prob::SciMLBase.SCCNonlinearProblem, alg; kwargs...)
	numscc = length(prob.probs)
	sols = Vector{SciMLBase.NonlinearSolution}(undef,numscc)
	cache = prob.allocate_cache(prob)
	for i in 1:numscc
		uprev = reduce(vcat,sols[1:i-1],init=eltype(eltype(prob.probs[i].u0))[])
		prob.probs[i].p[1] .= uprev
		prob.probs[i].p[2] = cache
		explictfun![i](cache,sols[1:i-1])
		sols[i] = solve(prob.probs[i], alg)
	end
	prob.free_cache(cache)
	SciMLBase.build_solution(prob, alg, reduce(vcat,sols), reduce(vcat,getproperty.(sols,:resid)),
							 original = sols)
end 

allocate_cache(prob) = zeros(3)
sccprob = SCCNonlinearProblem((prob1,prob2,prob3))

scc_sol = solve(sccprob, NewtonRaphson(autodiff=AutoFiniteDiff()))
sol ≈ scc_sol

The function wrappers are rough but we can specialize differently.