diff --git a/src/functionwrapper.jl b/src/functionwrapper.jl index efaf14f..6ded6f0 100644 --- a/src/functionwrapper.jl +++ b/src/functionwrapper.jl @@ -9,7 +9,7 @@ end (g::SDEDiffusionTermWrapper{false})(u, p, t) = g.g(u, g.h, p, t) struct SDEFunctionWrapper{iip, F, G, H, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, TWt, TPJ, GG, - TCV} <: DiffEqBase.AbstractRODEFunction{iip} + TCV, ID, S} <: DiffEqBase.AbstractRODEFunction{iip} f::F g::G h::H @@ -26,6 +26,8 @@ struct SDEFunctionWrapper{iip, F, G, H, TMM, Ta, Tt, TJ, JVP, VJP, JP, SP, TW, T paramjac::TPJ ggprime::GG colorvec::TCV + initialization_data::ID + sys::S end (f::SDEFunctionWrapper{true})(du, u, p, t) = f.f(du, u, f.h, p, t) @@ -53,17 +55,9 @@ function wrap_functions_and_history(f::SDDEFunction, g, h) typeof(f.analytic), typeof(f.tgrad), typeof(jac), typeof(f.jvp), typeof(f.vjp), typeof(f.jac_prototype), typeof(f.sparsity), typeof(f.Wfact), typeof(f.Wfact_t), typeof(f.paramjac), - typeof(f.ggprime), typeof(f.colorvec)}(f.f, gwh, h, - f.mass_matrix, - f.analytic, - f.tgrad, jac, - f.jvp, f.vjp, - f.jac_prototype, - f.sparsity, - f.Wfact, - f.Wfact_t, - f.paramjac, - f.ggprime, - f.colorvec), + typeof(f.ggprime), typeof(f.colorvec), typeof(f.initialization_data), + typeof(f.sys)}(f.f, gwh, h, f.mass_matrix, f.analytic, f.tgrad, jac, + f.jvp, f.vjp, f.jac_prototype, f.sparsity, f.Wfact, f.Wfact_t, f.paramjac, + f.ggprime, f.colorvec, f.initialization_data, f.sys), gwh end