diff --git a/ext/SciMLOperatorsStaticArraysCoreExt.jl b/ext/SciMLOperatorsStaticArraysCoreExt.jl index a0c55186..4eea16fb 100644 --- a/ext/SciMLOperatorsStaticArraysCoreExt.jl +++ b/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 \ No newline at end of file +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 diff --git a/src/basic.jl b/src/basic.jl index ddf837e8..1257588a 100644 --- a/src/basic.jl +++ b/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,6 +384,7 @@ 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) @@ -378,6 +392,7 @@ for op in (:+, :-) @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 """ diff --git a/src/batch.jl b/src/batch.jl index 076e8c47..ea6e8d52 100644 --- a/src/batch.jl +++ b/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,) diff --git a/src/func.jl b/src/func.jl index fb3a8e4d..e1b2cadc 100644 --- a/src/func.jl +++ b/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) diff --git a/src/interface.jl b/src/interface.jl index 59156ded..fee504e8 100644 --- a/src/interface.jl +++ b/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 ### diff --git a/src/matrix.jl b/src/matrix.jl index b5369ebd..6fcc10c6 100644 --- a/src/matrix.jl +++ b/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, α, diff --git a/src/scalar.jl b/src/scalar.jl index ed185025..9b47fc15 100644 --- a/src/scalar.jl +++ b/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)) diff --git a/test/basic.jl b/test/basic.jl index 8c8c72bc..2f9de506 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -1,4 +1,4 @@ -using SciMLOperators, LinearAlgebra +using SciMLOperators, LinearAlgebra, SparseArrays using Random using SciMLOperators: IdentityOperator, @@ -138,6 +138,13 @@ end @test ldiv!(op, u) ≈ (α * D) \ v end +function apply_op!(H, du, u, p, t) + H(du, u, p, t) + return nothing +end + +test_apply_noalloc(H, du, u, p, t) = @test (@allocations apply_op!(H, du, u, p, t)) == 0 + @testset "AddedOperator" begin A = rand(N, N) |> MatrixOperator B = rand(N, N) |> MatrixOperator @@ -184,6 +191,45 @@ end for op in L.ops @test !isa(op, AddedOperator) end + + # Allocations Tests + + @allocations apply_op!(op, v, u, (), 1.0) # warmup + test_apply_noalloc(op, v, u, (), 1.0) + + ## 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 + + @allocations apply_op!(H_sparse, du, u, p, t) # warmup + @allocations apply_op!(H_dense, du, u, p, t) # warmup + test_apply_noalloc(H_sparse, du, u, p, t) + test_apply_noalloc(H_dense, du, u, p, t) + end end @testset "ComposedOperator" begin