# When plotting Optimization.jl callbacks with Plots.jl, use just one figure in VS Code plots paned

In the callback function below `cb = function(θ,l,pred)` … how can I have only the current prediction appear on a plot with each iteration to illustrate the progress of fitting?
I have tried putting `fig = plot()` inside the callback, but that opens a new figure for each iteration, which is also what I’d like to avoid.

``````using ModelingToolkit, MethodOfLines, OrdinaryDiffEq, DomainSets
using Plots
using DifferentialEquations, Optimization, Zygote, OptimizationOptimJL
using SciMLSensitivity
using Statistics
gr()
@parameters t x
p = @parameters Dn, Dp
@variables u(..) v(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2

eqs  = [Dt(u(t,x)) ~ Dn * Dxx(u(t,x)) + u(t,x)*v(t,x),
Dt(v(t,x)) ~ Dp * Dxx(v(t,x)) - u(t,x)*v(t,x)]
bcs = [u(0,x) ~ sin(pi*x/2),
v(0,x) ~ sin(pi*x/2),
u(t,0) ~ 0.0, Dx(u(t,1)) ~ 0.0,
v(t,0) ~ 0.0, Dx(v(t,1)) ~ 0.0]

domains = [t ∈ Interval(0.0,1.0),
x ∈ Interval(0.0,1.0)]

@named pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)],[Dn=>0.5, Dp=>2.0])

discretization = MOLFiniteDifference([x=>0.1],t)

pde_prob = discretize(pdesys,discretization)

sol = solve(pde_prob, Tsit5(), saveat=1.0, abstol=1e-16)
meas = sol.u[2] .+ 0.01*randn(18)

idcs = Int.(ModelingToolkit.varmap_to_vars([p[1]=>1, p[2]=>2], p))

function loss(θ)
new_pde_prob = remake(pde_prob, p=θ[idcs])
sol = solve(new_pde_prob, Tsit5(), saveat=1.0, abstol=1e-16)
return sqrt(mean( (sol.u[2] .- meas).^2 )), Array(sol)
end

fig = plot(legend=false)
cb = function(θ,l,pred)
println("loss = ", l, "; θ = ", θ)
plot!(fig, pred[:,2])
plot!(fig,meas)
display(fig)
false
end

guess = [1.0,1.0]

opt_prob = OptimizationProblem(opt_func, guess, lb=[0.1,0.1],ub=[3.0,3.0])

res = Optimization.solve(opt_prob, ParticleSwarm(),callback=cb)
``````
1 Like

The problem is the line `display(fig)` in the definition of `cb`. Here is a much more minimal example that yields the same misbehavior:

``````using Plots

fig = plot()

cb = function()
plot!(fig, rand(2), rand(2))
display(fig)                         # !
end

function calculate()
for _ in 1:20
cb()
end
end

res = calculate()
``````

The solution is to remove the indicated line. It also helps to wrap everything in a `main()` guard and return `fig` at the end, if the figure is all you want:

``````using Plots

function main()
fig = plot()

function cb()    # More idiomatic than `cb = function()` IMO
plot!(fig, rand(2), rand(2))
end

function calculate()
for _ in 1:20
cb()
end
end

calculate()      # Returns `nothing` anyway

return fig
end

main()
``````
1 Like

Thanks! I edited this part of the post:

how can I have only the current prediction appear on a plot with each iteration to illustrate the progress of fitting?

I’m looking for a way to vizualize each iteration live rather generating a plot of each iteration all at once at the end. That’s why I added `display(fig)` inside `cb`. The goal is to:

• Plot each iteration on a single plot (without generating new figures in the plot pane)
• Plot only the most recent iteration (removing or overwriting the previous curve)

And yeah, a simple loop is definitely a more minimal example!

seems like the right solution for this. Please chime in on that issue if you have any suggestions for the API/requirements/comments!

2 Likes