SciMLOperators operator evaluation interface allowing different defining and action vectors

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLOperators"
 uuid = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 authors = ["Vedant Puri <[email protected]>"]
-version = "0.3.13"
+version = "0.4.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/SciMLOperators.jl

@@ -26,21 +26,29 @@ Subtypes of `AbstractSciMLOperator` represent linear, nonlinear,
 time-dependent operators acting on vectors, or matrix column-vectors.
 A lazy operator algebra is also defined for `AbstractSciMLOperator`s.
 
+# Mathematical Notation
+
+An `AbstractSciMLOperator` ``L`` is an operator which is used to represent
+the following type of equation:
+
+```math
+w = L(u,p,t)[v]
+```
+
+where `L[v]` is the operator application of ``L`` on the vector ``v``. 
+
 # Interface
 
-An `AbstractSciMLOperator` can be called like a function. This behaves
-like multiplication by the linear operator represented by the
-`AbstractSciMLOperator`. Possible signatures are
+An `AbstractSciMLOperator` can be called  like a function in the following ways:
 
-- `L(v, u, p, t)` for in-place operator evaluation
-- `v = L(u, p, t)` for out-of-place operator evaluation
+- `L(v, u, p, t)` - Out-of-place application where `v` is the action vector and `u` is the update vector
+- `L(w, v, u, p, t)` - In-place application where `w` is the destination, `v` is the action vector, and `u` is the update vector
+- `L(w, v, u, p, t, α, β)` - In-place application with scaling: `w = α*(L*v) + β*w`
 
-Operator evaluation methods update its coefficients with `(u, p, t)`
-information using the `update_coefficients(!)` method. The methods
-are exported and can be called as follows:
+Operator state can be updated separately from application:
 
-- `update_coefficients!(L, u, p, t)` for out-of-place operator update
-- `L = update_coefficients(L, u, p, t)` for in-place operator update
+- `update_coefficients!(L, u, p, t)` for in-place operator update
+- `L = update_coefficients(L, u, p, t)` for out-of-place operator update
 
 SciMLOperators also overloads `Base.*`, `LinearAlgebra.mul!`,
 `LinearAlgebra.ldiv!` for operator evaluation without updating operator state.

src/basic.jl

@@ -69,6 +69,27 @@ function LinearAlgebra.ldiv!(ii::IdentityOperator, u::AbstractVecOrMat)
     u
 end
 
+# Out-of-place: v is action vector, u is update vector
+function (ii::IdentityOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    @assert size(v, 1) == ii.len
+    update_coefficients(ii, u, p, t; kwargs...)
+    copy(v)
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (ii::IdentityOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    @assert size(v, 1) == ii.len
+    update_coefficients!(ii, u, p, t; kwargs...)
+    copy!(w, v)
+end
+
+# In-place with scaling: w = α*(ii*v) + β*w
+function (ii::IdentityOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    @assert size(v, 1) == ii.len
+    update_coefficients!(ii, u, p, t; kwargs...)
+    mul!(w, I, v, α, β)
+end
+
 # operator fusion with identity returns operator itself
 for op in (:*, :∘)
     @eval function Base.$op(ii::IdentityOperator, A::AbstractSciMLOperator)
@@ -146,6 +167,26 @@ function LinearAlgebra.mul!(v::AbstractVecOrMat,
     lmul!(β, v)
 end
 
+# Out-of-place: v is action vector, u is update vector
+function (nn::NullOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    @assert size(v, 1) == nn.len
+    zero(v)
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (nn::NullOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    @assert size(v, 1) == nn.len
+    lmul!(false, w)
+    w
+end
+
+# In-place with scaling: w = α*(nn*v) + β*w
+function (nn::NullOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    @assert size(v, 1) == nn.len
+    lmul!(β, w)
+    w
+end
+
 # operator fusion, composition
 for op in (:*, :∘)
     @eval function Base.$op(nn::NullOperator, A::AbstractSciMLOperator)
@@ -218,6 +259,26 @@ for T in SCALINGNUMBERTYPES
     end
 end
 
+# Special cases for constant scalars. These simplify the structure when applicable
+for T in SCALINGNUMBERTYPES[2:end]
+    @eval function Base.:*(α::$T, L::ScaledOperator)
+        isconstant(L.λ) && return ScaledOperator(α * L.λ, L.L)
+        return ScaledOperator(L.λ, α * L.L) # Try to propagate the rule
+    end
+    @eval function Base.:*(L::ScaledOperator, α::$T)
+        isconstant(L.λ) && return ScaledOperator(α * L.λ, L.L)
+        return ScaledOperator(L.λ, α * L.L) # Try to propagate the rule
+    end
+    @eval function Base.:*(α::$T, L::MatrixOperator)
+        isconstant(L) && return MatrixOperator(α * L.A)
+        return ScaledOperator(α, L) # Going back to the generic case
+    end
+    @eval function Base.:*(L::MatrixOperator, α::$T)
+        isconstant(L) && return MatrixOperator(α * L.A)
+        return ScaledOperator(α, L) # Going back to the generic case
+    end
+end
+
 Base.:-(L::AbstractSciMLOperator) = ScaledOperator(-true, L)
 Base.:+(L::AbstractSciMLOperator) = L
 
@@ -316,6 +377,43 @@ function LinearAlgebra.ldiv!(L::ScaledOperator, u::AbstractVecOrMat)
     ldiv!(L.L, u)
 end
 
+# Out-of-place: v is action vector, u is update vector
+function (L::ScaledOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    L = update_coefficients(L, u, p, t; kwargs...)
+    if iszero(L.λ)
+        return zero(v)
+    else
+        return L.λ * (L.L * v)
+    end
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (L::ScaledOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    if iszero(L.λ)
+        lmul!(false, w)
+        return w
+    else
+        a = convert(Number, L.λ)
+        mul!(w, L.L, v, a, false)
+        return w
+    end
+end
+
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::ScaledOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    if iszero(L.λ)
+        lmul!(β, w)
+        return w
+    else
+        a = convert(Number, L.λ * α)
+        mul!(w, L.L, v, a, β)
+        return w
+    end
+end
+
+
 """
 Lazy operator addition
 
@@ -518,6 +616,35 @@ end
         v
     end
 end
+# Out-of-place: v is action vector, u is update vector
+function (L::AddedOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    L = update_coefficients(L, u, p, t; kwargs...)
+    sum(op -> iszero(op) ? zero(v) : op(v, u, p, t; kwargs...), L.ops)
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (L::AddedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    L.ops[1](w, v, u, p, t; kwargs...)
+    for i in 2:length(L.ops)
+        if !iszero(L.ops[i])
+            L.ops[i](w, v, u, p, t, 1.0, 1.0; kwargs...)
+        end
+    end
+    w
+end
+
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::AddedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    lmul!(β, w)
+    for op in L.ops
+        if !iszero(op)
+            op(w, v, u, p, t, α, 1.0; kwargs...)
+        end
+    end
+    w
+end
 
 """
     Lazy operator composition
@@ -772,6 +899,41 @@ function LinearAlgebra.ldiv!(L::ComposedOperator, u::AbstractVecOrMat)
     u
 end
 
+# Out-of-place: v is action vector, u is update vector
+function (L::ComposedOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    L = update_coefficients(L, u, p, t; kwargs...)
+    result = v
+    for op in reverse(L.ops)
+        result = op(result, u, p, t; kwargs...)
+    end
+    result
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (L::ComposedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    @assert iscached(L) "Cache needs to be set up for ComposedOperator. Call cache_operator(L, u) first."
+
+    vecs = (w, L.cache[1:(end-1)]..., v)
+    for i in reverse(1:length(L.ops))
+        L.ops[i](vecs[i], vecs[i+1], u, p, t; kwargs...)
+    end
+    w
+end
+
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::ComposedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    @assert iscached(L) "Cache needs to be set up for ComposedOperator. Call cache_operator(L, u) first."
+
+    cache = L.cache[end]
+    copy!(cache, w)
+
+    L(w, v, u, p, t; kwargs...)
+    lmul!(α, w)
+    axpy!(β, cache, w)
+end
+
 """
     Lazy Operator Inverse
 """
@@ -889,4 +1051,29 @@ function LinearAlgebra.ldiv!(L::InvertedOperator, u::AbstractVecOrMat)
     copy!(L.cache, u)
     mul!(u, L.L, L.cache)
 end
+
+# Out-of-place: v is action vector, u is update vector
+function (L::InvertedOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    L = update_coefficients(L, u, p, t; kwargs...)
+    L.L \ v
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (L::InvertedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    ldiv!(w, L.L, v)
+    w
+end
+
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::InvertedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    @assert iscached(L) "Cache needs to be set up for InvertedOperator. Call cache_operator(L, u) first."
+
+    copy!(L.cache, w)
+    ldiv!(w, L.L, v)
+    lmul!(α, w)
+    axpy!(β, L.cache, w)
+    w
+end
 #

src/batch.jl

@@ -153,4 +153,25 @@ function LinearAlgebra.ldiv!(L::BatchedDiagonalOperator, u::AbstractVecOrMat)
 
     u
 end
+
+function (L::BatchedDiagonalOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    L = update_coefficients(L, u, p, t; kwargs...)
+    L.diag .* v
+end
+
+function (L::BatchedDiagonalOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    w .= L.diag .* v
+    return w
+end
+
+function (L::BatchedDiagonalOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+    if β == 0
+        w .= α .* (L.diag .* v)
+    else
+        w .= α .* (L.diag .* v) .+ β .* w
+    end    
+    return w
+end
 #

src/func.jl

@@ -782,4 +782,68 @@ end
 function LinearAlgebra.ldiv!(L::FunctionOperator{false}, u::AbstractArray)
     @error "LinearAlgebra.ldiv! not defined for out-of-place $L"
 end
+
+# Out-of-place: v is action vector, u is update vector
+function (L::FunctionOperator)(v::AbstractArray, u, p, t; kwargs...)
+    L = update_coefficients(L, u, p, t; kwargs...)
+    _sizecheck(L, v, nothing)
+    V, _, vec_output = _unvec(L, v, nothing)
+
+    # Apply the operator to action vector v after updating with u
+    if L.traits.outofplace
+        result = L.op(V, L.p, L.t; L.traits.kwargs...)
+        return vec_output ? vec(result) : result
+    else
+        # For operators without out-of-place methods, use their in-place methods with a temporary
+        Co = similar(V)
+        L.op(Co, V, L.p, L.t; L.traits.kwargs...)
+        return vec_output ? vec(Co) : Co
+    end
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (L::FunctionOperator)(w::AbstractArray, v::AbstractArray, u, p, t; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+
+    # Check dimensions
+    _sizecheck(L, v, w)
+    V, W, _ = _unvec(L, v, w)
+
+    # Apply the operator in-place to action vector v after updating with u
+    if L.traits.isinplace
+        L.op(W, V, L.p, L.t; L.traits.kwargs...)
+    else
+        # For operators without in-place methods, use their out-of-place methods
+        result = L.op(V, L.p, L.t; L.traits.kwargs...)
+        copyto!(W, result)
+    end
+
+    return w
+end
+
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::FunctionOperator)(w::AbstractArray, v::AbstractArray, u, p, t, α, β; kwargs...)
+    update_coefficients!(L, u, p, t; kwargs...)
+
+    # Check dimensions
+    _sizecheck(L, v, w)
+    V, W, _ = _unvec(L, v, w)
+
+    # Apply the operator in-place to action vector v with scaling
+    if L.traits.isinplace && L.traits.has_mul5
+        # Direct 5-arg mul! if supported
+        L.op(W, V, L.p, L.t, α, β; L.traits.kwargs...)
+    elseif L.traits.isinplace
+        # Use temporary for regular in-place
+        temp = copy(W)
+        L.op(W, V, L.p, L.t; L.traits.kwargs...)
+        axpby!(β, temp, α, W)
+    else
+        # Out-of-place with scaling
+        result = L.op(V, L.p, L.t; L.traits.kwargs...)
+        axpby!(β, W, α, result)
+    end
+
+    return w
+end
 #