Require output_writer in diagnostics

src/diagnostics/default_diagnostics.jl

@@ -5,37 +5,38 @@
 # level interfaces, add them here. Feel free to include extra files.
 
 """
-    default_diagnostics(model)
+    default_diagnostics(model, output_writer)
 
 
-Return a list of `ScheduledDiagnostic`s associated with the given `model`.
+Return a list of `ScheduledDiagnostic`s associated with the given `model` that use
+`output_write` to write to disk.
 
 """
-function default_diagnostics(model::AtmosModel)
+function default_diagnostics(model::AtmosModel; output_writer)
     # Unfortunately, [] is not treated nicely in a map (we would like it to be "excluded"),
     # so we need to manually filter out the submodels that don't have defaults associated
     # to
     non_empty_fields = filter(
-        x -> default_diagnostics(getfield(model, x)) != [],
+        x -> default_diagnostics(getfield(model, x); output_writer) != [],
         fieldnames(AtmosModel),
     )
 
     # We use a map because we want to ensure that diagnostics is a well defined type, not
     # Any. This reduces latency.
     return vcat(
-        core_default_diagnostics(),
+        core_default_diagnostics(output_writer),
         map(non_empty_fields) do field
-            default_diagnostics(getfield(model, field))
+            default_diagnostics(getfield(model, field); output_writer)
         end...,
     )
 end
 
 # Base case: if we call default_diagnostics on something that we don't have information
 # about, we get nothing back (to be specific, we get an empty list, so that we can assume
 # that all the default_diagnostics return the same type). This is used by
-# default_diagnostics(model::AtmosModel), so that we can ignore defaults for submodels
-# that have no given defaults.
-default_diagnostics(_) = []
+# default_diagnostics(model::AtmosModel; output_writer), so that we can ignore defaults for
+# submodels that have no given defaults.
+default_diagnostics(submodel; output_writer) = []
 
 
 """
@@ -69,96 +70,95 @@ function average_pre_output_hook!(accum, counter)
 end
 
 """
-    daily_maxs(short_names...; output_writer = HDF5Writer())
+    daily_maxs(short_names...; output_writer)
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily max for the given variables.
 """
-daily_maxs(short_names...; output_writer = HDF5Writer()) =
+daily_maxs(short_names...; output_writer) =
     common_diagnostics(24 * 60 * 60, max, output_writer, short_names...)
 """
     daily_max(short_names; output_writer = HDF5Writer())
 
 
 Return a `ScheduledDiagnostics` that computes the daily max for the given variable.
 """
-daily_max(short_names; output_writer = HDF5Writer()) =
+daily_max(short_names; output_writer) =
     daily_maxs(short_names; output_writer)[1]
 
 """
-    daily_mins(short_names...; output_writer = HDF5Writer())
+    daily_mins(short_names...; output_writer)
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily min for the given variables.
 """
-daily_mins(short_names...; output_writer = HDF5Writer()) =
+daily_mins(short_names...; output_writer) =
     common_diagnostics(24 * 60 * 60, min, output_writer, short_names...)
 """
     daily_min(short_names; output_writer = HDF5Writer())
 
 
 Return a `ScheduledDiagnostics` that computes the daily min for the given variable.
 """
-daily_min(short_names; output_writer = HDF5Writer()) =
+daily_min(short_names; output_writer) =
     daily_mins(short_names; output_writer)[1]
 
 """
-    daily_averages(short_names...; output_writer = HDF5Writer())
+    daily_averages(short_names...; output_writer)
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily average for the given variables.
 """
 # An average is just a sum with a normalization before output
-daily_averages(short_names...; output_writer = HDF5Writer()) =
-    common_diagnostics(
-        24 * 60 * 60,
-        (+),
-        output_writer,
-        short_names...;
-        pre_output_hook! = average_pre_output_hook!,
-    )
+daily_averages(short_names...; output_writer) = common_diagnostics(
+    24 * 60 * 60,
+    (+),
+    output_writer,
+    short_names...;
+    pre_output_hook! = average_pre_output_hook!,
+)
 """
-    daily_average(short_names; output_writer = HDF5Writer())
+    daily_average(short_names; output_writer)
 
 
 Return a `ScheduledDiagnostics` that compute the daily average for the given variable.
 """
 # An average is just a sum with a normalization before output
-daily_average(short_names; output_writer = HDF5Writer()) =
+daily_average(short_names; output_writer) =
     daily_averages(short_names; output_writer)[1]
 """
-    hourly_maxs(short_names...; output_writer = HDF5Writer())
+    hourly_maxs(short_names...; output_writer)
 
 
 Return a list of `ScheduledDiagnostics` that compute the hourly max for the given variables.
 """
-hourly_maxs(short_names...; output_writer = HDF5Writer()) =
+hourly_maxs(short_names...; output_writer) =
     common_diagnostics(60 * 60, max, output_writer, short_names...)
 
 """
-    hourly_max(short_names; output_writer = HDF5Writer())
+    hourly_max(short_names; output_writer)
 
 
-Return a `ScheduledDiagnostics` that computse the hourly max for the given variable.
+Return a `ScheduledDiagnostics` that computes the hourly max for the given variable.
 """
-hourly_max(short_names...; output_writer = HDF5Writer()) =
-    hourly_maxs(short_names)[1]
+hourly_max(short_names...; output_writer) =
+    hourly_maxs(short_names; output_writer)[1]
 
 """
-    hourly_mins(short_names...; output_writer = HDF5Writer())
+    hourly_mins(short_names...; output_writer)
 
 
 Return a list of `ScheduledDiagnostics` that compute the hourly min for the given variables.
 """
-hourly_mins(short_names...; output_writer = HDF5Writer()) =
+hourly_mins(short_names...; output_writer) =
     common_diagnostics(60 * 60, min, output_writer, short_names...)
 """
-    hourly_mins(short_names...; output_writer = HDF5Writer())
+    hourly_mins(short_names...; output_writer)
 
 
 Return a `ScheduledDiagnostics` that computes the hourly min for the given variable.
 """
-hourly_min(short_names; output_writer = HDF5Writer()) =
+hourly_min(short_names; output_writer) =
     hourly_mins(short_names; output_writer)[1]
 
 # An average is just a sum with a normalization before output
@@ -168,81 +168,84 @@ hourly_min(short_names; output_writer = HDF5Writer()) =
 
 Return a list of `ScheduledDiagnostics` that compute the hourly average for the given variables.
 """
-hourly_averages(short_names...; output_writer = HDF5Writer()) =
-    common_diagnostics(
-        60 * 60,
-        (+),
-        output_writer,
-        short_names...;
-        pre_output_hook! = average_pre_output_hook!,
-    )
+hourly_averages(short_names...; output_writer) = common_diagnostics(
+    60 * 60,
+    (+),
+    output_writer,
+    short_names...;
+    pre_output_hook! = average_pre_output_hook!,
+)
 
 """
-    hourly_average(short_names...; output_writer = HDF5Writer())
+    hourly_average(short_names...; output_writer)
 
 
 Return a `ScheduledDiagnostics` that computes the hourly average for the given variable.
 """
-hourly_average(short_names; output_writer = HDF5Writer()) =
+hourly_average(short_names; output_writer) =
     hourly_averages(short_names; output_writer)[1]
 
 # Include all the subdefaults
 
 ########
 # Core #
 ########
-function core_default_diagnostics()
+function core_default_diagnostics(output_writer)
     core_diagnostics = ["ts", "ta", "thetaa", "pfull", "rhoa", "ua", "va", "wa"]
 
     return [
-        daily_averages(core_diagnostics...)...,
-        daily_max("ts"),
-        daily_min("ts"),
+        daily_averages(core_diagnostics...; output_writer)...,
+        daily_max("ts"; output_writer),
+        daily_min("ts"; output_writer),
     ]
 end
 
 ###############
 # Energy form #
 ###############
-function default_diagnostics(::TotalEnergy)
+function default_diagnostics(::TotalEnergy; output_writer)
     total_energy_diagnostics = ["hfes"]
 
-    return [daily_averages(total_energy_diagnostics...)...]
+    return [daily_averages(total_energy_diagnostics...; output_writer)...]
 end
 
 
 ##################
 # Moisture model #
 ##################
 function default_diagnostics(
-    ::T,
+    ::T;
+    output_writer,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
     moist_diagnostics = ["hur", "hus", "hussfc", "evspsbl"]
 
-    return [daily_averages(moist_diagnostics...)...]
+    return [daily_averages(moist_diagnostics...; output_writer)...]
 end
 
 #######################
 # Precipitation model #
 #######################
-function default_diagnostics(::Microphysics0Moment)
+function default_diagnostics(::Microphysics0Moment; output_writer)
     precip_diagnostics = ["pr"]
 
-    return [daily_averages(precip_diagnostics...)...]
+    return [daily_averages(precip_diagnostics...; output_writer)...]
 end
 
 ##################
 # Radiation mode #
 ##################
-function default_diagnostics(::RRTMGPI.AbstractRRTMGPMode)
+function default_diagnostics(::RRTMGPI.AbstractRRTMGPMode; output_writer)
     allsky_diagnostics = ["rsd", "rsu", "rld", "rlu"]
 
-    return [daily_averages(allsky_diagnostics...)...]
+    return [daily_averages(allsky_diagnostics...; output_writer)...]
 end
 
 
-function default_diagnostics(::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics)
+function default_diagnostics(
+    ::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics;
+    output_writer,
+)
     clear_diagnostics = ["rsdcs", "rsucs", "rldcs", "rlucs"]
 
-    return [daily_averages(clear_diagnostics...)...]
+    return [daily_averages(clear_diagnostics...; output_writer)...]
 end

src/solver/type_getters.jl

@@ -633,13 +633,17 @@ function get_diagnostics(parsed_args, atmos_model)
         "average" => ((+), CAD.average_pre_output_hook!),
     )
 
+    hdf5_writer = CAD.HDF5Writer()
+    netcdf_writer = CAD.NetCDFWriter()
+    writers = (hdf5_writer, netcdf_writer)
+
     # The default writer is HDF5
     ALLOWED_WRITERS = Dict(
-        "nothing" => CAD.HDF5Writer(),
-        "h5" => CAD.HDF5Writer(),
-        "hdf5" => CAD.HDF5Writer(),
-        "nc" => CAD.NetCDFWriter(),
-        "netcdf" => CAD.NetCDFWriter(),
+        "nothing" => hdf5_writer,
+        "h5" => hdf5_writer,
+        "hdf5" => hdf5_writer,
+        "nc" => netcdf_writer,
+        "netcdf" => netcdf_writer,
     )
 
     diagnostics = map(yaml_diagnostics) do yaml_diag
@@ -699,9 +703,16 @@ function get_diagnostics(parsed_args, atmos_model)
     end
 
     if parsed_args["output_default_diagnostics"]
-        return [CAD.default_diagnostics(atmos_model)..., diagnostics...]
+        return [
+            CAD.default_diagnostics(
+                atmos_model;
+                output_writer = netcdf_writer,
+            )...,
+            diagnostics...,
+        ],
+        writers
     else
-        return collect(diagnostics)
+        return collect(diagnostics), writers
     end
 end
 
@@ -847,8 +858,10 @@ function get_integrator(config::AtmosConfig)
 
     # Initialize diagnostics
     s = @timed_str begin
-        diagnostics = get_diagnostics(config.parsed_args, atmos)
+        diagnostics, writers = get_diagnostics(config.parsed_args, atmos)
     end
+    # Add writers to the cache, so that we can close the files in solve_atmos!
+    p = merge(p, (; writers = writers))
     @info "initializing diagnostics: $s"
 
     # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,