Introducing output helpers

examples/two_dimensional_turbulence.jl

@@ -106,11 +106,7 @@ s = sqrt(u^2 + v^2)
 
 # We pass these operations to an output writer below to calculate and output them during the simulation.
 filename = "two_dimensional_turbulence"
-
-simulation.output_writers[:fields] = JLD2OutputWriter(model, (; ω, s),
-                                                      schedule = TimeInterval(0.6),
-                                                      filename = filename * ".jld2",
-                                                      overwrite_existing = true)
+output!(simulation, (; ω, s); schedule=TimeInterval(0.6), filename)
 
 # ## Running the simulation
 #

src/OutputWriters/jld2_output_writer.jl

@@ -28,17 +28,17 @@ ext(::Type{JLD2OutputWriter}) = ".jld2"
 
 """
     JLD2OutputWriter(model, outputs; filename, schedule,
-                              dir = ".",
-                          indices = (:, :, :),
-                       with_halos = false,
-                       array_type = Array{Float64},
-                   file_splitting = NoFileSplitting(),
-               overwrite_existing = false,
-                             init = noinit,
-                        including = [:grid, :coriolis, :buoyancy, :closure],
-                          verbose = false,
-                             part = 1,
-                          jld2_kw = Dict{Symbol, Any}())
+                     dir = ".",
+                     indices = (:, :, :),
+                     with_halos = false,
+                     array_type = Array{Float64},
+                     file_splitting = NoFileSplitting(),
+                     overwrite_existing = false,
+                     init = noinit,
+                     including = [:grid, :coriolis, :buoyancy, :closure],
+                     verbose = false,
+                     part = 1,
+                     jld2_kw = Dict{Symbol, Any}())
 
 Construct a `JLD2OutputWriter` for an Oceananigans `model` that writes `label, output` pairs
 in `outputs` to a JLD2 file.
@@ -163,17 +163,17 @@ JLD2OutputWriter scheduled on TimeInterval(20 minutes):
 ```
 """
 function JLD2OutputWriter(model, outputs; filename, schedule,
-                                   dir = ".",
-                               indices = (:, :, :),
-                            with_halos = false,
-                            array_type = Array{Float64},
-                        file_splitting = NoFileSplitting(),
-                    overwrite_existing = false,
-                                  init = noinit,
-                             including = default_included_properties(model),
-                               verbose = false,
-                                  part = 1,
-                               jld2_kw = Dict{Symbol, Any}())
+                          dir = ".",
+                          indices = (:, :, :),
+                          with_halos = true,
+                          array_type = Array{Float64},
+                          file_splitting = NoFileSplitting(),
+                          overwrite_existing = true,
+                          init = noinit,
+                          including = default_included_properties(model),
+                          verbose = false,
+                          part = 1,
+                          jld2_kw = Dict{Symbol, Any}())
 
     mkpath(dir)
     filename = auto_extension(filename, ".jld2")
@@ -283,7 +283,7 @@ function write_output!(writer::JLD2OutputWriter, model)
         verbose && @info "Writing JLD2 output $(keys(writer.outputs)) to $path..."
 
         start_time, old_filesize = time_ns(), filesize(writer.filepath)
-        jld2output!(writer.filepath, model.clock.iteration, model.clock.time, data, writer.jld2_kw)
+        write_jld2_output!(writer.filepath, model.clock.iteration, model.clock.time, data, writer.jld2_kw)
         end_time, new_filesize = time_ns(), filesize(writer.filepath)
 
         verbose && @info @sprintf("Writing done: time=%s, size=%s, Δsize=%s",
@@ -296,14 +296,14 @@ function write_output!(writer::JLD2OutputWriter, model)
 end
 
 """
-    jld2output!(path, iter, time, data, kwargs)
+    write_jld2_output!(path, iter, time, data, kwargs)
 
 Write the (name, value) pairs in `data`, including the simulation
 `time`, to the JLD2 file at `path` in the `timeseries` group,
 stamping them with `iter` and using `kwargs` when opening
 the JLD2 file.
 """
-function jld2output!(path, iter, time, data, kwargs)
+function write_jld2_output!(path, iter, time, data, kwargs)
     jldopen(path, "r+"; kwargs...) do file
         file["timeseries/t/$iter"] = time
         for name in keys(data)

src/Simulations/Simulations.jl

@@ -6,6 +6,7 @@ export run!
 export Callback, add_callback!
 export iteration
 export stopwatch
+export JLD2_output!
 
 using Oceananigans.Models
 using Oceananigans.Diagnostics
@@ -20,9 +21,23 @@ using OrderedCollections: OrderedDict
 
 import Base: show
 
+function unique_name(prefix::Symbol, existing_names)
+    if !(prefix ∈ existing_names)
+        name = prefix
+    else # make it unique
+        n = 1
+        while Symbol(prefix, n) ∈ existing_names
+            n += 1
+        end
+        name = Symbol(prefix, n)
+    end
+    return name
+end
+
 include("callback.jl")
 include("simulation.jl")
 include("run.jl")
 include("time_step_wizard.jl")
+include("output_helpers.jl")
 
 end # module

src/Simulations/callback.jl

@@ -71,27 +71,8 @@ Callback(wta::WindowedTimeAverage, schedule; kw...) =
     throw(ArgumentError("Schedule must be inferred from WindowedTimeAverage. 
                         Use Callback(windowed_time_average)"))
 
-struct GenericName end
-
-function unique_callback_name(name, existing_names)
-    if name ∈ existing_names
-        return Symbol(:another_, name)
-    else
-        return name
-    end
-end
-
-function unique_callback_name(::GenericName, existing_names)
-    prefix = :callback # yeah, that's generic
-
-    # Find a unique one
-    n = 1
-    while Symbol(prefix, n) ∈ existing_names
-        n += 1
-    end
-
-    return Symbol(prefix, n)
-end
+struct GenericCallbackName end
+unique_name(::GenericCallbackName, existing) = unique_name(:callback, existing)
 
 """
     add_callback!(simulation, callback::Callback; name = GenericName(), callback_kw...)
@@ -109,8 +90,8 @@ already exists.
 
 The `callback` (which contains a schedule) can also be supplied directly.
 """
-function add_callback!(simulation, callback::Callback; name = GenericName())
-    name = unique_callback_name(name, keys(simulation.callbacks))
+function add_callback!(simulation, callback::Callback; name = GenericCallbackName())
+    name = unique_name(GenericCallbackName(), keys(simulation.callbacks))
     simulation.callbacks[name] = callback
     return nothing
 end

src/Simulations/output_helpers.jl

@@ -0,0 +1,25 @@
+#####
+##### Lend a helping hand
+#####
+
+struct GenericJLD2Name end
+unique_name(::GenericJLD2Name, existing) = unique_name(:jld2, existing)
+
+struct JLD2Format end
+
+"""
+    output!(simulation, outputs [, format=JLD2Format()]; kw...)
+
+"""
+output!(simulation, outputs; kw...) = output!(simulation, outputs, JLD2Format(); kw...)
+
+function output!(simulation, outputs, ::JLD2Format; kw...)
+    if !(name ∈ keys(kw))
+        name = GenericJLD2Name()        
+    end
+    name = unique_name(name, keys(simulation.output_writers))
+    ow = JLD2OutputWriter(simulation.model, outputs; kw...)
+    simulation.output_writers[name] = ow
+    return nothing
+end
+