Operators in work, modificataions -> operator update interface, operator application and evaluation interfaces, tests

Author: divital-coder

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

@@ -1,7 +1,7 @@
 """
 $(README)
 """
-module SciMLOperators# Temporary fix while we debug method overwriting issues
+module SciMLOperators
 
 using DocStringExtensions
 

src/basic.jl

@@ -218,6 +218,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
 

src/batch.jl

@@ -82,12 +82,12 @@ function LinearAlgebra.ishermitian(L::BatchedDiagonalOperator)
 end
 LinearAlgebra.isposdef(L::BatchedDiagonalOperator) = isposdef(Diagonal(vec(L.diag)))
 
-function update_coefficients(L::BatchedDiagonalOperator, u_update, p, t; kwargs...)
-    @reset L.diag = L.update_func(L.diag, u_update, p, t; kwargs...)
+function update_coefficients(L::BatchedDiagonalOperator, u, p, t; kwargs...)
+    @reset L.diag = L.update_func(L.diag, u, p, t; kwargs...)
 end
 
-function update_coefficients!(L::BatchedDiagonalOperator, u_update, p, t; kwargs...)
-    L.update_func!(L.diag, u_update, p, t; kwargs...)
+function update_coefficients!(L::BatchedDiagonalOperator, u, p, t; kwargs...)
+    L.update_func!(L.diag, u, p, t; kwargs...)
 
     nothing
 end

src/func.jl

@@ -357,7 +357,7 @@ end
 end
 @inline __and_val(vs...) = mapreduce(_unwrap_val, *, vs)
 
-function update_coefficients(L::FunctionOperator, u_update, p, t; kwargs...)
+function update_coefficients(L::FunctionOperator, u, p, t; kwargs...)
 
     # update p, t
     L = set_p(L, p)
@@ -370,14 +370,14 @@ function update_coefficients(L::FunctionOperator, u_update, p, t; kwargs...)
 
     isconstant(L) && return L
 
-    L = set_op(L, update_coefficients(L.op, u_update, p, t; filtered_kwargs...))
-    L = set_op_adjoint(L, update_coefficients(L.op_adjoint, u_update, p, t; filtered_kwargs...))
-    L = set_op_inverse(L, update_coefficients(L.op_inverse, u_update, p, t; filtered_kwargs...))
+    L = set_op(L, update_coefficients(L.op, u, p, t; filtered_kwargs...))
+    L = set_op_adjoint(L, update_coefficients(L.op_adjoint, u, p, t; filtered_kwargs...))
+    L = set_op_inverse(L, update_coefficients(L.op_inverse, u, p, t; filtered_kwargs...))
     L = set_op_adjoint_inverse(L,
-        update_coefficients(L.op_adjoint_inverse, u_update, p, t; filtered_kwargs...))
+        update_coefficients(L.op_adjoint_inverse, u, p, t; filtered_kwargs...))
 end
 
-function update_coefficients!(L::FunctionOperator, u_update, p, t; kwargs...)
+function update_coefficients!(L::FunctionOperator, u, p, t; kwargs...)
 
     # update p, t
     L.p = p
@@ -390,7 +390,7 @@ function update_coefficients!(L::FunctionOperator, u_update, p, t; kwargs...)
     isconstant(L) && return
 
     for op in getops(L)
-        update_coefficients!(op, u_update, p, t; filtered_kwargs...)
+        update_coefficients!(op, u, p, t; filtered_kwargs...)
     end
 
     nothing
@@ -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
 #

src/interface.jl

@@ -108,204 +108,25 @@ end
 ###
 # operator evaluation interface
 ###
-
-abstract type OperatorMethodTag end
-struct OutOfPlaceTag <: OperatorMethodTag end
-struct InPlaceTag <: OperatorMethodTag end
-struct ScaledInPlaceTag <: OperatorMethodTag end
-
-"""
-$SIGNATURES
-
-Apply the operator L to the vector x, after updating the coefficients of L using u_update.
-
-This method is out-of-place, i.e., it allocates a new vector for the result.
-
-# Arguments
-- `x`: The vector to which the operator is applied
-- `p`: Parameter object
-- `t`: Time parameter
-- `u_update`: Vector used to update the operator coefficients (defaults to `x`)
-- `kwargs...`: Additional keyword arguments for the update function
-
-# Returns
-The result of applying the operator to x
-
-# Example
-```julia
-L = MatrixOperator(zeros(4,4); update_func = some_update_func)
-x = rand(4)
-p = some_params
-t = 1.0
-u_update = rand(4)  # Some reference state for updating coefficients
-
-# Update using u_update, then apply operator to x
-v = L(x, p, t, u_update)
-```
-"""
-function (L::AbstractSciMLOperator{T})(x, p, t, u_update=x; kwargs...) where {T}
-    _check_device_match(x, u_update)
-    _check_size_compatibility(L, u_update, x)
-    update_coefficients(L, u_update, p, t; kwargs...) * x
+# 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
 end
-
-"""
-$SIGNATURES
-
-Apply the operator L to the vector u in-place, storing the result in du,
-after updating the coefficients of L using u_update.
-
-# Arguments
-- `du`: The output vector where the result is stored
-- `u`: The vector to which the operator is applied
-- `p`: Parameter object
-- `t`: Time parameter
-- `u_update`: Vector used to update the operator coefficients (defaults to `u`)
-- `kwargs...`: Additional keyword arguments for the update function
-
-# Example
-```julia
-L = MatrixOperator(zeros(4,4); update_func = some_update_func)
-u = rand(4)
-du = similar(u)
-p = some_params
-t = 1.0
-u_update = rand(4)  # Some reference state for updating coefficients
-
-# Update using u_update, then apply operator to u, storing in du
-L(du, u, p, t, u_update)
-```
-"""
-function (L::AbstractSciMLOperator{T})(du::AbstractArray, u::AbstractArray, p, t, u_update=u; kwargs...) where {T}
-    _check_device_match(du, u, u_update)
-    _check_size_compatibility(L, u_update, u, du)
-    update_coefficients!(L, u_update, p, t; kwargs...)
-    mul!(du, L, u)
-    return du  # Explicitly return du
+# In-place: w is destination, v is action vector, u is update vector
+function (L::AbstractSciMLOperator)(w, v, u, p, t; kwargs...)
+    (update_coefficients!(L, u, p, t; kwargs...); mul!(w, L, v))
 end
-
-"""
-$SIGNATURES
-
-Apply the operator L to vector u with scaling factors α and β, computing du = α*L*u + β*du,
-after updating the coefficients of L using u_update.
-
-# Arguments
-- `du`: The output vector where the result is accumulated
-- `u`: The vector to which the operator is applied
-- `p`: Parameter object
-- `t`: Time parameter
-- `α`: Scaling factor for L*u
-- `β`: Scaling factor for the existing value in du
-- `u_update`: Vector used to update the operator coefficients (defaults to `u`)
-- `kwargs...`: Additional keyword arguments for the update function
-
-# Example
-```julia
-L = MatrixOperator(zeros(4,4); update_func = some_update_func)
-u = rand(4)
-du = rand(4)
-p = some_params
-t = 1.0
-α = 2.0
-β = 1.0
-u_update = rand(4)  # Some reference state for updating coefficients
-
-# Compute du = α*L*u + β*du
-L(du, u, p, t, α, β, u_update)
-```
-"""
-function (L::AbstractSciMLOperator{T})(du::AbstractArray, u::AbstractArray, p, t, α, β, u_update=u; kwargs...) where {T}
-    _check_device_match(du, u, u_update)
-    _check_size_compatibility(L, u_update, u, du)
-    update_coefficients!(L, u_update, p, t; kwargs...)
-    mul!(du, L, u, α, β)
-    return du  # Explicitly return du
+# In-place with scaling: w = α*(L*v) + β*w
+function (L::AbstractSciMLOperator)(w, v, u, p, t, α, β; kwargs...)
+    (update_coefficients!(L, u, p, t; kwargs...); mul!(w, L, v, α, β))
 end
 
-function (L::AbstractSciMLOperator)(du::Number, u::Number, p, t, args...; kwargs...)
-    msg = """Nonallocating L(v, u, p, t) type methods are not available for
+function (L::AbstractSciMLOperator)(w::Number, v::Number, u, p, t, args...; kwargs...)
+    msg = """Nonallocating L(w, v, u, p, t) type methods are not available for
     subtypes of `Number`."""
     throw(ArgumentError(msg))
 end
 
-"""
-@private
-
-Check that all vectors are on the same device (CPU/GPU).
-This function is a no-op in the standard implementation but can be
-extended by packages that provide GPU support.
-"""
-function _check_device_match(args...)
-    # Default implementation - no device checking in base package
-    # This would be extended by GPU-supporting packages
-    nothing
-end
-
-"""
-@private
-
-Verify that the sizes of vectors are compatible with the operator and with each other.
-"""
-function _check_size_compatibility(L::AbstractSciMLOperator, u_update, u, du=nothing)
-    # Special case for scalar operators which have size() = ()
-    if L isa AbstractSciMLScalarOperator
-        # Scalar operators can operate on any size inputs
-        # Just check batch dimensions if present
-        if u isa AbstractMatrix && u_update isa AbstractMatrix
-            if size(u, 2) != size(u_update, 2)
-                throw(DimensionMismatch(
-                    "Batch dimension of u ($(size(u, 2))) must match batch dimension of u_update ($(size(u_update, 2)))"))
-            end
-        end
-
-        if du !== nothing && u isa AbstractMatrix && du isa AbstractMatrix
-            if size(u, 2) != size(du, 2)
-                throw(DimensionMismatch(
-                    "Batch dimension of u ($(size(u, 2))) must match batch dimension of du ($(size(du, 2)))"))
-            end
-        end
-
-        return nothing
-    end
-
-    # For regular operators with dimensions
-    # Verify u_update has compatible size for updating operator
-    if size(u_update, 1) != size(L, 2)
-        throw(DimensionMismatch(
-            "Size of u_update ($(size(u_update, 1))) must match the input dimension of operator ($(size(L, 2)))"))
-    end
-
-    # Verify u has compatible size for operator application
-    if size(u, 1) != size(L, 2)
-        throw(DimensionMismatch(
-            "Size of u ($(size(u, 1))) must match the input dimension of operator ($(size(L, 2)))"))
-    end
-
-    # If du is provided, verify it has compatible size for storing the result
-    if du !== nothing && size(du, 1) != size(L, 1)
-        throw(DimensionMismatch(
-            "Size of du ($(size(du, 1))) must match the output dimension of operator ($(size(L, 1)))"))
-    end
-
-    # Verify batch dimensions match if present
-    if u isa AbstractMatrix && u_update isa AbstractMatrix
-        if size(u, 2) != size(u_update, 2)
-            throw(DimensionMismatch(
-                "Batch dimension of u ($(size(u, 2))) must match batch dimension of u_update ($(size(u_update, 2)))"))
-        end
-    end
-
-    if du !== nothing && u isa AbstractMatrix && du isa AbstractMatrix
-        if size(u, 2) != size(du, 2)
-            throw(DimensionMismatch(
-                "Batch dimension of u ($(size(u, 2))) must match batch dimension of du ($(size(du, 2)))"))
-        end
-    end
-
-    nothing
-end
-
 ###
 # operator caching interface
 ###