sde build_function

Author: shashi

src/build_function.jl

@@ -224,6 +224,7 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
                          skipzeros = outputidxs===nothing,
 						 fillzeros = skipzeros && !(typeof(rhss)<:SparseMatrixCSC),
 						 parallel=SerialForm(), kwargs...)
+    conv = conv ∘ rm_calls_with_iv
 	if multithread isa Bool
 		@warn("multithraded is deprecated for the parallel argument. See the documentation.")
 		parallel = multithread ? MultithreadedForm() : SerialForm()
@@ -466,6 +467,10 @@ function vars_to_pairs(name,vs)
     [term_to_symbol(value(vs))], [name]
 end
 
+function rm_calls_with_iv(expr)
+    Rewriters.Prewalk(Rewriters.Chain([@rule((~f::(x->x isa Sym))(~t::(x->x isa Sym)) => Sym{symtype((~f)(~t))}((term_to_symbol(~f))))]))(value(expr))
+end
+
 get_varnumber(varop::Operation,vars::Vector{Operation}) =  findfirst(x->isequal(x,varop),vars)
 get_varnumber(varop::Operation,vars::Vector{<:Variable})  =  findfirst(x->isequal(x,varop.op),vars)
 
@@ -486,7 +491,7 @@ end
 
 function numbered_expr(de::ModelingToolkit.Equation,args...;varordering = args[1],
                        lhsname=gensym("du"),rhsnames=[gensym("MTK") for i in 1:length(args)],offset=0)
-    i = findfirst(x->isequal(x isa Variable ? term_to_symbol(x) : term_to_symbol(x.op),term_to_symbol(var_from_nested_derivative(de.lhs)[1])),varordering)
+    i = findfirst(x->isequal(x isa Sym ? term_to_symbol(x) : term_to_symbol(x.op),term_to_symbol(var_from_nested_derivative(de.lhs)[1])),varordering)
     :($lhsname[$(i+offset)] = $(numbered_expr(de.rhs,args...;offset=offset,
 											  varordering = varordering,
 											  lhsname = lhsname,

src/systems/abstractsystem.jl

@@ -272,3 +272,4 @@ function (f::AbstractSysToExpr)(O::Operation)
     return build_expr(:call, Any[O.op; f.(O.args)])
 end
 (f::AbstractSysToExpr)(x) = toexpr(x)
+

src/systems/diffeqs/modelingtoolkitize.jl

@@ -39,15 +39,17 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem)
     prob.f isa DiffEqBase.AbstractParameterizedFunction &&
                             return (prob.f.sys, prob.f.sys.states, prob.f.sys.ps)
     @parameters t
-    vars = reshape([Variable(:x, i)(t) for i in eachindex(prob.u0)],size(prob.u0))
+    var(x, i) = Sym{FnType{Tuple{symtype(t)}, Number}}(nameof(Variable(:x, i)))
+    vars = reshape([var(:x, i)(t) for i in eachindex(prob.u0)],size(prob.u0))
     params = prob.p isa DiffEqBase.NullParameters ? [] :
-             reshape([Variable(:α,i)() for i in eachindex(prob.p)],size(prob.p))
+             reshape([Variable(:α,i) for i in eachindex(prob.p)],size(prob.p))
     @derivatives D'~t
 
     rhs = [D(var) for var in vars]
 
     if DiffEqBase.isinplace(prob)
         lhs = similar(vars, Any)
+
         prob.f(lhs, vars, params, t)
 
         if DiffEqBase.is_diagonal_noise(prob)

src/systems/diffeqs/odesystem.jl

@@ -81,8 +81,8 @@ end
 var_from_nested_derivative(x, i=0) = (missing, missing)
 var_from_nested_derivative(x::Term,i=0) = x.op isa Differential ? var_from_nested_derivative(x.args[1],i+1) : (x,i)
 
-iv_from_nested_derivative(x) = x.op isa Differential ? iv_from_nested_derivative(x.args[1]) : x.args[1]
-iv_from_nested_derivative(x::Constant) = missing
+iv_from_nested_derivative(x::Term) = x.op isa Differential ? iv_from_nested_derivative(x.args[1]) : x.args[1]
+iv_from_nested_derivative(x) = missing
 
 vars(exprs::Term) = vars([exprs])
 vars(exprs) = foldl(vars!, exprs; init = Set())

src/systems/diffeqs/sdesystem.jl

@@ -30,35 +30,35 @@ struct SDESystem <: AbstractODESystem
     """The expressions defining the drift term."""
     eqs::Vector{Equation}
     """The expressions defining the diffusion term."""
-    noiseeqs::AbstractArray{Operation}
+    noiseeqs::AbstractArray
     """Independent variable."""
-    iv::Variable
+    iv::Sym
     """Dependent (state) variables."""
-    states::Vector{Variable}
+    states::Vector
     """Parameter variables."""
-    ps::Vector{Variable}
-    pins::Vector{Variable}
-    observed::Vector{Equation}
+    ps::Vector
+    pins::Vector
+    observed::Vector
     """
     Time-derivative matrix. Note: this field will not be defined until
     [`calculate_tgrad`](@ref) is called on the system.
     """
-    tgrad::RefValue{Vector{Expression}}
+    tgrad::RefValue
     """
     Jacobian matrix. Note: this field will not be defined until
     [`calculate_jacobian`](@ref) is called on the system.
     """
-    jac::RefValue{Matrix{Expression}}
+    jac::RefValue
     """
     `Wfact` matrix. Note: this field will not be defined until
     [`generate_factorized_W`](@ref) is called on the system.
     """
-    Wfact::RefValue{Matrix{Expression}}
+    Wfact::RefValue
     """
     `Wfact_t` matrix. Note: this field will not be defined until
     [`generate_factorized_W`](@ref) is called on the system.
     """
-    Wfact_t::RefValue{Matrix{Expression}}
+    Wfact_t::RefValue
     """
     Name: the name of the system
     """
@@ -70,24 +70,22 @@ struct SDESystem <: AbstractODESystem
 end
 
 function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
-                   pins = Variable[],
-                   observed = Operation[],
+                   pins = [],
+                   observed = [],
                    systems = SDESystem[],
                    name = gensym(:SDESystem))
-    iv′ = convert(Variable,iv)
-    dvs′ = convert.(Variable,dvs)
-    ps′ = convert.(Variable,ps)
-    tgrad = RefValue(Vector{Expression}(undef, 0))
-    jac = RefValue(Matrix{Expression}(undef, 0, 0))
-    Wfact   = RefValue(Matrix{Expression}(undef, 0, 0))
-    Wfact_t = RefValue(Matrix{Expression}(undef, 0, 0))
+    iv′ = value(iv)
+    dvs′ = value.(dvs)
+    ps′ = value.(ps)
+    tgrad = RefValue(Vector{Num}(undef, 0))
+    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
+    Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
+    Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
     SDESystem(deqs, neqs, iv′, dvs′, ps′, pins, observed, tgrad, jac, Wfact, Wfact_t, name, systems)
 end
 
 function generate_diffusion_function(sys::SDESystem, dvs = sys.states, ps = sys.ps; kwargs...)
-    dvs′ = convert.(Variable,dvs)
-    ps′ = convert.(Variable,ps)
-    return build_function(sys.noiseeqs, dvs′, ps′, sys.iv;
+    return build_function(sys.noiseeqs, dvs, ps, sys.iv;
                           conv = ODEToExpr(sys),kwargs...)
 end