function op working

Author: vpuri3

src/func.jl

@@ -2,7 +2,7 @@
 """
     Matrix free operators (given by a function)
 """
-mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,K,C} <: AbstractSciMLOperator{T}
+mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSciMLOperator{T}
     """ Function with signature op(u, p, t) and (if isinplace) op(du, u, p, t) """
     op::F
     """ Adjoint operator"""
@@ -17,8 +17,8 @@ mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,K,C}
     p::P
     """ Time """
     t::Tt
-    """ Keyword arguments """
-    kwargs::K
+    """ kwargs """
+    kwargs::Dict{Symbol,Any} # TODO move inside traits later
     """ Cache """
     cache::C
 
@@ -30,7 +30,7 @@ mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,K,C}
                               traits,
                               p,
                               t,
-                              accepted_kwargs,
+                              kwargs,
                               cache
                              )
 
@@ -51,7 +51,6 @@ mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,K,C}
             typeof(traits),
             typeof(p),
             typeof(t),
-            typeof(accepted_kwargs),
             typeof(cache),
            }(
              op,
@@ -61,7 +60,7 @@ mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,K,C}
              traits,
              p,
              t,
-             accepted_kwargs,
+             kwargs,
              cache,
             )
     end
@@ -107,7 +106,7 @@ function FunctionOperator(op,
 
                           p=nothing,
                           t::Union{Number,Nothing}=nothing,
-                          accepted_kwargs = (),
+                          accepted_kwargs::NTuple{N,Symbol} = (),
 
                           ifcache::Bool = true,
 
@@ -118,13 +117,14 @@ 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))
     sz = (size(output, 1), size(input, 1))
     T  = isnothing(T) ? promote_type(eltypes...) : T
     t  = isnothing(t) ? zero(real(T)) : t
+    kwargs = Dict{Symbol, Any}()
 
     isinplace = if isnothing(isinplace)
         static_hasmethod(op, typeof((output, input, p, t)))
@@ -188,6 +188,7 @@ function FunctionOperator(op,
               T = T,
               size = sz,
               eltypes = eltypes,
+              accepted_kwargs = accepted_kwargs,
              )
 
     L = FunctionOperator(
@@ -198,7 +199,7 @@ function FunctionOperator(op,
                          traits,
                          p,
                          t,
-                         normalize_kwargs(accepted_kwargs),
+                         kwargs,
                          cache
                         )
 
@@ -214,9 +215,10 @@ function update_coefficients(L::FunctionOperator, u, p, t; kwargs...)
     @set! L.p = p
     @set! L.t = t
 
-    isconstant(L) && return L
+    filtered_kwargs = get_filtered_kwargs(kwargs, L.traits.accepted_kwargs)
+    @set! L.kwargs = Dict(filtered_kwargs)
 
-    filtered_kwargs = (kwarg => kwargs[kwarg] for kwarg in L.kwargs if haskey(kwargs, kwarg))
+    isconstant(L) && return L
 
     @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...)
@@ -226,17 +228,18 @@ end
 
 function update_coefficients!(L::FunctionOperator, u, p, t; kwargs...)
 
-    isconstant(L) && return
+    L.p = p
+    L.t = t
+
+    filtered_kwargs = get_filtered_kwargs(kwargs, L.traits.accepted_kwargs)
+    L.kwargs = Dict(filtered_kwargs)
 
-    filtered_kwargs = (kwarg => kwargs[kwarg] for kwarg in L.kwargs if haskey(kwargs, kwarg))
+    isconstant(L) && return
 
     for op in getops(L)
         update_coefficients!(op, u, p, t; filtered_kwargs...)
     end
 
-    L.p = p
-    L.t = t
-
     L
 end
 
@@ -383,8 +386,13 @@ has_ldiv!(L::FunctionOperator{iip}) where{iip} = iip & !(L.op_inverse isa Nothin
 # 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; L.kwargs...)
+function Base.:*(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip}
+    L.op(u, L.p, L.t; L.kwargs...)
+end
+
+function Base.:\(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip}
+    L.op_inverse(u, L.p, L.t; L.kwargs...)
+end
 
 function Base.:*(L::FunctionOperator{true,false}, u::AbstractVecOrMat)
     _, co = L.cache

src/utils.jl

@@ -18,13 +18,16 @@ struct FilterKwargs{F,K}
     f::F
     accepted_kwargs::K
 end
-function (f_filter::FilterKwargs)(args...; kwargs...)
-    # Filter keyword arguments to those accepted by function.
-    # Avoid throwing errors here if a keyword argument is not provided: defer this to the function call for a more readable error.
-    filtered_kwargs = (kwarg => kwargs[kwarg] for kwarg in f_filter.accepted_kwargs if haskey(kwargs, kwarg))
-    f_filter.f(args...; filtered_kwargs...)
+
+# Filter keyword arguments to those accepted by function.
+# Avoid throwing errors here if a keyword argument is not provided: defer
+# this to the function call for a more readable error.
+function get_filtered_kwargs(kwargs::Base.Pairs, accepted_kwargs::NTuple{N,Symbol}) where{N}
+    (kw => kwargs[kw] for kw in accepted_kwargs if haskey(kwargs, kw))
+end
+
+function (f::FilterKwargs)(args...; kwargs...)
+    filtered_kwargs = get_filtered_kwargs(kwargs, f.accepted_kwargs)
+    f.f(args...; filtered_kwargs...)
 end
-# automatically convert NamedTuple's, etc. to a normalized kwargs representation (i.e. Base.Pairs) 
-normalize_kwargs(; kwargs...) = kwargs
-normalize_kwargs(kwargs) = normalize_kwargs(; kwargs...)
 #

test/func.jl

@@ -107,10 +107,11 @@ end
     scale = rand()
 
     # Accept a kwarg "scale" in operator action
-    f(du,u,p,t; scale) = mul!(du, Diagonal(p*t*scale), u)
-    f(u, p, t; scale) = Diagonal(p * t * scale) * u
+    f(du,u,p,t; scale = 1.0) = mul!(du, Diagonal(p*t*scale), u)
+    f(u, p, t; scale = 1.0) = Diagonal(p * t * scale) * u
 
-    L = FunctionOperator(f, u, u; p=zero(p), t=zero(t), accepted_kwargs=(;scale=zero(scale)))
+    L = FunctionOperator(f, u, u; p=zero(p), t=zero(t),
+                         accepted_kwargs = (:scale,))
 
     ans = @. u * p * t * scale 
     @test L(u,p,t; scale) ≈ ans

test/scalar.jl

@@ -151,7 +151,7 @@ end
     # Test scalar operator which expects keyword argument to update,
     # modeled in the style of a DiffEq W-operator.
     γ = ScalarOperator(0.0; update_func = (args...; dtgamma) -> dtgamma,
-                       accepted_kwargs=(:dtgamma,))
+                       accepted_kwargs = (:dtgamma,))
 
     dtgamma = rand()
     @test γ(u,p,t; dtgamma) ≈ dtgamma * u

test/total.jl

@@ -86,7 +86,7 @@ end
     C = rand(N,N)
     # Introduce update function for D dependent on kwarg "matrix"
     D = MatrixOperator(zeros(N,N); update_func=(A, u, p, t; matrix) -> (A .= p*t*matrix), 
-                       accepted_kwargs=(:matrix,))
+                       accepted_kwargs = (:matrix,))
 
     u = rand(N2,K)
     p = rand()
@@ -99,11 +99,9 @@ end
     T1 = ⊗(A, B)
     T2 = ⊗(C, D)
 
-    # Introduce update function for D1
-    D1  = DiagonalOperator(zeros(N2); update_func=(d, u, p, t) -> (d .= p))
-    # Introduce update funcion for D2 dependent on kwarg "diag" 
-    D2  = DiagonalOperator(zeros(N2); update_func=(d, u, p, t; diag) -> (d .= p*t*diag),
-                           accepted_kwargs=(:diag,))
+    D1  = DiagonalOperator(zeros(N2); update_func = (d, u, p, t) -> p)
+    D2  = DiagonalOperator(zeros(N2); update_func = (d, u, p, t; diag) -> p*t*diag,
+                           accepted_kwargs = (:diag,))
 
     TT = [T1, T2]
     DD = Diagonal([D1, D2])