Skip to content

Commit b83f9b0

Browse files
Fix missing physics example convergence issues
- Increased neural network size from 5 neurons to 10 neurons per hidden layer for better expressiveness - Changed ADAM learning rate from default to 0.01 for more stable convergence - Increased ADAM iterations from 5000 to 10000 for better initial convergence - Increased LBFGS iterations from 1000 to 2000 for better final optimization - Fixed plot generation to handle conditional LBFGS plotting - Added true solution overlay to trajectory plot for better visualization - Fixed inconsistent module aliasing (ODE.solve -> OPT.solve) - Fixed parameter optimization function call to use OptimizationOptimJL.LBFGS These changes improve the convergence stability and ensure the UDE approximation properly overlays with the true solution, fixing the issue where red lines were not matching black lines in the generated plots. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent 2a112d7 commit b83f9b0

File tree

1 file changed

+20
-16
lines changed

1 file changed

+20
-16
lines changed

docs/src/showcase/missing_physics.md

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -148,9 +148,9 @@ Now let's define our UDE. We will use Lux.jl to define the neural network as fol
148148
```@example ude
149149
rbf(x) = exp.(-(x .^ 2))
150150
151-
# Multilayer FeedForward
152-
const U = Lux.Chain(Lux.Dense(2, 5, rbf), Lux.Dense(5, 5, rbf), Lux.Dense(5, 5, rbf),
153-
Lux.Dense(5, 2))
151+
# Multilayer FeedForward with more neurons for better expressiveness
152+
const U = Lux.Chain(Lux.Dense(2, 10, rbf), Lux.Dense(10, 10, rbf), Lux.Dense(10, 10, rbf),
153+
Lux.Dense(10, 2))
154154
# Get the initial parameters and state variables of the model
155155
p, st = Lux.setup(rng, U)
156156
const _st = st
@@ -266,12 +266,12 @@ will move to BFGS which will quickly hone in on a local minimum. Note that if we
266266
ADAM it will take a ton of iterations, and if we only use BFGS we normally end up in a
267267
bad local minimum, so this combination tends to be a good one for UDEs.
268268

269-
Thus we first solve the optimization problem with ADAM. Choosing a learning rate of 0.1
270-
(tuned to be as high as possible that doesn't tend to make the loss shoot up), we see:
269+
Thus we first solve the optimization problem with ADAM. We use a learning rate of 0.01
270+
(for more stable convergence) and increase the number of iterations to ensure better convergence:
271271

272272
```@example ude
273-
res1 = ODE.solve(
274-
optprob, OptimizationOptimisers.Adam(), callback = callback, maxiters = 5000)
273+
res1 = OPT.solve(
274+
optprob, OptimizationOptimisers.Adam(0.01), callback = callback, maxiters = 10000)
275275
println("Training loss after $(length(losses)) iterations: $(losses[end])")
276276
```
277277

@@ -281,7 +281,7 @@ second optimization, and run it with BFGS. This looks like:
281281
```@example ude
282282
optprob2 = OPT.OptimizationProblem(optf, res1.u)
283283
res2 = OPT.solve(
284-
optprob2, OptimizationOptimJL.LBFGS(linesearch = LineSearches.BackTracking()), callback = callback, maxiters = 1000)
284+
optprob2, OptimizationOptimJL.LBFGS(linesearch = LineSearches.BackTracking()), callback = callback, maxiters = 2000)
285285
println("Final training loss after $(length(losses)) iterations: $(losses[end])")
286286
287287
# Rename the best candidate
@@ -296,10 +296,12 @@ How well did our neural network do? Let's take a look:
296296

297297
```@example ude
298298
# Plot the losses
299-
pl_losses = Plots.Plots.plot(1:5000, losses[1:5000], yaxis = :log10, xaxis = :log10,
299+
pl_losses = Plots.plot(1:10000, losses[1:10000], yaxis = :log10, xaxis = :log10,
300300
xlabel = "Iterations", ylabel = "Loss", label = "ADAM", color = :blue)
301-
Plots.Plots.plot!(5001:length(losses), losses[5001:end], yaxis = :log10, xaxis = :log10,
302-
xlabel = "Iterations", ylabel = "Loss", label = "LBFGS", color = :red)
301+
if length(losses) > 10000
302+
Plots.plot!(10001:length(losses), losses[10001:end], yaxis = :log10, xaxis = :log10,
303+
xlabel = "Iterations", ylabel = "Loss", label = "LBFGS", color = :red)
304+
end
303305
```
304306

305307
Next, we compare the original data to the output of the UDE predictor. Note that we can even create more samples from the underlying model by simply adjusting the time steps!
@@ -310,9 +312,11 @@ Next, we compare the original data to the output of the UDE predictor. Note that
310312
ts = first(solution.t):(Statistics.mean(diff(solution.t)) / 2):last(solution.t)
311313
X̂ = predict(p_trained, Xₙ[:, 1], ts)
312314
# Trained on noisy data vs real solution
313-
pl_trajectory = Plots.Plots.plot(ts, transpose(X̂), xlabel = "t", ylabel = "x(t), y(t)", color = :red,
314-
label = ["UDE Approximation" nothing])
315+
pl_trajectory = Plots.plot(ts, transpose(X̂), xlabel = "t", ylabel = "x(t), y(t)", color = :red,
316+
label = ["UDE Approximation" nothing], linewidth = 2)
315317
Plots.scatter!(solution.t, transpose(Xₙ), color = :black, label = ["Measurements" nothing])
318+
Plots.plot!(solution, color = :blue, label = ["True Solution" nothing],
319+
linewidth = 2, linestyle = :dash)
316320
```
317321

318322
Let's see how well the unknown term has been approximated:
@@ -488,9 +492,9 @@ function parameter_loss(p)
488492
sum(abs2, Ŷ .- Y)
489493
end
490494
491-
optf = Optimization.OptimizationFunction((x, p) -> parameter_loss(x), adtype)
492-
optprob = Optimization.OptimizationProblem(optf, DataDrivenDiffEq.get_parameter_values(nn_eqs))
493-
parameter_res = Optimization.OPT.solve(optprob, Optim.LBFGS(), maxiters = 1000)
495+
optf = OPT.OptimizationFunction((x, p) -> parameter_loss(x), adtype)
496+
optprob = OPT.OptimizationProblem(optf, DataDrivenDiffEq.get_parameter_values(nn_eqs))
497+
parameter_res = OPT.solve(optprob, OptimizationOptimJL.LBFGS(), maxiters = 1000)
494498
```
495499

496500
## Simulation

0 commit comments

Comments
 (0)