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

Specific Domains Visualization
, ,
1

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
ad_type = Optimization.AutoZygote()
opt_func = OptimizationFunction((x,p)->loss(x), ad_type)

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
2

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
3

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!

4

https://github.com/julia-vscode/julia-vscode/pull/2940

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