Should we add a MovieMaker to Oceananigans?

[tomchor]

This issue was originally posted in tomchor/Oceanostics.jl#186.

The basic idea is that, in the early stages of assembling a simulation, we tend to spend a lot of time (relatively) coding the output writing and plotting/animating parts of the simulation. Most of that code, however, is pretty similar across simulations: setting up a writer with a couple of variables, then reading it, creating a figure, plotting heatmaps, and writing the animation loop. Lots of copy-paste. And then if we want to add other variables to the plot, several lines have to be added/modified to accommodate for that.

I think this process can be automated. As @glwagner suggested, perhaps the best way would be to avoid writing altogether and create a plotter which updates a figure periodically as the simulation is running.

Here's a quick snippet of a proof of concept I tried:

struct Plotter
    fig
    io
    func
end

function Plotter(func; fig=nothing, io=nothing, kwargs...)
    if fig isa Nothing
        fig = Figure()
        ax1 = Axis(fig[1, 1])
        ax2 = Axis(fig[1, 2])
    end
    if io isa Nothing
        io = VideoStream(fig, kwargs...)
    end
    return Plotter(fig, io, func)
end

(plotter::Plotter)(simulation) = plotter.func(simulation, plotter.fig, plotter.io)

function update_plot(sim, fig, io)
    heatmap!(fig[1, 1], ω; colormap = :balance, colorrange = (-2, 2))
    heatmap!(fig[1, 2], s; colormap = :balance, colorrange = (-2, 2))
    recordframe!(io)
end

plotter = Plotter(update_plot)
add_callback!(simulation, plotter, TimeInterval(0.6))

Applied to the 2D turbulence example in the docs this produces:

test.mp4

Is there interest in implementing something like this here? If so I'll close the issue on Oceanostics and open a PR here.

CC @glwagner @jbisits

[glwagner]

Before proof of concept, can you propose the user API that the proof of concept achieves? I can't quite see that clearly from the implementation.

[tomchor]

The way I thought this could work is to start with a flexible interface where users could write their own plotting function, like:

fig = Figure()
ax1 = Axis(fig[1, 1])
ax2 = Axis(fig[1, 2])

function plotting_function(sim, fig, io)
    heatmap!(fig[1, 1], u)
    heatmap!(fig[1, 2], s)
    recordframe!(io)
end

plotter = Plotter(plotting_function, figure=fig, filename="test.mp4")
add_callback!(simulation, plotter, TimeInterval(1))

# Do other stuff
run!(simulation)
# Animation plotted!

But eventually we might want to also pre-code some common plots, like panel grids, where users can just pass a list of variables like:

plotter = PanelGridPlotter((u, v, w, b), name="panels.mp4")
add_callback!(simulation, plotter, TimeInterval(1))
run!(simulation)

[glwagner]

Ok, that's interesting. In your example, I'm not sure how fig and io can be passed into plotting_function, because Plotter doesn't have that information. Should we design it so that users have to create the figures, or so that the figures are created automagically? I think one could be excused for simply implementing the panel plotter directly, since this feature seems oriented towards quick and dirty plots as is?

One could mirror the syntax of the output writer with filename and indices (the latter being necessary for 3D situations, since we don't support auto-plotting of 3D data yet)

[tomchor]

Ok, that's interesting. In your example, I'm not sure how fig and io can be passed into plotting_function, because Plotter doesn't have that information.

Ah, thanks for catching that. I meant to pass fig to Plotter in the first example since I created fig manually there. In the second example I think it'd be pretty easy to just create fig when instantiating Plotter. io would probably always be created when instantiating Plotter since most people aren't familiar with the call to VideoStream, but in the end it's pretty simple to use it.

Should we design it so that users have to create the figures, or so that the figures are created automagically? I think one could be excused for simply implementing the panel plotter directly, since this feature seems oriented towards quick and dirty plots as is?

I think both are fair. In order to make it flexible I think we'd need to allow for users to create a figure themselves, but the quick-and-dirty approach of creating everything when Plotter is instantiated would be super useful.

One could mirror the syntax of the output writer with filename and indices (the latter being necessary for 3D situations, since we don't support auto-plotting of 3D data yet)

And yes, that's exactly what I was thinking. The proof of concept above only works because it's a 2D model. In fact, here's the full code if you're interested in playing around with it. (Although I think we should make 3D plots work as well at some point!)

[glwagner]

Ah sorry, I still don't really grasp what you're proposing.

Not sure if this is different from what you're saying, but here is one idea:

fields_to_plot = merge(model.velocities, model.tracers)
movie_maker = MovieMaker(fields_to_plot, indices=(:, 1, :), filename="test.mp4")

then movie_maker can be used as a callback. The callback schedule determines the interval between frames. Additional kwargs to MovieMaker could be used to pre-make the axis, maybe something like

fig = Figure()
axes = (u = Axis(fig[1, 1]), T = Axis(fig[1, 2]))
fields = (u=model.velocities.u, T=model.tracers.T) # the keys of axes, fields have to match
movie_maker = MovieMaker(fields, indices=(:, 1, :), filename="test.mp4", figure=fig, axes=axes)

I think for an interface with a custom plotting function one only needs two arguments, since the figure and axes have to be captured by a callback?

# create figure, axes etc
function makeplot(sim, io)
    heatmap!(axu, u)
    heatmap!(axs, s)
    recordframe!(io)
end

movie_maker = MovieMaker(makeplot, filename="test.mp4")

I think there isn't really a point to passing fig, axes between MovieMaker and also into makeplot since either way, it has to be created before?

Also I was wondering if there's scope for plotting a single snapshot... but that might be easier so not necessary to think about here.

[tomchor]

Ah sorry, I still don't really grasp what you're proposing.

Not sure if this is different from what you're saying, but here is one idea:

fields_to_plot = merge(model.velocities, model.tracers)
movie_maker = MovieMaker(fields_to_plot, indices=(:, 1, :), filename="test.mp4")

then movie_maker can be used as a callback. The callback schedule determines the interval between frames.

Yes, that's pretty much what I was proposing. Sorry I wasn't able to be more clear.

Additional kwargs to MovieMaker could be used to pre-make the axis, maybe something like

fig = Figure()
axes = (u = Axis(fig[1, 1]), T = Axis(fig[1, 2]))
fields = (u=model.velocities.u, T=model.tracers.T) # the keys of axes, fields have to match
movie_maker = MovieMaker(fields, indices=(:, 1, :), filename="test.mp4", figure=fig, axes=axes)

I think for an interface with a custom plotting function one only needs two arguments, since the figure and axes have to be captured by a callback?

I'm not sure about this one. I gotta think a bit more about this. As a rule here though, I think it's good if we can make things as simple as possible on the user side, since the main use for this feature would be more of a quick check on the simulations.

I do think we need to make room for kwargs to VideoStream (which creates io), since that'll dictate some things that may be important like the frame rate, video compression, etc.

Also I was wondering if there's scope for plotting a single snapshot... but that might be easier so not necessary to think about here.

I guess we could set it up in a way that, if the user chooses, instead of recordframe!ing each snapshot into a video, we can just plot each snapshot (dictated by schedule) into a separate figure (which some sort of suffix to indicate when it was plotted) and leave it at that. Although idk how useful that would be since you can always print a frame from a video after the fact if you truly just want a figure.

[glwagner]

well this seems easy enough so I support trying in the Makie extension!

On naming, I think if one creates "plots" then it should be a Plotter but if we are greating animations then the name should reflect that.

[tomchor]

I agree with the naming :)

Also, note that this technically doesn't need the Makie extension to work. A user could define a plotting function like:

function makeplot(sim, io)
    heatmap!(ax1, Array(u[1, :, :]))
    recordframe!(io)
end

Given that, should we still include this with the Makie extension?

Although on the other hand, it does need Makie to be imported and used, which will activate load the extension regardless if I understand correctly?

[glwagner]

You'd need the Makie extension to auto-create io right?

[tomchor]

You'd need the Makie extension to auto-create io right?

I need to be using Makie to do it. But I'm not sure if that means I need the extension. Maybe this confusion stems from me not knowing the terms to be using here exactly.

[tomchor]

@glwagner if that's okay with you, I'll create a PR soon to try and implement this. I just have a question: for the video to actually be written to disk, I need to issue save(plotter.filename, plotter.io) after the last time step of a simulation. Can we currently schedule that with a callback? Or should I just code that step separately?

[glwagner]

That's the first motivation I've heard to have a finalize! interface for callbacks, in addition to initialize! (which we already have)

[tomchor]

That's the first motivation I've heard to have a finalize! interface for callbacks, in addition to initialize! (which we already have)

I think it would also be useful for writing a checkpoint at the end of a simulation, no?

[glwagner]

Adding finalize! here: #4116

[glwagner]

That's the first motivation I've heard to have a finalize! interface for callbacks, in addition to initialize! (which we already have)

I think it would also be useful for writing a checkpoint at the end of a simulation, no?

that seems like a good idea too