Merge pull request #1758 from SciML/myb/scimlfun

Forward `sys` to `SciMLFunction`s

Author: YingboMa

Project.toml

@@ -68,7 +68,7 @@ NonlinearSolve = "0.3.8"
 RecursiveArrayTools = "2.3"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.4.3, 0.5"
-SciMLBase = "1.26.2"
+SciMLBase = "1.49"
 Setfield = "0.7, 0.8, 1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"

src/systems/diffeqs/abstractodesystem.jl

@@ -339,6 +339,7 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
         nothing
     end
     ODEFunction{iip}(f,
+                     sys = sys,
                      jac = _jac === nothing ? nothing : _jac,
                      tgrad = _tgrad === nothing ? nothing : _tgrad,
                      mass_matrix = _M,
@@ -427,6 +428,7 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
     end
 
     DAEFunction{iip}(f,
+                     sys = sys,
                      jac = _jac === nothing ? nothing : _jac,
                      syms = Symbol.(dvs),
                      jac_prototype = jac_prototype,

src/systems/diffeqs/sdesystem.jl

@@ -406,6 +406,7 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = states(sys),
 
     sts = states(sys)
     SDEFunction{iip}(f, g,
+                     sys = sys,
                      jac = _jac === nothing ? nothing : _jac,
                      tgrad = _tgrad === nothing ? nothing : _tgrad,
                      Wfact = _Wfact === nothing ? nothing : _Wfact,

src/systems/discrete_system/discrete_system.jl

@@ -198,7 +198,7 @@ function SciMLBase.DiscreteProblem(sys::DiscreteSystem, u0map, tspan,
                               expression_module = eval_module)
     f_oop, _ = (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen)
     f(u, p, iv) = f_oop(u, p, iv)
-    fd = DiscreteFunction(f, syms = Symbol.(dvs))
+    fd = DiscreteFunction(f; syms = Symbol.(dvs), sys = sys)
     DiscreteProblem(fd, u0, tspan, p; kwargs...)
 end
 

src/systems/jumps/jumpsystem.jl

@@ -283,7 +283,7 @@ function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,
         end
     end
 
-    df = DiscreteFunction{true, true}(f, syms = Symbol.(states(sys)),
+    df = DiscreteFunction{true, true}(f; syms = Symbol.(states(sys)), sys = sys,
                                       observed = observedfun)
     DiscreteProblem(df, u0, tspan, p; kwargs...)
 end

src/systems/nonlinear/nonlinearsystem.jl

@@ -233,6 +233,7 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sys
     end
 
     NonlinearFunction{iip}(f,
+                           sys = sys,
                            jac = _jac === nothing ? nothing : _jac,
                            jac_prototype = sparse ?
                                            similar(calculate_jacobian(sys, sparse = sparse),

src/systems/optimization/optimizationsystem.jl

@@ -231,6 +231,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         end
 
         _f = DiffEqBase.OptimizationFunction{iip}(f,
+                                                  sys = sys,
                                                   SciMLBase.NoAD();
                                                   grad = _grad,
                                                   hess = _hess,
@@ -244,6 +245,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                                   cons_expr = cons_expr)
     else
         _f = DiffEqBase.OptimizationFunction{iip}(f,
+                                                  sys = sys,
                                                   SciMLBase.NoAD();
                                                   grad = _grad,
                                                   hess = _hess,

test/discretesystem.jl

@@ -32,6 +32,7 @@ u0 = [S => 990.0, I => 10.0, R => 0.0]
 p = [β => 0.05, c => 10.0, γ => 0.25, δt => 0.1, nsteps => 400]
 tspan = (0.0, ModelingToolkit.value(substitute(nsteps, p))) # value function (from Symbolics) is used to convert a Num to Float64
 prob_map = DiscreteProblem(sys, u0, tspan, p)
+@test prob_map.f.sys === sys
 
 # Solution
 using OrdinaryDiffEq

test/nonlinearsystem.jl

@@ -79,6 +79,7 @@ sH = calculate_hessian(ns)
       getfield.(sparse.(sH), :rowval)
 
 prob = NonlinearProblem(ns, ones(3), ones(3))
+@test prob.f.sys === ns
 sol = solve(prob, NewtonRaphson())
 @test sol.u[1] ≈ sol.u[2]
 

test/odesystem.jl

@@ -213,6 +213,7 @@ p2 = (k₁ => 0.04,
       k₃ => 1e4)
 tspan = (0.0, 100000.0)
 prob1 = ODEProblem(sys, u0, tspan, p)
+@test prob1.f.sys === sys
 prob12 = ODEProblem(sys, u0, tspan, [0.04, 3e7, 1e4])
 prob13 = ODEProblem(sys, u0, tspan, (0.04, 3e7, 1e4))
 prob14 = ODEProblem(sys, u0, tspan, p2)