From 6a6853ddf53c53c321c16393682b786e82ef5a8f Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Wed, 25 Sep 2024 20:01:11 -0600 Subject: [PATCH] Introducing output helpers --- examples/two_dimensional_turbulence.jl | 6 +-- src/OutputWriters/jld2_output_writer.jl | 50 ++++++++++++------------- src/Simulations/Simulations.jl | 15 ++++++++ src/Simulations/callback.jl | 27 ++----------- src/Simulations/output_helpers.jl | 25 +++++++++++++ 5 files changed, 70 insertions(+), 53 deletions(-) create mode 100644 src/Simulations/output_helpers.jl diff --git a/examples/two_dimensional_turbulence.jl b/examples/two_dimensional_turbulence.jl index e6807b963d..b554a63740 100644 --- a/examples/two_dimensional_turbulence.jl +++ b/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 # diff --git a/src/OutputWriters/jld2_output_writer.jl b/src/OutputWriters/jld2_output_writer.jl index 50a0a31937..367c5dc1ca 100644 --- a/src/OutputWriters/jld2_output_writer.jl +++ b/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) diff --git a/src/Simulations/Simulations.jl b/src/Simulations/Simulations.jl index 9dc720b815..ffb0893a7d 100644 --- a/src/Simulations/Simulations.jl +++ b/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 diff --git a/src/Simulations/callback.jl b/src/Simulations/callback.jl index acc40486f1..ad1f81cf60 100644 --- a/src/Simulations/callback.jl +++ b/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 diff --git a/src/Simulations/output_helpers.jl b/src/Simulations/output_helpers.jl new file mode 100644 index 0000000000..d157bb7439 --- /dev/null +++ b/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 +