Skip to content

Commit

Permalink
format scc
Browse files Browse the repository at this point in the history
  • Loading branch information
@ChrisRackauckas
ChrisRackauckas committed Nov 15, 2024
1 parent 1fb625e commit b328dcc
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 42 deletions.
39 changes: 21 additions & 18 deletions lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,30 +4,33 @@ import SciMLBase
import CommonSolve

function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg; kwargs...)
numscc = length(prob.probs)
sols = [SciMLBase.build_solution(prob, nothing, prob.u0, convert(eltype(prob.u0),NaN)*prob.u0) for prob in prob.probs]
u = reduce(vcat,[prob.u0 for prob in prob.probs])
resid = copy(u)
numscc = length(prob.probs)
sols = [SciMLBase.build_solution(
prob, nothing, prob.u0, convert(eltype(prob.u0), NaN) * prob.u0)
for prob in prob.probs]
u = reduce(vcat, [prob.u0 for prob in prob.probs])
resid = copy(u)

lasti = 1
for i in 1:numscc
prob.explictfuns![i](parameter_values(prob.probs[i])[1],sols)
sol = SciMLBase.solve(prob.probs[i], alg; kwargs...)
_sol = SciMLBase.build_solution(prob.probs[i], nothing, sol.u, sol.resid, retcode = sol.retcode)
sols[i] = _sol
lasti = i
for i in 1:numscc
prob.explictfuns![i](parameter_values(prob.probs[i])[1], sols)
sol = SciMLBase.solve(prob.probs[i], alg; kwargs...)
_sol = SciMLBase.build_solution(
prob.probs[i], nothing, sol.u, sol.resid, retcode = sol.retcode)
sols[i] = _sol
lasti = i
if !SciMLBase.successful_retcode(_sol)
break
end
end

# TODO: fix allocations with a lazy concatenation
u .= reduce(vcat,sols)
resid .= reduce(vcat,getproperty.(sols,:resid))
end

retcode = sols[lasti].retcode
# TODO: fix allocations with a lazy concatenation
u .= reduce(vcat, sols)
resid .= reduce(vcat, getproperty.(sols, :resid))

SciMLBase.build_solution(prob, alg, u, resid; retcode, original = sols)
retcode = sols[lasti].retcode

SciMLBase.build_solution(prob, alg, u, resid; retcode, original = sols)
end

end
end
55 changes: 31 additions & 24 deletions lib/SCCNonlinearSolve/test/core_tests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,13 @@ include("../../../common/common_rootfind_testing.jl")
end

@testitem "Manual SCC" setup=[CoreRootfindTesting] tags=[:core] begin
function f(du,u,p)
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[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
Expand All @@ -21,60 +21,67 @@ end
u0 = zeros(2)
cache = zeros(3)

function f1(du,u,(cache,p))
function f1(du, u, (cache, p))
du[1] = cos(u[2]) - u[1]
du[2] = sin(u[1] + u[2]) + u[2]
end
explicitfun1(cache,sols) = nothing
prob1 = NonlinearProblem(NonlinearFunction{true, SciMLBase.NoSpecialize}(f1), zeros(2),(cache,nothing))
explicitfun1(cache, sols) = nothing
prob1 = NonlinearProblem(
NonlinearFunction{true, SciMLBase.NoSpecialize}(f1), zeros(2), (cache, nothing))
sol1 = solve(prob1, NewtonRaphson())

function f2(du,u,(cache,p))
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
explicitfun2(cache,sols) = nothing
prob2 = NonlinearProblem(NonlinearFunction{true, SciMLBase.NoSpecialize}(f2), zeros(3),(cache,nothing))
explicitfun2(cache, sols) = nothing
prob2 = NonlinearProblem(
NonlinearFunction{true, SciMLBase.NoSpecialize}(f2), zeros(3), (cache, nothing))
sol2 = solve(prob2, NewtonRaphson())

function f3(du,u,(cache,p))
function f3(du, u, (cache, p))
du[1] = cache[1] + 2.0u[1] + 2.5u[2] + 1.5u[3]
du[2] = cache[2] + 4.0u[1] - 1.5u[2] + 1.5u[3]
du[3] = cache[3] + + u[1] - u[2] - u[3]
du[3] = cache[3] + +u[1] - u[2] - u[3]
end
prob3 = NonlinearProblem(NonlinearFunction{true, SciMLBase.NoSpecialize}(f3), zeros(3),(cache,nothing))
function explicitfun3(cache,sols)
prob3 = NonlinearProblem(
NonlinearFunction{true, SciMLBase.NoSpecialize}(f3), zeros(3), (cache, nothing))
function explicitfun3(cache, sols)
cache[1] = sols[1][1] + sols[1][2] + sols[2][1] + sols[2][2] + sols[2][3]
cache[2] = sols[1][1] + sols[1][2] + sols[2][1] + 2.0sols[2][2] + sols[2][3]
cache[3] = sols[1][1] + 2.0sols[1][2] + 3.0sols[2][1] + 5.0sols[2][2] + 6.0sols[2][3]
cache[3] = sols[1][1] + 2.0sols[1][2] + 3.0sols[2][1] + 5.0sols[2][2] +
6.0sols[2][3]
end
explicitfun3(cache,[sol1,sol2])
explicitfun3(cache, [sol1, sol2])
sol3 = solve(prob3, NewtonRaphson())
manualscc = [sol1; sol2; sol3]

function SciMLBase.solve(prob::SciMLBase.SCCNonlinearProblem, alg; kwargs...)
numscc = length(prob.probs)
sols = [SciMLBase.build_solution(prob, nothing, prob.u0, convert(eltype(u0),NaN)*prob.u0) for prob in prob.probs]
u = reduce(vcat,[prob.u0 for prob in prob.probs])
sols = [SciMLBase.build_solution(
prob, nothing, prob.u0, convert(eltype(u0), NaN) * prob.u0)
for prob in prob.probs]
u = reduce(vcat, [prob.u0 for prob in prob.probs])
resid = copy(u)

for i in 1:numscc
prob.explictfuns![i](prob.probs[i].p[1],sols)
prob.explictfuns![i](prob.probs[i].p[1], sols)
sol = solve(prob.probs[i], alg)
_sol = SciMLBase.build_solution(prob.probs[i], nothing, sol.u, sol.resid, retcode = sol.retcode)
_sol = SciMLBase.build_solution(
prob.probs[i], nothing, sol.u, sol.resid, retcode = sol.retcode)
sols[i] = _sol
end

# TODO: fix allocations with a lazy concatenation
u .= reduce(vcat,sols)
resid .= reduce(vcat,getproperty.(sols,:resid))
u .= reduce(vcat, sols)
resid .= reduce(vcat, getproperty.(sols, :resid))

SciMLBase.build_solution(prob, alg, u, resid, original = sols)
end

sccprob = SciMLBase.SCCNonlinearProblem([prob1,prob2,prob3], SciMLBase.Void{Any}.([explicitfun1,explicitfun2,explicitfun3]))
sccprob = SciMLBase.SCCNonlinearProblem([prob1, prob2, prob3],
SciMLBase.Void{Any}.([explicitfun1, explicitfun2, explicitfun3]))
scc_sol = solve(sccprob, NewtonRaphson())
@test sol ≈ manualscc ≈ scc_sol

end
end

0 comments on commit b328dcc

Please sign in to comment.