Merge pull request #184 from jbisits/joey-bettervelocitysaving

Better option for velocity saving
jbisits · Aug 23, 2024 · 2eec150 · 2eec150
2 parents f4441f3 + 1186202
commit 2eec150
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "TwoLayerDirectNumericalShenanigans"
 uuid = "40aaee9f-3595-48be-b36c-f1067009652f"
 authors = ["Josef Bisits <[email protected]>"]
-version = "0.6.4"
+version = "0.7.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

diff --git a/src/twolayerdns.jl b/src/twolayerdns.jl
@@ -152,16 +152,16 @@ there is no `Checkpointer`.
 """
 function TLDNS_simulation_setup(tldns::TwoLayerDNS, Δt::Number,
                                 stop_time::Number, save_schedule::Number,
-                                save_custom_output!::Function=no_custom_output!;
+                                save_custom_output!::Function=no_custom_output!,
+                                save_velocities!::Function=no_velocities!;
                                 save_file = :netcdf,
                                 output_path = SIMULATION_PATH,
                                 checkpointer_time_interval = nothing,
                                 cfl = 0.75,
                                 diffusive_cfl = 0.75,
                                 max_change = 1.2,
                                 max_Δt = 1e-1,
-                                overwrite_saved_output = true,
-                                save_velocities = false)
+                                overwrite_saved_output = true)
 
     model = tldns.model
     simulation = Simulation(model; Δt, stop_time)
@@ -180,7 +180,7 @@ function TLDNS_simulation_setup(tldns::TwoLayerDNS, Δt::Number,
     save_tracers!(simulation, model, save_info)
 
     # model velocities
-    save_velocities ? save_velocities!(simulation, model, save_info) : nothing
+    save_velocities!(simulation, model, save_info)
 
     # Custom saved output
     save_custom_output!(simulation, tldns, save_info)
@@ -251,11 +251,11 @@ end
 save_tracers!(simulation, model, save_info::Tuple) =
     save_tracers!(simulation, model, save_info...)
 """
-    function save_velocities!(simulation, model, save_schedule, save_file, output_dir,
+    function save_all_velocities!(simulation, model, save_schedule, save_file, output_dir,
                               overwrite_saved_output)
 Save `model.velocities` during a `Simulation` using an `OutputWriter`.
 """
-function save_velocities!(simulation, model, save_schedule, save_file, output_dir,
+function save_all_velocities!(simulation, model, save_schedule, save_file, output_dir,
                           overwrite_saved_output)
 
     u, v, w = model.velocities
@@ -277,8 +277,39 @@ function save_velocities!(simulation, model, save_schedule, save_file, output_di
     return nothing
 
 end
-save_velocities!(simulation, model, save_info::Tuple) =
-    save_velocities!(simulation, model, save_info...)
+save_all_velocities!(simulation, model, save_info::Tuple) =
+    save_all_velocities!(simulation, model, save_info...)
+"""
+    function save_vertical_velocities!(simulation, model, save_schedule, save_file, output_dir,
+                                    overwrite_saved_output)
+Only save vertical velocity.
+"""
+function save_vertical_velocities!(simulation, model, save_schedule, save_file, output_dir,
+                                    overwrite_saved_output)
+
+    w = model.velocities.w
+    velocities = Dict("w" => w)
+
+    simulation.output_writers[:velocities] =
+    save_file == :netcdf ? NetCDFOutputWriter(model, velocities;
+                                filename = "velocities",
+                                dir = output_dir,
+                                overwrite_existing = overwrite_saved_output,
+                                schedule = TimeInterval(save_schedule)
+                                ) :
+                JLD2OutputWriter(model, velocities;
+                                filename = "velocities",
+                                dir = output_dir,
+                                schedule = TimeInterval(save_schedule),
+                                overwrite_existing = overwrite_saved_output)
+
+    return nothing
+
+end
+save_vertical_velocities!(simulation, model, save_info::Tuple) =
+    save_vertical_velocities!(simulation, model, save_info...)
+"Default"
+no_velocities!(simulation, model, save_info...) = nothing
 """
     function save_computed_output!(simulation, tldns, save_schedule, save_file, output_dir,
                                    overwrite_saved_output, reference_gp_height)