Improve mul!, AddedOperator, and update_coefficients! to remove memory allocations

ext/SciMLOperatorsStaticArraysCoreExt.jl

@@ -1,11 +1,11 @@
-module SciMLOperatorsStaticArraysCoreExt
-
-import SciMLOperators
-import StaticArraysCore
-
-function Base.copyto!(L::SciMLOperators.MatrixOperator,
-    rhs::Base.Broadcast.Broadcasted{<:StaticArraysCore.StaticArrayStyle})
-(copyto!(L.A, rhs); L)
-end
-
-end #module
+module SciMLOperatorsStaticArraysCoreExt
+
+import SciMLOperators
+import StaticArraysCore
+
+function Base.copyto!(L::SciMLOperators.MatrixOperator,
+        rhs::Base.Broadcast.Broadcasted{<:StaticArraysCore.StaticArrayStyle})
+    (copyto!(L.A, rhs); L)
+end
+
+end #module

src/basic.jl

@@ -201,7 +201,7 @@ end
 
 for T in SCALINGNUMBERTYPES
     @eval function ScaledOperator(λ::$T, L::ScaledOperator)
-        λ = ScalarOperator(λ) * L.λ
+        λ = λ * L.λ
         ScaledOperator(λ, L.L)
     end
 
@@ -250,7 +250,7 @@ function update_coefficients!(L::ScaledOperator, u, p, t)
     update_coefficients!(L.L, u, p, t)
     update_coefficients!(L.λ, u, p, t)
 
-    L
+    nothing
 end
 
 getops(L::ScaledOperator) = (L.λ, L.L)
@@ -288,13 +288,14 @@ end
 Base.:*(L::ScaledOperator, u::AbstractVecOrMat) = L.λ * (L.L * u)
 Base.:\(L::ScaledOperator, u::AbstractVecOrMat) = L.λ \ (L.L \ u)
 
-function LinearAlgebra.mul!(v::AbstractVecOrMat, L::ScaledOperator, u::AbstractVecOrMat)
+@inline function LinearAlgebra.mul!(
+        v::AbstractVecOrMat, L::ScaledOperator, u::AbstractVecOrMat)
     iszero(L.λ) && return lmul!(false, v)
     a = convert(Number, L.λ)
     mul!(v, L.L, u, a, false)
 end
 
-function LinearAlgebra.mul!(v::AbstractVecOrMat,
+@inline function LinearAlgebra.mul!(v::AbstractVecOrMat,
         L::ScaledOperator,
         u::AbstractVecOrMat,
         α,
@@ -326,22 +327,34 @@ struct AddedOperator{T,
 
     function AddedOperator(ops)
         @assert !isempty(ops)
+        _check_AddedOperator_sizes(ops)
         T = promote_type(eltype.(ops)...)
         new{T, typeof(ops)}(ops)
     end
 end
 
 function AddedOperator(ops::AbstractSciMLOperator...)
-    sz = size(first(ops))
-    for op in ops[2:end]
-        @assert size(op)==sz "Dimension mismatch: cannot add operators of
-      sizes $(sz), and $(size(op))."
-    end
     AddedOperator(ops)
 end
 
 AddedOperator(L::AbstractSciMLOperator) = L
 
+@generated function _check_AddedOperator_sizes(ops::Tuple)
+    ops_types = ops.parameters
+    N = length(ops_types)
+    sz_expr_list = ()
+    sz_expr = :(sz = size(first(ops)))
+    for i in 2:N
+        sz_expr_list = (sz_expr_list..., :(size(ops[$i]) == sz))
+    end
+
+    quote
+        $sz_expr
+        @assert all(tuple($(sz_expr_list...))) "Dimension mismatch: cannot add operators of different sizes."
+        nothing
+    end
+end
+
 # constructors
 Base.:+(A::AbstractSciMLOperator, B::AbstractMatrix) = A + MatrixOperator(B)
 Base.:+(A::AbstractMatrix, B::AbstractSciMLOperator) = MatrixOperator(A) + B
@@ -371,13 +384,15 @@ for op in (:+, :-)
         for LT in SCALINGCOMBINETYPES
             @eval function Base.$op(L::$LT, λ::$T)
                 @assert issquare(L)
+                iszero(λ) && return L
                 N = size(L, 1)
                 Id = IdentityOperator(N)
                 AddedOperator(L, $op(λ) * Id)
             end
 
             @eval function Base.$op(λ::$T, L::$LT)
                 @assert issquare(L)
+                iszero(λ) && return $op(L)
                 N = size(L, 1)
                 Id = IdentityOperator(N)
                 AddedOperator(λ * Id, $op(L))
@@ -386,6 +401,23 @@ for op in (:+, :-)
     end
 end
 
+for T in SCALINGNUMBERTYPES[2:end]
+    @eval function Base.:*(λ::$T, L::AddedOperator)
+        ops = map(op -> λ * op, L.ops)
+        AddedOperator(ops)
+    end
+
+    @eval function Base.:*(L::AddedOperator, λ::$T)
+        ops = map(op -> λ * op, L.ops)
+        AddedOperator(ops)
+    end
+
+    @eval function Base.:/(L::AddedOperator, λ::$T)
+        ops = map(op -> op / λ, L.ops)
+        AddedOperator(ops)
+    end
+end
+
 function Base.convert(::Type{AbstractMatrix}, L::AddedOperator)
     sum(op -> convert(AbstractMatrix, op), L.ops)
 end
@@ -422,16 +454,32 @@ function update_coefficients(L::AddedOperator, u, p, t)
     @reset L.ops = ops
 end
 
+@generated function update_coefficients!(L::AddedOperator, u, p, t)
+    ops_types = L.parameters[2].parameters
+    N = length(ops_types)
+    quote
+        Base.@nexprs $N i->begin
+            update_coefficients!(L.ops[i], u, p, t)
+        end
+
+        nothing
+    end
+end
+
 getops(L::AddedOperator) = L.ops
 islinear(L::AddedOperator) = all(islinear, getops(L))
 Base.iszero(L::AddedOperator) = all(iszero, getops(L))
 has_adjoint(L::AddedOperator) = all(has_adjoint, L.ops)
 
-function cache_internals(L::AddedOperator, u::AbstractVecOrMat)
-    for i in 1:length(L.ops)
-        @reset L.ops[i] = cache_operator(L.ops[i], u)
+@generated function cache_internals(L::AddedOperator, u::AbstractVecOrMat)
+    ops_types = L.parameters[2].parameters
+    N = length(ops_types)
+    quote
+        Base.@nexprs $N i->begin
+            @reset L.ops[i] = cache_operator(L.ops[i], u)
+        end
+        L
     end
-    L
 end
 
 getindex(L::AddedOperator, i::Int) = sum(op -> op[i], L.ops)
@@ -441,26 +489,33 @@ function Base.:*(L::AddedOperator, u::AbstractVecOrMat)
     sum(op -> iszero(op) ? zero(u) : op * u, L.ops)
 end
 
-function LinearAlgebra.mul!(v::AbstractVecOrMat, L::AddedOperator, u::AbstractVecOrMat)
-    mul!(v, first(L.ops), u)
-    for op in L.ops[2:end]
-        iszero(op) && continue
-        mul!(v, op, u, true, true)
+@generated function LinearAlgebra.mul!(
+        v::AbstractVecOrMat, L::AddedOperator, u::AbstractVecOrMat)
+    ops_types = L.parameters[2].parameters
+    N = length(ops_types)
+    quote
+        mul!(v, L.ops[1], u)
+        Base.@nexprs $(N - 1) i->begin
+            mul!(v, L.ops[i + 1], u, true, true)
+        end
+        v
     end
-    v
 end
 
-function LinearAlgebra.mul!(v::AbstractVecOrMat,
+@generated function LinearAlgebra.mul!(v::AbstractVecOrMat,
         L::AddedOperator,
         u::AbstractVecOrMat,
         α,
         β)
-    lmul!(β, v)
-    for op in L.ops
-        iszero(op) && continue
-        mul!(v, op, u, α, true)
+    ops_types = L.parameters[2].parameters
+    N = length(ops_types)
+    quote
+        lmul!(β, v)
+        Base.@nexprs $(N) i->begin
+            mul!(v, L.ops[i], u, α, true)
+        end
+        v
     end
-    v
 end
 
 """

src/batch.jl

@@ -88,6 +88,8 @@ end
 
 function update_coefficients!(L::BatchedDiagonalOperator, u, p, t; kwargs...)
     L.update_func!(L.diag, u, p, t; kwargs...)
+
+    nothing
 end
 
 getops(L::BatchedDiagonalOperator) = (L.diag,)

src/func.jl

@@ -382,7 +382,7 @@ function update_coefficients!(L::FunctionOperator, u, p, t; kwargs...)
         update_coefficients!(op, u, p, t; filtered_kwargs...)
     end
 
-    L
+    nothing
 end
 
 function iscached(L::FunctionOperator)

src/interface.jl

@@ -98,11 +98,11 @@ L * u
 """
 update_coefficients!(L, u, p, t; kwargs...) = nothing
 
+# We cannot use @generate because we don't know the type structure of L in advance
 function update_coefficients!(L::AbstractSciMLOperator, u, p, t; kwargs...)
-    for op in getops(L)
-        update_coefficients!(op, u, p, t; kwargs...)
-    end
-    L
+    foreach(op -> update_coefficients!(op, u, p, t; kwargs...), getops(L))
+
+    nothing
 end
 
 ###

src/matrix.jl

@@ -161,6 +161,8 @@ end
 
 function update_coefficients!(L::MatrixOperator, u, p, t; kwargs...)
     L.update_func!(L.A, u, p, t; kwargs...)
+
+    nothing
 end
 
 # TODO - add tests for MatrixOperator indexing
@@ -194,10 +196,11 @@ end
 # operator application
 Base.:*(L::MatrixOperator, u::AbstractVecOrMat) = L.A * u
 Base.:\(L::MatrixOperator, u::AbstractVecOrMat) = L.A \ u
-function LinearAlgebra.mul!(v::AbstractVecOrMat, L::MatrixOperator, u::AbstractVecOrMat)
+@inline function LinearAlgebra.mul!(
+        v::AbstractVecOrMat, L::MatrixOperator, u::AbstractVecOrMat)
     mul!(v, L.A, u)
 end
-function LinearAlgebra.mul!(v::AbstractVecOrMat,
+@inline function LinearAlgebra.mul!(v::AbstractVecOrMat,
         L::MatrixOperator,
         u::AbstractVecOrMat,
         α,

src/scalar.jl

@@ -191,10 +191,12 @@ has_ldiv!(α::ScalarOperator) = has_ldiv(α)
 
 function update_coefficients!(L::ScalarOperator, u, p, t; kwargs...)
     L.val = L.update_func(L.val, u, p, t; kwargs...)
+    nothing
 end
 
 function update_coefficients(L::ScalarOperator, u, p, t; kwargs...)
-    @reset L.val = L.update_func(L.val, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    L
 end
 
 """
@@ -313,6 +315,26 @@ for op in (:*, :∘)
     end
 end
 
+# Different methods for constant ScalarOperators
+for T in SCALINGNUMBERTYPES[2:end]
+    @eval function Base.:*(α::ScalarOperator, x::$T)
+        if isconstant(α)
+            T2 = promote_type($T, eltype(α))
+            return ScalarOperator(convert(T2, α) * x, α.update_func)
+        else
+            return ComposedScalarOperator(α, ScalarOperator(x))
+        end
+    end
+    @eval function Base.:*(x::$T, α::ScalarOperator)
+        if isconstant(α)
+            T2 = promote_type($T, eltype(α))
+            return ScalarOperator(convert(T2, α) * x, α.update_func)
+        else
+            return ComposedScalarOperator(ScalarOperator(x), α)
+        end
+    end
+end
+
 function Base.convert(T::Type{<:Number}, α::ComposedScalarOperator)
     iszero(α) && return zero(T)
     prod(convert.(T, α.ops))

test/basic.jl

@@ -1,4 +1,4 @@
-using SciMLOperators, LinearAlgebra
+using SciMLOperators, LinearAlgebra, SparseArrays
 using Random
 
 using SciMLOperators: IdentityOperator,
@@ -138,6 +138,11 @@ end
     @test ldiv!(op, u) ≈ (α * D) \ v
 end
 
+function apply_op!(H, du, u, p, t)
+    H(du, u, p, t)
+    return nothing
+end
+
 @testset "AddedOperator" begin
     A = rand(N, N) |> MatrixOperator
     B = rand(N, N) |> MatrixOperator
@@ -184,6 +189,48 @@ end
     for op in L.ops
         @test !isa(op, AddedOperator)
     end
+
+    # Allocations Tests
+
+    allocs_tot = @allocations apply_op!(op, v, u, (), 1.0) # warmup
+    allocs_tot = @allocations apply_op!(op, v, u, (), 1.0)
+    @test allocs_tot <= 1 # BenchmarkTools.jl returns 0 instead of 1
+
+    ## Time-Dependent Coefficients
+
+    for T in (Float32, Float64, ComplexF32, ComplexF64)
+        N = 100
+        A1_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
+        A2_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
+        A3_sparse = MatrixOperator(sprand(T, N, N, 5 / N))
+
+        A1_dense = MatrixOperator(rand(T, N, N))
+        A2_dense = MatrixOperator(rand(T, N, N))
+        A3_dense = MatrixOperator(rand(T, N, N))
+
+        coeff1(a, u, p, t) = sin(p.ω * t)
+        coeff2(a, u, p, t) = cos(p.ω * t)
+        coeff3(a, u, p, t) = sin(p.ω * t) * cos(p.ω * t)
+
+        c1 = ScalarOperator(rand(T), coeff1)
+        c2 = ScalarOperator(rand(T), coeff2)
+        c3 = ScalarOperator(rand(T), coeff3)
+
+        H_sparse = c1 * A1_sparse + c2 * A2_sparse + c3 * A3_sparse
+        H_dense = c1 * A1_dense + c2 * A2_dense + c3 * A3_dense
+
+        u = rand(T, N)
+        du = similar(u)
+        p = (ω = 0.1,)
+        t = 0.1
+
+        allocs_sparse = @allocations apply_op!(H_sparse, du, u, p, t) # warmup
+        allocs_sparse = @allocations apply_op!(H_sparse, du, u, p, t)
+        allocs_dense = @allocations apply_op!(H_dense, du, u, p, t) # warmup
+        allocs_dense = @allocations apply_op!(H_dense, du, u, p, t)
+        @test allocs_sparse <= 1 # BenchmarkTools.jl returns 0 instead of 1
+        @test allocs_dense <= 1 # BenchmarkTools.jl returns 0 instead of 1
+    end
 end
 
 @testset "ComposedOperator" begin