Merge pull request #63 from SciML/smc/update

Update to MTK@v10 and simplify connectors

Author: SebastianM-C

.github/workflows/Downgrade.yml

@@ -16,14 +16,14 @@ jobs:
   test:
     runs-on: ${{ matrix.os }}
     env:
-      GROUP: ${{ matrix.group }}    
+      GROUP: ${{ matrix.group }}
     strategy:
       fail-fast: false
       matrix:
         group:
           - Core
         version:
-          - '1'
+          - '1.10'
         os:
           - ubuntu-latest
           - macos-latest
@@ -40,4 +40,4 @@ jobs:
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
       - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-runtest@v1

CHANGELOG.md

@@ -0,0 +1,6 @@
+# ModelingToolkitNeuralNets v2
+
+## Breaking changes
+
+- The `NeuralNetworkBlock` no longer uses `RealInputArray` & `RealOutputArray`,
+the ports are now `inputs` and `outputs` and they are normal vector variables.

Project.toml

@@ -1,41 +1,41 @@
 name = "ModelingToolkitNeuralNets"
 uuid = "f162e290-f571-43a6-83d9-22ecc16da15f"
 authors = ["Sebastian Micluța-Câmpeanu <[email protected]> and contributors"]
-version = "1.7.0"
+version = "2.0.0"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
-ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
 Aqua = "0.8"
-ComponentArrays = "0.15.11"
-DifferentiationInterface = "0.6"
-ForwardDiff = "0.10.36"
+ComponentArrays = "0.15.28"
+DifferentiationInterface = "0.6, 0.7"
+ForwardDiff = "0.10.36, 1"
 IntervalSets = "0.7.10"
-JET = "0.8, 0.9"
-Lux = "1"
-LuxCore = "1"
-ModelingToolkit = "9.64"
-ModelingToolkitStandardLibrary = "2.7"
-Optimization = "3.24, 4"
-OptimizationOptimisers = "0.2.1, 0.3"
-OrdinaryDiffEq = "6.74"
+JET = "0.8, 0.9, 0.10"
+Lux = "1.14"
+LuxCore = "1.2"
+ModelingToolkit = "10"
+ModelingToolkitStandardLibrary = "2.24"
+Optimization = "4"
+OptimizationOptimisers = "0.3"
+OrdinaryDiffEqVerner = "1.2"
 Random = "1.10"
 SafeTestsets = "0.1"
 SciMLSensitivity = "7.72"
 SciMLStructures = "1.1.0"
 StableRNGs = "1"
-SymbolicIndexingInterface = "0.3.15"
-Symbolics = "6.36"
+Statistics = "1.10"
+SymbolicIndexingInterface = "0.3.41"
+Symbolics = "6.43"
 Test = "1.10"
-Zygote = "0.6.73"
+Zygote = "0.6.73, 0.7"
 julia = "1.10"
 
 [extras]
@@ -45,14 +45,19 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
+ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "JET", "Test", "OrdinaryDiffEq", "DifferentiationInterface", "SciMLSensitivity", "Zygote", "ForwardDiff", "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures", "StableRNGs", "SymbolicIndexingInterface"]
+test = ["Aqua", "JET", "Test", "OrdinaryDiffEqVerner", "DifferentiationInterface",
+    "SciMLSensitivity", "Zygote", "ForwardDiff", "ModelingToolkitStandardLibrary",
+    "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures",
+    "StableRNGs", "Statistics", "SymbolicIndexingInterface"]

README.md

@@ -12,3 +12,11 @@ ModelingToolkitNeuralNets.jl is a package to create neural network blocks define
 ## Tutorials and Documentation
 
 For information on using the package, [see the stable documentation](https://docs.sciml.ai/ModelingToolkitNeuralNets/stable/). Use the [in-development documentation](https://docs.sciml.ai/ModelingToolkitNeuralNets/dev/) for the version of the documentation, which contains the unreleased features.
+
+## Breaking changes in v2
+
+The `NeuralNetworkBlock` no longer uses `RealInputArray` & `RealOutputArray`,
+the ports are now `inputs` and `outputs` and they are normal vector variables.
+This simplifies the usage a bit and removes the need for the ModelingToolkitStandardLibrary dependency.
+
+This version also moves to ModelingToolkit@v10.

docs/Project.toml

@@ -10,13 +10,17 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
+[sources]
+ModelingToolkitNeuralNets = {path = ".."}
+
 [compat]
 Documenter = "1.3"
 Lux = "1"
-ModelingToolkit = "9.9"
-ModelingToolkitNeuralNets = "1"
+ModelingToolkit = "10"
+ModelingToolkitNeuralNets = "2"
 ModelingToolkitStandardLibrary = "2.7"
 Optimization = "3.24, 4.0"
 OptimizationOptimisers = "0.2.1, 0.3"
@@ -25,6 +29,3 @@ Plots = "1"
 SciMLStructures = "1.1.0"
 StableRNGs = "1"
 SymbolicIndexingInterface = "0.3.15"
-
-[sources]
-ModelingToolkitNeuralNets = { path = ".." }

docs/src/api.md

@@ -2,4 +2,5 @@
 
 ```@docs
 NeuralNetworkBlock
+SymbolicNeuralNetwork
 ```

docs/src/friction.md

@@ -18,6 +18,7 @@ using OptimizationOptimisers: Adam
 using SciMLStructures
 using SciMLStructures: Tunable
 using SymbolicIndexingInterface
+using Statistics
 using StableRNGs
 using Lux
 using Plots
@@ -49,16 +50,16 @@ function friction_true()
     eqs = [
         Dt(y) ~ Fu - friction(y)
     ]
-    return ODESystem(eqs, t, name = :friction_true)
+    return System(eqs, t, name = :friction_true)
 end
 ```
 
 Now that we have defined the model, we will simulate it from 0 to 0.1 seconds.
 
 ```@example friction
-model_true = structural_simplify(friction_true())
-prob_true = ODEProblem(model_true, [], (0, 0.1), [])
-sol_ref = solve(prob_true, Rodas4(); saveat = 0.001)
+model_true = mtkcompile(friction_true())
+prob_true = ODEProblem(model_true, [], (0, 0.1))
+sol_ref = solve(prob_true, Vern7(); saveat = 0.001)
 ```
 
 Let's plot it.
@@ -81,28 +82,23 @@ Now, we will try to learn the same friction model using a neural network. We wil
 function friction_ude(Fu)
     @variables y(t) = 0.0
     @constants Fu = Fu
-    @named nn_in = RealInputArray(nin = 1)
-    @named nn_out = RealOutputArray(nout = 1)
-    eqs = [Dt(y) ~ Fu - nn_in.u[1]
-           y ~ nn_out.u[1]]
-    return ODESystem(eqs, t, name = :friction, systems = [nn_in, nn_out])
-end
 
-Fu = 120.0
-model = friction_ude(Fu)
+    chain = Lux.Chain(
+        Lux.Dense(1 => 10, Lux.mish, use_bias = false),
+        Lux.Dense(10 => 10, Lux.mish, use_bias = false),
+        Lux.Dense(10 => 1, use_bias = false)
+    )
+    @named nn = NeuralNetworkBlock(1, 1; chain = chain, rng = StableRNG(1111))
 
-chain = Lux.Chain(
-    Lux.Dense(1 => 10, Lux.mish, use_bias = false),
-    Lux.Dense(10 => 10, Lux.mish, use_bias = false),
-    Lux.Dense(10 => 1, use_bias = false)
-)
-@named nn = NeuralNetworkBlock(1, 1; chain = chain, rng = StableRNG(1111))
+    eqs = [Dt(y) ~ Fu - nn.outputs[1]
+           y ~ nn.inputs[1]]
+    return System(eqs, t, name = :friction, systems = [nn])
+end
 
-eqs = [connect(model.nn_in, nn.output)
-       connect(model.nn_out, nn.input)]
+Fu = 120.0
 
-ude_sys = complete(ODESystem(eqs, t, systems = [model, nn], name = :ude_sys))
-sys = structural_simplify(ude_sys)
+ude_sys = friction_ude(Fu)
+sys = mtkcompile(ude_sys)
 ```
 
 ## Optimization Setup
@@ -114,22 +110,19 @@ 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)
-    loss = zero(eltype(x))
-    for i in eachindex(new_sol.u)
-        loss += sum(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i)))
-    end
+    new_sol = solve(new_prob, Vern7(), saveat = ts, abstol = 1e-8, reltol = 1e-8)
+
     if SciMLBase.successful_retcode(new_sol)
-        loss
+        mean(abs2.(reduce(hcat, get_vars(new_sol)) .- reduce(hcat, get_refs(sol_ref))))
     else
         Inf
     end
 end
 
-of = OptimizationFunction{true}(loss, AutoForwardDiff())
+of = OptimizationFunction(loss, AutoForwardDiff())
 
-prob = ODEProblem(sys, [], (0, 0.1), [])
-get_vars = getu(sys, [sys.friction.y])
+prob = ODEProblem(sys, [], (0, 0.1))
+get_vars = getu(sys, [sys.y])
 get_refs = getu(model_true, [model_true.y])
 set_x = setp_oop(sys, sys.nn.p)
 x0 = default_values(sys)[sys.nn.p]
@@ -150,31 +143,31 @@ We now have a trained neural network! We can check whether running the simulatio
 ```@example friction
 res_p = set_x(prob, res.u)
 res_prob = remake(prob, p = res_p)
-res_sol = solve(res_prob, Rodas4(), saveat = sol_ref.t)
+res_sol = solve(res_prob, Vern7(), saveat = sol_ref.t)
 @test first.(sol_ref.u)≈first.(res_sol.u) rtol=1e-3 #hide
-@test friction.(first.(sol_ref.u))≈(getindex.(res_sol[sys.nn.output.u], 1)) rtol=1e-1 #hide
+@test friction.(first.(sol_ref.u))≈(getindex.(res_sol[sys.nn.outputs], 1)) rtol=1e-1 #hide
 nothing #hide
 ```
 
 Also, it would be interesting to check the simulation before the training to get an idea of the starting point of the network.
 
 ```@example friction
-initial_sol = solve(prob, Rodas4(), saveat = sol_ref.t)
+initial_sol = solve(prob, Vern7(), saveat = sol_ref.t)
 ```
 
 Now we plot it.
 
 ```@example friction
 scatter(sol_ref, idxs = [model_true.y], label = "ground truth velocity")
-plot!(res_sol, idxs = [sys.friction.y], label = "velocity after training")
-plot!(initial_sol, idxs = [sys.friction.y], label = "velocity before training")
+plot!(res_sol, idxs = [sys.y], label = "velocity after training")
+plot!(initial_sol, idxs = [sys.y], label = "velocity before training")
 ```
 
 It matches the data well! Let's also check the predictions for the friction force and whether the network learnt the friction model or not.
 
 ```@example friction
 scatter(sol_ref.t, friction.(first.(sol_ref.u)), label = "ground truth friction")
-plot!(res_sol.t, getindex.(res_sol[sys.nn.output.u], 1),
+plot!(res_sol.t, getindex.(res_sol[sys.nn.outputs], 1),
     label = "friction from neural network")
 ```
 

docs/src/index.md

@@ -23,7 +23,7 @@ Pkg.add("ModelingToolkitNeuralNets")
 
   - See the [SciML Style Guide](https://github.com/SciML/SciMLStyle) for common coding practices and other style decisions.
   - There are a few community forums:
-    
+
       + The #diffeq-bridged and #sciml-bridged channels in the
         [Julia Slack](https://julialang.org/slack/)
       + The #diffeq-bridged and #sciml-bridged channels in the

src/ModelingToolkitNeuralNets.jl

@@ -1,10 +1,9 @@
 module ModelingToolkitNeuralNets
 
-using ModelingToolkit: @parameters, @named, ODESystem, t_nounits
+using ModelingToolkit: @parameters, @named, @variables, System, t_nounits
 using IntervalSets: var".."
-using ModelingToolkitStandardLibrary.Blocks: RealInputArray, RealOutputArray
 using Symbolics: Symbolics, @register_array_symbolic, @wrapped
-using LuxCore: stateless_apply
+using LuxCore: stateless_apply, outputsize
 using Lux: Lux
 using Random: Xoshiro
 using ComponentArrays: ComponentArray
@@ -21,7 +20,7 @@ include("utils.jl")
         eltype = Float64,
         name)
 
-Create an `ODESystem` with a neural network inside.
+Create a component neural network as a `System`.
 """
 function NeuralNetworkBlock(; n_input = 1, n_output = 1,
         chain = multi_layer_feed_forward(n_input, n_output),
@@ -31,19 +30,21 @@ function NeuralNetworkBlock(; n_input = 1, n_output = 1,
         name)
     ca = ComponentArray{eltype}(init_params)
 
-    @parameters p[1:length(ca)] = Vector(ca)
+    @parameters p[1:length(ca)]=Vector(ca) [tunable = true]
     @parameters T::typeof(typeof(ca))=typeof(ca) [tunable = false]
-    @parameters lux_model::typeof(chain) = chain
+    @parameters lux_model::typeof(chain)=chain [tunable = false]
 
-    @named input = RealInputArray(nin = n_input)
-    @named output = RealOutputArray(nout = n_output)
+    @variables inputs(t_nounits)[1:n_input] [input = true]
+    @variables outputs(t_nounits)[1:n_output] [output = true]
 
-    out = stateless_apply(lux_model, input.u, lazyconvert(T, p))
+    expected_outsz = only(outputsize(chain, inputs, rng))
+    msg = "The outputsize of the given Lux network ($expected_outsz) does not match `n_output = $n_output`"
+    @assert n_output==expected_outsz msg
 
-    eqs = [output.u ~ out]
+    eqs = [outputs ~ stateless_apply(lux_model, inputs, lazyconvert(T, p))]
 
-    ude_comp = ODESystem(
-        eqs, t_nounits, [], [lux_model, p, T]; systems = [input, output], name)
+    ude_comp = System(
+        eqs, t_nounits, [inputs, outputs], [lux_model, p, T]; name)
     return ude_comp
 end