Support observed function building for split parameters

Author: YingboMa

src/systems/diffeqs/abstractodesystem.jl

@@ -397,32 +397,64 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = s
 
     obs = observed(sys)
     observedfun = if steady_state
-        let sys = sys, dict = Dict()
+        let sys = sys, dict = Dict(), ps = ps
             function generated_observed(obsvar, args...)
                 obs = get!(dict, value(obsvar)) do
                     build_explicit_observed_function(sys, obsvar)
                 end
                 if args === ()
                     let obs = obs
-                        (u, p, t = Inf) -> obs(u, p, t)
+                        (u, p, t = Inf) -> if ps isa Tuple
+                            obs(u, p..., t)
+                        else
+                            obs(u, p, t)
+                        end
                     end
                 else
-                    length(args) == 2 ? obs(args..., Inf) : obs(args...)
+                    if ps isa Tuple
+                        if length(args) == 2
+                            u, p = args
+                            obs(u, p..., Inf)
+                        else
+                            u, p, t = args
+                            obs(u, p..., t)
+                        end
+                    else
+                        if length(args) == 2
+                            u, p = args
+                            obs(u, p, Inf)
+                        else
+                            u, p, t = args
+                            obs(u, p, t)
+                        end
+                    end
                 end
             end
         end
     else
-        let sys = sys, dict = Dict()
+        let sys = sys, dict = Dict(), ps = ps
             function generated_observed(obsvar, args...)
                 obs = get!(dict, value(obsvar)) do
-                    build_explicit_observed_function(sys, obsvar; checkbounds = checkbounds)
+                    build_explicit_observed_function(sys,
+                        obsvar;
+                        checkbounds = checkbounds,
+                        ps)
                 end
                 if args === ()
                     let obs = obs
-                        (u, p, t) -> obs(u, p, t)
+                        (u, p, t) -> if ps isa Tuple
+                            obs(u, p..., t)
+                        else
+                            obs(u, p, t)
+                        end
                     end
                 else
-                    obs(args...)
+                    if ps isa Tuple # split parameters
+                        u, p, t = args
+                        obs(u, p..., t)
+                    else
+                        obs(args...)
+                    end
                 end
             end
         end

src/systems/diffeqs/odesystem.jl

@@ -314,6 +314,7 @@ function build_explicit_observed_function(sys, ts;
     output_type = Array,
     checkbounds = true,
     drop_expr = drop_expr,
+    ps = paramteres(sys),
     throw = true)
     if (isscalar = !(ts isa AbstractVector))
         ts = [ts]
@@ -389,13 +390,17 @@ function build_explicit_observed_function(sys, ts;
     if inputs !== nothing
         pars = setdiff(pars, inputs) # Inputs have been converted to parameters by io_preprocessing, remove those from the parameter list
     end
-    ps = DestructuredArgs(pars, inbounds = !checkbounds)
+    if ps isa Tuple
+        ps = DestructuredArgs.(ps, inbounds = !checkbounds)
+    else
+        ps = (DestructuredArgs(ps, inbounds = !checkbounds),)
+    end
     dvs = DestructuredArgs(states(sys), inbounds = !checkbounds)
     if inputs === nothing
-        args = [dvs, ps, ivs...]
+        args = [dvs, ps..., ivs...]
     else
         ipts = DestructuredArgs(inputs, inbounds = !checkbounds)
-        args = [dvs, ipts, ps, ivs...]
+        args = [dvs, ipts, ps..., ivs...]
     end
     pre = get_postprocess_fbody(sys)
 

test/split_parameters.jl

@@ -59,8 +59,8 @@ eqs = [D(y) ~ dy * a
     D(dy) ~ ddy * b
     ddy ~ sin(t) * c]
 
-@named sys = ODESystem(eqs, t, vars, pars)
-sys = structural_simplify(sys)
+@named model = ODESystem(eqs, t, vars, pars)
+sys = structural_simplify(model)
 
 tspan = (0.0, t_end)
 prob = ODEProblem(sys, [], tspan, [])
@@ -76,8 +76,7 @@ prob = ODEProblem(sys, [], tspan, []; tofloat = false)
 @test prob.p isa Tuple{Vector{Float64}, Vector{Int64}}
 sol = solve(prob, ImplicitEuler());
 @test sol.retcode == ReturnCode.Success
-
-
+sol[states(model)]
 
 # -------------------------  Observables
 
@@ -94,4 +93,4 @@ sys = structural_simplify(model)
 prob = ODEProblem(sys, Pair[int.x => 0.0], (0.0, 1.0))
 sol = solve(prob, Rodas4())
 @test sol.retcode == ReturnCode.Success
-sol[absb.output.u]
+sol[absb.output.u]