Merge pull request #100 from jbisits/jib-checkkeypresetnfunction

Check if keys are present in dataset before saving

src/staircase_background.jl

@@ -153,35 +153,40 @@ function save_background_state!(simulation, model, background_state)
     compute!(σ_background)
     σ_background_array = Array(interior(σ_background, :, :, :))
 
-    if simulation.output_writers[:tracers] isa NetCDFOutputWriter
-
-        NCDataset(simulation.output_writers[:tracers].filepath, "a") do ds
-            defVar(ds, "S_background", S_background_array, ("xC", "yC", "zC"),
-                attrib = Dict("longname" => "Background field for salinity",
-                                "units" => "gkg⁻¹"))
-            defVar(ds, "T_background", T_background_array, ("xC", "yC", "zC"),
-                attrib =  Dict("longname" => "Background field for temperature",
-                                "units" => "°C"))
-        end
-
-        NCDataset(simulation.output_writers[:computed_output].filepath, "a") do ds
-            defVar(ds, "σ_background", σ_background_array, ("xC", "yC", "zC"),
-                attrib = Dict("longname" => "Background field for potential density (0dbar) computed from the `S` and `T` background fields",
-                                "units" => "kgm⁻³"))
-        end
-
-    elseif simulation.output_writers[:tracers] isa JLD2OutputWriter
+    # Here need only check if one key is present as if one is all are
+    if !check_key_present(simulation, :tracers, "S_background")
+
+        ow = simulation.output_writers
+        if ow[:tracers] isa NetCDFOutputWriter
+
+            NCDataset(ow[:tracers].filepath, "a") do ds
+                defVar(ds, "S_background", S_background_array, ("xC", "yC", "zC"),
+                    attrib = Dict("longname" => "Background field for salinity",
+                                    "units" => "gkg⁻¹"))
+                defVar(ds, "T_background", T_background_array, ("xC", "yC", "zC"),
+                    attrib =  Dict("longname" => "Background field for temperature",
+                                    "units" => "°C"))
+            end
+
+            NCDataset(ow[:computed_output].filepath, "a") do ds
+                defVar(ds, "σ_background", σ_background_array, ("xC", "yC", "zC"),
+                    attrib = Dict("longname" => "Background field for potential density (0dbar) computed from the `S` and `T` background fields",
+                                    "units" => "kgm⁻³"))
+            end
+
+        elseif ow[:tracers] isa JLD2OutputWriter
+
+            jldopen(ow[:tracers].filepath, "a+") do f
+                f["S_background"] = S_background_array
+                f["T_background"] = T_background_array
+            end
+
+            jldopen(ow[:computed_output].filepath, "a+") do f
+                f["σ_background"] = σ_background_array
+            end
 
-        jldopen(simulation.output_writers[:tracers].filepath, "a+") do f
-            f["S_background"] = S_background_array
-            f["T_background"] = T_background_array
-        end
-
-        jldopen(simulation.output_writers[:computed_output].filepath, "a+") do f
-            f["σ_background"] = σ_background_array
         end
 
     end
-
     return nothing
 end

src/staircase_diagnostics.jl

@@ -29,8 +29,15 @@ function compute_R_ρ!(computed_output::AbstractString, tracers::AbstractString,
     R_ρ = @. (0.5 * (ρ_f - ρ_u) + 0.5 * (ρ_l - ρ_g)) / (0.5 * (ρ_f - ρ_l) + 0.5 * (ρ_u - ρ_g))
 
     NCDataset(computed_output, "a") do ds2
-        defVar(ds2, "R_ρ", R_ρ, ("time",),
-                attrib = Dict("long_name" => "Density ratio"))
+        if haskey(ds2, "R_ρ")
+            # rename so if picking up can saved the whole array
+            renameVar(ds2, "R_ρ", "R_ρ_prior_cp")
+            defVar(ds2, "R_ρ", R_ρ, ("time",),
+                    attrib = Dict("long_name" => "Density ratio"))
+        else
+            defVar(ds2, "R_ρ", R_ρ, ("time",),
+                    attrib = Dict("long_name" => "Density ratio"))
+        end
     end
 
     close(ds)

src/staircase_simulation.jl

@@ -279,6 +279,29 @@ function checkpointer_setup!(simulation, sdns, output_dir,
 
 end
 checkpointer_setup!(simulation, sdns, output_dir, checkpointer_time_interval::Nothing) = nothing
+
+"""
+    function check_key_present(simulation, name::Symbol, key)
+Check if `key` is already present in `simlulation.output_writers[name]`. Return `true` if
+`key` is present and `false` otherwise.
+"""
+function check_key_present(simulation, name::Symbol, key)
+
+    ow = simulation.output_writers[name]
+    key_present = false
+
+    if ow isa NetCDFOutputWriter
+        NCDataset(ow.filepath) do ds
+            key_present = haskey(ds, key)
+        end
+    elseif ow isa JLD2OutputWriter
+        jldopen(ow.filepath, "a+") do f
+            key_present = haskey(f, key)
+        end
+    end
+
+    return key_present
+end
 """
     function simulation_progress(sim)
 Useful progress messaging for simulation runs. This function is from an