finish and test SDESystem

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -1,16 +1,9 @@
 module ModelingToolkit
 
-export Operation, Expression
-export calculate_jacobian, generate_jacobian, generate_function
-export independent_variables, dependent_variables, parameters
-export simplified_expr, eval_function
-export @register, @I
-export modelingtoolkitize
-
-
 using DiffEqBase, Distributed
 using StaticArrays, LinearAlgebra, SparseArrays
 using Latexify
+using MacroTools
 
 using MacroTools
 import MacroTools: splitdef, combinedef, postwalk, striplines
@@ -92,6 +85,7 @@ include("utils.jl")
 include("direct.jl")
 include("domains.jl")
 include("systems/diffeqs/odesystem.jl")
+include("systems/diffeqs/sdesystem.jl")
 include("systems/diffeqs/abstractodesystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
 include("systems/diffeqs/modelingtoolkitize.jl")
@@ -100,4 +94,22 @@ include("systems/pde/pdesystem.jl")
 include("latexify_recipes.jl")
 include("build_function.jl")
 
+export ODESystem, ODEFunction
+export SDESystem, SDEFunction
+export NonlinearSystem
+export ode_order_lowering
+export PDESystem
+export Differential, expand_derivatives, @derivatives
+export IntervalDomain, ProductDomain, ⊗, CircleDomain
+export Equation, ConstrainedEquation
+export simplify_constants
+
+export Operation, Expression
+export calculate_jacobian, generate_jacobian, generate_function, generate_diffusion_function
+export independent_variables, dependent_variables, parameters
+export simplified_expr, eval_function
+export @register, @I
+export modelingtoolkitize
+export Variable, @variables, @parameters
+
 end # module

src/differentials.jl

@@ -1,6 +1,3 @@
-export Differential, expand_derivatives, @derivatives
-
-
 """
 $(TYPEDEF)
 

src/domains.jl

@@ -1,5 +1,3 @@
-export IntervalDomain, ProductDomain, ⊗, CircleDomain
-
 abstract type AbstractDomain{T,N} end
 
 struct VarDomainPairing

src/equations.jl

@@ -1,6 +1,3 @@
-export Equation, ConstrainedEquation
-
-
 """
 $(TYPEDEF)
 

src/simplify.jl

@@ -1,6 +1,3 @@
-export simplify_constants
-
-
 function simplify_constants(O::Operation, shorten_tree)
     while true
         O′ = _simplify_constants(O, shorten_tree)

src/systems/diffeqs/abstractodesystem.jl

@@ -21,7 +21,7 @@ function calculate_jacobian(sys::AbstractODESystem)
 end
 
 struct ODEToExpr
-    sys::ODESystem
+    sys::AbstractODESystem
 end
 function (f::ODEToExpr)(O::Operation)
     if isa(O.op, Variable)

src/systems/diffeqs/first_order_transform.jl

@@ -1,6 +1,3 @@
-export ode_order_lowering
-
-
 function lower_varname(var::Variable, idv, order)
     order == 0 && return var
     name = Symbol(var.name, :_, string(idv.name)^order)

src/systems/diffeqs/odesystem.jl

@@ -1,5 +1,3 @@
-export ODESystem, ODEFunction
-
 isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
 
 function flatten_differential(O::Operation)

src/systems/diffeqs/sdesystem.jl

@@ -51,3 +51,60 @@ function generate_diffusion_function(sys::SDESystem, dvs = sys.dvs, ps = sys.ps,
     ps′ = [clean(p) for p ∈ ps]
     return build_function(sys.noiseeqs, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
 end
+
+"""
+$(SIGNATURES)
+
+Create an `SDEFunction` from the [`SDESystem`](@ref). The arguments `dvs` and `ps`
+are used to set the order of the dependent variable and parameter vectors,
+respectively.
+"""
+function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.dvs, ps = sys.ps;
+                                     version = nothing, tgrad=false,
+                                     jac = false, Wfact = false) where {iip}
+    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
+    g_oop,g_iip = generate_diffusion_function(sys, dvs, ps, Val{false})
+
+    f(u,p,t) = f_oop(u,p,t)
+    f(du,u,p,t) = f_iip(du,u,p,t)
+    g(u,p,t) = g_oop(u,p,t)
+    g(du,u,p,t) = g_iip(du,u,p,t)
+
+    if tgrad
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false})
+        _tgrad(u,p,t) = tgrad_oop(u,p,t)
+        _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
+    else
+        _tgrad = nothing
+    end
+
+    if jac
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
+        _jac(u,p,t) = jac_oop(u,p,t)
+        _jac(J,u,p,t) = jac_iip(J,u,p,t)
+    else
+        _jac = nothing
+    end
+
+    if Wfact
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false})
+        Wfact_oop, Wfact_iip = tmp_Wfact
+        Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
+        _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)
+        _Wfact(W,u,p,dtgamma,t) = Wfact_iip(W,u,p,dtgamma,t)
+        _Wfact_t(u,p,dtgamma,t) = Wfact_oop_t(u,p,dtgamma,t)
+        _Wfact_t(W,u,p,dtgamma,t) = Wfact_iip_t(W,u,p,dtgamma,t)
+    else
+        _Wfact,_Wfact_t = nothing,nothing
+    end
+
+    SDEFunction{iip}(f,g,jac=_jac,
+                      tgrad = _tgrad,
+                      Wfact = _Wfact,
+                      Wfact_t = _Wfact_t,
+                      syms = Symbol.(sys.dvs))
+end
+
+function DiffEqBase.SDEFunction(sys::SDESystem, args...; kwargs...)
+    SDEFunction{true}(sys, args...; kwargs...)
+end

src/systems/nonlinear/nonlinear_system.jl

@@ -1,5 +1,3 @@
-export NonlinearSystem
-
 struct NLEq
     rhs::Expression
 end