test: fix initialization, account for Initial parameters in test

Project.toml

@@ -16,7 +16,7 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Aqua = "0.8"
 ComponentArrays = "0.15.11"
 ForwardDiff = "0.10.36"
-JET = "0.8"
+JET = "0.8, 0.9"
 Lux = "1"
 LuxCore = "1"
 ModelingToolkit = "9.9"

docs/src/friction.md

@@ -110,8 +110,8 @@ sys = structural_simplify(ude_sys)
 We now setup the loss function and the optimization loop.
 
 ```@example friction
-function loss(x, (prob, sol_ref, get_vars, get_refs))
-    new_p = SciMLStructures.replace(Tunable(), prob.p, x)
+function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
+    new_p = set_x(prob, x)
     new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0))
     ts = sol_ref.t
     new_sol = solve(new_prob, Rodas4(), saveat = ts, abstol = 1e-8, reltol = 1e-8)
@@ -131,14 +131,15 @@ of = OptimizationFunction{true}(loss, AutoForwardDiff())
 prob = ODEProblem(sys, [], (0, 0.1), [])
 get_vars = getu(sys, [sys.friction.y])
 get_refs = getu(model_true, [model_true.y])
-x0 = reduce(vcat, getindex.((default_values(sys),), tunable_parameters(sys)))
+set_x = setp_oop(sys, sys.nn.p)
+x0 = default_values(sys)[sys.nn.p]
 
 cb = (opt_state, loss) -> begin
     @info "step $(opt_state.iter), loss: $loss"
     return false
 end
 
-op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs))
+op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs, set_x))
 res = solve(op, Adam(5e-3); maxiters = 10000, callback = cb)
 ```
 
@@ -147,7 +148,7 @@ res = solve(op, Adam(5e-3); maxiters = 10000, callback = cb)
 We now have a trained neural network! We can check whether running the simulation of the model embedded with the neural network matches the data or not.
 
 ```@example friction
-res_p = SciMLStructures.replace(Tunable(), prob.p, res.u)
+res_p = set_x(prob, res.u)
 res_prob = remake(prob, p = res_p)
 res_sol = solve(res_prob, Rodas4(), saveat = sol_ref.t)
 @test first.(sol_ref.u)≈first.(res_sol.u) rtol=1e-3 #hide

test/lotka_volterra.jl

@@ -8,7 +8,7 @@ using SymbolicIndexingInterface
 using Optimization
 using OptimizationOptimisers: Adam
 using SciMLStructures
-using SciMLStructures: Tunable
+using SciMLStructures: Tunable, canonicalize
 using ForwardDiff
 using StableRNGs
 
@@ -51,7 +51,7 @@ eqs = [connect(model.nn_in, nn.output)
 
 ude_sys = complete(ODESystem(
     eqs, ModelingToolkit.t_nounits, systems = [model, nn],
-    name = :ude_sys, defaults = [nn.input.u => [0.0, 0.0]]))
+    name = :ude_sys))
 
 sys = structural_simplify(ude_sys)
 
@@ -61,13 +61,14 @@ model_true = structural_simplify(lotka_true())
 prob_true = ODEProblem{true, SciMLBase.FullSpecialize}(model_true, [], (0, 1.0), [])
 sol_ref = solve(prob_true, Rodas4())
 
-x0 = reduce(vcat, getindex.((default_values(sys),), tunable_parameters(sys)))
+x0 = default_values(sys)[nn.p]
 
 get_vars = getu(sys, [sys.lotka.x, sys.lotka.y])
 get_refs = getu(model_true, [model_true.x, model_true.y])
+set_x = setp_oop(sys, nn.p)
 
-function loss(x, (prob, sol_ref, get_vars, get_refs))
-    new_p = SciMLStructures.replace(Tunable(), prob.p, x)
+function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
+    new_p = set_x(prob, x)
     new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0))
     ts = sol_ref.t
     new_sol = solve(new_prob, Rodas4(), saveat = ts)
@@ -87,14 +88,14 @@ end
 
 of = OptimizationFunction{true}(loss, AutoForwardDiff())
 
-ps = (prob, sol_ref, get_vars, get_refs);
+ps = (prob, sol_ref, get_vars, get_refs, set_x);
 
 @test_call target_modules=(ModelingToolkitNeuralNets,) loss(x0, ps)
 @test_opt target_modules=(ModelingToolkitNeuralNets,) loss(x0, ps)
 
 @test all(.!isnan.(ForwardDiff.gradient(Base.Fix2(of, ps), x0)))
 
-op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs))
+op = OptimizationProblem(of, x0, ps)
 
 # using Plots
 
@@ -114,7 +115,7 @@ res = solve(op, Adam(), maxiters = 5000)#, callback = plot_cb)
 
 @test res.objective < 1
 
-res_p = SciMLStructures.replace(Tunable(), prob.p, res.u)
+res_p = set_x(prob, res.u)
 res_prob = remake(prob, p = res_p)
 res_sol = solve(res_prob, Rodas4(), saveat = sol_ref.t)