Merge pull request #34 from SciML/smc/macro

Compatibility with `@mtkmodel`

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkitNeuralNets"
 uuid = "f162e290-f571-43a6-83d9-22ecc16da15f"
 authors = ["Sebastian Micluța-Câmpeanu <sebastian.mc95@proton.me> and contributors"]
-version = "1.1.0"
+version = "1.2.0"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"

src/ModelingToolkitNeuralNets.jl

@@ -13,7 +13,7 @@ export NeuralNetworkBlock, multi_layer_feed_forward
 include("utils.jl")
 
 """
-    NeuralNetworkBlock(n_input = 1, n_output = 1;
+    NeuralNetworkBlock(; n_input = 1, n_output = 1,
         chain = multi_layer_feed_forward(n_input, n_output),
         rng = Xoshiro(0),
         init_params = Lux.initialparameters(rng, chain),
@@ -22,8 +22,7 @@ include("utils.jl")
 
 Create an `ODESystem` with a neural network inside.
 """
-function NeuralNetworkBlock(n_input = 1,
-        n_output = 1;
+function NeuralNetworkBlock(; n_input = 1, n_output = 1,
         chain = multi_layer_feed_forward(n_input, n_output),
         rng = Xoshiro(0),
         init_params = Lux.initialparameters(rng, chain),
@@ -46,6 +45,12 @@ function NeuralNetworkBlock(n_input = 1,
     return ude_comp
 end
 
+# added to avoid a breaking change from moving n_input & n_output in kwargs
+# https://github.com/SciML/ModelingToolkitNeuralNets.jl/issues/32
+function NeuralNetworkBlock(n_input, n_output = 1; kwargs...)
+    NeuralNetworkBlock(; n_input, n_output, kwargs...)
+end
+
 function lazyconvert(T, x::Symbolics.Arr)
     Symbolics.array_term(convert, T, x, size = size(x))
 end

test/lotka_volterra.jl

@@ -114,7 +114,7 @@ res = solve(op, Adam(), maxiters = 5000)#, callback = plot_cb)
 
 @test res.objective < 1
 
-res_p = SciMLStructures.replace(Tunable(), prob.p, res)
+res_p = SciMLStructures.replace(Tunable(), prob.p, res.u)
 res_prob = remake(prob, p = res_p)
 res_sol = solve(res_prob, Rodas4(), saveat = sol_ref.t)
 

test/macro.jl

@@ -0,0 +1,41 @@
+using ModelingToolkit, Symbolics
+using ModelingToolkit: t_nounits as t, D_nounits as D
+using OrdinaryDiffEq
+using ModelingToolkitNeuralNets
+using ModelingToolkitStandardLibrary.Blocks
+using Lux
+
+@mtkmodel Friction_UDE begin
+    @variables begin
+        y(t) = 0.0
+    end
+    @parameters begin
+        Fu
+    end
+    @components begin
+        nn_in = RealInputArray(nin = 1)
+        nn_out = RealOutputArray(nout = 1)
+    end
+    @equations begin
+        D(y) ~ Fu - nn_in.u[1]
+        y ~ nn_out.u[1]
+    end
+end
+
+@mtkmodel TestFriction_UDE begin
+    @components begin
+        friction_ude = Friction_UDE(Fu = 120.0)
+        nn = NeuralNetworkBlock(n_input = 1, n_output = 1)
+    end
+    @equations begin
+        connect(friction_ude.nn_in, nn.output)
+        connect(friction_ude.nn_out, nn.input)
+    end
+end
+
+@mtkbuild sys = TestFriction_UDE()
+
+prob = ODEProblem(sys, [], (0, 1.0), [])
+sol = solve(prob, Rodas4())
+
+@test SciMLBase.successful_retcode(sol)

test/runtests.jl

@@ -5,4 +5,5 @@ using SafeTestsets
 @testset verbose=true "ModelingToolkitNeuralNets.jl" begin
     @safetestset "QA" include("qa.jl")
     @safetestset "Basic" include("lotka_volterra.jl")
+    @safetestset "MTK model macro compatibility" include("macro.jl")
 end