Write output and plot animation/figure automatically

[tomchor]

As I'm in the early stages of assembling a simulation, I always find myself spending a lot of time in writing the output and plotting and animation. Most of that time I'm doing the same thing: setting up a writer with a couple of variables, then reading it, creating a figure, plotting heatmaps, and writing the animation loop. A lot of it a copy-paste. And if I want to other variables to plot, I have to modify many lines to accommodate for that.

I think this process can be automated. Obviously we won't be able to (and shouldn't aim to) reproduce all the flexibility we have when doing each these steps separately, but I think we can create something that can be used for something like 80% of use cases.

[glwagner]

Yes! I think this issue belongs in Oceananigans.

You can plot Fields and AbstractOperations directly with the Makie extension.

If you want to create an animation and plot multiple fields, then we need something new I agree.

There is an open PR in Oceananigans to make output easier. Partly, this addresses the issue that output is a little clunky:

CliMA/Oceananigans.jl#3793

However, I think it would be nice to have a feature that avoids output all together, and simply builds a FieldTimeSeries or FieldDataset directly. My idea is to build a new type called DatasetCollector with syntax like

collector = DatasetCollector(model, outputs)
add_callback!(collector, IterationInterval(10))

One could then use the data stored in the collector to generate an animation. We can also consider a helper function,

add_dataset_collector!(simulation, outputs, kw...)

[tomchor]

Yes! I think this issue belongs in Oceananigans.

You can plot Fields and AbstractOperations directly with the Makie extension.

If you want to create an animation and plot multiple fields, then we need something new I agree.

There is an open PR in Oceananigans to make output easier. Partly, this addresses the issue that output is a little clunky:

CliMA/Oceananigans.jl#3793

This is nice, but I was talking about something like:

psimulation = PlottedSimulation(simulation, outputs = (u, w, vorticity), indices = (1, :, :), a_few_more_kwargs...)
run!(psimulation)

and then bam: output files are generated, as well as figures/animations at the end, automatically. This would avoid the boilerplate of creating an output (which you sometimes want anyway), then importing Makie, getting the filename, reading the dataset, assembling the figure, title, panels, etc., and then writing the loop.

If you think that fits better in Oceananigans or that this isn't useful, I'll close this issue here. But I do think such a functionality would maybe we better off in Oceanostics.

However, I think it would be nice to have a feature that avoids output all together, and simply builds a FieldTimeSeries or FieldDataset directly.

I do agree that this would be useful. But most of the time I do want files written to investigate the output later.

[glwagner]

True but what about the same functionality with

simulation = Simulation(simulation)
add_plotter!(simulation, outputs = (u, w, vorticity), indices = (1, :, :), a_few_more_kwargs...)
run!(simulation)

This avoids the wrapper (which will add a lot of complexity because you have to extend a lot of functions ) and achieves the same thing I think.

[tomchor]

That's a pretty good idea. So basically the plotter would be like a callback?

[glwagner]

Yes, a callable object basically

[glwagner]

Implementing this will also drive developing of the plotting extension since it will rely on that

[jbisits]

For models I have been running I wrote an unregistered julia package to better keep track of everything (e.g. StaircaseShenanigans.jl) which contains plotting functions I want so I need only write them once (as I too can spend a lot of time copying and pasting!!) then I call them at the end of a simulation scripts.

The way I did this was to define the plotting functions in the src code then add methods to them in the Makie extension.

This way I just add the plotting I want at the end of a simulation script

using StaricaseShenanigans

## setup simulation...

run!(simulation)

## Produce animations
reduced_path = findlast('/', simulation.output_writers[:computed_output].filepath)
animation_path = simulation.output_writers[:computed_output].filepath[1:(reduced_path-1)]
cd(animation_path)
@info "Producing animations"
using CairoMakie
animate_density_anomaly(simulation.output_writers[:computed_output].filepath, "σ")
animate_tracers_anomaly(simulation.output_writers[:tracers].filepath)
animate_density(simulation.output_writers[:computed_output].filepath, "σ")
animate_tracers(simulation.output_writers[:tracers].filepath)

This is not as sophisticated as using a Callback as the simulation runs to produce the animations (as I did not really need to be but this is a great idea!), and does not avoid all the boilerplate code writing (though has greatly reduced my amount of copying and pasting), but perhaps something along these lines that is better and more flexible could be added to Oceananigans (if this is the kind of thing you are thinking of).

So you could have

simulation = Simulation(simulation)
add_plotter!(simulation, outputs = (u, w, vorticity), indices = (1, :, :), a_few_more_kwargs...)
run!(simulation)

where add_plotter might call a function animate_field for u, v and vorticity at indices e.g.

function animate_field(u, indices)
  # produce animation as simulation runs
  return nothing
end

The examples in GeophysicalFlows.jl are animated on the fly so perhaps could use some ideas from there.

I am happy to look into this further for Oceananigans if you want!