EnsembleProblem progress bar

[ograsdijk]

Currently EnsembleProblem does not have progress bars, i.e. using progress = true within solve does not result in a progress bar. To get around this I currently use the following to get access to a progress bar:

using ProgressMeter
using Distributed

addprocs(6)
@everywhere using DifferentialEquations

n_trajectories = 100
progress = Progress(n_trajectories)
const channel = RemoteChannel(()->Channel{Int}(n_trajectories))

@everywhere function prob_func(prob,i,repeat)
    remake(prob,u0=rand()*prob.u0)
end

@everywhere function f(u,p,t)
    sleep(0.01)
    1.01u
end

prob = ODEProblem(f,0.5,(0.0,1.0))

@sync begin
    @async begin
	    tasksdone = 0
        while tasksdone < n_trajectories
	        tasksdone += take!(channel)
	        update!(progress, tasksdone)
        end
    end
    @async begin
        @distributed for i =1:n_trajectories
            sol = solve(prob_func(prob, i, false), Tsit5())
            put!(channel, 1)
        end
    end
end

It would be nice if there progress bars were also functional for EnsembleProblem. Even something like the following:

using ProgressMeter
using Distributed

addprocs(6)
@everywhere using DifferentialEquations

n_trajectories = 100
progress = Progress(n_trajectories)
const channel = RemoteChannel(()->Channel{Int}(n_trajectories))

@everywhere function prob_func(prob,i,repeat)
    remake(prob,u0=rand()*prob.u0)
end

@everywhere function output_func(sol, i)
    put!(channel, 1)
    sol, false
end

@everywhere function f(u,p,t)
    sleep(0.01)
    1.01u
end

prob = ODEProblem(f,0.5,(0.0,1.0))
ensembleproblem = EnsembleProblem(prob, prob_func = prob_func, output_func = output_func)

@sync begin
    @async begin
        tasksdone = 0
        while tasksdone < n_trajectories
	        tasksdone += take!(channel)
	        update!(progress, tasksdone)
        end
    end
    @async begin
        sim = solve(ensembleproblem, Tsit5(), EnsembleDistributed(), trajectories=n_trajectories)
    end
end

would be nice but I can't pass the RemoteChannel to the EnsembleProblem.

[ograsdijk]

If you add @everywhere const channel = $channel after defining the channel on the main process the second example with EnsembleProblem also has a functioning progress bar.

[JordiBolibar]

Is there any progress on an integrated solution within DifferentialEquations.jl to this problem? I'm also very interested in making this work with DiffEqFlux.jl, in order to train UDEs in parallel and being able to track the progress.

[ChrisRackauckas]

There isn't one right now, and I don't know the logger systems well enough myself to work on it. @pfitzseb probably knows what the right way to do this is?

[JordiBolibar]

I have tried a workaround using pmap and manually implementing the ensemble for each batch, but using @showprogress with that only works with DifferentialEquations.jl. When using it with AD (e.g. with Zygote) with DiffEqFlux.jl it crashes during the backpropagation.