fix tests

Author: ChrisRackauckas

lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl

@@ -3,6 +3,7 @@ module SCCNonlinearSolve
 import SciMLBase
 import CommonSolve
 import SymbolicIndexingInterface
+import SciMLBase: NonlinearProblem, NonlinearLeastSquaresProblem, LinearProblem
 
 """
     SCCAlg(; nlalg = nothing, linalg = nothing)
@@ -25,41 +26,42 @@ function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem; kwargs...)
     CommonSolve.solve(prob, SCCAlg(nothing, nothing); kwargs...)
 end
 
+probvec(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}) = prob.u0
+probvec(prob::LinearProblem) = prob.b
+
+iteratively_build_sols(alg, sols; kwargs...) = sols
+
+function iteratively_build_sols(alg, sols, (prob, explicitfun), args...; kwargs...)
+    explicitfun(
+        SymbolicIndexingInterface.parameter_values(prob), sols)
+    
+    _sol = if prob isa SciMLBase.LinearProblem
+        sol = SciMLBase.solve(prob, alg.linalg; kwargs...)
+        SciMLBase.build_linear_solution(
+            alg.linalg, sol.u, nothing, nothing, retcode = sol.retcode)
+    else
+        sol = SciMLBase.solve(prob, alg.nlalg; kwargs...)
+        SciMLBase.build_solution(
+            prob, nothing, sol.u, sol.resid, retcode = sol.retcode)
+    end
+
+    iteratively_build_sols(alg, (sols..., _sol), args...)
+end
+
 function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg::SCCAlg; 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)
-
-    lasti = 1
-    for i in 1:numscc
-        prob.explictfuns![i](
-            SymbolicIndexingInterface.parameter_values(prob.probs[i]), sols)
-        
-        if prob.probs[i] isa SciMLBase.LinearProblem
-            sol = SciMLBase.solve(prob.probs[i], alg.linalg; kwargs...)
-            _sol = SciMLBase.build_solution(
-                prob.probs[i], nothing, sol.u, zero(sol.u), retcode = sol.retcode)
-        else
-            sol = SciMLBase.solve(prob.probs[i], alg.nlalg; kwargs...)
-            _sol = SciMLBase.build_solution(
-                prob.probs[i], nothing, sol.u, sol.resid, retcode = sol.retcode)
-        end
-        
-        sols[i] = _sol
-        lasti = i
-        if !SciMLBase.successful_retcode(_sol)
-            break
-        end
-    end
+    sols = iteratively_build_sols(alg, (), zip(prob.probs, prob.explicitfuns!)...; kwargs...) 
 
     # 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))
 
-    retcode = sols[lasti].retcode
+    retcode = if !all(SciMLBase.successful_retcode, sols)
+        idx = findfirst(!SciMLBase.successful_retcode, sols)
+        sols[idx].retcode
+    else
+        SciMLBase.ReturnCode.Success
+    end
 
     SciMLBase.build_solution(prob, alg, u, resid; retcode, original = sols)
 end

lib/SCCNonlinearSolve/test/core_tests.jl

@@ -47,16 +47,16 @@ end
     b3 = p  # b will be updated by explicitfun3
     prob3 = LinearProblem(A3, b3, zeros(3))
     function explicitfun3(p, sols)
-        p[1] = sols[1][1] + sols[1][2] + sols[2][1] + sols[2][2] + sols[2][3]
-        p[2] = sols[1][1] + sols[1][2] + sols[2][1] + 2.0sols[2][2] + sols[2][3]
-        p[3] = sols[1][1] + 2.0sols[1][2] + 3.0sols[2][1] + 5.0sols[2][2] +
-               6.0sols[2][3]
+        p[1] = -(sols[1][1] + sols[1][2] + sols[2][1] + sols[2][2] + sols[2][3])
+        p[2] = -(sols[1][1] + sols[1][2] + sols[2][1] + 2.0sols[2][2] + sols[2][3])
+        p[3] = -(sols[1][1] + 2.0sols[1][2] + 3.0sols[2][1] + 5.0sols[2][2] +
+               6.0sols[2][3])
     end
     explicitfun3(p, [sol1, sol2])
     sol3 = solve(prob3)  # LinearProblem uses default linear solver
-    manualscc = [sol1; sol2; sol3]
+    manualscc = reduce(vcat,(sol1, sol2, sol3))
 
-    sccprob = SciMLBase.SCCNonlinearProblem([prob1, prob2, prob3],
+    sccprob = SciMLBase.SCCNonlinearProblem((prob1, prob2, prob3),
         SciMLBase.Void{Any}.([explicitfun1, explicitfun2, explicitfun3]))
 
     # Test with SCCAlg that handles both nonlinear and linear problems