From 93bb6216a79cb5a69410e94c9158d5385b07fff7 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 11 Dec 2024 17:12:11 +0530 Subject: [PATCH 01/13] feat: add generic `remake(::AbstractSciMLFunction)` --- src/remake.jl | 81 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/src/remake.jl b/src/remake.jl index 53e1e2911..022310119 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -99,6 +99,87 @@ function remake( _remake_internal(prob; kwargs..., p) end +""" + $(TYPEDSIGNATURES) + +A utility function which merges two `NamedTuple`s `a` and `b`, assuming that the +keys of `a` are a subset of those of `b`. Values in `b` take priority over those +in `a`, except if they are `nothing`. Keys not present in `a` are assumed to have +a value of `nothing`. +""" +function _similar_namedtuple_merge_ignore_nothing(a::NamedTuple, b::NamedTuple) + ks = fieldnames(typeof(b)) + return NamedTuple{ks}(ntuple(Val(length(ks))) do i + something(get(b, ks[i], nothing), get(a, ks[i], nothing), Some(nothing)) + end) +end + +""" + remake(func::AbstractSciMLFunction; f = missing, g = missing, f2 = missing, kwargs...) + +`remake` the given `func`. Return an `AbstractSciMLFunction` of the same kind, `isinplace` and +`specialization` as `func`. Retain the properties of `func`, except those that are overridden +by keyword arguments. For stochastic functions (e.g. `SDEFunction`) the `g` keyword argument +is used to override `func.g`. For split functions (e.g. `SplitFunction`) the `f2` keyword +argument is used to override `func.f2`, and `f` is used for `func.f1`. If +`f isa AbstractSciMLFunction`, properties of `f` will override those of `func` (but not ones +provided via keyword arguments). Properties of `f` that are `nothing` will fall back to those +in `func` (unless provided via keyword arguments). If `f` is a different type of +`AbstractSciMLFunction` from `func`, the returned function will be of the kind of `f`. +""" +function remake(func::AbstractSciMLFunction; f = missing, g = missing, f2 = missing, kwargs...) + # retain iip and spec of original function + iip = isinplace(func) + spec = specialization(func) + # retain properties of original function + props = getproperties(func) + + if f === missing + # if no `f` is provided, create the same type of SciMLFunction + T = parameterless_type(func) + f = func.f + elseif f isa AbstractSciMLFunction + # if `f` is a SciMLFunction, create that type + T = parameterless_type(f) + # properties of `f` take priority over those in the existing `func` + # ignore properties of `f` which are `nothing` but present in `func` + props = _similar_namedtuple_merge_ignore_nothing(props, getproperties(f)) + f = f.f + else + # if `f` is provided but not a SciMLFunction, create the same type + T = parameterless_type(func) + end + + # minor hack to avoid breaking MTK, since prior to ~9.57 in `remake_initialization_data` + # it creates a `NonlinearFunction` inside a `NonlinearFunction`. Just recursively unwrap + # in this case and forget about properties. + while f isa AbstractSciMLFunction + f = f.f + end + + props = @delete props.f + + if isdefined(func, :g) + # For SDEs/SDDEs where `g` is not a keyword + g = coalesce(g, func.g) + + props = @delete props.g + T{iip, spec}(f, g; props..., kwargs...) + elseif isdefined(func, :f2) + # For SplitFunction + # we don't do the same thing as `g`, because for SDEs `g` is + # stored in the problem as well, whereas for Split ODEs etc + # f2 is a part of the function. Thus, if the user provides + # a SciMLFunction for `f` which contains `f2` we use that. + f2 = coalesce(f2, get(props, :f2, missing), func.f2) + + props = @delete props.f2 + T{iip, spec}(f, f2; props..., kwargs...) + else + T{iip, spec}(f; props..., kwargs...) + end +end + """ remake(prob::ODEProblem; f = missing, u0 = missing, tspan = missing, p = missing, kwargs = missing, _kwargs...) From 6c313f8073e91f05990fe12f0126c918d82e234e Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 11 Dec 2024 17:12:39 +0530 Subject: [PATCH 02/13] refactor: use new `remake(::SciMLFunction)` in `remake(::ODEProblem)` --- src/remake.jl | 47 +++++++++++------------------------------------ 1 file changed, 11 insertions(+), 36 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index 022310119..f001d1666 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -216,53 +216,28 @@ function remake(prob::ODEProblem; f = missing, initialization_data = nothing end - if f === missing - if specialization(prob.f) === FunctionWrapperSpecialize - ptspan = promote_tspan(tspan) - if iip - _f = ODEFunction{iip, FunctionWrapperSpecialize}( - wrapfun_iip( - unwrapped_f(prob.f.f), - (newu0, newu0, newp, - ptspan[1])); initialization_data) - else - _f = ODEFunction{iip, FunctionWrapperSpecialize}( - wrapfun_oop( - unwrapped_f(prob.f.f), - (newu0, newp, - ptspan[1])); initialization_data) - end - else - _f = prob.f - if __has_initialization_data(_f) - props = getproperties(_f) - @reset props.initialization_data = initialization_data - props = values(props) - _f = parameterless_type(_f){iip, specialization(_f), map(typeof, props)...}(props...) - end - end - elseif f isa AbstractODEFunction - _f = f - elseif specialization(prob.f) === FunctionWrapperSpecialize + f = coalesce(f, prob.f) + f = remake(prob.f; f, initialization_data) + + if specialization(f) === FunctionWrapperSpecialize ptspan = promote_tspan(tspan) if iip - _f = ODEFunction{iip, FunctionWrapperSpecialize}(wrapfun_iip(f, - (newu0, newu0, newp, - ptspan[1]))) + f = remake( + f; f = wrapfun_iip(unwrapped_f(f.f), (newu0, newu0, newp, ptspan[1]))) else - _f = ODEFunction{iip, FunctionWrapperSpecialize}(wrapfun_oop(f, - (newu0, newp, ptspan[1]))) + f = remake( + f; f = wrapfun_oop(unwrapped_f(f.f), (newu0, newu0, newp, ptspan[1]))) end else - _f = ODEFunction{isinplace(prob), specialization(prob.f)}(f) + _f = f end prob = if kwargs === missing - ODEProblem{isinplace(prob)}( + ODEProblem{iip}( _f, newu0, tspan, newp, prob.problem_type; prob.kwargs..., _kwargs...) else - ODEProblem{isinplace(prob)}(_f, newu0, tspan, newp, prob.problem_type; kwargs...) + ODEProblem{iip}(_f, newu0, tspan, newp, prob.problem_type; kwargs...) end if lazy_initialization === nothing From c4551ddc311899b28195b56721d0de7189b3f6e4 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 11 Dec 2024 17:20:19 +0530 Subject: [PATCH 03/13] refactor: remove old `remake(::Union{SDEFunction, SDDEFunction})` --- src/remake.jl | 36 ------------------------------------ 1 file changed, 36 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index f001d1666..373083f69 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -451,42 +451,6 @@ function remake(prob::SDEProblem; return prob end -""" - remake(func::SDEFunction; f = missing, g = missing, - mass_matrix = missing, analytic = missing, kwargs...) - -Remake the given `SDEFunction`. -""" -function remake(func::Union{SDEFunction, SDDEFunction}; - f = missing, - g = missing, - mass_matrix = missing, - analytic = missing, - sys = missing, - kwargs...) - props = getproperties(func) - props = @delete props.f - props = @delete props.g - @reset props.mass_matrix = coalesce(mass_matrix, func.mass_matrix) - @reset props.analytic = coalesce(analytic, func.analytic) - @reset props.sys = coalesce(sys, func.sys) - - if f === missing - f = func.f - end - - if g === missing - g = func.g - end - - if f isa AbstractSciMLFunction - f = f.f - end - - T = func isa SDEFunction ? SDEFunction : SDDEFunction - return T{isinplace(func)}(f, g; props..., kwargs...) -end - function remake(prob::DDEProblem; f = missing, h = missing, u0 = missing, tspan = missing, p = missing, constant_lags = missing, dependent_lags = missing, order_discontinuity_t0 = missing, From d814e2c98e159e726a029c110295c61737b1523f Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 11 Dec 2024 17:21:04 +0530 Subject: [PATCH 04/13] refactor: remove old `remake(::DDEFunction)` --- src/remake.jl | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index 373083f69..d6c4b4eba 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -517,28 +517,6 @@ function remake(prob::DDEProblem; f = missing, h = missing, u0 = missing, return prob end -function remake(func::DDEFunction; - f = missing, - mass_matrix = missing, - analytic = missing, - sys = missing, - kwargs...) - props = getproperties(func) - props = @delete props.f - @reset props.mass_matrix = coalesce(mass_matrix, func.mass_matrix) - @reset props.analytic = coalesce(analytic, func.analytic) - @reset props.sys = coalesce(sys, func.sys) - - if f === missing - f = func.f - end - if f isa AbstractSciMLFunction - f = f.f - end - - return DDEFunction{isinplace(func)}(f; props..., kwargs...) -end - function remake(prob::SDDEProblem; f = missing, g = missing, From 9d05f66ee7fbfcdbd682a86973cffdbc2c584b4a Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 11 Dec 2024 17:21:23 +0530 Subject: [PATCH 05/13] refactor: remove old `remake(::NonlinearFunction)` --- src/remake.jl | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index d6c4b4eba..ac02adccf 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -704,6 +704,7 @@ function remake(prob::NonlinearProblem; initialization_data = nothing end + f = coalesce(f, prob.f) f = remake(prob.f; f, initialization_data) if problem_type === missing @@ -735,22 +736,6 @@ function remake(prob::NonlinearProblem; return prob end -function remake(func::NonlinearFunction; - f = missing, - kwargs...) - props = getproperties(func) - props = @delete props.f - - if f === missing - f = func.f - end - if f isa AbstractSciMLFunction - f = f.f - end - - return NonlinearFunction{isinplace(func)}(f; props..., kwargs...) -end - """ remake(prob::NonlinearLeastSquaresProblem; f = missing, u0 = missing, p = missing, kwargs = missing, _kwargs...) @@ -773,6 +758,7 @@ function remake(prob::NonlinearLeastSquaresProblem; f = missing, u0 = missing, p initialization_data = nothing end + f = coalesce(f, prob.f) f = remake(prob.f; f, initialization_data) prob = if kwargs === missing From 14feee948e475c6d8d0bbc556f68f63e4954421e Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Wed, 11 Dec 2024 17:22:10 +0530 Subject: [PATCH 06/13] test: add test to ensure `remake` retains properties of `f` passed to it --- test/remake_tests.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/remake_tests.jl b/test/remake_tests.jl index 0a9b156be..3fdf797ff 100644 --- a/test/remake_tests.jl +++ b/test/remake_tests.jl @@ -372,3 +372,14 @@ end prob = ODEProblem(ODEFunction(foo; sys), [1.5, 2.5], (0.0, 1.0), [3.5, 4.5]) @test_nowarn remake(prob; u0 = [:x => nothing], p = [:a => nothing]) end + +@testset "retain properties of `SciMLFunction` passed to `remake`" begin + u0 = [1.0; 2.0; 3.0] + p = [10.0, 20.0, 30.0] + sys = SymbolCache([:x, :y, :z], [:a, :b, :c], :t) + fn = NonlinearFunction(nllorenz!; sys, resid_prototype = zeros(Float64, 3)) + prob = NonlinearProblem(fn, u0, p) + fn2 = NonlinearFunction(nllorenz!; resid_prototype = zeros(Float32, 3)) + prob2 = remake(prob; f = fn2) + @test prob2.f.resid_prototype isa Vector{Float32} +end From 723abd772868848276b6f279bb4c8b446f593d7d Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 12 Dec 2024 11:59:41 +0530 Subject: [PATCH 07/13] fix: fix `remake` for split functions --- src/remake.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index ac02adccf..412e5a011 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -137,14 +137,14 @@ function remake(func::AbstractSciMLFunction; f = missing, g = missing, f2 = miss if f === missing # if no `f` is provided, create the same type of SciMLFunction T = parameterless_type(func) - f = func.f + f = isdefined(func, :f) ? func.f : func.f1 elseif f isa AbstractSciMLFunction # if `f` is a SciMLFunction, create that type T = parameterless_type(f) # properties of `f` take priority over those in the existing `func` # ignore properties of `f` which are `nothing` but present in `func` props = _similar_namedtuple_merge_ignore_nothing(props, getproperties(f)) - f = f.f + f = isdefined(f, :f) ? f.f : f.f1 else # if `f` is provided but not a SciMLFunction, create the same type T = parameterless_type(func) @@ -154,10 +154,11 @@ function remake(func::AbstractSciMLFunction; f = missing, g = missing, f2 = miss # it creates a `NonlinearFunction` inside a `NonlinearFunction`. Just recursively unwrap # in this case and forget about properties. while f isa AbstractSciMLFunction - f = f.f + f = isdefined(f, :f) ? f.f : f.f1 end props = @delete props.f + props = @delete props.f1 if isdefined(func, :g) # For SDEs/SDDEs where `g` is not a keyword From c4a4692eff91fd7193576d71e51037df54564b69 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 12 Dec 2024 14:14:11 +0530 Subject: [PATCH 08/13] feat: add `is_split_function` and `split_function_f_wrapper` traits --- src/scimlfunctions.jl | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/scimlfunctions.jl b/src/scimlfunctions.jl index c18868c8f..c1fde4650 100644 --- a/src/scimlfunctions.jl +++ b/src/scimlfunctions.jl @@ -4601,6 +4601,20 @@ has_Wfact_t(f::JacobianWrapper) = has_Wfact_t(f.f) has_paramjac(f::JacobianWrapper) = has_paramjac(f.f) has_colorvec(f::JacobianWrapper) = has_colorvec(f.f) +is_split_function(x) = is_split_function(typeof(x)) +is_split_function(::Type) = false +function is_split_function(::Type{T}) where {T <: Union{ + SplitFunction, SplitSDEFunction, DynamicalODEFunction, + DynamicalDDEFunction, DynamicalSDEFunction}} + true +end + +split_function_f_wrapper(::Type{<:SplitFunction}) = ODEFunction +split_function_f_wrapper(::Type{<:SplitSDEFunction}) = SDEFunction +split_function_f_wrapper(::Type{<:DynamicalODEFunction}) = ODEFunction +split_function_f_wrapper(::Type{<:DynamicalDDEFunction}) = DDEFunction +split_function_f_wrapper(::Type{<:DynamicalSDEFunction}) = DDEFunction + ######### Additional traits islinear(::AbstractDiffEqFunction) = false From d04a6eb57ed315a9fe9c2580d9a34f899425a820 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 12 Dec 2024 14:14:28 +0530 Subject: [PATCH 09/13] fix: better handle split functions in `remake(::AbstractSciMLFunction)` --- src/remake.jl | 42 +++++++++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index 412e5a011..46cbc7765 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -122,19 +122,23 @@ end by keyword arguments. For stochastic functions (e.g. `SDEFunction`) the `g` keyword argument is used to override `func.g`. For split functions (e.g. `SplitFunction`) the `f2` keyword argument is used to override `func.f2`, and `f` is used for `func.f1`. If -`f isa AbstractSciMLFunction`, properties of `f` will override those of `func` (but not ones -provided via keyword arguments). Properties of `f` that are `nothing` will fall back to those -in `func` (unless provided via keyword arguments). If `f` is a different type of -`AbstractSciMLFunction` from `func`, the returned function will be of the kind of `f`. +`f isa AbstractSciMLFunction` and `func` is not a split function, properties of `f` will +override those of `func` (but not ones provided via keyword arguments). Properties of `f` that +are `nothing` will fall back to those in `func` (unless provided via keyword arguments). If +`f` is a different type of `AbstractSciMLFunction` from `func`, the returned function will be +of the kind of `f` unless `func` is a split function. If `func` is a split function, `f` and +`f2` will be wrapped in the appropriate `AbstractSciMLFunction` type with the same `isinplace` +and `specialization` as `func`. """ -function remake(func::AbstractSciMLFunction; f = missing, g = missing, f2 = missing, kwargs...) +function remake( + func::AbstractSciMLFunction; f = missing, g = missing, f2 = missing, kwargs...) # retain iip and spec of original function iip = isinplace(func) spec = specialization(func) # retain properties of original function props = getproperties(func) - if f === missing + if f === missing || is_split_function(func) # if no `f` is provided, create the same type of SciMLFunction T = parameterless_type(func) f = isdefined(func, :f) ? func.f : func.f1 @@ -153,32 +157,36 @@ function remake(func::AbstractSciMLFunction; f = missing, g = missing, f2 = miss # minor hack to avoid breaking MTK, since prior to ~9.57 in `remake_initialization_data` # it creates a `NonlinearFunction` inside a `NonlinearFunction`. Just recursively unwrap # in this case and forget about properties. - while f isa AbstractSciMLFunction + while !is_split_function(T) && f isa AbstractSciMLFunction f = isdefined(f, :f) ? f.f : f.f1 end props = @delete props.f props = @delete props.f1 - if isdefined(func, :g) - # For SDEs/SDDEs where `g` is not a keyword - g = coalesce(g, func.g) + if is_split_function(T) + f = f isa AbstractSciMLOperator ? f : split_function_f_wrapper(T){iip, spec}(f) + end - props = @delete props.g - T{iip, spec}(f, g; props..., kwargs...) - elseif isdefined(func, :f2) + args = (f,) + if is_split_function(T) # For SplitFunction # we don't do the same thing as `g`, because for SDEs `g` is # stored in the problem as well, whereas for Split ODEs etc # f2 is a part of the function. Thus, if the user provides # a SciMLFunction for `f` which contains `f2` we use that. f2 = coalesce(f2, get(props, :f2, missing), func.f2) - + f2 = split_function_f_wrapper(T){iip, spec}(f2) props = @delete props.f2 - T{iip, spec}(f, f2; props..., kwargs...) - else - T{iip, spec}(f; props..., kwargs...) + args = (args..., f2) + end + if isdefined(func, :g) + # For SDEs/SDDEs where `g` is not a keyword + g = coalesce(g, func.g) + props = @delete props.g + args = (args..., g) end + T{iip, spec}(args...; props..., kwargs...) end """ From 8b490b8301b7db7a4c80222f2145e18459436194 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 12 Dec 2024 14:16:16 +0530 Subject: [PATCH 10/13] refactor: format --- src/initialization.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/initialization.jl b/src/initialization.jl index ddf9a96d0..50dd35f44 100644 --- a/src/initialization.jl +++ b/src/initialization.jl @@ -124,11 +124,13 @@ function evaluate_f( return _evaluate_f(integrator, f, isinplace, integrator.du, u, p, t) end -function evaluate_f(integrator::AbstractDDEIntegrator, prob::AbstractDDEProblem, f, isinplace, u, p, t) +function evaluate_f( + integrator::AbstractDDEIntegrator, prob::AbstractDDEProblem, f, isinplace, u, p, t) return _evaluate_f(integrator, f, isinplace, u, get_history_function(integrator), p, t) end -function evaluate_f(integrator::AbstractSDDEIntegrator, prob::AbstractSDDEProblem, f, isinplace, u, p, t) +function evaluate_f(integrator::AbstractSDDEIntegrator, + prob::AbstractSDDEProblem, f, isinplace, u, p, t) return _evaluate_f(integrator, f, isinplace, u, get_history_function(integrator), p, t) end From 8e32e90e6b02779033e4559b000f335bcef6fdcc Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Fri, 13 Dec 2024 13:03:34 +0530 Subject: [PATCH 11/13] fix: handle `cache` kwarg of `SplitFunction`/`DynamicalSDEFunction` in `remake` --- src/remake.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/remake.jl b/src/remake.jl index 46cbc7765..76876d514 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -164,6 +164,9 @@ function remake( props = @delete props.f props = @delete props.f1 + # for DynamicalSDEFunction and SplitFunction + props = @delete props.cache + if is_split_function(T) f = f isa AbstractSciMLOperator ? f : split_function_f_wrapper(T){iip, spec}(f) end From 6d4374c460f8429df9a9c66ee0248939b77eeab9 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Fri, 13 Dec 2024 14:46:32 +0530 Subject: [PATCH 12/13] fix: fix undefined variable in `remake(::ODEProblem)` --- src/remake.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/remake.jl b/src/remake.jl index 76876d514..cf985ae86 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -240,16 +240,14 @@ function remake(prob::ODEProblem; f = missing, f = remake( f; f = wrapfun_oop(unwrapped_f(f.f), (newu0, newu0, newp, ptspan[1]))) end - else - _f = f end prob = if kwargs === missing ODEProblem{iip}( - _f, newu0, tspan, newp, prob.problem_type; prob.kwargs..., + f, newu0, tspan, newp, prob.problem_type; prob.kwargs..., _kwargs...) else - ODEProblem{iip}(_f, newu0, tspan, newp, prob.problem_type; kwargs...) + ODEProblem{iip}(f, newu0, tspan, newp, prob.problem_type; kwargs...) end if lazy_initialization === nothing From d455a5a6077429f05e49aed3d8bce66fd315bc87 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Fri, 13 Dec 2024 14:47:06 +0530 Subject: [PATCH 13/13] fix: retain cache in appropriate split functions --- src/SciMLBase.jl | 2 +- src/remake.jl | 22 ++++++++++++++-------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/SciMLBase.jl b/src/SciMLBase.jl index 2e30d27a5..30a6c0b57 100644 --- a/src/SciMLBase.jl +++ b/src/SciMLBase.jl @@ -22,7 +22,7 @@ import FunctionWrappersWrappers import RuntimeGeneratedFunctions import EnumX import ADTypes: ADTypes, AbstractADType -import Accessors: @set, @reset, @delete +import Accessors: @set, @reset, @delete, @insert using Expronicon.ADT: @match using Reexport diff --git a/src/remake.jl b/src/remake.jl index cf985ae86..31d013964 100644 --- a/src/remake.jl +++ b/src/remake.jl @@ -164,22 +164,28 @@ function remake( props = @delete props.f props = @delete props.f1 - # for DynamicalSDEFunction and SplitFunction - props = @delete props.cache - - if is_split_function(T) - f = f isa AbstractSciMLOperator ? f : split_function_f_wrapper(T){iip, spec}(f) - end - args = (f,) if is_split_function(T) + # for DynamicalSDEFunction and SplitFunction + if isdefined(props, :cache) + props = @insert props._func_cache = props.cache + props = @delete props.cache + end + + # `f1` and `f2` are wrapped in another SciMLFunction, unless they're + # already wrapped in the appropriate type or are an `AbstractSciMLOperator` + if !(f isa Union{AbstractSciMLOperator, split_function_f_wrapper(T)}) + f = split_function_f_wrapper(T){iip, spec}(f) + end # For SplitFunction # we don't do the same thing as `g`, because for SDEs `g` is # stored in the problem as well, whereas for Split ODEs etc # f2 is a part of the function. Thus, if the user provides # a SciMLFunction for `f` which contains `f2` we use that. f2 = coalesce(f2, get(props, :f2, missing), func.f2) - f2 = split_function_f_wrapper(T){iip, spec}(f2) + if !(f2 isa Union{AbstractSciMLOperator, split_function_f_wrapper(T)}) + f2 = split_function_f_wrapper(T){iip, spec}(f2) + end props = @delete props.f2 args = (args..., f2) end