Provide eval_expression=false option in ODEFunction, for very specific use cases that need to avoid world age issues arising from evaluating expressions.

Author: dpad

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)

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()