Merge pull request #143 from gaurav-arya/ag-kwargs

Propagate kwargs through update_coefficients!

Author: ChrisRackauckas

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 (by default, none are passed).
+
+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,19 +1,20 @@
 #
 """
-    BatchedDiagonalOperator(diag, [; update_func])
+    BatchedDiagonalOperator(diag; update_func, update_func!, 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::AbstractArray, u, p, t; <accepted kwarg fields>) -> [modifies diag]
 """
 struct BatchedDiagonalOperator{T,D,F,F!} <: AbstractSciMLOperator{T}
     diag::D
     update_func::F
     update_func!::F!
 
     function BatchedDiagonalOperator(diag::AbstractArray, update_func, update_func!)
+
         new{
             eltype(diag),
             typeof(diag),
@@ -25,15 +26,16 @@ struct BatchedDiagonalOperator{T,D,F,F!} <: AbstractSciMLOperator{T}
     end
 end
 
-function BatchedDiagonalOperator(diag::AbstractArray;
-                                 update_func = DEFAULT_UPDATE_FUNC,
-                                 update_func! = DEFAULT_UPDATE_FUNC)
-    BatchedDiagonalOperator(diag, update_func, update_func!)
-end
+function DiagonalOperator(u::AbstractArray;
+                          update_func = DEFAULT_UPDATE_FUNC,
+                          update_func! = DEFAULT_UPDATE_FUNC,
+                          accepted_kwargs = nothing
+                         )
+
+    update_func  = preprocess_update_func(update_func , accepted_kwargs)
+    update_func! = preprocess_update_func(update_func!, accepted_kwargs)
 
-function DiagonalOperator(u::AbstractArray; update_func = DEFAULT_UPDATE_FUNC,
-                          update_func! = DEFAULT_UPDATE_FUNC)
-    BatchedDiagonalOperator(u; update_func = update_func, update_func! = update_func!)
+    BatchedDiagonalOperator(u, update_func, update_func!)
 end
 
 # traits
@@ -46,38 +48,39 @@ 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
 
-function update_coefficients(L::BatchedDiagonalOperator,u,p,t)
-    @set! L.diag = L.update_func(L.diag,u,p,t)
+LinearAlgebra.issymmetric(L::BatchedDiagonalOperator) = true
+function LinearAlgebra.ishermitian(L::BatchedDiagonalOperator)
+    if isreal(L)
+        true
+    else
+        vec(L.diag) |> Diagonal |> ishermitian
+    end
+end
+LinearAlgebra.isposdef(L::BatchedDiagonalOperator) = isposdef(Diagonal(vec(L.diag)))
+
+function update_coefficients(L::BatchedDiagonalOperator,u ,p, t; kwargs...)
+    @set! L.diag = L.update_func(L.diag, u, p, t; kwargs...)
+end
+
+function update_coefficients!(L::BatchedDiagonalOperator, u, p, t; kwargs...)
+    L.update_func!(L.diag, u, p, t; kwargs...)
 end
-update_coefficients!(L::BatchedDiagonalOperator,u,p,t) = (L.update_func!(L.diag,u,p,t); L)
 
 getops(L::BatchedDiagonalOperator) = (L.diag,)
 
 function isconstant(L::BatchedDiagonalOperator)
-    L.update_func == L.update_func! == DEFAULT_UPDATE_FUNC
+    update_func_isconstant(L.update_func) & update_func_isconstant(L.update_func!)
 end
 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)
 
-LinearAlgebra.issymmetric(L::BatchedDiagonalOperator) = true
-function LinearAlgebra.ishermitian(L::BatchedDiagonalOperator)
-    if isreal(L)
-        true
-    else
-        d = vec(L.diag)
-        D = Diagonal(d)
-        ishermitian(d)
-    end
-end
-LinearAlgebra.isposdef(L::BatchedDiagonalOperator) = isposdef(Diagonal(vec(L.diag)))
-
 # operator application
 Base.:*(L::BatchedDiagonalOperator, u::AbstractVecOrMat) = L.diag .* u
 Base.:\(L::BatchedDiagonalOperator, u::AbstractVecOrMat) = L.diag .\ u

src/func.jl

@@ -84,6 +84,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;
@@ -101,6 +102,7 @@ function FunctionOperator(op,
 
                           p=nothing,
                           t::Union{Number,Nothing}=nothing,
+                          accepted_kwargs::NTuple{N,Symbol} = (),
 
                           ifcache::Bool = true,
 
@@ -111,7 +113,7 @@ function FunctionOperator(op,
                           issymmetric::Bool = false,
                           ishermitian::Bool = false,
                           isposdef::Bool = false,
-                         )
+                         ) where{N}
 
     # store eltype of input/output for caching with ComposedOperator.
     eltypes = eltype.((input, output))
@@ -181,6 +183,8 @@ function FunctionOperator(op,
               T = T,
               size = sz,
               eltypes = eltypes,
+              accepted_kwargs = accepted_kwargs,
+              kwargs = Dict{Symbol, Any}(),
              )
 
     L = FunctionOperator(
@@ -191,7 +195,7 @@ function FunctionOperator(op,
                          traits,
                          p,
                          t,
-                         cache,
+                         cache
                         )
 
     if ifcache & isnothing(L.cache)
@@ -201,36 +205,40 @@ function FunctionOperator(op,
     L
 end
 
-function update_coefficients(L::FunctionOperator, u, p, t)
-
-    if isconstant(L)
-        return L
-    end
-
-    @set! L.op = update_coefficients(L.op, u, p, t)
-    @set! L.op_adjoint = update_coefficients(L.op_adjoint, u, p, t)
-    @set! L.op_inverse = update_coefficients(L.op_inverse, u, p, t)
-    @set! L.op_adjoint_inverse = update_coefficients(L.op_adjoint_inverse, u, p, t)
+function update_coefficients(L::FunctionOperator, u, p, t; kwargs...)
 
+    # update p, t
     @set! L.p = p
     @set! L.t = t
 
-    L
-end
+    # filter and update kwargs
+    filtered_kwargs = get_filtered_kwargs(kwargs, L.traits.accepted_kwargs)
+    @set! L.traits.kwargs = Dict{Symbol, Any}(filtered_kwargs)
 
-function update_coefficients!(L::FunctionOperator, u, p, t)
+    isconstant(L) && return L
 
-    if isconstant(L)
-        return L
-    end
+    @set! L.op = update_coefficients(L.op, u, p, t; filtered_kwargs...)
+    @set! L.op_adjoint = update_coefficients(L.op_adjoint, u, p, t; filtered_kwargs...)
+    @set! L.op_inverse = update_coefficients(L.op_inverse, u, p, t; filtered_kwargs...)
+    @set! L.op_adjoint_inverse = update_coefficients(L.op_adjoint_inverse, u, p, t; filtered_kwargs...)
+end
 
-    for op in getops(L)
-        update_coefficients!(op, u, p, t)
-    end
+function update_coefficients!(L::FunctionOperator, u, p, t; kwargs...)
 
+    # update p, t
     L.p = p
     L.t = t
 
+    # filter and update kwargs
+    filtered_kwargs = get_filtered_kwargs(kwargs, L.traits.accepted_kwargs)
+    L.traits = (; L.traits..., kwargs = Dict{Symbol, Any}(filtered_kwargs))
+
+    isconstant(L) && return
+
+    for op in getops(L)
+        update_coefficients!(op, u, p, t; filtered_kwargs...)
+    end
+
     L
 end
 
@@ -267,9 +275,6 @@ function Base.adjoint(L::FunctionOperator)
     @set! traits.size = reverse(size(L))
     @set! traits.eltypes = reverse(traits.eltypes)
 
-    p = L.p
-    t = L.t
-
     cache = if iscached(L)
         cache = reverse(L.cache)
     else
@@ -281,8 +286,8 @@ function Base.adjoint(L::FunctionOperator)
                      op_inverse,
                      op_adjoint_inverse,
                      traits,
-                     p,
-                     t,
+                     L.p,
+                     L.t,
                      cache,
                     )
 end
@@ -310,9 +315,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
@@ -324,8 +326,8 @@ function Base.inv(L::FunctionOperator)
                      op_inverse,
                      op_adjoint_inverse,
                      traits,
-                     p,
-                     t,
+                     L.p,
+                     L.t,
                      cache,
                     )
 end
@@ -353,8 +355,8 @@ function LinearAlgebra.opnorm(L::FunctionOperator, p)
       argument. E.g., `(p::Real) -> p == Inf ? 100 : error("only Inf norm is
       defined")`
     """)
-    opn = L.opnorm
-    return opn isa Number ? opn : L.opnorm(p)
+    opn = L.traits.opnorm
+    return opn isa Number ? opn : L.traits.opnorm(p)
 end
 LinearAlgebra.issymmetric(L::FunctionOperator) = L.traits.issymmetric
 LinearAlgebra.ishermitian(L::FunctionOperator) = L.traits.ishermitian
@@ -373,31 +375,36 @@ end
 islinear(L::FunctionOperator) = L.traits.islinear
 isconstant(L::FunctionOperator) = L.traits.isconstant
 has_adjoint(L::FunctionOperator) = !(L.op_adjoint isa Nothing)
-has_mul(L::FunctionOperator{iip}) where{iip} = true
-has_mul!(L::FunctionOperator{iip}) where{iip} = iip
+has_mul(::FunctionOperator{iip}) where{iip} = true
+has_mul!(::FunctionOperator{iip}) where{iip} = iip
 has_ldiv(L::FunctionOperator{iip}) where{iip} = !(L.op_inverse isa Nothing)
 has_ldiv!(L::FunctionOperator{iip}) where{iip} = iip & !(L.op_inverse isa Nothing)
 
 # TODO - FunctionOperator, Base.conj, transpose
 
 # 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)
+function Base.:*(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip}
+    L.op(u, L.p, L.t; L.traits.kwargs...)
+end
+
+function Base.:\(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip}
+    L.op_inverse(u, L.p, L.t; L.traits.kwargs...)
+end
 
 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.traits.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.traits.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.traits.kwargs...)
 end
 
 function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{false}, u::AbstractVecOrMat, args...)
@@ -414,11 +421,11 @@ function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true, oop,
 end
 
 function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true, oop, true}, u::AbstractVecOrMat, α, β) where{oop}
-    L.op(v, u, L.p, L.t, α, β)
+    L.op(v, u, L.p, L.t, α, β; L.traits.kwargs...)
 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.traits.kwargs...)
 end
 
 function LinearAlgebra.ldiv!(L::FunctionOperator{true}, u::AbstractVecOrMat)