From 4d17fce563b34d84d23930a75ba6e768e0015ba7 Mon Sep 17 00:00:00 2001 From: James Gardner Date: Fri, 5 Feb 2021 15:08:22 +0000 Subject: [PATCH] Add DynamicalSDEFunction and Problem --- src/SciMLBase.jl | 1 + src/problems/sde_problems.jl | 32 ++++++++++++++ src/scimlfunctions.jl | 84 ++++++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+) diff --git a/src/SciMLBase.jl b/src/SciMLBase.jl index 536547443..dea670d5b 100644 --- a/src/SciMLBase.jl +++ b/src/SciMLBase.jl @@ -543,6 +543,7 @@ export ODEProblem, ODESolution export DynamicalODEFunction, DynamicalODEProblem, SecondOrderODEProblem, SplitFunction, SplitODEProblem export SplitSDEProblem +export DynamicalSDEFunction, DynamicalSDEProblem export RODEProblem, RODESolution, SDEProblem export DAEProblem, DAESolution export DDEProblem diff --git a/src/problems/sde_problems.jl b/src/problems/sde_problems.jl index 59e769bca..91bed0c64 100644 --- a/src/problems/sde_problems.jl +++ b/src/problems/sde_problems.jl @@ -83,3 +83,35 @@ function SplitSDEProblem{iip}(f::SplitSDEFunction,g,u0,tspan,p=NullParameters(); end SDEProblem(_f,g,u0,tspan,p;kwargs...) end + +""" +$(TYPEDEF) +""" +abstract type AbstractDynamicalSDEProblem end + +""" +$(TYPEDEF) +""" +struct DynamicalSDEProblem{iip} <: AbstractDynamicalSDEProblem end + +function DynamicalSDEProblem(f1,f2,g,v0,u0,tspan,p=NullParameters();kwargs...) + DynamicalSDEProblem(DynamicalSDEFunction(f1,f2,g),g,v0,u0,tspan,p;kwargs...) +end + +DynamicalSDEProblem(f::DynamicalSDEFunction,g,v0,u0,tspan,p=NullParameters();kwargs...) = + DynamicalSDEProblem{isinplace(f)}(f,g,u0,v0,tspan,p;kwargs...) + +function DynamicalSDEProblem{iip}(f1,f2,g,v0,u0,tspan,p=NullParameters();kwargs...) where iip + DynamicalSDEProblem(DynamicalSDEFunction(f1,f2,g),g,v0,u0,tspan,p;kwargs...) +end +function DynamicalSDEProblem{iip}(f::DynamicalSDEFunction,g,v0,u0,tspan,p=NullParameters(); + func_cache=nothing,kwargs...) where iip + if f.cache === nothing && iip + cache = similar(u0) + _f = DynamicalSDEFunction{iip}(f.f1, f.f2, f.g; mass_matrix=f.mass_matrix, + _func_cache=cache, analytic=f.analytic) + else + _f = f + end + SDEProblem(_f,g,ArrayPartition(v0,u0),tspan,p;kwargs...) +end \ No newline at end of file diff --git a/src/scimlfunctions.jl b/src/scimlfunctions.jl index 43909632a..4d518ef7e 100644 --- a/src/scimlfunctions.jl +++ b/src/scimlfunctions.jl @@ -194,6 +194,30 @@ struct SplitSDEFunction{iip,F1,F2,G,TMM,C,Ta,Tt,TJ,JVP,VJP,JP,SP,TW,TWt,TPJ,S,TC colorvec::TCV end +""" +$(TYPEDEF) +""" +struct DynamicalSDEFunction{iip,F1,F2,G,TMM,C,Ta,Tt,TJ,JVP,VJP,JP,SP,TW,TWt,TPJ,S,TCV} <: AbstractSDEFunction{iip} + # This is a direct copy of the SplitSDEFunction, maybe it's not necessary and the above can be used instead. + f1::F1 + f2::F2 + g::G + mass_matrix::TMM + cache::C + analytic::Ta + tgrad::Tt + jac::TJ + jvp::JVP + vjp::VJP + jac_prototype::JP + sparsity::SP + Wfact::TW + Wfact_t::TWt + paramjac::TPJ + syms::S + colorvec::TCV +end + """ $(TYPEDEF) """ @@ -745,6 +769,66 @@ SplitSDEFunction{iip}(f1,f2, g; kwargs...) where iip = SplitSDEFunction{iip,RECOMPILE_BY_DEFAULT}(SDEFunction(f1,g), SDEFunction{iip}(f2,g), g; kwargs...) SplitSDEFunction(f::SplitSDEFunction; kwargs...) = f +@add_kwonly function DynamicalSDEFunction(f1,f2,g,mass_matrix,cache,analytic,tgrad,jac,jvp,vjp, + jac_prototype,Wfact,Wfact_t,paramjac, + syms,colorvec) + f1 = typeof(f1) <: AbstractDiffEqOperator ? f1 : SDEFunction(f1) + f2 = SDEFunction(f2) + DynamicalSDEFunction{isinplace(f2),typeof(f1),typeof(f2),typeof(g),typeof(mass_matrix), + typeof(cache),typeof(analytic),typeof(tgrad),typeof(jac),typeof(jvp),typeof(vjp), + typeof(Wfact),typeof(Wfact_t),typeof(paramjac),typeof(syms), + typeof(colorvec)}(f1,f2,g,mass_matrix,cache,analytic,tgrad,jac, + jac_prototype,Wfact,Wfact_t,paramjac,syms,colorvec) +end + +function DynamicalSDEFunction{iip,true}(f1,f2,g; mass_matrix=I, + _func_cache=nothing,analytic=nothing, + tgrad = nothing, + jac = nothing, + jac_prototype=nothing, + sparsity=jac_prototype, + jvp=nothing, + vjp=nothing, + Wfact = nothing, + Wfact_t = nothing, + paramjac = nothing, + syms = nothing, + colorvec = nothing) where iip + DynamicalSDEFunction{iip,typeof(f1),typeof(f2),typeof(g), + typeof(mass_matrix),typeof(_func_cache), + typeof(analytic), + typeof(tgrad),typeof(jac),typeof(jvp),typeof(vjp),typeof(jac_prototype),typeof(sparsity), + typeof(Wfact),typeof(Wfact_t),typeof(paramjac),typeof(syms), + typeof(colorvec)}(f1,f2,g,mass_matrix,_func_cache,analytic, + tgrad,jac,jvp,vjp,jac_prototype,sparsity, + Wfact,Wfact_t,paramjac,syms,colorvec) +end +function DynamicalSDEFunction{iip,false}(f1,f2,g; mass_matrix=I, + _func_cache=nothing,analytic=nothing, + tgrad = nothing, + jac = nothing, + jvp=nothing, + vjp=nothing, + jac_prototype=nothing, + sparsity=jac_prototype, + Wfact = nothing, + Wfact_t = nothing, + paramjac = nothing, + syms = nothing, + colorvec = nothing) where iip + DynamicalSDEFunction{iip,Any,Any,Any,Any,Any, + Any,Any,Any,Any, + Any,Any,Any,Any,Any}( + f1,f2,g,mass_matrix,_func_cache,analytic, + tgrad,jac,jvp,vjp,jac_prototype,sparsity, + Wfact,Wfact_t,paramjac,syms,colorvec) +end +# Here I changed `isinplace(f2, 4) -> isinplace(f2, 5)` to allow for extra arguments for dynamical functions. +DynamicalSDEFunction(f1,f2,g; kwargs...) = DynamicalSDEFunction{isinplace(f2, 5)}(f1, f2, g; kwargs...) +DynamicalSDEFunction{iip}(f1,f2, g; kwargs...) where iip = +DynamicalSDEFunction{iip,RECOMPILE_BY_DEFAULT}(SDEFunction(f1,g), SDEFunction{iip}(f2,g), g; kwargs...) +DynamicalSDEFunction(f::DynamicalSDEFunction; kwargs...) = f + function RODEFunction{iip,true}(f; mass_matrix=I, analytic=nothing,