Merge pull request #477 from dpad/allow-gg-ODEfunctions

ChrisRackauckas · web-flow · commit 388eb290bd04 · 2020-06-24T09:04:00.000-04:00
Provide as_expression=false option to ODEs/SDEs, for specific use cases that need to avoid world age issues (generally not needed)
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -142,33 +142,36 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      ps = parameters(sys), u0 = nothing;
                                      version = nothing, tgrad=false,
                                      jac = false, Wfact = false,
-                                     sparse = false,
+                                     sparse = false, eval_expression = true,
                                      kwargs...) where {iip}
 
-    f_oop,f_iip = ModelingToolkit.eval.(generate_function(sys, dvs, ps; expression=Val{true}, kwargs...))
+    f_gen = generate_function(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+    f_oop,f_iip = eval_expression ? ModelingToolkit.eval.(f_gen) : f_gen
     f(u,p,t) = f_oop(u,p,t)
     f(du,u,p,t) = f_iip(du,u,p,t)
 
     if tgrad
-        tgrad_oop,tgrad_iip = ModelingToolkit.eval.(generate_tgrad(sys, dvs, ps; expression=Val{true}, kwargs...))
+        tgrad_gen = generate_tgrad(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+        tgrad_oop,tgrad_iip = eval_expression ? ModelingToolkit.eval.(tgrad_gen) : tgrad_gen
         _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 = ModelingToolkit.eval.(generate_jacobian(sys, dvs, ps; sparse = sparse, expression=Val{true}, kwargs...))
+        jac_gen = generate_jacobian(sys, dvs, ps; sparse = sparse, expression=Val{eval_expression}, kwargs...)
+        jac_oop,jac_iip = eval_expression ? ModelingToolkit.eval.(jac_gen) : jac_gen
         _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; expression=Val{true}, kwargs...)
-        Wfact_oop, Wfact_iip =  ModelingToolkit.eval.(tmp_Wfact)
-        Wfact_oop_t, Wfact_iip_t =  ModelingToolkit.eval.(tmp_Wfact_t)
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+        Wfact_oop, Wfact_iip =  eval_expression ? ModelingToolkit.eval.(tmp_Wfact) : tmp_Wfact
+        Wfact_oop_t, Wfact_iip_t =  eval_expression ? ModelingToolkit.eval.(tmp_Wfact_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)
@@ -288,14 +291,15 @@ function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     jac = false, Wfact = false,
                                     checkbounds = false, sparse = false,
                                     linenumbers = true, parallel=SerialForm(),
+                                    eval_expression = true,
                                     kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
     u0 = varmap_to_vars(u0map,dvs)
     p = varmap_to_vars(parammap,ps)
     f = ODEFunction{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
                         linenumbers=linenumbers,parallel=parallel,
-                        sparse=sparse)
+                        sparse=sparse,eval_expression=eval_expression)
     ODEProblem{iip}(f,u0,tspan,p;kwargs...)
 end
 
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
@@ -101,25 +101,29 @@ respectively.
 function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.states, ps = sys.ps,
                                      u0 = nothing;
                                      version = nothing, tgrad=false, sparse = false,
-                                     jac = false, Wfact = false, kwargs...) where {iip}
-    f_oop,f_iip = ModelingToolkit.eval.(generate_function(sys, dvs, ps; expression=Val{true}, kwargs...))
-    g_oop,g_iip = ModelingToolkit.eval.(generate_diffusion_function(sys, dvs, ps; expression=Val{true}, kwargs...))
+                                     jac = false, Wfact = false, eval_expression = true, kwargs...) where {iip}
+    f_gen = generate_function(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+    f_oop,f_iip = eval_expression ? ModelingToolkit.eval.(f_gen) : f_gen
+    g_gen = generate_diffusion_function(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+    g_oop,g_iip = eval_expression ? ModelingToolkit.eval.(g_gen) : g_gen
 
     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 = ModelingToolkit.eval.(generate_tgrad(sys, dvs, ps; expression=Val{true}, kwargs...))
+        tgrad_gen = generate_tgrad(sys, dvs, ps; expression=Val{eval_expression}, kwargs...)
+        tgrad_oop,tgrad_iip = eval_expression ? ModelingToolkit.eval.(tgrad_gen) : tgrad_gen
         _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 = ModelingToolkit.eval.(generate_jacobian(sys, dvs, ps; expression=Val{true}, sparse=sparse, kwargs...))
+        jac_gen = generate_jacobian(sys, dvs, ps; expression=Val{eval_expression}, sparse=sparse, kwargs...)
+        jac_oop,jac_iip = eval_expression ? ModelingToolkit.eval.(jac_gen) : jac_gen
         _jac(u,p,t) = jac_oop(u,p,t)
         _jac(J,u,p,t) = jac_iip(J,u,p,t)
     else
@@ -128,8 +132,8 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.states, ps = sys.
 
     if Wfact
         tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true; expression=Val{true}, kwargs...)
-        Wfact_oop, Wfact_iip =  ModelingToolkit.eval.(tmp_Wfact)
-        Wfact_oop_t, Wfact_iip_t =  ModelingToolkit.eval.(tmp_Wfact_t)
+        Wfact_oop, Wfact_iip = eval_expression ? ModelingToolkit.eval.(tmp_Wfact) : tmp_Wfact
+        Wfact_oop_t, Wfact_iip_t = eval_expression ? ModelingToolkit.eval.(tmp_Wfact_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)
@@ -253,6 +257,7 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,parammap=DiffEqBa
                                     checkbounds = false, sparse = false,
                                     sparsenoise = sparse,
                                     linenumbers = true, parallel=SerialForm(),
+                                    eval_expression = true,
                                     kwargs...) where iip
 
     dvs = states(sys)
@@ -262,7 +267,7 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,parammap=DiffEqBa
     f = SDEFunction{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,Wfact=Wfact,
                          checkbounds=checkbounds,
                          linenumbers=linenumbers,parallel=parallel,
-                         sparse=sparse)
+                         sparse=sparse, eval_expression=eval_expression)
     if typeof(sys.noiseeqs) <: AbstractVector
         noise_rate_prototype = nothing
     elseif sparsenoise
diff --git a/test/function_registration.jl b/test/function_registration.jl
@@ -3,6 +3,7 @@
 # appropriately calls the registered functions, whether the call is
 # qualified (with a module name) or not.
 
+
 # TEST: Function registration in a module.
 # ------------------------------------------------
 module MyModule
@@ -20,9 +21,11 @@ module MyModule
     sys = ODESystem([eq], t, [u], [x])
     fun = ODEFunction(sys)
 
-    @test fun([0.5], [5.0], 0.) == [30.0]
+    u0 = 5.0
+    @test fun([0.5], [u0], 0.) == [do_something(u0) * 2]
 end
 
+
 # TEST: Function registration in a nested module.
 # ------------------------------------------------
 module MyModule2
@@ -41,10 +44,12 @@ module MyModule2
         sys = ODESystem([eq], t, [u], [x])
         fun = ODEFunction(sys)
 
-        @test fun([0.5], [3.0], 0.) == [46.0]
+        u0 = 3.0
+        @test fun([0.5], [u0], 0.) == [do_something_2(u0) * 2]
     end
 end
 
+
 # TEST: Function registration outside any modules.
 # ------------------------------------------------
 using ModelingToolkit, DiffEqBase, LinearAlgebra, Test
@@ -61,9 +66,12 @@ eq  = Dt(u) ~ do_something_3(x) + (@__MODULE__).do_something_3(x)
 sys = ODESystem([eq], t, [u], [x])
 fun = ODEFunction(sys)
 
-@test fun([0.5], [7.0], 0.) == [74.0]
+u0 = 7.0
+@test fun([0.5], [u0], 0.) == [do_something_3(u0) * 2]
+
 
-# derivative
+# TEST: Function registration works with derivatives.
+# ---------------------------------------------------
 foo(x, y) = sin(x) * cos(y)
 @parameters t; @variables x(t) y(t) z(t); @derivatives D'~t;
 @register foo(x, y)
@@ -74,3 +82,29 @@ expr = foo(x, y)
 ModelingToolkit.derivative(::typeof(foo), (x, y), ::Val{1}) = cos(x) * cos(y) # derivative w.r.t. the first argument
 ModelingToolkit.derivative(::typeof(foo), (x, y), ::Val{2}) = -sin(x) * sin(y) # derivative w.r.t. the second argument
 @test isequal(expand_derivatives(D(foo(x, y))), expand_derivatives(D(sin(x) * cos(y))))
+
+
+# TEST: Function registration run from inside a function.
+# -------------------------------------------------------
+# This tests that we can get around the world age issue by falling back to
+# GeneralizedGenerated instead of function expressions.
+# Might be useful in cases where someone wants to define functions that build
+# up and use ODEFunctions given some parameters.
+function do_something_4(a)
+    a + 30
+end
+@register do_something_4(a)
+function build_ode()
+    @parameters t x
+    @variables u(t)
+    @derivatives Dt'~t
+    eq  = Dt(u) ~ do_something_4(x) + (@__MODULE__).do_something_4(x)
+    sys = ODESystem([eq], t, [u], [x])
+    fun = ODEFunction(sys, eval_expression=false)
+end
+function run_test()
+    fun = build_ode()
+    u0 = 10.0
+    @test fun([0.5], [u0], 0.) == [do_something_4(u0) * 2]
+end
+run_test()
diff --git a/test/sdesystem.jl b/test/sdesystem.jl
@@ -70,3 +70,11 @@ prob = SDEProblem(de,u0map,(0.0,100.0),parammap,sparsenoise=true)
 @test size(prob.noise_rate_prototype) == (3,3)
 @test prob.noise_rate_prototype isa SparseMatrixCSC
 sol = solve(prob,EM(),dt=0.001)
+
+# Test eval_expression=false
+function test_SDEFunction_no_eval()
+    # Need to test within a function scope to trigger world age issues
+    f = SDEFunction(de, eval_expression=false)
+    @test f([1.0,0.0,0.0], (10.0,26.0,2.33), (0.0,100.0)) ≈ [-10.0, 26.0, 0.0]
+end
+test_SDEFunction_no_eval()
