Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions .JuliaFormatter.toml

This file was deleted.

14 changes: 10 additions & 4 deletions .github/workflows/FormatCheck.yml
Original file line number Diff line number Diff line change
@@ -1,13 +1,19 @@
name: "Format Check"
name: format-check

on:
push:
branches:
- 'master'
- 'main'
- 'release-'
tags: '*'
pull_request:

jobs:
format-check:
name: "Format Check"
uses: "SciML/.github/.github/workflows/format-check.yml@v1"
runic:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: fredrikekre/runic-action@v1
with:
version: '1'
22 changes: 14 additions & 8 deletions docs/make.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,24 +6,30 @@ ENV["JULIA_DEBUG"] = "Documenter"
cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml", force = true)
cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true)

DocMeta.setdocmeta!(ModelingToolkitNeuralNets, :DocTestSetup,
:(using ModelingToolkitNeuralNets); recursive = true)
DocMeta.setdocmeta!(
ModelingToolkitNeuralNets, :DocTestSetup,
:(using ModelingToolkitNeuralNets); recursive = true
)

makedocs(;
modules = [ModelingToolkitNeuralNets],
authors = "Sebastian Micluța-Câmpeanu <sebastian.mc95@proton.me> and contributors",
sitename = "ModelingToolkitNeuralNets.jl",
format = Documenter.HTML(assets = ["assets/favicon.ico"],
canonical = "https://docs.sciml.ai/ModelingToolkitNeuralNets.jl/stable/"),
format = Documenter.HTML(
assets = ["assets/favicon.ico"],
canonical = "https://docs.sciml.ai/ModelingToolkitNeuralNets.jl/stable/"
),
clean = true,
doctest = false,
linkcheck = true,
pages = [
"Home" => "index.md",
"Tutorials" => ["NeuralNetworkBlock" => "nnblock.md"
"Friction Model" => "friction.md"
"Symbolic UDE Creation" => "symbolic_ude_tutorial.md"],
"API" => "api.md"
"Tutorials" => [
"NeuralNetworkBlock" => "nnblock.md"
"Friction Model" => "friction.md"
"Symbolic UDE Creation" => "symbolic_ude_tutorial.md"
],
"API" => "api.md",
]
)

Expand Down
33 changes: 19 additions & 14 deletions src/ModelingToolkitNeuralNets.jl
Original file line number Diff line number Diff line change
Expand Up @@ -22,40 +22,43 @@ include("utils.jl")

Create a component neural network as a `System`.
"""
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),
eltype = Float64,
name)
name
)
ca = ComponentArray{eltype}(init_params)

@parameters p[1:length(ca)]=Vector(ca) [tunable = true]
@parameters T::typeof(typeof(ca))=typeof(ca) [tunable = false]
@parameters lux_model::typeof(chain)=chain [tunable = false]
@parameters p[1:length(ca)] = Vector(ca) [tunable = true]
@parameters T::typeof(typeof(ca)) = typeof(ca) [tunable = false]
@parameters lux_model::typeof(chain) = chain [tunable = false]

@variables inputs(t_nounits)[1:n_input] [input = true]
@variables outputs(t_nounits)[1:n_output] [output = true]

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
@assert n_output == expected_outsz msg

eqs = [outputs ~ stateless_apply(lux_model, inputs, lazyconvert(T, p))]

ude_comp = System(
eqs, t_nounits, [inputs, outputs], [lux_model, p, T]; name)
eqs, t_nounits, [inputs, outputs], [lux_model, p, T]; name
)
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...)
return NeuralNetworkBlock(; n_input, n_output, kwargs...)
end

function lazyconvert(T, x::Symbolics.Arr)
Symbolics.array_term(convert, T, x, size = size(x))
return Symbolics.array_term(convert, T, x, size = size(x))
end

"""
Expand Down Expand Up @@ -96,13 +99,15 @@ where `sys` is a system (e.g. `ODESystem`) that contains `NN`, `input` is a vect

To get the underlying Lux model you can use `get_network(defaults(sys)[sys.NN])` or
"""
function SymbolicNeuralNetwork(; n_input = 1, n_output = 1,
function SymbolicNeuralNetwork(;
n_input = 1, n_output = 1,
chain = multi_layer_feed_forward(n_input, n_output),
rng = Xoshiro(0),
init_params = Lux.initialparameters(rng, chain),
nn_name = :NN,
nn_p_name = :p,
eltype = Float64)
eltype = Float64
)
ca = ComponentArray{eltype}(init_params)
wrapper = StatelessApplyWrapper(chain, typeof(ca))

Expand All @@ -118,16 +123,16 @@ struct StatelessApplyWrapper{NN}
end

function (wrapper::StatelessApplyWrapper)(input::AbstractArray, nn_p::AbstractVector)
stateless_apply(get_network(wrapper), input, convert(wrapper.T, nn_p))
return stateless_apply(get_network(wrapper), input, convert(wrapper.T, nn_p))
end

function (wrapper::StatelessApplyWrapper)(input::Number, nn_p::AbstractVector)
wrapper([input], nn_p)
return wrapper([input], nn_p)
end

function Base.show(io::IO, m::MIME"text/plain", wrapper::StatelessApplyWrapper)
printstyled(io, "LuxCore.stateless_apply wrapper for:\n", color = :gray)
show(io, m, get_network(wrapper))
return show(io, m, get_network(wrapper))
end

get_network(wrapper::StatelessApplyWrapper) = wrapper.lux_model
Expand Down
19 changes: 12 additions & 7 deletions src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,18 +6,23 @@ Create a Lux.jl `Chain` for use in [`NeuralNetworkBlock`](@ref)s. The weights of
are multiplied by the `initial_scaling_factor` in order to make the initial contribution
of the network small and thus help with achieving a stable starting position for the training.
"""
function multi_layer_feed_forward(; n_input, n_output, width::Int = 4,
depth::Int = 1, activation = tanh, use_bias = true, initial_scaling_factor = 1e-8)
Lux.Chain(
function multi_layer_feed_forward(;
n_input, n_output, width::Int = 4,
depth::Int = 1, activation = tanh, use_bias = true, initial_scaling_factor = 1.0e-8
)
return Lux.Chain(
Lux.Dense(n_input, width, activation; use_bias),
[Lux.Dense(width, width, activation; use_bias) for _ in 1:(depth)]...,
Lux.Dense(width, n_output;
Lux.Dense(
width, n_output;
init_weight = (
rng, a...) -> initial_scaling_factor *
Lux.kaiming_uniform(rng, a...), use_bias)
rng, a...,
) -> initial_scaling_factor *
Lux.kaiming_uniform(rng, a...), use_bias
)
)
end

function multi_layer_feed_forward(n_input, n_output; kwargs...)
multi_layer_feed_forward(; n_input, n_output, kwargs...)
return multi_layer_feed_forward(; n_input, n_output, kwargs...)
end
47 changes: 25 additions & 22 deletions test/lotka_volterra.jl
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@ using Statistics
using Lux

function lotka_ude(chain)
@variables t x(t)=3.1 y(t)=1.5
@parameters α=1.3 [tunable = false] δ=1.8 [tunable = false]
@variables t x(t) = 3.1 y(t) = 1.5
@parameters α = 1.3 [tunable = false] δ = 1.8 [tunable = false]
Dt = ModelingToolkit.D_nounits

@named nn = NeuralNetworkBlock(2, 2; chain, rng = StableRNG(42))
Expand All @@ -27,20 +27,21 @@ function lotka_ude(chain)
Dt(x) ~ α * x + nn.outputs[1],
Dt(y) ~ -δ * y + nn.outputs[2],
nn.inputs[1] ~ x,
nn.inputs[2] ~ y
nn.inputs[2] ~ y,
]
return System(
eqs, ModelingToolkit.t_nounits, name = :lotka, systems = [nn])
eqs, ModelingToolkit.t_nounits, name = :lotka, systems = [nn]
)
end

function lotka_true()
@variables t x(t)=3.1 y(t)=1.5
@parameters α=1.3 [tunable = false] β=0.9 γ=0.8 δ=1.8 [tunable = false]
@variables t x(t) = 3.1 y(t) = 1.5
@parameters α = 1.3 [tunable = false] β = 0.9 γ = 0.8 δ = 1.8 [tunable = false]
Dt = ModelingToolkit.D_nounits

eqs = [
Dt(x) ~ α * x - β * x * y,
Dt(y) ~ -δ * y + γ * x * y
Dt(y) ~ -δ * y + γ * x * y,
]
return System(eqs, ModelingToolkit.t_nounits, name = :lotka_true)
end
Expand All @@ -58,7 +59,7 @@ prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0, 5.0))

model_true = mtkcompile(lotka_true())
prob_true = ODEProblem{true, SciMLBase.FullSpecialize}(model_true, [], (0, 5.0))
sol_ref = solve(prob_true, Vern9(), abstol = 1e-8, reltol = 1e-8)
sol_ref = solve(prob_true, Vern9(), abstol = 1.0e-8, reltol = 1.0e-8)

ts = range(0, 5.0, length = 21)
data = reduce(hcat, sol_ref(ts, idxs = [model_true.x, model_true.y]).u)
Expand All @@ -71,9 +72,9 @@ set_x = setsym_oop(sys, sys.nn.p)
function loss(x, (prob, get_vars, data, ts, set_x))
new_u0, new_p = set_x(prob, x)
new_prob = remake(prob, p = new_p, u0 = new_u0)
new_sol = solve(new_prob, Vern9(), abstol = 1e-8, reltol = 1e-8, saveat = ts)
new_sol = solve(new_prob, Vern9(), abstol = 1.0e-8, reltol = 1.0e-8, saveat = ts)

if SciMLBase.successful_retcode(new_sol)
return if SciMLBase.successful_retcode(new_sol)
mean(abs2.(reduce(hcat, get_vars(new_sol)) .- data))
else
Inf
Expand All @@ -84,8 +85,8 @@ of = OptimizationFunction{true}(loss, AutoZygote())

ps = (prob, get_vars, data, ts, set_x);

@test_call target_modules=(ModelingToolkitNeuralNets,) loss(x0, ps)
@test_opt target_modules=(ModelingToolkitNeuralNets,) loss(x0, ps)
@test_call target_modules = (ModelingToolkitNeuralNets,) loss(x0, ps)
@test_opt target_modules = (ModelingToolkitNeuralNets,) loss(x0, ps)

∇l1 = DifferentiationInterface.gradient(Base.Fix2(of, ps), AutoForwardDiff(), x0)
∇l2 = DifferentiationInterface.gradient(Base.Fix2(of, ps), AutoFiniteDiff(), x0)
Expand All @@ -94,7 +95,7 @@ ps = (prob, get_vars, data, ts, set_x);
@test all(.!isnan.(∇l1))
@test !iszero(∇l1)

@test ∇l1∇l2 rtol=1e-4
@test ∇l1∇l2 rtol = 1.0e-4
@test ∇l1 ≈ ∇l3

op = OptimizationProblem(of, x0, ps)
Expand All @@ -114,14 +115,14 @@ op = OptimizationProblem(of, x0, ps)
# false
# end

res = solve(op, Adam(1e-3), maxiters = 25_000)#, callback = plot_cb)
res = solve(op, Adam(1.0e-3), maxiters = 25_000) #, callback = plot_cb)

display(res.stats)
@test res.objective < 1.5e-4

u0, p = set_x(prob, res.u)
res_prob = remake(prob; u0, p)
res_sol = solve(res_prob, Vern9(), abstol = 1e-8, reltol = 1e-8, saveat = ts)
res_sol = solve(res_prob, Vern9(), abstol = 1.0e-8, reltol = 1.0e-8, saveat = ts)

@test SciMLBase.successful_retcode(res_sol)
@test mean(abs2.(reduce(hcat, get_vars(res_sol)) .- data)) ≈ res.objective
Expand All @@ -131,30 +132,32 @@ res_sol = solve(res_prob, Vern9(), abstol = 1e-8, reltol = 1e-8, saveat = ts)
# plot!(res_sol, idxs = [sys.x, sys.y])

function lotka_ude2()
@variables t x(t)=3.1 y(t)=1.5 pred(t)[1:2]
@parameters α=1.3 [tunable = false] δ=1.8 [tunable = false]
@variables t x(t) = 3.1 y(t) = 1.5 pred(t)[1:2]
@parameters α = 1.3 [tunable = false] δ = 1.8 [tunable = false]
chain = multi_layer_feed_forward(2, 2; width = 5, initial_scaling_factor = 1)
NN, p = SymbolicNeuralNetwork(; chain, n_input = 2, n_output = 2, rng = StableRNG(42))
Dt = ModelingToolkit.D_nounits

eqs = [pred ~ NN([x, y], p)
Dt(x) ~ α * x + pred[1]
Dt(y) ~ -δ * y + pred[2]]
eqs = [
pred ~ NN([x, y], p)
Dt(x) ~ α * x + pred[1]
Dt(y) ~ -δ * y + pred[2]
]
return System(eqs, ModelingToolkit.t_nounits, name = :lotka)
end

sys2 = mtkcompile(lotka_ude2())

prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys2, [], (0, 5.0))

sol = solve(prob, Vern9(), abstol = 1e-10, reltol = 1e-8)
sol = solve(prob, Vern9(), abstol = 1.0e-10, reltol = 1.0e-8)

@test SciMLBase.successful_retcode(sol)

set_x2 = setsym_oop(sys2, sys2.p)
ps2 = (prob, get_vars, data, ts, set_x2);
op2 = OptimizationProblem(of, x0, ps2)

res2 = solve(op2, Adam(1e-3), maxiters = 25_000)
res2 = solve(op2, Adam(1.0e-3), maxiters = 25_000)

@test res.u ≈ res2.u
2 changes: 1 addition & 1 deletion test/qa.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ using ModelingToolkitNeuralNets
using Aqua
using JET

@testset verbose=true "Code quality (Aqua.jl)" begin
@testset verbose = true "Code quality (Aqua.jl)" begin
Aqua.find_persistent_tasks_deps(ModelingToolkitNeuralNets)
Aqua.test_ambiguities(ModelingToolkitNeuralNets, recursive = false)
Aqua.test_deps_compat(ModelingToolkitNeuralNets)
Expand Down
18 changes: 10 additions & 8 deletions test/reported_issues.jl
Original file line number Diff line number Diff line change
Expand Up @@ -13,16 +13,17 @@ using OrdinaryDiffEqVerner
)

sym_nn,
θ = SymbolicNeuralNetwork(;
nn_p_name = :θ, chain, n_input = 1, n_output = 1, rng = StableRNG(42))
θ = SymbolicNeuralNetwork(;
nn_p_name = :θ, chain, n_input = 1, n_output = 1, rng = StableRNG(42)
)

# Test that scalar dispatch works (fix for issue #83)
# Previously required: sym_nn([Y], θ)[1]
# Now can use: sym_nn(Y, θ)[1]
Dt = ModelingToolkit.D_nounits
eqs_ude = [
Dt(X) ~ sym_nn(Y, θ)[1] - d * X,
Dt(Y) ~ X - d * Y
Dt(Y) ~ X - d * Y,
]

@named sys = System(eqs_ude, ModelingToolkit.t_nounits)
Expand All @@ -35,14 +36,14 @@ using OrdinaryDiffEqVerner
(0.0, 1.0)
)

sol = solve(prob, Vern9(), abstol = 1e-8, reltol = 1e-8)
sol = solve(prob, Vern9(), abstol = 1.0e-8, reltol = 1.0e-8)

@test SciMLBase.successful_retcode(sol)

# Also test that the old array syntax still works
eqs_ude_old = [
Dt(X) ~ sym_nn([Y], θ)[1] - d * X,
Dt(Y) ~ X - d * Y
Dt(Y) ~ X - d * Y,
]

@named sys_old = System(eqs_ude_old, ModelingToolkit.t_nounits)
Expand All @@ -54,7 +55,7 @@ using OrdinaryDiffEqVerner
(0.0, 1.0)
)

sol_old = solve(prob_old, Vern9(), abstol = 1e-8, reltol = 1e-8)
sol_old = solve(prob_old, Vern9(), abstol = 1.0e-8, reltol = 1.0e-8)

@test SciMLBase.successful_retcode(sol_old)

Expand All @@ -80,8 +81,9 @@ end
nn_name = :custom_nn_name
nn_p_name = :custom_nn_p_name
NN, NN_p = SymbolicNeuralNetwork(;
chain, n_input = 1, n_output = 1, rng, nn_name, nn_p_name)
chain, n_input = 1, n_output = 1, rng, nn_name, nn_p_name
)

@test ModelingToolkit.getname(NN)==nn_name broken=true # :nn_name # Should be :custom_nn_name
@test ModelingToolkit.getname(NN) == nn_name broken = true # :nn_name # Should be :custom_nn_name
@test ModelingToolkit.getname(NN_p) == nn_p_name
end
Loading
Loading