Propagate kwargs through update_coefficients!

docs/src/interface.md

@@ -55,3 +55,26 @@ the proof to affine operators, so then ``exp(A*t)*v`` operations via Krylov meth
 affine as well, and all sorts of things. Thus affine operators have no matrix representation but they 
 are still compatible with essentially any Krylov method which would otherwise be compatible with
 matrix-free representations, hence their support in the SciMLOperators interface.
+
+## Note about keyword arguments to `update_coefficients!`
+
+In rare cases, an operator may be used in a context where additional state is expected to be provided
+to `update_coefficients!` beyond `u`, `p`, and `t`. In this case, the operator may accept this additional
+state through arbitrary keyword arguments to `update_coefficients!`. When the caller provides these, they will be recursively propagated downwards through composed operators just like `u`, `p`, and `t`, and provided to the operator.
+For the [premade SciMLOperators](premade_operators.md), one can specify the keyword arguments used by an operator with an `accepted_kwargs` argument that defaults to an empty tuple.
+
+In the below example, we create an operator that gleefully ignores `u`, `p`, and `t` and uses its own special scaling.
+```@example
+using SciMLOperators
+
+γ = ScalarOperator(0.0; update_func=(a, u, p, t; my_special_scaling) -> my_special_scaling,
+                   accepted_kwargs=(:my_special_scaling,))
+
+# Update coefficients, then apply operator
+update_coefficients!(γ, nothing, nothing, nothing; my_special_scaling=7.0)
+@show γ * [2.0]
+
+# Use operator application form
+@show γ([2.0], nothing, nothing; my_special_scaling = 5.0)
+nothing # hide
+```

src/batch.jl

@@ -1,33 +1,35 @@
 #
 """
-    BatchedDiagonalOperator(diag, [; update_func])
+    BatchedDiagonalOperator(diag; update_func=nothing, accepted_kwargs=())
 
 Represents a time-dependent elementwise scaling (diagonal-scaling) operation.
 Acts on `AbstractArray`s of the same size as `diag`. The update function is called
 by `update_coefficients!` and is assumed to have the following signature:
 
-    update_func(diag::AbstractVector,u,p,t) -> [modifies diag]
+    update_func(diag::AbstractVector,u,p,t; <accepted kwarg fields>) -> [modifies diag]
 """
 struct BatchedDiagonalOperator{T,D,F} <: AbstractSciMLOperator{T}
     diag::D
     update_func::F
 
     function BatchedDiagonalOperator(
                                      diag::AbstractArray;
-                                     update_func=DEFAULT_UPDATE_FUNC
+                                     update_func=nothing,
+                                     accepted_kwargs=()
                                     )
+        _update_func = preprocess_update_func(update_func, accepted_kwargs)
         new{
             eltype(diag),
             typeof(diag),
-            typeof(update_func)
+            typeof(_update_func)
            }(
-             diag, update_func,
+             diag, _update_func,
             )
     end
 end
 
-function DiagonalOperator(u::AbstractArray; update_func=DEFAULT_UPDATE_FUNC)
-    BatchedDiagonalOperator(u; update_func=update_func)
+function DiagonalOperator(u::AbstractArray; update_func=nothing, accepted_kwargs=())
+    BatchedDiagonalOperator(u; update_func, accepted_kwargs)
 end
 
 # traits
@@ -40,7 +42,7 @@ function Base.conj(L::BatchedDiagonalOperator) # TODO - test this thoroughly
     update_func = if isreal(L)
         L.update_func
     else
-        (L,u,p,t) -> conj(L.update_func(conj(L.diag),u,p,t))
+        (L,u,p,t; kwargs...) -> conj(L.update_func(conj(L.diag),u,p,t; kwargs...))
     end
     BatchedDiagonalOperator(diag; update_func=update_func)
 end
@@ -57,15 +59,15 @@ function LinearAlgebra.ishermitian(L::BatchedDiagonalOperator)
 end
 LinearAlgebra.isposdef(L::BatchedDiagonalOperator) = isposdef(Diagonal(vec(L.diag)))
 
-isconstant(L::BatchedDiagonalOperator) = L.update_func == DEFAULT_UPDATE_FUNC
+isconstant(L::BatchedDiagonalOperator) = update_func_isconstant(L.update_func)
 islinear(::BatchedDiagonalOperator) = true
 has_adjoint(L::BatchedDiagonalOperator) = true
 has_ldiv(L::BatchedDiagonalOperator) = all(x -> !iszero(x), L.diag)
 has_ldiv!(L::BatchedDiagonalOperator) = has_ldiv(L)
 
 getops(L::BatchedDiagonalOperator) = (L.diag,)
 
-update_coefficients!(L::BatchedDiagonalOperator,u,p,t) = (L.update_func(L.diag,u,p,t); nothing)
+update_coefficients!(L::BatchedDiagonalOperator,u,p,t; kwargs...) = (L.update_func(L.diag,u,p,t; kwargs...); nothing)
 
 # operator application
 Base.:*(L::BatchedDiagonalOperator, u::AbstractVecOrMat) = L.diag .* u

src/func.jl

@@ -2,7 +2,7 @@
 """
     Matrix free operators (given by a function)
 """
-mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSciMLOperator{T}
+mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,K,C} <: AbstractSciMLOperator{T}
     """ Function with signature op(u, p, t) and (if isinplace) op(du, u, p, t) """
     op::F
     """ Adjoint operator"""
@@ -17,6 +17,8 @@ mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: Abst
     p::P
     """ Time """
     t::Tt
+    """ Keyword arguments """
+    kwargs::K
     """ Cache """
     cache::C
 
@@ -28,6 +30,7 @@ mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: Abst
                               traits,
                               p,
                               t,
+                              accepted_kwargs,
                               cache
                              )
 
@@ -46,6 +49,7 @@ mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: Abst
             typeof(traits),
             typeof(p),
             typeof(t),
+            typeof(accepted_kwargs),
             typeof(cache),
            }(
              op,
@@ -55,6 +59,7 @@ mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: Abst
              traits,
              p,
              t,
+             accepted_kwargs,
              cache,
             )
     end
@@ -82,6 +87,7 @@ function FunctionOperator(op,
     FunctionOperator(op, input, output; kwargs...)
 end
 
+# TODO: document constructor and revisit design as needed (e.g. for "accepted_kwargs")
 function FunctionOperator(op,
                           input::AbstractVecOrMat,
                           output::AbstractVecOrMat =  input;
@@ -96,6 +102,7 @@ function FunctionOperator(op,
 
                           p=nothing,
                           t::Union{Number,Nothing}=nothing,
+                          accepted_kwargs=(),
 
                           ifcache::Bool = true,
 
@@ -169,17 +176,18 @@ function FunctionOperator(op,
                          traits,
                          p,
                          t,
-                         cache,
+                         normalize_kwargs(accepted_kwargs),
+                         cache
                         )
 
     ifcache ? cache_operator(L, input, output) : L
 end
 
-function update_coefficients(L::FunctionOperator, u, p, t)
-    op = update_coefficients(L.op, u, p, t)
-    op_adjoint = update_coefficients(L.op_adjoint, u, p, t)
-    op_inverse = update_coefficients(L.op_inverse, u, p, t)
-    op_adjoint_inverse = update_coefficients(L.op_adjoint_inverse, u, p, t)
+function update_coefficients(L::FunctionOperator, u, p, t; kwargs...)
+    op = update_coefficients(L.op, u, p, t; kwargs...)
+    op_adjoint = update_coefficients(L.op_adjoint, u, p, t; kwargs...)
+    op_inverse = update_coefficients(L.op_inverse, u, p, t; kwargs...)
+    op_adjoint_inverse = update_coefficients(L.op_adjoint_inverse, u, p, t; kwargs...)
 
     FunctionOperator(op,
                      op_adjoint,
@@ -188,18 +196,20 @@ function update_coefficients(L::FunctionOperator, u, p, t)
                      L.traits,
                      p,
                      t,
+                     accepted_kwargs=kwargs,
                      L.cache
                     )
 end
 
-function update_coefficients!(L::FunctionOperator, u, p, t)
+function update_coefficients!(L::FunctionOperator, u, p, t; kwargs...)
     ops = getops(L)
     for op in ops
-        update_coefficients!(op, u, p, t)
+        update_coefficients!(op, u, p, t; kwargs...)
     end
 
     L.p = p
     L.t = t
+    L.kwargs = kwargs
 
     nothing
 end
@@ -229,9 +239,6 @@ function Base.adjoint(L::FunctionOperator)
     traits = L.traits
     @set! traits.size = reverse(size(L))
 
-    p = L.p
-    t = L.t
-
     cache = if iscached(L)
         cache = reverse(L.cache)
     else
@@ -243,8 +250,9 @@ function Base.adjoint(L::FunctionOperator)
                      op_inverse,
                      op_adjoint_inverse,
                      traits,
-                     p,
-                     t,
+                     L.p,
+                     L.t,
+                     L.kwargs,
                      cache,
                     )
 end
@@ -271,9 +279,6 @@ function Base.inv(L::FunctionOperator)
         (p::Real) -> 1 / traits.opnorm(p)
     end
 
-    p = L.p
-    t = L.t
-
     cache = if iscached(L)
         cache = reverse(L.cache)
     else
@@ -285,8 +290,9 @@ function Base.inv(L::FunctionOperator)
                      op_inverse,
                      op_adjoint_inverse,
                      traits,
-                     p,
-                     t,
+                     L.p,
+                     L.t,
+                     L.kwargs,
                      cache,
                     )
 end
@@ -326,22 +332,22 @@ has_ldiv!(L::FunctionOperator{iip}) where{iip} = iip & !(L.op_inverse isa Nothin
 
 # operator application
 Base.:*(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip} = L.op(u, L.p, L.t)
-Base.:\(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip} = L.op_inverse(u, L.p, L.t)
+Base.:\(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip} = L.op_inverse(u, L.p, L.t; L.kwargs...)
 
 function Base.:*(L::FunctionOperator{true,false}, u::AbstractVecOrMat)
     _, co = L.cache
     du = zero(co)
-    L.op(du, u, L.p, L.t)
+    L.op(du, u, L.p, L.t; L.kwargs...)
 end
 
 function Base.:\(L::FunctionOperator{true,false}, u::AbstractVecOrMat)
     ci, _ = L.cache
     du = zero(ci)
-    L.op_inverse(du, u, L.p, L.t)
+    L.op_inverse(du, u, L.p, L.t; L.kwargs...)
 end
 
 function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true}, u::AbstractVecOrMat)
-    L.op(v, u, L.p, L.t)
+    L.op(v, u, L.p, L.t; L.kwargs...)
 end
 
 function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{false}, u::AbstractVecOrMat, args...)
@@ -358,7 +364,7 @@ function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true}, u::A
 end
 
 function LinearAlgebra.ldiv!(v::AbstractVecOrMat, L::FunctionOperator{true}, u::AbstractVecOrMat)
-    L.op_inverse(v, u, L.p, L.t)
+    L.op_inverse(v, u, L.p, L.t; L.kwargs...)
 end
 
 function LinearAlgebra.ldiv!(L::FunctionOperator{true}, u::AbstractVecOrMat)

src/interface.jl

@@ -15,19 +15,33 @@ out-of-place form B = update_coefficients(A,u,p,t).
 """
 function (::AbstractSciMLOperator) end
 
+# Utilities for update functions
 DEFAULT_UPDATE_FUNC(A,u,p,t) = A
+function preprocess_update_func(update_func, accepted_kwargs)
+    update_func = (update_func === nothing) ? DEFAULT_UPDATE_FUNC : update_func
+    # accepted_kwargs can be passed as nothing to indicate that we should not filter 
+    # (e.g. if the function already accepts all kwargs...). 
+    return (accepted_kwargs === nothing) ? update_func : FilterKwargs(update_func, accepted_kwargs)
+end
+function update_func_isconstant(update_func)
+    if update_func isa FilterKwargs
+        return update_func.f == DEFAULT_UPDATE_FUNC
+    else
+        return update_func == DEFAULT_UPDATE_FUNC
+    end
+end
 
-update_coefficients!(L,u,p,t) = nothing
-update_coefficients(L,u,p,t) = L
-function update_coefficients!(L::AbstractSciMLOperator, u, p, t)
+update_coefficients!(L,u,p,t; kwargs...) = nothing
+update_coefficients(L,u,p,t; kwargs...) = L
+function update_coefficients!(L::AbstractSciMLOperator, u, p, t; kwargs...)
     for op in getops(L)
-        update_coefficients!(op, u, p, t)
+        update_coefficients!(op, u, p, t; kwargs...)
     end
     nothing
 end
 
-(L::AbstractSciMLOperator)(u, p, t) = (update_coefficients!(L, u, p, t); L * u)
-(L::AbstractSciMLOperator)(du, u, p, t) = (update_coefficients!(L, u, p, t); mul!(du, L, u))
+(L::AbstractSciMLOperator)(u, p, t; kwargs...) = (update_coefficients!(L, u, p, t; kwargs...); L * u)
+(L::AbstractSciMLOperator)(du, u, p, t; kwargs...) = (update_coefficients!(L, u, p, t; kwargs...); mul!(du, L, u))
 
 ###
 # caching interface