Implement generated functions

Author: albertomercurio

src/basic.jl

@@ -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)
@@ -327,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
@@ -372,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))
@@ -440,24 +454,32 @@ function update_coefficients(L::AddedOperator, u, p, t)
     @reset L.ops = ops
 end
 
-function update_coefficients!(L::AddedOperator, u, p, t)
-    for op in L.ops
-        update_coefficients!(op, u, p, t)
-    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
 
-    L
+        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)
@@ -467,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

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
 
 """