Skip to content

Commit e3b4b9e

Browse files
author
SciML Bot
committed
Apply JuliaFormatter to fix code formatting
- Applied JuliaFormatter with SciML style guide - Formatted 4 files 🤖 Generated by OrgMaintenanceScripts.jl
1 parent bc48e5f commit e3b4b9e

File tree

4 files changed

+80
-65
lines changed

4 files changed

+80
-65
lines changed

CHANGELOG.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,5 @@
22

33
## Breaking changes
44

5-
- The `NeuralNetworkBlock` no longer uses `RealInputArray` & `RealOutputArray`,
6-
the ports are now `inputs` and `outputs` and they are normal vector variables.
5+
- The `NeuralNetworkBlock` no longer uses `RealInputArray` & `RealOutputArray`,
6+
the ports are now `inputs` and `outputs` and they are normal vector variables.

docs/src/nnblock.md

Lines changed: 73 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@
22

33
[`ModelingToolkitNeuralNets`](https://github.com/SciML/ModelingToolkitNeuralNets.jl) provides 2 main interfaces for representing neural networks symbolically:
44

5-
* The [`NeuralNetworkBlock`](@ref), which represents the neural network as a block component
6-
* The [`SymbolicNeuralNetwork`](@ref), which represents the neural network via callable parameters
5+
- The [`NeuralNetworkBlock`](@ref), which represents the neural network as a block component
6+
- The [`SymbolicNeuralNetwork`](@ref), which represents the neural network via callable parameters
77

88
This tutorial will introduce the [`NeuralNetworkBlock`](@ref). This representation is useful in the context of hierarchical acausal component-based model.
99

@@ -46,13 +46,13 @@ input_f(t) = (1+sin(0.005 * t^2))/2
4646
C2 = 15
4747
end
4848
@components begin
49-
input = Blocks.TimeVaryingFunction(f=input_f)
50-
source = PrescribedHeatFlow(T_ref=373.15)
51-
plate = HeatCapacitor(C=C1, T=273.15)
52-
pot = HeatCapacitor(C=C2, T=273.15)
53-
conduction = ThermalConductor(G=1)
54-
air = ThermalConductor(G=0.1)
55-
env = FixedTemperature(T=293.15)
49+
input = Blocks.TimeVaryingFunction(f = input_f)
50+
source = PrescribedHeatFlow(T_ref = 373.15)
51+
plate = HeatCapacitor(C = C1, T = 273.15)
52+
pot = HeatCapacitor(C = C2, T = 273.15)
53+
conduction = ThermalConductor(G = 1)
54+
air = ThermalConductor(G = 0.1)
55+
env = FixedTemperature(T = 293.15)
5656
Tsensor = TemperatureSensor()
5757
end
5858
@equations begin
@@ -70,11 +70,11 @@ end
7070
C2 = 15
7171
end
7272
@components begin
73-
input = Blocks.TimeVaryingFunction(f=input_f)
74-
source = PrescribedHeatFlow(T_ref=373.15)
75-
pot = HeatCapacitor(C=C2, T=273.15)
76-
air = ThermalConductor(G=0.1)
77-
env = FixedTemperature(T=293.15)
73+
input = Blocks.TimeVaryingFunction(f = input_f)
74+
source = PrescribedHeatFlow(T_ref = 373.15)
75+
pot = HeatCapacitor(C = C2, T = 273.15)
76+
air = ThermalConductor(G = 0.1)
77+
env = FixedTemperature(T = 293.15)
7878
Tsensor = TemperatureSensor()
7979
end
8080
@equations begin
@@ -91,11 +91,11 @@ end
9191
## solve and plot the temperature of the pot in the 2 systems
9292
9393
prob1 = ODEProblem(sys1, Pair[], (0, 100.0))
94-
sol1 = solve(prob1, Tsit5(), reltol=1e-6)
94+
sol1 = solve(prob1, Tsit5(), reltol = 1e-6)
9595
prob2 = ODEProblem(sys2, Pair[], (0, 100.0))
96-
sol2 = solve(prob2, Tsit5(), reltol=1e-6)
97-
plot(sol1, idxs=sys1.pot.T, label="pot.T in original system")
98-
plot!(sol2, idxs=sys1.pot.T, label="pot.T in simplified system")
96+
sol2 = solve(prob2, Tsit5(), reltol = 1e-6)
97+
plot(sol1, idxs = sys1.pot.T, label = "pot.T in original system")
98+
plot!(sol2, idxs = sys1.pot.T, label = "pot.T in simplified system")
9999
```
100100

101101
If we take a closer look at the 2 models, the original system has 2 unknowns,
@@ -105,6 +105,7 @@ unknowns(sys1)
105105
```
106106

107107
while the simplified system only has 1 unknown
108+
108109
```@example potplate
109110
unknowns(sys2)
110111
```
@@ -127,12 +128,13 @@ always output positive numbers for positive inputs, so this also makes physical
127128
begin
128129
n_input = 2
129130
n_output = 1
130-
chain = multi_layer_feed_forward(; n_input, n_output, depth=1, width=4, activation=Lux.swish)
131+
chain = multi_layer_feed_forward(;
132+
n_input, n_output, depth = 1, width = 4, activation = Lux.swish)
131133
end
132134
@components begin
133135
port_a = HeatPort()
134136
port_b = HeatPort()
135-
nn = NeuralNetworkBlock(; n_input, n_output, chain, rng=StableRNG(1337))
137+
nn = NeuralNetworkBlock(; n_input, n_output, chain, rng = StableRNG(1337))
136138
end
137139
@parameters begin
138140
T0 = 273.15
@@ -160,11 +162,11 @@ end
160162
C2 = 15
161163
end
162164
@components begin
163-
input = Blocks.TimeVaryingFunction(f=input_f)
164-
source = PrescribedHeatFlow(T_ref=373.15)
165-
pot = HeatCapacitor(C=C2, T=273.15)
166-
air = ThermalConductor(G=0.1)
167-
env = FixedTemperature(T=293.15)
165+
input = Blocks.TimeVaryingFunction(f = input_f)
166+
source = PrescribedHeatFlow(T_ref = 373.15)
167+
pot = HeatCapacitor(C = C2, T = 273.15)
168+
air = ThermalConductor(G = 0.1)
169+
env = FixedTemperature(T = 293.15)
168170
Tsensor = TemperatureSensor()
169171
thermal_nn = ThermalNN()
170172
end
@@ -184,7 +186,7 @@ sys3 = mtkcompile(model)
184186
# Let's check that we can successfully simulate the system in the
185187
# initial state
186188
prob3 = ODEProblem(sys3, Pair[], (0, 100.0))
187-
sol3 = solve(prob3, Tsit5(), abstol=1e-6, reltol=1e-6)
189+
sol3 = solve(prob3, Tsit5(), abstol = 1e-6, reltol = 1e-6)
188190
@assert SciMLBase.successful_retcode(sol3)
189191
```
190192

@@ -209,23 +211,28 @@ x0 = prob3.ps[tp]
209211
210212
oop_update = setsym_oop(prob3, tp);
211213
212-
plot_cb = (opt_state, loss) -> begin
214+
plot_cb = (opt_state,
215+
loss) -> begin
213216
opt_state.iter % 1000 ≠ 0 && return false
214217
@info "step $(opt_state.iter), loss: $loss"
215218
216219
(new_u0, new_p) = oop_update(prob3, opt_state.u)
217-
new_prob = remake(prob3, u0=new_u0, p=new_p)
218-
sol = solve(new_prob, Tsit5(), abstol=1e-8, reltol=1e-8)
219-
220-
plt = plot(sol, layout=(2,3), idxs=[
221-
sys3.thermal_nn.nn.inputs[1], sys3.thermal_nn.x,
222-
sys3.thermal_nn.nn.outputs[1], sys3.thermal_nn.port_b.T,
223-
sys3.pot.T, sys3.pot.port.Q_flow],
224-
size=(950,800))
225-
plot!(plt, sol1, idxs=[
226-
(sys1.conduction.port_a.T-273.15)/10, sys1.conduction.port_a.T,
227-
sys1.conduction.port_a.Q_flow, sys1.conduction.port_b.T,
228-
sys1.pot.T, sys1.pot.port.Q_flow])
220+
new_prob = remake(prob3, u0 = new_u0, p = new_p)
221+
sol = solve(new_prob, Tsit5(), abstol = 1e-8, reltol = 1e-8)
222+
223+
plt = plot(sol,
224+
layout = (2, 3),
225+
idxs = [
226+
sys3.thermal_nn.nn.inputs[1], sys3.thermal_nn.x,
227+
sys3.thermal_nn.nn.outputs[1], sys3.thermal_nn.port_b.T,
228+
sys3.pot.T, sys3.pot.port.Q_flow],
229+
size = (950, 800))
230+
plot!(plt,
231+
sol1,
232+
idxs = [
233+
(sys1.conduction.port_a.T-273.15)/10, sys1.conduction.port_a.T,
234+
sys1.conduction.port_a.Q_flow, sys1.conduction.port_b.T,
235+
sys1.pot.T, sys1.pot.port.Q_flow])
229236
display(plt)
230237
false
231238
end
@@ -236,7 +243,8 @@ function cost(x, opt_ps)
236243
u0, p = oop_update(prob, x)
237244
new_prob = remake(prob; u0, p)
238245
239-
new_sol = solve(new_prob, Tsit5(), saveat=ts, abstol=1e-8, reltol=1e-8, verbose=false, sensealg=GaussAdjoint())
246+
new_sol = solve(new_prob, Tsit5(), saveat = ts, abstol = 1e-8,
247+
reltol = 1e-8, verbose = false, sensealg = GaussAdjoint())
240248
241249
!SciMLBase.successful_retcode(new_sol) && return Inf
242250
@@ -249,35 +257,41 @@ data = sol1[sys1.pot.T]
249257
get_T = getsym(prob3, sys3.pot.T)
250258
opt_ps = (prob3, oop_update, data, sol1.t, get_T);
251259
252-
op = OptimizationProblem(of, x0, opt_ps,)
260+
op = OptimizationProblem(of, x0, opt_ps)
253261
254-
res = solve(op, Adam(); maxiters=10_000, callback=plot_cb)
262+
res = solve(op, Adam(); maxiters = 10_000, callback = plot_cb)
255263
op2 = OptimizationProblem(of, res.u, opt_ps)
256-
res2 = solve(op2, LBFGS(linesearch=BackTracking()); maxiters=2000, callback=plot_cb)
264+
res2 = solve(op2, LBFGS(linesearch = BackTracking()); maxiters = 2000, callback = plot_cb)
257265
258266
(new_u0, new_p) = oop_update(prob3, res2.u)
259-
new_prob1 = remake(prob3, u0=new_u0, p=new_p)
260-
new_sol1 = solve(new_prob1, Tsit5(), abstol=1e-6, reltol=1e-6)
261-
262-
plt = plot(new_sol1, layout=(2,3), idxs=[
263-
sys3.thermal_nn.nn.inputs[1], sys3.thermal_nn.x,
264-
sys3.thermal_nn.nn.outputs[1], sys3.thermal_nn.port_b.T,
265-
sys3.pot.T, sys3.pot.port.Q_flow],
266-
size=(950,800))
267-
plot!(plt, sol1, idxs=[
268-
(sys1.conduction.port_a.T-273.15)/10, sys1.conduction.port_a.T,
269-
sys1.conduction.port_a.Q_flow, sys1.conduction.port_b.T,
270-
sys1.pot.T, sys1.pot.port.Q_flow], ls=:dash)
267+
new_prob1 = remake(prob3, u0 = new_u0, p = new_p)
268+
new_sol1 = solve(new_prob1, Tsit5(), abstol = 1e-6, reltol = 1e-6)
269+
270+
plt = plot(new_sol1,
271+
layout = (2, 3),
272+
idxs = [
273+
sys3.thermal_nn.nn.inputs[1], sys3.thermal_nn.x,
274+
sys3.thermal_nn.nn.outputs[1], sys3.thermal_nn.port_b.T,
275+
sys3.pot.T, sys3.pot.port.Q_flow],
276+
size = (950, 800))
277+
plot!(plt,
278+
sol1,
279+
idxs = [
280+
(sys1.conduction.port_a.T-273.15)/10, sys1.conduction.port_a.T,
281+
sys1.conduction.port_a.Q_flow, sys1.conduction.port_b.T,
282+
sys1.pot.T, sys1.pot.port.Q_flow],
283+
ls = :dash)
271284
```
272285

273286
As we can see from the final plot, the neural network fits very well and not only the training data fits, but also the rest of the
274287
predictions of the system match the original system. Let us also compare against the predictions of the incomplete system:
275288

276289
```@example potplate
277-
plot(sol1, label=["original sys: pot T" "original sys: plate T"], lw=3)
278-
plot!(sol3; idxs=[sys3.pot.T], label="untrained UDE", lw=2.5)
279-
plot!(sol2; idxs=[sys2.pot.T], label="incomplete sys: pot T", lw=2.5)
280-
plot!(new_sol1; idxs=[sys3.pot.T, sys3.thermal_nn.x], label="trained UDE", ls=:dash, lw=2.5)
290+
plot(sol1, label = ["original sys: pot T" "original sys: plate T"], lw = 3)
291+
plot!(sol3; idxs = [sys3.pot.T], label = "untrained UDE", lw = 2.5)
292+
plot!(sol2; idxs = [sys2.pot.T], label = "incomplete sys: pot T", lw = 2.5)
293+
plot!(new_sol1; idxs = [sys3.pot.T, sys3.thermal_nn.x],
294+
label = "trained UDE", ls = :dash, lw = 2.5)
281295
```
282296

283297
Now that our neural network is trained, we can go a step further and use [`SymbolicRegression.jl`](https://github.com/MilesCranmer/SymbolicRegression.jl) to find

src/utils.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,9 @@ function multi_layer_feed_forward(; n_input, n_output, width::Int = 4,
1212
Lux.Dense(n_input, width, activation; use_bias),
1313
[Lux.Dense(width, width, activation; use_bias) for _ in 1:(depth)]...,
1414
Lux.Dense(width, n_output;
15-
init_weight = (rng, a...) -> initial_scaling_factor *
16-
Lux.kaiming_uniform(rng, a...), use_bias)
15+
init_weight = (
16+
rng, a...) -> initial_scaling_factor *
17+
Lux.kaiming_uniform(rng, a...), use_bias)
1718
)
1819
end
1920

test/lotka_volterra.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ using Statistics
1818

1919
function lotka_ude()
2020
@variables t x(t)=3.1 y(t)=1.5
21-
@parameters α=1.3 [tunable = false] δ=1.8 [tunable = false]
21+
@parameters α=1.3 [tunable=false] δ=1.8 [tunable=false]
2222
Dt = ModelingToolkit.D_nounits
2323

2424
chain = multi_layer_feed_forward(2, 2)
@@ -125,7 +125,7 @@ res_sol = solve(res_prob, Vern9())
125125

126126
function lotka_ude2()
127127
@variables t x(t)=3.1 y(t)=1.5 pred(t)[1:2]
128-
@parameters α=1.3 [tunable = false] δ=1.8 [tunable = false]
128+
@parameters α=1.3 [tunable=false] δ=1.8 [tunable=false]
129129
chain = multi_layer_feed_forward(2, 2)
130130
NN, p = SymbolicNeuralNetwork(; chain, n_input = 2, n_output = 2, rng = StableRNG(42))
131131
Dt = ModelingToolkit.D_nounits

0 commit comments

Comments
 (0)