Updated interface for operators, addition of new and modified tests

Author: divital-coder

src/SciMLOperators.jl

@@ -28,19 +28,16 @@ A lazy operator algebra is also defined for `AbstractSciMLOperator`s.
 
 # 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,29 @@ 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
+    update_coefficients(nn, u, p, t; kwargs...)
+    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
+    update_coefficients!(nn, u, p, t; kwargs...)
+    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
+    update_coefficients!(nn, u, p, t; kwargs...)
+    lmul!(β, w)
+    w
+end
+
 # operator fusion, composition
 for op in (:*, :∘)
     @eval function Base.$op(nn::NullOperator, A::AbstractSciMLOperator)
@@ -336,6 +380,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
 
@@ -538,6 +619,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
@@ -792,6 +902,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
 """
@@ -909,4 +1054,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/interface.jl

@@ -43,7 +43,7 @@ $(UPDATE_COEFFS_WARNING)
 # Example
 
 ```
-using SciMLOperator
+using SciMLOperators
 
 mat_update_func = (A, u, p, t; scale = 1.0) -> p * p' * scale * t
 
@@ -56,8 +56,12 @@ u = rand(4)
 p = rand(4)
 t = 1.0
 
+# Update the operator and apply it to `u`
 L = update_coefficients(L, u, p, t; scale = 2.0)
-L * u
+result = L * u
+
+# Or use the interface which separrates the update from the application
+result = L(u, u, p, t; scale = 2.0)
 ```
 
 """
@@ -108,6 +112,7 @@ end
 ###
 # operator evaluation interface
 ###
+
 # Out-of-place: v is action vector, u is update vector
 function (L::AbstractSciMLOperator)(v, u, p, t; kwargs...)
     update_coefficients(L, u, p, t; kwargs...) * v

src/left.jl

@@ -157,4 +157,81 @@ for (op, LType, VType) in ((:adjoint, :AdjointOperator, :AbstractAdjointVecOrMat
         u
     end
 end
+
+
+# For AdjointOperator
+# Out-of-place: v is action vector, u is update vector
+function (L::AdjointOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    # Adjoint operator applied to v means L.L' * v
+    # For matrices: (A')v = (v'A)'
+    # This means we need to compute L.L(v', u, p, t)' 
+    # Reshape v to match the adjoint operator's expected size
+    adjv = reshape(v', size(L.L, 1), :)
+    result = L.L(adjv, u, p, t; kwargs...)
+    return reshape(result', size(L, 1), :)
+end
+
+# In-place: w is destination, v is action vector, u is update vector
+function (L::AdjointOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    # Need temporary storage for adjoint operations
+    temp_v = reshape(v', size(L.L, 1), :)
+    temp_w = similar(temp_v)
+    
+    # Apply the internal operator
+    L.L(temp_w, temp_v, u, p, t; kwargs...)
+    
+    # Copy back to w with adjoint
+    w .= reshape(temp_w', size(L, 1), :)
+    return w
+end
+
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::AdjointOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    # Handle scaling of existing w
+    if β != 1.0
+        lmul!(β, w)
+    end
+    
+    # Need temporary storage for adjoint operations
+    temp_v = reshape(v', size(L.L, 1), :)
+    temp_w = similar(temp_v)
+    
+    # Apply the internal operator
+    L.L(temp_w, temp_v, u, p, t, 1.0, 0.0; kwargs...)
+    
+    # Add α * result' to w
+    w .+= α .* reshape(temp_w', size(L, 1), :)
+    return w
+end
+
+# For TransposedOperator
+function (L::TransposedOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
+    transv = reshape(v', size(L.L, 1), :)
+    result = L.L(transv, u, p, t; kwargs...)
+    return reshape(result', size(L, 1), :)
+end
+
+function (L::TransposedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
+    temp_v = reshape(v', size(L.L, 1), :)
+    temp_w = similar(temp_v)
+    
+    L.L(temp_w, temp_v, u, p, t; kwargs...)
+    
+    w .= reshape(temp_w', size(L, 1), :)
+    return w
+end
+
+function (L::TransposedOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
+    if β != 1.0
+        lmul!(β, w)
+    end
+    
+    temp_v = reshape(v', size(L.L, 1), :)
+    temp_w = similar(temp_v)
+    
+    L.L(temp_w, temp_v, u, p, t, 1.0, 0.0; kwargs...)
+    
+    w .+= α .* reshape(temp_w', size(L, 1), :)
+    return w
+end
 #