Support LinearProblems in SCCs

Author: ChrisRackauckas

lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl

@@ -4,11 +4,28 @@ import SciMLBase
 import CommonSolve
 import SymbolicIndexingInterface
 
+"""
+    SCCAlg(; nlalg = nothing, linalg = nothing)
+
+Algorithm for solving Strongly Connected Component (SCC) problems containing
+both nonlinear and linear subproblems.
+
+### Keyword Arguments
+- `nlalg`: Algorithm to use for solving NonlinearProblem components
+- `linalg`: Algorithm to use for solving LinearProblem components
+"""
+struct SCCAlg{N, L}
+    nlalg::N
+    linalg::L
+end
+
+SCCAlg(; nlalg = nothing, linalg = nothing) = SCCAlg(nlalg, linalg)
+
 function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem; kwargs...)
-    CommonSolve.solve(prob, nothing; kwargs...)
+    CommonSolve.solve(prob, SCCAlg(nothing, nothing); kwargs...)
 end
 
-function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg; kwargs...)
+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)
@@ -20,9 +37,17 @@ function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg; kwargs...)
     for i in 1:numscc
         prob.explictfuns![i](
             SymbolicIndexingInterface.parameter_values(prob.probs[i]), sols)
-        sol = SciMLBase.solve(prob.probs[i], alg; kwargs...)
-        _sol = SciMLBase.build_solution(
-            prob.probs[i], nothing, sol.u, sol.resid, retcode = sol.retcode)
+        
+        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)
@@ -39,4 +64,5 @@ function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg; kwargs...)
     SciMLBase.build_solution(prob, alg, u, resid; retcode, original = sols)
 end
 
+
 end

lib/SCCNonlinearSolve/test/core_tests.jl

@@ -41,30 +41,33 @@ end
         NonlinearFunction{true, SciMLBase.NoSpecialize}(f2), zeros(3), p)
     sol2 = solve(prob2, NewtonRaphson())
 
-    function f3(du, u, p)
-        du[1] = p[1] + 2.0u[1] + 2.5u[2] + 1.5u[3]
-        du[2] = p[2] + 4.0u[1] - 1.5u[2] + 1.5u[3]
-        du[3] = p[3] + +u[1] - u[2] - u[3]
-    end
-    prob3 = NonlinearProblem(
-        NonlinearFunction{true, SciMLBase.NoSpecialize}(f3), zeros(3), p)
+    # Convert f3 to a LinearProblem since it's linear in u
+    # du = Au + b where A is the coefficient matrix and b is from parameters
+    A3 = [2.0 2.5 1.5; 4.0 -1.5 1.5; 1.0 -1.0 -1.0]
+    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]
     end
     explicitfun3(p, [sol1, sol2])
-    sol3 = solve(prob3, NewtonRaphson())
+    sol3 = solve(prob3)  # LinearProblem uses default linear solver
     manualscc = [sol1; sol2; sol3]
 
     sccprob = SciMLBase.SCCNonlinearProblem([prob1, prob2, prob3],
         SciMLBase.Void{Any}.([explicitfun1, explicitfun2, explicitfun3]))
-    scc_sol = solve(sccprob, NewtonRaphson())
+    
+    # Test with SCCAlg that handles both nonlinear and linear problems
+    using SCCNonlinearSolve
+    scc_alg = SCCNonlinearSolve.SCCAlg(nlalg = NewtonRaphson(), linalg = nothing)
+    scc_sol = solve(sccprob, scc_alg)
     @test sol ≈ manualscc ≈ scc_sol
 
     import NonlinearSolve # Required for Default
 
-    scc_sol = solve(sccprob)
-    @test sol ≈ manualscc ≈ scc_sol
+    # Test default interface
+    scc_sol_default = solve(sccprob)
+    @test sol ≈ manualscc ≈ scc_sol_default
 end