From 5efe29256ee7c120213990f3238fd7e3a2d61af5 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 08:28:17 -0700
Subject: [PATCH 01/73] Add DiagnosticVariable

---
 docs/make.jl                   |  1 +
 docs/src/diagnostics.md        | 53 +++++++++++++++++++++++++++++++++
 src/ClimaAtmos.jl              |  2 ++
 src/diagnostics/Diagnostics.jl | 54 ++++++++++++++++++++++++++++++++++
 4 files changed, 110 insertions(+)
 create mode 100644 docs/src/diagnostics.md
 create mode 100644 src/diagnostics/Diagnostics.jl

diff --git a/docs/make.jl b/docs/make.jl
index 3f067559883..30fa9b15cb9 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -26,6 +26,7 @@ makedocs(
         "Contributor Guide" => "contributor_guide.md",
         "Equations" => "equations.md",
         "EDMF Equations" => "edmf_equations.md",
+        "Diagnostics" => "diagnostics.md",
         "Diagnostic EDMF Equations" => "diagnostic_edmf_equations.md",
         "Gravity Wave Drag Parameterizations" => "gravity_wave.md",
         "Radiative Equilibrium" => "radiative_equilibrium.md",
diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
new file mode 100644
index 00000000000..8a81145935a
--- /dev/null
+++ b/docs/src/diagnostics.md
@@ -0,0 +1,53 @@
+# Computing and saving diagnostics
+
+## I want to add a new diagnostic variable
+
+Diagnostic variables are represented in `ClimaAtmos` with a `DiagnosticVariable`
+`struct`. Fundamentally, a `DiagnosticVariable` contains metadata about the
+variable, and a function that computes it from the state.
+
+### Metadata
+
+The metadata we currently support is `short_name`, `long_name`, `units`,
+`comments`. This metadata is relevant mainly in the context of how the variable
+is output. Therefore, it is responsibility of the `output_writer` (see
+`ScheduledDiagnostic`) to handle the metadata properly. The `output_writer`s
+provided by `ClimaAtmos` use this metadata.
+
+In `ClimaAtmos`, we follow the convention that:
+
+- `short_name` is the name used to identify the variable in the output files and
+                in the file names. It is short, but descriptive. We identify
+                diagnostics by their short name, so the diagnostics defined by
+                `ClimaAtmos` have to have unique `long_name`s. We follow the
+                Coupled Model Intercomparison Project (CMIP) convetions.
+
+- `long_name`: Name used to identify the variable in the input files.
+
+- `units`: Physical units of the variable.
+
+- `comments`: More verbose explanation of what the variable is, or comments related to how
+              it is defined or computed.
+
+In `ClimaAtmos`, we try to follow [this Google
+spreadsheet](https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx)
+for variable naming (except for the `short_names`, which we prefer being more
+descriptive).
+
+### Compute function
+
+The other piece of information needed to specify a `DiagnosticVariable` is a
+function `compute_from_integrator`. Schematically, a `compute_from_integrator` has to look like
+```julia
+function compute_from_integrator(integrator, out)
+    # FIXME: Remove this line when ClimaCore implements the broadcasting to enable this
+    out .= # Calculcations with the state (= integrator.u) and the parameters (= integrator.p)
+end
+```
+Diagnostics are implemented as callbacks function which pass the `integrator`
+object (from `OrdinaryDiffEq`) to `compute_from_integrator`.
+
+`compute_from_integrator` also takes a second argument, `out`, which is used to
+avoid extra memory allocations (which hurt performance). If `out` is `nothing`,
+and new area of memory is allocated. If `out` is a `ClimaCore.Field`, the
+operation is done in-place without additional memory allocations.
diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
index 07f6847372a..1f612e69a7f 100644
--- a/src/ClimaAtmos.jl
+++ b/src/ClimaAtmos.jl
@@ -117,4 +117,6 @@ include(joinpath("solver", "solve.jl"))
 
 include(joinpath("parameters", "create_parameters.jl"))
 
+include(joinpath("diagnostics", "Diagnostics.jl"))
+
 end # module
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
new file mode 100644
index 00000000000..9f9095c1ae1
--- /dev/null
+++ b/src/diagnostics/Diagnostics.jl
@@ -0,0 +1,54 @@
+# Diagnostics.jl
+#
+# This file contains:
+#
+# - The definition of what a DiagnosticVariable is. Morally, a DiagnosticVariable is a
+#   variable we know how to compute from the state. We attach more information to it for
+#   documentation and to reference to it with its long name. DiagnosticVariables can exist
+#   irrespective of the existence of an actual simulation that is being run.
+
+"""
+    DiagnosticVariable
+
+
+A recipe to compute a diagnostic variable from the state, along with some useful metadata.
+
+The primary use for `DiagnosticVariable`s is to be embedded in a `ScheduledDiagnostic` to
+compute diagnostics while the simulation is running.
+
+The metadata is used exclusively by the `output_writer` in the `ScheduledDiagnostic`. It is
+responsibility of the `output_writer` to follow the conventions about the meaning of the
+metadata and their use.
+
+In `ClimaAtmos`, we roughly follow the naming conventions listed in this file:
+https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
+
+Keyword arguments
+=================
+
+- `short_name`: Name used to identify the variable in the output files and in the file
+                names. Short but descriptive. `ClimaAtmos` follows the CMIP conventions and
+                the diagnostics are identified by the short name.
+
+- `long_name`: Name used to identify the variable in the input files.
+
+- `units`: Physical units of the variable.
+
+- `comments`: More verbose explanation of what the variable is, or comments related to how
+              it is defined or computed.
+
+- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
+                             It has to take two arguments: the `integrator`, and a
+                             pre-allocated area of memory where to write the result of the
+                             computation. It the no pre-allocated area is available, a new
+                             one will be allocated. To avoid extra allocations, this
+                             function should perform the calculation in-place (i.e., using
+                             `.=`).
+"""
+Base.@kwdef struct DiagnosticVariable{T <: AbstractString, T2}
+    short_name::T
+    long_name::T
+    units::T
+    comments::T
+    compute_from_integrator::T2
+end

From e702b2e2495053715be9200390045a0bb450f9c4 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 09:39:30 -0700
Subject: [PATCH 02/73] Add ALL_DIAGNOSTICS

---
 docs/src/diagnostics.md             | 15 ++++++
 src/diagnostics/Diagnostics.jl      | 73 ++++++++++++++++++++++++++++-
 src/diagnostics/core_diagnostics.jl | 15 ++++++
 3 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 src/diagnostics/core_diagnostics.jl

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 8a81145935a..8a04d25e22a 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -51,3 +51,18 @@ object (from `OrdinaryDiffEq`) to `compute_from_integrator`.
 avoid extra memory allocations (which hurt performance). If `out` is `nothing`,
 and new area of memory is allocated. If `out` is a `ClimaCore.Field`, the
 operation is done in-place without additional memory allocations.
+
+### Adding to the `ALL_DIAGNOSTICS` dictionary
+
+`ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable` in the
+`ALL_DIAGNOSTICS` dictionary. `ALL_DIAGNOSTICS` maps a `long_name` with the
+corresponding `DiagnosticVariable`.
+
+If you are extending `ClimaAtmos` and want to add a new diagnostic variable to
+`ALL_DIAGNOSTICS`, go ahead and look at the files we `include` in
+`diagnostics/Diagnostics.jl`. You can add more diagnostics in those files or add
+a new one. We provide a convenience function, `add_diagnostic_variable!` to add
+new `DiagnosticVariable`s to the `ALL_DIAGNOSTICS` dictionary.
+`add_diagnostic_variable!` take the same arguments as the constructor for
+`DiagnosticVariable`, but also performs additional checks. So, use
+`add_diagnostic_variable!` instead of editing the `ALL_DIAGNOSTICS` directly.
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 9f9095c1ae1..850a266ddeb 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -5,7 +5,13 @@
 # - The definition of what a DiagnosticVariable is. Morally, a DiagnosticVariable is a
 #   variable we know how to compute from the state. We attach more information to it for
 #   documentation and to reference to it with its long name. DiagnosticVariables can exist
-#   irrespective of the existence of an actual simulation that is being run.
+#   irrespective of the existence of an actual simulation that is being run. ClimaAtmos
+#   comes with several diagnostics already defined (in the `ALL_DIAGNOSTICS` dictionary).
+#
+# - A dictionary `ALL_DIAGNOSTICS` with all the diagnostics we know how to compute, keyed
+#   over their long name. If you want to add more diagnostics, look at the included files.
+#   You can add your own file if you want to define several new diagnostics that are
+#   conceptually related.
 
 """
     DiagnosticVariable
@@ -52,3 +58,68 @@ Base.@kwdef struct DiagnosticVariable{T <: AbstractString, T2}
     comments::T
     compute_from_integrator::T2
 end
+
+# ClimaAtmos diagnostics
+
+const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
+
+"""
+
+    add_diagnostic_variable!(; short_name,
+                               long_name,
+                               units,
+                               description,
+                               compute_from_integrator)
+
+
+Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates the state of
+`ClimaAtmos.ALL_DIAGNOSTICS`).
+
+If possible, please follow the naming scheme outline in
+https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
+
+Keyword arguments
+=================
+
+
+- `short_name`: Name used to identify the variable in the output files and in the file
+                names. Short but descriptive. `ClimaAtmos` diagnostics are identified by the
+                short name. We follow the Coupled Model Intercomparison Project conventions.
+
+- `long_name`: Name used to identify the variable in the input files.
+
+- `units`: Physical units of the variable.
+
+- `comments`: More verbose explanation of what the variable is, or comments related to how
+              it is defined or computed.
+
+- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
+                             It has to take two arguments: the `integrator`, and a
+                             pre-allocated area of memory where to write the result of the
+                             computation. It the no pre-allocated area is available, a new
+                             one will be allocated. To avoid extra allocations, this
+                             function should perform the calculation in-place (i.e., using
+                             `.=`).
+
+"""
+function add_diagnostic_variable!(;
+    short_name,
+    long_name,
+    units,
+    comments,
+    compute_from_integrator
+)
+    haskey(ALL_DIAGNOSTICS, short_name) &&
+        error("diagnostic $short_name already defined")
+
+    ALL_DIAGNOSTICS[short_name] = DiagnosticVariable(;
+        short_name,
+        long_name,
+        units,
+        comments,
+        compute_from_integrator
+    )
+end
+
+# Do you want to define more diagnostics? Add them here
+include("core_diagnostics.jl")
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
new file mode 100644
index 00000000000..2397a937f5c
--- /dev/null
+++ b/src/diagnostics/core_diagnostics.jl
@@ -0,0 +1,15 @@
+# This file is included in Diagnostics.jl
+
+# FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
+# physics. Please fix this!
+add_diagnostic_variable!(
+    short_name = "air_density",
+    long_name = "air_density",
+    units = "kg m^-3",
+    comments = "Density of air, a prognostic variable",
+    compute_from_integrator = (integrator, out) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        return deepcopy(integrator.u.c.ρ)
+    end,
+)

From 0a9dd4b58cfe560ba303d003d5ad379147f4b464 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 13:37:21 -0700
Subject: [PATCH 03/73] Add ScheduledDiagnosticIterations

---
 docs/src/diagnostics.md        |  87 ++++++++++++++++++++
 src/diagnostics/Diagnostics.jl | 146 +++++++++++++++++++++++++++++++++
 2 files changed, 233 insertions(+)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 8a04d25e22a..e5aef1b82af 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -1,5 +1,92 @@
 # Computing and saving diagnostics
 
+## I want to compute and output a diagnostic variable
+
+### The low-level interface
+
+Diagnostics are computed and output through callbacks to the main integrator.
+`ClimaAtmos` produces the list of callbacks from a ordered list of
+`ScheduledDiagnostic`s.
+
+A `ScheduledDiagnostic` is an instruction on how to compute and output a given
+`DiagnosticVariable` (see below), along with specific choices regarding
+reductions, compute/output frequencies, and so on. It can be point-wise in space
+and time, or can be the result of a reduction in a period that is defined by
+`output_every` (e.g., the daily average temperature).
+
+More specifically, a `ScheduledDiagnostic` contains the following pieces of data
+
+- `variable`: The diagnostic variable that has to be computed and output.
+- `output_every`: Frequency of how often to save the results to disk.
+- `output_writer`: Function that controls out to save the computed diagnostic
+  variable to disk.
+- `reduction_time_func`: If not `nothing`, the `ScheduledDiagnostic` receives an
+  area of scratch space `acc` where to accumulate partial results. Then, at
+  every `compute_every`, `reduction_time_func` is computed between the
+  previously stored value in `acc` and the new value. This implements a running
+  reduction. For example, if `reduction_time_func = max`, the space `acc` will
+  hold the running maxima of the diagnostic. `acc` is reset after output.
+- `reduction_space_func`: NOT IMPLEMENTED YET
+- `compute_every`: Run the computations every `compute_every`. This is not
+  particularly useful for point-wise diagnostics, where we enforce that
+  `compute_every` = `output_every`. `compute_every` has to evenly divide
+  `output_every`.
+- `pre_output_hook!`: Function that has to be run before saving to disk for
+   reductions (mostly used to implement averages). The function
+   `pre_output_hook!` is called with two arguments: the value accumulated during
+   the reduction, and the number of times the diagnostic was computed from the
+   last time it was output.
+
+To implement operations like the arithmetic average, the `reduction_time_func`
+has to be chosen as `sum`, and a `pre_output_hook!` that renormalize `acc` by
+the number of samples has to be provided. `pre_output_hook!` should mutate the
+accumulator in place. The return value of `pre_output_hook!` is discarded. An
+example of `pre_output_hook!` to compute the arithmetic average is
+`pre_output_hook!(acc, N) = @. acc = acc / N`.
+
+For custom reductions, it is necessary to also specify the identity of operation
+by defining a new method to `identity_of_reduction`. `identity_of_reduction` is
+a function that takes a `Val{op}` argument, where `op` is the operation for
+which we want to define the identity. For instance, for the `max`,
+`identity_of_reduction` would be `identity_of_reduction(::Val{max}) = -Inf`. The
+identities known to `ClimaAtmos` are defined in the
+`diagnostics/reduction_identities.jl` file. The identity is needed to ensure
+that we have a neutral state for the accumulators that are used in the
+reductions.
+
+A key entry in a `ScheduledDiagnostic` is the `output_writer`, the function
+responsible for saving the output to disk. `output_writer` is called with three
+arguments: the value that has to be output, the `ScheduledDiagnostic`, and the
+integrator. Internally, the integrator contains extra information (such as the
+current timestep), so the `output_writer` can implement arbitrarily complex
+behaviors. It is responsibility of the `output_writer` to properly use the
+provided information for meaningful output. `ClimaAtmos` provides functions that
+return `output_writer`s for standard operations (e.g., for writing to HDF5
+files).
+
+There are two flavors of `ScheduledDiagnostic`s:
+`ScheduledDiagnosticIterations`, and `ScheduledDiagnosticTime`. The main
+difference between the two is the domain of definition of their frequencies,
+which is measured in timesteps for the first one, and in seconds for the other
+one. `ScheduledDiagnosticTime`s offer a more natural way to set up physically
+meaningful reductions (e.g., we want a daily average). However, it is not clear
+what to do when the period does not line up with the timestep. What is the
+expected behavior when we want a daily average but our timestep is of 10 hours?
+There are multiple possible answer to this question. In `ClimaAtmos`, we enforce
+that all the periods are multiples of the timestep. With this restriction, a
+`ScheduledDiagnosticTime` can be translated to a
+`ScheduledDiagnosticIterations`, where the problem is perfectly represented (in
+this sense, one can think of `ScheduledDiagnosticIterations` as being as
+internal representation and as `ScheduledDiagnosticTime` being the user-facing
+one).
+
+`ScheduledDiagnosticTime` behave like as `ScheduledDiagnosticIterations`, with
+the exception that they can take a special value for `compute_every`, which can
+be set to `:timestep` (the symbol) to ensure that the diagnostic is computed at
+the end of every integration step. This is particularly convenient when defining
+default diagnostics, as they should be largely independent on the choice of the
+specific simulation being run (and its timestep).
+
 ## I want to add a new diagnostic variable
 
 Diagnostic variables are represented in `ClimaAtmos` with a `DiagnosticVariable`
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 850a266ddeb..9f976de1e71 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -12,6 +12,26 @@
 #   over their long name. If you want to add more diagnostics, look at the included files.
 #   You can add your own file if you want to define several new diagnostics that are
 #   conceptually related.
+#
+# - The definition of what a ScheduledDiagnostics is. Morally, a ScheduledDiagnostics is a
+#   DiagnosticVariable we want to compute in a given simulation. For example, it could be
+#   the temperature averaged over a day. We can have multiple ScheduledDiagnostics for the
+#   same DiagnosticVariable (e.g., daily and monthly average temperatures).
+#
+#   We provide two types of ScheduledDiagnostics: ScheduledDiagnosticIterations and
+#   ScheduledDiagnosticTime, with the difference being only in what domain the recurrence
+#   time is defined (are we doing something at every N timesteps or every T seconds?). It is
+#   much cleaner and simpler to work with ScheduledDiagnosticIterations because iterations
+#   are well defined and consistent. On the other hand, working in the time domain requires
+#   dealing with what happens when the timestep is not lined up with the output period.
+#   Possible solutions to this problem include: uneven output, interpolation, or restricting
+#   the user from picking specific combinations of timestep/output period. In the current
+#   implementation, we choose the third option. So, ScheduledDiagnosticTime is provided
+#   because it is the physically interesting quantity. If we know what is the timestep, we
+#   can convert between the two and check if the diagnostics are well-posed in terms of the
+#   relationship between the periods and the timesteps. In some sense, you can think of
+#   ScheduledDiagnosticIterations as an internal representation and ScheduledDiagnosticTime
+#   as the external interface.
 
 """
     DiagnosticVariable
@@ -123,3 +143,129 @@ end
 
 # Do you want to define more diagnostics? Add them here
 include("core_diagnostics.jl")
+
+
+# ScheduledDiagnostics
+
+# NOTE: The definitions of ScheduledDiagnosticTime and ScheduledDiagnosticIterations are
+# nearly identical except for the fact that one is assumed to use units of seconds the other
+# units of integration steps. However, we allow for this little repetition of code to avoid
+# adding an extra layer of abstraction just to deal with these two objects (some people say
+# that "duplication is better than over-abstraction"). Most users will only work with
+# ScheduledDiagnosticTime. (It would be nice to have defaults fields in abstract types, as
+# proposed in 2013 in https://github.com/JuliaLang/julia/issues/4935) Having two distinct
+# types allow us to implement different checks and behaviors (e.g., we allow
+# ScheduledDiagnosticTime to have placeholders values for {compute, output}_every so that we
+# can plug the timestep in it).
+
+struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
+    variable::DiagnosticVariable
+    output_every::T1
+    output_writer::OW
+    reduction_time_func::F1
+    reduction_space_func::F2
+    compute_every::T2
+    pre_output_hook!::PO
+
+    """
+        ScheduledDiagnosticIterations(; variable::DiagnosticVariable,
+                                        output_every,
+                                        output_writer,
+                                        reduction_time_func = nothing,
+                                        reduction_space_func = nothing,
+                                        compute_every = isa_reduction ? 1 : output_every,
+                                        pre_output_hook! = (accum, count) -> nothing)
+
+
+    A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence
+    defined by the number of iterations, with an optional reduction applied to it (e.g., compute
+    the maximum temperature over the course of every 10 timesteps). This object is turned into
+    two callbacks (one for computing and the other for output) and executed by the integrator.
+
+    Keyword arguments
+    =================
+
+    - `variable`: The diagnostic variable that has to be computed and output.
+
+    - `output_every`: Save the results to disk every `output_every` iterations.
+
+    - `output_writer`: Function that controls out to save the computed diagnostic variable to
+                       disk. `output_writer` has to take three arguments: the value that has to
+                       be output, the `ScheduledDiagnostic`, and the integrator. Internally, the
+                       integrator contains extra information (such as the current timestep). It
+                       is responsibility of the `output_writer` to properly use the provided
+                       information for meaningful output.
+
+    - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of
+                             scratch space `acc` where to accumulate partial results. Then, at
+                             every `compute_every`, `reduction_time_func` is computed between
+                             the previously stored value in `acc` and the new value. This
+                             implements a running reduction. For example, if
+                             `reduction_time_func = max`, the space `acc` will hold the running
+                             maxima of the diagnostic. To implement operations like the
+                             arithmetic average, the `reduction_time_func` has to be chosen as
+                             `sum`, and a `pre_output_hook!` that renormalize `acc` by the
+                             number of samples has to be provided. For custom reductions, it is
+                             necessary to also specify the identity of operation by defining a
+                             new method to `identity_of_reduction`.
+
+    - `reduction_space_func`: NOT IMPLEMENTED YET
+
+    - `compute_every`: Run the computations every `compute_every` iterations. This is not
+                       particularly useful for point-wise diagnostics, where we enforce that
+                       `compute_every` = `output_every`. For time reductions, `compute_every` is
+                       set to 1 (compute at every timestep) by default. `compute_every` has to
+                       evenly divide `output_every`.
+
+    - `pre_output_hook!`: Function that has to be run before saving to disk for reductions
+                          (mostly used to implement averages). The function `pre_output_hook!`
+                          is called with two arguments: the value accumulated during the
+                          reduction, and the number of times the diagnostic was computed from
+                          the last time it was output. `pre_output_hook!` should mutate the
+                          accumulator in place. The return value of `pre_output_hook!` is
+                          discarded. An example of `pre_output_hook!` to compute the arithmetic
+                          average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
+
+    """
+    function ScheduledDiagnosticIterations(;
+        variable::DiagnosticVariable,
+        output_every,
+        output_writer,
+        reduction_time_func = nothing,
+        reduction_space_func = nothing,
+        compute_every = isnothing(reduction_time_func) ? output_every : 1,
+        pre_output_hook! = (accum, count) -> nothing,
+    )
+
+        # We provide an inner constructor to enforce some constraints
+
+        output_every % compute_every == 0 || error(
+            "output_every should be multiple of compute_every for variable $(variable.long_name)",
+        )
+
+        isa_reduction = !isnothing(reduction_time_func)
+
+        # If it is not a reduction, we compute only when we output
+        if !isa_reduction && compute_every != output_every
+            @warn "output_every != compute_every for $(variable.long_name), changing compute_every to match"
+            compute_every = output_every
+        end
+
+        T1 = typeof(output_every)
+        T2 = typeof(compute_every)
+        OW = typeof(output_writer)
+        F1 = typeof(reduction_time_func)
+        F2 = typeof(reduction_space_func)
+        PO = typeof(pre_output_hook!)
+
+        new{T1, T2, OW, F1, F2, PO}(
+            variable,
+            output_every,
+            output_writer,
+            reduction_time_func,
+            reduction_space_func,
+            compute_every,
+            pre_output_hook!,
+        )
+    end
+end

From 4ff287f4cc38c1a543c85a527560f9f4fc1600c1 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 14:09:42 -0700
Subject: [PATCH 04/73] Add reduction_identities.jl

---
 src/diagnostics/Diagnostics.jl          |  7 +++++++
 src/diagnostics/reduction_identities.jl | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 src/diagnostics/reduction_identities.jl

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 9f976de1e71..59ce0d48110 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -32,6 +32,10 @@
 #   relationship between the periods and the timesteps. In some sense, you can think of
 #   ScheduledDiagnosticIterations as an internal representation and ScheduledDiagnosticTime
 #   as the external interface.
+#
+# - This file also also include several other files, including (but not limited to):
+#   - core_diagnostics.jl
+#   - reduction_identities.jl
 
 """
     DiagnosticVariable
@@ -269,3 +273,6 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         )
     end
 end
+
+# We define all the known identities in reduction_identities.jl
+include("reduction_identities.jl")
diff --git a/src/diagnostics/reduction_identities.jl b/src/diagnostics/reduction_identities.jl
new file mode 100644
index 00000000000..e75d15957b5
--- /dev/null
+++ b/src/diagnostics/reduction_identities.jl
@@ -0,0 +1,19 @@
+# This file defines group identities (technically _monoid identities_) and is included in
+# Diagnostics.jl.
+#
+# Several diagnostics require performing reductions, such as taking the maximum or the
+# average. Since it is not feasible to store all the lists of all the intermediate values,
+# we aggregate the results in specific storage areas (e.g., we take
+# max(max(max(max(t1, t2), t3), t4), t5) instead of max(t1, t2, t3, t4, t5)
+# In this, it is convenient to preallocate the space where we want to accumulate the
+# intermediate. However, we have to fill the space with something that does not affect the
+# reduction. This, by definition, is the identity of the operation. The identity of the operation
+# + is 0 because x + 0 = x for every x.
+#
+# We have to know the identity for every operation we want to support. Of course, users are
+# welcome to define their own by adding new methods to identity_of_reduction.
+
+identity_of_reduction(::Val{max}) = -Inf
+identity_of_reduction(::Val{min}) = +Inf
+identity_of_reduction(::Val{+}) = 0
+identity_of_reduction(::Val{*}) = 1

From f6ad16405220d972ff659b38badf0464df4af1f0 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 13:47:04 -0700
Subject: [PATCH 05/73] Add ScheduledDiagnosticTime

---
 src/diagnostics/Diagnostics.jl | 122 +++++++++++++++++++++++++++++++++
 1 file changed, 122 insertions(+)

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 59ce0d48110..99e6eec9a82 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -274,5 +274,127 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     end
 end
 
+
+struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
+    variable::DiagnosticVariable
+    output_every::T1
+    output_writer::OW
+    reduction_time_func::F1
+    reduction_space_func::F2
+    compute_every::T2
+    pre_output_hook!::PO
+
+    """
+        ScheduledDiagnosticTime(; variable::DiagnosticVariable,
+                                            output_every,
+                                            output_writer,
+                                            reduction_time_func = nothing,
+                                            reduction_space_func = nothing,
+                                            compute_every = isa_reduction ? :timestep : output_every,
+                                            pre_output_hook! = (accum, count) -> nothing)
+
+
+    A `DiagnosticVariable` that has to be computed and output during a simulation with a
+    cadence defined by how many seconds in simulation time, with an optional reduction
+    applied to it (e.g., compute the maximum temperature over the course of every day). This
+    object is turned into a `ScheduledDiagnosticIterations`, which is turned into two
+    callbacks (one for computing and the other for output) and executed by the integrator.
+
+    Keyword arguments
+    =================
+
+    - `variable`: The diagnostic variable that has to be computed and output.
+
+    - `output_every`: Save the results to disk every `output_every` seconds.
+
+    - `output_writer`: Function that controls out to save the computed diagnostic variable to
+                       disk. `output_writer` has to take three arguments: the value that has to
+                       be output, the `ScheduledDiagnostic`, and the integrator. Internally, the
+                       integrator contains extra information (such as the current timestep). It
+                       is responsibility of the `output_writer` to properly use the provided
+                       information for meaningful output.
+
+    - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of
+                             scratch space `acc` where to accumulate partial results. Then, at
+                             every `compute_every`, `reduction_time_func` is computed between
+                             the previously stored value in `acc` and the new value. This
+                             implements a running reduction. For example, if
+                             `reduction_time_func = max`, the space `acc` will hold the running
+                             maxima of the diagnostic. To implement operations like the
+                             arithmetic average, the `reduction_time_func` has to be chosen as
+                             `sum`, and a `pre_output_hook!` that renormalize `acc` by the
+                             number of samples has to be provided. For custom reductions, it is
+                             necessary to also specify the identity of operation by defining a
+                             new method to `identity_of_reduction`.
+
+    - `reduction_space_func`: NOT IMPLEMENTED YET
+
+    - `compute_every`: Run the computations every `compute_every` seconds. This is not
+                       particularly useful for point-wise diagnostics, where we enforce that
+                       `compute_every` = `output_every`. For time reductions,
+                       `compute_every` is set to `:timestep` (compute at every timestep) by
+                       default. `compute_every` has to evenly divide `output_every`.
+                       `compute_every` can take the special symbol `:timestep` which is a
+                       placeholder for the timestep of the simulation to which this
+                       `ScheduledDiagnostic` is attached.
+
+    - `pre_output_hook!`: Function that has to be run before saving to disk for reductions
+                          (mostly used to implement averages). The function `pre_output_hook!`
+                          is called with two arguments: the value accumulated during the
+                          reduction, and the number of times the diagnostic was computed from
+                          the last time it was output. `pre_output_hook!` should mutate the
+                          accumulator in place. The return value of `pre_output_hook!` is
+                          discarded. An example of `pre_output_hook!` to compute the arithmetic
+                          average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
+
+    """
+    function ScheduledDiagnosticTime(;
+        variable::DiagnosticVariable,
+        output_every,
+        output_writer,
+        reduction_time_func = nothing,
+        reduction_space_func = nothing,
+        compute_every = isnothing(reduction_time_func) ? output_every :
+                        :timestep,
+        pre_output_hook! = (accum, count) -> nothing,
+    )
+
+        # We provide an inner constructor to enforce some constraints
+
+        # compute_every could be a Symbol (:timestep). We process this that when we process
+        # the list of diagnostics
+        if !isa(compute_every, Symbol)
+            output_every % compute_every == 0 || error(
+                "output_every should be multiple of compute_every for variable $(variable.long_name)",
+            )
+        end
+
+        isa_reduction = !isnothing(reduction_time_func)
+
+        # If it is not a reduction, we compute only when we output
+        if !isa_reduction && compute_every != output_every
+            @warn "output_every != compute_every for $(variable.long_name), changing compute_every to match"
+            compute_every = output_every
+        end
+
+        T1 = typeof(output_every)
+        T2 = typeof(compute_every)
+        OW = typeof(output_writer)
+        F1 = typeof(reduction_time_func)
+        F2 = typeof(reduction_space_func)
+        PO = typeof(pre_output_hook!)
+
+        new{T1, T2, OW, F1, F2, PO}(
+            variable,
+            output_every,
+            output_writer,
+            reduction_time_func,
+            reduction_space_func,
+            compute_every,
+            pre_output_hook!,
+        )
+    end
+end
+
 # We define all the known identities in reduction_identities.jl
 include("reduction_identities.jl")

From 705d732486d72bca36a86dad6adf03374fd5c0f5 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 13:53:06 -0700
Subject: [PATCH 06/73] Add more constructors for ScheduledDiagnostics

---
 docs/src/diagnostics.md        |  4 +++
 src/diagnostics/Diagnostics.jl | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index e5aef1b82af..f525a389552 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -87,6 +87,10 @@ the end of every integration step. This is particularly convenient when defining
 default diagnostics, as they should be largely independent on the choice of the
 specific simulation being run (and its timestep).
 
+Given a timestep `dt`, a `ScheduledDiagnosticIterations` can be obtained from a
+`ScheduledDiagnosticTime` `sd` simply by calling
+``ScheduledDiagnosticIterations(sd, dt)`.
+
 ## I want to add a new diagnostic variable
 
 Diagnostic variables are represented in `ClimaAtmos` with a `DiagnosticVariable`
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 99e6eec9a82..1d01729e4b7 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -396,5 +396,58 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
     end
 end
 
+"""
+    ScheduledDiagnosticIterations(sd_time::ScheduledDiagnosticTime, Δt)
+
+
+Create a `ScheduledDiagnosticIterations` given a `ScheduledDiagnosticTime` and a timestep
+`Δt`. In this, ensure that `compute_every` and `output_every` are meaningful for the given
+timestep.
+
+"""
+
+function ScheduledDiagnosticIterations(
+    sd_time::ScheduledDiagnosticTime,
+    Δt::T,
+) where {T}
+
+    # If we have the timestep, we can convert time in seconds into iterations
+
+    # if compute_every is :timestep, then we want to compute after every iterations
+    compute_every =
+        sd_time.compute_every == :timestep ? 1 : sd_time.compute_every / Δt
+    output_every = sd_time.output_every / Δt
+
+    isinteger(output_every) || error(
+        "output_every should be multiple of the timestep for variable $(sd_time.variable.long_name)",
+    )
+    isinteger(compute_every) || error(
+        "compute_every should be multiple of the timestep for variable $(sd_time.variable.long_name)",
+    )
+
+    ScheduledDiagnosticIterations(;
+        sd_time.variable,
+        output_every = convert(Int, output_every),
+        sd_time.output_writer,
+        sd_time.reduction_time_func,
+        sd_time.reduction_space_func,
+        compute_every = convert(Int, compute_every),
+        sd_time.pre_output_hook!,
+    )
+end
+
+# We provide also a companion constructor for ScheduledDiagnosticIterations which returns
+# itself (without copy) when called with a timestep.
+#
+# This is so that we can assume that
+# ScheduledDiagnosticIterations(ScheduledDiagnostic{Time, Iterations}, Δt)
+# always returns a valid ScheduledDiagnosticIterations
+ScheduledDiagnosticIterations(
+    sd::ScheduledDiagnosticIterations,
+    _Δt::T,
+) where {T} = sd
+
 # We define all the known identities in reduction_identities.jl
 include("reduction_identities.jl")
+
+

From 034c3a64ff48c6dc16d03eff65acace0cbaa616f Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 14:24:12 -0700
Subject: [PATCH 07/73] Add get_callbacks_from_diagnostics

---
 src/diagnostics/Diagnostics.jl | 109 +++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 1d01729e4b7..a98808c5355 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -33,6 +33,12 @@
 #   ScheduledDiagnosticIterations as an internal representation and ScheduledDiagnosticTime
 #   as the external interface.
 #
+# - A function to convert a list of ScheduledDiagnosticIterations into a list of
+#   AtmosCallbacks. This function takes three arguments: the list of diagnostics and two
+#   dictionaries that map each scheduled diagnostic to an area of memory where to save the
+#   result and where to keep track of how many times the function was called (so that we
+#   can compute stuff like averages).
+#
 # - This file also also include several other files, including (but not limited to):
 #   - core_diagnostics.jl
 #   - reduction_identities.jl
@@ -451,3 +457,106 @@ ScheduledDiagnosticIterations(
 include("reduction_identities.jl")
 
 
+
+"""
+    get_callbacks_from_diagnostics(diagnostics, storage, counters)
+
+
+Translate a list of diagnostics into a list of callbacks.
+
+Positional arguments
+=====================
+
+- `diagnostics`: List of `ScheduledDiagnosticIterations` that have to be converted to
+                 callbacks. We want to have `ScheduledDiagnosticIterations` here so that we
+                 can define callbacks that occur at the end of every N integration steps.
+
+- `storage`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially
+             pre-allocated area of memory where to accumulate/save results.
+
+- `counters`: Dictionary that maps a given `ScheduledDiagnosticIterations` to the counter
+              that tracks how many times the given diagnostics was computed from the last
+              time it was output to disk.
+
+"""
+function get_callbacks_from_diagnostics(diagnostics, storage, counters)
+    # We have two types of callbacks: to compute and accumulate diagnostics, and to dump
+    # them to disk. Note that our callbacks do not contain any branching
+
+    # storage is used to pre-allocate memory and to accumulate partial results for those
+    # diagnostics that perform reductions.
+
+    callbacks = Any[]
+
+    for diag in diagnostics
+        variable = diag.variable
+        isa_reduction = !isnothing(diag.reduction_time_func)
+
+        # reduction is used below. If we are not given a reduction_time_func, we just want
+        # to move the computed quantity to its storage (so, we return the second argument,
+        # which that will be the newly computed one). If we have a reduction, we apply it
+        # point-wise
+        reduction = isa_reduction ? diag.reduction_time_func : (_, y) -> y
+
+        # If we have a reduction, we have to reset the accumulator to its neutral state. (If
+        # we don't have a reduction, we don't have to do anything)
+        #
+        # ClimaAtmos defines methods for identity_of_reduction for standard
+        # reduction_time_func in reduction_identities.jl
+        reset_accumulator! =
+            isa_reduction ?
+            () -> begin
+                # identity_of_reduction works by dispatching over Val{operation}
+                identity =
+                    identity_of_reduction(Val(diag.reduction_time_func))
+                # We also need to make sure that we are consistent with the types
+                float_type = eltype(storage[diag])
+                identity_ft = convert(float_type, identity)
+                storage[diag] .= identity_ft
+            end : () -> nothing
+
+        compute_callback =
+            integrator -> begin
+                # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
+                value = variable.compute_from_integrator(integrator, nothing)
+                storage[diag] .= reduction.(storage[diag], value)
+                counters[diag] += 1
+                return nothing
+            end
+
+        output_callback =
+            integrator -> begin
+                # Any operations we have to perform before writing to output?
+                # Here is where we would divide by N to obtain an arithmetic average
+                diag.pre_output_hook!(storage[diag], counters[diag])
+
+                # Write to disk
+                diag.output_writer(storage[diag], diag, integrator)
+
+                reset_accumulator!()
+                counters[diag] = 0
+                return nothing
+            end
+
+        # Here we have skip_first = true. This is important because we are going to manually
+        # call all the callbacks so that we can verify that they are meaningful for the
+        # model under consideration (and they don't have bugs).
+        append!(
+            callbacks,
+            [
+                call_every_n_steps(
+                    compute_callback,
+                    diag.compute_every,
+                    skip_first = true,
+                ),
+                call_every_n_steps(
+                    output_callback,
+                    diag.output_every,
+                    skip_first = true,
+                ),
+            ],
+        )
+    end
+
+    return callbacks
+end

From 33c690ca2f7160d6a29f21971387eba78223c1ed Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 15:01:43 -0700
Subject: [PATCH 08/73] Add HDF5Writer

---
 docs/src/diagnostics.md    |  8 ++++--
 src/ClimaAtmos.jl          |  1 +
 src/diagnostics/Writers.jl | 58 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 src/diagnostics/Writers.jl

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index f525a389552..647f0ae3bdc 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -61,8 +61,12 @@ integrator. Internally, the integrator contains extra information (such as the
 current timestep), so the `output_writer` can implement arbitrarily complex
 behaviors. It is responsibility of the `output_writer` to properly use the
 provided information for meaningful output. `ClimaAtmos` provides functions that
-return `output_writer`s for standard operations (e.g., for writing to HDF5
-files).
+return `output_writer`s for standard operations. The main one is currently
+`HDF5Writer`, which should be enough for most use cases. To use it, just
+initialize a `ClimaAtmos.HDF5Writer` object with your choice of configuration
+and pass it as a `output_writer` argument to the `ScheduledDiagnostic`. More
+information about the options supported by `ClimaAtmos.HDF5Writer` is available
+in its constructor.
 
 There are two flavors of `ScheduledDiagnostic`s:
 `ScheduledDiagnosticIterations`, and `ScheduledDiagnosticTime`. The main
diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
index 1f612e69a7f..705b19a7bf0 100644
--- a/src/ClimaAtmos.jl
+++ b/src/ClimaAtmos.jl
@@ -118,5 +118,6 @@ include(joinpath("solver", "solve.jl"))
 include(joinpath("parameters", "create_parameters.jl"))
 
 include(joinpath("diagnostics", "Diagnostics.jl"))
+include(joinpath("diagnostics", "Writers.jl"))
 
 end # module
diff --git a/src/diagnostics/Writers.jl b/src/diagnostics/Writers.jl
new file mode 100644
index 00000000000..ab5963c56ac
--- /dev/null
+++ b/src/diagnostics/Writers.jl
@@ -0,0 +1,58 @@
+# Writers.jl
+#
+# This file defines function-generating functions for output_writers for diagnostics. The
+# writers come with opinionated defaults.
+
+
+"""
+    HDF5Writer()
+
+
+Save a `ScheduledDiagnostic` to a HDF5 file inside the `output_dir` of the simulation.
+
+
+TODO: This is a very barebone HDF5Writer.
+
+We need to implement the following features/options:
+- Toggle for write new files/append
+- Checks for existing files
+- Check for new subfolders that have to be created
+- More meaningful naming conventions (keeping in mind that we can have multiple variables
+  with different reductions)
+- All variables in one file/each variable in its own file
+- All timesteps in one file/each timestep in its own file
+- Writing the correct attributes
+- Overriding simulation.output_dir (e.g., if the path starts with /)
+- ...more features/options
+
+"""
+
+function HDF5Writer()
+    # output_drivers are called with the three arguments: the value, the ScheduledDiagnostic,
+    # and the integrator
+    function write_to_hdf5(value, diagnostic, integrator)
+        var = diagnostic.variable
+        time = lpad(integrator.t, 10, "0")
+        red = isnothing(diagnostic.reduction_time_func) ? "" : diagnostic.reduction_time_func
+
+        output_path = joinpath(
+            integrator.p.simulation.output_dir,
+            "$(var.short_name)_$(red)_$(time).h5",
+        )
+
+        hdfwriter = InputOutput.HDF5Writer(output_path, integrator.p.comms_ctx)
+        InputOutput.HDF5.write_attribute(hdfwriter.file, "time", time)
+        InputOutput.HDF5.write_attribute(
+            hdfwriter.file,
+            "long_name",
+            var.long_name,
+        )
+        InputOutput.write!(
+            hdfwriter,
+            Fields.FieldVector(; Symbol(var.short_name) => value),
+            "diagnostics",
+        )
+        Base.close(hdfwriter)
+        return nothing
+    end
+end

From 29bdb28182b24569c7b1bed92503872a28b86b2c Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <sbozzolator@gmail.com>
Date: Fri, 1 Sep 2023 15:44:33 -0700
Subject: [PATCH 09/73] Add default_diagnostics.jl

---
 docs/src/diagnostics.md                    | 40 ++++++++++++++++
 src/diagnostics/Diagnostics.jl             |  3 ++
 src/diagnostics/default_diagnostics.jl     | 53 ++++++++++++++++++++++
 src/diagnostics/defaults/moisture_model.jl | 18 ++++++++
 4 files changed, 114 insertions(+)
 create mode 100644 src/diagnostics/default_diagnostics.jl
 create mode 100644 src/diagnostics/defaults/moisture_model.jl

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 647f0ae3bdc..5929b2c1c64 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -2,6 +2,46 @@
 
 ## I want to compute and output a diagnostic variable
 
+### From a script
+
+The simplest way to get started with diagnostics is to use the defaults for your
+atmospheric model. `ClimaAtmos` defines a function `get_default_diagnostic`. You
+can execute this function on an `AtmosModel` or on any of its fields to obtain a
+list of diagnostics ready to be passed to the simulation. So, for example
+
+```julia
+
+model = ClimaAtmos.AtmosModel(..., moisture_model = ClimaAtmos.DryModel(), ...)
+
+diagnostics = ClimaAtmos.get_default_diagnostics(model)
+# => List of diagnostics that include the ones specified for the DryModel
+```
+
+Technically, the diagnostics are represented as `ScheduledDiagnostic` objects,
+which contain information about what variable has to be computed, how often,
+where to save it, and so on (read below for more information on this). You can
+construct your own lists of `ScheduledDiagnostic`s starting from the variables
+defined by `ClimaAtmos`. The diagnostics that `ClimaAtmos` knows how to compute
+are collected in a global dictionary called `ALL_DIAGNOSTICS`. The variables in
+`ALL_DIAGNOSTICS` are identified with a long and unique name, so that you can
+access them directly. One way to do so is by using the provided convenience
+functions for common operations, e.g., continuing the previous example
+
+```julia
+
+push!(diagnostics, get_daily_max("air_density", "air_temperature"))
+```
+
+Now `diagnostics` will also contain the instructions to compute the daily
+maximum of `air_density` and `air_temperature`.
+
+**TODO: Add link to table with known diagnostics**
+
+If you are using `ClimaAtmos` with a script-based interface, you have access to
+the complete flexibility in your diagnostics. Read the section about the
+low-level interface to see how to implement custom diagnostics, reductions, or
+writers.
+
 ### The low-level interface
 
 Diagnostics are computed and output through callbacks to the main integrator.
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index a98808c5355..ce3b3e6d929 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -41,6 +41,7 @@
 #
 # - This file also also include several other files, including (but not limited to):
 #   - core_diagnostics.jl
+#   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
 #   - reduction_identities.jl
 
 """
@@ -154,6 +155,8 @@ end
 # Do you want to define more diagnostics? Add them here
 include("core_diagnostics.jl")
 
+# Default diagnostics and higher level interfaces
+include("default_diagnostics.jl")
 
 # ScheduledDiagnostics
 
diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
new file mode 100644
index 00000000000..b3b534a57a6
--- /dev/null
+++ b/src/diagnostics/default_diagnostics.jl
@@ -0,0 +1,53 @@
+# This file is included by Diagnostics.jl and defines all the defaults for various models. A
+# model here is either a global AtmosModel, or small (sub)models (e.g., DryModel()).
+#
+# If you are developing new models, add your defaults here. If you want to add more high
+# level interfaces, add them here. Feel free to include extra files.
+
+"""
+    get_default_diagnostics(model)
+
+
+Return a list of `ScheduledDiagnostic`s associated with the given `model`.
+
+"""
+function get_default_diagnostics(model::AtmosModel)
+    # TODO: Probably not the most elegant way to do this...
+    defaults = Any[]
+
+    for field in fieldnames(AtmosModel)
+        def_model = get_default_diagnostics(getfield(model, field))
+        append!(defaults, def_model)
+    end
+
+    return defaults
+end
+
+# Base case: if we call get_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 get_default_diagnostics return the same type). This is used by
+# get_default_diagnostics(model::AtmosModel), so that we can ignore defaults for submodels
+# that have no given defaults.
+get_default_diagnostics(_) = []
+
+"""
+    get_daily_max(long_names...; output_writer = HDF5Writer())
+
+
+Return a list of `ScheduledDiagnostics` that compute the daily max for the given variables.
+"""
+function get_daily_max(long_names...; output_writer = HDF5Writer())
+    # TODO: Add mechanism to print out reasonable error on variables that are not in ALL_DIAGNOSTICS
+    return [
+        ScheduledDiagnosticTime(
+            variable = ALL_DIAGNOSTICS[long_name],
+            compute_every = :timestep,
+            output_every = 86400, # seconds
+            reduction_time_func = max,
+            output_writer = output_writer,
+        ) for long_name in long_names
+    ]
+end
+
+# Include all the subdefaults
+include("defaults/moisture_model.jl")
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
new file mode 100644
index 00000000000..14e55e664c2
--- /dev/null
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -0,0 +1,18 @@
+# FIXME: Gabriele added this as an example. Put something meaningful here!
+function get_default_diagnostics(::T) where {T <: DryModel}
+    return [
+        ScheduledDiagnosticTime(
+            variable = ALL_DIAGNOSTICS["air_density"],
+            compute_every = :timestep,
+            output_every = 86400, # seconds
+            reduction_time_func = max,
+            output_writer = HDF5Writer(),
+        ),
+        ScheduledDiagnosticIterations(
+            variable = ALL_DIAGNOSTICS["air_density"],
+            compute_every = 1,
+            output_every = 1, # iteration
+            output_writer = HDF5Writer(),
+        ),
+    ]
+end

From cd0eb75c8e7c391c4260d41506d42878572a6331 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Tue, 5 Sep 2023 08:40:33 -0700
Subject: [PATCH 10/73] Add design note of pre_output_hook!

This note is mostly to explain why pre_output_hook! is there and why it
takes the argument that it takes. Hopefully, it will inform future
developers interested in expanding this section of the code.
---
 src/diagnostics/Diagnostics.jl | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index ce3b3e6d929..b9bd5647262 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -171,6 +171,19 @@ include("default_diagnostics.jl")
 # ScheduledDiagnosticTime to have placeholders values for {compute, output}_every so that we
 # can plug the timestep in it).
 
+# Design note: pre_output_hook!
+#
+# One of our key requirements is to be able to compute arithmetic averages. Unfortunately,
+# computing an arithmetic average requires keeping track of how many elements we are summing
+# up. pre_output_hook! was introduced so that we can think of an average as a sum coupled
+# with division, and perform the division (by the number of elements) before output.
+# pre_output_hook! could be used for other operations, but we decided to keep it simple and
+# target directly the most important use case for us.
+#
+# This choice restricts what reductions can be performed. For example, it is not possible to
+# have a geometric average. If more complex reduction are needed, this mechanism has to be
+# changed.
+
 struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     variable::DiagnosticVariable
     output_every::T1

From b3c4ef1af3cac9f66bb15c2e5cc8f31f4a59375b Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 11:20:19 -0700
Subject: [PATCH 11/73] Add basic support for diagnostics

---
 src/solver/type_getters.jl | 56 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 54 insertions(+), 2 deletions(-)

diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 09a2e116049..6429cb22f77 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -533,10 +533,9 @@ function get_callbacks(parsed_args, simulation, atmos, params)
         )
     end
 
-    return SciMLBase.CallbackSet(callbacks...)
+    return callbacks
 end
 
-
 function get_cache(
     Y,
     parsed_args,
@@ -742,8 +741,44 @@ function get_integrator(config::AtmosConfig)
         callback = get_callbacks(config.parsed_args, simulation, atmos, params)
     end
     @info "get_callbacks: $s"
+
+    # Initialize diagnostics
+    @info "Initializing diagnostics"
+
+    diagnostics = get_default_diagnostics(atmos)
+
+    # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,
+    # ensuring that there is consistency in the timestep and the periods and translating
+    # those periods that depended on the timestep
+    diagnostics_iterations =
+        [ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics]
+
+    # For diagnostics that perform reductions, the storage is used as an accumulator, for
+    # the other ones it is still defined to avoid allocating new space every time.
+    diagnostic_storage = Dict()
+    diagnostic_counters = Dict()
+
+    # NOTE: The diagnostics_callbacks are not called at the initial timestep
+    diagnostics_callbacks = get_callbacks_from_diagnostics(
+        diagnostics_iterations,
+        diagnostic_storage,
+        diagnostic_counters,
+    )
+
+    # We need to ensure the precomputed quantities are indeed precomputed
+    # TODO: Remove this when we can assume that the precomputed_quantities are in sync with the state
+    sync_precomputed = call_every_n_steps(
+        (int) -> set_precomputed_quantities!(int.u, int.p, int.t),
+    )
+
+    callback = SciMLBase.CallbackSet(
+        callback...,
+        sync_precomputed,
+        diagnostics_callbacks...,
+    )
     @info "n_steps_per_cycle_per_cb: $(n_steps_per_cycle_per_cb(callback, simulation.dt))"
     @info "n_steps_per_cycle: $(n_steps_per_cycle(callback, simulation.dt))"
+
     tspan = (t_start, simulation.t_end)
     s = @timed_str begin
         integrator_args, integrator_kwargs = args_integrator(
@@ -760,5 +795,22 @@ function get_integrator(config::AtmosConfig)
         integrator = SciMLBase.init(integrator_args...; integrator_kwargs...)
     end
     @info "init integrator: $s"
+
+    for diag in diagnostics_iterations
+        variable = diag.variable
+        try
+            # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+            diagnostic_storage[diag] =
+                variable.compute_from_integrator(integrator, nothing)
+            diagnostic_counters[diag] = 1
+            # If it is not a reduction, call the output writer as well
+            if isnothing(diag.reduction_time_func)
+                diag.output_writer(diagnostic_storage[diag], diag, integrator)
+            end
+        catch e
+            error("Could not compute diagnostic $(variable.long_name): $e")
+        end
+    end
+
     return integrator
 end

From d57ed17651493c07d145028201078308e7c7f48a Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 6 Sep 2023 08:22:24 -0700
Subject: [PATCH 12/73] Add constructor for SDTime from SDIterations

---
 src/diagnostics/Diagnostics.jl | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index b9bd5647262..43596c50955 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -458,6 +458,37 @@ function ScheduledDiagnosticIterations(
     )
 end
 
+"""
+    ScheduledDiagnosticTime(sd_time::ScheduledDiagnosticIterations, Δt)
+
+
+Create a `ScheduledDiagnosticTime` given a `ScheduledDiagnosticIterations` and a timestep
+`Δt`.
+
+"""
+function ScheduledDiagnosticTime(
+    sd_time::ScheduledDiagnosticIterations,
+    Δt::T,
+) where {T}
+
+    # If we have the timestep, we can convert time in iterations to seconds
+
+    # if compute_every is :timestep, then we want to compute after every iterations
+    compute_every =
+        sd_time.compute_every == 1 ? :timestep : sd_time.compute_every * Δt
+    output_every = sd_time.output_every * Δt
+
+    ScheduledDiagnosticTime(;
+        sd_time.variable,
+        output_every,
+        sd_time.output_writer,
+        sd_time.reduction_time_func,
+        sd_time.reduction_space_func,
+        compute_every,
+        sd_time.pre_output_hook!,
+    )
+end
+
 # We provide also a companion constructor for ScheduledDiagnosticIterations which returns
 # itself (without copy) when called with a timestep.
 #

From 16beb0b73fa4357f7c2a193273f6e2bb2375db21 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 6 Sep 2023 08:36:39 -0700
Subject: [PATCH 13/73] Add descriptive_name for ScheduledDiagnostics

---
 src/diagnostics/Writers.jl | 83 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/src/diagnostics/Writers.jl b/src/diagnostics/Writers.jl
index ab5963c56ac..d4707af8987 100644
--- a/src/diagnostics/Writers.jl
+++ b/src/diagnostics/Writers.jl
@@ -4,6 +4,89 @@
 # writers come with opinionated defaults.
 
 
+"""
+    descriptive_name(sd_t::ScheduledDiagnosticTime)
+
+
+Return a compact, unique-ish, identifier for the given `ScheduledDiagnosticTime` `sd_t`.
+
+We split the period in seconds into days, hours, minutes, seconds. In most cases
+(with standard periods), this output will look something like
+`air_density_1d_max`. This function can be used for filenames.
+
+The name is not unique to the `ScheduledDiagnosticTime` because it ignores other
+parameters such as whether there is a reduction in space or the compute
+frequency.
+
+"""
+function descriptive_name(sd_t::ScheduledDiagnosticTime)
+    var = "$(sd_t.variable.short_name)"
+
+    isa_reduction = !isnothing(sd_t.reduction_time_func)
+
+    if isa_reduction
+        red = "$(sd_t.reduction_time_func)"
+
+        # Convert period from seconds to days, hours, minutes, seconds
+        period = ""
+
+        days, rem_seconds = divrem(sd_t.output_every, 24 * 60 * 60)
+        hours, rem_seconds = divrem(rem_seconds, 60 * 60)
+        minutes, seconds = divrem(rem_seconds, 60)
+
+        if days > 0
+            period *= "$(days)d_"
+        end
+        if hours > 0
+            period *= "$(hours)h_"
+        end
+        if minutes > 0
+            period *= "$(minutes)m_"
+        end
+        if seconds > 0
+            period *= "$(seconds)s_"
+        end
+
+        suffix = period * red
+    else
+        # Not a reduction
+        suffix = "inst"
+    end
+    return "$(var)_$(suffix)"
+end
+
+"""
+    descriptive_name(sd_i::ScheduledDiagnosticIterations, [Δt])
+
+
+Return a compact, unique-ish, identifier for the given
+`ScheduledDiagnosticIterations` `sd_i`.
+
+If the timestep `Δt` is provided, convert the steps into seconds. In this case,
+the output will look like `air_density_1d_max`. Otherwise, the output will look
+like `air_density_100it_max`. This function can be used for filenames.
+
+The name is not unique to the `ScheduledDiagnosticIterations` because it ignores
+other parameters such as whether there is a reduction in space or the compute
+frequency.
+
+"""
+function descriptive_name(sd_i::ScheduledDiagnosticIterations,
+                          Δt = nothing)
+
+    if !isnothing(Δt)
+        # Convert iterations into time
+        return descriptive_name(ScheduledDiagnosticTime(sd_i, Δt))
+    else
+        var = "$(sd_i.variable.short_name)"
+        suffix =
+            isnothing(sd_i.reduction_time_func) ? "inst" :
+            "$(sd_i.output_every)it_(sd_i.reduction_time_func)"
+        return "$(var)_$(suffix)"
+    end
+end
+
+
 """
     HDF5Writer()
 

From 800a4839a93d784aa6c2665cd5e34f609f961fe1 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 6 Sep 2023 08:36:43 -0700
Subject: [PATCH 14/73] Use descriptive_name for filenames and errors

---
 src/diagnostics/Diagnostics.jl       | 35 +++++++++---
 src/diagnostics/Writers.jl           | 83 ++++++++++------------------
 src/diagnostics/diagnostics_utils.jl | 61 ++++++++++++++++++++
 3 files changed, 115 insertions(+), 64 deletions(-)
 create mode 100644 src/diagnostics/diagnostics_utils.jl

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 43596c50955..a033be0870b 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -43,6 +43,7 @@
 #   - core_diagnostics.jl
 #   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
 #   - reduction_identities.jl
+#   - diagnostic_utils.jl
 
 """
     DiagnosticVariable
@@ -138,7 +139,7 @@ function add_diagnostic_variable!(;
     long_name,
     units,
     comments,
-    compute_from_integrator
+    compute_from_integrator,
 )
     haskey(ALL_DIAGNOSTICS, short_name) &&
         error("diagnostic $short_name already defined")
@@ -148,7 +149,7 @@ function add_diagnostic_variable!(;
         long_name,
         units,
         comments,
-        compute_from_integrator
+        compute_from_integrator,
     )
 end
 
@@ -264,16 +265,18 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     )
 
         # We provide an inner constructor to enforce some constraints
+        descriptive_name =
+            get_descriptive_name(variable, output_every, reduction_time_func)
 
         output_every % compute_every == 0 || error(
-            "output_every should be multiple of compute_every for variable $(variable.long_name)",
+            "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
         )
 
         isa_reduction = !isnothing(reduction_time_func)
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(variable.long_name), changing compute_every to match"
+            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -382,12 +385,18 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
     )
 
         # We provide an inner constructor to enforce some constraints
+        descriptive_name = get_descriptive_name(
+            variable,
+            output_every,
+            reduction_time_func;
+            units_are_seconds = false,
+        )
 
         # compute_every could be a Symbol (:timestep). We process this that when we process
         # the list of diagnostics
         if !isa(compute_every, Symbol)
             output_every % compute_every == 0 || error(
-                "output_every should be multiple of compute_every for variable $(variable.long_name)",
+                "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
             )
         end
 
@@ -395,7 +404,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(variable.long_name), changing compute_every to match"
+            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -440,11 +449,17 @@ function ScheduledDiagnosticIterations(
         sd_time.compute_every == :timestep ? 1 : sd_time.compute_every / Δt
     output_every = sd_time.output_every / Δt
 
+    descriptive_name = get_descriptive_name(
+        sd_time.variable,
+        sd_time.output_every,
+        sd_time.reduction_time_func,
+    )
+
     isinteger(output_every) || error(
-        "output_every should be multiple of the timestep for variable $(sd_time.variable.long_name)",
+        "output_every should be multiple of the timestep for diagnostic $(descriptive_name)",
     )
     isinteger(compute_every) || error(
-        "compute_every should be multiple of the timestep for variable $(sd_time.variable.long_name)",
+        "compute_every should be multiple of the timestep for diagnostic $(descriptive_name)",
     )
 
     ScheduledDiagnosticIterations(;
@@ -500,10 +515,12 @@ ScheduledDiagnosticIterations(
     _Δt::T,
 ) where {T} = sd
 
+
 # We define all the known identities in reduction_identities.jl
 include("reduction_identities.jl")
 
-
+# Helper functions
+include("diagnostics_utils.jl")
 
 """
     get_callbacks_from_diagnostics(diagnostics, storage, counters)
diff --git a/src/diagnostics/Writers.jl b/src/diagnostics/Writers.jl
index d4707af8987..d4b0fc4a231 100644
--- a/src/diagnostics/Writers.jl
+++ b/src/diagnostics/Writers.jl
@@ -3,9 +3,8 @@
 # This file defines function-generating functions for output_writers for diagnostics. The
 # writers come with opinionated defaults.
 
-
 """
-    descriptive_name(sd_t::ScheduledDiagnosticTime)
+    get_descriptive_name(sd_t::ScheduledDiagnosticTime)
 
 
 Return a compact, unique-ish, identifier for the given `ScheduledDiagnosticTime` `sd_t`.
@@ -19,44 +18,15 @@ parameters such as whether there is a reduction in space or the compute
 frequency.
 
 """
-function descriptive_name(sd_t::ScheduledDiagnosticTime)
-    var = "$(sd_t.variable.short_name)"
-
-    isa_reduction = !isnothing(sd_t.reduction_time_func)
-
-    if isa_reduction
-        red = "$(sd_t.reduction_time_func)"
-
-        # Convert period from seconds to days, hours, minutes, seconds
-        period = ""
-
-        days, rem_seconds = divrem(sd_t.output_every, 24 * 60 * 60)
-        hours, rem_seconds = divrem(rem_seconds, 60 * 60)
-        minutes, seconds = divrem(rem_seconds, 60)
-
-        if days > 0
-            period *= "$(days)d_"
-        end
-        if hours > 0
-            period *= "$(hours)h_"
-        end
-        if minutes > 0
-            period *= "$(minutes)m_"
-        end
-        if seconds > 0
-            period *= "$(seconds)s_"
-        end
-
-        suffix = period * red
-    else
-        # Not a reduction
-        suffix = "inst"
-    end
-    return "$(var)_$(suffix)"
-end
+get_descriptive_name(sd_t::ScheduledDiagnosticTime) = get_descriptive_name(
+    sd_t.variable,
+    sd_t.output_every,
+    sd_t.reduction_time_func;
+    units_are_seconds = true,
+)
 
 """
-    descriptive_name(sd_i::ScheduledDiagnosticIterations, [Δt])
+    get_descriptive_name(sd_i::ScheduledDiagnosticIterations[, Δt])
 
 
 Return a compact, unique-ish, identifier for the given
@@ -71,20 +41,20 @@ other parameters such as whether there is a reduction in space or the compute
 frequency.
 
 """
-function descriptive_name(sd_i::ScheduledDiagnosticIterations,
-                          Δt = nothing)
-
-    if !isnothing(Δt)
-        # Convert iterations into time
-        return descriptive_name(ScheduledDiagnosticTime(sd_i, Δt))
-    else
-        var = "$(sd_i.variable.short_name)"
-        suffix =
-            isnothing(sd_i.reduction_time_func) ? "inst" :
-            "$(sd_i.output_every)it_(sd_i.reduction_time_func)"
-        return "$(var)_$(suffix)"
-    end
-end
+get_descriptive_name(sd_i::ScheduledDiagnosticIterations, Δt::Nothing) =
+    get_descriptive_name(
+        sd_t.variable,
+        sd_t.output_every,
+        sd_t.reduction_time_func;
+        units_are_seconds = false,
+    )
+get_descriptive_name(sd_i::ScheduledDiagnosticIterations, Δt::T) where {T} =
+    get_descriptive_name(
+        sd_i.variable,
+        sd_i.output_every * Δt,
+        sd_i.reduction_time_func;
+        units_are_seconds = true,
+    )
 
 
 """
@@ -115,12 +85,15 @@ function HDF5Writer()
     # and the integrator
     function write_to_hdf5(value, diagnostic, integrator)
         var = diagnostic.variable
-        time = lpad(integrator.t, 10, "0")
-        red = isnothing(diagnostic.reduction_time_func) ? "" : diagnostic.reduction_time_func
+        time = integrator.t
+
+        # diagnostic here is a ScheduledDiagnosticIteration. If we want to obtain a
+        # descriptive name (e.g., something with "daily"), we have to pass the timestep as
+        # well
 
         output_path = joinpath(
             integrator.p.simulation.output_dir,
-            "$(var.short_name)_$(red)_$(time).h5",
+            "$(get_descriptive_name(diagnostic, integrator.p.simulation.dt))_$time.h5",
         )
 
         hdfwriter = InputOutput.HDF5Writer(output_path, integrator.p.comms_ctx)
diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl
new file mode 100644
index 00000000000..9f37dedfb60
--- /dev/null
+++ b/src/diagnostics/diagnostics_utils.jl
@@ -0,0 +1,61 @@
+# diagnostic_utils.jl
+#
+# This file contains:
+# - get_descriptive_name: to condense ScheduledDiagnostic information into few characters.
+
+
+"""
+    get_descriptive_name(variable::DiagnosticVariable,
+                         output_every,
+                         reduction_time_func;
+                         units_are_seconds = true)
+
+
+Return a compact, unique-ish, identifier generated from the given information.
+
+`output_every` is interpreted as in seconds if `units_are_seconds` is `true`. Otherwise, it
+is interpreted as in units of number of iterations.
+
+This function is useful for filenames and error messages.
+
+ """
+
+function get_descriptive_name(
+    variable::DiagnosticVariable,
+    output_every,
+    reduction_time_func;
+    units_are_seconds = true,
+)
+    var = "$(variable.short_name)"
+    isa_reduction = !isnothing(reduction_time_func)
+
+    if units_are_seconds
+        if isa_reduction
+            red = "$(reduction_time_func)"
+
+            # Convert period from seconds to days, hours, minutes, seconds
+            period = ""
+
+            days, rem_seconds = divrem(output_every, 24 * 60 * 60)
+            hours, rem_seconds = divrem(rem_seconds, 60 * 60)
+            minutes, seconds = divrem(rem_seconds, 60)
+
+            days > 0 && (period *= "$(days)d_")
+            hours > 0 && (period *= "$(hours)h_")
+            minutes > 0 && (period *= "$(minutes)m_")
+            seconds > 0 && (period *= "$(seconds)s_")
+
+            suffix = period * red
+        else
+            # Not a reduction
+            suffix = "inst"
+        end
+    else
+        if isa_reduction
+            suffix = "$(output_every)it_(reduction_time_func)"
+        else
+            suffix = "inst"
+        end
+    end
+    return "$(var)_$(suffix)"
+end

From 42be42dcc8e69834a8343d8ae6154799b46736f9 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 6 Sep 2023 14:30:27 -0700
Subject: [PATCH 15/73] Add automatically generated list of diagnostics

---
 .gitignore                    |  3 +++
 docs/make.jl                  |  3 +++
 docs/make_diagnostic_table.jl | 28 ++++++++++++++++++++++++++++
 docs/src/diagnostics.md       |  3 ++-
 4 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 docs/make_diagnostic_table.jl

diff --git a/.gitignore b/.gitignore
index f1c4fd2abd4..8f3dc90f697 100644
--- a/.gitignore
+++ b/.gitignore
@@ -22,6 +22,9 @@ deps/src/
 docs/build/
 docs/site/
 
+# File generated by make_available_diagnostics.jl
+docs/src/available_diagnostics.md
+
 # File generated by Pkg, the package manager, based on a corresponding Project.toml
 # It records a fixed state of all packages used by the project. As such, it should not be
 # committed for packages, but should be committed for applications that require a static
diff --git a/docs/make.jl b/docs/make.jl
index 30fa9b15cb9..a2d06ebd901 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -8,6 +8,8 @@ disable_logging(Base.CoreLogging.Info) # Hide doctest's `@info` printing
 doctest(ClimaAtmos)
 disable_logging(Base.CoreLogging.BelowMinLevel) # Re-enable all logging
 
+include("make_diagnostic_table.jl")
+
 makedocs(
     CitationBibliography(joinpath(@__DIR__, "bibliography.bib")),
     modules = [ClimaAtmos],
@@ -27,6 +29,7 @@ makedocs(
         "Equations" => "equations.md",
         "EDMF Equations" => "edmf_equations.md",
         "Diagnostics" => "diagnostics.md",
+        "Available Diagnostics" => "available_diagnostics.md",
         "Diagnostic EDMF Equations" => "diagnostic_edmf_equations.md",
         "Gravity Wave Drag Parameterizations" => "gravity_wave.md",
         "Radiative Equilibrium" => "radiative_equilibrium.md",
diff --git a/docs/make_diagnostic_table.jl b/docs/make_diagnostic_table.jl
new file mode 100644
index 00000000000..f3aedf2e994
--- /dev/null
+++ b/docs/make_diagnostic_table.jl
@@ -0,0 +1,28 @@
+import ClimaAtmos as CA
+
+# Read all the diagnostics we know how to compute, and print them into a
+# markdown table that is later compiled into the docs
+
+# basename(pwd()) if the code is run from inside the docs folder. If we don't
+# have that, we will assume that we are running from the project root. If this
+# code is run from anywhere but these two places, mkdocs will fail to find
+# availbale_diagnostics.md
+prefix = basename(pwd()) == "docs" ? "" : "docs/"
+
+out_path = "$(prefix)src/available_diagnostics.md"
+
+open(out_path, "w") do file
+
+    write(file, "# Available diagnostic variables\n\n")
+
+    write(file, "| Short name | Long name  | Units | Comments |\n")
+    write(file, "|---|---|---|---|\n")
+
+    for d in values(CA.ALL_DIAGNOSTICS)
+        write(file, "| `$(d.short_name)` ")
+        write(file, "| `$(d.long_name)` ")
+        write(file, "| $(d.units) ")
+        write(file, "| $(d.comments)|\n")
+    end
+end
+@info "Written $out_path"
diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 5929b2c1c64..31ed4771429 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -35,7 +35,8 @@ push!(diagnostics, get_daily_max("air_density", "air_temperature"))
 Now `diagnostics` will also contain the instructions to compute the daily
 maximum of `air_density` and `air_temperature`.
 
-**TODO: Add link to table with known diagnostics**
+The diagnostics that are built-in `ClimaAtmos` are collected in [Available
+diagnostic variables](@ref).
 
 If you are using `ClimaAtmos` with a script-based interface, you have access to
 the complete flexibility in your diagnostics. Read the section about the

From daca016051156d5ad6bfb3876954df9451920934 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 7 Sep 2023 13:44:11 -0700
Subject: [PATCH 16/73] Capture the 'average' reduction in naming

---
 src/diagnostics/Diagnostics.jl       | 12 ++++++++---
 src/diagnostics/Writers.jl           |  6 ++++--
 src/diagnostics/diagnostics_utils.jl | 30 +++++++++++++++++-----------
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index a033be0870b..af539236067 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -265,8 +265,12 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     )
 
         # We provide an inner constructor to enforce some constraints
-        descriptive_name =
-            get_descriptive_name(variable, output_every, reduction_time_func)
+        descriptive_name = get_descriptive_name(
+            variable,
+            output_every,
+            pre_output_hook!,
+            reduction_time_func,
+        )
 
         output_every % compute_every == 0 || error(
             "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
@@ -388,7 +392,8 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         descriptive_name = get_descriptive_name(
             variable,
             output_every,
-            reduction_time_func;
+            reduction_time_func,
+            pre_output_hook!;
             units_are_seconds = false,
         )
 
@@ -453,6 +458,7 @@ function ScheduledDiagnosticIterations(
         sd_time.variable,
         sd_time.output_every,
         sd_time.reduction_time_func,
+        sd_time.pre_output_hook!,
     )
 
     isinteger(output_every) || error(
diff --git a/src/diagnostics/Writers.jl b/src/diagnostics/Writers.jl
index d4b0fc4a231..3a7d0062636 100644
--- a/src/diagnostics/Writers.jl
+++ b/src/diagnostics/Writers.jl
@@ -45,14 +45,16 @@ get_descriptive_name(sd_i::ScheduledDiagnosticIterations, Δt::Nothing) =
     get_descriptive_name(
         sd_t.variable,
         sd_t.output_every,
-        sd_t.reduction_time_func;
+        sd_t.reduction_time_func,
+        sd_t.pre_output_hook!;
         units_are_seconds = false,
     )
 get_descriptive_name(sd_i::ScheduledDiagnosticIterations, Δt::T) where {T} =
     get_descriptive_name(
         sd_i.variable,
         sd_i.output_every * Δt,
-        sd_i.reduction_time_func;
+        sd_i.reduction_time_func,
+        sd_i.pre_output_hook!;
         units_are_seconds = true,
     )
 
diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl
index 9f37dedfb60..1863f802678 100644
--- a/src/diagnostics/diagnostics_utils.jl
+++ b/src/diagnostics/diagnostics_utils.jl
@@ -7,7 +7,8 @@
 """
     get_descriptive_name(variable::DiagnosticVariable,
                          output_every,
-                         reduction_time_func;
+                         reduction_time_func,
+                         pre_output_hook!;
                          units_are_seconds = true)
 
 
@@ -23,15 +24,25 @@ This function is useful for filenames and error messages.
 function get_descriptive_name(
     variable::DiagnosticVariable,
     output_every,
-    reduction_time_func;
+    reduction_time_func,
+    pre_output_hook!;
     units_are_seconds = true,
 )
     var = "$(variable.short_name)"
     isa_reduction = !isnothing(reduction_time_func)
 
-    if units_are_seconds
-        if isa_reduction
-            red = "$(reduction_time_func)"
+
+    if isa_reduction
+        red = "$(reduction_time_func)"
+
+        # Let's check if we are computing the average. Note that this might slip under the
+        # radar if the user passes their own pre_output_hook!.
+        if reduction_time_func == (+) &&
+           pre_output_hook! == average_pre_output_hook!
+            red = "average"
+        end
+
+        if units_are_seconds
 
             # Convert period from seconds to days, hours, minutes, seconds
             period = ""
@@ -47,15 +58,10 @@ function get_descriptive_name(
 
             suffix = period * red
         else
-            # Not a reduction
-            suffix = "inst"
-        end
-    else
-        if isa_reduction
             suffix = "$(output_every)it_(reduction_time_func)"
-        else
-            suffix = "inst"
         end
+    else
+        suffix = "inst"
     end
     return "$(var)_$(suffix)"
 end

From 10b1201a32f4cbd8d35fa07cc2f493228fe774f5 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 7 Sep 2023 13:44:36 -0700
Subject: [PATCH 17/73] Add helper function for diagnostics

---
 docs/src/diagnostics.md                    |  7 +-
 src/diagnostics/default_diagnostics.jl     | 93 +++++++++++++++++++---
 src/diagnostics/defaults/moisture_model.jl | 34 ++++----
 3 files changed, 104 insertions(+), 30 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 31ed4771429..392a168c072 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -79,11 +79,12 @@ More specifically, a `ScheduledDiagnostic` contains the following pieces of data
    last time it was output.
 
 To implement operations like the arithmetic average, the `reduction_time_func`
-has to be chosen as `sum`, and a `pre_output_hook!` that renormalize `acc` by
-the number of samples has to be provided. `pre_output_hook!` should mutate the
+has to be chosen as `+`, and a `pre_output_hook!` that renormalize `acc` by the
+number of samples has to be provided. `pre_output_hook!` should mutate the
 accumulator in place. The return value of `pre_output_hook!` is discarded. An
 example of `pre_output_hook!` to compute the arithmetic average is
-`pre_output_hook!(acc, N) = @. acc = acc / N`.
+`pre_output_hook!(acc, N) = @. acc = acc / N`. `ClimaAtmos` provides an alias to
+the function needed to compute averages `ClimaAtmos.average_pre_output_hook!`.
 
 For custom reductions, it is necessary to also specify the identity of operation
 by defining a new method to `identity_of_reduction`. `identity_of_reduction` is
diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index b3b534a57a6..59fedef52da 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -30,24 +30,93 @@ end
 # that have no given defaults.
 get_default_diagnostics(_) = []
 
+
+"""
+    produce_common_diagnostic_function(period, reduction)
+
+
+Helper function to define functions like `get_daily_max`.
+"""
+function produce_common_diagnostic_function(
+    period,
+    reduction;
+    pre_output_hook! = (accum, count) -> nothing,
+)
+    return (long_names...; output_writer = HDF5Writer()) -> begin
+        [
+            ScheduledDiagnosticTime(
+                variable = ALL_DIAGNOSTICS[long_name],
+                compute_every = :timestep,
+                output_every = period, # seconds
+                reduction_time_func = reduction,
+                output_writer = output_writer,
+                pre_output_hook! = pre_output_hook!,
+            ) for long_name in long_names
+        ]
+    end
+end
+
+function average_pre_output_hook!(accum, counter)
+    @. accum = accum / counter
+    nothing
+end
+
 """
     get_daily_max(long_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily max for the given variables.
 """
-function get_daily_max(long_names...; output_writer = HDF5Writer())
-    # TODO: Add mechanism to print out reasonable error on variables that are not in ALL_DIAGNOSTICS
-    return [
-        ScheduledDiagnosticTime(
-            variable = ALL_DIAGNOSTICS[long_name],
-            compute_every = :timestep,
-            output_every = 86400, # seconds
-            reduction_time_func = max,
-            output_writer = output_writer,
-        ) for long_name in long_names
-    ]
-end
+get_daily_max = produce_common_diagnostic_function(24 * 60 * 60, max)
+"""
+    get_daily_min(long_names...; output_writer = HDF5Writer())
+
+
+Return a list of `ScheduledDiagnostics` that compute the daily min for the given variables.
+"""
+get_daily_min = produce_common_diagnostic_function(24 * 60 * 60, min)
+"""
+    get_daily_average(long_names...; output_writer = HDF5Writer())
+
+
+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
+get_daily_average = produce_common_diagnostic_function(
+    24 * 60 * 60,
+    (+);
+    pre_output_hook! = average_pre_output_hook!,
+)
+
+"""
+    get_hourly_max(long_names...; output_writer = HDF5Writer())
+
+
+Return a list of `ScheduledDiagnostics` that compute the hourly max for the given variables.
+"""
+get_hourly_max = produce_common_diagnostic_function(60 * 60, max)
+
+"""
+    get_hourly_min(long_names...; output_writer = HDF5Writer())
+
+
+Return a list of `ScheduledDiagnostics` that compute the hourly min for the given variables.
+"""
+get_hourly_min = produce_common_diagnostic_function(60 * 60, min)
+
+"""
+    get_daily_average(long_names...; output_writer = HDF5Writer())
+
+
+Return a list of `ScheduledDiagnostics` that compute the hourly average for the given variables.
+"""
+
+# An average is just a sum with a normalization before output
+get_hourly_average = produce_common_diagnostic_function(
+    60 * 60,
+    (+);
+    pre_output_hook! = average_pre_output_hook!,
+)
 
 # Include all the subdefaults
 include("defaults/moisture_model.jl")
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
index 14e55e664c2..14b26d0eb75 100644
--- a/src/diagnostics/defaults/moisture_model.jl
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -1,18 +1,22 @@
 # FIXME: Gabriele added this as an example. Put something meaningful here!
 function get_default_diagnostics(::T) where {T <: DryModel}
-    return [
-        ScheduledDiagnosticTime(
-            variable = ALL_DIAGNOSTICS["air_density"],
-            compute_every = :timestep,
-            output_every = 86400, # seconds
-            reduction_time_func = max,
-            output_writer = HDF5Writer(),
-        ),
-        ScheduledDiagnosticIterations(
-            variable = ALL_DIAGNOSTICS["air_density"],
-            compute_every = 1,
-            output_every = 1, # iteration
-            output_writer = HDF5Writer(),
-        ),
-    ]
+    return vcat(
+        get_daily_average("air_density"),
+        get_hourly_max("air_density"),
+        [
+            ScheduledDiagnosticTime(
+                variable = ALL_DIAGNOSTICS["air_density"],
+                compute_every = :timestep,
+                output_every = 86400, # seconds
+                reduction_time_func = min,
+                output_writer = HDF5Writer(),
+            ),
+            ScheduledDiagnosticIterations(
+                variable = ALL_DIAGNOSTICS["air_density"],
+                compute_every = 1,
+                output_every = 1, # iteration
+                output_writer = HDF5Writer(),
+            ),
+        ],
+    )
 end

From 1fd4d427f0f50872e639084a045beda26c5a4f3f Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 09:16:55 -0700
Subject: [PATCH 18/73] Add example and docs on model-dependent diagnostic

---
 docs/src/diagnostics.md                 | 31 +++++++++++++++++++
 src/diagnostics/Diagnostics.jl          |  2 ++
 src/diagnostics/core_diagnostics.jl     | 39 +++++++++++++++++++++++
 src/diagnostics/turbconv_diagnostics.jl | 41 +++++++++++++++++++++++++
 4 files changed, 113 insertions(+)
 create mode 100644 src/diagnostics/turbconv_diagnostics.jl

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 392a168c072..d99d84306c0 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -189,6 +189,37 @@ avoid extra memory allocations (which hurt performance). If `out` is `nothing`,
 and new area of memory is allocated. If `out` is a `ClimaCore.Field`, the
 operation is done in-place without additional memory allocations.
 
+If your diagnostic depends on the details of the model, we recommend using
+additional functions so that the correct one can be found through dispatching.
+For instance, if you want to compute relative humidity, which does not make
+sense for dry simulations, you should define the functions
+
+```julia
+function compute_relative_humidity_from_integrator(
+    integrator,
+    out,
+    moisture_model::T,
+)
+    error("Cannot compute relative_humidity with moisture_model = $T")
+end
+
+function compute_relative_humidity_from_integrator(
+    integrator,
+    out,
+    moisture_model::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    thermo_params = CAP.thermodynamics_params(integrator.p.params)
+    return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
+end
+```
+
+This will return the correct relative humidity and throw informative errors when
+it cannot be computed. We could specialize
+`compute_relative_humidity_from_integrator` further if the relative humidity
+were computed differently for `EquilMoistModel` and `NonEquilMoistModel`.
+
 ### Adding to the `ALL_DIAGNOSTICS` dictionary
 
 `ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable` in the
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index af539236067..af0c9146641 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -41,6 +41,7 @@
 #
 # - This file also also include several other files, including (but not limited to):
 #   - core_diagnostics.jl
+#   - turbconv_diagnostics.jl
 #   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
 #   - reduction_identities.jl
 #   - diagnostic_utils.jl
@@ -155,6 +156,7 @@ end
 
 # Do you want to define more diagnostics? Add them here
 include("core_diagnostics.jl")
+include("turbconv_diagnostics.jl")
 
 # Default diagnostics and higher level interfaces
 include("default_diagnostics.jl")
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 2397a937f5c..1bffcf2b40c 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -1,5 +1,7 @@
 # This file is included in Diagnostics.jl
 
+# Rho
+
 # FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
 # physics. Please fix this!
 add_diagnostic_variable!(
@@ -13,3 +15,40 @@ add_diagnostic_variable!(
         return deepcopy(integrator.u.c.ρ)
     end,
 )
+
+# Relative humidity
+
+function compute_relative_humidity_from_integrator(
+    integrator,
+    out,
+    moisture_model::T,
+) where {T}
+    error("Cannot compute relative_humidity with moisture_model = $T")
+end
+
+function compute_relative_humidity_from_integrator(
+    integrator,
+    out,
+    moisture_model::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    thermo_params = CAP.thermodynamics_params(integrator.p.params)
+    return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
+end
+
+# FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
+# physics. Please fix this!
+add_diagnostic_variable!(
+    short_name = "relative_humidity",
+    long_name = "relative_humidity",
+    units = "",
+    comments = "Relative Humidity",
+    compute_from_integrator = (integrator, out) -> begin
+        return compute_relative_humidity_from_integrator(
+            integrator,
+            out,
+            integrator.p.atmos.moisture_model,
+        )
+    end,
+)
diff --git a/src/diagnostics/turbconv_diagnostics.jl b/src/diagnostics/turbconv_diagnostics.jl
new file mode 100644
index 00000000000..9b2c8c15449
--- /dev/null
+++ b/src/diagnostics/turbconv_diagnostics.jl
@@ -0,0 +1,41 @@
+# This file is included in Diagnostics.jl
+
+# TKE
+
+# This is an example of how to compute the same diagnostic variable differently depending on
+# the model. This is also exposed to the user, which could define their own
+# compute_tke_from_integrator.
+
+function compute_tke_from_integrator(integrator, out, ::EDMFX)
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    return deepcopy(integrator.p.ᶜspecific⁰.tke)
+end
+
+function compute_tke_from_integrator(integrator, out, ::DiagnosticEDMFX)
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    return deepcopy(integrator.p.tke⁰)
+end
+
+function compute_tke_from_integrator(
+    integrator,
+    out,
+    turbconv_model::T,
+) where {T}
+    error("Cannot compute tke with turbconv_model = $T")
+end
+
+# FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
+# physics. Please fix this!
+add_diagnostic_variable!(
+    short_name = "tke",
+    long_name = "turbolent_kinetic_energy",
+    units = "J",
+    comments = "Turbolent Kinetic Energy",
+    compute_from_integrator = (integrator, out) -> compute_tke_from_integrator(
+        integrator,
+        out,
+        integrator.p.atmos.turbconv_model,
+    ),
+)

From 0b85ac9a46fc6b96d017093ddcc95fb82a483cd7 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:12:22 -0700
Subject: [PATCH 19/73] Make Diagnostics a submodule

---
 docs/make_diagnostic_table.jl              |   2 +-
 src/ClimaAtmos.jl                          |   5 +-
 src/diagnostics/Diagnostics.jl             | 659 +--------------------
 src/diagnostics/diagnostic.jl              | 635 ++++++++++++++++++++
 src/diagnostics/{Writers.jl => writers.jl} |   0
 src/solver/type_getters.jl                 |   7 +-
 6 files changed, 670 insertions(+), 638 deletions(-)
 create mode 100644 src/diagnostics/diagnostic.jl
 rename src/diagnostics/{Writers.jl => writers.jl} (100%)

diff --git a/docs/make_diagnostic_table.jl b/docs/make_diagnostic_table.jl
index f3aedf2e994..fe27fbb75cf 100644
--- a/docs/make_diagnostic_table.jl
+++ b/docs/make_diagnostic_table.jl
@@ -18,7 +18,7 @@ open(out_path, "w") do file
     write(file, "| Short name | Long name  | Units | Comments |\n")
     write(file, "|---|---|---|---|\n")
 
-    for d in values(CA.ALL_DIAGNOSTICS)
+    for d in values(CA.Diagnostics.ALL_DIAGNOSTICS)
         write(file, "| `$(d.short_name)` ")
         write(file, "| `$(d.long_name)` ")
         write(file, "| $(d.units) ")
diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
index 705b19a7bf0..c822fe3b92c 100644
--- a/src/ClimaAtmos.jl
+++ b/src/ClimaAtmos.jl
@@ -110,6 +110,8 @@ include(joinpath("prognostic_equations", "limited_tendencies.jl"))
 
 include(joinpath("callbacks", "callbacks.jl"))
 
+include(joinpath("diagnostics", "Diagnostics.jl"))
+
 include(joinpath("solver", "model_getters.jl")) # high-level (using parsed_args) model getters
 include(joinpath("solver", "type_getters.jl"))
 include(joinpath("solver", "yaml_helper.jl"))
@@ -117,7 +119,4 @@ include(joinpath("solver", "solve.jl"))
 
 include(joinpath("parameters", "create_parameters.jl"))
 
-include(joinpath("diagnostics", "Diagnostics.jl"))
-include(joinpath("diagnostics", "Writers.jl"))
-
 end # module
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index af0c9146641..36069cea02e 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -1,634 +1,31 @@
-# Diagnostics.jl
-#
-# This file contains:
-#
-# - The definition of what a DiagnosticVariable is. Morally, a DiagnosticVariable is a
-#   variable we know how to compute from the state. We attach more information to it for
-#   documentation and to reference to it with its long name. DiagnosticVariables can exist
-#   irrespective of the existence of an actual simulation that is being run. ClimaAtmos
-#   comes with several diagnostics already defined (in the `ALL_DIAGNOSTICS` dictionary).
-#
-# - A dictionary `ALL_DIAGNOSTICS` with all the diagnostics we know how to compute, keyed
-#   over their long name. If you want to add more diagnostics, look at the included files.
-#   You can add your own file if you want to define several new diagnostics that are
-#   conceptually related.
-#
-# - The definition of what a ScheduledDiagnostics is. Morally, a ScheduledDiagnostics is a
-#   DiagnosticVariable we want to compute in a given simulation. For example, it could be
-#   the temperature averaged over a day. We can have multiple ScheduledDiagnostics for the
-#   same DiagnosticVariable (e.g., daily and monthly average temperatures).
-#
-#   We provide two types of ScheduledDiagnostics: ScheduledDiagnosticIterations and
-#   ScheduledDiagnosticTime, with the difference being only in what domain the recurrence
-#   time is defined (are we doing something at every N timesteps or every T seconds?). It is
-#   much cleaner and simpler to work with ScheduledDiagnosticIterations because iterations
-#   are well defined and consistent. On the other hand, working in the time domain requires
-#   dealing with what happens when the timestep is not lined up with the output period.
-#   Possible solutions to this problem include: uneven output, interpolation, or restricting
-#   the user from picking specific combinations of timestep/output period. In the current
-#   implementation, we choose the third option. So, ScheduledDiagnosticTime is provided
-#   because it is the physically interesting quantity. If we know what is the timestep, we
-#   can convert between the two and check if the diagnostics are well-posed in terms of the
-#   relationship between the periods and the timesteps. In some sense, you can think of
-#   ScheduledDiagnosticIterations as an internal representation and ScheduledDiagnosticTime
-#   as the external interface.
-#
-# - A function to convert a list of ScheduledDiagnosticIterations into a list of
-#   AtmosCallbacks. This function takes three arguments: the list of diagnostics and two
-#   dictionaries that map each scheduled diagnostic to an area of memory where to save the
-#   result and where to keep track of how many times the function was called (so that we
-#   can compute stuff like averages).
-#
-# - This file also also include several other files, including (but not limited to):
-#   - core_diagnostics.jl
-#   - turbconv_diagnostics.jl
-#   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
-#   - reduction_identities.jl
-#   - diagnostic_utils.jl
-
-"""
-    DiagnosticVariable
-
-
-A recipe to compute a diagnostic variable from the state, along with some useful metadata.
-
-The primary use for `DiagnosticVariable`s is to be embedded in a `ScheduledDiagnostic` to
-compute diagnostics while the simulation is running.
-
-The metadata is used exclusively by the `output_writer` in the `ScheduledDiagnostic`. It is
-responsibility of the `output_writer` to follow the conventions about the meaning of the
-metadata and their use.
-
-In `ClimaAtmos`, we roughly follow the naming conventions listed in this file:
-https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
-
-Keyword arguments
-=================
-
-- `short_name`: Name used to identify the variable in the output files and in the file
-                names. Short but descriptive. `ClimaAtmos` follows the CMIP conventions and
-                the diagnostics are identified by the short name.
-
-- `long_name`: Name used to identify the variable in the input files.
-
-- `units`: Physical units of the variable.
-
-- `comments`: More verbose explanation of what the variable is, or comments related to how
-              it is defined or computed.
-
-- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
-                             It has to take two arguments: the `integrator`, and a
-                             pre-allocated area of memory where to write the result of the
-                             computation. It the no pre-allocated area is available, a new
-                             one will be allocated. To avoid extra allocations, this
-                             function should perform the calculation in-place (i.e., using
-                             `.=`).
-"""
-Base.@kwdef struct DiagnosticVariable{T <: AbstractString, T2}
-    short_name::T
-    long_name::T
-    units::T
-    comments::T
-    compute_from_integrator::T2
-end
-
-# ClimaAtmos diagnostics
-
-const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
-
-"""
-
-    add_diagnostic_variable!(; short_name,
-                               long_name,
-                               units,
-                               description,
-                               compute_from_integrator)
-
-
-Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates the state of
-`ClimaAtmos.ALL_DIAGNOSTICS`).
-
-If possible, please follow the naming scheme outline in
-https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
-
-Keyword arguments
-=================
-
-
-- `short_name`: Name used to identify the variable in the output files and in the file
-                names. Short but descriptive. `ClimaAtmos` diagnostics are identified by the
-                short name. We follow the Coupled Model Intercomparison Project conventions.
-
-- `long_name`: Name used to identify the variable in the input files.
-
-- `units`: Physical units of the variable.
-
-- `comments`: More verbose explanation of what the variable is, or comments related to how
-              it is defined or computed.
-
-- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
-                             It has to take two arguments: the `integrator`, and a
-                             pre-allocated area of memory where to write the result of the
-                             computation. It the no pre-allocated area is available, a new
-                             one will be allocated. To avoid extra allocations, this
-                             function should perform the calculation in-place (i.e., using
-                             `.=`).
-
-"""
-function add_diagnostic_variable!(;
-    short_name,
-    long_name,
-    units,
-    comments,
-    compute_from_integrator,
+module Diagnostics
+
+import ClimaCore: InputOutput
+import ClimaCore: Fields
+
+import ..AtmosModel
+import ..call_every_n_steps
+
+# moisture_model
+import ..DryModel
+import ..EquilMoistModel
+import ..NonEquilMoistModel
+
+# turbconv_model
+import ..EDMFX
+import ..DiagnosticEDMFX
+
+# Abbreviations (following utils/abbreviations.jl)
+const curlₕ = Operators.Curl()
+const CT3 = Geometry.Contravariant3Vector
+const ᶜinterp = Operators.InterpolateF2C()
+# TODO: Implement proper extrapolation instead of simply reusing the first
+# interior value at the surface.
+const ᶠinterp = Operators.InterpolateC2F(
+    bottom = Operators.Extrapolate(),
+    top = Operators.Extrapolate(),
 )
-    haskey(ALL_DIAGNOSTICS, short_name) &&
-        error("diagnostic $short_name already defined")
-
-    ALL_DIAGNOSTICS[short_name] = DiagnosticVariable(;
-        short_name,
-        long_name,
-        units,
-        comments,
-        compute_from_integrator,
-    )
-end
-
-# Do you want to define more diagnostics? Add them here
-include("core_diagnostics.jl")
-include("turbconv_diagnostics.jl")
-
-# Default diagnostics and higher level interfaces
-include("default_diagnostics.jl")
-
-# ScheduledDiagnostics
-
-# NOTE: The definitions of ScheduledDiagnosticTime and ScheduledDiagnosticIterations are
-# nearly identical except for the fact that one is assumed to use units of seconds the other
-# units of integration steps. However, we allow for this little repetition of code to avoid
-# adding an extra layer of abstraction just to deal with these two objects (some people say
-# that "duplication is better than over-abstraction"). Most users will only work with
-# ScheduledDiagnosticTime. (It would be nice to have defaults fields in abstract types, as
-# proposed in 2013 in https://github.com/JuliaLang/julia/issues/4935) Having two distinct
-# types allow us to implement different checks and behaviors (e.g., we allow
-# ScheduledDiagnosticTime to have placeholders values for {compute, output}_every so that we
-# can plug the timestep in it).
-
-# Design note: pre_output_hook!
-#
-# One of our key requirements is to be able to compute arithmetic averages. Unfortunately,
-# computing an arithmetic average requires keeping track of how many elements we are summing
-# up. pre_output_hook! was introduced so that we can think of an average as a sum coupled
-# with division, and perform the division (by the number of elements) before output.
-# pre_output_hook! could be used for other operations, but we decided to keep it simple and
-# target directly the most important use case for us.
-#
-# This choice restricts what reductions can be performed. For example, it is not possible to
-# have a geometric average. If more complex reduction are needed, this mechanism has to be
-# changed.
-
-struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
-    variable::DiagnosticVariable
-    output_every::T1
-    output_writer::OW
-    reduction_time_func::F1
-    reduction_space_func::F2
-    compute_every::T2
-    pre_output_hook!::PO
-
-    """
-        ScheduledDiagnosticIterations(; variable::DiagnosticVariable,
-                                        output_every,
-                                        output_writer,
-                                        reduction_time_func = nothing,
-                                        reduction_space_func = nothing,
-                                        compute_every = isa_reduction ? 1 : output_every,
-                                        pre_output_hook! = (accum, count) -> nothing)
-
-
-    A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence
-    defined by the number of iterations, with an optional reduction applied to it (e.g., compute
-    the maximum temperature over the course of every 10 timesteps). This object is turned into
-    two callbacks (one for computing and the other for output) and executed by the integrator.
-
-    Keyword arguments
-    =================
-
-    - `variable`: The diagnostic variable that has to be computed and output.
-
-    - `output_every`: Save the results to disk every `output_every` iterations.
-
-    - `output_writer`: Function that controls out to save the computed diagnostic variable to
-                       disk. `output_writer` has to take three arguments: the value that has to
-                       be output, the `ScheduledDiagnostic`, and the integrator. Internally, the
-                       integrator contains extra information (such as the current timestep). It
-                       is responsibility of the `output_writer` to properly use the provided
-                       information for meaningful output.
-
-    - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of
-                             scratch space `acc` where to accumulate partial results. Then, at
-                             every `compute_every`, `reduction_time_func` is computed between
-                             the previously stored value in `acc` and the new value. This
-                             implements a running reduction. For example, if
-                             `reduction_time_func = max`, the space `acc` will hold the running
-                             maxima of the diagnostic. To implement operations like the
-                             arithmetic average, the `reduction_time_func` has to be chosen as
-                             `sum`, and a `pre_output_hook!` that renormalize `acc` by the
-                             number of samples has to be provided. For custom reductions, it is
-                             necessary to also specify the identity of operation by defining a
-                             new method to `identity_of_reduction`.
-
-    - `reduction_space_func`: NOT IMPLEMENTED YET
-
-    - `compute_every`: Run the computations every `compute_every` iterations. This is not
-                       particularly useful for point-wise diagnostics, where we enforce that
-                       `compute_every` = `output_every`. For time reductions, `compute_every` is
-                       set to 1 (compute at every timestep) by default. `compute_every` has to
-                       evenly divide `output_every`.
-
-    - `pre_output_hook!`: Function that has to be run before saving to disk for reductions
-                          (mostly used to implement averages). The function `pre_output_hook!`
-                          is called with two arguments: the value accumulated during the
-                          reduction, and the number of times the diagnostic was computed from
-                          the last time it was output. `pre_output_hook!` should mutate the
-                          accumulator in place. The return value of `pre_output_hook!` is
-                          discarded. An example of `pre_output_hook!` to compute the arithmetic
-                          average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
-
-    """
-    function ScheduledDiagnosticIterations(;
-        variable::DiagnosticVariable,
-        output_every,
-        output_writer,
-        reduction_time_func = nothing,
-        reduction_space_func = nothing,
-        compute_every = isnothing(reduction_time_func) ? output_every : 1,
-        pre_output_hook! = (accum, count) -> nothing,
-    )
-
-        # We provide an inner constructor to enforce some constraints
-        descriptive_name = get_descriptive_name(
-            variable,
-            output_every,
-            pre_output_hook!,
-            reduction_time_func,
-        )
-
-        output_every % compute_every == 0 || error(
-            "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
-        )
-
-        isa_reduction = !isnothing(reduction_time_func)
-
-        # If it is not a reduction, we compute only when we output
-        if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
-            compute_every = output_every
-        end
-
-        T1 = typeof(output_every)
-        T2 = typeof(compute_every)
-        OW = typeof(output_writer)
-        F1 = typeof(reduction_time_func)
-        F2 = typeof(reduction_space_func)
-        PO = typeof(pre_output_hook!)
-
-        new{T1, T2, OW, F1, F2, PO}(
-            variable,
-            output_every,
-            output_writer,
-            reduction_time_func,
-            reduction_space_func,
-            compute_every,
-            pre_output_hook!,
-        )
-    end
-end
-
-
-struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
-    variable::DiagnosticVariable
-    output_every::T1
-    output_writer::OW
-    reduction_time_func::F1
-    reduction_space_func::F2
-    compute_every::T2
-    pre_output_hook!::PO
-
-    """
-        ScheduledDiagnosticTime(; variable::DiagnosticVariable,
-                                            output_every,
-                                            output_writer,
-                                            reduction_time_func = nothing,
-                                            reduction_space_func = nothing,
-                                            compute_every = isa_reduction ? :timestep : output_every,
-                                            pre_output_hook! = (accum, count) -> nothing)
-
-
-    A `DiagnosticVariable` that has to be computed and output during a simulation with a
-    cadence defined by how many seconds in simulation time, with an optional reduction
-    applied to it (e.g., compute the maximum temperature over the course of every day). This
-    object is turned into a `ScheduledDiagnosticIterations`, which is turned into two
-    callbacks (one for computing and the other for output) and executed by the integrator.
-
-    Keyword arguments
-    =================
-
-    - `variable`: The diagnostic variable that has to be computed and output.
-
-    - `output_every`: Save the results to disk every `output_every` seconds.
-
-    - `output_writer`: Function that controls out to save the computed diagnostic variable to
-                       disk. `output_writer` has to take three arguments: the value that has to
-                       be output, the `ScheduledDiagnostic`, and the integrator. Internally, the
-                       integrator contains extra information (such as the current timestep). It
-                       is responsibility of the `output_writer` to properly use the provided
-                       information for meaningful output.
-
-    - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of
-                             scratch space `acc` where to accumulate partial results. Then, at
-                             every `compute_every`, `reduction_time_func` is computed between
-                             the previously stored value in `acc` and the new value. This
-                             implements a running reduction. For example, if
-                             `reduction_time_func = max`, the space `acc` will hold the running
-                             maxima of the diagnostic. To implement operations like the
-                             arithmetic average, the `reduction_time_func` has to be chosen as
-                             `sum`, and a `pre_output_hook!` that renormalize `acc` by the
-                             number of samples has to be provided. For custom reductions, it is
-                             necessary to also specify the identity of operation by defining a
-                             new method to `identity_of_reduction`.
-
-    - `reduction_space_func`: NOT IMPLEMENTED YET
-
-    - `compute_every`: Run the computations every `compute_every` seconds. This is not
-                       particularly useful for point-wise diagnostics, where we enforce that
-                       `compute_every` = `output_every`. For time reductions,
-                       `compute_every` is set to `:timestep` (compute at every timestep) by
-                       default. `compute_every` has to evenly divide `output_every`.
-                       `compute_every` can take the special symbol `:timestep` which is a
-                       placeholder for the timestep of the simulation to which this
-                       `ScheduledDiagnostic` is attached.
-
-    - `pre_output_hook!`: Function that has to be run before saving to disk for reductions
-                          (mostly used to implement averages). The function `pre_output_hook!`
-                          is called with two arguments: the value accumulated during the
-                          reduction, and the number of times the diagnostic was computed from
-                          the last time it was output. `pre_output_hook!` should mutate the
-                          accumulator in place. The return value of `pre_output_hook!` is
-                          discarded. An example of `pre_output_hook!` to compute the arithmetic
-                          average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
-
-    """
-    function ScheduledDiagnosticTime(;
-        variable::DiagnosticVariable,
-        output_every,
-        output_writer,
-        reduction_time_func = nothing,
-        reduction_space_func = nothing,
-        compute_every = isnothing(reduction_time_func) ? output_every :
-                        :timestep,
-        pre_output_hook! = (accum, count) -> nothing,
-    )
-
-        # We provide an inner constructor to enforce some constraints
-        descriptive_name = get_descriptive_name(
-            variable,
-            output_every,
-            reduction_time_func,
-            pre_output_hook!;
-            units_are_seconds = false,
-        )
-
-        # compute_every could be a Symbol (:timestep). We process this that when we process
-        # the list of diagnostics
-        if !isa(compute_every, Symbol)
-            output_every % compute_every == 0 || error(
-                "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
-            )
-        end
-
-        isa_reduction = !isnothing(reduction_time_func)
-
-        # If it is not a reduction, we compute only when we output
-        if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
-            compute_every = output_every
-        end
-
-        T1 = typeof(output_every)
-        T2 = typeof(compute_every)
-        OW = typeof(output_writer)
-        F1 = typeof(reduction_time_func)
-        F2 = typeof(reduction_space_func)
-        PO = typeof(pre_output_hook!)
-
-        new{T1, T2, OW, F1, F2, PO}(
-            variable,
-            output_every,
-            output_writer,
-            reduction_time_func,
-            reduction_space_func,
-            compute_every,
-            pre_output_hook!,
-        )
-    end
-end
-
-"""
-    ScheduledDiagnosticIterations(sd_time::ScheduledDiagnosticTime, Δt)
-
-
-Create a `ScheduledDiagnosticIterations` given a `ScheduledDiagnosticTime` and a timestep
-`Δt`. In this, ensure that `compute_every` and `output_every` are meaningful for the given
-timestep.
-
-"""
-
-function ScheduledDiagnosticIterations(
-    sd_time::ScheduledDiagnosticTime,
-    Δt::T,
-) where {T}
-
-    # If we have the timestep, we can convert time in seconds into iterations
-
-    # if compute_every is :timestep, then we want to compute after every iterations
-    compute_every =
-        sd_time.compute_every == :timestep ? 1 : sd_time.compute_every / Δt
-    output_every = sd_time.output_every / Δt
-
-    descriptive_name = get_descriptive_name(
-        sd_time.variable,
-        sd_time.output_every,
-        sd_time.reduction_time_func,
-        sd_time.pre_output_hook!,
-    )
-
-    isinteger(output_every) || error(
-        "output_every should be multiple of the timestep for diagnostic $(descriptive_name)",
-    )
-    isinteger(compute_every) || error(
-        "compute_every should be multiple of the timestep for diagnostic $(descriptive_name)",
-    )
-
-    ScheduledDiagnosticIterations(;
-        sd_time.variable,
-        output_every = convert(Int, output_every),
-        sd_time.output_writer,
-        sd_time.reduction_time_func,
-        sd_time.reduction_space_func,
-        compute_every = convert(Int, compute_every),
-        sd_time.pre_output_hook!,
-    )
-end
-
-"""
-    ScheduledDiagnosticTime(sd_time::ScheduledDiagnosticIterations, Δt)
-
-
-Create a `ScheduledDiagnosticTime` given a `ScheduledDiagnosticIterations` and a timestep
-`Δt`.
-
-"""
-function ScheduledDiagnosticTime(
-    sd_time::ScheduledDiagnosticIterations,
-    Δt::T,
-) where {T}
-
-    # If we have the timestep, we can convert time in iterations to seconds
-
-    # if compute_every is :timestep, then we want to compute after every iterations
-    compute_every =
-        sd_time.compute_every == 1 ? :timestep : sd_time.compute_every * Δt
-    output_every = sd_time.output_every * Δt
-
-    ScheduledDiagnosticTime(;
-        sd_time.variable,
-        output_every,
-        sd_time.output_writer,
-        sd_time.reduction_time_func,
-        sd_time.reduction_space_func,
-        compute_every,
-        sd_time.pre_output_hook!,
-    )
-end
-
-# We provide also a companion constructor for ScheduledDiagnosticIterations which returns
-# itself (without copy) when called with a timestep.
-#
-# This is so that we can assume that
-# ScheduledDiagnosticIterations(ScheduledDiagnostic{Time, Iterations}, Δt)
-# always returns a valid ScheduledDiagnosticIterations
-ScheduledDiagnosticIterations(
-    sd::ScheduledDiagnosticIterations,
-    _Δt::T,
-) where {T} = sd
-
-
-# We define all the known identities in reduction_identities.jl
-include("reduction_identities.jl")
-
-# Helper functions
-include("diagnostics_utils.jl")
-
-"""
-    get_callbacks_from_diagnostics(diagnostics, storage, counters)
-
-
-Translate a list of diagnostics into a list of callbacks.
-
-Positional arguments
-=====================
-
-- `diagnostics`: List of `ScheduledDiagnosticIterations` that have to be converted to
-                 callbacks. We want to have `ScheduledDiagnosticIterations` here so that we
-                 can define callbacks that occur at the end of every N integration steps.
-
-- `storage`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially
-             pre-allocated area of memory where to accumulate/save results.
-
-- `counters`: Dictionary that maps a given `ScheduledDiagnosticIterations` to the counter
-              that tracks how many times the given diagnostics was computed from the last
-              time it was output to disk.
-
-"""
-function get_callbacks_from_diagnostics(diagnostics, storage, counters)
-    # We have two types of callbacks: to compute and accumulate diagnostics, and to dump
-    # them to disk. Note that our callbacks do not contain any branching
-
-    # storage is used to pre-allocate memory and to accumulate partial results for those
-    # diagnostics that perform reductions.
-
-    callbacks = Any[]
-
-    for diag in diagnostics
-        variable = diag.variable
-        isa_reduction = !isnothing(diag.reduction_time_func)
-
-        # reduction is used below. If we are not given a reduction_time_func, we just want
-        # to move the computed quantity to its storage (so, we return the second argument,
-        # which that will be the newly computed one). If we have a reduction, we apply it
-        # point-wise
-        reduction = isa_reduction ? diag.reduction_time_func : (_, y) -> y
-
-        # If we have a reduction, we have to reset the accumulator to its neutral state. (If
-        # we don't have a reduction, we don't have to do anything)
-        #
-        # ClimaAtmos defines methods for identity_of_reduction for standard
-        # reduction_time_func in reduction_identities.jl
-        reset_accumulator! =
-            isa_reduction ?
-            () -> begin
-                # identity_of_reduction works by dispatching over Val{operation}
-                identity =
-                    identity_of_reduction(Val(diag.reduction_time_func))
-                # We also need to make sure that we are consistent with the types
-                float_type = eltype(storage[diag])
-                identity_ft = convert(float_type, identity)
-                storage[diag] .= identity_ft
-            end : () -> nothing
-
-        compute_callback =
-            integrator -> begin
-                # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
-                value = variable.compute_from_integrator(integrator, nothing)
-                storage[diag] .= reduction.(storage[diag], value)
-                counters[diag] += 1
-                return nothing
-            end
-
-        output_callback =
-            integrator -> begin
-                # Any operations we have to perform before writing to output?
-                # Here is where we would divide by N to obtain an arithmetic average
-                diag.pre_output_hook!(storage[diag], counters[diag])
-
-                # Write to disk
-                diag.output_writer(storage[diag], diag, integrator)
-
-                reset_accumulator!()
-                counters[diag] = 0
-                return nothing
-            end
-
-        # Here we have skip_first = true. This is important because we are going to manually
-        # call all the callbacks so that we can verify that they are meaningful for the
-        # model under consideration (and they don't have bugs).
-        append!(
-            callbacks,
-            [
-                call_every_n_steps(
-                    compute_callback,
-                    diag.compute_every,
-                    skip_first = true,
-                ),
-                call_every_n_steps(
-                    output_callback,
-                    diag.output_every,
-                    skip_first = true,
-                ),
-            ],
-        )
-    end
 
-    return callbacks
+include("diagnostic.jl")
+include("writers.jl")
 end
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
new file mode 100644
index 00000000000..4e7c0c22e08
--- /dev/null
+++ b/src/diagnostics/diagnostic.jl
@@ -0,0 +1,635 @@
+# diagnostic.jl
+#
+# This file contains:
+#
+# - The definition of what a DiagnosticVariable is. Morally, a DiagnosticVariable is a
+#   variable we know how to compute from the state. We attach more information to it for
+#   documentation and to reference to it with its short name. DiagnosticVariables can exist
+#   irrespective of the existence of an actual simulation that is being run. ClimaAtmos
+#   comes with several diagnostics already defined (in the `ALL_DIAGNOSTICS` dictionary).
+#
+# - A dictionary `ALL_DIAGNOSTICS` with all the diagnostics we know how to compute, keyed
+#   over their short name. If you want to add more diagnostics, look at the included files.
+#   You can add your own file if you want to define several new diagnostics that are
+#   conceptually related.
+#
+# - The definition of what a ScheduledDiagnostics is. Morally, a ScheduledDiagnostics is a
+#   DiagnosticVariable we want to compute in a given simulation. For example, it could be
+#   the temperature averaged over a day. We can have multiple ScheduledDiagnostics for the
+#   same DiagnosticVariable (e.g., daily and monthly average temperatures).
+#
+#   We provide two types of ScheduledDiagnostics: ScheduledDiagnosticIterations and
+#   ScheduledDiagnosticTime, with the difference being only in what domain the recurrence
+#   time is defined (are we doing something at every N timesteps or every T seconds?). It is
+#   much cleaner and simpler to work with ScheduledDiagnosticIterations because iterations
+#   are well defined and consistent. On the other hand, working in the time domain requires
+#   dealing with what happens when the timestep is not lined up with the output period.
+#   Possible solutions to this problem include: uneven output, interpolation, or restricting
+#   the user from picking specific combinations of timestep/output period. In the current
+#   implementation, we choose the third option. So, ScheduledDiagnosticTime is provided
+#   because it is the physically interesting quantity. If we know what is the timestep, we
+#   can convert between the two and check if the diagnostics are well-posed in terms of the
+#   relationship between the periods and the timesteps. In some sense, you can think of
+#   ScheduledDiagnosticIterations as an internal representation and ScheduledDiagnosticTime
+#   as the external interface.
+#
+# - A function to convert a list of ScheduledDiagnosticIterations into a list of
+#   AtmosCallbacks. This function takes three arguments: the list of diagnostics and two
+#   dictionaries that map each scheduled diagnostic to an area of memory where to save the
+#   result and where to keep track of how many times the function was called (so that we
+#   can compute stuff like averages).
+#
+# - This file also also include several other files, including (but not limited to):
+#   - core_diagnostics.jl
+#   - turbconv_diagnostics.jl
+#   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
+#   - reduction_identities.jl
+#   - diagnostic_utils.jl
+
+
+"""
+    DiagnosticVariable
+
+
+A recipe to compute a diagnostic variable from the state, along with some useful metadata.
+
+The primary use for `DiagnosticVariable`s is to be embedded in a `ScheduledDiagnostic` to
+compute diagnostics while the simulation is running.
+
+The metadata is used exclusively by the `output_writer` in the `ScheduledDiagnostic`. It is
+responsibility of the `output_writer` to follow the conventions about the meaning of the
+metadata and their use.
+
+In `ClimaAtmos`, we roughly follow the naming conventions listed in this file:
+https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
+
+Keyword arguments
+=================
+
+- `short_name`: Name used to identify the variable in the output files and in the file
+                names. Short but descriptive. `ClimaAtmos` follows the CMIP conventions and
+                the diagnostics are identified by the short name.
+
+- `long_name`: Name used to identify the variable in the output files.
+
+- `units`: Physical units of the variable.
+
+- `comments`: More verbose explanation of what the variable is, or comments related to how
+              it is defined or computed.
+
+- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
+                             It has to take two arguments: the `integrator`, and a
+                             pre-allocated area of memory where to write the result of the
+                             computation. It the no pre-allocated area is available, a new
+                             one will be allocated. To avoid extra allocations, this
+                             function should perform the calculation in-place (i.e., using
+                             `.=`).
+"""
+Base.@kwdef struct DiagnosticVariable{T <: AbstractString, T2}
+    short_name::T
+    long_name::T
+    units::T
+    comments::T
+    compute_from_integrator::T2
+end
+
+# ClimaAtmos diagnostics
+
+const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
+
+"""
+
+    add_diagnostic_variable!(; short_name,
+                               long_name,
+                               units,
+                               description,
+                               compute_from_integrator)
+
+
+Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates the state of
+`ClimaAtmos.ALL_DIAGNOSTICS`).
+
+If possible, please follow the naming scheme outline in
+https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
+
+Keyword arguments
+=================
+
+
+- `short_name`: Name used to identify the variable in the output files and in the file
+                names. Short but descriptive. `ClimaAtmos` diagnostics are identified by the
+                short name. We follow the Coupled Model Intercomparison Project conventions.
+
+- `long_name`: Name used to identify the variable in the output files.
+
+- `units`: Physical units of the variable.
+
+- `comments`: More verbose explanation of what the variable is, or comments related to how
+              it is defined or computed.
+
+- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
+                             It has to take two arguments: the `integrator`, and a
+                             pre-allocated area of memory where to write the result of the
+                             computation. It the no pre-allocated area is available, a new
+                             one will be allocated. To avoid extra allocations, this
+                             function should perform the calculation in-place (i.e., using
+                             `.=`).
+
+"""
+function add_diagnostic_variable!(;
+    short_name,
+    long_name,
+    units,
+    comments,
+    compute_from_integrator,
+)
+    haskey(ALL_DIAGNOSTICS, short_name) &&
+        error("diagnostic $short_name already defined")
+
+    ALL_DIAGNOSTICS[short_name] = DiagnosticVariable(;
+        short_name,
+        long_name,
+        units,
+        comments,
+        compute_from_integrator,
+    )
+end
+
+# Do you want to define more diagnostics? Add them here
+include("core_diagnostics.jl")
+include("turbconv_diagnostics.jl")
+
+# Default diagnostics and higher level interfaces
+include("default_diagnostics.jl")
+
+# ScheduledDiagnostics
+
+# NOTE: The definitions of ScheduledDiagnosticTime and ScheduledDiagnosticIterations are
+# nearly identical except for the fact that one is assumed to use units of seconds the other
+# units of integration steps. However, we allow for this little repetition of code to avoid
+# adding an extra layer of abstraction just to deal with these two objects (some people say
+# that "duplication is better than over-abstraction"). Most users will only work with
+# ScheduledDiagnosticTime. (It would be nice to have defaults fields in abstract types, as
+# proposed in 2013 in https://github.com/JuliaLang/julia/issues/4935) Having two distinct
+# types allow us to implement different checks and behaviors (e.g., we allow
+# ScheduledDiagnosticTime to have placeholders values for {compute, output}_every so that we
+# can plug the timestep in it).
+
+# Design note: pre_output_hook!
+#
+# One of our key requirements is to be able to compute arithmetic averages. Unfortunately,
+# computing an arithmetic average requires keeping track of how many elements we are summing
+# up. pre_output_hook! was introduced so that we can think of an average as a sum coupled
+# with division, and perform the division (by the number of elements) before output.
+# pre_output_hook! could be used for other operations, but we decided to keep it simple and
+# target directly the most important use case for us.
+#
+# This choice restricts what reductions can be performed. For example, it is not possible to
+# have a geometric average. If more complex reduction are needed, this mechanism has to be
+# changed.
+
+struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
+    variable::DiagnosticVariable
+    output_every::T1
+    output_writer::OW
+    reduction_time_func::F1
+    reduction_space_func::F2
+    compute_every::T2
+    pre_output_hook!::PO
+
+    """
+        ScheduledDiagnosticIterations(; variable::DiagnosticVariable,
+                                        output_every,
+                                        output_writer,
+                                        reduction_time_func = nothing,
+                                        reduction_space_func = nothing,
+                                        compute_every = isa_reduction ? 1 : output_every,
+                                        pre_output_hook! = (accum, count) -> nothing)
+
+
+    A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence
+    defined by the number of iterations, with an optional reduction applied to it (e.g., compute
+    the maximum temperature over the course of every 10 timesteps). This object is turned into
+    two callbacks (one for computing and the other for output) and executed by the integrator.
+
+    Keyword arguments
+    =================
+
+    - `variable`: The diagnostic variable that has to be computed and output.
+
+    - `output_every`: Save the results to disk every `output_every` iterations.
+
+    - `output_writer`: Function that controls out to save the computed diagnostic variable to
+                       disk. `output_writer` has to take three arguments: the value that has to
+                       be output, the `ScheduledDiagnostic`, and the integrator. Internally, the
+                       integrator contains extra information (such as the current timestep). It
+                       is responsibility of the `output_writer` to properly use the provided
+                       information for meaningful output.
+
+    - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of
+                             scratch space `acc` where to accumulate partial results. Then, at
+                             every `compute_every`, `reduction_time_func` is computed between
+                             the previously stored value in `acc` and the new value. This
+                             implements a running reduction. For example, if
+                             `reduction_time_func = max`, the space `acc` will hold the running
+                             maxima of the diagnostic. To implement operations like the
+                             arithmetic average, the `reduction_time_func` has to be chosen as
+                             `sum`, and a `pre_output_hook!` that renormalize `acc` by the
+                             number of samples has to be provided. For custom reductions, it is
+                             necessary to also specify the identity of operation by defining a
+                             new method to `identity_of_reduction`.
+
+    - `reduction_space_func`: NOT IMPLEMENTED YET
+
+    - `compute_every`: Run the computations every `compute_every` iterations. This is not
+                       particularly useful for point-wise diagnostics, where we enforce that
+                       `compute_every` = `output_every`. For time reductions, `compute_every` is
+                       set to 1 (compute at every timestep) by default. `compute_every` has to
+                       evenly divide `output_every`.
+
+    - `pre_output_hook!`: Function that has to be run before saving to disk for reductions
+                          (mostly used to implement averages). The function `pre_output_hook!`
+                          is called with two arguments: the value accumulated during the
+                          reduction, and the number of times the diagnostic was computed from
+                          the last time it was output. `pre_output_hook!` should mutate the
+                          accumulator in place. The return value of `pre_output_hook!` is
+                          discarded. An example of `pre_output_hook!` to compute the arithmetic
+                          average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
+
+    """
+    function ScheduledDiagnosticIterations(;
+        variable::DiagnosticVariable,
+        output_every,
+        output_writer,
+        reduction_time_func = nothing,
+        reduction_space_func = nothing,
+        compute_every = isnothing(reduction_time_func) ? output_every : 1,
+        pre_output_hook! = (accum, count) -> nothing,
+    )
+
+        # We provide an inner constructor to enforce some constraints
+        descriptive_name = get_descriptive_name(
+            variable,
+            output_every,
+            pre_output_hook!,
+            reduction_time_func,
+        )
+
+        output_every % compute_every == 0 || error(
+            "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
+        )
+
+        isa_reduction = !isnothing(reduction_time_func)
+
+        # If it is not a reduction, we compute only when we output
+        if !isa_reduction && compute_every != output_every
+            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
+            compute_every = output_every
+        end
+
+        T1 = typeof(output_every)
+        T2 = typeof(compute_every)
+        OW = typeof(output_writer)
+        F1 = typeof(reduction_time_func)
+        F2 = typeof(reduction_space_func)
+        PO = typeof(pre_output_hook!)
+
+        new{T1, T2, OW, F1, F2, PO}(
+            variable,
+            output_every,
+            output_writer,
+            reduction_time_func,
+            reduction_space_func,
+            compute_every,
+            pre_output_hook!,
+        )
+    end
+end
+
+
+struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
+    variable::DiagnosticVariable
+    output_every::T1
+    output_writer::OW
+    reduction_time_func::F1
+    reduction_space_func::F2
+    compute_every::T2
+    pre_output_hook!::PO
+
+    """
+        ScheduledDiagnosticTime(; variable::DiagnosticVariable,
+                                            output_every,
+                                            output_writer,
+                                            reduction_time_func = nothing,
+                                            reduction_space_func = nothing,
+                                            compute_every = isa_reduction ? :timestep : output_every,
+                                            pre_output_hook! = (accum, count) -> nothing)
+
+
+    A `DiagnosticVariable` that has to be computed and output during a simulation with a
+    cadence defined by how many seconds in simulation time, with an optional reduction
+    applied to it (e.g., compute the maximum temperature over the course of every day). This
+    object is turned into a `ScheduledDiagnosticIterations`, which is turned into two
+    callbacks (one for computing and the other for output) and executed by the integrator.
+
+    Keyword arguments
+    =================
+
+    - `variable`: The diagnostic variable that has to be computed and output.
+
+    - `output_every`: Save the results to disk every `output_every` seconds.
+
+    - `output_writer`: Function that controls out to save the computed diagnostic variable to
+                       disk. `output_writer` has to take three arguments: the value that has to
+                       be output, the `ScheduledDiagnostic`, and the integrator. Internally, the
+                       integrator contains extra information (such as the current timestep). It
+                       is responsibility of the `output_writer` to properly use the provided
+                       information for meaningful output.
+
+    - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of
+                             scratch space `acc` where to accumulate partial results. Then, at
+                             every `compute_every`, `reduction_time_func` is computed between
+                             the previously stored value in `acc` and the new value. This
+                             implements a running reduction. For example, if
+                             `reduction_time_func = max`, the space `acc` will hold the running
+                             maxima of the diagnostic. To implement operations like the
+                             arithmetic average, the `reduction_time_func` has to be chosen as
+                             `sum`, and a `pre_output_hook!` that renormalize `acc` by the
+                             number of samples has to be provided. For custom reductions, it is
+                             necessary to also specify the identity of operation by defining a
+                             new method to `identity_of_reduction`.
+
+    - `reduction_space_func`: NOT IMPLEMENTED YET
+
+    - `compute_every`: Run the computations every `compute_every` seconds. This is not
+                       particularly useful for point-wise diagnostics, where we enforce that
+                       `compute_every` = `output_every`. For time reductions,
+                       `compute_every` is set to `:timestep` (compute at every timestep) by
+                       default. `compute_every` has to evenly divide `output_every`.
+                       `compute_every` can take the special symbol `:timestep` which is a
+                       placeholder for the timestep of the simulation to which this
+                       `ScheduledDiagnostic` is attached.
+
+    - `pre_output_hook!`: Function that has to be run before saving to disk for reductions
+                          (mostly used to implement averages). The function `pre_output_hook!`
+                          is called with two arguments: the value accumulated during the
+                          reduction, and the number of times the diagnostic was computed from
+                          the last time it was output. `pre_output_hook!` should mutate the
+                          accumulator in place. The return value of `pre_output_hook!` is
+                          discarded. An example of `pre_output_hook!` to compute the arithmetic
+                          average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
+
+    """
+    function ScheduledDiagnosticTime(;
+        variable::DiagnosticVariable,
+        output_every,
+        output_writer,
+        reduction_time_func = nothing,
+        reduction_space_func = nothing,
+        compute_every = isnothing(reduction_time_func) ? output_every :
+                        :timestep,
+        pre_output_hook! = (accum, count) -> nothing,
+    )
+
+        # We provide an inner constructor to enforce some constraints
+        descriptive_name = get_descriptive_name(
+            variable,
+            output_every,
+            reduction_time_func,
+            pre_output_hook!;
+            units_are_seconds = false,
+        )
+
+        # compute_every could be a Symbol (:timestep). We process this that when we process
+        # the list of diagnostics
+        if !isa(compute_every, Symbol)
+            output_every % compute_every == 0 || error(
+                "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
+            )
+        end
+
+        isa_reduction = !isnothing(reduction_time_func)
+
+        # If it is not a reduction, we compute only when we output
+        if !isa_reduction && compute_every != output_every
+            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
+            compute_every = output_every
+        end
+
+        T1 = typeof(output_every)
+        T2 = typeof(compute_every)
+        OW = typeof(output_writer)
+        F1 = typeof(reduction_time_func)
+        F2 = typeof(reduction_space_func)
+        PO = typeof(pre_output_hook!)
+
+        new{T1, T2, OW, F1, F2, PO}(
+            variable,
+            output_every,
+            output_writer,
+            reduction_time_func,
+            reduction_space_func,
+            compute_every,
+            pre_output_hook!,
+        )
+    end
+end
+
+"""
+    ScheduledDiagnosticIterations(sd_time::ScheduledDiagnosticTime, Δt)
+
+
+Create a `ScheduledDiagnosticIterations` given a `ScheduledDiagnosticTime` and a timestep
+`Δt`. In this, ensure that `compute_every` and `output_every` are meaningful for the given
+timestep.
+
+"""
+
+function ScheduledDiagnosticIterations(
+    sd_time::ScheduledDiagnosticTime,
+    Δt::T,
+) where {T}
+
+    # If we have the timestep, we can convert time in seconds into iterations
+
+    # if compute_every is :timestep, then we want to compute after every iterations
+    compute_every =
+        sd_time.compute_every == :timestep ? 1 : sd_time.compute_every / Δt
+    output_every = sd_time.output_every / Δt
+
+    descriptive_name = get_descriptive_name(
+        sd_time.variable,
+        sd_time.output_every,
+        sd_time.reduction_time_func,
+        sd_time.pre_output_hook!,
+    )
+
+    isinteger(output_every) || error(
+        "output_every should be multiple of the timestep for diagnostic $(descriptive_name)",
+    )
+    isinteger(compute_every) || error(
+        "compute_every should be multiple of the timestep for diagnostic $(descriptive_name)",
+    )
+
+    ScheduledDiagnosticIterations(;
+        sd_time.variable,
+        output_every = convert(Int, output_every),
+        sd_time.output_writer,
+        sd_time.reduction_time_func,
+        sd_time.reduction_space_func,
+        compute_every = convert(Int, compute_every),
+        sd_time.pre_output_hook!,
+    )
+end
+
+"""
+    ScheduledDiagnosticTime(sd_time::ScheduledDiagnosticIterations, Δt)
+
+
+Create a `ScheduledDiagnosticTime` given a `ScheduledDiagnosticIterations` and a timestep
+`Δt`.
+
+"""
+function ScheduledDiagnosticTime(
+    sd_time::ScheduledDiagnosticIterations,
+    Δt::T,
+) where {T}
+
+    # If we have the timestep, we can convert time in iterations to seconds
+
+    # if compute_every is :timestep, then we want to compute after every iterations
+    compute_every =
+        sd_time.compute_every == 1 ? :timestep : sd_time.compute_every * Δt
+    output_every = sd_time.output_every * Δt
+
+    ScheduledDiagnosticTime(;
+        sd_time.variable,
+        output_every,
+        sd_time.output_writer,
+        sd_time.reduction_time_func,
+        sd_time.reduction_space_func,
+        compute_every,
+        sd_time.pre_output_hook!,
+    )
+end
+
+# We provide also a companion constructor for ScheduledDiagnosticIterations which returns
+# itself (without copy) when called with a timestep.
+#
+# This is so that we can assume that
+# ScheduledDiagnosticIterations(ScheduledDiagnostic{Time, Iterations}, Δt)
+# always returns a valid ScheduledDiagnosticIterations
+ScheduledDiagnosticIterations(
+    sd::ScheduledDiagnosticIterations,
+    _Δt::T,
+) where {T} = sd
+
+
+# We define all the known identities in reduction_identities.jl
+include("reduction_identities.jl")
+
+# Helper functions
+include("diagnostics_utils.jl")
+
+"""
+    get_callbacks_from_diagnostics(diagnostics, storage, counters)
+
+
+Translate a list of diagnostics into a list of callbacks.
+
+Positional arguments
+=====================
+
+- `diagnostics`: List of `ScheduledDiagnosticIterations` that have to be converted to
+                 callbacks. We want to have `ScheduledDiagnosticIterations` here so that we
+                 can define callbacks that occur at the end of every N integration steps.
+
+- `storage`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially
+             pre-allocated area of memory where to accumulate/save results.
+
+- `counters`: Dictionary that maps a given `ScheduledDiagnosticIterations` to the counter
+              that tracks how many times the given diagnostics was computed from the last
+              time it was output to disk.
+
+"""
+function get_callbacks_from_diagnostics(diagnostics, storage, counters)
+    # We have two types of callbacks: to compute and accumulate diagnostics, and to dump
+    # them to disk. Note that our callbacks do not contain any branching
+
+    # storage is used to pre-allocate memory and to accumulate partial results for those
+    # diagnostics that perform reductions.
+
+    callbacks = Any[]
+
+    for diag in diagnostics
+        variable = diag.variable
+        isa_reduction = !isnothing(diag.reduction_time_func)
+
+        # reduction is used below. If we are not given a reduction_time_func, we just want
+        # to move the computed quantity to its storage (so, we return the second argument,
+        # which that will be the newly computed one). If we have a reduction, we apply it
+        # point-wise
+        reduction = isa_reduction ? diag.reduction_time_func : (_, y) -> y
+
+        # If we have a reduction, we have to reset the accumulator to its neutral state. (If
+        # we don't have a reduction, we don't have to do anything)
+        #
+        # ClimaAtmos defines methods for identity_of_reduction for standard
+        # reduction_time_func in reduction_identities.jl
+        reset_accumulator! =
+            isa_reduction ?
+            () -> begin
+                # identity_of_reduction works by dispatching over Val{operation}
+                identity =
+                    identity_of_reduction(Val(diag.reduction_time_func))
+                # We also need to make sure that we are consistent with the types
+                float_type = eltype(storage[diag])
+                identity_ft = convert(float_type, identity)
+                storage[diag] .= identity_ft
+            end : () -> nothing
+
+        compute_callback =
+            integrator -> begin
+                # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
+                value = variable.compute_from_integrator(integrator, nothing)
+                storage[diag] .= reduction.(storage[diag], value)
+                counters[diag] += 1
+                return nothing
+            end
+
+        output_callback =
+            integrator -> begin
+                # Any operations we have to perform before writing to output?
+                # Here is where we would divide by N to obtain an arithmetic average
+                diag.pre_output_hook!(storage[diag], counters[diag])
+
+                # Write to disk
+                diag.output_writer(storage[diag], diag, integrator)
+
+                reset_accumulator!()
+                counters[diag] = 0
+                return nothing
+            end
+
+        # Here we have skip_first = true. This is important because we are going to manually
+        # call all the callbacks so that we can verify that they are meaningful for the
+        # model under consideration (and they don't have bugs).
+        append!(
+            callbacks,
+            [
+                call_every_n_steps(
+                    compute_callback,
+                    diag.compute_every,
+                    skip_first = true,
+                ),
+                call_every_n_steps(
+                    output_callback,
+                    diag.output_every,
+                    skip_first = true,
+                ),
+            ],
+        )
+    end
+
+    return callbacks
+end
diff --git a/src/diagnostics/Writers.jl b/src/diagnostics/writers.jl
similarity index 100%
rename from src/diagnostics/Writers.jl
rename to src/diagnostics/writers.jl
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 6429cb22f77..0447d8dadcc 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -6,6 +6,7 @@ using Interpolations
 import ClimaCore: InputOutput, Meshes, Spaces
 import ClimaAtmos.RRTMGPInterface as RRTMGPI
 import ClimaAtmos as CA
+import .Diagnostics as CAD
 import LinearAlgebra
 import ClimaCore.Fields
 import OrdinaryDiffEq as ODE
@@ -745,13 +746,13 @@ function get_integrator(config::AtmosConfig)
     # Initialize diagnostics
     @info "Initializing diagnostics"
 
-    diagnostics = get_default_diagnostics(atmos)
+    diagnostics = CAD.get_default_diagnostics(atmos)
 
     # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,
     # ensuring that there is consistency in the timestep and the periods and translating
     # those periods that depended on the timestep
     diagnostics_iterations =
-        [ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics]
+        [CAD.ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics]
 
     # For diagnostics that perform reductions, the storage is used as an accumulator, for
     # the other ones it is still defined to avoid allocating new space every time.
@@ -759,7 +760,7 @@ function get_integrator(config::AtmosConfig)
     diagnostic_counters = Dict()
 
     # NOTE: The diagnostics_callbacks are not called at the initial timestep
-    diagnostics_callbacks = get_callbacks_from_diagnostics(
+    diagnostics_callbacks = CAD.get_callbacks_from_diagnostics(
         diagnostics_iterations,
         diagnostic_storage,
         diagnostic_counters,

From bbba03e3edd8b992ae1782f7978b5e3e9360a6ca Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:20:05 -0700
Subject: [PATCH 20/73] Fix short/long names in documentation

---
 docs/src/diagnostics.md | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index d99d84306c0..1a61310117c 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -23,13 +23,13 @@ where to save it, and so on (read below for more information on this). You can
 construct your own lists of `ScheduledDiagnostic`s starting from the variables
 defined by `ClimaAtmos`. The diagnostics that `ClimaAtmos` knows how to compute
 are collected in a global dictionary called `ALL_DIAGNOSTICS`. The variables in
-`ALL_DIAGNOSTICS` are identified with a long and unique name, so that you can
-access them directly. One way to do so is by using the provided convenience
+`ALL_DIAGNOSTICS` are identified with by the short and unique name, so that you
+can access them directly. One way to do so is by using the provided convenience
 functions for common operations, e.g., continuing the previous example
 
 ```julia
 
-push!(diagnostics, get_daily_max("air_density", "air_temperature"))
+push!(diagnostics, daily_max("air_density", "air_temperature"))
 ```
 
 Now `diagnostics` will also contain the instructions to compute the daily
@@ -156,10 +156,10 @@ In `ClimaAtmos`, we follow the convention that:
 - `short_name` is the name used to identify the variable in the output files and
                 in the file names. It is short, but descriptive. We identify
                 diagnostics by their short name, so the diagnostics defined by
-                `ClimaAtmos` have to have unique `long_name`s. We follow the
-                Coupled Model Intercomparison Project (CMIP) convetions.
+                `ClimaAtmos` have to have unique `short_name`s. We follow the
+                Coupled Model Intercomparison Project (CMIP) conventions.
 
-- `long_name`: Name used to identify the variable in the input files.
+- `long_name`: Name used to describe the variable in the output file.
 
 - `units`: Physical units of the variable.
 
@@ -168,8 +168,10 @@ In `ClimaAtmos`, we follow the convention that:
 
 In `ClimaAtmos`, we try to follow [this Google
 spreadsheet](https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx)
-for variable naming (except for the `short_names`, which we prefer being more
-descriptive).
+for variable naming, with a `long_name` the does not have spaces and capital
+letters. [Standard
+names](http://cfconventions.org/Data/cf-standard-names/current/build/cf-standard-name-table.html)
+are not used.
 
 ### Compute function
 
@@ -223,7 +225,7 @@ were computed differently for `EquilMoistModel` and `NonEquilMoistModel`.
 ### Adding to the `ALL_DIAGNOSTICS` dictionary
 
 `ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable` in the
-`ALL_DIAGNOSTICS` dictionary. `ALL_DIAGNOSTICS` maps a `long_name` with the
+`ALL_DIAGNOSTICS` dictionary. `ALL_DIAGNOSTICS` maps a `short_name` with the
 corresponding `DiagnosticVariable`.
 
 If you are extending `ClimaAtmos` and want to add a new diagnostic variable to

From b977817905a153d60d78487e23a27a085e1bcd98 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:28:09 -0700
Subject: [PATCH 21/73] Rename compute_from_integrator

compute_from_integrator! now follows the julia conventions (it mutates
the first argument and it comes with a bang !)
---
 docs/src/diagnostics.md                 | 25 ++++++++++++++++---------
 src/diagnostics/core_diagnostics.jl     | 25 +++++++++++++------------
 src/diagnostics/diagnostic.jl           | 14 +++++++-------
 src/diagnostics/turbconv_diagnostics.jl | 20 ++++++++++----------
 src/solver/type_getters.jl              |  7 ++++---
 5 files changed, 50 insertions(+), 41 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 1a61310117c..0aebe1580dd 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -176,17 +176,17 @@ are not used.
 ### Compute function
 
 The other piece of information needed to specify a `DiagnosticVariable` is a
-function `compute_from_integrator`. Schematically, a `compute_from_integrator` has to look like
+function `compute_from_integrator!`. Schematically, a `compute_from_integrator!` has to look like
 ```julia
-function compute_from_integrator(integrator, out)
+function compute_from_integrator!(out, integrator)
     # FIXME: Remove this line when ClimaCore implements the broadcasting to enable this
     out .= # Calculcations with the state (= integrator.u) and the parameters (= integrator.p)
 end
 ```
 Diagnostics are implemented as callbacks function which pass the `integrator`
-object (from `OrdinaryDiffEq`) to `compute_from_integrator`.
+object (from `OrdinaryDiffEq`) to `compute_from_integrator!`.
 
-`compute_from_integrator` also takes a second argument, `out`, which is used to
+`compute_from_integrator!` also takes a second argument, `out`, which is used to
 avoid extra memory allocations (which hurt performance). If `out` is `nothing`,
 and new area of memory is allocated. If `out` is a `ClimaCore.Field`, the
 operation is done in-place without additional memory allocations.
@@ -197,17 +197,17 @@ For instance, if you want to compute relative humidity, which does not make
 sense for dry simulations, you should define the functions
 
 ```julia
-function compute_relative_humidity_from_integrator(
-    integrator,
+function compute_relative_humidity_from_integrator!(
     out,
+    integrator,
     moisture_model::T,
-)
+) where {T}
     error("Cannot compute relative_humidity with moisture_model = $T")
 end
 
-function compute_relative_humidity_from_integrator(
-    integrator,
+function compute_relative_humidity_from_integrator!(
     out,
+    integrator,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
@@ -215,6 +215,13 @@ function compute_relative_humidity_from_integrator(
     thermo_params = CAP.thermodynamics_params(integrator.p.params)
     return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
 end
+
+compute_relative_humidity_from_integrator!(out, integrator) =
+    compute_relative_humidity_from_integrator!(
+        out,
+        integrator,
+        integrator.p.atmos,
+    )
 ```
 
 This will return the correct relative humidity and throw informative errors when
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 1bffcf2b40c..72a02112267 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -9,7 +9,7 @@ add_diagnostic_variable!(
     long_name = "air_density",
     units = "kg m^-3",
     comments = "Density of air, a prognostic variable",
-    compute_from_integrator = (integrator, out) -> begin
+    compute_from_integrator! = (out, integrator) -> begin
         # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
         # We will want: out .= integrator.u.c.ρ
         return deepcopy(integrator.u.c.ρ)
@@ -18,17 +18,17 @@ add_diagnostic_variable!(
 
 # Relative humidity
 
-function compute_relative_humidity_from_integrator(
-    integrator,
+function compute_relative_humidity_from_integrator!(
     out,
+    integrator,
     moisture_model::T,
 ) where {T}
     error("Cannot compute relative_humidity with moisture_model = $T")
 end
 
-function compute_relative_humidity_from_integrator(
-    integrator,
+function compute_relative_humidity_from_integrator!(
     out,
+    integrator,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
@@ -37,6 +37,13 @@ function compute_relative_humidity_from_integrator(
     return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
 end
 
+compute_relative_humidity_from_integrator!(out, integrator) =
+    compute_relative_humidity_from_integrator!(
+        out,
+        integrator,
+        integrator.p.atmos,
+    )
+
 # FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
 # physics. Please fix this!
 add_diagnostic_variable!(
@@ -44,11 +51,5 @@ add_diagnostic_variable!(
     long_name = "relative_humidity",
     units = "",
     comments = "Relative Humidity",
-    compute_from_integrator = (integrator, out) -> begin
-        return compute_relative_humidity_from_integrator(
-            integrator,
-            out,
-            integrator.p.atmos.moisture_model,
-        )
-    end,
+    compute_from_integrator! = compute_relative_humidity_from_integrator!,
 )
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 4e7c0c22e08..657e1c8080a 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -77,7 +77,7 @@ Keyword arguments
 - `comments`: More verbose explanation of what the variable is, or comments related to how
               it is defined or computed.
 
-- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
+- `compute_from_integrator!`: Function that compute the diagnostic variable from the state.
                              It has to take two arguments: the `integrator`, and a
                              pre-allocated area of memory where to write the result of the
                              computation. It the no pre-allocated area is available, a new
@@ -90,7 +90,7 @@ Base.@kwdef struct DiagnosticVariable{T <: AbstractString, T2}
     long_name::T
     units::T
     comments::T
-    compute_from_integrator::T2
+    compute_from_integrator!::T2
 end
 
 # ClimaAtmos diagnostics
@@ -103,7 +103,7 @@ const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
                                long_name,
                                units,
                                description,
-                               compute_from_integrator)
+                               compute_from_integrator!)
 
 
 Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates the state of
@@ -127,7 +127,7 @@ Keyword arguments
 - `comments`: More verbose explanation of what the variable is, or comments related to how
               it is defined or computed.
 
-- `compute_from_integrator`: Function that compute the diagnostic variable from the state.
+- `compute_from_integrator!`: Function that compute the diagnostic variable from the state.
                              It has to take two arguments: the `integrator`, and a
                              pre-allocated area of memory where to write the result of the
                              computation. It the no pre-allocated area is available, a new
@@ -141,7 +141,7 @@ function add_diagnostic_variable!(;
     long_name,
     units,
     comments,
-    compute_from_integrator,
+    compute_from_integrator!,
 )
     haskey(ALL_DIAGNOSTICS, short_name) &&
         error("diagnostic $short_name already defined")
@@ -151,7 +151,7 @@ function add_diagnostic_variable!(;
         long_name,
         units,
         comments,
-        compute_from_integrator,
+        compute_from_integrator!,
     )
 end
 
@@ -591,7 +591,7 @@ function get_callbacks_from_diagnostics(diagnostics, storage, counters)
         compute_callback =
             integrator -> begin
                 # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
-                value = variable.compute_from_integrator(integrator, nothing)
+                value = variable.compute_from_integrator!(nothing, integrator)
                 storage[diag] .= reduction.(storage[diag], value)
                 counters[diag] += 1
                 return nothing
diff --git a/src/diagnostics/turbconv_diagnostics.jl b/src/diagnostics/turbconv_diagnostics.jl
index 9b2c8c15449..bd2797b5f6c 100644
--- a/src/diagnostics/turbconv_diagnostics.jl
+++ b/src/diagnostics/turbconv_diagnostics.jl
@@ -6,21 +6,24 @@
 # the model. This is also exposed to the user, which could define their own
 # compute_tke_from_integrator.
 
-function compute_tke_from_integrator(integrator, out, ::EDMFX)
+function compute_tke_from_integrator!(out, integrator, ::EDMFX)
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
     return deepcopy(integrator.p.ᶜspecific⁰.tke)
 end
 
-function compute_tke_from_integrator(integrator, out, ::DiagnosticEDMFX)
+function compute_tke_from_integrator!(out, integrator, ::DiagnosticEDMFX)
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-    return deepcopy(integrator.p.tke⁰)
+    return copy(integrator.p.tke⁰)
 end
 
-function compute_tke_from_integrator(
-    integrator,
+compute_tke_from_integrator!(out, integrator) =
+    compute_tke_from_integrator!(out, integrator, integrator.p.atmos)
+
+function compute_tke_from_integrator!(
     out,
+    integrator,
     turbconv_model::T,
 ) where {T}
     error("Cannot compute tke with turbconv_model = $T")
@@ -33,9 +36,6 @@ add_diagnostic_variable!(
     long_name = "turbolent_kinetic_energy",
     units = "J",
     comments = "Turbolent Kinetic Energy",
-    compute_from_integrator = (integrator, out) -> compute_tke_from_integrator(
-        integrator,
-        out,
-        integrator.p.atmos.turbconv_model,
-    ),
+    compute_from_integrator! = (integrator, out) ->
+        compute_tke_from_integrator!,
 )
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 0447d8dadcc..5921f01377d 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -751,8 +751,9 @@ function get_integrator(config::AtmosConfig)
     # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,
     # ensuring that there is consistency in the timestep and the periods and translating
     # those periods that depended on the timestep
-    diagnostics_iterations =
-        [CAD.ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics]
+    diagnostics_iterations = [
+        CAD.ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics
+    ]
 
     # For diagnostics that perform reductions, the storage is used as an accumulator, for
     # the other ones it is still defined to avoid allocating new space every time.
@@ -802,7 +803,7 @@ function get_integrator(config::AtmosConfig)
         try
             # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
             diagnostic_storage[diag] =
-                variable.compute_from_integrator(integrator, nothing)
+                variable.compute_from_integrator!(nothing, integrator)
             diagnostic_counters[diag] = 1
             # If it is not a reduction, call the output writer as well
             if isnothing(diag.reduction_time_func)

From 21cd34b1ad4310160434a07db2283a493629d763 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:30:08 -0700
Subject: [PATCH 22/73] Switch deepcopies to copies

---
 src/diagnostics/core_diagnostics.jl     | 2 +-
 src/diagnostics/turbconv_diagnostics.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 72a02112267..e4cec1d41d9 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -12,7 +12,7 @@ add_diagnostic_variable!(
     compute_from_integrator! = (out, integrator) -> begin
         # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
         # We will want: out .= integrator.u.c.ρ
-        return deepcopy(integrator.u.c.ρ)
+        return copy(integrator.u.c.ρ)
     end,
 )
 
diff --git a/src/diagnostics/turbconv_diagnostics.jl b/src/diagnostics/turbconv_diagnostics.jl
index bd2797b5f6c..38e556d6ac4 100644
--- a/src/diagnostics/turbconv_diagnostics.jl
+++ b/src/diagnostics/turbconv_diagnostics.jl
@@ -9,7 +9,7 @@
 function compute_tke_from_integrator!(out, integrator, ::EDMFX)
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-    return deepcopy(integrator.p.ᶜspecific⁰.tke)
+    return copy(integrator.p.ᶜspecific⁰.tke)
 end
 
 function compute_tke_from_integrator!(out, integrator, ::DiagnosticEDMFX)

From a229222ac870bd8d672ce6b5ee8430ddbb496301 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:30:48 -0700
Subject: [PATCH 23/73] Do not parametrize string in DiagnosticVariable

---
 src/diagnostics/diagnostic.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 657e1c8080a..eeb8049564c 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -85,12 +85,12 @@ Keyword arguments
                              function should perform the calculation in-place (i.e., using
                              `.=`).
 """
-Base.@kwdef struct DiagnosticVariable{T <: AbstractString, T2}
-    short_name::T
-    long_name::T
-    units::T
-    comments::T
-    compute_from_integrator!::T2
+Base.@kwdef struct DiagnosticVariable{T}
+    short_name::String
+    long_name::String
+    units::String
+    comments::String
+    compute_from_integrator!::T
 end
 
 # ClimaAtmos diagnostics

From 24aaa5c941471fb793a8b7775f4ae46f837856f2 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:31:30 -0700
Subject: [PATCH 24/73] Use more proper English

---
 src/diagnostics/diagnostic.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index eeb8049564c..70231e43892 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -2,7 +2,7 @@
 #
 # This file contains:
 #
-# - The definition of what a DiagnosticVariable is. Morally, a DiagnosticVariable is a
+# - The definition of what a DiagnosticVariable is. Conceptually, a DiagnosticVariable is a
 #   variable we know how to compute from the state. We attach more information to it for
 #   documentation and to reference to it with its short name. DiagnosticVariables can exist
 #   irrespective of the existence of an actual simulation that is being run. ClimaAtmos
@@ -13,7 +13,7 @@
 #   You can add your own file if you want to define several new diagnostics that are
 #   conceptually related.
 #
-# - The definition of what a ScheduledDiagnostics is. Morally, a ScheduledDiagnostics is a
+# - The definition of what a ScheduledDiagnostics is. Conceptually, a ScheduledDiagnostics is a
 #   DiagnosticVariable we want to compute in a given simulation. For example, it could be
 #   the temperature averaged over a day. We can have multiple ScheduledDiagnostics for the
 #   same DiagnosticVariable (e.g., daily and monthly average temperatures).

From fff3a933e29319f0761f279a3240233fc999d33e Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 15:52:49 -0700
Subject: [PATCH 25/73] Add name field to ScheduledDiagnostics

---
 src/diagnostics/diagnostic.jl | 80 ++++++++++++++++++++---------------
 src/diagnostics/writers.jl    | 58 +------------------------
 2 files changed, 46 insertions(+), 92 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 70231e43892..17ced28e5d8 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -162,6 +162,9 @@ include("turbconv_diagnostics.jl")
 # Default diagnostics and higher level interfaces
 include("default_diagnostics.jl")
 
+# Helper functions
+include("diagnostics_utils.jl")
+
 # ScheduledDiagnostics
 
 # NOTE: The definitions of ScheduledDiagnosticTime and ScheduledDiagnosticIterations are
@@ -196,6 +199,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     reduction_space_func::F2
     compute_every::T2
     pre_output_hook!::PO
+    name::String
 
     """
         ScheduledDiagnosticIterations(; variable::DiagnosticVariable,
@@ -204,7 +208,8 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                                         reduction_time_func = nothing,
                                         reduction_space_func = nothing,
                                         compute_every = isa_reduction ? 1 : output_every,
-                                        pre_output_hook! = (accum, count) -> nothing)
+                                        pre_output_hook! = (accum, count) -> nothing,
+                                        name = descriptive_name(self) )
 
 
     A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence
@@ -256,6 +261,11 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                           discarded. An example of `pre_output_hook!` to compute the arithmetic
                           average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
 
+   - `name`: A descriptive name for this particular diagnostic. If none is provided, one
+             will be generated mixing the short name of the variable, the reduction, and the
+             period of the reduction.
+
+
     """
     function ScheduledDiagnosticIterations(;
         variable::DiagnosticVariable,
@@ -265,25 +275,26 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         reduction_space_func = nothing,
         compute_every = isnothing(reduction_time_func) ? output_every : 1,
         pre_output_hook! = (accum, count) -> nothing,
-    )
-
-        # We provide an inner constructor to enforce some constraints
-        descriptive_name = get_descriptive_name(
+        name = get_descriptive_name(
             variable,
             output_every,
-            pre_output_hook!,
             reduction_time_func,
-        )
+            pre_output_hook!;
+            units_are_seconds = false,
+        ),
+    )
+
+        # We provide an inner constructor to enforce some constraints
 
         output_every % compute_every == 0 || error(
-            "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
+            "output_every should be multiple of compute_every for diagnostic $(name)",
         )
 
         isa_reduction = !isnothing(reduction_time_func)
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
+            @warn "output_every != compute_every for $(name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -302,6 +313,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
             reduction_space_func,
             compute_every,
             pre_output_hook!,
+            name,
         )
     end
 end
@@ -315,15 +327,17 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
     reduction_space_func::F2
     compute_every::T2
     pre_output_hook!::PO
+    name::String
 
     """
         ScheduledDiagnosticTime(; variable::DiagnosticVariable,
-                                            output_every,
-                                            output_writer,
-                                            reduction_time_func = nothing,
-                                            reduction_space_func = nothing,
-                                            compute_every = isa_reduction ? :timestep : output_every,
-                                            pre_output_hook! = (accum, count) -> nothing)
+                                  output_every,
+                                  output_writer,
+                                  reduction_time_func = nothing,
+                                  reduction_space_func = nothing,
+                                  compute_every = isa_reduction ? :timestep : output_every,
+                                  pre_output_hook! = (accum, count) -> nothing,
+                                  name = descriptive_name(self))
 
 
     A `DiagnosticVariable` that has to be computed and output during a simulation with a
@@ -379,6 +393,9 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                           discarded. An example of `pre_output_hook!` to compute the arithmetic
                           average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
 
+   - `name`: A descriptive name for this particular diagnostic. If none is provided, one
+             will be generated mixing the short name of the variable, the reduction, and the
+             period of the reduction.
     """
     function ScheduledDiagnosticTime(;
         variable::DiagnosticVariable,
@@ -389,22 +406,22 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         compute_every = isnothing(reduction_time_func) ? output_every :
                         :timestep,
         pre_output_hook! = (accum, count) -> nothing,
-    )
-
-        # We provide an inner constructor to enforce some constraints
-        descriptive_name = get_descriptive_name(
+        name = get_descriptive_name(
             variable,
             output_every,
             reduction_time_func,
             pre_output_hook!;
-            units_are_seconds = false,
-        )
+            units_are_seconds = true,
+        ),
+    )
+
+        # We provide an inner constructor to enforce some constraints
 
         # compute_every could be a Symbol (:timestep). We process this that when we process
         # the list of diagnostics
         if !isa(compute_every, Symbol)
             output_every % compute_every == 0 || error(
-                "output_every should be multiple of compute_every for diagnostic $(descriptive_name)",
+                "output_every should be multiple of compute_every for diagnostic $(name)",
             )
         end
 
@@ -412,7 +429,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(descriptive_name), changing compute_every to match"
+            @warn "output_every != compute_every for $(name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -431,6 +448,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
             reduction_space_func,
             compute_every,
             pre_output_hook!,
+            name,
         )
     end
 end
@@ -457,18 +475,11 @@ function ScheduledDiagnosticIterations(
         sd_time.compute_every == :timestep ? 1 : sd_time.compute_every / Δt
     output_every = sd_time.output_every / Δt
 
-    descriptive_name = get_descriptive_name(
-        sd_time.variable,
-        sd_time.output_every,
-        sd_time.reduction_time_func,
-        sd_time.pre_output_hook!,
-    )
-
     isinteger(output_every) || error(
-        "output_every should be multiple of the timestep for diagnostic $(descriptive_name)",
+        "output_every should be multiple of the timestep for diagnostic $(sd_time.name)",
     )
     isinteger(compute_every) || error(
-        "compute_every should be multiple of the timestep for diagnostic $(descriptive_name)",
+        "compute_every should be multiple of the timestep for diagnostic $(sd_time.name)",
     )
 
     ScheduledDiagnosticIterations(;
@@ -479,6 +490,7 @@ function ScheduledDiagnosticIterations(
         sd_time.reduction_space_func,
         compute_every = convert(Int, compute_every),
         sd_time.pre_output_hook!,
+        sd_time.name,
     )
 end
 
@@ -510,6 +522,7 @@ function ScheduledDiagnosticTime(
         sd_time.reduction_space_func,
         compute_every,
         sd_time.pre_output_hook!,
+        sd_time.name,
     )
 end
 
@@ -528,9 +541,6 @@ ScheduledDiagnosticIterations(
 # We define all the known identities in reduction_identities.jl
 include("reduction_identities.jl")
 
-# Helper functions
-include("diagnostics_utils.jl")
-
 """
     get_callbacks_from_diagnostics(diagnostics, storage, counters)
 
diff --git a/src/diagnostics/writers.jl b/src/diagnostics/writers.jl
index 3a7d0062636..5aa04901145 100644
--- a/src/diagnostics/writers.jl
+++ b/src/diagnostics/writers.jl
@@ -3,62 +3,6 @@
 # This file defines function-generating functions for output_writers for diagnostics. The
 # writers come with opinionated defaults.
 
-"""
-    get_descriptive_name(sd_t::ScheduledDiagnosticTime)
-
-
-Return a compact, unique-ish, identifier for the given `ScheduledDiagnosticTime` `sd_t`.
-
-We split the period in seconds into days, hours, minutes, seconds. In most cases
-(with standard periods), this output will look something like
-`air_density_1d_max`. This function can be used for filenames.
-
-The name is not unique to the `ScheduledDiagnosticTime` because it ignores other
-parameters such as whether there is a reduction in space or the compute
-frequency.
-
-"""
-get_descriptive_name(sd_t::ScheduledDiagnosticTime) = get_descriptive_name(
-    sd_t.variable,
-    sd_t.output_every,
-    sd_t.reduction_time_func;
-    units_are_seconds = true,
-)
-
-"""
-    get_descriptive_name(sd_i::ScheduledDiagnosticIterations[, Δt])
-
-
-Return a compact, unique-ish, identifier for the given
-`ScheduledDiagnosticIterations` `sd_i`.
-
-If the timestep `Δt` is provided, convert the steps into seconds. In this case,
-the output will look like `air_density_1d_max`. Otherwise, the output will look
-like `air_density_100it_max`. This function can be used for filenames.
-
-The name is not unique to the `ScheduledDiagnosticIterations` because it ignores
-other parameters such as whether there is a reduction in space or the compute
-frequency.
-
-"""
-get_descriptive_name(sd_i::ScheduledDiagnosticIterations, Δt::Nothing) =
-    get_descriptive_name(
-        sd_t.variable,
-        sd_t.output_every,
-        sd_t.reduction_time_func,
-        sd_t.pre_output_hook!;
-        units_are_seconds = false,
-    )
-get_descriptive_name(sd_i::ScheduledDiagnosticIterations, Δt::T) where {T} =
-    get_descriptive_name(
-        sd_i.variable,
-        sd_i.output_every * Δt,
-        sd_i.reduction_time_func,
-        sd_i.pre_output_hook!;
-        units_are_seconds = true,
-    )
-
-
 """
     HDF5Writer()
 
@@ -95,7 +39,7 @@ function HDF5Writer()
 
         output_path = joinpath(
             integrator.p.simulation.output_dir,
-            "$(get_descriptive_name(diagnostic, integrator.p.simulation.dt))_$time.h5",
+            "$(diagnostic.name)_$time.h5",
         )
 
         hdfwriter = InputOutput.HDF5Writer(output_path, integrator.p.comms_ctx)

From 36a525f417f8cd333994225fa1f13bbe7676c967 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 16:26:43 -0700
Subject: [PATCH 26/73] Add numeric values to error messages

---
 src/diagnostics/diagnostic.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 17ced28e5d8..76b3709e622 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -287,14 +287,14 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         # We provide an inner constructor to enforce some constraints
 
         output_every % compute_every == 0 || error(
-            "output_every should be multiple of compute_every for diagnostic $(name)",
+            "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(name)",
         )
 
         isa_reduction = !isnothing(reduction_time_func)
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(name), changing compute_every to match"
+            @warn "output_every ($output_every) != compute_every ($compute_every) for $(name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -421,7 +421,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         # the list of diagnostics
         if !isa(compute_every, Symbol)
             output_every % compute_every == 0 || error(
-                "output_every should be multiple of compute_every for diagnostic $(name)",
+                "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(name)",
             )
         end
 
@@ -429,7 +429,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every != compute_every for $(name), changing compute_every to match"
+            @warn "output_every ($output_every) != compute_every ($compute_every) for $(name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -476,10 +476,10 @@ function ScheduledDiagnosticIterations(
     output_every = sd_time.output_every / Δt
 
     isinteger(output_every) || error(
-        "output_every should be multiple of the timestep for diagnostic $(sd_time.name)",
+        "output_every ($output_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.name)",
     )
     isinteger(compute_every) || error(
-        "compute_every should be multiple of the timestep for diagnostic $(sd_time.name)",
+        "compute_every ($compute_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.name)",
     )
 
     ScheduledDiagnosticIterations(;

From d53cf574936c897ca69400e7f2b26f809ba5fd79 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 16:28:13 -0700
Subject: [PATCH 27/73] Do not use Val for functions

---
 docs/src/diagnostics.md                 | 8 ++++----
 src/diagnostics/diagnostic.jl           | 3 +--
 src/diagnostics/reduction_identities.jl | 8 ++++----
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 0aebe1580dd..544937ba169 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -88,10 +88,10 @@ the function needed to compute averages `ClimaAtmos.average_pre_output_hook!`.
 
 For custom reductions, it is necessary to also specify the identity of operation
 by defining a new method to `identity_of_reduction`. `identity_of_reduction` is
-a function that takes a `Val{op}` argument, where `op` is the operation for
-which we want to define the identity. For instance, for the `max`,
-`identity_of_reduction` would be `identity_of_reduction(::Val{max}) = -Inf`. The
-identities known to `ClimaAtmos` are defined in the
+a function that takes a `op` argument, where `op` is the operation for which we
+want to define the identity. For instance, for the `max`,
+`identity_of_reduction` would be `identity_of_reduction(::typeof{max}) = -Inf`.
+The identities known to `ClimaAtmos` are defined in the
 `diagnostics/reduction_identities.jl` file. The identity is needed to ensure
 that we have a neutral state for the accumulators that are used in the
 reductions.
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 76b3709e622..dcc6f213864 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -589,9 +589,8 @@ function get_callbacks_from_diagnostics(diagnostics, storage, counters)
         reset_accumulator! =
             isa_reduction ?
             () -> begin
-                # identity_of_reduction works by dispatching over Val{operation}
                 identity =
-                    identity_of_reduction(Val(diag.reduction_time_func))
+                    identity_of_reduction(diag.reduction_time_func)
                 # We also need to make sure that we are consistent with the types
                 float_type = eltype(storage[diag])
                 identity_ft = convert(float_type, identity)
diff --git a/src/diagnostics/reduction_identities.jl b/src/diagnostics/reduction_identities.jl
index e75d15957b5..0fce5e58845 100644
--- a/src/diagnostics/reduction_identities.jl
+++ b/src/diagnostics/reduction_identities.jl
@@ -13,7 +13,7 @@
 # We have to know the identity for every operation we want to support. Of course, users are
 # welcome to define their own by adding new methods to identity_of_reduction.
 
-identity_of_reduction(::Val{max}) = -Inf
-identity_of_reduction(::Val{min}) = +Inf
-identity_of_reduction(::Val{+}) = 0
-identity_of_reduction(::Val{*}) = 1
+identity_of_reduction(::typeof(max)) = -Inf
+identity_of_reduction(::typeof(min)) = +Inf
+identity_of_reduction(::typeof(+)) = 0
+identity_of_reduction(::typeof(*)) = 1

From 6a0dd0b211b909b1c67711601bd9d5fe7062c662 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 16:31:29 -0700
Subject: [PATCH 28/73] Use push! instead of append!

---
 src/diagnostics/diagnostic.jl | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index dcc6f213864..c6a64f6af04 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -623,20 +623,18 @@ function get_callbacks_from_diagnostics(diagnostics, storage, counters)
         # Here we have skip_first = true. This is important because we are going to manually
         # call all the callbacks so that we can verify that they are meaningful for the
         # model under consideration (and they don't have bugs).
-        append!(
+        push!(
             callbacks,
-            [
-                call_every_n_steps(
-                    compute_callback,
-                    diag.compute_every,
-                    skip_first = true,
-                ),
-                call_every_n_steps(
-                    output_callback,
-                    diag.output_every,
-                    skip_first = true,
-                ),
-            ],
+            call_every_n_steps(
+                compute_callback,
+                diag.compute_every,
+                skip_first=true,
+            ),
+            call_every_n_steps(
+                output_callback,
+                diag.output_every,
+                skip_first=true,
+            ),
         )
     end
 

From c821638cbc6bf101d6635736cb7dcc70c6de4721 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 17:14:25 -0700
Subject: [PATCH 29/73] Add diagnostic accumulators

---
 src/diagnostics/diagnostic.jl | 91 +++++++++++++++++++++++------------
 src/solver/type_getters.jl    |  9 +++-
 2 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index c6a64f6af04..198a8452bef 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -541,6 +541,34 @@ ScheduledDiagnosticIterations(
 # We define all the known identities in reduction_identities.jl
 include("reduction_identities.jl")
 
+# Helper functions for the callbacks:
+# - reset_accumulator!
+# - accumulate!
+
+# When the reduction is nothing, do nothing
+reset_accumulator!(_, reduction_time_func::Nothing) = nothing
+
+# If we have a reduction, we have to reset the accumulator to its neutral state. (If we
+# don't have a reduction, we don't have to do anything)
+#
+# ClimaAtmos defines methods for identity_of_reduction for standard reduction_time_func in
+# reduction_identities.jl
+function reset_accumulator!(diag_accumulator, reduction_time_func)
+    # identity_of_reduction works by dispatching over operation
+    identity = identity_of_reduction(reduction_time_func)
+    float_type = eltype(diag_accumulator)
+    identity_ft = convert(float_type, identity)
+    diag_accumulator .= identity_ft
+end
+
+# When the reduction is nothing, we do not need to accumulate anything
+accumulate!(_, _, reduction_time_func::Nothing) = nothing
+
+# When we have a reduction, apply it between the accumulated value one
+function accumulate!(diag_accumulator, diag_storage, reduction_time_func)
+    diag_accumulator .= reduction_time_func.(diag_accumulator, diag_storage)
+end
+
 """
     get_callbacks_from_diagnostics(diagnostics, storage, counters)
 
@@ -555,14 +583,22 @@ Positional arguments
                  can define callbacks that occur at the end of every N integration steps.
 
 - `storage`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially
-             pre-allocated area of memory where to accumulate/save results.
+             pre-allocated area of memory where to save the newly computed results.
+
+- `accumulator`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially
+                 pre-allocated area of memory where to accumulate results.
 
 - `counters`: Dictionary that maps a given `ScheduledDiagnosticIterations` to the counter
               that tracks how many times the given diagnostics was computed from the last
               time it was output to disk.
 
 """
-function get_callbacks_from_diagnostics(diagnostics, storage, counters)
+function get_callbacks_from_diagnostics(
+    diagnostics,
+    storage,
+    accumulators,
+    counters,
+)
     # We have two types of callbacks: to compute and accumulate diagnostics, and to dump
     # them to disk. Note that our callbacks do not contain any branching
 
@@ -573,41 +609,33 @@ function get_callbacks_from_diagnostics(diagnostics, storage, counters)
 
     for diag in diagnostics
         variable = diag.variable
-        isa_reduction = !isnothing(diag.reduction_time_func)
-
-        # reduction is used below. If we are not given a reduction_time_func, we just want
-        # to move the computed quantity to its storage (so, we return the second argument,
-        # which that will be the newly computed one). If we have a reduction, we apply it
-        # point-wise
-        reduction = isa_reduction ? diag.reduction_time_func : (_, y) -> y
-
-        # If we have a reduction, we have to reset the accumulator to its neutral state. (If
-        # we don't have a reduction, we don't have to do anything)
-        #
-        # ClimaAtmos defines methods for identity_of_reduction for standard
-        # reduction_time_func in reduction_identities.jl
-        reset_accumulator! =
-            isa_reduction ?
-            () -> begin
-                identity =
-                    identity_of_reduction(diag.reduction_time_func)
-                # We also need to make sure that we are consistent with the types
-                float_type = eltype(storage[diag])
-                identity_ft = convert(float_type, identity)
-                storage[diag] .= identity_ft
-            end : () -> nothing
 
         compute_callback =
             integrator -> begin
                 # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
-                value = variable.compute_from_integrator!(nothing, integrator)
-                storage[diag] .= reduction.(storage[diag], value)
+                variable.compute_from_integrator!(storage[diag], integrator)
+
+                # accumulator[diag] is not defined for non-reductions
+                diag_accumulator = get(accumulators, diag, nothing)
+
+                accumulate!(
+                    diag_accumulator,
+                    storage[diag],
+                    diag.reduction_time_func,
+                )
                 counters[diag] += 1
                 return nothing
             end
 
         output_callback =
             integrator -> begin
+                # Move accumulated value to storage so that we can output it (for
+                # reductions). This provides a unified interface to pre_output_hook! and
+                # output, at the cost of an additional copy. If this copy turns out to be
+                # too expensive, we can move the if statement below.
+                isnothing(diag.reduction_time_func) ||
+                    (storage[diag] .= accumulators[diag])
+
                 # Any operations we have to perform before writing to output?
                 # Here is where we would divide by N to obtain an arithmetic average
                 diag.pre_output_hook!(storage[diag], counters[diag])
@@ -615,7 +643,10 @@ function get_callbacks_from_diagnostics(diagnostics, storage, counters)
                 # Write to disk
                 diag.output_writer(storage[diag], diag, integrator)
 
-                reset_accumulator!()
+                # accumulator[diag] is not defined for non-reductions
+                diag_accumulator = get(accumulators, diag, nothing)
+
+                reset_accumulator!(diag_accumulator, diag.reduction_time_func)
                 counters[diag] = 0
                 return nothing
             end
@@ -628,12 +659,12 @@ function get_callbacks_from_diagnostics(diagnostics, storage, counters)
             call_every_n_steps(
                 compute_callback,
                 diag.compute_every,
-                skip_first=true,
+                skip_first = true,
             ),
             call_every_n_steps(
                 output_callback,
                 diag.output_every,
-                skip_first=true,
+                skip_first = true,
             ),
         )
     end
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 5921f01377d..7ae51aaf5df 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -755,15 +755,17 @@ function get_integrator(config::AtmosConfig)
         CAD.ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics
     ]
 
-    # For diagnostics that perform reductions, the storage is used as an accumulator, for
-    # the other ones it is still defined to avoid allocating new space every time.
+    # For diagnostics that perform reductions, the storage is used for the values computed
+    # at each call. Reductions also save the accumulated value in in diagnostic_accumulators.
     diagnostic_storage = Dict()
+    diagnostic_accumulators = Dict()
     diagnostic_counters = Dict()
 
     # NOTE: The diagnostics_callbacks are not called at the initial timestep
     diagnostics_callbacks = CAD.get_callbacks_from_diagnostics(
         diagnostics_iterations,
         diagnostic_storage,
+        diagnostic_accumulators,
         diagnostic_counters,
     )
 
@@ -808,6 +810,9 @@ function get_integrator(config::AtmosConfig)
             # If it is not a reduction, call the output writer as well
             if isnothing(diag.reduction_time_func)
                 diag.output_writer(diagnostic_storage[diag], diag, integrator)
+            else
+                # Add to the accumulator
+                diagnostic_accumulators[diag] = copy(diagnostic_storage[diag])
             end
         catch e
             error("Could not compute diagnostic $(variable.long_name): $e")

From 6d07133f686a78d295e707cf7ebe1698bc735b1c Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 8 Sep 2023 17:25:15 -0700
Subject: [PATCH 30/73] Refactor produce_common_diagnostic_function

Refactor closure in favor of a function with more arguments
---
 src/diagnostics/default_diagnostics.jl     | 60 +++++++++++++---------
 src/diagnostics/defaults/moisture_model.jl |  3 +-
 2 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index 59fedef52da..9b39525bf15 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -37,23 +37,23 @@ get_default_diagnostics(_) = []
 
 Helper function to define functions like `get_daily_max`.
 """
-function produce_common_diagnostic_function(
+function common_diagnostics(
     period,
-    reduction;
+    reduction,
+    output_writer,
+    short_names...;
     pre_output_hook! = (accum, count) -> nothing,
 )
-    return (long_names...; output_writer = HDF5Writer()) -> begin
-        [
-            ScheduledDiagnosticTime(
-                variable = ALL_DIAGNOSTICS[long_name],
-                compute_every = :timestep,
-                output_every = period, # seconds
-                reduction_time_func = reduction,
-                output_writer = output_writer,
-                pre_output_hook! = pre_output_hook!,
-            ) for long_name in long_names
-        ]
-    end
+    return [
+        ScheduledDiagnosticTime(
+            variable = ALL_DIAGNOSTICS[short_name],
+            compute_every = :timestep,
+            output_every = period, # seconds
+            reduction_time_func = reduction,
+            output_writer = output_writer,
+            pre_output_hook! = pre_output_hook!,
+        ) for short_name in short_names
+    ]
 end
 
 function average_pre_output_hook!(accum, counter)
@@ -67,14 +67,17 @@ end
 
 Return a list of `ScheduledDiagnostics` that compute the daily max for the given variables.
 """
-get_daily_max = produce_common_diagnostic_function(24 * 60 * 60, max)
+daily_max(long_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(24 * 60 * 60, max, output_writer, long_names...)
+
 """
     get_daily_min(long_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily min for the given variables.
 """
-get_daily_min = produce_common_diagnostic_function(24 * 60 * 60, min)
+daily_min(long_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(24 * 60 * 60, min, output_writer, long_names...)
 """
     get_daily_average(long_names...; output_writer = HDF5Writer())
 
@@ -82,9 +85,11 @@ get_daily_min = produce_common_diagnostic_function(24 * 60 * 60, min)
 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
-get_daily_average = produce_common_diagnostic_function(
+daily_average(long_names...; output_writer = HDF5Writer()) = common_diagnostics(
     24 * 60 * 60,
-    (+);
+    (+),
+    output_writer,
+    long_names...;
     pre_output_hook! = average_pre_output_hook!,
 )
 
@@ -94,7 +99,8 @@ get_daily_average = produce_common_diagnostic_function(
 
 Return a list of `ScheduledDiagnostics` that compute the hourly max for the given variables.
 """
-get_hourly_max = produce_common_diagnostic_function(60 * 60, max)
+hourly_max(long_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(60 * 60, max, output_writer, long_names...)
 
 """
     get_hourly_min(long_names...; output_writer = HDF5Writer())
@@ -102,7 +108,8 @@ get_hourly_max = produce_common_diagnostic_function(60 * 60, max)
 
 Return a list of `ScheduledDiagnostics` that compute the hourly min for the given variables.
 """
-get_hourly_min = produce_common_diagnostic_function(60 * 60, min)
+hourly_min(long_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(60 * 60, min, output_writer, long_names...)
 
 """
     get_daily_average(long_names...; output_writer = HDF5Writer())
@@ -112,11 +119,14 @@ Return a list of `ScheduledDiagnostics` that compute the hourly average for the
 """
 
 # An average is just a sum with a normalization before output
-get_hourly_average = produce_common_diagnostic_function(
-    60 * 60,
-    (+);
-    pre_output_hook! = average_pre_output_hook!,
-)
+hourly_average(long_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(
+        60 * 60,
+        (+),
+        output_writer,
+        long_names...;
+        pre_output_hook! = average_pre_output_hook!,
+    )
 
 # Include all the subdefaults
 include("defaults/moisture_model.jl")
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
index 14b26d0eb75..fbf01d3c05c 100644
--- a/src/diagnostics/defaults/moisture_model.jl
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -1,8 +1,7 @@
 # FIXME: Gabriele added this as an example. Put something meaningful here!
 function get_default_diagnostics(::T) where {T <: DryModel}
     return vcat(
-        get_daily_average("air_density"),
-        get_hourly_max("air_density"),
+        daily_average("air_density"),
         [
             ScheduledDiagnosticTime(
                 variable = ALL_DIAGNOSTICS["air_density"],

From 9515540c930810c98d3bc1571421c25856bae4b9 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 21 Sep 2023 11:07:15 -0700
Subject: [PATCH 31/73] Add flag to output default diagnostics

---
 config/default_configs/default_config.yml | 3 +++
 src/solver/type_getters.jl                | 4 +++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
index 46ed8faadff..d7c8a435e73 100644
--- a/config/default_configs/default_config.yml
+++ b/config/default_configs/default_config.yml
@@ -249,3 +249,6 @@ override_τ_precip:
 log_params:
   help: "Log parameters to file [`false` (default), `true`]"
   value: false
+output_default_diagnostics:
+  help: "Output the default diagnostics associated to the selected atmospheric model"
+  value: true
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 7ae51aaf5df..723f0b324cf 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -746,7 +746,9 @@ function get_integrator(config::AtmosConfig)
     # Initialize diagnostics
     @info "Initializing diagnostics"
 
-    diagnostics = CAD.get_default_diagnostics(atmos)
+    diagnostics =
+        config.parsed_args["output_default_diagnostics"] ?
+        CAD.get_default_diagnostics(atmos) : []
 
     # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,
     # ensuring that there is consistency in the timestep and the periods and translating

From 80e95db19ea4e2aaf609e1c94c8ccfd365044ea6 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 08:47:13 -0700
Subject: [PATCH 32/73] Accept nothing as pre_output_hook!

---
 src/diagnostics/default_diagnostics.jl |  2 +-
 src/diagnostics/diagnostic.jl          | 22 ++++++++++++++++++----
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index 9b39525bf15..452df573093 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -42,7 +42,7 @@ function common_diagnostics(
     reduction,
     output_writer,
     short_names...;
-    pre_output_hook! = (accum, count) -> nothing,
+    pre_output_hook! = nothing,
 )
     return [
         ScheduledDiagnosticTime(
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 198a8452bef..3286a2f4208 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -208,7 +208,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                                         reduction_time_func = nothing,
                                         reduction_space_func = nothing,
                                         compute_every = isa_reduction ? 1 : output_every,
-                                        pre_output_hook! = (accum, count) -> nothing,
+                                        pre_output_hook! = nothing,
                                         name = descriptive_name(self) )
 
 
@@ -274,7 +274,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         reduction_time_func = nothing,
         reduction_space_func = nothing,
         compute_every = isnothing(reduction_time_func) ? output_every : 1,
-        pre_output_hook! = (accum, count) -> nothing,
+        pre_output_hook! = nothing,
         name = get_descriptive_name(
             variable,
             output_every,
@@ -298,6 +298,13 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
             compute_every = output_every
         end
 
+        # pre_output_hook! has to be a function, but it is much more intuitive to specify
+        # `nothing` when we want nothing to happen. Here, we convert the nothing keyword
+        # into a function that does nothing
+        if isnothing(pre_output_hook!)
+            pre_output_hook! = (accum, count) -> nothing
+        end
+
         T1 = typeof(output_every)
         T2 = typeof(compute_every)
         OW = typeof(output_writer)
@@ -336,7 +343,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                                   reduction_time_func = nothing,
                                   reduction_space_func = nothing,
                                   compute_every = isa_reduction ? :timestep : output_every,
-                                  pre_output_hook! = (accum, count) -> nothing,
+                                  pre_output_hook! = nothing,
                                   name = descriptive_name(self))
 
 
@@ -405,7 +412,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         reduction_space_func = nothing,
         compute_every = isnothing(reduction_time_func) ? output_every :
                         :timestep,
-        pre_output_hook! = (accum, count) -> nothing,
+        pre_output_hook! = nothing,
         name = get_descriptive_name(
             variable,
             output_every,
@@ -433,6 +440,13 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
             compute_every = output_every
         end
 
+        # pre_output_hook! has to be a function, but it is much more intuitive to specify
+        # `nothing` when we want nothing to happen. Here, we convert the nothing keyword
+        # into a function that does nothing
+        if isnothing(pre_output_hook!)
+            pre_output_hook! = (accum, count) -> nothing
+        end
+
         T1 = typeof(output_every)
         T2 = typeof(compute_every)
         OW = typeof(output_writer)

From 96c085dd4ff66f94d41b73b23be880b5a20a6cff Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 11 Sep 2023 10:32:31 -0700
Subject: [PATCH 33/73] Add YAML reader for diagnostics

---
 docs/src/diagnostics.md    | 24 +++++++++++
 src/solver/type_getters.jl | 86 ++++++++++++++++++++++++++++++++++++--
 src/solver/yaml_helper.jl  |  7 ++++
 3 files changed, 113 insertions(+), 4 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 544937ba169..1a04a479d7e 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -2,6 +2,30 @@
 
 ## I want to compute and output a diagnostic variable
 
+### From a YAML file
+
+If you configure your simulation with YAML files, there are two options that are
+important to know about. When `output_default_diagnostics` is set to `true`, the
+default diagnostics for the given atmospheric model will be output. Note that
+they might be incompatible with your simulation (e.g., you want to output
+hourly maxima when the timestep is 4 hours).
+
+Second, you can specify the diagnostics you want to output directly in the
+`diagnostics` section of your YAML file. For instance:
+```
+diagnostics:
+  - short_name: air_density
+    name: a_name
+    period: 3hours
+  - reduction_time: average
+    short_name: air_density
+    period: 12hours
+```
+This adds two diagnostics (both for `air_density`). The `period` keyword
+identifies the period over which to compute the reduction and how often to save
+to disk. `name` is optional, and if provided, it identifies the name of the
+output file.
+
 ### From a script
 
 The simplest way to get started with diagnostics is to use the defaults for your
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 723f0b324cf..f31634fd01e 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -607,6 +607,87 @@ function get_simulation(config::AtmosConfig, comms_ctx)
     return sim
 end
 
+function get_diagnostics(parsed_args, atmos_model)
+
+    diagnostics =
+        parsed_args["output_default_diagnostics"] ?
+        CAD.get_default_diagnostics(atmos_model) : Any[]
+
+    # We either get the diagnostics section in the YAML file, or we return an empty
+    # dictionary (which will result in an empty list being created by the map below)
+    yaml_diagnostics = get(parsed_args, "diagnostics", Dict())
+
+    # ALLOWED_REDUCTIONS is the collection of reductions we support. The keys are the
+    # strings that have to be provided in the YAML file. The values are tuples with the
+    # function that has to be passed to reduction_time_func and the one that has to passed
+    # to pre_output_hook!
+
+    # We make "nothing" a string so that we can accept also the word "nothing", in addition
+    # to the absence of the value
+    #
+    # NOTE: Everything has to be lowercase in ALLOWED_REDUCTIONS (so that we can match
+    # "max" and "Max")
+    ALLOWED_REDUCTIONS = Dict(
+        "nothing" => (nothing, nothing), # nothing is: just dump the variable
+        "max" => (max, nothing),
+        "min" => (min, nothing),
+        "average" => ((+), CAD.average_pre_output_hook!),
+    )
+
+    for yaml_diag in yaml_diagnostics
+        # Return "nothing" if "reduction_time" is not in the YAML block
+        #
+        # We also normalize everything to lowercase, so that can accept "max" but
+        # also "Max"
+        reduction_time_yaml =
+            lowercase(get(yaml_diag, "reduction_time", "nothing"))
+
+        if !haskey(ALLOWED_REDUCTIONS, reduction_time_yaml)
+            error("reduction $reduction_time_yaml not implemented")
+        else
+            reduction_time_func, pre_output_hook! =
+                ALLOWED_REDUCTIONS[reduction_time_yaml]
+        end
+
+        name = get(yaml_diag, "name", nothing)
+
+        haskey(yaml_diag, "period") ||
+            error("period keyword required for diagnostics")
+
+        period_seconds = time_to_seconds(yaml_diag["period"])
+
+        if isnothing(name)
+            name = CAD.get_descriptive_name(
+                CAD.ALL_DIAGNOSTICS[yaml_diag["short_name"]],
+                period_seconds,
+                reduction_time_func,
+                pre_output_hook!,
+            )
+        end
+
+        if isnothing(reduction_time_func)
+            compute_every = period_seconds
+        else
+            compute_every = :timestep
+        end
+
+        push!(
+            diagnostics,
+            CAD.ScheduledDiagnosticTime(
+                variable = CAD.ALL_DIAGNOSTICS[yaml_diag["short_name"]],
+                output_every = period_seconds,
+                compute_every = compute_every,
+                reduction_time_func = reduction_time_func,
+                pre_output_hook! = pre_output_hook!,
+                output_writer = CAD.HDF5Writer(),
+                name = name,
+            ),
+        )
+    end
+
+    return diagnostics
+end
+
 function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback)
     (; atmos, simulation) = p
     (; dt) = simulation
@@ -745,10 +826,7 @@ function get_integrator(config::AtmosConfig)
 
     # Initialize diagnostics
     @info "Initializing diagnostics"
-
-    diagnostics =
-        config.parsed_args["output_default_diagnostics"] ?
-        CAD.get_default_diagnostics(atmos) : []
+    diagnostics = get_diagnostics(config.parsed_args, atmos)
 
     # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,
     # ensuring that there is consistency in the timestep and the periods and translating
diff --git a/src/solver/yaml_helper.jl b/src/solver/yaml_helper.jl
index 399532a30ce..800ac198583 100644
--- a/src/solver/yaml_helper.jl
+++ b/src/solver/yaml_helper.jl
@@ -74,6 +74,13 @@ function override_default_config(config_dict::AbstractDict;)
         v = config_dict[k]
         config[k] = isnothing(default_config[k]) ? v : default_type(v)
     end
+
+    # The "diagnostics" entry is a more complex type that doesn't fit the schema described in
+    # the previous lines. So, we manually add it.
+    if haskey(config_dict, "diagnostics")
+        config["diagnostics"] = config_dict["diagnostics"]
+    end
+
     return config
 end
 

From 3a401eb8b75a5cef216ce70bcdaaa5a716584186 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 11 Sep 2023 12:04:48 -0700
Subject: [PATCH 34/73] Clean up daily/hourly functions

---
 src/diagnostics/default_diagnostics.jl     | 97 +++++++++++++++++-----
 src/diagnostics/defaults/moisture_model.jl | 34 ++++----
 2 files changed, 90 insertions(+), 41 deletions(-)

diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index 452df573093..dcecc48e600 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -62,71 +62,122 @@ function average_pre_output_hook!(accum, counter)
 end
 
 """
-    get_daily_max(long_names...; output_writer = HDF5Writer())
+    daily_maxs(short_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily max for the given variables.
 """
-daily_max(long_names...; output_writer = HDF5Writer()) =
-    common_diagnostics(24 * 60 * 60, max, output_writer, long_names...)
+daily_maxs(short_names...; output_writer = HDF5Writer()) =
+    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.
 """
-    get_daily_min(long_names...; output_writer = HDF5Writer())
+daily_max(short_names; output_writer = HDF5Writer()) =
+    daily_maxs(short_names; output_writer)[1]
+
+"""
+    daily_mins(short_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the daily min for the given variables.
 """
-daily_min(long_names...; output_writer = HDF5Writer()) =
-    common_diagnostics(24 * 60 * 60, min, output_writer, long_names...)
+daily_mins(short_names...; output_writer = HDF5Writer()) =
+    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_mins(short_names; output_writer)[1]
+
 """
-    get_daily_average(long_names...; output_writer = HDF5Writer())
+    daily_averages(short_names...; output_writer = HDF5Writer())
 
 
 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_average(long_names...; output_writer = HDF5Writer()) = common_diagnostics(
-    24 * 60 * 60,
-    (+),
-    output_writer,
-    long_names...;
-    pre_output_hook! = average_pre_output_hook!,
-)
+daily_averages(short_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(
+        24 * 60 * 60,
+        (+),
+        output_writer,
+        short_names...;
+        pre_output_hook! = average_pre_output_hook!,
+    )
+"""
+    daily_average(short_names; output_writer = HDF5Writer())
 
+
+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_averages(short_names; output_writer)[1]
 """
-    get_hourly_max(long_names...; output_writer = HDF5Writer())
+    hourly_maxs(short_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the hourly max for the given variables.
 """
-hourly_max(long_names...; output_writer = HDF5Writer()) =
-    common_diagnostics(60 * 60, max, output_writer, long_names...)
+hourly_maxs(short_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(60 * 60, max, output_writer, short_names...)
+
+"""
+    hourly_max(short_names; output_writer = HDF5Writer())
+
+
+Return a `ScheduledDiagnostics` that computse the hourly max for the given variable.
+"""
+hourly_max(short_names...; output_writer = HDF5Writer()) =
+    hourly_maxs(short_names)[1]
 
 """
-    get_hourly_min(long_names...; output_writer = HDF5Writer())
+    hourly_mins(short_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the hourly min for the given variables.
 """
-hourly_min(long_names...; output_writer = HDF5Writer()) =
-    common_diagnostics(60 * 60, min, output_writer, long_names...)
+hourly_mins(short_names...; output_writer = HDF5Writer()) =
+    common_diagnostics(60 * 60, min, output_writer, short_names...)
+"""
+    hourly_mins(short_names...; output_writer = HDF5Writer())
+
+
+Return a `ScheduledDiagnostics` that computes the hourly min for the given variable.
+"""
+hourly_min(short_names; output_writer = HDF5Writer()) =
+    hourly_mins(short_names; output_writer)[1]
 
 """
-    get_daily_average(long_names...; output_writer = HDF5Writer())
+    hourly_averages(short_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the hourly average for the given variables.
 """
 
 # An average is just a sum with a normalization before output
-hourly_average(long_names...; output_writer = HDF5Writer()) =
+hourly_averages(short_names...; output_writer = HDF5Writer()) =
     common_diagnostics(
         60 * 60,
         (+),
         output_writer,
-        long_names...;
+        short_names...;
         pre_output_hook! = average_pre_output_hook!,
     )
+"""
+    hourly_average(short_names...; output_writer = HDF5Writer())
+
+
+Return a `ScheduledDiagnostics` that computes the hourly average for the given variable.
+"""
+hourly_average(short_names; output_writer = HDF5Writer()) =
+    hourly_averages(short_names; output_writer)[1]
 
 # Include all the subdefaults
 include("defaults/moisture_model.jl")
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
index fbf01d3c05c..d4aed5a3f53 100644
--- a/src/diagnostics/defaults/moisture_model.jl
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -1,21 +1,19 @@
 # FIXME: Gabriele added this as an example. Put something meaningful here!
 function get_default_diagnostics(::T) where {T <: DryModel}
-    return vcat(
-        daily_average("air_density"),
-        [
-            ScheduledDiagnosticTime(
-                variable = ALL_DIAGNOSTICS["air_density"],
-                compute_every = :timestep,
-                output_every = 86400, # seconds
-                reduction_time_func = min,
-                output_writer = HDF5Writer(),
-            ),
-            ScheduledDiagnosticIterations(
-                variable = ALL_DIAGNOSTICS["air_density"],
-                compute_every = 1,
-                output_every = 1, # iteration
-                output_writer = HDF5Writer(),
-            ),
-        ],
-    )
+    return [
+        daily_averages("air_density")...,
+        ScheduledDiagnosticTime(
+            variable = ALL_DIAGNOSTICS["air_density"],
+            compute_every = :timestep,
+            output_every = 86400, # seconds
+            reduction_time_func = min,
+            output_writer = HDF5Writer(),
+        ),
+        ScheduledDiagnosticIterations(
+            variable = ALL_DIAGNOSTICS["air_density"],
+            compute_every = 1,
+            output_every = 1, # iteration
+            output_writer = HDF5Writer(),
+        ),
+    ]
 end

From dfe4508111f5b103fcbcc1011d4e8c71e0a327cb Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 11 Sep 2023 12:07:50 -0700
Subject: [PATCH 35/73] Change name to output_name in diagnostics

---
 docs/src/diagnostics.md       |  6 +++--
 src/diagnostics/diagnostic.jl | 45 ++++++++++++++++++-----------------
 src/diagnostics/writers.jl    |  2 +-
 src/solver/type_getters.jl    |  2 +-
 4 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 1a04a479d7e..8c95c889867 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -15,7 +15,7 @@ Second, you can specify the diagnostics you want to output directly in the
 ```
 diagnostics:
   - short_name: air_density
-    name: a_name
+    output_name: a_name
     period: 3hours
   - reduction_time: average
     short_name: air_density
@@ -23,7 +23,7 @@ diagnostics:
 ```
 This adds two diagnostics (both for `air_density`). The `period` keyword
 identifies the period over which to compute the reduction and how often to save
-to disk. `name` is optional, and if provided, it identifies the name of the
+to disk. `output_name` is optional, and if provided, it identifies the name of the
 output file.
 
 ### From a script
@@ -101,6 +101,8 @@ More specifically, a `ScheduledDiagnostic` contains the following pieces of data
    `pre_output_hook!` is called with two arguments: the value accumulated during
    the reduction, and the number of times the diagnostic was computed from the
    last time it was output.
+- `output_name`: A descriptive name that can be used by the `output_writer`. If
+   not provided, a default one is generated.
 
 To implement operations like the arithmetic average, the `reduction_time_func`
 has to be chosen as `+`, and a `pre_output_hook!` that renormalize `acc` by the
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 3286a2f4208..4887535e70b 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -199,7 +199,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     reduction_space_func::F2
     compute_every::T2
     pre_output_hook!::PO
-    name::String
+    output_name::String
 
     """
         ScheduledDiagnosticIterations(; variable::DiagnosticVariable,
@@ -209,7 +209,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                                         reduction_space_func = nothing,
                                         compute_every = isa_reduction ? 1 : output_every,
                                         pre_output_hook! = nothing,
-                                        name = descriptive_name(self) )
+                                        output_name = descriptive_name(self) )
 
 
     A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence
@@ -261,9 +261,9 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                           discarded. An example of `pre_output_hook!` to compute the arithmetic
                           average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
 
-   - `name`: A descriptive name for this particular diagnostic. If none is provided, one
-             will be generated mixing the short name of the variable, the reduction, and the
-             period of the reduction.
+   - `output_name`: A descriptive name for this particular diagnostic. If none is provided,
+                    one will be generated mixing the short name of the variable, the
+                    reduction, and the period of the reduction.
 
 
     """
@@ -275,7 +275,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         reduction_space_func = nothing,
         compute_every = isnothing(reduction_time_func) ? output_every : 1,
         pre_output_hook! = nothing,
-        name = get_descriptive_name(
+        output_name = get_descriptive_name(
             variable,
             output_every,
             reduction_time_func,
@@ -287,14 +287,14 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         # We provide an inner constructor to enforce some constraints
 
         output_every % compute_every == 0 || error(
-            "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(name)",
+            "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_name)",
         )
 
         isa_reduction = !isnothing(reduction_time_func)
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every ($output_every) != compute_every ($compute_every) for $(name), changing compute_every to match"
+            @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -320,7 +320,7 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
             reduction_space_func,
             compute_every,
             pre_output_hook!,
-            name,
+            output_name,
         )
     end
 end
@@ -334,7 +334,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
     reduction_space_func::F2
     compute_every::T2
     pre_output_hook!::PO
-    name::String
+    output_name::String
 
     """
         ScheduledDiagnosticTime(; variable::DiagnosticVariable,
@@ -344,7 +344,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                                   reduction_space_func = nothing,
                                   compute_every = isa_reduction ? :timestep : output_every,
                                   pre_output_hook! = nothing,
-                                  name = descriptive_name(self))
+                                  output_name = descriptive_name(self))
 
 
     A `DiagnosticVariable` that has to be computed and output during a simulation with a
@@ -400,10 +400,11 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                           discarded. An example of `pre_output_hook!` to compute the arithmetic
                           average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
 
-   - `name`: A descriptive name for this particular diagnostic. If none is provided, one
-             will be generated mixing the short name of the variable, the reduction, and the
-             period of the reduction.
+   - `output_name`: A descriptive name for this particular diagnostic. If none is provided,
+                    one will be generated mixing the short name of the variable, the
+                    reduction, and the period of the reduction.
     """
+
     function ScheduledDiagnosticTime(;
         variable::DiagnosticVariable,
         output_every,
@@ -413,7 +414,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         compute_every = isnothing(reduction_time_func) ? output_every :
                         :timestep,
         pre_output_hook! = nothing,
-        name = get_descriptive_name(
+        output_name = get_descriptive_name(
             variable,
             output_every,
             reduction_time_func,
@@ -428,7 +429,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         # the list of diagnostics
         if !isa(compute_every, Symbol)
             output_every % compute_every == 0 || error(
-                "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(name)",
+                "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_name)",
             )
         end
 
@@ -436,7 +437,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every ($output_every) != compute_every ($compute_every) for $(name), changing compute_every to match"
+            @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -462,7 +463,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
             reduction_space_func,
             compute_every,
             pre_output_hook!,
-            name,
+            output_name,
         )
     end
 end
@@ -490,10 +491,10 @@ function ScheduledDiagnosticIterations(
     output_every = sd_time.output_every / Δt
 
     isinteger(output_every) || error(
-        "output_every ($output_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.name)",
+        "output_every ($output_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
     )
     isinteger(compute_every) || error(
-        "compute_every ($compute_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.name)",
+        "compute_every ($compute_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
     )
 
     ScheduledDiagnosticIterations(;
@@ -504,7 +505,7 @@ function ScheduledDiagnosticIterations(
         sd_time.reduction_space_func,
         compute_every = convert(Int, compute_every),
         sd_time.pre_output_hook!,
-        sd_time.name,
+        sd_time.output_name,
     )
 end
 
@@ -536,7 +537,7 @@ function ScheduledDiagnosticTime(
         sd_time.reduction_space_func,
         compute_every,
         sd_time.pre_output_hook!,
-        sd_time.name,
+        sd_time.output_name,
     )
 end
 
diff --git a/src/diagnostics/writers.jl b/src/diagnostics/writers.jl
index 5aa04901145..6a1becbd9b2 100644
--- a/src/diagnostics/writers.jl
+++ b/src/diagnostics/writers.jl
@@ -39,7 +39,7 @@ function HDF5Writer()
 
         output_path = joinpath(
             integrator.p.simulation.output_dir,
-            "$(diagnostic.name)_$time.h5",
+            "$(diagnostic.output_name)_$time.h5",
         )
 
         hdfwriter = InputOutput.HDF5Writer(output_path, integrator.p.comms_ctx)
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index f31634fd01e..52fe7267dd2 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -680,7 +680,7 @@ function get_diagnostics(parsed_args, atmos_model)
                 reduction_time_func = reduction_time_func,
                 pre_output_hook! = pre_output_hook!,
                 output_writer = CAD.HDF5Writer(),
-                name = name,
+                output_name = name,
             ),
         )
     end

From bf20885a6738e4917c028efa67b7c95b41e5a3a5 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 13 Sep 2023 12:24:41 -0700
Subject: [PATCH 36/73] Add preliminary support to writing to NetCDF

---
 docs/src/diagnostics.md        |   6 ++
 src/diagnostics/Diagnostics.jl |   3 +-
 src/diagnostics/writers.jl     | 103 ++++++++++++++++++++++++++++++++-
 src/solver/type_getters.jl     |  19 +++++-
 4 files changed, 127 insertions(+), 4 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 8c95c889867..ab36d908244 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -17,15 +17,21 @@ diagnostics:
   - short_name: air_density
     output_name: a_name
     period: 3hours
+    writer: nc
   - reduction_time: average
     short_name: air_density
     period: 12hours
+    writer: h5
 ```
 This adds two diagnostics (both for `air_density`). The `period` keyword
 identifies the period over which to compute the reduction and how often to save
 to disk. `output_name` is optional, and if provided, it identifies the name of the
 output file.
 
+The default `writer` is HDF5. If `writer` is `nc` or `netcdf`, the output is
+remapped non-conservatively on a Cartesian grid and saved to a NetCDF file.
+Currently, only 3D fields on cubed spheres are supported.
+
 ### From a script
 
 The simplest way to get started with diagnostics is to use the defaults for your
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 36069cea02e..36918382a4f 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -1,7 +1,6 @@
 module Diagnostics
 
-import ClimaCore: InputOutput
-import ClimaCore: Fields
+import ClimaCore: Fields, Geometry, InputOutput, Meshes, Spaces
 
 import ..AtmosModel
 import ..call_every_n_steps
diff --git a/src/diagnostics/writers.jl b/src/diagnostics/writers.jl
index 6a1becbd9b2..ac54b62da96 100644
--- a/src/diagnostics/writers.jl
+++ b/src/diagnostics/writers.jl
@@ -3,6 +3,9 @@
 # This file defines function-generating functions for output_writers for diagnostics. The
 # writers come with opinionated defaults.
 
+import ClimaCore.Remapping: interpolate_array
+import NCDatasets
+
 """
     HDF5Writer()
 
@@ -10,7 +13,8 @@
 Save a `ScheduledDiagnostic` to a HDF5 file inside the `output_dir` of the simulation.
 
 
-TODO: This is a very barebone HDF5Writer.
+TODO: This is a very barebone HDF5Writer. Do not consider this implementation as the "final
+word".
 
 We need to implement the following features/options:
 - Toggle for write new files/append
@@ -58,3 +62,100 @@ function HDF5Writer()
         return nothing
     end
 end
+
+
+"""
+    NetCDFWriter()
+
+
+Save a `ScheduledDiagnostic` to a NetCDF file inside the `output_dir` of the simulation by
+performing a pointwise (non-conservative) remapping first. This writer does not work on
+distributed simulations.
+
+
+TODO: This is a very barebone NetCDFWriter. This writer only supports 3D fields on cubed
+spheres at the moment. Do not consider this implementation as the "final word".
+
+We need to implement the following features/options:
+- Toggle for write new files/append
+- Checks for existing files
+- Check for new subfolders that have to be created
+- More meaningful naming conventions (keeping in mind that we can have multiple variables
+  with different reductions)
+- All variables in one file/each variable in its own file
+- All timesteps in one file/each timestep in its own file
+- Writing the correct attributes
+- Overriding simulation.output_dir (e.g., if the path starts with /)
+- ...more features/options
+
+"""
+function NetCDFWriter(;
+    num_points_longitude = 100,
+    num_points_latitude = 100,
+    num_points_altitude = 100,
+    compression_level = 9,
+)
+    function write_to_netcdf(field, diagnostic, integrator)
+        var = diagnostic.variable
+        time = integrator.t
+
+        # TODO: Automatically figure out details on the remapping depending on the given
+        # field
+
+        # Let's support only cubed spheres for the moment
+        typeof(axes(field).horizontal_space.topology.mesh) <:
+        Meshes.AbstractCubedSphere ||
+            error("NetCDF writer supports only cubed sphere at the moment")
+
+        # diagnostic here is a ScheduledDiagnosticIteration. If we want to obtain a
+        # descriptive name (e.g., something with "daily"), we have to pass the timestep as
+        # well
+        output_path = joinpath(
+            integrator.p.simulation.output_dir,
+            "$(diagnostic.output_name)_$time.nc",
+        )
+
+        vert_domain = axes(field).vertical_topology.mesh.domain
+        z_min, z_max = vert_domain.coord_min.z, vert_domain.coord_max.z
+
+        FT = Spaces.undertype(axes(field))
+
+        longpts = range(
+            Geometry.LongPoint(-FT(180.0)),
+            Geometry.LongPoint(FT(180.0)),
+            length = num_points_longitude,
+        )
+        latpts = range(
+            Geometry.LatPoint(-FT(80.0)),
+            Geometry.LatPoint(FT(80.0)),
+            length = num_points_latitude,
+        )
+        zpts = range(
+            Geometry.ZPoint(FT(z_min)),
+            Geometry.ZPoint(FT(z_max)),
+            length = num_points_altitude,
+        )
+
+        nc = NCDatasets.Dataset(output_path, "c")
+
+        NCDatasets.defDim(nc, "lon", num_points_longitude)
+        NCDatasets.defDim(nc, "lat", num_points_latitude)
+        NCDatasets.defDim(nc, "z", num_points_altitude)
+
+        nc.attrib["long_name"] = var.long_name
+        nc.attrib["units"] = var.units
+        nc.attrib["comments"] = var.comments
+
+        v = NCDatasets.defVar(
+            nc,
+            "$(var.short_name)",
+            FT,
+            ("lon", "lat", "z"),
+            deflatelevel = compression_level,
+        )
+
+        v[:, :, :] = interpolate_array(field, longpts, latpts, zpts)
+
+        close(nc)
+    end
+end
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 52fe7267dd2..03ad26121c7 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -634,6 +634,15 @@ function get_diagnostics(parsed_args, atmos_model)
         "average" => ((+), CAD.average_pre_output_hook!),
     )
 
+    # The default writer is HDF5
+    ALLOWED_WRITERS = Dict(
+        "nothing" => CAD.HDF5Writer(),
+        "h5" => CAD.HDF5Writer(),
+        "hdf5" => CAD.HDF5Writer(),
+        "nc" => CAD.NetCDFWriter(),
+        "netcdf" => CAD.NetCDFWriter(),
+    )
+
     for yaml_diag in yaml_diagnostics
         # Return "nothing" if "reduction_time" is not in the YAML block
         #
@@ -649,6 +658,14 @@ function get_diagnostics(parsed_args, atmos_model)
                 ALLOWED_REDUCTIONS[reduction_time_yaml]
         end
 
+        writer_ext = lowercase(get(yaml_diag, "writer", "nothing"))
+
+        if !haskey(ALLOWED_WRITERS, writer_ext)
+            error("writer $writer_ext not implemented")
+        else
+            writer = ALLOWED_WRITERS[writer_ext]
+        end
+
         name = get(yaml_diag, "name", nothing)
 
         haskey(yaml_diag, "period") ||
@@ -679,7 +696,7 @@ function get_diagnostics(parsed_args, atmos_model)
                 compute_every = compute_every,
                 reduction_time_func = reduction_time_func,
                 pre_output_hook! = pre_output_hook!,
-                output_writer = CAD.HDF5Writer(),
+                output_writer = writer,
                 output_name = name,
             ),
         )

From ff662044f67156237c2c861c5f17a66177016601 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Tue, 12 Sep 2023 10:09:03 -0700
Subject: [PATCH 37/73] Fix argument to compute_relative_humidty

---
 src/diagnostics/core_diagnostics.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index e4cec1d41d9..792e22f46b9 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -41,7 +41,7 @@ compute_relative_humidity_from_integrator!(out, integrator) =
     compute_relative_humidity_from_integrator!(
         out,
         integrator,
-        integrator.p.atmos,
+        integrator.p.atmos.moisture_model,
     )
 
 # FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the

From 0cd590197c3dc408235aea0fdbcd1a3252ade630 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Tue, 12 Sep 2023 13:42:36 -0700
Subject: [PATCH 38/73] Pass state, cache, and time to diagnostics

---
 docs/src/diagnostics.md                 | 37 +++++++++++++++----------
 src/diagnostics/core_diagnostics.jl     | 32 ++++++++++++---------
 src/diagnostics/diagnostic.jl           | 26 ++++++++---------
 src/diagnostics/turbconv_diagnostics.jl | 22 ++++++---------
 src/solver/type_getters.jl              |  8 ++++--
 5 files changed, 68 insertions(+), 57 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index ab36d908244..4297f5dde00 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -208,17 +208,18 @@ are not used.
 ### Compute function
 
 The other piece of information needed to specify a `DiagnosticVariable` is a
-function `compute_from_integrator!`. Schematically, a `compute_from_integrator!` has to look like
+function `compute!`. Schematically, a `compute!` has to look like
 ```julia
-function compute_from_integrator!(out, integrator)
+function compute!(out, state, cache, time)
     # FIXME: Remove this line when ClimaCore implements the broadcasting to enable this
-    out .= # Calculcations with the state (= integrator.u) and the parameters (= integrator.p)
+    out .= # Calculcations with the state and the cache
 end
 ```
-Diagnostics are implemented as callbacks function which pass the `integrator`
-object (from `OrdinaryDiffEq`) to `compute_from_integrator!`.
 
-`compute_from_integrator!` also takes a second argument, `out`, which is used to
+Diagnostics are implemented as callbacks functions which pass the `state`,
+`cache`, and `time` from the integrator to `compute!`.
+
+`compute!` also takes a second argument, `out`, which is used to
 avoid extra memory allocations (which hurt performance). If `out` is `nothing`,
 and new area of memory is allocated. If `out` is a `ClimaCore.Field`, the
 operation is done in-place without additional memory allocations.
@@ -229,17 +230,21 @@ For instance, if you want to compute relative humidity, which does not make
 sense for dry simulations, you should define the functions
 
 ```julia
-function compute_relative_humidity_from_integrator!(
+function compute_relative_humidity!(
     out,
-    integrator,
+    state,
+    cache,
+    time,
     moisture_model::T,
 ) where {T}
     error("Cannot compute relative_humidity with moisture_model = $T")
 end
 
-function compute_relative_humidity_from_integrator!(
+function compute_relative_humidity!(
     out,
-    integrator,
+    state,
+    cache,
+    time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
@@ -248,17 +253,19 @@ function compute_relative_humidity_from_integrator!(
     return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
 end
 
-compute_relative_humidity_from_integrator!(out, integrator) =
-    compute_relative_humidity_from_integrator!(
+compute_relative_humidity!(out, integrator) =
+    compute_relative_humidity!(
         out,
-        integrator,
-        integrator.p.atmos,
+        state,
+        cache,
+        time,
+        cache.atmos,
     )
 ```
 
 This will return the correct relative humidity and throw informative errors when
 it cannot be computed. We could specialize
-`compute_relative_humidity_from_integrator` further if the relative humidity
+`compute_relative_humidity` further if the relative humidity
 were computed differently for `EquilMoistModel` and `NonEquilMoistModel`.
 
 ### Adding to the `ALL_DIAGNOSTICS` dictionary
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 792e22f46b9..1e31126f97b 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -9,39 +9,45 @@ add_diagnostic_variable!(
     long_name = "air_density",
     units = "kg m^-3",
     comments = "Density of air, a prognostic variable",
-    compute_from_integrator! = (out, integrator) -> begin
+    compute! = (out, state, cache, time) -> begin
         # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
         # We will want: out .= integrator.u.c.ρ
-        return copy(integrator.u.c.ρ)
+        return copy(state.c.ρ)
     end,
 )
 
 # Relative humidity
 
-function compute_relative_humidity_from_integrator!(
+function compute_relative_humidity!(
     out,
-    integrator,
+    state,
+    cache,
+    time,
     moisture_model::T,
 ) where {T}
     error("Cannot compute relative_humidity with moisture_model = $T")
 end
 
-function compute_relative_humidity_from_integrator!(
+function compute_relative_humidity!(
     out,
-    integrator,
+    state,
+    cache,
+    time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-    thermo_params = CAP.thermodynamics_params(integrator.p.params)
-    return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
+    thermo_params = CAP.thermodynamics_params(cache.params)
+    return TD.relative_humidity.(thermo_params, cache.ᶜts)
 end
 
-compute_relative_humidity_from_integrator!(out, integrator) =
-    compute_relative_humidity_from_integrator!(
+compute_relative_humidity!(out, state, cache, time) =
+    compute_relative_humidity!(
         out,
-        integrator,
-        integrator.p.atmos.moisture_model,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
     )
 
 # FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
@@ -51,5 +57,5 @@ add_diagnostic_variable!(
     long_name = "relative_humidity",
     units = "",
     comments = "Relative Humidity",
-    compute_from_integrator! = compute_relative_humidity_from_integrator!,
+    compute! = compute_relative_humidity!,
 )
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 4887535e70b..065492af313 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -77,7 +77,7 @@ Keyword arguments
 - `comments`: More verbose explanation of what the variable is, or comments related to how
               it is defined or computed.
 
-- `compute_from_integrator!`: Function that compute the diagnostic variable from the state.
+- `compute!`: Function that compute the diagnostic variable from the state.
                              It has to take two arguments: the `integrator`, and a
                              pre-allocated area of memory where to write the result of the
                              computation. It the no pre-allocated area is available, a new
@@ -90,7 +90,7 @@ Base.@kwdef struct DiagnosticVariable{T}
     long_name::String
     units::String
     comments::String
-    compute_from_integrator!::T
+    compute!::T
 end
 
 # ClimaAtmos diagnostics
@@ -103,7 +103,7 @@ const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
                                long_name,
                                units,
                                description,
-                               compute_from_integrator!)
+                               compute!)
 
 
 Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates the state of
@@ -127,7 +127,7 @@ Keyword arguments
 - `comments`: More verbose explanation of what the variable is, or comments related to how
               it is defined or computed.
 
-- `compute_from_integrator!`: Function that compute the diagnostic variable from the state.
+- `compute!`: Function that compute the diagnostic variable from the state.
                              It has to take two arguments: the `integrator`, and a
                              pre-allocated area of memory where to write the result of the
                              computation. It the no pre-allocated area is available, a new
@@ -141,18 +141,13 @@ function add_diagnostic_variable!(;
     long_name,
     units,
     comments,
-    compute_from_integrator!,
+    compute!,
 )
     haskey(ALL_DIAGNOSTICS, short_name) &&
         error("diagnostic $short_name already defined")
 
-    ALL_DIAGNOSTICS[short_name] = DiagnosticVariable(;
-        short_name,
-        long_name,
-        units,
-        comments,
-        compute_from_integrator!,
-    )
+    ALL_DIAGNOSTICS[short_name] =
+        DiagnosticVariable(; short_name, long_name, units, comments, compute!)
 end
 
 # Do you want to define more diagnostics? Add them here
@@ -628,7 +623,12 @@ function get_callbacks_from_diagnostics(
         compute_callback =
             integrator -> begin
                 # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
-                variable.compute_from_integrator!(storage[diag], integrator)
+                variable.compute!(
+                    storage[diag],
+                    integrator.u,
+                    integrator.p,
+                    integrator.t,
+                )
 
                 # accumulator[diag] is not defined for non-reductions
                 diag_accumulator = get(accumulators, diag, nothing)
diff --git a/src/diagnostics/turbconv_diagnostics.jl b/src/diagnostics/turbconv_diagnostics.jl
index 38e556d6ac4..a0084cdf048 100644
--- a/src/diagnostics/turbconv_diagnostics.jl
+++ b/src/diagnostics/turbconv_diagnostics.jl
@@ -4,28 +4,23 @@
 
 # This is an example of how to compute the same diagnostic variable differently depending on
 # the model. This is also exposed to the user, which could define their own
-# compute_tke_from_integrator.
+# compute_tke.
 
-function compute_tke_from_integrator!(out, integrator, ::EDMFX)
+function compute_tke!(out, state, cache, time, ::EDMFX)
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-    return copy(integrator.p.ᶜspecific⁰.tke)
+    return copy(cache.ᶜspecific⁰.tke)
 end
 
-function compute_tke_from_integrator!(out, integrator, ::DiagnosticEDMFX)
+function compute_tke!(out, state, cache, time, ::DiagnosticEDMFX)
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-    return copy(integrator.p.tke⁰)
+    return copy(cache.tke⁰)
 end
 
-compute_tke_from_integrator!(out, integrator) =
-    compute_tke_from_integrator!(out, integrator, integrator.p.atmos)
+compute_tke!(out, state, cache, time) = compute_tke!(out, state, cache.atmos)
 
-function compute_tke_from_integrator!(
-    out,
-    integrator,
-    turbconv_model::T,
-) where {T}
+function compute_tke!(out, state, cache, time, turbconv_model::T) where {T}
     error("Cannot compute tke with turbconv_model = $T")
 end
 
@@ -36,6 +31,5 @@ add_diagnostic_variable!(
     long_name = "turbolent_kinetic_energy",
     units = "J",
     comments = "Turbolent Kinetic Energy",
-    compute_from_integrator! = (integrator, out) ->
-        compute_tke_from_integrator!,
+    compute! = compute_tke!,
 )
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 03ad26121c7..2c19c6c41fa 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -901,8 +901,12 @@ function get_integrator(config::AtmosConfig)
         variable = diag.variable
         try
             # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-            diagnostic_storage[diag] =
-                variable.compute_from_integrator!(nothing, integrator)
+            diagnostic_storage[diag] = variable.compute!(
+                nothing,
+                integrator.u,
+                integrator.p,
+                integrator.t,
+            )
             diagnostic_counters[diag] = 1
             # If it is not a reduction, call the output writer as well
             if isnothing(diag.reduction_time_func)

From f91774bb7b60c85ca13670f4bfc7b55b2b8a8356 Mon Sep 17 00:00:00 2001
From: LenkaNovak <lenka@caltech.edu>
Date: Wed, 13 Sep 2023 15:48:07 -0700
Subject: [PATCH 39/73] add core diagnostics

---
 src/diagnostics/core_diagnostics.jl | 411 ++++++++++++++++++++++++++--
 1 file changed, 392 insertions(+), 19 deletions(-)

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 1e31126f97b..afdf675be9a 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -1,12 +1,23 @@
 # This file is included in Diagnostics.jl
 
-# Rho
+# general helper for undefined functions for a particular model (avoids some repetition)
+function compute_variable!(
+    variable::Val,
+    out,
+    state,
+    cache,
+    time,
+    model::T,
+) where {T}
+    error("Cannot compute $variable with model = $T")
+end
 
-# FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
-# physics. Please fix this!
+###
+# Rho (3d)
+###
 add_diagnostic_variable!(
     short_name = "air_density",
-    long_name = "air_density",
+    long_name = "Air Density",
     units = "kg m^-3",
     comments = "Density of air, a prognostic variable",
     compute! = (out, state, cache, time) -> begin
@@ -16,19 +27,160 @@ add_diagnostic_variable!(
     end,
 )
 
-# Relative humidity
+###
+# U velocity (3d)
+###
+add_diagnostic_variable!(
+    short_name = "eastward_wind",
+    long_name = "Eastward Wind",
+    units = "m s^-1",
+    comments = "Eastward (zonal) wind component, a prognostic variable",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        return copy(Geometry.UVector.(cache.ᶜu))
+    end,
+)
+
+###
+# V velocity (3d)
+###
+add_diagnostic_variable!(
+    short_name = "northward_wind",
+    long_name = "Northward Wind",
+    units = "m s^-1",
+    comments = "Northward (meridional) wind component, a prognostic variable",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        return copy(Geometry.VVector.(cache.ᶜu))
+    end,
+)
+
+###
+# W velocity (3d)
+###
+# TODO: may want to convert to omega (Lagrangian pressure tendency) as standard output,
+# but this is probably more useful for now
+add_diagnostic_variable!(
+    short_name = "vertical_wind",
+    long_name = "Vertical Wind",
+    units = "m s^-1",
+    comments = "Vertical wind component, a prognostic variable",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        return copy(Geometry.WVector.(cache.ᶜu))
+    end,
+)
+
+###
+# Temperature (3d)
+###
+add_diagnostic_variable!(
+    short_name = "air_temperature",
+    long_name = "Air Temperature",
+    units = "K",
+    comments = "Temperature of air",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        thermo_params = CAP.thermodynamics_params(cache.params)
+        return copy(TD.air_temperature.(thermo_params, cache.ᶜts))
+    end,
+)
+
+###
+# Potential temperature (3d)
+###
+add_diagnostic_variable!(
+    short_name = "air_potential_temperature",
+    long_name = "Air potential temperature",
+    units = "K",
+    comments = "Potential temperature of air",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        thermo_params = CAP.thermodynamics_params(cache.params)
+        return copy(TD.dry_pottemp.(thermo_params, cache.ᶜts))
+    end,
+)
+
+###
+# Air pressure (3d)
+###
+add_diagnostic_variable!(
+    short_name = "air_pressure",
+    long_name = "Air pressure",
+    units = "Pa",
+    comments = "Pressure of air",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        return copy(cache.ᶜp)
+    end,
+)
+
+###
+# Vorticity (3d)
+###
+add_diagnostic_variable!(
+    short_name = "atmosphere_relative_vorticity",
+    long_name = "Vertical component of relative vorticity",
+    units = "s^-1",
+    comments = "Vertical component of relative vorticity",
+    compute! = (out, state, cache, time) -> begin
+        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+        # We will want: out .= integrator.u.c.ρ
+        ᶜvort = @. Geometry.WVector(curlₕ(state.c.uₕ))
+        if cache.do_dss
+            Spaces.weighted_dss!(ᶜvort)
+        end
+        return copy(ᶜvort)
+    end,
+)
+
 
-function compute_relative_humidity!(
+###
+# Relative humidity (3d)
+###
+function compute_variable!(
+    Val{:relative_humidity},
     out,
     state,
     cache,
     time,
     moisture_model::T,
-) where {T}
-    error("Cannot compute relative_humidity with moisture_model = $T")
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    thermo_params = CAP.thermodynamics_params(cache.params)
+    return TD.relative_humidity.(thermo_params, cache.ᶜts)
 end
 
-function compute_relative_humidity!(
+compute_relative_humidity!(out, state, cache, time) =
+    compute_variable!(
+        Val(:relative_humidity),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
+    )
+
+add_diagnostic_variable!(
+    short_name = "Relative Humidity",
+    long_name = "relative_humidity",
+    units = "",
+    comments = "Total amount of water vapor in the air relative to the amount achievable by saturation at the current temperature",
+    compute! = compute_relative_humidity!,
+)
+
+###
+# Total specific humidity (3d)
+###
+function compute_variable!(
+    ::Val{:specific_humidity},
     out,
     state,
     cache,
@@ -38,11 +190,12 @@ function compute_relative_humidity!(
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
     thermo_params = CAP.thermodynamics_params(cache.params)
-    return TD.relative_humidity.(thermo_params, cache.ᶜts)
+    return TD.total_specific_humidity.(thermo_params, cache.ᶜts)
 end
 
-compute_relative_humidity!(out, state, cache, time) =
-    compute_relative_humidity!(
+compute_specific_humidity!(out, state, cache, time) =
+    compute_variable!(
+        Val(:specific_humidity),
         out,
         state,
         cache,
@@ -50,12 +203,232 @@ compute_relative_humidity!(out, state, cache, time) =
         cache.atmos.moisture_model,
     )
 
-# FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
-# physics. Please fix this!
 add_diagnostic_variable!(
-    short_name = "relative_humidity",
-    long_name = "relative_humidity",
-    units = "",
-    comments = "Relative Humidity",
-    compute! = compute_relative_humidity!,
+    short_name = "Specific_humidity",
+    long_name = "Specific Humidity",
+    units = "kg kg^-1",
+    comments = "Mass of all water phases per mass of air, a prognostic variable",
+    compute! = compute_specific_humidity!,
+)
+
+###
+# Surface specific humidity (2d) - although this could be collapsed into the 3d version + reduction, it probably makes sense to keep it separate, since it's really the property of the surface
+###
+function compute_variable!(
+    ::Val{:surface_specific_humidity},
+    out,
+    state,
+    cache,
+    time,
+    moisture_model::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    thermo_params = CAP.thermodynamics_params(cache.params)
+    return TD.total_specific_humidity.(thermo_params, cache.sfc_conditions.ts)
+end
+
+compute_surface_specific_humidity!(out, state, cache, time) =
+    compute_variable!(
+        Val(:surface_specific_humidity),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
+    )
+
+add_diagnostic_variable!(
+    short_name = "surface_specific_humidity",
+    long_name = "Surface Specific Humidity",
+    units = "kg kg^-1",
+    comments = "Mass of all water phases per mass of air in the near-surface layer",
+    compute! = compute_surface_specific_humidity!,
+)
+
+###
+# Surface temperature (2d)
+###
+add_diagnostic_variable!(
+    short_name = "surface_temperature",
+    long_name = "Surface Temperature",
+    units = "K",
+    comments = "Temperature of the surface",
+    compute! = TD.air_temperature.(thermo_params, cache.sfc_conditions.ts),
+)
+
+###
+# Eastward surface drag component (2d)
+###
+function drag_vector!(state, cache)
+    sfc_local_geometry =
+    Fields.level(Fields.local_geometry_field(state.f), Fields.half)
+    surface_ct3_unit =
+        CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    (; ρ_flux_uₕ) = cache.sfc_conditions
+    sfc_flux_momentum =
+        Geometry.UVVector.(
+            adjoint.(ρ_flux_uₕ ./ Spaces.level(ᶠinterp.(state.c.ρ), half)) .*
+            surface_ct3_unit
+    )
+end
+
+function compute_variable!(
+    ::Val{:eastward_drag},
+    out,
+    state,
+    cache,
+    time,
+    model::T,
+) where {T <: TotalEnergy}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+
+    return drag_vector!(state, cache).components.data.:1
+end
+
+compute_eastward_drag!(out, state, cache, time) =
+    compute_variable!(
+        Val(:eastward_drag),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.energy_form,
+    )
+
+add_diagnostic_variable!(
+    short_name = "eastward_drag",
+    long_name = "Eastward component of the surface drag",
+    units = "kg m^-2 s^-2",
+    comments = "Eastward component of the surface drag",
+    compute! = compute_eastward_drag!,
+)
+
+###
+# Northward surface drag component (2d)
+###
+function compute_variable!(
+    ::Val{:northward_drag},
+    out,
+    state,
+    cache,
+    time,
+    model::T,
+) where {T <: TotalEnergy}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+
+    return drag_vector!(state, cache).components.data.:2
+end
+
+compute_northward_drag!(out, state, cache, time) =
+    compute_variable!(
+        Val(:northward_drag),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.energy_form,
+    )
+
+add_diagnostic_variable!(
+    short_name = "northward_drag",
+    long_name = "Northward component of the surface drag",
+    units = "kg m^-2 s^-2",
+    comments = "Northward component of the surface drag",
+    compute! = compute_northward_drag!,
+)
+
+###
+# Surface energy flux (2d) - TODO: this may need to be split into sensible and latent heat fluxes
+###
+function compute_variable!(
+    ::Val{:surface_energy_flux},
+    out,
+    state,
+    cache,
+    time,
+    model::T,
+) where {T <: TotalEnergy}
+    (;ρ_flux_h_tot) = cache.sfc_conditions
+    sfc_local_geometry =
+        Fields.level(Fields.local_geometry_field(state.f), Fields.half)
+    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    return dot.(ρ_flux_h_tot, surface_ct3_unit)
+end
+
+compute_surface_energy_flux!(out, state, cache, time) =
+    compute_variable!(
+        Val(:surface_energy_flux),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.energy_form,
+    )
+
+add_diagnostic_variable!(
+    short_name = "surface_energy_flux",
+    long_name = "Surface energy flux",
+    units = "W m^-2",
+    comments =  "energy flux at the surface",
+    compute! = compute_surface_energy_flux!,
 )
+
+###
+# Surface evaporation (2d)
+###
+function compute_variable!(
+    ::Val{:surface_evaporation},
+    out,
+    state,
+    cache,
+    time,
+    model::T,
+) where {T <: TotalEnergy}
+
+    compute_variable!(
+        Val(:surface_evaporation),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
+    )
+end
+
+function compute_variable!(
+    ::Val{:surface_evaporation},
+    out,
+    state,
+    cache,
+    time,
+    model::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    (;ρ_flux_q_tot) = cache.sfc_conditions
+    sfc_local_geometry =
+        Fields.level(Fields.local_geometry_field(state.f), Fields.half)
+    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+    return dot.(ρ_flux_q_tot, surface_ct3_unit)
+end
+
+compute_surface_evaporation!(out, state, cache, time) =
+    compute_variable!(
+        Val(:surface_evaporation),
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.energy_form,
+    )
+
+add_diagnostic_variable!(
+    short_name = "surface_evaporation",
+    long_name = "Surface evaporation",
+    units = "kg s^-1 m^-2",
+    comments =  "evaporation at the surface",
+    compute! = compute_surface_evaporation!,
+)
+
+# as required, all 3d variables will be sliced to X? (TBD) levels

From ae7a2c6ef85ac0a66e4fba3ce6f5fa584c106065 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 13 Sep 2023 17:34:21 -0700
Subject: [PATCH 40/73] Fix diagnostics, and store them correctly

---
 docs/src/diagnostics.md             |  28 ++-
 src/diagnostics/Diagnostics.jl      |  12 +-
 src/diagnostics/core_diagnostics.jl | 292 ++++++++++++++--------------
 src/diagnostics/diagnostic.jl       |   2 +-
 4 files changed, 187 insertions(+), 147 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 4297f5dde00..42074c0588c 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -212,7 +212,7 @@ function `compute!`. Schematically, a `compute!` has to look like
 ```julia
 function compute!(out, state, cache, time)
     # FIXME: Remove this line when ClimaCore implements the broadcasting to enable this
-    out .= # Calculcations with the state and the cache
+    out .= # Calculations with the state and the cache
 end
 ```
 
@@ -253,13 +253,13 @@ function compute_relative_humidity!(
     return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
 end
 
-compute_relative_humidity!(out, integrator) =
+compute_relative_humidity!(out, state, cache, time) =
     compute_relative_humidity!(
         out,
         state,
         cache,
         time,
-        cache.atmos,
+        cache.atmos.moisture_model,
     )
 ```
 
@@ -268,6 +268,28 @@ it cannot be computed. We could specialize
 `compute_relative_humidity` further if the relative humidity
 were computed differently for `EquilMoistModel` and `NonEquilMoistModel`.
 
+In `ClimaAtmos`, we define some helper functions to produce error messages, so
+the above code can be written as
+```julia
+compute_relative_humidity!(out, state, cache, time) =
+    compute_relative_humidity!(out, state, cache, time, cache.atmos.moisture_model)
+compute_relative_humidity!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("relative_humidity", model)
+
+function compute_relative_humidity!(
+    out,
+    state,
+    cache,
+    time,
+    ::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    thermo_params = CAP.thermodynamics_params(cache.params)
+    return TD.relative_humidity.(thermo_params, cache.ᶜts)
+end
+```
+
 ### Adding to the `ALL_DIAGNOSTICS` dictionary
 
 `ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable` in the
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 36918382a4f..88341c05131 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -1,15 +1,25 @@
 module Diagnostics
 
-import ClimaCore: Fields, Geometry, InputOutput, Meshes, Spaces
+import LinearAlgebra: dot
+
+import ClimaCore: Fields, Geometry, InputOutput, Meshes, Spaces, Operators
+import ClimaCore.Utilities: half
+import Thermodynamics as TD
 
 import ..AtmosModel
 import ..call_every_n_steps
+import ..Parameters as CAP
+
+import ..unit_basis_vector_data
 
 # moisture_model
 import ..DryModel
 import ..EquilMoistModel
 import ..NonEquilMoistModel
 
+# energy_form
+import ..TotalEnergy
+
 # turbconv_model
 import ..EDMFX
 import ..DiagnosticEDMFX
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index afdf675be9a..f67c3e667d2 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -1,16 +1,50 @@
 # This file is included in Diagnostics.jl
-
-# general helper for undefined functions for a particular model (avoids some repetition)
-function compute_variable!(
-    variable::Val,
-    out,
-    state,
-    cache,
-    time,
-    model::T,
-) where {T}
-    error("Cannot compute $variable with model = $T")
-end
+#
+# README: Adding a new core diagnostic:
+#
+# In addition to the metadata (names, comments, ...), the most important step in adding a
+# new DiagnosticVariable is defining its compute! function. `compute!` has to take four
+# arguments: (out, state, cache, time), and as to write the diagnostic in place into the
+# `out` variable.
+#
+# Often, it is possible to compute certain diagnostics only for specific models (e.g.,
+# humidity for moist models). For that, it is convenient to adopt the following pattern:
+#
+# 1. Define a catch base function that does the computation we want to do for the case we know
+# how to handle, for example
+#
+# function compute_relative_humidity!(
+#     out,
+#     state,
+#     cache,
+#     time,
+#     ::T,
+# ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+#     thermo_params = CAP.thermodynamics_params(cache.params)
+#     out .= TD.relative_humidity.(thermo_params, cache.ᶜts)
+# end
+#
+# 2. Define a function that has the correct signature and calls this function
+#
+# compute_relative_humidity!(out, state, cache, time) =
+#     compute_relative_humidity!(out, state, cache, time, cache.atmos.moisture_model)
+#
+# 3. Define a function that returns an error when the model is incorrect
+#
+# compute_relative_humidity!(_, _, _, _, model::T) where {T} =
+#     error_diagnostic_variable("relative_humidity", model)
+#
+# We can also output a specific error message
+#
+# compute_relative_humidity!(_, _, _, _, model::T) where {T} =
+#     error_diagnostic_variable("relative humidity makes sense only for moist models")
+
+# General helper functions for undefined diagnostics for a particular model
+error_diagnostic_variable(
+    message = "Cannot compute $variable with model = $T",
+) = error(message)
+error_diagnostic_variable(variable, model::T) where {T} =
+    error_diagnostic_variable("Cannot compute $variable with model = $T")
 
 ###
 # Rho (3d)
@@ -30,6 +64,9 @@ add_diagnostic_variable!(
 ###
 # U velocity (3d)
 ###
+
+# TODO: This velocity might not be defined (e.g., in a column model). Add dispatch to catch
+# that.
 add_diagnostic_variable!(
     short_name = "eastward_wind",
     long_name = "Eastward Wind",
@@ -45,6 +82,9 @@ add_diagnostic_variable!(
 ###
 # V velocity (3d)
 ###
+
+# TODO: This velocity might not be defined (e.g., in a column model). Add dispatch to catch
+# that.
 add_diagnostic_variable!(
     short_name = "northward_wind",
     long_name = "Northward Wind",
@@ -62,6 +102,7 @@ add_diagnostic_variable!(
 ###
 # TODO: may want to convert to omega (Lagrangian pressure tendency) as standard output,
 # but this is probably more useful for now
+#
 add_diagnostic_variable!(
     short_name = "vertical_wind",
     long_name = "Vertical Wind",
@@ -86,7 +127,7 @@ add_diagnostic_variable!(
         # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
         # We will want: out .= integrator.u.c.ρ
         thermo_params = CAP.thermodynamics_params(cache.params)
-        return copy(TD.air_temperature.(thermo_params, cache.ᶜts))
+        return TD.air_temperature.(thermo_params, cache.ᶜts)
     end,
 )
 
@@ -102,7 +143,7 @@ add_diagnostic_variable!(
         # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
         # We will want: out .= integrator.u.c.ρ
         thermo_params = CAP.thermodynamics_params(cache.params)
-        return copy(TD.dry_pottemp.(thermo_params, cache.ᶜts))
+        return TD.dry_pottemp.(thermo_params, cache.ᶜts)
     end,
 )
 
@@ -144,8 +185,18 @@ add_diagnostic_variable!(
 ###
 # Relative humidity (3d)
 ###
-function compute_variable!(
-    Val{:relative_humidity},
+compute_relative_humidity!(out, state, cache, time) =
+    compute_relative_humidity!(
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
+    )
+compute_relative_humidity!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("relative_humidity", model)
+
+function compute_relative_humidity!(
     out,
     state,
     cache,
@@ -158,19 +209,9 @@ function compute_variable!(
     return TD.relative_humidity.(thermo_params, cache.ᶜts)
 end
 
-compute_relative_humidity!(out, state, cache, time) =
-    compute_variable!(
-        Val(:relative_humidity),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-    )
-
 add_diagnostic_variable!(
-    short_name = "Relative Humidity",
-    long_name = "relative_humidity",
+    short_name = "relative_humidity",
+    long_name = "Relative Humidity",
     units = "",
     comments = "Total amount of water vapor in the air relative to the amount achievable by saturation at the current temperature",
     compute! = compute_relative_humidity!,
@@ -179,8 +220,18 @@ add_diagnostic_variable!(
 ###
 # Total specific humidity (3d)
 ###
-function compute_variable!(
-    ::Val{:specific_humidity},
+compute_specific_humidity!(out, state, cache, time) =
+    compute_specific_humidity!(
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
+    )
+compute_specific_humidity!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("specific_humidity", model)
+
+function compute_specific_humidity!(
     out,
     state,
     cache,
@@ -193,18 +244,8 @@ function compute_variable!(
     return TD.total_specific_humidity.(thermo_params, cache.ᶜts)
 end
 
-compute_specific_humidity!(out, state, cache, time) =
-    compute_variable!(
-        Val(:specific_humidity),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-    )
-
 add_diagnostic_variable!(
-    short_name = "Specific_humidity",
+    short_name = "specific_humidity",
     long_name = "Specific Humidity",
     units = "kg kg^-1",
     comments = "Mass of all water phases per mass of air, a prognostic variable",
@@ -212,10 +253,20 @@ add_diagnostic_variable!(
 )
 
 ###
-# Surface specific humidity (2d) - although this could be collapsed into the 3d version + reduction, it probably makes sense to keep it separate, since it's really the property of the surface
+# Surface specific humidity (2d)
 ###
-function compute_variable!(
-    ::Val{:surface_specific_humidity},
+compute_surface_specific_humidity!(out, state, cache, time) =
+    compute_surface_specific_humidity!(
+        out,
+        state,
+        cache,
+        time,
+        cache.atmos.moisture_model,
+    )
+compute_surface_specific_humidity!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("surface_specific_humidity", model)
+
+function compute_surface_specific_humidity!(
     out,
     state,
     cache,
@@ -228,16 +279,6 @@ function compute_variable!(
     return TD.total_specific_humidity.(thermo_params, cache.sfc_conditions.ts)
 end
 
-compute_surface_specific_humidity!(out, state, cache, time) =
-    compute_variable!(
-        Val(:surface_specific_humidity),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-    )
-
 add_diagnostic_variable!(
     short_name = "surface_specific_humidity",
     long_name = "Surface Specific Humidity",
@@ -249,54 +290,52 @@ add_diagnostic_variable!(
 ###
 # Surface temperature (2d)
 ###
+function compute_surface_temperature!(out, state, cache, time)
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    thermo_params = CAP.thermodynamics_params(cache.params)
+    return TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
+end
+
 add_diagnostic_variable!(
     short_name = "surface_temperature",
     long_name = "Surface Temperature",
     units = "K",
     comments = "Temperature of the surface",
-    compute! = TD.air_temperature.(thermo_params, cache.sfc_conditions.ts),
+    compute! = compute_surface_temperature!,
 )
 
 ###
 # Eastward surface drag component (2d)
 ###
-function drag_vector!(state, cache)
+compute_eastward_drag!(out, state, cache, time) =
+    compute_eastward_drag!(out, state, cache, time, cache.atmos.energy_form)
+compute_eastward_drag!(_, _, _, _, energy_form::T) where {T} =
+    error_diagnostic_variable("eastward_drag", energy_form)
+
+function drag_vector(state, cache)
     sfc_local_geometry =
-    Fields.level(Fields.local_geometry_field(state.f), Fields.half)
-    surface_ct3_unit =
-        CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
+        Fields.level(Fields.local_geometry_field(state.f), Fields.half)
+    surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
     (; ρ_flux_uₕ) = cache.sfc_conditions
-    sfc_flux_momentum =
-        Geometry.UVVector.(
-            adjoint.(ρ_flux_uₕ ./ Spaces.level(ᶠinterp.(state.c.ρ), half)) .*
-            surface_ct3_unit
+    return Geometry.UVVector.(
+        adjoint.(ρ_flux_uₕ ./ Spaces.level(ᶠinterp.(state.c.ρ), half)) .*
+        surface_ct3_unit
     )
 end
 
-function compute_variable!(
-    ::Val{:eastward_drag},
+function compute_eastward_drag!(
     out,
     state,
     cache,
     time,
-    model::T,
+    energy_form::T,
 ) where {T <: TotalEnergy}
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-
-    return drag_vector!(state, cache).components.data.:1
+    return drag_vector(state, cache).components.data.:1
 end
 
-compute_eastward_drag!(out, state, cache, time) =
-    compute_variable!(
-        Val(:eastward_drag),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.energy_form,
-    )
-
 add_diagnostic_variable!(
     short_name = "eastward_drag",
     long_name = "Eastward component of the surface drag",
@@ -308,30 +347,23 @@ add_diagnostic_variable!(
 ###
 # Northward surface drag component (2d)
 ###
-function compute_variable!(
-    ::Val{:northward_drag},
+compute_northward_drag!(out, state, cache, time) =
+    compute_northward_drag!(out, state, cache, time, cache.atmos.energy_form)
+compute_northward_drag!(_, _, _, _, energy_form::T) where {T} =
+    error_diagnostic_variable("northward_drag", energy_form)
+
+function compute_northward_drag!(
     out,
     state,
     cache,
     time,
-    model::T,
+    energy_form::T,
 ) where {T <: TotalEnergy}
     # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
     # We will want: out .= integrator.u.c.ρ
-
-    return drag_vector!(state, cache).components.data.:2
+    return drag_vector(state, cache).components.data.:2
 end
 
-compute_northward_drag!(out, state, cache, time) =
-    compute_variable!(
-        Val(:northward_drag),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.energy_form,
-    )
-
 add_diagnostic_variable!(
     short_name = "northward_drag",
     long_name = "Northward component of the surface drag",
@@ -343,92 +375,68 @@ add_diagnostic_variable!(
 ###
 # Surface energy flux (2d) - TODO: this may need to be split into sensible and latent heat fluxes
 ###
-function compute_variable!(
-    ::Val{:surface_energy_flux},
-    out,
-    state,
-    cache,
-    time,
-    model::T,
-) where {T <: TotalEnergy}
-    (;ρ_flux_h_tot) = cache.sfc_conditions
+function compute_surface_energy_flux!(out, state, cache, time)
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    (; ρ_flux_h_tot) = cache.sfc_conditions
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
     surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
     return dot.(ρ_flux_h_tot, surface_ct3_unit)
 end
 
-compute_surface_energy_flux!(out, state, cache, time) =
-    compute_variable!(
-        Val(:surface_energy_flux),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.energy_form,
-    )
-
 add_diagnostic_variable!(
     short_name = "surface_energy_flux",
     long_name = "Surface energy flux",
     units = "W m^-2",
-    comments =  "energy flux at the surface",
+    comments = "Energy flux at the surface",
     compute! = compute_surface_energy_flux!,
 )
 
 ###
 # Surface evaporation (2d)
 ###
-function compute_variable!(
-    ::Val{:surface_evaporation},
-    out,
-    state,
-    cache,
-    time,
-    model::T,
-) where {T <: TotalEnergy}
-
-    compute_variable!(
-        Val(:surface_evaporation),
+compute_surface_evaporation!(out, state, cache, time) =
+    compute_surface_evaporation!(
         out,
         state,
         cache,
         time,
         cache.atmos.moisture_model,
+        cache.atmos.energy_form,
     )
-end
+compute_surface_evaporation!(
+    _,
+    _,
+    _,
+    _,
+    moisture_model::T1,
+    energy_form::T2,
+) where {T1, T2} = error_diagnostic_variable(
+    "Can only compute surface_evaporation with energy_form = TotalEnergy() and with a moist model",
+)
 
-function compute_variable!(
-    ::Val{:surface_evaporation},
+function compute_surface_evaporation!(
     out,
     state,
     cache,
     time,
-    model::T,
+    moisture_model::T,
+    energy_form::TotalEnergy,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    (;ρ_flux_q_tot) = cache.sfc_conditions
+    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+    # We will want: out .= integrator.u.c.ρ
+    (; ρ_flux_q_tot) = cache.sfc_conditions
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
     surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
     return dot.(ρ_flux_q_tot, surface_ct3_unit)
 end
 
-compute_surface_evaporation!(out, state, cache, time) =
-    compute_variable!(
-        Val(:surface_evaporation),
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.energy_form,
-    )
-
 add_diagnostic_variable!(
     short_name = "surface_evaporation",
     long_name = "Surface evaporation",
     units = "kg s^-1 m^-2",
-    comments =  "evaporation at the surface",
+    comments = "evaporation at the surface",
     compute! = compute_surface_evaporation!,
 )
-
-# as required, all 3d variables will be sliced to X? (TBD) levels
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 065492af313..8aa258868cb 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -623,7 +623,7 @@ function get_callbacks_from_diagnostics(
         compute_callback =
             integrator -> begin
                 # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
-                variable.compute!(
+                storage[diag] .= variable.compute!(
                     storage[diag],
                     integrator.u,
                     integrator.p,

From 0864fefc42ad723c564b0db36651ce917fde3e03 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 08:58:10 -0700
Subject: [PATCH 41/73] Remove newlines between docstrings and functions

---
 src/diagnostics/default_diagnostics.jl | 4 ++--
 src/diagnostics/diagnostic.jl          | 2 --
 src/diagnostics/diagnostics_utils.jl   | 3 +--
 src/diagnostics/writers.jl             | 1 -
 4 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index dcecc48e600..7800990e375 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -154,14 +154,13 @@ Return a `ScheduledDiagnostics` that computes the hourly min for the given varia
 hourly_min(short_names; output_writer = HDF5Writer()) =
     hourly_mins(short_names; output_writer)[1]
 
+# An average is just a sum with a normalization before output
 """
     hourly_averages(short_names...; output_writer = HDF5Writer())
 
 
 Return a list of `ScheduledDiagnostics` that compute the hourly average for the given variables.
 """
-
-# An average is just a sum with a normalization before output
 hourly_averages(short_names...; output_writer = HDF5Writer()) =
     common_diagnostics(
         60 * 60,
@@ -170,6 +169,7 @@ hourly_averages(short_names...; output_writer = HDF5Writer()) =
         short_names...;
         pre_output_hook! = average_pre_output_hook!,
     )
+
 """
     hourly_average(short_names...; output_writer = HDF5Writer())
 
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 8aa258868cb..629d5b10673 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -399,7 +399,6 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                     one will be generated mixing the short name of the variable, the
                     reduction, and the period of the reduction.
     """
-
     function ScheduledDiagnosticTime(;
         variable::DiagnosticVariable,
         output_every,
@@ -472,7 +471,6 @@ Create a `ScheduledDiagnosticIterations` given a `ScheduledDiagnosticTime` and a
 timestep.
 
 """
-
 function ScheduledDiagnosticIterations(
     sd_time::ScheduledDiagnosticTime,
     Δt::T,
diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl
index 1863f802678..a403ecbbebf 100644
--- a/src/diagnostics/diagnostics_utils.jl
+++ b/src/diagnostics/diagnostics_utils.jl
@@ -19,8 +19,7 @@ is interpreted as in units of number of iterations.
 
 This function is useful for filenames and error messages.
 
- """
-
+"""
 function get_descriptive_name(
     variable::DiagnosticVariable,
     output_every,
diff --git a/src/diagnostics/writers.jl b/src/diagnostics/writers.jl
index ac54b62da96..3390f8f82a8 100644
--- a/src/diagnostics/writers.jl
+++ b/src/diagnostics/writers.jl
@@ -29,7 +29,6 @@ We need to implement the following features/options:
 - ...more features/options
 
 """
-
 function HDF5Writer()
     # output_drivers are called with the three arguments: the value, the ScheduledDiagnostic,
     # and the integrator

From 1e6b01e00df1932afaecf03b0c751703c8ac7d87 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 13 Sep 2023 17:55:02 -0700
Subject: [PATCH 42/73] Avoid unnecessary specialization

---
 src/diagnostics/defaults/moisture_model.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
index d4aed5a3f53..f1083f0c77b 100644
--- a/src/diagnostics/defaults/moisture_model.jl
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -1,5 +1,5 @@
 # FIXME: Gabriele added this as an example. Put something meaningful here!
-function get_default_diagnostics(::T) where {T <: DryModel}
+function get_default_diagnostics(::DryModel)
     return [
         daily_averages("air_density")...,
         ScheduledDiagnosticTime(

From ba68487b426d6c7d6540427cda2f099a63feb96e Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 08:26:21 -0700
Subject: [PATCH 43/73] Move ClimaCore patch outside of abbreviations.jl

We want to use `abbreviations.jl` in submodules too. Therefore, we have
to avoid any definition like the one had before this commit (otherwise,
we get a redefinition).
---
 src/solver/types.jl        | 8 ++++++++
 src/utils/abbreviations.jl | 8 --------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/solver/types.jl b/src/solver/types.jl
index c0137a7bf43..598cefa675b 100644
--- a/src/solver/types.jl
+++ b/src/solver/types.jl
@@ -466,3 +466,11 @@ function AtmosConfig(;
     return AtmosConfig{FT, TD, PA, C}(toml_dict, config, comms_ctx)
 end
 Base.eltype(::AtmosConfig{FT}) where {FT} = FT
+
+# In order to specify C2F operator boundary conditions with 0 instead of FT(0),
+# we need to tell ClimaCore how to convert AxisTensor components from Int to FT.
+# TODO: Move this monkey patch to ClimaCore in the next release.
+using ClimaCore.Geometry: AxisTensor, components
+AxisTensor{T′, N, A, S′}(a::AxisTensor{T, N, A, S}) where {T, N, A, S, T′, S′} =
+    AxisTensor(axes(a), S′(components(a)))
+Base.convert(::Type{AT}, a::AxisTensor) where {AT <: AxisTensor} = AT(a)
diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl
index a374cbd3960..9db0375ccea 100644
--- a/src/utils/abbreviations.jl
+++ b/src/utils/abbreviations.jl
@@ -75,11 +75,3 @@ const ᶜdivᵥ_stencil = Operators.Operator2Stencil(ᶜdivᵥ)
 const ᶜadvdivᵥ_stencil = Operators.Operator2Stencil(ᶜadvdivᵥ)
 const ᶠinterp_stencil = Operators.Operator2Stencil(ᶠinterp)
 const ᶠgradᵥ_stencil = Operators.Operator2Stencil(ᶠgradᵥ)
-
-# In order to specify C2F operator boundary conditions with 0 instead of FT(0),
-# we need to tell ClimaCore how to convert AxisTensor components from Int to FT.
-# TODO: Move this monkey patch to ClimaCore in the next release.
-using ClimaCore.Geometry: AxisTensor, components
-AxisTensor{T′, N, A, S′}(a::AxisTensor{T, N, A, S}) where {T, N, A, S, T′, S′} =
-    AxisTensor(axes(a), S′(components(a)))
-Base.convert(::Type{AT}, a::AxisTensor) where {AT <: AxisTensor} = AT(a)

From 2e7cc69551cffb071d2c45198d82835c82bbdf53 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 08:30:25 -0700
Subject: [PATCH 44/73] Use abbreviations.jl instead of definining our own

---
 src/diagnostics/Diagnostics.jl | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 88341c05131..9646f654cc2 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -24,16 +24,8 @@ import ..TotalEnergy
 import ..EDMFX
 import ..DiagnosticEDMFX
 
-# Abbreviations (following utils/abbreviations.jl)
-const curlₕ = Operators.Curl()
-const CT3 = Geometry.Contravariant3Vector
-const ᶜinterp = Operators.InterpolateF2C()
-# TODO: Implement proper extrapolation instead of simply reusing the first
-# interior value at the surface.
-const ᶠinterp = Operators.InterpolateC2F(
-    bottom = Operators.Extrapolate(),
-    top = Operators.Extrapolate(),
-)
+# We need the abbreviations for symbols like curl, grad, and so on
+include(joinpath("..", "utils", "abbreviations.jl"))
 
 include("diagnostic.jl")
 include("writers.jl")

From f8cf05f507784ba190f5582b039b612e445a58f0 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 08:38:38 -0700
Subject: [PATCH 45/73] Print error values in seconds instead of iteration

Prior to this commit, the error message printed out the computed
output/computed iterations as opposed to seconds, which made the error
message confusing
---
 src/diagnostics/diagnostic.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 629d5b10673..9f375c05273 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -484,10 +484,10 @@ function ScheduledDiagnosticIterations(
     output_every = sd_time.output_every / Δt
 
     isinteger(output_every) || error(
-        "output_every ($output_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
+        "output_every ($(sd_time.output_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
     )
     isinteger(compute_every) || error(
-        "compute_every ($compute_every) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
+        "compute_every ($(sd_time.compute_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
     )
 
     ScheduledDiagnosticIterations(;

From e7507763953c1815f13efa3b2763b04d1e5d4e4c Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 08:53:26 -0700
Subject: [PATCH 46/73] Make ALL_DIAGNOSTICS an implementation detail

Currently, the diagnostic variables are stored in a dictionary. We might
want to be able to change the internal representation in the future. For
that, we provide an interface to the database of known diagnostic
variables through `set_diagnostic_variable!` and
`get_diagnostic_variable`. If we want to change the internal
representation in the future, we will be able to do so easily.
---
 docs/src/diagnostics.md                    | 29 ++++++++++------------
 src/diagnostics/default_diagnostics.jl     |  2 +-
 src/diagnostics/defaults/moisture_model.jl |  4 +--
 src/diagnostics/diagnostic.jl              | 17 ++++++++++++-
 src/solver/type_getters.jl                 |  4 +--
 5 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 42074c0588c..8fc7b4a50d0 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -51,11 +51,11 @@ Technically, the diagnostics are represented as `ScheduledDiagnostic` objects,
 which contain information about what variable has to be computed, how often,
 where to save it, and so on (read below for more information on this). You can
 construct your own lists of `ScheduledDiagnostic`s starting from the variables
-defined by `ClimaAtmos`. The diagnostics that `ClimaAtmos` knows how to compute
-are collected in a global dictionary called `ALL_DIAGNOSTICS`. The variables in
-`ALL_DIAGNOSTICS` are identified with by the short and unique name, so that you
-can access them directly. One way to do so is by using the provided convenience
-functions for common operations, e.g., continuing the previous example
+defined by `ClimaAtmos`. The `DiagnosticVariable`s in `ClimaAtmos` are
+identified with by the short and unique name, so that you can access them
+directly with the function `get_diagnostic_variable`. One way to do so is by
+using the provided convenience functions for common operations, e.g., continuing
+the previous example
 
 ```julia
 
@@ -290,17 +290,14 @@ function compute_relative_humidity!(
 end
 ```
 
-### Adding to the `ALL_DIAGNOSTICS` dictionary
+### The `ClimaAtmos` `DiagnosticVariable`s
 
-`ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable` in the
-`ALL_DIAGNOSTICS` dictionary. `ALL_DIAGNOSTICS` maps a `short_name` with the
-corresponding `DiagnosticVariable`.
-
-If you are extending `ClimaAtmos` and want to add a new diagnostic variable to
-`ALL_DIAGNOSTICS`, go ahead and look at the files we `include` in
+`ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable`, index
+with their `short_name`s. If you are extending `ClimaAtmos` and want to add a
+new diagnostic variable, go ahead and look at the files we `include` in
 `diagnostics/Diagnostics.jl`. You can add more diagnostics in those files or add
 a new one. We provide a convenience function, `add_diagnostic_variable!` to add
-new `DiagnosticVariable`s to the `ALL_DIAGNOSTICS` dictionary.
-`add_diagnostic_variable!` take the same arguments as the constructor for
-`DiagnosticVariable`, but also performs additional checks. So, use
-`add_diagnostic_variable!` instead of editing the `ALL_DIAGNOSTICS` directly.
+new `DiagnosticVariable`s. `add_diagnostic_variable!` take the same arguments as
+the constructor for `DiagnosticVariable`, but also performs additional checks.
+Similarly, if you want to retrieve a diagnostic from `ALL_DIAGNOSTICS`, use the
+`get_diagnostic_variable`function.
diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index 7800990e375..d13d8f306a6 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -46,7 +46,7 @@ function common_diagnostics(
 )
     return [
         ScheduledDiagnosticTime(
-            variable = ALL_DIAGNOSTICS[short_name],
+            variable = get_diagnostic_variable(short_name),
             compute_every = :timestep,
             output_every = period, # seconds
             reduction_time_func = reduction,
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
index f1083f0c77b..a58129f9a75 100644
--- a/src/diagnostics/defaults/moisture_model.jl
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -3,14 +3,14 @@ function get_default_diagnostics(::DryModel)
     return [
         daily_averages("air_density")...,
         ScheduledDiagnosticTime(
-            variable = ALL_DIAGNOSTICS["air_density"],
+            variable = get_diagnostic_variable("air_density"),
             compute_every = :timestep,
             output_every = 86400, # seconds
             reduction_time_func = min,
             output_writer = HDF5Writer(),
         ),
         ScheduledDiagnosticIterations(
-            variable = ALL_DIAGNOSTICS["air_density"],
+            variable = get_diagnostic_variable("air_density"),
             compute_every = 1,
             output_every = 1, # iteration
             output_writer = HDF5Writer(),
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 9f375c05273..380ec12e476 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -11,7 +11,8 @@
 # - A dictionary `ALL_DIAGNOSTICS` with all the diagnostics we know how to compute, keyed
 #   over their short name. If you want to add more diagnostics, look at the included files.
 #   You can add your own file if you want to define several new diagnostics that are
-#   conceptually related.
+#   conceptually related. The dictionary `ALL_DIAGNOSTICS` should be considered an
+#   implementation detail.
 #
 # - The definition of what a ScheduledDiagnostics is. Conceptually, a ScheduledDiagnostics is a
 #   DiagnosticVariable we want to compute in a given simulation. For example, it could be
@@ -150,6 +151,20 @@ function add_diagnostic_variable!(;
         DiagnosticVariable(; short_name, long_name, units, comments, compute!)
 end
 
+"""
+
+    get_diagnostic_variable!(short_name)
+
+
+Return a `DiagnosticVariable` from its `short_name`, if it exists.
+"""
+function get_diagnostic_variable(short_name)
+    haskey(ALL_DIAGNOSTICS, short_name) ||
+        error("diagnostic $short_name does not exist")
+
+    return ALL_DIAGNOSTICS[short_name]
+end
+
 # Do you want to define more diagnostics? Add them here
 include("core_diagnostics.jl")
 include("turbconv_diagnostics.jl")
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 2c19c6c41fa..b0da4371403 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -675,7 +675,7 @@ function get_diagnostics(parsed_args, atmos_model)
 
         if isnothing(name)
             name = CAD.get_descriptive_name(
-                CAD.ALL_DIAGNOSTICS[yaml_diag["short_name"]],
+                CAD.get_diagnostic_variable(yaml_diag["short_name"]),
                 period_seconds,
                 reduction_time_func,
                 pre_output_hook!,
@@ -691,7 +691,7 @@ function get_diagnostics(parsed_args, atmos_model)
         push!(
             diagnostics,
             CAD.ScheduledDiagnosticTime(
-                variable = CAD.ALL_DIAGNOSTICS[yaml_diag["short_name"]],
+                variable = CAD.get_diagnostic_variable(yaml_diag["short_name"]),
                 output_every = period_seconds,
                 compute_every = compute_every,
                 reduction_time_func = reduction_time_func,

From 26e730bd8a76770ad85e6a97d52f34487bfec6e8 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 14 Sep 2023 11:39:49 -0700
Subject: [PATCH 47/73] Avoid allocations of diagnostics in callbacks

---
 docs/src/diagnostics.md                 |  27 +++--
 src/diagnostics/core_diagnostics.jl     | 136 +++++++++++++++---------
 src/diagnostics/diagnostic.jl           |   5 +-
 src/diagnostics/turbconv_diagnostics.jl |  35 ------
 src/solver/type_getters.jl              |   3 +-
 5 files changed, 107 insertions(+), 99 deletions(-)
 delete mode 100644 src/diagnostics/turbconv_diagnostics.jl

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 8fc7b4a50d0..4296f295f72 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -211,10 +211,16 @@ The other piece of information needed to specify a `DiagnosticVariable` is a
 function `compute!`. Schematically, a `compute!` has to look like
 ```julia
 function compute!(out, state, cache, time)
-    # FIXME: Remove this line when ClimaCore implements the broadcasting to enable this
-    out .= # Calculations with the state and the cache
+    if isnothing(out)
+        return ... # Calculations with the state and the cache
+    else
+        out .= ... # Calculations with the state and the cache
+    end
 end
 ```
+The first time `compute!` is called, the function has to allocate memory and
+return its output. All the subsequent times, `out` will be the pre-allocated
+area of memory, so the function has to write the new value in place.
 
 Diagnostics are implemented as callbacks functions which pass the `state`,
 `cache`, and `time` from the integrator to `compute!`.
@@ -247,10 +253,11 @@ function compute_relative_humidity!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
-    thermo_params = CAP.thermodynamics_params(integrator.p.params)
-    return TD.relative_humidity.(thermo_params, integrator.p.ᶜts)
+    if isnothing(out)
+        return TD.relative_humidity.(thermo_params, cache.ᶜts)
+    else
+        out .= TD.relative_humidity.(thermo_params, cache.ᶜts)
+    end
 end
 
 compute_relative_humidity!(out, state, cache, time) =
@@ -283,10 +290,12 @@ function compute_relative_humidity!(
     time,
     ::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     thermo_params = CAP.thermodynamics_params(cache.params)
-    return TD.relative_humidity.(thermo_params, cache.ᶜts)
+    if isnothing(out)
+        return TD.relative_humidity.(thermo_params, cache.ᶜts)
+    else
+        out .= TD.relative_humidity.(thermo_params, cache.ᶜts)
+    end
 end
 ```
 
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index f67c3e667d2..fde5da6779d 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -55,9 +55,11 @@ add_diagnostic_variable!(
     units = "kg m^-3",
     comments = "Density of air, a prognostic variable",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
-        return copy(state.c.ρ)
+        if isnothing(out)
+            return copy(state.c.ρ)
+        else
+            out .= state.c.ρ
+        end
     end,
 )
 
@@ -73,9 +75,11 @@ add_diagnostic_variable!(
     units = "m s^-1",
     comments = "Eastward (zonal) wind component, a prognostic variable",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
-        return copy(Geometry.UVector.(cache.ᶜu))
+        if isnothing(out)
+            return copy(Geometry.UVector.(cache.ᶜu))
+        else
+            out .= Geometry.UVector.(cache.ᶜu)
+        end
     end,
 )
 
@@ -91,9 +95,11 @@ add_diagnostic_variable!(
     units = "m s^-1",
     comments = "Northward (meridional) wind component, a prognostic variable",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
-        return copy(Geometry.VVector.(cache.ᶜu))
+        if isnothing(out)
+            return copy(Geometry.VVector.(cache.ᶜu))
+        else
+            out .= Geometry.VVector.(cache.ᶜu)
+        end
     end,
 )
 
@@ -109,9 +115,11 @@ add_diagnostic_variable!(
     units = "m s^-1",
     comments = "Vertical wind component, a prognostic variable",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
-        return copy(Geometry.WVector.(cache.ᶜu))
+        if isnothing(out)
+            return copy(Geometry.WVector.(cache.ᶜu))
+        else
+            out .= Geometry.WVector.(cache.ᶜu)
+        end
     end,
 )
 
@@ -124,10 +132,12 @@ add_diagnostic_variable!(
     units = "K",
     comments = "Temperature of air",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
         thermo_params = CAP.thermodynamics_params(cache.params)
-        return TD.air_temperature.(thermo_params, cache.ᶜts)
+        if isnothing(out)
+            return TD.air_temperature.(thermo_params, cache.ᶜts)
+        else
+            out .= TD.air_temperature.(thermo_params, cache.ᶜts)
+        end
     end,
 )
 
@@ -140,10 +150,12 @@ add_diagnostic_variable!(
     units = "K",
     comments = "Potential temperature of air",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
         thermo_params = CAP.thermodynamics_params(cache.params)
-        return TD.dry_pottemp.(thermo_params, cache.ᶜts)
+        if isnothing(out)
+            return TD.dry_pottemp.(thermo_params, cache.ᶜts)
+        else
+            out .= TD.dry_pottemp.(thermo_params, cache.ᶜts)
+        end
     end,
 )
 
@@ -156,9 +168,11 @@ add_diagnostic_variable!(
     units = "Pa",
     comments = "Pressure of air",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
-        return copy(cache.ᶜp)
+        if isnothing(out)
+            return copy(cache.ᶜp)
+        else
+            out .= cache.ᶜp
+        end
     end,
 )
 
@@ -171,13 +185,14 @@ add_diagnostic_variable!(
     units = "s^-1",
     comments = "Vertical component of relative vorticity",
     compute! = (out, state, cache, time) -> begin
-        # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-        # We will want: out .= integrator.u.c.ρ
         ᶜvort = @. Geometry.WVector(curlₕ(state.c.uₕ))
-        if cache.do_dss
-            Spaces.weighted_dss!(ᶜvort)
+        # We need to ensure smoothness, so we call DSS
+        Spaces.weighted_dss!(ᶜvort)
+        if isnothing(out)
+            return copy(ᶜvort)
+        else
+            out .= ᶜvort
         end
-        return copy(ᶜvort)
     end,
 )
 
@@ -203,10 +218,12 @@ function compute_relative_humidity!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     thermo_params = CAP.thermodynamics_params(cache.params)
-    return TD.relative_humidity.(thermo_params, cache.ᶜts)
+    if isnothing(out)
+        return TD.relative_humidity.(thermo_params, cache.ᶜts)
+    else
+        out .= TD.relative_humidity.(thermo_params, cache.ᶜts)
+    end
 end
 
 add_diagnostic_variable!(
@@ -238,10 +255,12 @@ function compute_specific_humidity!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     thermo_params = CAP.thermodynamics_params(cache.params)
-    return TD.total_specific_humidity.(thermo_params, cache.ᶜts)
+    if isnothing(out)
+        return TD.total_specific_humidity.(thermo_params, cache.ᶜts)
+    else
+        out .= TD.total_specific_humidity.(thermo_params, cache.ᶜts)
+    end
 end
 
 add_diagnostic_variable!(
@@ -273,10 +292,16 @@ function compute_surface_specific_humidity!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     thermo_params = CAP.thermodynamics_params(cache.params)
-    return TD.total_specific_humidity.(thermo_params, cache.sfc_conditions.ts)
+    if isnothing(out)
+        return TD.total_specific_humidity.(
+            thermo_params,
+            cache.sfc_conditions.ts,
+        )
+    else
+        out .=
+            TD.total_specific_humidity.(thermo_params, cache.sfc_conditions.ts)
+    end
 end
 
 add_diagnostic_variable!(
@@ -291,10 +316,12 @@ add_diagnostic_variable!(
 # Surface temperature (2d)
 ###
 function compute_surface_temperature!(out, state, cache, time)
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     thermo_params = CAP.thermodynamics_params(cache.params)
-    return TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
+    if isnothing(out)
+        return TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
+    else
+        out .= TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
+    end
 end
 
 add_diagnostic_variable!(
@@ -331,9 +358,11 @@ function compute_eastward_drag!(
     time,
     energy_form::T,
 ) where {T <: TotalEnergy}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
-    return drag_vector(state, cache).components.data.:1
+    if isnothing(out)
+        return drag_vector(state, cache).components.data.:1
+    else
+        out .= drag_vector(state, cache).components.data.:1
+    end
 end
 
 add_diagnostic_variable!(
@@ -359,9 +388,11 @@ function compute_northward_drag!(
     time,
     energy_form::T,
 ) where {T <: TotalEnergy}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
-    return drag_vector(state, cache).components.data.:2
+    if isnothing(out)
+        return drag_vector(state, cache).components.data.:2
+    else
+        out .= drag_vector(state, cache).components.data.:2
+    end
 end
 
 add_diagnostic_variable!(
@@ -376,13 +407,15 @@ add_diagnostic_variable!(
 # Surface energy flux (2d) - TODO: this may need to be split into sensible and latent heat fluxes
 ###
 function compute_surface_energy_flux!(out, state, cache, time)
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     (; ρ_flux_h_tot) = cache.sfc_conditions
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
     surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
-    return dot.(ρ_flux_h_tot, surface_ct3_unit)
+    if isnothing(out)
+        return dot.(ρ_flux_h_tot, surface_ct3_unit)
+    else
+        out .= dot.(ρ_flux_h_tot, surface_ct3_unit)
+    end
 end
 
 add_diagnostic_variable!(
@@ -424,13 +457,16 @@ function compute_surface_evaporation!(
     moisture_model::T,
     energy_form::TotalEnergy,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
     (; ρ_flux_q_tot) = cache.sfc_conditions
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
     surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
-    return dot.(ρ_flux_q_tot, surface_ct3_unit)
+
+    if isnothing(out)
+        return dot.(ρ_flux_q_tot, surface_ct3_unit)
+    else
+        out .= dot.(ρ_flux_q_tot, surface_ct3_unit)
+    end
 end
 
 add_diagnostic_variable!(
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 380ec12e476..264a0fd99cf 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -42,7 +42,6 @@
 #
 # - This file also also include several other files, including (but not limited to):
 #   - core_diagnostics.jl
-#   - turbconv_diagnostics.jl
 #   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
 #   - reduction_identities.jl
 #   - diagnostic_utils.jl
@@ -167,7 +166,6 @@ end
 
 # Do you want to define more diagnostics? Add them here
 include("core_diagnostics.jl")
-include("turbconv_diagnostics.jl")
 
 # Default diagnostics and higher level interfaces
 include("default_diagnostics.jl")
@@ -635,8 +633,7 @@ function get_callbacks_from_diagnostics(
 
         compute_callback =
             integrator -> begin
-                # FIXME: Change when ClimaCore overrides .= for us to avoid multiple allocations
-                storage[diag] .= variable.compute!(
+                variable.compute!(
                     storage[diag],
                     integrator.u,
                     integrator.p,
diff --git a/src/diagnostics/turbconv_diagnostics.jl b/src/diagnostics/turbconv_diagnostics.jl
deleted file mode 100644
index a0084cdf048..00000000000
--- a/src/diagnostics/turbconv_diagnostics.jl
+++ /dev/null
@@ -1,35 +0,0 @@
-# This file is included in Diagnostics.jl
-
-# TKE
-
-# This is an example of how to compute the same diagnostic variable differently depending on
-# the model. This is also exposed to the user, which could define their own
-# compute_tke.
-
-function compute_tke!(out, state, cache, time, ::EDMFX)
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
-    return copy(cache.ᶜspecific⁰.tke)
-end
-
-function compute_tke!(out, state, cache, time, ::DiagnosticEDMFX)
-    # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
-    # We will want: out .= integrator.u.c.ρ
-    return copy(cache.tke⁰)
-end
-
-compute_tke!(out, state, cache, time) = compute_tke!(out, state, cache.atmos)
-
-function compute_tke!(out, state, cache, time, turbconv_model::T) where {T}
-    error("Cannot compute tke with turbconv_model = $T")
-end
-
-# FIXME: Gabriele wrote this as an example. Gabriele doesn't know anything about the
-# physics. Please fix this!
-add_diagnostic_variable!(
-    short_name = "tke",
-    long_name = "turbolent_kinetic_energy",
-    units = "J",
-    comments = "Turbolent Kinetic Energy",
-    compute! = compute_tke!,
-)
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index b0da4371403..34d81c41818 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -900,7 +900,8 @@ function get_integrator(config::AtmosConfig)
     for diag in diagnostics_iterations
         variable = diag.variable
         try
-            # FIXME: Avoid extra allocations when ClimaCore overloads .= for this use case
+            # The first time we call compute! we use its return value. All the subsequent
+            # times (in the callbacks), we will write the result in place
             diagnostic_storage[diag] = variable.compute!(
                 nothing,
                 integrator.u,

From 02ea9259e8572269750d442b6294396ebc80f6a6 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 15 Sep 2023 14:57:14 -0700
Subject: [PATCH 48/73] Fix element type for more complex Fields

If we a Vector Field, `eltype` will return `Vector`, not float. So, we
have to be more creative with how to get the float_type for general
cases.
---
 src/diagnostics/diagnostic.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 264a0fd99cf..12e3441e2fc 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -577,7 +577,7 @@ reset_accumulator!(_, reduction_time_func::Nothing) = nothing
 function reset_accumulator!(diag_accumulator, reduction_time_func)
     # identity_of_reduction works by dispatching over operation
     identity = identity_of_reduction(reduction_time_func)
-    float_type = eltype(diag_accumulator)
+    float_type = Spaces.undertype(axes((diag_accumulator)))
     identity_ft = convert(float_type, identity)
     diag_accumulator .= identity_ft
 end

From 86be6fc0ac5770ce590cbd53ae73c4446a56544a Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 15 Sep 2023 15:56:20 -0700
Subject: [PATCH 49/73] Allow accumulators to take Vector values

.= does not work with Vector Fields, but parent() returns the underlying
data, so we can directly fill the array with the desired value
---
 src/diagnostics/diagnostic.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 12e3441e2fc..07af85616e0 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -579,7 +579,7 @@ function reset_accumulator!(diag_accumulator, reduction_time_func)
     identity = identity_of_reduction(reduction_time_func)
     float_type = Spaces.undertype(axes((diag_accumulator)))
     identity_ft = convert(float_type, identity)
-    diag_accumulator .= identity_ft
+    fill!(parent(diag_accumulator), identity_ft)
 end
 
 # When the reduction is nothing, we do not need to accumulate anything

From 9491af217b4ec9cfc2ab855aa79f6644089cbd2c Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Fri, 15 Sep 2023 17:17:45 -0700
Subject: [PATCH 50/73] Modify and add some core diagnostics

---
 docs/src/diagnostics.md             |  17 +-
 src/diagnostics/Diagnostics.jl      |   6 +
 src/diagnostics/core_diagnostics.jl | 545 +++++++++++++++++++++-------
 3 files changed, 423 insertions(+), 145 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 4296f295f72..149b5f9566a 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -14,16 +14,16 @@ Second, you can specify the diagnostics you want to output directly in the
 `diagnostics` section of your YAML file. For instance:
 ```
 diagnostics:
-  - short_name: air_density
+  - short_name: rhoa
     output_name: a_name
     period: 3hours
     writer: nc
   - reduction_time: average
-    short_name: air_density
+    short_name: rhoa
     period: 12hours
     writer: h5
 ```
-This adds two diagnostics (both for `air_density`). The `period` keyword
+This adds two diagnostics (both for `rhoa`). The `period` keyword
 identifies the period over which to compute the reduction and how often to save
 to disk. `output_name` is optional, and if provided, it identifies the name of the
 output file.
@@ -188,8 +188,7 @@ In `ClimaAtmos`, we follow the convention that:
 - `short_name` is the name used to identify the variable in the output files and
                 in the file names. It is short, but descriptive. We identify
                 diagnostics by their short name, so the diagnostics defined by
-                `ClimaAtmos` have to have unique `short_name`s. We follow the
-                Coupled Model Intercomparison Project (CMIP) conventions.
+                `ClimaAtmos` have to have unique `short_name`s.
 
 - `long_name`: Name used to describe the variable in the output file.
 
@@ -198,12 +197,8 @@ In `ClimaAtmos`, we follow the convention that:
 - `comments`: More verbose explanation of what the variable is, or comments related to how
               it is defined or computed.
 
-In `ClimaAtmos`, we try to follow [this Google
-spreadsheet](https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx)
-for variable naming, with a `long_name` the does not have spaces and capital
-letters. [Standard
-names](http://cfconventions.org/Data/cf-standard-names/current/build/cf-standard-name-table.html)
-are not used.
+In `ClimaAtmos`, we follow the [CMIP6 MIP table](https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb)
+for short names and long names where available. Standard names in the table are not used.
 
 ### Compute function
 
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 9646f654cc2..29a02c6d6c1 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -20,6 +20,12 @@ import ..NonEquilMoistModel
 # energy_form
 import ..TotalEnergy
 
+# precip_model
+import ..Microphysics0Moment
+
+# radiation
+import ClimaAtmos.RRTMGPInterface as RRTMGPI
+
 # turbconv_model
 import ..EDMFX
 import ..DiagnosticEDMFX
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index fde5da6779d..9394db8ebaf 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -13,7 +13,7 @@
 # 1. Define a catch base function that does the computation we want to do for the case we know
 # how to handle, for example
 #
-# function compute_relative_humidity!(
+# function compute_hur!(
 #     out,
 #     state,
 #     cache,
@@ -26,17 +26,17 @@
 #
 # 2. Define a function that has the correct signature and calls this function
 #
-# compute_relative_humidity!(out, state, cache, time) =
-#     compute_relative_humidity!(out, state, cache, time, cache.atmos.moisture_model)
+# compute_hur!(out, state, cache, time) =
+#     compute_hur!(out, state, cache, time, cache.atmos.moisture_model)
 #
 # 3. Define a function that returns an error when the model is incorrect
 #
-# compute_relative_humidity!(_, _, _, _, model::T) where {T} =
+# compute_hur!(_, _, _, _, model::T) where {T} =
 #     error_diagnostic_variable("relative_humidity", model)
 #
 # We can also output a specific error message
 #
-# compute_relative_humidity!(_, _, _, _, model::T) where {T} =
+# compute_hur!(_, _, _, _, model::T) where {T} =
 #     error_diagnostic_variable("relative humidity makes sense only for moist models")
 
 # General helper functions for undefined diagnostics for a particular model
@@ -50,10 +50,10 @@ error_diagnostic_variable(variable, model::T) where {T} =
 # Rho (3d)
 ###
 add_diagnostic_variable!(
-    short_name = "air_density",
+    short_name = "rhoa",
     long_name = "Air Density",
     units = "kg m^-3",
-    comments = "Density of air, a prognostic variable",
+    comments = "Density of air",
     compute! = (out, state, cache, time) -> begin
         if isnothing(out)
             return copy(state.c.ρ)
@@ -66,19 +66,16 @@ add_diagnostic_variable!(
 ###
 # U velocity (3d)
 ###
-
-# TODO: This velocity might not be defined (e.g., in a column model). Add dispatch to catch
-# that.
 add_diagnostic_variable!(
-    short_name = "eastward_wind",
+    short_name = "ua",
     long_name = "Eastward Wind",
     units = "m s^-1",
-    comments = "Eastward (zonal) wind component, a prognostic variable",
+    comments = "Eastward (zonal) wind component",
     compute! = (out, state, cache, time) -> begin
         if isnothing(out)
-            return copy(Geometry.UVector.(cache.ᶜu))
+            return copy(Geometry.UVector.(cache.ᶜu).components.data.:1)
         else
-            out .= Geometry.UVector.(cache.ᶜu)
+            out .= Geometry.UVector.(cache.ᶜu).components.data.:1
         end
     end,
 )
@@ -86,19 +83,16 @@ add_diagnostic_variable!(
 ###
 # V velocity (3d)
 ###
-
-# TODO: This velocity might not be defined (e.g., in a column model). Add dispatch to catch
-# that.
 add_diagnostic_variable!(
-    short_name = "northward_wind",
+    short_name = "va",
     long_name = "Northward Wind",
     units = "m s^-1",
-    comments = "Northward (meridional) wind component, a prognostic variable",
+    comments = "Northward (meridional) wind component",
     compute! = (out, state, cache, time) -> begin
         if isnothing(out)
-            return copy(Geometry.VVector.(cache.ᶜu))
+            return copy(Geometry.VVector.(cache.ᶜu).components.data.:1)
         else
-            out .= Geometry.VVector.(cache.ᶜu)
+            out .= Geometry.VVector.(cache.ᶜu).components.data.:1
         end
     end,
 )
@@ -110,15 +104,15 @@ add_diagnostic_variable!(
 # but this is probably more useful for now
 #
 add_diagnostic_variable!(
-    short_name = "vertical_wind",
-    long_name = "Vertical Wind",
+    short_name = "wa",
+    long_name = "Upward Air Velocity",
     units = "m s^-1",
-    comments = "Vertical wind component, a prognostic variable",
+    comments = "Vertical wind component",
     compute! = (out, state, cache, time) -> begin
         if isnothing(out)
-            return copy(Geometry.WVector.(cache.ᶜu))
+            return copy(Geometry.WVector.(cache.ᶜu).components.data.:1)
         else
-            out .= Geometry.WVector.(cache.ᶜu)
+            out .= Geometry.WVector.(cache.ᶜu).components.data.:1
         end
     end,
 )
@@ -127,7 +121,7 @@ add_diagnostic_variable!(
 # Temperature (3d)
 ###
 add_diagnostic_variable!(
-    short_name = "air_temperature",
+    short_name = "ta",
     long_name = "Air Temperature",
     units = "K",
     comments = "Temperature of air",
@@ -145,8 +139,8 @@ add_diagnostic_variable!(
 # Potential temperature (3d)
 ###
 add_diagnostic_variable!(
-    short_name = "air_potential_temperature",
-    long_name = "Air potential temperature",
+    short_name = "thetaa",
+    long_name = "Air Potential Temperature",
     units = "K",
     comments = "Potential temperature of air",
     compute! = (out, state, cache, time) -> begin
@@ -163,8 +157,8 @@ add_diagnostic_variable!(
 # Air pressure (3d)
 ###
 add_diagnostic_variable!(
-    short_name = "air_pressure",
-    long_name = "Air pressure",
+    short_name = "pfull",
+    long_name = "Pressure at Model Full-Levels",
     units = "Pa",
     comments = "Pressure of air",
     compute! = (out, state, cache, time) -> begin
@@ -180,18 +174,18 @@ add_diagnostic_variable!(
 # Vorticity (3d)
 ###
 add_diagnostic_variable!(
-    short_name = "atmosphere_relative_vorticity",
-    long_name = "Vertical component of relative vorticity",
+    short_name = "rv",
+    long_name = "Relative Vorticity",
     units = "s^-1",
     comments = "Vertical component of relative vorticity",
     compute! = (out, state, cache, time) -> begin
-        ᶜvort = @. Geometry.WVector(curlₕ(state.c.uₕ))
+        vort = @. Geometry.WVector(curlₕ(state.c.uₕ)).components.data.:1
         # We need to ensure smoothness, so we call DSS
-        Spaces.weighted_dss!(ᶜvort)
+        Spaces.weighted_dss!(vort)
         if isnothing(out)
-            return copy(ᶜvort)
+            return copy(vort)
         else
-            out .= ᶜvort
+            out .= vort
         end
     end,
 )
@@ -200,18 +194,12 @@ add_diagnostic_variable!(
 ###
 # Relative humidity (3d)
 ###
-compute_relative_humidity!(out, state, cache, time) =
-    compute_relative_humidity!(
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-    )
-compute_relative_humidity!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("relative_humidity", model)
+compute_hur!(out, state, cache, time) =
+    compute_hur!(out, state, cache, time, cache.atmos.moisture_model)
+compute_hur!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("hur", model)
 
-function compute_relative_humidity!(
+function compute_hur!(
     out,
     state,
     cache,
@@ -227,28 +215,22 @@ function compute_relative_humidity!(
 end
 
 add_diagnostic_variable!(
-    short_name = "relative_humidity",
+    short_name = "hur",
     long_name = "Relative Humidity",
     units = "",
     comments = "Total amount of water vapor in the air relative to the amount achievable by saturation at the current temperature",
-    compute! = compute_relative_humidity!,
+    compute! = compute_hur!,
 )
 
 ###
 # Total specific humidity (3d)
 ###
-compute_specific_humidity!(out, state, cache, time) =
-    compute_specific_humidity!(
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-    )
-compute_specific_humidity!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("specific_humidity", model)
+compute_hus!(out, state, cache, time) =
+    compute_hus!(out, state, cache, time, cache.atmos.moisture_model)
+compute_hus!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("hus", model)
 
-function compute_specific_humidity!(
+function compute_hus!(
     out,
     state,
     cache,
@@ -264,28 +246,22 @@ function compute_specific_humidity!(
 end
 
 add_diagnostic_variable!(
-    short_name = "specific_humidity",
+    short_name = "hus",
     long_name = "Specific Humidity",
     units = "kg kg^-1",
-    comments = "Mass of all water phases per mass of air, a prognostic variable",
-    compute! = compute_specific_humidity!,
+    comments = "Mass of all water phases per mass of air",
+    compute! = compute_hus!,
 )
 
 ###
 # Surface specific humidity (2d)
 ###
-compute_surface_specific_humidity!(out, state, cache, time) =
-    compute_surface_specific_humidity!(
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-    )
-compute_surface_specific_humidity!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("surface_specific_humidity", model)
+compute_hussfc!(out, state, cache, time) =
+    compute_hussfc!(out, state, cache, time, cache.atmos.moisture_model)
+compute_hussfc!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("hussfc", model)
 
-function compute_surface_specific_humidity!(
+function compute_hussfc!(
     out,
     state,
     cache,
@@ -305,59 +281,51 @@ function compute_surface_specific_humidity!(
 end
 
 add_diagnostic_variable!(
-    short_name = "surface_specific_humidity",
+    short_name = "hussfc",
     long_name = "Surface Specific Humidity",
     units = "kg kg^-1",
-    comments = "Mass of all water phases per mass of air in the near-surface layer",
-    compute! = compute_surface_specific_humidity!,
+    comments = "Mass of all water phases per mass of air in the layer infinitely close to the surface",
+    compute! = compute_hussfc!,
 )
 
 ###
 # Surface temperature (2d)
 ###
-function compute_surface_temperature!(out, state, cache, time)
-    thermo_params = CAP.thermodynamics_params(cache.params)
-    if isnothing(out)
-        return TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
-    else
-        out .= TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
-    end
-end
-
 add_diagnostic_variable!(
-    short_name = "surface_temperature",
+    short_name = "ts",
     long_name = "Surface Temperature",
     units = "K",
-    comments = "Temperature of the surface",
-    compute! = compute_surface_temperature!,
+    comments = "Temperature of the lower boundary of the atmosphere",
+    compute! = (out, state, cache, time) -> begin
+        thermo_params = CAP.thermodynamics_params(cache.params)
+        if isnothing(out)
+            return TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
+        else
+            out .= TD.air_temperature.(thermo_params, cache.sfc_conditions.ts)
+        end
+    end,
 )
 
 ###
 # Eastward surface drag component (2d)
 ###
-compute_eastward_drag!(out, state, cache, time) =
-    compute_eastward_drag!(out, state, cache, time, cache.atmos.energy_form)
-compute_eastward_drag!(_, _, _, _, energy_form::T) where {T} =
-    error_diagnostic_variable("eastward_drag", energy_form)
-
 function drag_vector(state, cache)
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
     surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
     (; ρ_flux_uₕ) = cache.sfc_conditions
     return Geometry.UVVector.(
-        adjoint.(ρ_flux_uₕ ./ Spaces.level(ᶠinterp.(state.c.ρ), half)) .*
+        adjoint.(ρ_flux_uₕ) .*
         surface_ct3_unit
     )
 end
 
-function compute_eastward_drag!(
+function compute_tauu!(
     out,
     state,
     cache,
     time,
-    energy_form::T,
-) where {T <: TotalEnergy}
+)
     if isnothing(out)
         return drag_vector(state, cache).components.data.:1
     else
@@ -366,28 +334,22 @@ function compute_eastward_drag!(
 end
 
 add_diagnostic_variable!(
-    short_name = "eastward_drag",
-    long_name = "Eastward component of the surface drag",
-    units = "kg m^-2 s^-2",
+    short_name = "tauu",
+    long_name = "Surface Downward Eastward Wind Stress",
+    units = "Pa",
     comments = "Eastward component of the surface drag",
-    compute! = compute_eastward_drag!,
+    compute! = compute_tauu!,
 )
 
 ###
 # Northward surface drag component (2d)
 ###
-compute_northward_drag!(out, state, cache, time) =
-    compute_northward_drag!(out, state, cache, time, cache.atmos.energy_form)
-compute_northward_drag!(_, _, _, _, energy_form::T) where {T} =
-    error_diagnostic_variable("northward_drag", energy_form)
-
-function compute_northward_drag!(
+function compute_tauv!(
     out,
     state,
     cache,
     time,
-    energy_form::T,
-) where {T <: TotalEnergy}
+)
     if isnothing(out)
         return drag_vector(state, cache).components.data.:2
     else
@@ -396,17 +358,22 @@ function compute_northward_drag!(
 end
 
 add_diagnostic_variable!(
-    short_name = "northward_drag",
-    long_name = "Northward component of the surface drag",
-    units = "kg m^-2 s^-2",
+    short_name = "tauv",
+    long_name = "Surface Downward Northward Wind Stress",
+    units = "Pa",
     comments = "Northward component of the surface drag",
-    compute! = compute_northward_drag!,
+    compute! = compute_tauv!,
 )
 
 ###
 # Surface energy flux (2d) - TODO: this may need to be split into sensible and latent heat fluxes
 ###
-function compute_surface_energy_flux!(out, state, cache, time)
+compute_hfes!(out, state, cache, time) =
+    compute_hfes!(out, state, cache, time, cache.atmos.energy_form)
+compute_hfes!(_, _, _, _, energy_form::T) where {T} =
+    error_diagnostic_variable("hfes", energy_form)
+
+function compute_hfes!(out, state, cache, time, energy_form::TotalEnergy)
     (; ρ_flux_h_tot) = cache.sfc_conditions
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
@@ -419,26 +386,25 @@ function compute_surface_energy_flux!(out, state, cache, time)
 end
 
 add_diagnostic_variable!(
-    short_name = "surface_energy_flux",
-    long_name = "Surface energy flux",
+    short_name = "hfes",
+    long_name = "Surface Upward Energy Flux",
     units = "W m^-2",
     comments = "Energy flux at the surface",
-    compute! = compute_surface_energy_flux!,
+    compute! = compute_hfes!,
 )
 
 ###
 # Surface evaporation (2d)
 ###
-compute_surface_evaporation!(out, state, cache, time) =
-    compute_surface_evaporation!(
-        out,
-        state,
-        cache,
-        time,
-        cache.atmos.moisture_model,
-        cache.atmos.energy_form,
-    )
-compute_surface_evaporation!(
+compute_evspsbl!(out, state, cache, time) = compute_evspsbl!(
+    out,
+    state,
+    cache,
+    time,
+    cache.atmos.moisture_model,
+    cache.atmos.energy_form,
+)
+compute_evspsbl!(
     _,
     _,
     _,
@@ -449,7 +415,7 @@ compute_surface_evaporation!(
     "Can only compute surface_evaporation with energy_form = TotalEnergy() and with a moist model",
 )
 
-function compute_surface_evaporation!(
+function compute_evspsbl!(
     out,
     state,
     cache,
@@ -470,9 +436,320 @@ function compute_surface_evaporation!(
 end
 
 add_diagnostic_variable!(
-    short_name = "surface_evaporation",
-    long_name = "Surface evaporation",
-    units = "kg s^-1 m^-2",
+    short_name = "evspsbl",
+    long_name = "Evaporation Including Sublimation and Transpiration",
+    units = "kg m^-2 s^-1",
     comments = "evaporation at the surface",
-    compute! = compute_surface_evaporation!,
+    compute! = compute_evspsbl!,
+)
+
+###
+# Precipitation (2d) - TODO: change to kg m^-2 s^-1
+###
+compute_pr!(out, state, cache, time) =
+    compute_pr!(out, state, cache, time, cache.atmos.precip_model)
+compute_pr!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("pr", model)
+
+function compute_pr!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::T,
+) where {T <: Microphysics0Moment}
+    thermo_params = CAP.thermodynamics_params(cache.params)
+    if isnothing(out)
+        return cache.col_integrated_rain .+ cache.col_integrated_snow
+    else
+        out .= cache.col_integrated_rain .+ cache.col_integrated_snow
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "pr",
+    long_name = "Precipitation",
+    units = "m s^-1",
+    comments = "Total precipitation including rain and snow",
+    compute! = compute_pr!,
+)
+
+###
+# Donwelling shortwave radiation (3d)
+###
+compute_rsd!(out, state, cache, time) =
+    compute_rsd!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsd!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rsd", model)
+
+function compute_rsd!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    sw_flux_dn = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_sw_flux_dn),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return sw_flux_dn
+    else
+        out .= sw_flux_dn
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsd",
+    long_name = "Downwelling Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling shortwave radiation",
+    compute! = compute_rsd!,
+)
+
+###
+# Upwelling shortwave radiation (3d)
+###
+compute_rsu!(out, state, cache, time) =
+    compute_rsu!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsu!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rsu", model)
+
+function compute_rsu!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    sw_flux_up = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_sw_flux_up),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return sw_flux_up
+    else
+        out .= sw_flux_up
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsu",
+    long_name = "Upwelling Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling shortwave radiation",
+    compute! = compute_rsu!,
+)
+
+###
+# Downwelling longwave radiation (3d)
+###
+compute_rld!(out, state, cache, time) =
+    compute_rld!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rld!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rld", model)
+
+function compute_rld!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    lw_flux_dn = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_lw_flux_dn),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return lw_flux_dn
+    else
+        out .= lw_flux_dn
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rld",
+    long_name = "Downwelling Longwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling longwave radiation",
+    compute! = compute_rld!,
+)
+
+###
+# Upwelling longwave radiation (3d)
+###
+compute_rlu!(out, state, cache, time) =
+    compute_rlu!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rlu!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rlu", model)
+
+function compute_rlu!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    lw_flux_up = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_lw_flux_up),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return lw_flux_up
+    else
+        out .= lw_flux_up
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rlu",
+    long_name = "Upwelling Longwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling longwave radiation",
+    compute! = compute_rlu!,
+)
+
+###
+# Donwelling clear sky shortwave radiation (3d)
+###
+compute_rsdcs!(out, state, cache, time) =
+    compute_rsdcs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsdcs!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rsdcs", model)
+
+function compute_rsdcs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
+    FT = eltype(state)
+    clear_sw_flux_dn = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_clear_sw_flux_dn),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return clear_sw_flux_dn
+    else
+        out .= clear_sw_flux_dn
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsdcs",
+    long_name = "Downwelling Clear-Sky Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling clear sky shortwave radiation",
+    compute! = compute_rsdcs!,
+)
+
+###
+# Upwelling clear sky shortwave radiation (3d)
+###
+compute_rsucs!(out, state, cache, time) =
+    compute_rsucs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsucs!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rsucs", model)
+
+function compute_rsucs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
+    FT = eltype(state)
+    clear_sw_flux_up = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_clear_sw_flux_up),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return clear_sw_flux_up
+    else
+        out .= clear_sw_flux_up
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsucs",
+    long_name = "Upwelling Clear-Sky Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling clear sky shortwave radiation",
+    compute! = compute_rsucs!,
+)
+
+###
+# Downwelling clear sky longwave radiation (3d)
+###
+compute_rldcs!(out, state, cache, time) =
+    compute_rldcs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rldcs!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rldcs", model)
+
+function compute_rldcs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
+    FT = eltype(state)
+    clear_lw_flux_dn = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_clear_lw_flux_dn),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return clear_lw_flux_dn
+    else
+        out .= clear_lw_flux_dn
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rldcs",
+    long_name = "Downwelling Clear-Sky Longwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling clear sky longwave radiation",
+    compute! = compute_rldcs!,
+)
+
+###
+# Upwelling clear sky longwave radiation (3d)
+###
+compute_rlucs!(out, state, cache, time) =
+    compute_rlucs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rlucs!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rlucs", model)
+
+function compute_rlucs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
+    FT = eltype(state)
+    clear_lw_flux_up = RRTMGPI.array2field(
+        FT.(cache.radiation_model.face_clear_lw_flux_up),
+        axes(state.f),
+    )
+    if isnothing(out)
+        return clear_lw_flux_up
+    else
+        out .= clear_lw_flux_up
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rlucs",
+    long_name = "Upwelling Clear-Sky Longwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling clear sky longwave radiation",
+    compute! = compute_rlucs!,
 )

From f63f15b8f807296365b7544244700574dd38514b Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 08:37:42 -0700
Subject: [PATCH 51/73] Remove allocation in drag_vector

---
 src/diagnostics/core_diagnostics.jl | 55 +++++++++++++----------------
 1 file changed, 24 insertions(+), 31 deletions(-)

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index 9394db8ebaf..b1e6215beaa 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -307,32 +307,41 @@ add_diagnostic_variable!(
 )
 
 ###
-# Eastward surface drag component (2d)
+# Eastward and northward surface drag component (2d)
 ###
-function drag_vector(state, cache)
+compute_tau!(_, _, _, _, energy_form::T) where {T} =
+    error_diagnostic_variable("tau", energy_form)
+
+function compute_tau!(out, state, cache, component, energy_form::TotalEnergy)
     sfc_local_geometry =
         Fields.level(Fields.local_geometry_field(state.f), Fields.half)
     surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry))
     (; ρ_flux_uₕ) = cache.sfc_conditions
-    return Geometry.UVVector.(
-        adjoint.(ρ_flux_uₕ) .*
-        surface_ct3_unit
-    )
-end
 
-function compute_tauu!(
-    out,
-    state,
-    cache,
-    time,
-)
     if isnothing(out)
-        return drag_vector(state, cache).components.data.:1
+        return getproperty(
+            Geometry.UVVector.(
+                adjoint.(ρ_flux_uₕ) .* surface_ct3_unit
+            ).components.data,
+            component,
+        )
     else
-        out .= drag_vector(state, cache).components.data.:1
+        out .= getproperty(
+            Geometry.UVVector.(
+                adjoint.(ρ_flux_uₕ) .* surface_ct3_unit
+            ).components.data,
+            component,
+        )
     end
+
+    return
 end
 
+compute_tauu!(out, state, cache, time) =
+    compute_tau!(out, state, cache, :1, cache.atmos.energy_form)
+compute_tauv!(out, state, cache, time) =
+    compute_tau!(out, state, cache, :2, cache.atmos.energy_form)
+
 add_diagnostic_variable!(
     short_name = "tauu",
     long_name = "Surface Downward Eastward Wind Stress",
@@ -341,22 +350,6 @@ add_diagnostic_variable!(
     compute! = compute_tauu!,
 )
 
-###
-# Northward surface drag component (2d)
-###
-function compute_tauv!(
-    out,
-    state,
-    cache,
-    time,
-)
-    if isnothing(out)
-        return drag_vector(state, cache).components.data.:2
-    else
-        out .= drag_vector(state, cache).components.data.:2
-    end
-end
-
 add_diagnostic_variable!(
     short_name = "tauv",
     long_name = "Surface Downward Northward Wind Stress",

From 17dcb20970c4621d0abcf5949a95cf056c4e7f4f Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 08:52:24 -0700
Subject: [PATCH 52/73] Clean up radiation_diagnostics.jl

---
 src/diagnostics/core_diagnostics.jl      | 298 +---------------------
 src/diagnostics/diagnostic.jl            |   3 +-
 src/diagnostics/radiation_diagnostics.jl | 299 +++++++++++++++++++++++
 3 files changed, 306 insertions(+), 294 deletions(-)
 create mode 100644 src/diagnostics/radiation_diagnostics.jl

diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index b1e6215beaa..f447e4d9665 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -237,11 +237,10 @@ function compute_hus!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    thermo_params = CAP.thermodynamics_params(cache.params)
     if isnothing(out)
-        return TD.total_specific_humidity.(thermo_params, cache.ᶜts)
+        return state.c.ρq_tot ./ state.c.ρ
     else
-        out .= TD.total_specific_humidity.(thermo_params, cache.ᶜts)
+        out .= state.c.ρq_tot ./ state.c.ρ
     end
 end
 
@@ -441,17 +440,10 @@ add_diagnostic_variable!(
 ###
 compute_pr!(out, state, cache, time) =
     compute_pr!(out, state, cache, time, cache.atmos.precip_model)
-compute_pr!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("pr", model)
+compute_pr!(_, _, _, _, precip_model::T) where {T} =
+    error_diagnostic_variable("pr", precip_model)
 
-function compute_pr!(
-    out,
-    state,
-    cache,
-    time,
-    precip_model::T,
-) where {T <: Microphysics0Moment}
-    thermo_params = CAP.thermodynamics_params(cache.params)
+function compute_pr!(out, state, cache, time, precip_model::Microphysics0Moment)
     if isnothing(out)
         return cache.col_integrated_rain .+ cache.col_integrated_snow
     else
@@ -466,283 +458,3 @@ add_diagnostic_variable!(
     comments = "Total precipitation including rain and snow",
     compute! = compute_pr!,
 )
-
-###
-# Donwelling shortwave radiation (3d)
-###
-compute_rsd!(out, state, cache, time) =
-    compute_rsd!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rsd!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rsd", model)
-
-function compute_rsd!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AbstractRRTMGPMode}
-    FT = eltype(state)
-    sw_flux_dn = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_sw_flux_dn),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return sw_flux_dn
-    else
-        out .= sw_flux_dn
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rsd",
-    long_name = "Downwelling Shortwave Radiation",
-    units = "W m^-2",
-    comments = "Downwelling shortwave radiation",
-    compute! = compute_rsd!,
-)
-
-###
-# Upwelling shortwave radiation (3d)
-###
-compute_rsu!(out, state, cache, time) =
-    compute_rsu!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rsu!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rsu", model)
-
-function compute_rsu!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AbstractRRTMGPMode}
-    FT = eltype(state)
-    sw_flux_up = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_sw_flux_up),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return sw_flux_up
-    else
-        out .= sw_flux_up
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rsu",
-    long_name = "Upwelling Shortwave Radiation",
-    units = "W m^-2",
-    comments = "Upwelling shortwave radiation",
-    compute! = compute_rsu!,
-)
-
-###
-# Downwelling longwave radiation (3d)
-###
-compute_rld!(out, state, cache, time) =
-    compute_rld!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rld!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rld", model)
-
-function compute_rld!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AbstractRRTMGPMode}
-    FT = eltype(state)
-    lw_flux_dn = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_lw_flux_dn),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return lw_flux_dn
-    else
-        out .= lw_flux_dn
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rld",
-    long_name = "Downwelling Longwave Radiation",
-    units = "W m^-2",
-    comments = "Downwelling longwave radiation",
-    compute! = compute_rld!,
-)
-
-###
-# Upwelling longwave radiation (3d)
-###
-compute_rlu!(out, state, cache, time) =
-    compute_rlu!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rlu!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rlu", model)
-
-function compute_rlu!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AbstractRRTMGPMode}
-    FT = eltype(state)
-    lw_flux_up = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_lw_flux_up),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return lw_flux_up
-    else
-        out .= lw_flux_up
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rlu",
-    long_name = "Upwelling Longwave Radiation",
-    units = "W m^-2",
-    comments = "Upwelling longwave radiation",
-    compute! = compute_rlu!,
-)
-
-###
-# Donwelling clear sky shortwave radiation (3d)
-###
-compute_rsdcs!(out, state, cache, time) =
-    compute_rsdcs!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rsdcs!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rsdcs", model)
-
-function compute_rsdcs!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
-    FT = eltype(state)
-    clear_sw_flux_dn = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_clear_sw_flux_dn),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return clear_sw_flux_dn
-    else
-        out .= clear_sw_flux_dn
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rsdcs",
-    long_name = "Downwelling Clear-Sky Shortwave Radiation",
-    units = "W m^-2",
-    comments = "Downwelling clear sky shortwave radiation",
-    compute! = compute_rsdcs!,
-)
-
-###
-# Upwelling clear sky shortwave radiation (3d)
-###
-compute_rsucs!(out, state, cache, time) =
-    compute_rsucs!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rsucs!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rsucs", model)
-
-function compute_rsucs!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
-    FT = eltype(state)
-    clear_sw_flux_up = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_clear_sw_flux_up),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return clear_sw_flux_up
-    else
-        out .= clear_sw_flux_up
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rsucs",
-    long_name = "Upwelling Clear-Sky Shortwave Radiation",
-    units = "W m^-2",
-    comments = "Upwelling clear sky shortwave radiation",
-    compute! = compute_rsucs!,
-)
-
-###
-# Downwelling clear sky longwave radiation (3d)
-###
-compute_rldcs!(out, state, cache, time) =
-    compute_rldcs!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rldcs!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rldcs", model)
-
-function compute_rldcs!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
-    FT = eltype(state)
-    clear_lw_flux_dn = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_clear_lw_flux_dn),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return clear_lw_flux_dn
-    else
-        out .= clear_lw_flux_dn
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rldcs",
-    long_name = "Downwelling Clear-Sky Longwave Radiation",
-    units = "W m^-2",
-    comments = "Downwelling clear sky longwave radiation",
-    compute! = compute_rldcs!,
-)
-
-###
-# Upwelling clear sky longwave radiation (3d)
-###
-compute_rlucs!(out, state, cache, time) =
-    compute_rlucs!(out, state, cache, time, cache.atmos.radiation_mode)
-compute_rlucs!(_, _, _, _, model::T) where {T} =
-    error_diagnostic_variable("rlucs", model)
-
-function compute_rlucs!(
-    out,
-    state,
-    cache,
-    time,
-    radiation_mode::T,
-) where {T <: RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics}
-    FT = eltype(state)
-    clear_lw_flux_up = RRTMGPI.array2field(
-        FT.(cache.radiation_model.face_clear_lw_flux_up),
-        axes(state.f),
-    )
-    if isnothing(out)
-        return clear_lw_flux_up
-    else
-        out .= clear_lw_flux_up
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "rlucs",
-    long_name = "Upwelling Clear-Sky Longwave Radiation",
-    units = "W m^-2",
-    comments = "Upwelling clear sky longwave radiation",
-    compute! = compute_rlucs!,
-)
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 07af85616e0..abd56a81b1b 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -41,7 +41,7 @@
 #   can compute stuff like averages).
 #
 # - This file also also include several other files, including (but not limited to):
-#   - core_diagnostics.jl
+#   - core_diagnostics.jl, radiation_diagnostics.jl
 #   - default_diagnostics.jl (which defines all the higher-level interfaces and defaults)
 #   - reduction_identities.jl
 #   - diagnostic_utils.jl
@@ -166,6 +166,7 @@ end
 
 # Do you want to define more diagnostics? Add them here
 include("core_diagnostics.jl")
+include("radiation_diagnostics.jl")
 
 # Default diagnostics and higher level interfaces
 include("default_diagnostics.jl")
diff --git a/src/diagnostics/radiation_diagnostics.jl b/src/diagnostics/radiation_diagnostics.jl
new file mode 100644
index 00000000000..867a6ce25c3
--- /dev/null
+++ b/src/diagnostics/radiation_diagnostics.jl
@@ -0,0 +1,299 @@
+# This file is included in Diagnostics.jl
+
+# Radiative fluxes
+
+###
+# Downwelling shortwave radiation (3d)
+###
+compute_rsd!(out, state, cache, time) =
+    compute_rsd!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsd!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rsd", radiation_mode)
+
+function compute_rsd!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_sw_flux_dn),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_sw_flux_dn),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsd",
+    long_name = "Downwelling Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling shortwave radiation",
+    compute! = compute_rsd!,
+)
+
+###
+# Upwelling shortwave radiation (3d)
+###
+compute_rsu!(out, state, cache, time) =
+    compute_rsu!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsu!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rsu", radiation_mode)
+
+function compute_rsu!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_sw_flux_up),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_sw_flux_up),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsu",
+    long_name = "Upwelling Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling shortwave radiation",
+    compute! = compute_rsu!,
+)
+
+###
+# Downwelling longwave radiation (3d)
+###
+compute_rld!(out, state, cache, time) =
+    compute_rld!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rld!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rld", radiation_mode)
+
+function compute_rld!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_lw_flux_dn),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_lw_flux_dn),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rld",
+    long_name = "Downwelling Longwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling longwave radiation",
+    compute! = compute_rld!,
+)
+
+###
+# Upwelling longwave radiation (3d)
+###
+compute_rlu!(out, state, cache, time) =
+    compute_rlu!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rlu!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rlu", radiation_mode)
+
+function compute_rlu!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::T,
+) where {T <: RRTMGPI.AbstractRRTMGPMode}
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_lw_flux_up),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_lw_flux_up),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rlu",
+    long_name = "Upwelling Longwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling longwave radiation",
+    compute! = compute_rlu!,
+)
+
+###
+# Downelling clear sky shortwave radiation (3d)
+###
+compute_rsdcs!(out, state, cache, time) =
+    compute_rsdcs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsdcs!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rsdcs", radiation_mode)
+
+function compute_rsdcs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics,
+)
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_sw_flux_dn),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_sw_flux_dn),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsdcs",
+    long_name = "Downwelling Clear-Sky Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling clear sky shortwave radiation",
+    compute! = compute_rsdcs!,
+)
+
+###
+# Upwelling clear sky shortwave radiation (3d)
+###
+compute_rsucs!(out, state, cache, time) =
+    compute_rsucs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rsucs!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rsucs", radiation_mode)
+
+function compute_rsucs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics,
+)
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_sw_flux_up),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_sw_flux_up),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rsucs",
+    long_name = "Upwelling Clear-Sky Shortwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling clear sky shortwave radiation",
+    compute! = compute_rsucs!,
+)
+
+###
+# Downwelling clear sky longwave radiation (3d)
+###
+compute_rldcs!(out, state, cache, time) =
+    compute_rldcs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rldcs!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rldcs", radiation_mode)
+
+function compute_rldcs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics,
+)
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_lw_flux_dn),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_lw_flux_dn),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rldcs",
+    long_name = "Downwelling Clear-Sky Longwave Radiation",
+    units = "W m^-2",
+    comments = "Downwelling clear sky longwave radiation",
+    compute! = compute_rldcs!,
+)
+
+###
+# Upwelling clear sky longwave radiation (3d)
+###
+compute_rlucs!(out, state, cache, time) =
+    compute_rlucs!(out, state, cache, time, cache.atmos.radiation_mode)
+compute_rlucs!(_, _, _, _, radiation_mode::T) where {T} =
+    error_diagnostic_variable("rlucs", radiation_mode)
+
+function compute_rlucs!(
+    out,
+    state,
+    cache,
+    time,
+    radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics,
+)
+    FT = eltype(state)
+    if isnothing(out)
+        return RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_lw_flux_up),
+            axes(state.f),
+        )
+    else
+        out .= RRTMGPI.array2field(
+            FT.(cache.radiation_model.face_clear_lw_flux_up),
+            axes(state.f),
+        )
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rlucs",
+    long_name = "Upwelling Clear-Sky Longwave Radiation",
+    units = "W m^-2",
+    comments = "Upwelling clear sky longwave radiation",
+    compute! = compute_rlucs!,
+)

From e6e901fdffe1758c556433d08f9c6b5296b5dc3b Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 10:58:46 -0700
Subject: [PATCH 53/73] Fix reduction name in suffix

---
 src/diagnostics/diagnostics_utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl
index a403ecbbebf..8746535159a 100644
--- a/src/diagnostics/diagnostics_utils.jl
+++ b/src/diagnostics/diagnostics_utils.jl
@@ -57,7 +57,7 @@ function get_descriptive_name(
 
             suffix = period * red
         else
-            suffix = "$(output_every)it_(reduction_time_func)"
+            suffix = "$(output_every)it_$(red)"
         end
     else
         suffix = "inst"

From c549ac8fb871aca506039ca4ed6fd3dee686808a Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 14:50:36 -0700
Subject: [PATCH 54/73] Add standard names and long ones for ScheduledDiag

---
 docs/make_diagnostic_table.jl            |  10 ++-
 docs/src/diagnostics.md                  |  13 ++-
 src/diagnostics/core_diagnostics.jl      |  14 ++++
 src/diagnostics/diagnostic.jl            | 101 +++++++++++++++++------
 src/diagnostics/diagnostics_utils.jl     |  76 +++++++++++++++--
 src/diagnostics/radiation_diagnostics.jl |   8 ++
 src/diagnostics/writers.jl               |  28 ++++---
 src/solver/type_getters.jl               |   4 +-
 8 files changed, 201 insertions(+), 53 deletions(-)

diff --git a/docs/make_diagnostic_table.jl b/docs/make_diagnostic_table.jl
index fe27fbb75cf..ee6d007ae73 100644
--- a/docs/make_diagnostic_table.jl
+++ b/docs/make_diagnostic_table.jl
@@ -15,12 +15,16 @@ open(out_path, "w") do file
 
     write(file, "# Available diagnostic variables\n\n")
 
-    write(file, "| Short name | Long name  | Units | Comments |\n")
-    write(file, "|---|---|---|---|\n")
+    write(
+        file,
+        "| Short name | Long name | Standard name | Units | Comments |\n",
+    )
+    write(file, "|---|---|---|---|---|\n")
 
     for d in values(CA.Diagnostics.ALL_DIAGNOSTICS)
         write(file, "| `$(d.short_name)` ")
-        write(file, "| `$(d.long_name)` ")
+        write(file, "| $(d.long_name) ")
+        write(file, "| `$(d.standard_name)` ")
         write(file, "| $(d.units) ")
         write(file, "| $(d.comments)|\n")
     end
diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 149b5f9566a..01213df87eb 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -107,8 +107,11 @@ More specifically, a `ScheduledDiagnostic` contains the following pieces of data
    `pre_output_hook!` is called with two arguments: the value accumulated during
    the reduction, and the number of times the diagnostic was computed from the
    last time it was output.
-- `output_name`: A descriptive name that can be used by the `output_writer`. If
-   not provided, a default one is generated.
+- `output_short_name`: A descriptive name that can be used by the
+   `output_writer`. If not provided, a default one is generated. It has to be
+   unique.
+- `output_long_name`: A descriptive name that can be used by the `output_writer`
+   as attribute. If not provided, a default one is generated.
 
 To implement operations like the arithmetic average, the `reduction_time_func`
 has to be chosen as `+`, and a `pre_output_hook!` that renormalize `acc` by the
@@ -190,7 +193,11 @@ In `ClimaAtmos`, we follow the convention that:
                 diagnostics by their short name, so the diagnostics defined by
                 `ClimaAtmos` have to have unique `short_name`s.
 
-- `long_name`: Name used to describe the variable in the output file.
+- `long_name`: Name used to describe the variable in the output file as attribute.
+
+- `standard_name`: Standard name, as in
+  [CF
+  conventions](http://cfconventions.org/Data/cf-standard-names/71/build/cf-standard-name-table.html)
 
 - `units`: Physical units of the variable.
 
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index f447e4d9665..45d070e732e 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -52,6 +52,7 @@ error_diagnostic_variable(variable, model::T) where {T} =
 add_diagnostic_variable!(
     short_name = "rhoa",
     long_name = "Air Density",
+    standard_name = "air_density",
     units = "kg m^-3",
     comments = "Density of air",
     compute! = (out, state, cache, time) -> begin
@@ -69,6 +70,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "ua",
     long_name = "Eastward Wind",
+    standard_name = "eastward_wind",
     units = "m s^-1",
     comments = "Eastward (zonal) wind component",
     compute! = (out, state, cache, time) -> begin
@@ -86,6 +88,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "va",
     long_name = "Northward Wind",
+    standard_name = "northward_wind",
     units = "m s^-1",
     comments = "Northward (meridional) wind component",
     compute! = (out, state, cache, time) -> begin
@@ -106,6 +109,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "wa",
     long_name = "Upward Air Velocity",
+    standard_name = "upward_air_velocity",
     units = "m s^-1",
     comments = "Vertical wind component",
     compute! = (out, state, cache, time) -> begin
@@ -123,6 +127,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "ta",
     long_name = "Air Temperature",
+    standard_name = "air_temperature",
     units = "K",
     comments = "Temperature of air",
     compute! = (out, state, cache, time) -> begin
@@ -141,6 +146,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "thetaa",
     long_name = "Air Potential Temperature",
+    standard_name = "air_potential_temperature",
     units = "K",
     comments = "Potential temperature of air",
     compute! = (out, state, cache, time) -> begin
@@ -176,6 +182,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "rv",
     long_name = "Relative Vorticity",
+    standard_name = "relative_vorticity",
     units = "s^-1",
     comments = "Vertical component of relative vorticity",
     compute! = (out, state, cache, time) -> begin
@@ -217,6 +224,7 @@ end
 add_diagnostic_variable!(
     short_name = "hur",
     long_name = "Relative Humidity",
+    standard_name = "relative_humidity",
     units = "",
     comments = "Total amount of water vapor in the air relative to the amount achievable by saturation at the current temperature",
     compute! = compute_hur!,
@@ -247,6 +255,7 @@ end
 add_diagnostic_variable!(
     short_name = "hus",
     long_name = "Specific Humidity",
+    standard_name = "specific_humidity",
     units = "kg kg^-1",
     comments = "Mass of all water phases per mass of air",
     compute! = compute_hus!,
@@ -282,6 +291,7 @@ end
 add_diagnostic_variable!(
     short_name = "hussfc",
     long_name = "Surface Specific Humidity",
+    standard_name = "specific_humidity",
     units = "kg kg^-1",
     comments = "Mass of all water phases per mass of air in the layer infinitely close to the surface",
     compute! = compute_hussfc!,
@@ -293,6 +303,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "ts",
     long_name = "Surface Temperature",
+    standard_name = "surface_temperature",
     units = "K",
     comments = "Temperature of the lower boundary of the atmosphere",
     compute! = (out, state, cache, time) -> begin
@@ -344,6 +355,7 @@ compute_tauv!(out, state, cache, time) =
 add_diagnostic_variable!(
     short_name = "tauu",
     long_name = "Surface Downward Eastward Wind Stress",
+    standard_name = "downward_eastward_stress",
     units = "Pa",
     comments = "Eastward component of the surface drag",
     compute! = compute_tauu!,
@@ -352,6 +364,7 @@ add_diagnostic_variable!(
 add_diagnostic_variable!(
     short_name = "tauv",
     long_name = "Surface Downward Northward Wind Stress",
+    standard_name = "downward_northward_stress",
     units = "Pa",
     comments = "Northward component of the surface drag",
     compute! = compute_tauv!,
@@ -454,6 +467,7 @@ end
 add_diagnostic_variable!(
     short_name = "pr",
     long_name = "Precipitation",
+    standard_name = "precipitation",
     units = "m s^-1",
     comments = "Total precipitation including rain and snow",
     compute! = compute_pr!,
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index abd56a81b1b..7d4f1bf5870 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -70,7 +70,10 @@ Keyword arguments
                 names. Short but descriptive. `ClimaAtmos` follows the CMIP conventions and
                 the diagnostics are identified by the short name.
 
-- `long_name`: Name used to identify the variable in the output files.
+- `long_name`: Name used to describe the variable in the output files.
+
+- `standard_name`: Standard name, as in
+                   http://cfconventions.org/Data/cf-standard-names/71/build/cf-standard-name-table.html
 
 - `units`: Physical units of the variable.
 
@@ -88,6 +91,7 @@ Keyword arguments
 Base.@kwdef struct DiagnosticVariable{T}
     short_name::String
     long_name::String
+    standard_name::String
     units::String
     comments::String
     compute!::T
@@ -101,6 +105,7 @@ const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}()
 
     add_diagnostic_variable!(; short_name,
                                long_name,
+                               standard_name,
                                units,
                                description,
                                compute!)
@@ -122,6 +127,9 @@ Keyword arguments
 
 - `long_name`: Name used to identify the variable in the output files.
 
+- `standard_name`: Standard name, as in
+                   http://cfconventions.org/Data/cf-standard-names/71/build/cf-standard-name-table.html
+
 - `units`: Physical units of the variable.
 
 - `comments`: More verbose explanation of what the variable is, or comments related to how
@@ -139,15 +147,22 @@ Keyword arguments
 function add_diagnostic_variable!(;
     short_name,
     long_name,
+    standard_name = "",
     units,
-    comments,
+    comments = "",
     compute!,
 )
     haskey(ALL_DIAGNOSTICS, short_name) &&
         error("diagnostic $short_name already defined")
 
-    ALL_DIAGNOSTICS[short_name] =
-        DiagnosticVariable(; short_name, long_name, units, comments, compute!)
+    ALL_DIAGNOSTICS[short_name] = DiagnosticVariable(;
+        short_name,
+        long_name,
+        standard_name,
+        units,
+        comments,
+        compute!,
+    )
 end
 
 """
@@ -208,7 +223,8 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
     reduction_space_func::F2
     compute_every::T2
     pre_output_hook!::PO
-    output_name::String
+    output_short_name::String
+    output_long_name::String
 
     """
         ScheduledDiagnosticIterations(; variable::DiagnosticVariable,
@@ -218,7 +234,8 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                                         reduction_space_func = nothing,
                                         compute_every = isa_reduction ? 1 : output_every,
                                         pre_output_hook! = nothing,
-                                        output_name = descriptive_name(self) )
+                                        output_short_name = descriptive_short_name(self),
+                                        output_short_name = descriptive_long_name(self))
 
 
     A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence
@@ -270,10 +287,15 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
                           discarded. An example of `pre_output_hook!` to compute the arithmetic
                           average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
 
-   - `output_name`: A descriptive name for this particular diagnostic. If none is provided,
-                    one will be generated mixing the short name of the variable, the
-                    reduction, and the period of the reduction.
+   - `output_short_name`: A descriptive name for this particular diagnostic. If none is
+                          provided, one will be generated mixing the short name of the
+                          variable, the reduction, and the period of the reduction.
+                          Normally, it has to be unique. In `ClimaAtmos`, we follow the CMIP
+                          conventions for this.
 
+   - `output_long_name`: A descriptive name for this particular diagnostic. If none is
+                         provided, one will be generated mixing the short name of the
+                         variable, the reduction, and the period of the reduction.
 
     """
     function ScheduledDiagnosticIterations(;
@@ -284,7 +306,14 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         reduction_space_func = nothing,
         compute_every = isnothing(reduction_time_func) ? output_every : 1,
         pre_output_hook! = nothing,
-        output_name = get_descriptive_name(
+        output_short_name = get_descriptive_short_name(
+            variable,
+            output_every,
+            reduction_time_func,
+            pre_output_hook!;
+            units_are_seconds = false,
+        ),
+        output_long_name = get_descriptive_long_name(
             variable,
             output_every,
             reduction_time_func,
@@ -296,14 +325,14 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         # We provide an inner constructor to enforce some constraints
 
         output_every % compute_every == 0 || error(
-            "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_name)",
+            "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_short_name)",
         )
 
         isa_reduction = !isnothing(reduction_time_func)
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_name), changing compute_every to match"
+            @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_short_name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -329,7 +358,8 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
             reduction_space_func,
             compute_every,
             pre_output_hook!,
-            output_name,
+            output_short_name,
+            output_long_name,
         )
     end
 end
@@ -343,7 +373,8 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
     reduction_space_func::F2
     compute_every::T2
     pre_output_hook!::PO
-    output_name::String
+    output_short_name::String
+    output_long_name::String
 
     """
         ScheduledDiagnosticTime(; variable::DiagnosticVariable,
@@ -353,7 +384,9 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                                   reduction_space_func = nothing,
                                   compute_every = isa_reduction ? :timestep : output_every,
                                   pre_output_hook! = nothing,
-                                  output_name = descriptive_name(self))
+                                  output_short_name = descriptive_short_name(self),
+                                  output_long_name = descriptive_long_name(self),
+                                  )
 
 
     A `DiagnosticVariable` that has to be computed and output during a simulation with a
@@ -409,9 +442,15 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
                           discarded. An example of `pre_output_hook!` to compute the arithmetic
                           average is `pre_output_hook!(acc, N) = @. acc = acc / N`.
 
-   - `output_name`: A descriptive name for this particular diagnostic. If none is provided,
-                    one will be generated mixing the short name of the variable, the
-                    reduction, and the period of the reduction.
+   - `output_short_name`: A descriptive name for this particular diagnostic. If none is
+                          provided, one will be generated mixing the short name of the
+                          variable, the reduction, and the period of the reduction.
+                          Normally, it has to be unique. In `ClimaAtmos`, we follow the CMIP
+                          conventions for this.
+
+   - `output_long_name`: A descriptive name for this particular diagnostic. If none is
+                         provided, one will be generated mixing the short name of the
+                         variable, the reduction, and the period of the reduction.
     """
     function ScheduledDiagnosticTime(;
         variable::DiagnosticVariable,
@@ -422,7 +461,14 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         compute_every = isnothing(reduction_time_func) ? output_every :
                         :timestep,
         pre_output_hook! = nothing,
-        output_name = get_descriptive_name(
+        output_short_name = get_descriptive_short_name(
+            variable,
+            output_every,
+            reduction_time_func,
+            pre_output_hook!;
+            units_are_seconds = true,
+        ),
+        output_long_name = get_descriptive_long_name(
             variable,
             output_every,
             reduction_time_func,
@@ -437,7 +483,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         # the list of diagnostics
         if !isa(compute_every, Symbol)
             output_every % compute_every == 0 || error(
-                "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_name)",
+                "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_short_name)",
             )
         end
 
@@ -445,7 +491,7 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
 
         # If it is not a reduction, we compute only when we output
         if !isa_reduction && compute_every != output_every
-            @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_name), changing compute_every to match"
+            @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_short_name), changing compute_every to match"
             compute_every = output_every
         end
 
@@ -471,7 +517,8 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
             reduction_space_func,
             compute_every,
             pre_output_hook!,
-            output_name,
+            output_short_name,
+            output_long_name,
         )
     end
 end
@@ -498,10 +545,10 @@ function ScheduledDiagnosticIterations(
     output_every = sd_time.output_every / Δt
 
     isinteger(output_every) || error(
-        "output_every ($(sd_time.output_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
+        "output_every ($(sd_time.output_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_short_name)",
     )
     isinteger(compute_every) || error(
-        "compute_every ($(sd_time.compute_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_name)",
+        "compute_every ($(sd_time.compute_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_short_name)",
     )
 
     ScheduledDiagnosticIterations(;
@@ -512,7 +559,8 @@ function ScheduledDiagnosticIterations(
         sd_time.reduction_space_func,
         compute_every = convert(Int, compute_every),
         sd_time.pre_output_hook!,
-        sd_time.output_name,
+        sd_time.output_short_name,
+        sd_time.output_long_name,
     )
 end
 
@@ -544,7 +592,8 @@ function ScheduledDiagnosticTime(
         sd_time.reduction_space_func,
         compute_every,
         sd_time.pre_output_hook!,
-        sd_time.output_name,
+        sd_time.output_short_name,
+        sd_time.output_long_name,
     )
 end
 
diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl
index 8746535159a..4e3e5204145 100644
--- a/src/diagnostics/diagnostics_utils.jl
+++ b/src/diagnostics/diagnostics_utils.jl
@@ -1,15 +1,16 @@
 # diagnostic_utils.jl
 #
 # This file contains:
-# - get_descriptive_name: to condense ScheduledDiagnostic information into few characters.
+# - get_descriptive_short_name: to condense ScheduledDiagnostic information into few characters.
+# - get_descriptive_long_name: to produce full names that are clearly human-understandable
 
 
 """
-    get_descriptive_name(variable::DiagnosticVariable,
-                         output_every,
-                         reduction_time_func,
-                         pre_output_hook!;
-                         units_are_seconds = true)
+    get_descriptive_short_name(variable::DiagnosticVariable,
+                               output_every,
+                               reduction_time_func,
+                               pre_output_hook!;
+                               units_are_seconds = true)
 
 
 Return a compact, unique-ish, identifier generated from the given information.
@@ -20,7 +21,7 @@ is interpreted as in units of number of iterations.
 This function is useful for filenames and error messages.
 
 """
-function get_descriptive_name(
+function get_descriptive_short_name(
     variable::DiagnosticVariable,
     output_every,
     reduction_time_func,
@@ -64,3 +65,64 @@ function get_descriptive_name(
     end
     return "$(var)_$(suffix)"
 end
+
+"""
+    get_descriptive_long_name(variable::DiagnosticVariable,
+                              output_every,
+                              reduction_time_func,
+                              pre_output_hook!;
+                              units_are_seconds = true)
+
+
+Return a verbose description of the given output variable.
+
+`output_every` is interpreted as in seconds if `units_are_seconds` is `true`. Otherwise, it
+is interpreted as in units of number of iterations.
+
+This function is useful for attributes in output files.
+
+"""
+function get_descriptive_long_name(
+    variable::DiagnosticVariable,
+    output_every,
+    reduction_time_func,
+    pre_output_hook!;
+    units_are_seconds = true,
+)
+    var = "$(variable.long_name)"
+    isa_reduction = !isnothing(reduction_time_func)
+
+    if isa_reduction
+        red = "$(reduction_time_func)"
+
+        # Let's check if we are computing the average. Note that this might slip under the
+        # radar if the user passes their own pre_output_hook!.
+        if reduction_time_func == (+) &&
+           pre_output_hook! == average_pre_output_hook!
+            red = "average"
+        end
+
+        if units_are_seconds
+
+            # Convert period from seconds to days, hours, minutes, seconds
+            period = ""
+
+            days, rem_seconds = divrem(output_every, 24 * 60 * 60)
+            hours, rem_seconds = divrem(rem_seconds, 60 * 60)
+            minutes, seconds = divrem(rem_seconds, 60)
+
+            days > 0 && (period *= "$(days) Day(s)")
+            hours > 0 && (period *= "$(hours) Hour(s)")
+            minutes > 0 && (period *= "$(minutes) Minute(s)")
+            seconds > 0 && (period *= "$(seconds) Second(s)")
+
+            period_str = period * red
+        else
+            period_str = "$(output_every) Iterations"
+        end
+        suffix = "$(red) within $(period_str)"
+    else
+        suffix = "Instantaneous"
+    end
+    return "$(var), $(suffix)"
+end
diff --git a/src/diagnostics/radiation_diagnostics.jl b/src/diagnostics/radiation_diagnostics.jl
index 867a6ce25c3..fcded7d982c 100644
--- a/src/diagnostics/radiation_diagnostics.jl
+++ b/src/diagnostics/radiation_diagnostics.jl
@@ -34,6 +34,7 @@ end
 add_diagnostic_variable!(
     short_name = "rsd",
     long_name = "Downwelling Shortwave Radiation",
+    standard_name = "surface_downwelling_shortwave_flux_in_air",
     units = "W m^-2",
     comments = "Downwelling shortwave radiation",
     compute! = compute_rsd!,
@@ -71,6 +72,7 @@ end
 add_diagnostic_variable!(
     short_name = "rsu",
     long_name = "Upwelling Shortwave Radiation",
+    standard_name = "surface_upwelling_shortwave_flux_in_air",
     units = "W m^-2",
     comments = "Upwelling shortwave radiation",
     compute! = compute_rsu!,
@@ -108,6 +110,7 @@ end
 add_diagnostic_variable!(
     short_name = "rld",
     long_name = "Downwelling Longwave Radiation",
+    standard_name = "surface_downwelling_longwave_flux_in_air",
     units = "W m^-2",
     comments = "Downwelling longwave radiation",
     compute! = compute_rld!,
@@ -145,6 +148,7 @@ end
 add_diagnostic_variable!(
     short_name = "rlu",
     long_name = "Upwelling Longwave Radiation",
+    standard_name = "surface_upwelling_longwave_flux_in_air",
     units = "W m^-2",
     comments = "Upwelling longwave radiation",
     compute! = compute_rlu!,
@@ -182,6 +186,7 @@ end
 add_diagnostic_variable!(
     short_name = "rsdcs",
     long_name = "Downwelling Clear-Sky Shortwave Radiation",
+    standard_name = "surface_downwelling_shortwave_flux_in_air_assuming_clear_sky",
     units = "W m^-2",
     comments = "Downwelling clear sky shortwave radiation",
     compute! = compute_rsdcs!,
@@ -219,6 +224,7 @@ end
 add_diagnostic_variable!(
     short_name = "rsucs",
     long_name = "Upwelling Clear-Sky Shortwave Radiation",
+    standard_name = "surface_upwelling_shortwave_flux_in_air_assuming_clear_sky",
     units = "W m^-2",
     comments = "Upwelling clear sky shortwave radiation",
     compute! = compute_rsucs!,
@@ -256,6 +262,7 @@ end
 add_diagnostic_variable!(
     short_name = "rldcs",
     long_name = "Downwelling Clear-Sky Longwave Radiation",
+    standard_name = "surface_downwelling_longwave_flux_in_air_assuming_clear_sky",
     units = "W m^-2",
     comments = "Downwelling clear sky longwave radiation",
     compute! = compute_rldcs!,
@@ -293,6 +300,7 @@ end
 add_diagnostic_variable!(
     short_name = "rlucs",
     long_name = "Upwelling Clear-Sky Longwave Radiation",
+    standard_name = "surface_upwelling_longwave_flux_in_air_assuming_clear_sky",
     units = "W m^-2",
     comments = "Upwelling clear sky longwave radiation",
     compute! = compute_rlucs!,
diff --git a/src/diagnostics/writers.jl b/src/diagnostics/writers.jl
index 3390f8f82a8..43de0c78257 100644
--- a/src/diagnostics/writers.jl
+++ b/src/diagnostics/writers.jl
@@ -42,21 +42,25 @@ function HDF5Writer()
 
         output_path = joinpath(
             integrator.p.simulation.output_dir,
-            "$(diagnostic.output_name)_$time.h5",
+            "$(diagnostic.output_short_name)_$(time).h5",
         )
 
         hdfwriter = InputOutput.HDF5Writer(output_path, integrator.p.comms_ctx)
-        InputOutput.HDF5.write_attribute(hdfwriter.file, "time", time)
-        InputOutput.HDF5.write_attribute(
-            hdfwriter.file,
-            "long_name",
-            var.long_name,
+        InputOutput.write!(hdfwriter, value, "$(diagnostic.output_short_name)")
+        attributes = Dict(
+            "time" => time,
+            "long_name" => diagnostic.output_long_name,
+            "variable_units" => var.units,
+            "standard_variable_name" => var.standard_name,
         )
-        InputOutput.write!(
-            hdfwriter,
-            Fields.FieldVector(; Symbol(var.short_name) => value),
-            "diagnostics",
+
+        # TODO: Use directly InputOutput functions
+        InputOutput.HDF5.h5writeattr(
+            hdfwriter.file.filename,
+            "fields/$(diagnostic.output_short_name)",
+            attributes,
         )
+
         Base.close(hdfwriter)
         return nothing
     end
@@ -111,7 +115,7 @@ function NetCDFWriter(;
         # well
         output_path = joinpath(
             integrator.p.simulation.output_dir,
-            "$(diagnostic.output_name)_$time.nc",
+            "$(diagnostic.output_short_name)_$time.nc",
         )
 
         vert_domain = axes(field).vertical_topology.mesh.domain
@@ -141,7 +145,7 @@ function NetCDFWriter(;
         NCDatasets.defDim(nc, "lat", num_points_latitude)
         NCDatasets.defDim(nc, "z", num_points_altitude)
 
-        nc.attrib["long_name"] = var.long_name
+        nc.attrib["long_name"] = diagnostic.output_long_name
         nc.attrib["units"] = var.units
         nc.attrib["comments"] = var.comments
 
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 34d81c41818..71f13ba9163 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -674,7 +674,7 @@ function get_diagnostics(parsed_args, atmos_model)
         period_seconds = time_to_seconds(yaml_diag["period"])
 
         if isnothing(name)
-            name = CAD.get_descriptive_name(
+            name = CAD.get_descriptive_short_name(
                 CAD.get_diagnostic_variable(yaml_diag["short_name"]),
                 period_seconds,
                 reduction_time_func,
@@ -697,7 +697,7 @@ function get_diagnostics(parsed_args, atmos_model)
                 reduction_time_func = reduction_time_func,
                 pre_output_hook! = pre_output_hook!,
                 output_writer = writer,
-                output_name = name,
+                output_short_name = name,
             ),
         )
     end

From 55b490038de7fa032c8617057380bc84f3a125ae Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 16:00:55 -0700
Subject: [PATCH 55/73] Reword "get_" functions

---
 docs/src/diagnostics.md                    |  6 ++---
 src/diagnostics/default_diagnostics.jl     | 16 ++++++-------
 src/diagnostics/defaults/moisture_model.jl |  2 +-
 src/diagnostics/diagnostic.jl              |  8 +++----
 src/diagnostics/diagnostics_utils.jl       | 28 +++++++++++-----------
 src/solver/type_getters.jl                 |  2 +-
 6 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index 01213df87eb..d652674be01 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -35,7 +35,7 @@ Currently, only 3D fields on cubed spheres are supported.
 ### From a script
 
 The simplest way to get started with diagnostics is to use the defaults for your
-atmospheric model. `ClimaAtmos` defines a function `get_default_diagnostic`. You
+atmospheric model. `ClimaAtmos` defines a function `default_diagnostic`. You
 can execute this function on an `AtmosModel` or on any of its fields to obtain a
 list of diagnostics ready to be passed to the simulation. So, for example
 
@@ -43,7 +43,7 @@ list of diagnostics ready to be passed to the simulation. So, for example
 
 model = ClimaAtmos.AtmosModel(..., moisture_model = ClimaAtmos.DryModel(), ...)
 
-diagnostics = ClimaAtmos.get_default_diagnostics(model)
+diagnostics = ClimaAtmos.default_diagnostics(model)
 # => List of diagnostics that include the ones specified for the DryModel
 ```
 
@@ -53,7 +53,7 @@ where to save it, and so on (read below for more information on this). You can
 construct your own lists of `ScheduledDiagnostic`s starting from the variables
 defined by `ClimaAtmos`. The `DiagnosticVariable`s in `ClimaAtmos` are
 identified with by the short and unique name, so that you can access them
-directly with the function `get_diagnostic_variable`. One way to do so is by
+directly with the function `diagnostic_variable`. One way to do so is by
 using the provided convenience functions for common operations, e.g., continuing
 the previous example
 
diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index d13d8f306a6..d2ca7a74152 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -5,37 +5,37 @@
 # level interfaces, add them here. Feel free to include extra files.
 
 """
-    get_default_diagnostics(model)
+    default_diagnostics(model)
 
 
 Return a list of `ScheduledDiagnostic`s associated with the given `model`.
 
 """
-function get_default_diagnostics(model::AtmosModel)
+function default_diagnostics(model::AtmosModel)
     # TODO: Probably not the most elegant way to do this...
     defaults = Any[]
 
     for field in fieldnames(AtmosModel)
-        def_model = get_default_diagnostics(getfield(model, field))
+        def_model = default_diagnostics(getfield(model, field))
         append!(defaults, def_model)
     end
 
     return defaults
 end
 
-# Base case: if we call get_default_diagnostics on something that we don't have information
+# 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 get_default_diagnostics return the same type). This is used by
-# get_default_diagnostics(model::AtmosModel), so that we can ignore defaults for submodels
+# 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.
-get_default_diagnostics(_) = []
+default_diagnostics(_) = []
 
 
 """
     produce_common_diagnostic_function(period, reduction)
 
 
-Helper function to define functions like `get_daily_max`.
+Helper function to define functions like `daily_max`.
 """
 function common_diagnostics(
     period,
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
index a58129f9a75..71133265740 100644
--- a/src/diagnostics/defaults/moisture_model.jl
+++ b/src/diagnostics/defaults/moisture_model.jl
@@ -1,5 +1,5 @@
 # FIXME: Gabriele added this as an example. Put something meaningful here!
-function get_default_diagnostics(::DryModel)
+function default_diagnostics(::DryModel)
     return [
         daily_averages("air_density")...,
         ScheduledDiagnosticTime(
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 7d4f1bf5870..312b9676935 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -306,14 +306,14 @@ struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO}
         reduction_space_func = nothing,
         compute_every = isnothing(reduction_time_func) ? output_every : 1,
         pre_output_hook! = nothing,
-        output_short_name = get_descriptive_short_name(
+        output_short_name = descriptive_short_name(
             variable,
             output_every,
             reduction_time_func,
             pre_output_hook!;
             units_are_seconds = false,
         ),
-        output_long_name = get_descriptive_long_name(
+        output_long_name = descriptive_long_name(
             variable,
             output_every,
             reduction_time_func,
@@ -461,14 +461,14 @@ struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO}
         compute_every = isnothing(reduction_time_func) ? output_every :
                         :timestep,
         pre_output_hook! = nothing,
-        output_short_name = get_descriptive_short_name(
+        output_short_name = descriptive_short_name(
             variable,
             output_every,
             reduction_time_func,
             pre_output_hook!;
             units_are_seconds = true,
         ),
-        output_long_name = get_descriptive_long_name(
+        output_long_name = descriptive_long_name(
             variable,
             output_every,
             reduction_time_func,
diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl
index 4e3e5204145..1580c66cbf9 100644
--- a/src/diagnostics/diagnostics_utils.jl
+++ b/src/diagnostics/diagnostics_utils.jl
@@ -1,16 +1,16 @@
 # diagnostic_utils.jl
 #
 # This file contains:
-# - get_descriptive_short_name: to condense ScheduledDiagnostic information into few characters.
-# - get_descriptive_long_name: to produce full names that are clearly human-understandable
+# - descriptive_short_name: to condense ScheduledDiagnostic information into few characters.
+# - descriptive_long_name: to produce full names that are clearly human-understandable
 
 
 """
-    get_descriptive_short_name(variable::DiagnosticVariable,
-                               output_every,
-                               reduction_time_func,
-                               pre_output_hook!;
-                               units_are_seconds = true)
+    descriptive_short_name(variable::DiagnosticVariable,
+                           output_every,
+                           reduction_time_func,
+                           pre_output_hook!;
+                           units_are_seconds = true)
 
 
 Return a compact, unique-ish, identifier generated from the given information.
@@ -21,7 +21,7 @@ is interpreted as in units of number of iterations.
 This function is useful for filenames and error messages.
 
 """
-function get_descriptive_short_name(
+function descriptive_short_name(
     variable::DiagnosticVariable,
     output_every,
     reduction_time_func,
@@ -67,11 +67,11 @@ function get_descriptive_short_name(
 end
 
 """
-    get_descriptive_long_name(variable::DiagnosticVariable,
-                              output_every,
-                              reduction_time_func,
-                              pre_output_hook!;
-                              units_are_seconds = true)
+    descriptive_long_name(variable::DiagnosticVariable,
+                          output_every,
+                          reduction_time_func,
+                          pre_output_hook!;
+                          units_are_seconds = true)
 
 
 Return a verbose description of the given output variable.
@@ -82,7 +82,7 @@ is interpreted as in units of number of iterations.
 This function is useful for attributes in output files.
 
 """
-function get_descriptive_long_name(
+function descriptive_long_name(
     variable::DiagnosticVariable,
     output_every,
     reduction_time_func,
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 71f13ba9163..01ab4309737 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -674,7 +674,7 @@ function get_diagnostics(parsed_args, atmos_model)
         period_seconds = time_to_seconds(yaml_diag["period"])
 
         if isnothing(name)
-            name = CAD.get_descriptive_short_name(
+            name = CAD.descriptive_short_name(
                 CAD.get_diagnostic_variable(yaml_diag["short_name"]),
                 period_seconds,
                 reduction_time_func,

From 8ff2f00fb4b9ddbdbd73d1ebfecef3ba9b897285 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 15:22:29 -0700
Subject: [PATCH 56/73] Add some default diagnostics

---
 src/diagnostics/default_diagnostics.jl     | 62 +++++++++++++++++++++-
 src/diagnostics/defaults/moisture_model.jl | 19 -------
 2 files changed, 61 insertions(+), 20 deletions(-)
 delete mode 100644 src/diagnostics/defaults/moisture_model.jl

diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index d2ca7a74152..ec0ff614aec 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -15,6 +15,8 @@ function default_diagnostics(model::AtmosModel)
     # TODO: Probably not the most elegant way to do this...
     defaults = Any[]
 
+    append!(defaults, core_default_diagnostics())
+
     for field in fieldnames(AtmosModel)
         def_model = default_diagnostics(getfield(model, field))
         append!(defaults, def_model)
@@ -180,4 +182,62 @@ hourly_average(short_names; output_writer = HDF5Writer()) =
     hourly_averages(short_names; output_writer)[1]
 
 # Include all the subdefaults
-include("defaults/moisture_model.jl")
+
+########
+# Core #
+########
+function core_default_diagnostics()
+    core_diagnostics = ["ts", "ta", "thetaa", "pfull", "rhoa", "ua", "va", "wa"]
+
+    return [
+        daily_averages(core_diagnostics...)...,
+        daily_max("ts"),
+        daily_min("ts"),
+    ]
+end
+
+###############
+# Energy form #
+###############
+function default_diagnostics(::TotalEnergy)
+    total_energy_diagnostics = ["hfes"]
+
+    return [daily_averages(total_energy_diagnostics...)...]
+end
+
+
+##################
+# Moisture model #
+##################
+function default_diagnostics(
+    ::T,
+) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
+    moist_diagnostics = ["hur", "hus", "hussfc", "evspsbl"]
+
+    return [daily_averages(moist_diagnostics...)...]
+end
+
+#######################
+# Precipitation model #
+#######################
+function default_diagnostics(::Microphysics0Moment)
+    precip_diagnostics = ["pr"]
+
+    return [daily_averages(precip_diagnostics...)...]
+end
+
+##################
+# Radiation mode #
+##################
+function default_diagnostics(::RRTMGPI.AbstractRRTMGPMode)
+    allsky_diagnostics = ["rsd", "rsu", "rld", "rlu"]
+
+    return [daily_averages(allsky_diagnostics...)...]
+end
+
+
+function default_diagnostics(::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics)
+    clear_diagnostics = ["rsdcs", "rsucs", "rldcs", "rlucs"]
+
+    return [daily_averages(clear_diagnostics...)...]
+end
diff --git a/src/diagnostics/defaults/moisture_model.jl b/src/diagnostics/defaults/moisture_model.jl
deleted file mode 100644
index 71133265740..00000000000
--- a/src/diagnostics/defaults/moisture_model.jl
+++ /dev/null
@@ -1,19 +0,0 @@
-# FIXME: Gabriele added this as an example. Put something meaningful here!
-function default_diagnostics(::DryModel)
-    return [
-        daily_averages("air_density")...,
-        ScheduledDiagnosticTime(
-            variable = get_diagnostic_variable("air_density"),
-            compute_every = :timestep,
-            output_every = 86400, # seconds
-            reduction_time_func = min,
-            output_writer = HDF5Writer(),
-        ),
-        ScheduledDiagnosticIterations(
-            variable = get_diagnostic_variable("air_density"),
-            compute_every = 1,
-            output_every = 1, # iteration
-            output_writer = HDF5Writer(),
-        ),
-    ]
-end

From cbd3783a5ecd3f5624cfaecc858ac97fd3b2f587 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 18 Sep 2023 17:15:00 -0700
Subject: [PATCH 57/73] Add timing information to diagnostics

---
 src/solver/type_getters.jl | 79 ++++++++++++++++++++++----------------
 1 file changed, 46 insertions(+), 33 deletions(-)

diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 01ab4309737..dd4cb6a3872 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -842,8 +842,10 @@ function get_integrator(config::AtmosConfig)
     @info "get_callbacks: $s"
 
     # Initialize diagnostics
-    @info "Initializing diagnostics"
-    diagnostics = get_diagnostics(config.parsed_args, atmos)
+    s = @timed_str begin
+        diagnostics = get_diagnostics(config.parsed_args, atmos)
+    end
+    @info "initializing diagnostics: $s"
 
     # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration,
     # ensuring that there is consistency in the timestep and the periods and translating
@@ -859,12 +861,15 @@ function get_integrator(config::AtmosConfig)
     diagnostic_counters = Dict()
 
     # NOTE: The diagnostics_callbacks are not called at the initial timestep
-    diagnostics_callbacks = CAD.get_callbacks_from_diagnostics(
-        diagnostics_iterations,
-        diagnostic_storage,
-        diagnostic_accumulators,
-        diagnostic_counters,
-    )
+    s = @timed_str begin
+        diagnostics_callbacks = CAD.get_callbacks_from_diagnostics(
+            diagnostics_iterations,
+            diagnostic_storage,
+            diagnostic_accumulators,
+            diagnostic_counters,
+        )
+    end
+    @info "Prepared diagnostic callbacks: $s"
 
     # We need to ensure the precomputed quantities are indeed precomputed
     # TODO: Remove this when we can assume that the precomputed_quantities are in sync with the state
@@ -872,11 +877,14 @@ function get_integrator(config::AtmosConfig)
         (int) -> set_precomputed_quantities!(int.u, int.p, int.t),
     )
 
-    callback = SciMLBase.CallbackSet(
-        callback...,
-        sync_precomputed,
-        diagnostics_callbacks...,
-    )
+    s = @timed_str begin
+        callback = SciMLBase.CallbackSet(
+            callback...,
+            sync_precomputed,
+            diagnostics_callbacks...,
+        )
+    end
+    @info "Prepared SciMLBase.CallbackSet callbacks: $s"
     @info "n_steps_per_cycle_per_cb: $(n_steps_per_cycle_per_cb(callback, simulation.dt))"
     @info "n_steps_per_cycle: $(n_steps_per_cycle(callback, simulation.dt))"
 
@@ -897,29 +905,34 @@ function get_integrator(config::AtmosConfig)
     end
     @info "init integrator: $s"
 
-    for diag in diagnostics_iterations
-        variable = diag.variable
-        try
-            # The first time we call compute! we use its return value. All the subsequent
-            # times (in the callbacks), we will write the result in place
-            diagnostic_storage[diag] = variable.compute!(
-                nothing,
-                integrator.u,
-                integrator.p,
-                integrator.t,
-            )
-            diagnostic_counters[diag] = 1
-            # If it is not a reduction, call the output writer as well
-            if isnothing(diag.reduction_time_func)
-                diag.output_writer(diagnostic_storage[diag], diag, integrator)
-            else
-                # Add to the accumulator
-                diagnostic_accumulators[diag] = copy(diagnostic_storage[diag])
+    s = @timed_str begin
+        for diag in diagnostics_iterations
+            variable = diag.variable
+            try
+                # The first time we call compute! we use its return value. All
+                # the subsequent times (in the callbacks), we will write the
+                # result in place
+                diagnostic_storage[diag] = variable.compute!(
+                    nothing,
+                    integrator.u,
+                    integrator.p,
+                    integrator.t,
+                )
+                diagnostic_counters[diag] = 1
+                # If it is not a reduction, call the output writer as well
+                if isnothing(diag.reduction_time_func)
+                    diag.output_writer(diagnostic_storage[diag], diag, integrator)
+                else
+                    # Add to the accumulator
+                    diagnostic_accumulators[diag] =
+                        copy(diagnostic_storage[diag])
+                end
+            catch e
+                error("Could not compute diagnostic $(variable.long_name): $e")
             end
-        catch e
-            error("Could not compute diagnostic $(variable.long_name): $e")
         end
     end
+    @info "Init diagnostics: $s"
 
     return integrator
 end

From 68fe9fd0bb67eda59a19539c78d826693b57b14e Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Tue, 19 Sep 2023 08:51:23 -0700
Subject: [PATCH 58/73] Reduce inference time in SciMLBase.CallbackSet

The generic constructor for SciMLBase.CallbackSet has to split callbacks
into discrete and continuous. This is not hard, but can introduce
significant latency. However, all the callbacks in ClimaAtmos are
discrete_callbacks, so we directly pass this information to the
constructor
---
 src/solver/type_getters.jl | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index dd4cb6a3872..9ef2548a238 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -877,16 +877,24 @@ function get_integrator(config::AtmosConfig)
         (int) -> set_precomputed_quantities!(int.u, int.p, int.t),
     )
 
+    # The generic constructor for SciMLBase.CallbackSet has to split callbacks into discrete
+    # and continuous. This is not hard, but can introduce significant latency. However, all
+    # the callbacks in ClimaAtmos are discrete_callbacks, so we directly pass this
+    # information to the constructor
+    continuous_callbacks = tuple()
+    discrete_callbacks = (callback...,
+                          sync_precomputed,
+                          diagnostics_callbacks...)
+
     s = @timed_str begin
-        callback = SciMLBase.CallbackSet(
-            callback...,
-            sync_precomputed,
-            diagnostics_callbacks...,
+        all_callbacks = SciMLBase.CallbackSet(
+            continuous_callbacks,
+            discrete_callbacks
         )
     end
     @info "Prepared SciMLBase.CallbackSet callbacks: $s"
-    @info "n_steps_per_cycle_per_cb: $(n_steps_per_cycle_per_cb(callback, simulation.dt))"
-    @info "n_steps_per_cycle: $(n_steps_per_cycle(callback, simulation.dt))"
+    @info "n_steps_per_cycle_per_cb: $(n_steps_per_cycle_per_cb(all_callbacks, simulation.dt))"
+    @info "n_steps_per_cycle: $(n_steps_per_cycle(all_callbacks, simulation.dt))"
 
     tspan = (t_start, simulation.t_end)
     s = @timed_str begin
@@ -896,7 +904,7 @@ function get_integrator(config::AtmosConfig)
             p,
             tspan,
             ode_algo,
-            callback,
+            all_callbacks,
         )
     end
 

From 3c689b7a1d0af0bd4449057e0a71bfa2e1eb9a1c Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 20 Sep 2023 13:51:21 -0700
Subject: [PATCH 59/73] Remove Any[]s

---
 src/diagnostics/default_diagnostics.jl | 25 ++++++++++-------
 src/diagnostics/diagnostic.jl          | 13 ++++-----
 src/solver/type_getters.jl             | 37 ++++++++++++--------------
 3 files changed, 37 insertions(+), 38 deletions(-)

diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl
index ec0ff614aec..8f2f6f54f31 100644
--- a/src/diagnostics/default_diagnostics.jl
+++ b/src/diagnostics/default_diagnostics.jl
@@ -12,17 +12,22 @@ Return a list of `ScheduledDiagnostic`s associated with the given `model`.
 
 """
 function default_diagnostics(model::AtmosModel)
-    # TODO: Probably not the most elegant way to do this...
-    defaults = Any[]
-
-    append!(defaults, core_default_diagnostics())
-
-    for field in fieldnames(AtmosModel)
-        def_model = default_diagnostics(getfield(model, field))
-        append!(defaults, def_model)
-    end
+    # 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)) != [],
+        fieldnames(AtmosModel),
+    )
 
-    return defaults
+    # 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(),
+        map(non_empty_fields) do field
+            default_diagnostics(getfield(model, field))
+        end...,
+    )
 end
 
 # Base case: if we call default_diagnostics on something that we don't have information
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 312b9676935..617dc3c148b 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -676,11 +676,8 @@ function get_callbacks_from_diagnostics(
     # storage is used to pre-allocate memory and to accumulate partial results for those
     # diagnostics that perform reductions.
 
-    callbacks = Any[]
-
-    for diag in diagnostics
+    callback_arrays = map(diagnostics) do diag
         variable = diag.variable
-
         compute_callback =
             integrator -> begin
                 variable.compute!(
@@ -729,8 +726,7 @@ function get_callbacks_from_diagnostics(
         # Here we have skip_first = true. This is important because we are going to manually
         # call all the callbacks so that we can verify that they are meaningful for the
         # model under consideration (and they don't have bugs).
-        push!(
-            callbacks,
+        return [
             call_every_n_steps(
                 compute_callback,
                 diag.compute_every,
@@ -741,8 +737,9 @@ function get_callbacks_from_diagnostics(
                 diag.output_every,
                 skip_first = true,
             ),
-        )
+        ]
     end
 
-    return callbacks
+    # We need to flatten to tuples
+    return vcat(callback_arrays...)
 end
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 9ef2548a238..6ca2879f713 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -609,13 +609,9 @@ end
 
 function get_diagnostics(parsed_args, atmos_model)
 
-    diagnostics =
-        parsed_args["output_default_diagnostics"] ?
-        CAD.get_default_diagnostics(atmos_model) : Any[]
-
-    # We either get the diagnostics section in the YAML file, or we return an empty
-    # dictionary (which will result in an empty list being created by the map below)
-    yaml_diagnostics = get(parsed_args, "diagnostics", Dict())
+    # We either get the diagnostics section in the YAML file, or we return an empty list
+    # (which will result in an empty list being created by the map below)
+    yaml_diagnostics = get(parsed_args, "diagnostics", [])
 
     # ALLOWED_REDUCTIONS is the collection of reductions we support. The keys are the
     # strings that have to be provided in the YAML file. The values are tuples with the
@@ -643,7 +639,7 @@ function get_diagnostics(parsed_args, atmos_model)
         "netcdf" => CAD.NetCDFWriter(),
     )
 
-    for yaml_diag in yaml_diagnostics
+    diagnostics = map(yaml_diagnostics) do yaml_diag
         # Return "nothing" if "reduction_time" is not in the YAML block
         #
         # We also normalize everything to lowercase, so that can accept "max" but
@@ -688,21 +684,22 @@ function get_diagnostics(parsed_args, atmos_model)
             compute_every = :timestep
         end
 
-        push!(
-            diagnostics,
-            CAD.ScheduledDiagnosticTime(
-                variable = CAD.get_diagnostic_variable(yaml_diag["short_name"]),
-                output_every = period_seconds,
-                compute_every = compute_every,
-                reduction_time_func = reduction_time_func,
-                pre_output_hook! = pre_output_hook!,
-                output_writer = writer,
-                output_short_name = name,
-            ),
+        return CAD.ScheduledDiagnosticTime(
+            variable = CAD.get_diagnostic_variable(yaml_diag["short_name"]),
+            output_every = period_seconds,
+            compute_every = compute_every,
+            reduction_time_func = reduction_time_func,
+            pre_output_hook! = pre_output_hook!,
+            output_writer = writer,
+            output_short_name = name,
         )
     end
 
-    return diagnostics
+    if parsed_args["output_default_diagnostics"]
+        return [CAD.default_diagnostics(atmos_model)..., diagnostics...]
+    else
+        return collect(diagnostics)
+    end
 end
 
 function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback)

From 2a400c9f662c3e9107ef8ca9b6e0e252e12c9693 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Wed, 20 Sep 2023 15:04:56 -0700
Subject: [PATCH 60/73] Move to orchestrate_callbacks

The current design leads to exploding compile times. The reason is
currently unknown. To avoid these problems, this new design adds one
single callback called every time step that calls the diagnostic
functions when it is their turn.
---
 docs/src/diagnostics.md           |  4 ++-
 src/callbacks/callback_helpers.jl |  3 +++
 src/diagnostics/Diagnostics.jl    |  4 ++-
 src/diagnostics/diagnostic.jl     | 15 ++---------
 src/solver/type_getters.jl        | 45 ++++++++++++++++++++++---------
 5 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
index d652674be01..b55ad74b690 100644
--- a/docs/src/diagnostics.md
+++ b/docs/src/diagnostics.md
@@ -77,7 +77,9 @@ writers.
 
 Diagnostics are computed and output through callbacks to the main integrator.
 `ClimaAtmos` produces the list of callbacks from a ordered list of
-`ScheduledDiagnostic`s.
+`ScheduledDiagnostic`s. These callbacks are orchestrated by a callback
+`orchestrate_diagnostics` that runs at the end of every step and calls all the
+diagnostic callbacks that are scheduled to be run at that step.
 
 A `ScheduledDiagnostic` is an instruction on how to compute and output a given
 `DiagnosticVariable` (see below), along with specific choices regarding
diff --git a/src/callbacks/callback_helpers.jl b/src/callbacks/callback_helpers.jl
index 3c028e91ea1..7430bf95508 100644
--- a/src/callbacks/callback_helpers.jl
+++ b/src/callbacks/callback_helpers.jl
@@ -97,3 +97,6 @@ function n_steps_per_cycle_per_cb(cbs::SciMLBase.CallbackSet, dt)
         end
     end
 end
+
+n_steps_per_cycle_per_cb_diagnostic(cbs) =
+    [callback_frequency(cb).n for cb in cbs]
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
index 29a02c6d6c1..17f3d4ff5f7 100644
--- a/src/diagnostics/Diagnostics.jl
+++ b/src/diagnostics/Diagnostics.jl
@@ -7,7 +7,9 @@ import ClimaCore.Utilities: half
 import Thermodynamics as TD
 
 import ..AtmosModel
-import ..call_every_n_steps
+import ..AtmosCallback
+import ..EveryNSteps
+
 import ..Parameters as CAP
 
 import ..unit_basis_vector_data
diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 617dc3c148b..51d6068188c 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -723,20 +723,9 @@ function get_callbacks_from_diagnostics(
                 return nothing
             end
 
-        # Here we have skip_first = true. This is important because we are going to manually
-        # call all the callbacks so that we can verify that they are meaningful for the
-        # model under consideration (and they don't have bugs).
         return [
-            call_every_n_steps(
-                compute_callback,
-                diag.compute_every,
-                skip_first = true,
-            ),
-            call_every_n_steps(
-                output_callback,
-                diag.output_every,
-                skip_first = true,
-            ),
+            AtmosCallback(compute_callback, EveryNSteps(diag.compute_every)),
+            AtmosCallback(output_callback, EveryNSteps(diag.output_every)),
         ]
     end
 
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 6ca2879f713..c6079c8ef78 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -852,14 +852,14 @@ function get_integrator(config::AtmosConfig)
     ]
 
     # For diagnostics that perform reductions, the storage is used for the values computed
-    # at each call. Reductions also save the accumulated value in in diagnostic_accumulators.
+    # at each call. Reductions also save the accumulated value in diagnostic_accumulators.
     diagnostic_storage = Dict()
     diagnostic_accumulators = Dict()
     diagnostic_counters = Dict()
 
     # NOTE: The diagnostics_callbacks are not called at the initial timestep
     s = @timed_str begin
-        diagnostics_callbacks = CAD.get_callbacks_from_diagnostics(
+        diagnostics_functions = CAD.get_callbacks_from_diagnostics(
             diagnostics_iterations,
             diagnostic_storage,
             diagnostic_accumulators,
@@ -868,8 +868,27 @@ function get_integrator(config::AtmosConfig)
     end
     @info "Prepared diagnostic callbacks: $s"
 
+    # It would be nice to just pass the callbacks to the integrator. However, this leads to
+    # a significant increase in compile time for reasons that are not known. For this
+    # reason, we only add one callback to the integrator, and this function takes care of
+    # executing the other callbacks. This single function is orchestrate_diagnostics
+
+    function orchestrate_diagnostics(integrator)
+        diagnostics_to_be_run =
+            filter(d -> integrator.step % d.cbf.n == 0, diagnostics_functions)
+
+        for diag_func in diagnostics_to_be_run
+            diag_func.f!(integrator)
+        end
+    end
+
+    diagnostic_callbacks =
+        call_every_n_steps(orchestrate_diagnostics, skip_first = true)
+
     # We need to ensure the precomputed quantities are indeed precomputed
-    # TODO: Remove this when we can assume that the precomputed_quantities are in sync with the state
+
+    # TODO: Remove this when we can assume that the precomputed_quantities are in sync with
+    # the state
     sync_precomputed = call_every_n_steps(
         (int) -> set_precomputed_quantities!(int.u, int.p, int.t),
     )
@@ -879,19 +898,21 @@ function get_integrator(config::AtmosConfig)
     # the callbacks in ClimaAtmos are discrete_callbacks, so we directly pass this
     # information to the constructor
     continuous_callbacks = tuple()
-    discrete_callbacks = (callback...,
-                          sync_precomputed,
-                          diagnostics_callbacks...)
+    discrete_callbacks = (callback..., sync_precomputed, diagnostic_callbacks)
 
     s = @timed_str begin
-        all_callbacks = SciMLBase.CallbackSet(
-            continuous_callbacks,
-            discrete_callbacks
-        )
+        all_callbacks =
+            SciMLBase.CallbackSet(continuous_callbacks, discrete_callbacks)
     end
     @info "Prepared SciMLBase.CallbackSet callbacks: $s"
-    @info "n_steps_per_cycle_per_cb: $(n_steps_per_cycle_per_cb(all_callbacks, simulation.dt))"
-    @info "n_steps_per_cycle: $(n_steps_per_cycle(all_callbacks, simulation.dt))"
+    steps_cycle_non_diag =
+        n_steps_per_cycle_per_cb(all_callbacks, simulation.dt)
+    steps_cycle_diag =
+        n_steps_per_cycle_per_cb_diagnostic(diagnostics_functions)
+    steps_cycle = lcm([steps_cycle_non_diag..., steps_cycle_diag...])
+    @info "n_steps_per_cycle_per_cb (non diagnostics): $steps_cycle_non_diag"
+    @info "n_steps_per_cycle_per_cb_diagnostic: $steps_cycle_diag"
+    @info "n_steps_per_cycle (non diagnostics): $steps_cycle"
 
     tspan = (t_start, simulation.t_end)
     s = @timed_str begin

From b913c8027b6572d64f62335c28b45f9a21c40ba4 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 21 Sep 2023 10:22:51 -0700
Subject: [PATCH 61/73] Make timestep evenly divide day in CI

---
 config/default_configs/default_config.yml                       | 2 +-
 config/model_configs/sphere_baroclinic_wave_rhoe.yml            | 2 +-
 config/model_configs/sphere_held_suarez_rhoe_hightop.yml        | 2 +-
 config/model_configs/sphere_held_suarez_rhotheta.yml            | 2 +-
 ...upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml | 2 +-
 config/perf_configs/gpu_baroclinic_wave_rhoe.yml                | 2 +-
 config/perf_configs/gpu_held_suarez_rhoe_hightop.yml            | 2 +-
 regression_tests/ref_counter.jl                                 | 2 +-
 8 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
index d7c8a435e73..20ec70dc9b5 100644
--- a/config/default_configs/default_config.yml
+++ b/config/default_configs/default_config.yml
@@ -251,4 +251,4 @@ log_params:
   value: false
 output_default_diagnostics:
   help: "Output the default diagnostics associated to the selected atmospheric model"
-  value: true
+  value: false
diff --git a/config/model_configs/sphere_baroclinic_wave_rhoe.yml b/config/model_configs/sphere_baroclinic_wave_rhoe.yml
index 1a0176c8df7..8febe78d17e 100644
--- a/config/model_configs/sphere_baroclinic_wave_rhoe.yml
+++ b/config/model_configs/sphere_baroclinic_wave_rhoe.yml
@@ -1,6 +1,6 @@
 dt_save_to_disk: "2days"
 regression_test: true
 initial_condition: "DryBaroclinicWave"
-dt: "580secs"
+dt: "400secs"
 t_end: "10days"
 job_id: "sphere_baroclinic_wave_rhoe"
diff --git a/config/model_configs/sphere_held_suarez_rhoe_hightop.yml b/config/model_configs/sphere_held_suarez_rhoe_hightop.yml
index 45f060b6f2f..89e86b61cdb 100644
--- a/config/model_configs/sphere_held_suarez_rhoe_hightop.yml
+++ b/config/model_configs/sphere_held_suarez_rhoe_hightop.yml
@@ -3,7 +3,7 @@ dt_save_to_disk: "4days"
 regression_test: true
 t_end: "8days"
 forcing: "held_suarez"
-dt: "500secs"
+dt: "400secs"
 z_elem: 25
 job_id: "sphere_held_suarez_rhoe_hightop"
 z_max: 45000.0
diff --git a/config/model_configs/sphere_held_suarez_rhotheta.yml b/config/model_configs/sphere_held_suarez_rhotheta.yml
index e80b5c63cc5..2223e1ee8cd 100644
--- a/config/model_configs/sphere_held_suarez_rhotheta.yml
+++ b/config/model_configs/sphere_held_suarez_rhotheta.yml
@@ -3,5 +3,5 @@ dt_save_to_disk: "10days"
 regression_test: true
 t_end: "20days"
 forcing: "held_suarez"
-dt: "500secs"
+dt: "400secs"
 job_id: "sphere_held_suarez_rhotheta"
diff --git a/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml b/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml
index df7f22f94eb..aa077d17487 100644
--- a/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml
+++ b/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml
@@ -1,7 +1,7 @@
 dt_save_to_disk: "5days"
 initial_condition: "MoistBaroclinicWave"
 max_newton_iters_ode: 4
-dt: "500secs"
+dt: "400secs"
 tracer_upwinding: zalesak
 t_end: "5days"
 ode_algo: "SSP333"
diff --git a/config/perf_configs/gpu_baroclinic_wave_rhoe.yml b/config/perf_configs/gpu_baroclinic_wave_rhoe.yml
index 90fdecc5e93..99349ca7bb8 100644
--- a/config/perf_configs/gpu_baroclinic_wave_rhoe.yml
+++ b/config/perf_configs/gpu_baroclinic_wave_rhoe.yml
@@ -1,5 +1,5 @@
 job_id: "gpu_baroclinic_wave_rhoe" 
-dt: "580secs" 
+dt: "400secs"
 t_end: "10days" 
 dt_save_to_disk: "2days" 
 initial_condition: "DryBaroclinicWave" 
diff --git a/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml b/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml
index 5eac11e6c55..b785b5d01fd 100644
--- a/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml
+++ b/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml
@@ -3,6 +3,6 @@ z_elem: 25
 dz_bottom: 300 
 forcing: "held_suarez" 
 job_id: "gpu_held_suarez_rhoe_hightop" 
-dt: "500secs" 
+dt: "400secs"
 t_end: "8days" 
 dt_save_to_disk: "4days"
diff --git a/regression_tests/ref_counter.jl b/regression_tests/ref_counter.jl
index a949a93dfcc..b0d73241cad 100644
--- a/regression_tests/ref_counter.jl
+++ b/regression_tests/ref_counter.jl
@@ -1 +1 @@
-128
+129

From 057099a5d502d0b56168fb03c58b62e7cfa68bc6 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 21 Sep 2023 10:59:58 -0700
Subject: [PATCH 62/73] Update link to diagnostic table

---
 src/diagnostics/diagnostic.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 51d6068188c..882fee3c6ef 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -61,7 +61,7 @@ responsibility of the `output_writer` to follow the conventions about the meanin
 metadata and their use.
 
 In `ClimaAtmos`, we roughly follow the naming conventions listed in this file:
-https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
+https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb
 
 Keyword arguments
 =================
@@ -115,7 +115,7 @@ Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates th
 `ClimaAtmos.ALL_DIAGNOSTICS`).
 
 If possible, please follow the naming scheme outline in
-https://docs.google.com/spreadsheets/d/1qUauozwXkq7r1g-L4ALMIkCNINIhhCPx
+https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb
 
 Keyword arguments
 =================

From c836a5b6e4a028840b07daab682eb827053c7f3e Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Thu, 21 Sep 2023 11:06:43 -0700
Subject: [PATCH 63/73] Add diagnostics to one CI model

---
 .../sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml        | 1 +
 1 file changed, 1 insertion(+)

diff --git a/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml b/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml
index 466b9944dfb..d758db4bee2 100644
--- a/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml
+++ b/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml
@@ -13,3 +13,4 @@ viscous_sponge: true
 job_id: "sphere_held_suarez_rhoe_equilmoist_hightop_sponge"
 moist: "equil"
 toml: [toml/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.toml]
+output_default_diagnostics: true

From 1f6862ac05640937ebb6e8e21c9e2362a7f1531c Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 22 Sep 2023 09:47:14 -0700
Subject: [PATCH 64/73] Add diagnostic flamegraph

---
 .buildkite/pipeline.yml                   |  8 ++++++++
 config/perf_configs/flame/diagnostics.yml | 11 +++++++++++
 perf/flame.jl                             | 11 ++++++-----
 3 files changed, 25 insertions(+), 5 deletions(-)
 create mode 100644 config/perf_configs/flame/diagnostics.yml

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 62ad092398a..245e9cf433b 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -1023,6 +1023,14 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: ":fire: Flame graph: perf target (diagnostics)"
+        command: >
+          julia --color=yes --project=perf perf/flame.jl
+          --config_file $PERF_CONFIG_PATH/flame/diagnostics.yml
+        artifact_paths: "flame_perf_diagnostics/*"
+        agents:
+          slurm_mem: 20GB
+
       # Inference
       - label: ":rocket: JET n-failures (inference)"
         command: >
diff --git a/config/perf_configs/flame/diagnostics.yml b/config/perf_configs/flame/diagnostics.yml
new file mode 100644
index 00000000000..5a80e59bdfb
--- /dev/null
+++ b/config/perf_configs/flame/diagnostics.yml
@@ -0,0 +1,11 @@
+job_id: "flame_perf_diagnostics"
+diagnostics:
+  - short_name: ua
+    period: 1secs
+    reduction_time: average
+  - short_name: va
+    period: 1secs
+    reduction_time: max
+  - short_name: ta
+    period: 1secs
+    reduction_time: max
diff --git a/perf/flame.jl b/perf/flame.jl
index a2490677a3a..9fe702cccce 100644
--- a/perf/flame.jl
+++ b/perf/flame.jl
@@ -54,11 +54,12 @@ allocs = @allocated SciMLBase.step!(integrator)
 @info "`allocs ($job_id)`: $(allocs)"
 
 allocs_limit = Dict()
-allocs_limit["flame_perf_target"] = 7968
-allocs_limit["flame_perf_target_tracers"] = 207600
-allocs_limit["flame_perf_target_edmfx"] = 274832
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 748176
-allocs_limit["flame_perf_target_edmf"] = 12031972048
+allocs_limit["flame_perf_target"] = 12864
+allocs_limit["flame_perf_target_tracers"] = 212496
+allocs_limit["flame_perf_target_edmfx"] = 304064
+allocs_limit["flame_perf_diagnostics"] = 3023192
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 754848
+allocs_limit["flame_perf_target_edmf"] = 12459299664
 allocs_limit["flame_perf_target_threaded"] = 6175664
 allocs_limit["flame_perf_target_callbacks"] = 45111592
 allocs_limit["flame_perf_gw"] = 4911463328

From bbeaf893ff0a6db395cf9437280f6c3d521f26c5 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 25 Sep 2023 08:42:36 -0700
Subject: [PATCH 65/73] Move import of Diagnostics

---
 src/ClimaAtmos.jl          | 1 +
 src/solver/type_getters.jl | 1 -
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
index c822fe3b92c..1f0c7fcac38 100644
--- a/src/ClimaAtmos.jl
+++ b/src/ClimaAtmos.jl
@@ -111,6 +111,7 @@ include(joinpath("prognostic_equations", "limited_tendencies.jl"))
 include(joinpath("callbacks", "callbacks.jl"))
 
 include(joinpath("diagnostics", "Diagnostics.jl"))
+import .Diagnostics as CAD
 
 include(joinpath("solver", "model_getters.jl")) # high-level (using parsed_args) model getters
 include(joinpath("solver", "type_getters.jl"))
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index c6079c8ef78..186b467d161 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -6,7 +6,6 @@ using Interpolations
 import ClimaCore: InputOutput, Meshes, Spaces
 import ClimaAtmos.RRTMGPInterface as RRTMGPI
 import ClimaAtmos as CA
-import .Diagnostics as CAD
 import LinearAlgebra
 import ClimaCore.Fields
 import OrdinaryDiffEq as ODE

From b78910ff54c3fa7bec35f56cc68d23865d31a638 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Mon, 25 Sep 2023 08:44:12 -0700
Subject: [PATCH 66/73] Add return to accumulate to help inference

---
 src/diagnostics/diagnostic.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl
index 882fee3c6ef..cd83ee7f7d9 100644
--- a/src/diagnostics/diagnostic.jl
+++ b/src/diagnostics/diagnostic.jl
@@ -638,6 +638,7 @@ accumulate!(_, _, reduction_time_func::Nothing) = nothing
 # When we have a reduction, apply it between the accumulated value one
 function accumulate!(diag_accumulator, diag_storage, reduction_time_func)
     diag_accumulator .= reduction_time_func.(diag_accumulator, diag_storage)
+    return nothing
 end
 
 """

From 5b12af4a73579ccace501c9e084757e40cacb02b Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Wed, 20 Sep 2023 16:06:48 -0700
Subject: [PATCH 67/73] add explicit vertical diffusion job

---
 .buildkite/pipeline.yml                            | 14 ++++++++++++++
 ...re_baroclinic_wave_rhoe_equilmoist_expvdiff.yml | 11 +++++++++++
 ...e_baroclinic_wave_rhoe_equilmoist_expvdiff.toml | 10 ++++++++++
 3 files changed, 35 insertions(+)
 create mode 100644 config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml
 create mode 100644 toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 62ad092398a..46c589363bd 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -246,6 +246,20 @@ steps:
           --job_id sphere_baroclinic_wave_rhoe_equilmoist
           --out_dir sphere_baroclinic_wave_rhoe_equilmoist
         artifact_paths: "sphere_baroclinic_wave_rhoe_equilmoist/*"
+      
+      - label: ":computer: no lim ARS baroclinic wave (ρe) equilmoist explicit vertdiff"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml
+
+          julia --color=yes --project=examples post_processing/remap/remap_pipeline.jl
+          --data_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff
+          --out_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff
+
+          julia --color=yes --project=examples post_processing/plot/plot_pipeline.jl
+          --nc_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff
+          --fig_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff --case_name aquaplanet
+        artifact_paths: "sphere_baroclinic_wave_rhoe_equilmoist_expvdiff/*"
 
       - label: ":computer: SSP zalesak tracer & energy upwind baroclinic wave (ρe_tot) equilmoist"
         command: >
diff --git a/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml
new file mode 100644
index 00000000000..45ca42e3b95
--- /dev/null
+++ b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml
@@ -0,0 +1,11 @@
+precip_model: "0M"
+vert_diff: "true"
+z_elem: 20
+dz_bottom: 100
+dt_save_to_disk: "12hours"
+initial_condition: "MoistBaroclinicWave"
+dt: "40secs"
+t_end: "12hours"
+job_id: "sphere_baroclinic_wave_rhoe_equilmoist_expvdiff"
+moist: "equil"
+toml: [toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml]
diff --git a/toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml b/toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml
new file mode 100644
index 00000000000..91de57be7f7
--- /dev/null
+++ b/toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml
@@ -0,0 +1,10 @@
+[C_H]
+alias = "C_H"
+value = 0.0
+type = "float"
+
+[C_E]
+alias = "C_E"
+value = 1
+type = "float"
+

From c9fa3d45d3182c9e15d61bdf82c2d37d31c70f7c Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Mon, 25 Sep 2023 17:45:23 -0700
Subject: [PATCH 68/73] reduce t_end in edmfx_trmm_box

---
 config/model_configs/edmfx_trmm_box.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/config/model_configs/edmfx_trmm_box.yml b/config/model_configs/edmfx_trmm_box.yml
index e9f1012fb32..62bbd4cfdf5 100644
--- a/config/model_configs/edmfx_trmm_box.yml
+++ b/config/model_configs/edmfx_trmm_box.yml
@@ -23,7 +23,7 @@ y_elem: 2
 z_elem: 82
 z_stretch: false
 dt: 1secs
-t_end: 6hours
+t_end: 2hours
 dt_save_to_disk: 10mins
 FLOAT_TYPE: "Float64"
 toml: [toml/edmfx_box.toml]

From bab10c9b243a213352ad9c909e5253b2622d472d Mon Sep 17 00:00:00 2001
From: Anna Jaruga <ajaruga@caltech.edu>
Date: Mon, 25 Sep 2023 11:22:31 -0700
Subject: [PATCH 69/73] Set the default adv test params, fix init adv test
 condition, add rho to plots

---
 config/model_configs/edmfx_adv_test_box.yml  | 6 ++++--
 perf/flame.jl                                | 2 +-
 post_processing/contours_and_profiles.jl     | 1 +
 src/callbacks/callbacks.jl                   | 1 +
 src/initial_conditions/initial_conditions.jl | 4 ++--
 toml/edmfx_box_advection.toml                | 2 +-
 6 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/config/model_configs/edmfx_adv_test_box.yml b/config/model_configs/edmfx_adv_test_box.yml
index 9a54b0105be..5ec3535702d 100644
--- a/config/model_configs/edmfx_adv_test_box.yml
+++ b/config/model_configs/edmfx_adv_test_box.yml
@@ -10,12 +10,14 @@ hyperdiff: "true"
 kappa_4: 1e8
 x_max: 1e4
 y_max: 1e4
-z_max: 3e4
+z_max: 5.5e4
 x_elem: 2
 y_elem: 2
-z_elem: 45
+z_elem: 63
 dz_bottom: 30.0
+dz_top: 3000.0
 dt: "10secs"
 t_end: "3600secs"
 dt_save_to_disk: "100secs"
+FLOAT_TYPE: "Float64"
 toml: [toml/edmfx_box_advection.toml]
diff --git a/perf/flame.jl b/perf/flame.jl
index 9fe702cccce..786aef4273d 100644
--- a/perf/flame.jl
+++ b/perf/flame.jl
@@ -61,7 +61,7 @@ allocs_limit["flame_perf_diagnostics"] = 3023192
 allocs_limit["flame_perf_target_diagnostic_edmfx"] = 754848
 allocs_limit["flame_perf_target_edmf"] = 12459299664
 allocs_limit["flame_perf_target_threaded"] = 6175664
-allocs_limit["flame_perf_target_callbacks"] = 45111592
+allocs_limit["flame_perf_target_callbacks"] = 46413904
 allocs_limit["flame_perf_gw"] = 4911463328
 
 if allocs < allocs_limit[job_id] * buffer
diff --git a/post_processing/contours_and_profiles.jl b/post_processing/contours_and_profiles.jl
index 080c6c59f3f..7b01ddec4a8 100644
--- a/post_processing/contours_and_profiles.jl
+++ b/post_processing/contours_and_profiles.jl
@@ -429,6 +429,7 @@ function contours_and_profiles(output_dir, ref_job_id = nothing)
         (all_categories, :w_velocity),
         (all_categories, :buoyancy),
         (all_categories, :specific_enthalpy),
+        (all_categories, :density),
     )
     if has_moisture
         profile_variables = (
diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
index f07039909a1..ff13a84e647 100644
--- a/src/callbacks/callbacks.jl
+++ b/src/callbacks/callbacks.jl
@@ -219,6 +219,7 @@ function common_diagnostics(p, ᶜu, ᶜts)
         potential_temperature = TD.dry_pottemp.(thermo_params, ᶜts),
         specific_enthalpy = TD.specific_enthalpy.(thermo_params, ᶜts),
         buoyancy = CAP.grav(p.params) .* (p.ᶜρ_ref .- ᶜρ) ./ ᶜρ,
+        density = TD.air_density.(thermo_params, ᶜts),
     )
     if !(p.atmos.moisture_model isa DryModel)
         diagnostics = (;
diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl
index 38762a61a9d..ca71174e567 100644
--- a/src/initial_conditions/initial_conditions.jl
+++ b/src/initial_conditions/initial_conditions.jl
@@ -470,10 +470,10 @@ Base.@kwdef struct MoistAdiabaticProfileEDMFX <: InitialCondition
 end
 
 draft_area(::Type{FT}) where {FT} =
-    z -> FT(0.5) * exp(-(z - FT(4000.0))^2 / 2 / FT(1000.0)^2)
+    z -> z < 0.7e4 ? FT(0.5) * exp(-(z - FT(4e3))^2 / 2 / FT(1e3)^2) : FT(0)
 
 edmfx_q_tot(::Type{FT}) where {FT} =
-    z -> FT(0.001) * exp(-(z - FT(4000.0))^2 / 2 / FT(1000.0)^2)
+    z -> z < 0.7e4 ? FT(1e-3) * exp(-(z - FT(4e3))^2 / 2 / FT(1e3)^2) : FT(0)
 
 function (initial_condition::MoistAdiabaticProfileEDMFX)(params)
     (; perturb) = initial_condition
diff --git a/toml/edmfx_box_advection.toml b/toml/edmfx_box_advection.toml
index d39ecbfaa8f..fbffbcba315 100644
--- a/toml/edmfx_box_advection.toml
+++ b/toml/edmfx_box_advection.toml
@@ -5,6 +5,6 @@ type = "float"
 
 [EDMF_min_area]
 alias = "min_area"
-value = 1.0e-2
+value = 1.0e-3
 type = "float"
 description = "Minimum area fraction per updraft. Parameter not described in the literature."

From 2e24cfac92bfac2b5a3389902667db06d06b384f Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Mon, 11 Sep 2023 12:25:48 -0700
Subject: [PATCH 70/73] add diffusion flux for momentum

---
 perf/flame.jl                                 |  4 +-
 .../diagnostic_edmf_precomputed_quantities.jl |  2 +-
 src/cache/edmf_precomputed_quantities.jl      |  2 +-
 src/cache/temporary_quantities.jl             |  4 ++
 src/prognostic_equations/edmfx_sgs_flux.jl    | 48 +++++++++++--------
 src/utils/utilities.jl                        | 26 +++++++++-
 6 files changed, 60 insertions(+), 26 deletions(-)

diff --git a/perf/flame.jl b/perf/flame.jl
index 786aef4273d..80ab8d826be 100644
--- a/perf/flame.jl
+++ b/perf/flame.jl
@@ -57,8 +57,8 @@ allocs_limit = Dict()
 allocs_limit["flame_perf_target"] = 12864
 allocs_limit["flame_perf_target_tracers"] = 212496
 allocs_limit["flame_perf_target_edmfx"] = 304064
-allocs_limit["flame_perf_diagnostics"] = 3023192
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 754848
+allocs_limit["flame_perf_diagnostics"] = 3024344
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 762784
 allocs_limit["flame_perf_target_edmf"] = 12459299664
 allocs_limit["flame_perf_target_threaded"] = 6175664
 allocs_limit["flame_perf_target_callbacks"] = 46413904
diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
index a496a80bc0f..80236a82b43 100644
--- a/src/cache/diagnostic_edmf_precomputed_quantities.jl
+++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -681,7 +681,7 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
     ᶠu⁰ = p.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
     ᶜstrain_rate = p.ᶜtemp_UVWxUVW
-    compute_strain_rate!(ᶜstrain_rate, ᶠu⁰)
+    compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.ᶜtemp_scalar
diff --git a/src/cache/edmf_precomputed_quantities.jl b/src/cache/edmf_precomputed_quantities.jl
index b4a517089a4..2276fff43bb 100644
--- a/src/cache/edmf_precomputed_quantities.jl
+++ b/src/cache/edmf_precomputed_quantities.jl
@@ -230,7 +230,7 @@ function set_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
     ᶠu⁰ = p.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
     ᶜstrain_rate = p.ᶜtemp_UVWxUVW
-    compute_strain_rate!(ᶜstrain_rate, ᶠu⁰)
+    compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.ᶜtemp_scalar
diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl
index c666866b903..e43a05acda0 100644
--- a/src/cache/temporary_quantities.jl
+++ b/src/cache/temporary_quantities.jl
@@ -54,6 +54,10 @@ function temporary_quantities(atmos, center_space, face_space)
             typeof(UVW(FT(0), FT(0), FT(0)) * UVW(FT(0), FT(0), FT(0))'),
             center_space,
         ), # ᶜstrain_rate
+        ᶠtemp_UVWxUVW = Fields.Field(
+            typeof(UVW(FT(0), FT(0), FT(0)) * UVW(FT(0), FT(0), FT(0))'),
+            face_space,
+        ), # ᶠstrain_rate
         # TODO: Remove this hack
         sfc_temp_C3 = Fields.Field(C3{FT}, Spaces.level(face_space, half)), # ρ_flux_χ
     )
diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl
index 5a9709ef9e1..355205f52f8 100644
--- a/src/prognostic_equations/edmfx_sgs_flux.jl
+++ b/src/prognostic_equations/edmfx_sgs_flux.jl
@@ -11,14 +11,14 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX)
     (; edmfx_upwinding, sfc_conditions) = p
     (; ᶠu³, ᶜh_tot, ᶜspecific) = p
     (; ᶠu³ʲs, ᶜh_totʲs, ᶜspecificʲs) = p
-    (; ᶜρa⁰, ᶠu³⁰, ᶜh_tot⁰, ᶜspecific⁰) = p
+    (; ᶜρa⁰, ᶠu³⁰, ᶜu⁰, ᶜh_tot⁰, ᶜspecific⁰) = p
     (; ᶜK_u, ᶜK_h) = p
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
     ᶠgradᵥ = Operators.GradientC2F()
 
     if p.atmos.edmfx_sgs_flux
-        # mass flux
+        # energy mass flux
         ᶠu³_diff_colidx = p.ᶠtemp_CT3[colidx]
         ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.ᶜtemp_scalar[colidx]
         for j in 1:n
@@ -46,7 +46,7 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX)
             edmfx_upwinding,
         )
 
-        # diffusive flux
+        # energy diffusive flux
         ᶠρaK_h = p.ᶠtemp_scalar
         @. ᶠρaK_h[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_h[colidx])
 
@@ -58,7 +58,7 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX)
             ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot⁰[colidx])))
 
         if !(p.atmos.moisture_model isa DryModel)
-            # mass flux
+            # specific humidity mass flux
             for j in 1:n
                 @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
                 @. ᶜq_tot_diff_colidx =
@@ -86,7 +86,7 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX)
                 edmfx_upwinding,
             )
 
-            # diffusive flux
+            # specific humidity diffusive flux
             ᶜdivᵥ_ρq_tot = Operators.DivergenceF2C(
                 top = Operators.SetValue(C3(FT(0))),
                 bottom = Operators.SetValue(
@@ -98,19 +98,23 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX)
             )
         end
 
-        # diffusive flux
+        # momentum diffusive flux
         ᶠρaK_u = p.ᶠtemp_scalar
         @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx])
+        ᶠstrain_rate = p.ᶠtemp_UVWxUVW
+        compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu⁰[colidx])
+        @. Yₜ.c.uₕ[colidx] -= C12(
+            ᶜdivᵥ(-(2 * ᶠρaK_u[colidx] * ᶠstrain_rate[colidx])) / Y.c.ρ[colidx],
+        )
+        # apply boundary condition for momentum flux
         ᶜdivᵥ_uₕ = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(FT(0)) ⊗ C12(FT(0), FT(0))),
             bottom = Operators.SetValue(sfc_conditions.ρ_flux_uₕ[colidx]),
         )
         @. Yₜ.c.uₕ[colidx] -=
-            ᶜdivᵥ_uₕ(-(ᶠρaK_u[colidx] * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx]
+            ᶜdivᵥ_uₕ(-(FT(0) * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx]
     end
 
-    # TODO: Add momentum mass flux
-
     # TODO: Add tracer flux
 
     return nothing
@@ -128,15 +132,15 @@ function edmfx_sgs_flux_tendency!(
     FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
     (; edmfx_upwinding, sfc_conditions) = p
-    (; ᶠu³, ᶜh_tot, ᶜspecific) = p
+    (; ᶠu³, ᶜu, ᶜh_tot, ᶜspecific) = p
     (; ᶜρaʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜq_totʲs) = p
-    (; ᶜh_tot⁰, ᶜK_u, ᶜK_h) = p
+    (; ᶜK_u, ᶜK_h) = p
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
     ᶠgradᵥ = Operators.GradientC2F()
 
     if p.atmos.edmfx_sgs_flux
-        # mass flux
+        # energy mass flux
         # TODO: check if there is contribution from the environment
         ᶠu³_diff_colidx = p.ᶠtemp_CT3[colidx]
         ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.ᶜtemp_scalar[colidx]
@@ -154,7 +158,7 @@ function edmfx_sgs_flux_tendency!(
             )
         end
 
-        # diffusive flux
+        # energy diffusive flux
         ᶠρaK_h = p.ᶠtemp_scalar
         @. ᶠρaK_h[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_h[colidx])
 
@@ -163,10 +167,10 @@ function edmfx_sgs_flux_tendency!(
             bottom = Operators.SetValue(sfc_conditions.ρ_flux_h_tot[colidx]),
         )
         @. Yₜ.c.ρe_tot[colidx] -=
-            ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot⁰[colidx])))
+            ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot[colidx])))
 
         if !(p.atmos.moisture_model isa DryModel)
-            # mass flux
+            # specific humidity mass flux
             for j in 1:n
                 @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
                 @. ᶜq_tot_diff_colidx =
@@ -182,7 +186,7 @@ function edmfx_sgs_flux_tendency!(
                 )
             end
 
-            # diffusive flux
+            # specific humidity diffusive flux
             ᶜdivᵥ_ρq_tot = Operators.DivergenceF2C(
                 top = Operators.SetValue(C3(FT(0))),
                 bottom = Operators.SetValue(
@@ -194,19 +198,23 @@ function edmfx_sgs_flux_tendency!(
             )
         end
 
-        # diffusive flux
+        # momentum diffusive flux
         ᶠρaK_u = p.ᶠtemp_scalar
         @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx])
+        ᶠstrain_rate = p.ᶠtemp_UVWxUVW
+        compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu[colidx])
+        @. Yₜ.c.uₕ[colidx] -= C12(
+            ᶜdivᵥ(-(2 * ᶠρaK_u[colidx] * ᶠstrain_rate[colidx])) / Y.c.ρ[colidx],
+        )
+        # apply boundary condition for momentum flux
         ᶜdivᵥ_uₕ = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(FT(0)) ⊗ C12(FT(0), FT(0))),
             bottom = Operators.SetValue(sfc_conditions.ρ_flux_uₕ[colidx]),
         )
         @. Yₜ.c.uₕ[colidx] -=
-            ᶜdivᵥ_uₕ(-(ᶠρaK_u[colidx] * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx]
+            ᶜdivᵥ_uₕ(-(FT(0) * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx]
     end
 
-    # TODO: Add momentum mass flux
-
     # TODO: Add tracer flux
 
     return nothing
diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl
index a6eb38ed21b..4d0a5b6c7f2 100644
--- a/src/utils/utilities.jl
+++ b/src/utils/utilities.jl
@@ -79,12 +79,12 @@ compute_kinetic!(κ::Fields.Field, Y::Fields.FieldVector) =
     compute_kinetic!(κ, Y.c.uₕ, Y.f.u₃)
 
 """
-    compute_strain_rate!(ϵ::Field, u::Field)
+    compute_strain_rate_center!(ϵ::Field, u::Field)
 
 Compute the strain_rate at cell centers, storing in `ϵ` from
 velocity at cell faces.
 """
-function compute_strain_rate!(ϵ::Fields.Field, u::Fields.Field)
+function compute_strain_rate_center!(ϵ::Fields.Field, u::Fields.Field)
     @assert eltype(u) <: C123
     axis_uvw = Geometry.UVWAxis()
     @. ϵ =
@@ -94,6 +94,28 @@ function compute_strain_rate!(ϵ::Fields.Field, u::Fields.Field)
         ) / 2
 end
 
+"""
+    compute_strain_rate_face!(ϵ::Field, u::Field)
+
+Compute the strain_rate at cell faces, storing in `ϵ` from
+velocity at cell centers.
+"""
+function compute_strain_rate_face!(ϵ::Fields.Field, u::Fields.Field)
+    @assert eltype(u) <: C123
+    ∇ᵥuvw_boundary =
+        Geometry.outer(Geometry.WVector(0), Geometry.UVWVector(0, 0, 0))
+    ᶠgradᵥ = Operators.GradientC2F(
+        bottom = Operators.SetGradient(∇ᵥuvw_boundary),
+        top = Operators.SetGradient(∇ᵥuvw_boundary),
+    )
+    axis_uvw = Geometry.UVWAxis()
+    @. ϵ =
+        (
+            Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u))) +
+            adjoint(Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u))))
+        ) / 2
+end
+
 """
     g³³_field(field)
 

From 70953a54e41368c34fee02fa056901e37d82ad80 Mon Sep 17 00:00:00 2001
From: nefrathenrici <nat.henrici@gmail.com>
Date: Fri, 22 Sep 2023 13:13:57 -0700
Subject: [PATCH 71/73] Fix scaling pipeline, change to use srun

---
 .buildkite/scaling/pipeline.sh          | 37 ++++++++-----------------
 post_processing/plot_scaling_results.jl | 13 ++-------
 2 files changed, 15 insertions(+), 35 deletions(-)

diff --git a/.buildkite/scaling/pipeline.sh b/.buildkite/scaling/pipeline.sh
index 28af1daa0e0..d579b6cf90f 100755
--- a/.buildkite/scaling/pipeline.sh
+++ b/.buildkite/scaling/pipeline.sh
@@ -65,32 +65,28 @@ for i in "${!resolutions[@]}"; do
 done
 
 # set up environment and agents
-cat << EOM
-env:
-  JULIA_VERSION: "1.9.3"
-  MPICH_VERSION: "4.0.0"
-  OPENMPI_VERSION: "4.1.1"
-  MPI_IMPL: "$mpi_impl"
-  CUDA_VERSION: "11.3"
-  OPENBLAS_NUM_THREADS: 1
-  CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: "true"
-
+cat << 'EOM'
 agents:
-  config: cpu
   queue: central
+  modules: julia/1.9.3 cuda/11.8 ucx/1.14.1_cuda-11.8 openmpi/4.1.5_cuda-11.8 hdf5/1.12.2-ompi415 nsight-systems/2023.2.1
+
+env:
+  JULIA_LOAD_PATH: "${JULIA_LOAD_PATH}:${BUILDKITE_BUILD_CHECKOUT_PATH}/.buildkite"
+  OPENBLAS_NUM_THREADS: 1
+  JULIA_NVTX_CALLBACKS: gc
+  OMPI_MCA_opal_warn_on_missing_libcuda: 0
+  JULIA_MAX_NUM_PRECOMPILE_FILES: 100
+  JULIA_CPU_TARGET: 'broadwell;skylake'
+  SLURM_KILL_BAD_EXIT: 1
 
 steps:
   - label: "init :computer:"
     key: "init_cpu_env"
     command:
-      - echo "--- Configure MPI"
-      - julia -e 'using Pkg; Pkg.add("MPIPreferences"); using MPIPreferences; use_system_binary()'
-
       - echo "--- Instantiate"
       - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
       - "julia --project=examples -e 'using Pkg; Pkg.precompile()'"
       - "julia --project=examples -e 'using Pkg; Pkg.status()'"
-
     agents:
       slurm_cpus_per_task: 8
     env:
@@ -129,11 +125,7 @@ if [[ "$profiling" == "enable" ]]; then
 else
     cpus_per_proc=1
 fi
-if [[ "$mpi_impl" == "mpich" ]]; then
-    launcher="srun --cpu-bind=cores"
-else
-    launcher="mpiexec --map-by node:PE=$cpus_per_proc --bind-to core"
-fi
+launcher="srun --cpu-bind=cores"
 
 if [[ "$res" == "low" ]]; then
     time="04:00:00"
@@ -160,7 +152,6 @@ cat << EOM
     - label: "$nprocs"
       key: "$job_id"
       command:
-        - "module load cuda/11.3 nsight-systems/2022.2.1"
         - "$launcher $command"
         - "find ${job_id} -iname '*.nsys-rep' -printf '%f\\\\n' | sort -V | jq --raw-input --slurp 'split(\"\n\") | .[0:-1] | {files: .} + {\"extension\": \"nsys-view\", \"version\": \"1.0\"}' > ${job_id}/${job_id}.nsys-view"
         - "find ${job_id} -iname '*.nsys-*' | sort -V | tar cvzf ${job_id}-nsys.tar.gz -T -"
@@ -170,8 +161,6 @@ cat << EOM
       env:
         CLIMACORE_DISTRIBUTED: "MPI"
       agents:
-        config: cpu
-        queue: central
         slurm_time: $time
 EOM
 
@@ -209,8 +198,6 @@ cat << EOM
         - "${res}-*.png"
         - "${res}-*.pdf"
       agents:
-        config: cpu
-        queue: central
         slurm_nodes: 1
         slurm_tasks_per_node: 1
 
diff --git a/post_processing/plot_scaling_results.jl b/post_processing/plot_scaling_results.jl
index d330e5854b6..0a979f12a0d 100644
--- a/post_processing/plot_scaling_results.jl
+++ b/post_processing/plot_scaling_results.jl
@@ -89,7 +89,8 @@ ax1 = Axis(
     xgridvisible = true,
     ygridvisible = false,
 )
-scatterlines!(ax1, nprocs_clima_atmos, sypd_clima_atmos)
+scatterlines!(nprocs_clima_atmos, sypd_clima_atmos)
+# Plot a second axis to display tick labels clearly
 ax1 = Axis(
     fig[1, 1],
     yaxisposition = :right,
@@ -103,15 +104,7 @@ ax1 = Axis(
     yscale = log10,
     ytickformat = "{:.2f}",
 )
-hlines!(
-    ax1,
-    sypd_clima_atmos,
-    xscale = log10,
-    yscale = log10,
-    color = :gray,
-    alpha = 0.5,
-    linestyle = :dash,
-)
+scatterlines!(nprocs_clima_atmos, sypd_clima_atmos)
 
 ax2 = Axis(
     fig[2, 1],

From 29817a94c6bd4b8721fb0e8f8777b0a9e68d5b45 Mon Sep 17 00:00:00 2001
From: nefrathenrici <nat.henrici@gmail.com>
Date: Tue, 26 Sep 2023 17:27:28 -0700
Subject: [PATCH 72/73] Move parameter logging

---
 src/parameters/create_parameters.jl | 7 +------
 src/solver/type_getters.jl          | 4 ++++
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/parameters/create_parameters.jl b/src/parameters/create_parameters.jl
index f8c790536f6..5c01f1e7bed 100644
--- a/src/parameters/create_parameters.jl
+++ b/src/parameters/create_parameters.jl
@@ -70,7 +70,7 @@ function create_parameter_set(config::AtmosConfig)
         TCP.TurbulenceConvectionParameters;
         subparam_structs = (; microphys_params, surf_flux_params),
     )
-    param_set = create_parameter_struct(
+    return create_parameter_struct(
         CAP.ClimaAtmosParameters;
         subparam_structs = (;
             thermodynamics_params = thermo_params,
@@ -81,9 +81,4 @@ function create_parameter_set(config::AtmosConfig)
             turbconv_params,
         ),
     )
-    if parsed_args["log_params"]
-        logfilepath = joinpath(@__DIR__, "$(parsed_args["job_id"])_$FT.toml")
-        CP.log_parameter_information(toml_dict, logfilepath)
-    end
-    return param_set
 end
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 186b467d161..befa7660e6a 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -787,6 +787,10 @@ function get_integrator(config::AtmosConfig)
     atmos = get_atmos(config, params)
     numerics = get_numerics(config.parsed_args)
     simulation = get_simulation(config, config.comms_ctx)
+    if config.parsed_args["log_params"]
+        filepath = joinpath(simulation.output_dir, "$(job_id)_parameters.toml")
+        CP.log_parameter_information(config.toml_dict, filepath)
+    end
     initial_condition = get_initial_condition(config.parsed_args)
     surface_setup = get_surface_setup(config.parsed_args)
 

From f8d2c598b8a438cf320692e6fdef2f13bed821f8 Mon Sep 17 00:00:00 2001
From: nefrathenrici <nat.henrici@gmail.com>
Date: Tue, 26 Sep 2023 17:27:42 -0700
Subject: [PATCH 73/73] Remove extra imports

---
 src/parameters/create_parameters.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/parameters/create_parameters.jl b/src/parameters/create_parameters.jl
index 5c01f1e7bed..b06cf796eca 100644
--- a/src/parameters/create_parameters.jl
+++ b/src/parameters/create_parameters.jl
@@ -7,9 +7,6 @@ import SurfaceFluxes.UniversalFunctions as UF
 import Insolation.Parameters as IP
 import Thermodynamics as TD
 import CloudMicrophysics as CM
-# TODO: Remove these imports?
-import ClimaCore
-import ClimaCore as CC
 
 function create_parameter_set(config::AtmosConfig)
     # Helper function that creates a parameter struct. If a struct has nested