Close output writers at the end of the simulation

Author: ph-kev

experiments/integrated/fluxnet/ozark_pft.jl

@@ -324,6 +324,8 @@ prob = SciMLBase.ODEProblem(
 
 sol = SciMLBase.solve(prob, ode_algo; dt = dt, callback = cb);
 
+ClimaLand.Diagnostics.close_output_writers(diags)
+
 # Extract model output from the saved diagnostics
 short_names_1D = [
     "sif", # SIF

experiments/integrated/fluxnet/run_fluxnet.jl

@@ -293,6 +293,8 @@ prob = SciMLBase.ODEProblem(
 
 sol = SciMLBase.solve(prob, ode_algo; dt = dt, callback = cb);
 
+ClimaLand.Diagnostics.close_output_writers(diags)
+
 # Extract model output from the saved diagnostics
 short_names_1D = [
     "sif", # SIF

experiments/integrated/global/global_soil_canopy.jl

@@ -244,6 +244,8 @@ sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
     callback = SciMLBase.CallbackSet(driver_cb, diag_cb),
 )
 
+ClimaLand.Diagnostics.close_output_writers(diags)
+
 # ClimaAnalysis
 if ClimaComms.iamroot(context)
     simdir = ClimaAnalysis.SimDir(outdir)

experiments/standalone/Bucket/bucket_era5.jl

@@ -213,6 +213,8 @@ sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
     callback = cb,
 );
 
+ClimaLand.Diagnostics.close_output_writers(diags)
+
 simdir = ClimaAnalysis.SimDir(output_dir)
 short_names = ["rn", "tsfc", "qsfc", "lhf", "shf", "wsoil", "wsfc", "ssfc"]
 for short_name in short_names

experiments/standalone/Bucket/global_bucket_function.jl

@@ -193,6 +193,8 @@ sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
     callback = SciMLBase.CallbackSet(driver_cb, diag_cb),
 )
 
+ClimaLand.Diagnostics.close_output_writers(diags)
+
 #### ClimaAnalysis ####
 
 # all

src/diagnostics/Diagnostics.jl

@@ -32,6 +32,8 @@ import ClimaDiagnostics.Writers: HDF5Writer, NetCDFWriter, DictWriter
 import ClimaCore.Fields: zeros, field_values
 import ClimaCore.Operators: column_integral_definite!
 
+export close_output_writers
+
 include("diagnostic.jl")
 
 end

src/diagnostics/diagnostic.jl

@@ -150,6 +150,22 @@ function diagnostic_as_vectors(writer::DictWriter, diagnostic; layer = 1)
     return times, vector_layer_n
 end
 
+"""
+    close_output_writers(diagnostics)
+
+Close the output writers in the `diagnostics`, an iterable of
+`ClimaDiagnostics.ScheduledDiagnostic` or `nothing`.
+
+This function should be called at the end of every simulation.
+"""
+function close_output_writers(diagnostics)
+    isnothing(diagnostics) && return nothing
+    for diagnostic in diagnostics
+        close(diagnostic.output_writer)
+    end
+    return nothing
+end
+
 # Do you want to define more diagnostics? Add them here
 include("land_compute_methods.jl")
 

src/simulations/Simulations.jl

@@ -10,6 +10,8 @@ using ClimaLand
 export step!, solve!, LandSimulation
 include("initial_conditions.jl")
 
+import ..Diagnostics: close_output_writers
+
 """
     LandSimulation{
         M <: ClimaLand.AbstractModel,
@@ -194,7 +196,14 @@ Advances the land simulation `landsim` forward from the initial to final time,
 updating `landsim` in place.
 """
 function solve!(landsim::LandSimulation)
-    SciMLBase.solve!(landsim._integrator)
+    try
+        SciMLBase.solve!(landsim._integrator)
+    catch ret_code
+        @error "ClimaLand simulation crashed. Stacktrace for failed simulation" exception =
+            (ret_code, catch_backtrace())
+    finally
+        close_output_writers(landsim.diagnostics)
+    end
 end