Support SciMLOperators in LinearSolve

Project.toml

@@ -17,6 +17,7 @@ Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -43,7 +44,8 @@ KrylovKit = "0.5, 0.6"
 Preferences = "1"
 RecursiveFactorization = "0.2.8"
 Reexport = "1"
-SciMLBase = "1.68"
+SciMLBase = "1.81.1"
+SciMLOperators = "0.1.19"
 Setfield = "0.7, 0.8, 1"
 SnoopPrecompile = "1"
 Sparspak = "0.3.6"

lib/LinearSolveCUDA/src/LinearSolveCUDA.jl

@@ -1,6 +1,7 @@
 module LinearSolveCUDA
 
 using CUDA, LinearAlgebra, LinearSolve, SciMLBase
+using SciMLBase: AbstractSciMLOperator
 
 struct CudaOffloadFactorization <: LinearSolve.AbstractFactorization end
 
@@ -17,12 +18,13 @@ function SciMLBase.solve(cache::LinearSolve.LinearCache, alg::CudaOffloadFactori
 end
 
 function LinearSolve.do_factorization(alg::CudaOffloadFactorization, A, b, u)
-    A isa Union{AbstractMatrix, SciMLBase.AbstractDiffEqOperator} ||
+    A isa Union{AbstractMatrix, AbstractSciMLOperator} ||
         error("LU is not defined for $(typeof(A))")
 
-    if A isa SciMLBase.AbstractDiffEqOperator
+    if A isa Union{MatrixOperator, DiffEqArrayOperator}
         A = A.A
     end
+
     fact = qr(CUDA.CuArray(A))
     return fact
 end

src/LinearSolve.jl

@@ -6,11 +6,11 @@ end
 using ArrayInterfaceCore
 using RecursiveFactorization
 using Base: cache_dependencies, Bool
-import Base: eltype, adjoint, inv
 using LinearAlgebra
-using IterativeSolvers: Identity
 using SparseArrays
-using SciMLBase: AbstractDiffEqOperator, AbstractLinearAlgorithm
+using SciMLBase: AbstractLinearAlgorithm
+using SciMLOperators
+using SciMLOperators: AbstractSciMLOperator, IdentityOperator
 using Setfield
 using UnPack
 using SuiteSparse
@@ -41,6 +41,15 @@ needs_concrete_A(alg::AbstractFactorization) = true
 needs_concrete_A(alg::AbstractKrylovSubspaceMethod) = false
 needs_concrete_A(alg::AbstractSolveFunction) = false
 
+# Util
+
+_isidentity_struct(A) = false
+_isidentity_struct(λ::Number) = isone(λ)
+_isidentity_struct(A::UniformScaling) = isone(A.λ)
+_isidentity_struct(::IterativeSolvers.Identity) = true
+_isidentity_struct(::SciMLOperators.IdentityOperator) = true
+_isidentity_struct(::SciMLBase.DiffEqIdentity) = true
+
 # Code
 
 const INCLUDE_SPARSE = Preferences.@load_preference("include_sparse", Base.USE_GPL_LIBS)
@@ -97,7 +106,7 @@ export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization,
        UMFPACKFactorization, KLUFactorization, FastLUFactorization, FastQRFactorization,
        SparspakFactorization, DiagonalFactorization
 
-export LinearSolveFunction
+export LinearSolveFunction, DirectLdiv!
 
 export KrylovJL, KrylovJL_CG, KrylovJL_MINRES, KrylovJL_GMRES,
        KrylovJL_BICGSTAB, KrylovJL_LSMR, KrylovJL_CRAIGMR,

src/common.jl

@@ -1,11 +1,11 @@
-struct OperatorAssumptions{issquare} end
+struct OperatorAssumptions{issq} end
 function OperatorAssumptions(issquare = nothing)
-    issquare = something(_unwrap_val(issquare), Nothing)
-    OperatorAssumptions{issquare}()
+    issq = something(_unwrap_val(issquare), Nothing)
+    OperatorAssumptions{issq}()
 end
-issquare(::OperatorAssumptions{issq}) where {issq} = issq
+SciMLOperators.issquare(::OperatorAssumptions{issq}) where {issq} = issq
 
-struct LinearCache{TA, Tb, Tu, Tp, Talg, Tc, Tl, Tr, Ttol, issquare}
+struct LinearCache{TA, Tb, Tu, Tp, Talg, Tc, Tl, Tr, Ttol, issq}
     A::TA
     b::Tb
     u::Tu
@@ -19,7 +19,7 @@ struct LinearCache{TA, Tb, Tu, Tp, Talg, Tc, Tl, Tr, Ttol, issquare}
     reltol::Ttol
     maxiters::Int
     verbose::Bool
-    assumptions::OperatorAssumptions{issquare}
+    assumptions::OperatorAssumptions{issq}
 end
 
 """
@@ -92,9 +92,9 @@ function SciMLBase.init(prob::LinearProblem, alg::Union{SciMLLinearSolveAlgorith
                         reltol = default_tol(eltype(prob.A)),
                         maxiters::Int = length(prob.b),
                         verbose::Bool = false,
-                        Pl = Identity(),
-                        Pr = Identity(),
-                        assumptions = OperatorAssumptions(),
+                        Pl = IdentityOperator{size(prob.A, 1)}(),
+                        Pr = IdentityOperator{size(prob.A, 2)}(),
+                        assumptions = OperatorAssumptions(issquare(prob.A)),
                         kwargs...)
     @unpack A, b, u0, p = prob
 

src/default.jl

@@ -2,26 +2,30 @@
 # For SciML algorithms already using `defaultalg`, all assume square matrix.
 defaultalg(A, b) = defaultalg(A, b, OperatorAssumptions(Val(true)))
 
-function defaultalg(A::DiffEqArrayOperator, b, assumptions::OperatorAssumptions)
+function defaultalg(A::Union{DiffEqArrayOperator, MatrixOperator}, b,
+                    assumptions::OperatorAssumptions)
     defaultalg(A.A, b, assumptions)
 end
 
 # Ambiguity handling
-function defaultalg(A::DiffEqArrayOperator, b, assumptions::OperatorAssumptions{nothing})
+function defaultalg(A::Union{DiffEqArrayOperator, MatrixOperator}, b,
+                    assumptions::OperatorAssumptions{nothing})
     defaultalg(A.A, b, assumptions)
 end
 
-function defaultalg(A::DiffEqArrayOperator, b, assumptions::OperatorAssumptions{false})
+function defaultalg(A::Union{DiffEqArrayOperator, MatrixOperator}, b,
+                    assumptions::OperatorAssumptions{false})
     defaultalg(A.A, b, assumptions)
 end
 
-function defaultalg(A::DiffEqArrayOperator, b, assumptions::OperatorAssumptions{true})
+function defaultalg(A::Union{DiffEqArrayOperator, MatrixOperator}, b,
+                    assumptions::OperatorAssumptions{true})
     defaultalg(A.A, b, assumptions)
 end
 
 function defaultalg(A, b, ::OperatorAssumptions{Nothing})
-    issquare = size(A, 1) == size(A, 2)
-    defaultalg(A, b, OperatorAssumptions(Val(issquare)))
+    issq = issquare(A)
+    defaultalg(A, b, OperatorAssumptions(Val(issq)))
 end
 
 function defaultalg(A::Tridiagonal, b, ::OperatorAssumptions{true})
@@ -33,10 +37,13 @@ end
 function defaultalg(A::SymTridiagonal, b, ::OperatorAssumptions{true})
     GenericFactorization(; fact_alg = ldlt!)
 end
-function defaultalg(A::Diagonal, b, ::OperatorAssumptions{true})
-    DiagonalFactorization()
+function defaultalg(A::Bidiagonal, b, ::OperatorAssumptions{true})
+    DirectLdiv!()
+end
+function defaultalg(A::Factorization, b, ::OperatorAssumptions{true})
+    DirectLdiv!()
 end
-function defaultalg(A::Diagonal, b, ::OperatorAssumptions{false})
+function defaultalg(A::Diagonal, b, ::OperatorAssumptions{true})
     DiagonalFactorization()
 end
 function defaultalg(A::Diagonal, b, ::OperatorAssumptions{Nothing})
@@ -75,18 +82,26 @@ function defaultalg(A, b::GPUArraysCore.AbstractGPUArray, ::OperatorAssumptions{
     end
 end
 
-function defaultalg(A::SciMLBase.AbstractDiffEqOperator, b,
-                    assumptions::OperatorAssumptions)
+function defaultalg(A::SciMLBase.AbstractSciMLOperator, b,
+                    assumptions::OperatorAssumptions{true})
+    if has_ldiv!(A)
+        return DirectLdiv!()
+    end
+
     KrylovJL_GMRES()
 end
 
 # Ambiguity handling
-function defaultalg(A::SciMLBase.AbstractDiffEqOperator, b,
+function defaultalg(A::SciMLBase.AbstractSciMLOperator, b,
                     assumptions::OperatorAssumptions{Nothing})
+    if has_ldiv!(A)
+        return DirectLdiv!()
+    end
+
     KrylovJL_GMRES()
 end
 
-function defaultalg(A::SciMLBase.AbstractDiffEqOperator, b,
+function defaultalg(A::SciMLBase.AbstractSciMLOperator, b,
                     assumptions::OperatorAssumptions{false})
     m, n = size(A)
     if m < n

src/iterative_wrappers.jl

@@ -150,8 +150,9 @@ function SciMLBase.solve(cache::LinearCache, alg::KrylovJL; kwargs...)
     M = cache.Pl
     N = cache.Pr
 
-    M = (M === Identity()) ? I : InvPreconditioner(M)
-    N = (N === Identity()) ? I : InvPreconditioner(N)
+    # use no-op preconditioner for Krylov.jl (LinearAlgebra.I) when M/N is identity
+    M = _isidentity_struct(M) ? I : M
+    N = _isidentity_struct(M) ? I : N
 
     atol = float(cache.abstol)
     rtol = float(cache.reltol)
@@ -160,7 +161,7 @@ function SciMLBase.solve(cache::LinearCache, alg::KrylovJL; kwargs...)
 
     args = (cache.cacheval, cache.A, cache.b)
     kwargs = (atol = atol, rtol = rtol, itmax = itmax, verbose = verbose,
-              history = true, alg.kwargs...)
+              ldiv = true, history = true, alg.kwargs...)
 
     if cache.cacheval isa Krylov.CgSolver
         N !== I &&
@@ -234,15 +235,15 @@ function init_cacheval(alg::IterativeSolversJL, A, b, u, Pl, Pr, maxiters::Int,
               alg.kwargs...)
 
     iterable = if alg.generate_iterator === IterativeSolvers.cg_iterator!
-        Pr !== Identity() &&
+        !_isidentity_struct(Pr) &&
             @warn "$(alg.generate_iterator) doesn't support right preconditioning"
         alg.generate_iterator(u, A, b, Pl;
                               kwargs...)
     elseif alg.generate_iterator === IterativeSolvers.gmres_iterable!
         alg.generate_iterator(u, A, b; Pl = Pl, Pr = Pr, restart = restart,
                               kwargs...)
     elseif alg.generate_iterator === IterativeSolvers.bicgstabl_iterator!
-        Pr !== Identity() &&
+        !_isidentity_struct(Pr) &&
             @warn "$(alg.generate_iterator) doesn't support right preconditioning"
         alg.generate_iterator(u, A, b, alg.args...; Pl = Pl,
                               abstol = abstol, reltol = reltol,

src/solve_function.jl

@@ -13,3 +13,12 @@ function SciMLBase.solve(cache::LinearCache, alg::LinearSolveFunction,
 
     return SciMLBase.build_linear_solution(alg, cache.u, nothing, cache)
 end
+
+struct DirectLdiv! <: AbstractSolveFunction end
+
+function SciMLBase.solve(cache::LinearCache, alg::DirectLdiv!, args...; kwargs...)
+    @unpack A, b, u = cache
+    ldiv!(u, A, b)
+
+    return SciMLBase.build_linear_solution(alg, cache.u, nothing, cache)
+end

test/basictests.jl

@@ -1,4 +1,5 @@
 using LinearSolve, LinearAlgebra, SparseArrays, MultiFloats, ForwardDiff
+using SciMLOperators
 using Test
 import Random
 
@@ -27,20 +28,20 @@ function test_interface(alg, prob1, prob2)
     b2 = prob2.b
     x2 = prob2.u0
 
-    y = solve(prob1, alg; cache_kwargs...)
-    @test A1 * y ≈ b1
+    sol = solve(prob1, alg; cache_kwargs...)
+    @test A1 * sol.u ≈ b1
 
     cache = SciMLBase.init(prob1, alg; cache_kwargs...) # initialize cache
-    y = solve(cache)
-    @test A1 * y ≈ b1
+    sol = solve(cache)
+    @test A1 * sol.u ≈ b1
 
-    cache = LinearSolve.set_A(cache, copy(A2))
-    y = solve(cache; cache_kwargs...)
-    @test A2 * y ≈ b1
+    cache = LinearSolve.set_A(cache, deepcopy(A2))
+    sol = solve(cache; cache_kwargs...)
+    @test A2 * sol.u ≈ b1
 
     cache = LinearSolve.set_b(cache, b2)
-    y = solve(cache; cache_kwargs...)
-    @test A2 * y ≈ b2
+    sol = solve(cache; cache_kwargs...)
+    @test A2 * sol.u ≈ b2
 
     return
 end
@@ -271,12 +272,14 @@ end
 
     @testset "Preconditioners" begin
         @testset "Vector Diagonal Preconditioner" begin
-            s = rand(n)
-            Pl, Pr = Diagonal(s), LinearSolve.InvPreconditioner(Diagonal(s))
-
             x = rand(n, n)
             y = rand(n, n)
 
+            s = rand(n)
+            Pl = Diagonal(s) |> MatrixOperator
+            Pr = Diagonal(s) |> MatrixOperator |> inv
+            Pr = cache_operator(Pr, x)
+
             mul!(y, Pl, x)
             @test y ≈ s .* x
             mul!(y, Pr, x)
@@ -346,33 +349,67 @@ end
     end
 
     @testset "Solve Function" begin
-        A1 = rand(n) |> Diagonal
+        A1 = rand(n) |> Diagonal |> Array
         b1 = rand(n)
         x1 = zero(b1)
-        A2 = rand(n) |> Diagonal
+        A2 = rand(n) |> Diagonal |> Array
         b2 = rand(n)
         x2 = zero(b1)
 
-        function sol_func(A, b, u, p, newA, Pl, Pr, solverdata; verbose = true, kwargs...)
-            if verbose == true
-                println("out-of-place solve")
+        @testset "LinearSolveFunction" begin
+            function sol_func(A, b, u, p, newA, Pl, Pr, solverdata; verbose = true,
+                              kwargs...)
+                if verbose == true
+                    println("out-of-place solve")
+                end
+                u = A \ b
             end
-            u = A \ b
-        end
 
-        function sol_func!(A, b, u, p, newA, Pl, Pr, solverdata; verbose = true, kwargs...)
-            if verbose == true
-                println("in-place solve")
+            function sol_func!(A, b, u, p, newA, Pl, Pr, solverdata; verbose = true,
+                               kwargs...)
+                if verbose == true
+                    println("in-place solve")
+                end
+                ldiv!(u, A, b)
+            end
+
+            prob1 = LinearProblem(A1, b1; u0 = x1)
+            prob2 = LinearProblem(A1, b1; u0 = x1)
+
+            for alg in (LinearSolveFunction(sol_func),
+                        LinearSolveFunction(sol_func!))
+                test_interface(alg, prob1, prob2)
             end
-            ldiv!(u, A, b)
         end
 
-        prob1 = LinearProblem(A1, b1; u0 = x1)
-        prob2 = LinearProblem(A1, b1; u0 = x1)
+        @testset "DirectLdiv!" begin
+            function get_operator(A, u)
+                F = lu(A)
 
-        for alg in (LinearSolveFunction(sol_func),
-                    LinearSolveFunction(sol_func!))
-            test_interface(alg, prob1, prob2)
+                function f(du, u, p, t)
+                    println("using FunctionOperator mul!")
+                    mul!(du, A, u)
+                end
+
+                function fi(du, u, p, t)
+                    println("using FunctionOperator ldiv!")
+                    ldiv!(du, F, u)
+                end
+
+                FunctionOperator(f, u, u; isinplace = true, op_inverse = fi)
+            end
+
+            op1 = get_operator(A1, x1 * 0)
+            op2 = get_operator(A2, x2 * 0)
+
+            prob1 = LinearProblem(op1, b1; u0 = x1)
+            prob2 = LinearProblem(op2, b2; u0 = x2)
+
+            @test LinearSolve.defaultalg(op1, x1) isa DirectLdiv!
+            @test LinearSolve.defaultalg(op2, x2) isa DirectLdiv!
+
+            test_interface(DirectLdiv!(), prob1, prob2)
+            test_interface(nothing, prob1, prob2)
         end
     end
 end # testset