Parallel plotting in Julia with Dagger.jl

dagger_plotting.jl

using Distributed

procs = Distributed.addprocs(16)

@everywhere begin
    using Printf
    using Dagger
    using PyCall
    using PyPlot
    using NCDatasets

    cmocean = pyimport("cmocean")

    PyPlot.ioff()  # Turn off interactive plotting

    ds = NCDataset("dry_convection.nc")

    frames = length(ds["time"])
    x, y, z = ds["xC"], ds["yC"], ds["zC"]
    Nx, Ny, Nz = length(ds["xC"]), length(ds["yC"]), length(ds["zC"])
end

@everywhere function plot_frame(frame)
    @info "[worker $(myid())] Frame $frame/$frames..."
    
    ρe = ds["ρe"][:, Int(Ny/2), :, frame] |> Array
    ρe₀ = ds["ρe₀"][:, Int(Ny/2), :] |> Array
    ρe′ = transpose(ρe .- ρe₀)

    filepath = @sprintf("dry_convection_%04d.png", frame)
    PyPlot.contourf(x, z, ρe′, vmin=0, vmax=400, cmap=cmocean.cm.ice, levels=20)
    PyPlot.savefig(filepath, dpi=200)
    PyPlot.close("all")

    return filepath
end

ctx = Dagger.Context()
addprocs!(ctx, procs)
thunk = Dagger.delayed(vcat)((Dagger.delayed(plot_frame)(n) for n in 1:frames)...)
job = @async collect(ctx, thunk)
@time result = fetch(job)

Distributed.workers() |> Distributed.rmprocs