fix: some bug fixes

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -165,7 +165,7 @@ function generate_custom_function(sys::AbstractSystem, exprs, dvs = unknowns(sys
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system.")
     end
-    p = reorder_parameters(sys, ps)
+    p = reorder_parameters(sys, unwrap.(ps))
     isscalar = !(exprs isa AbstractArray)
     if wrap_code === nothing
         wrap_code = isscalar ? identity : (identity, identity)
@@ -1701,6 +1701,7 @@ function linearization_function(sys::AbstractSystem, inputs,
     u0, _p, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
     ps = parameters(sys)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
+            @show p full_parameters(sys)
         p = MTKParameters(sys, p)
     else
         p = _p

src/systems/index_cache.jl

@@ -240,6 +240,7 @@ function reorder_parameters(sys::AbstractSystem, ps; kwargs...)
 end
 
 function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
+    isempty(ps) && return ()
     param_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
     for temp in ic.tunable_buffer_sizes)
     disc_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
@@ -281,6 +282,9 @@ function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
             end
         end
     end
+    if all(isempty, result)
+        return ()
+    end
     return result
 end
 

src/systems/optimization/constraints_system.jl

@@ -164,10 +164,11 @@ function calculate_jacobian(sys::ConstraintsSystem; sparse = false, simplify = f
 end
 
 function generate_jacobian(
-        sys::ConstraintsSystem, vs = unknowns(sys), ps = parameters(sys);
+        sys::ConstraintsSystem, vs = unknowns(sys), ps = full_parameters(sys);
         sparse = false, simplify = false, kwargs...)
     jac = calculate_jacobian(sys, sparse = sparse, simplify = simplify)
-    return build_function(jac, vs, ps; kwargs...)
+    p = reorder_parameters(sys, ps)
+    return build_function(jac, vs, p...; kwargs...)
 end
 
 function calculate_hessian(sys::ConstraintsSystem; sparse = false, simplify = false)
@@ -181,19 +182,20 @@ function calculate_hessian(sys::ConstraintsSystem; sparse = false, simplify = fa
     return hess
 end
 
-function generate_hessian(sys::ConstraintsSystem, vs = unknowns(sys), ps = parameters(sys);
+function generate_hessian(sys::ConstraintsSystem, vs = unknowns(sys), ps = full_parameters(sys);
         sparse = false, simplify = false, kwargs...)
     hess = calculate_hessian(sys, sparse = sparse, simplify = simplify)
-    return build_function(hess, vs, ps; kwargs...)
+    p = reorder_parameters(sys, ps)
+    return build_function(hess, vs, p...; kwargs...)
 end
 
 function generate_function(sys::ConstraintsSystem, dvs = unknowns(sys),
-        ps = parameters(sys);
+        ps = full_parameters(sys);
         kwargs...)
     lhss = generate_canonical_form_lhss(sys)
     pre, sol_states = get_substitutions_and_solved_unknowns(sys)
-
-    func = build_function(lhss, value.(dvs), value.(ps); postprocess_fbody = pre,
+    p = reorder_parameters(sys, value.(ps))
+    func = build_function(lhss, value.(dvs), p...; postprocess_fbody = pre,
         states = sol_states, kwargs...)
 
     cstr = constraints(sys)

src/systems/optimization/optimizationsystem.jl

@@ -361,6 +361,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
 
     if length(cstr) > 0
         @named cons_sys = ConstraintsSystem(cstr, dvs, ps)
+        cons_sys = complete(cons_sys)
         cons, lcons_, ucons_ = generate_function(cons_sys, checkbounds = checkbounds,
             linenumbers = linenumbers,
             expression = Val{false})

src/systems/parameter_buffer.jl

@@ -14,10 +14,10 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
     else
         error("Cannot create MTKParameters if system does not have index_cache")
     end
-    all_ps = Set(unwrap.(parameters(sys)))
-    union!(all_ps, default_toterm.(unwrap.(parameters(sys))))
+    all_ps = Set(unwrap.(full_parameters(sys)))
+    union!(all_ps, default_toterm.(unwrap.(full_parameters(sys))))
     if p isa Vector && !(eltype(p) <: Pair) && !isempty(p)
-        ps = parameters(sys)
+        ps = full_parameters(sys)
         length(p) == length(ps) || error("Invalid parameters")
         p = ps .=> p
     end
@@ -38,7 +38,7 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
 
     for (sym, _) in p
         if istree(sym) && operation(sym) === getindex &&
-           is_parameter(sys, first(arguments(sym)))
+            first(arguments(sym)) in all_ps
             error("Scalarized parameter values ($sym) are not supported. Instead of `[p[1] => 1.0, p[2] => 2.0]` use `[p => [1.0, 2.0]]`")
         end
     end
@@ -99,7 +99,7 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
             i, j = ic.dependent_idx[sym]
             dep_exprs.x[i][j] = wrap(val)
         end
-        p = reorder_parameters(ic, parameters(sys))
+        p = reorder_parameters(ic, full_parameters(sys))
         oop, iip = build_function(dep_exprs, p...)
         update_function_iip, update_function_oop = RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(iip),
         RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(oop)

src/systems/systems.jl

@@ -32,6 +32,16 @@ function structural_simplify(
         @set! newsys.parent = complete(sys; split)
     end
     newsys = complete(newsys; split)
+    if has_defaults(newsys) && (defs = get_defaults(newsys)) !== nothing
+        ks = collect(keys(defs))
+        for k in ks
+            if Symbolics.isarraysymbolic(k) && Symbolics.shape(k) !== Symbolics.Unknown()
+                for i in eachindex(k)
+                    defs[k[i]] = defs[k][i]
+                end
+            end
+        end
+    end
     if newsys′ isa Tuple
         idxs = [parameter_index(newsys, i) for i in io[1]]
         return newsys, idxs

test/initializationsystem.jl

@@ -355,7 +355,7 @@ prob = ODEProblem(sys, [sys.dx => 1], (0, 1)) # OK
 prob = ODEProblem(sys, [sys.ddx => -2], (0, 1), guesses = [sys.dx => 1])
 sol = solve(prob, Tsit5())
 @test SciMLBase.successful_retcode(sol)
-@test sol[1] == [1.0]
+@test sol.u[1] == [1.0]
 
 ## Late binding initialization_eqs
 
@@ -376,7 +376,7 @@ end
 prob = ODEProblem(sys, [], (0, 1), guesses = [sys.dx => 1])
 sol = solve(prob, Tsit5())
 @test SciMLBase.successful_retcode(sol)
-@test sol[1] == [1.0]
+@test sol.u[1] == [1.0]
 
 # Steady state initialization
 
@@ -417,16 +417,16 @@ tspan = (0.0, 10.0)
 
 prob = ODEProblem(simpsys, [D(x) => 0.0, y => 0.0], tspan, guesses = [x => 0.0])
 sol = solve(prob, Tsit5())
-@test sol[1] == [0.0, 0.0]
+@test sol.u[1] == [0.0, 0.0]
 
 # Initialize with an observed variable
 prob = ODEProblem(simpsys, [z => 0.0], tspan, guesses = [x => 2.0, y => 4.0])
 sol = solve(prob, Tsit5())
-@test sol[1] == [0.0, 0.0]
+@test sol.u[1] == [0.0, 0.0]
 
 prob = ODEProblem(simpsys, [z => 1.0, y => 1.0], tspan, guesses = [x => 2.0])
 sol = solve(prob, Tsit5())
-@test sol[1] == [0.0, 1.0]
+@test sol.u[1] == [0.0, 1.0]
 
-@test_broken @test_warn "Initialization system is underdetermined. 1 equations for 3 unknowns. Initialization will default to using least squares. To suppress this warning pass warn_initialize_determined = false." prob=ODEProblem(
-    simpsys, [], tspan, guesses = [x => 2.0])
+# This should warn, but logging tests can't be marked as broken
+@test_logs prob=ODEProblem(simpsys, [], tspan, guesses = [x => 2.0])

test/odesystem.jl

@@ -472,7 +472,7 @@ sts = @variables x(t)[1:3]=[1, 2, 3.0] y(t)=1.0
 ps = @parameters p[1:3] = [1, 2, 3]
 eqs = [collect(D.(x) .~ x)
        D(y) ~ norm(collect(x)) * y - x[1]]
-@named sys = ODESystem(eqs, t, [sts...;], [ps...;])
+@named sys = ODESystem(eqs, t, sts, ps)
 sys = structural_simplify(sys)
 @test isequal(@nonamespace(sys.x), x)
 @test isequal(@nonamespace(sys.y), y)