Add support for unevaluated compute! functions

`LazyBroadcast.jl` provides a way to return an unevaluated function.
This is useful in two cases:
1. reduce code verbosity to handle the `isnothing(out)` case
2. allow clustering all the broadcasted expressions in a single place

In turn, 2. is useful because it is the first step in fusing different
broadcasted calls.

This commit adds support for such functions.

Author: Sbozzolo

.buildkite/Project.toml

@@ -10,6 +10,7 @@ ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"

NEWS.md

@@ -4,6 +4,7 @@ v0.2.4
 -------
 
 - Add `@assign`.
+- Add support for lazy compute functions that return `Base.Broadcast.Broadcasted` objects. 
 
 v0.2.3
 -------

Project.toml

@@ -20,6 +20,7 @@ ClimaCore = "0.13.4, 0.14"
 ClimaTimeSteppers = "0.7.10"
 Dates = "1"
 Documenter = "1"
+LazyBroadcast = "0.1.3"
 JuliaFormatter = "1"
 NCDatasets = "0.13.1, 0.14"
 Profile = "1"
@@ -35,10 +36,11 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 ProfileCanvas = "efd6af41-a80b-495e-886c-e51b0c7d77a3"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "BenchmarkTools", "ClimaTimeSteppers", "Documenter", "JuliaFormatter", "Profile", "ProfileCanvas", "SafeTestsets", "Test"]
+test = ["Aqua", "BenchmarkTools", "ClimaTimeSteppers", "Documenter", "JuliaFormatter", "LazyBroadcast", "Profile", "ProfileCanvas", "SafeTestsets", "Test"]

docs/src/developer_guide.md

@@ -201,6 +201,27 @@ add_diagnostic_variable!(
 )
 ```
 
+When writing compute functions, consider using [`@assign`](@ref) to simplify
+your code and, when possible, making your expression lazy with
+[LazyBroadcast.jl](https://github.com/CliMA/LazyBroadcast.jl) to further improve
+clarity and performance. To do that, add `LazyBroadcast` to your dependencies
+and import `@lazy`. The previous example would look like:
+
+```julia
+###
+# Density (3d)
+###
+add_diagnostic_variable!(
+    short_name = "rhoa",
+    long_name = "Air Density",
+    standard_name = "air_density",
+    units = "kg m^-3",
+    compute! = (out, state, cache, time) -> begin
+        return @lazy @. state.c.ρ
+    end,
+)
+```
+
 It is a good idea to put safeguards in place to ensure that your users will not
 be allowed to call diagnostics that do not make sense for the simulation they
 are running. If your package has a notion of `Model` that is stored in `p`, you
@@ -224,11 +245,7 @@ function compute_hus!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    if isnothing(out)
-        return state.c.ρq_tot ./ state.c.ρ
-    else
-        out .= state.c.ρq_tot ./ state.c.ρ
-    end
+    @assign out state.c.ρq_tot ./ state.c.ρ
 end
 
 add_diagnostic_variable!(

docs/src/index.md

@@ -23,5 +23,6 @@ the Developer guide page.
 - Allow users to define arbitrary new diagnostics;
 - Trigger diagnostics on arbitrary conditions;
 - Save output to HDF5 or NetCDF files, or a dictionary in memory;
-
+- Work with lazy expressions (such as the ones produced by
+  [LazyBroadcast.jl](https://github.com/CliMA/LazyBroadcast.jl)).
 

docs/src/user_guide.md

@@ -58,7 +58,8 @@ var = DiagnosticVariable(;
 `compute_ta!` is the key function here. It determines how the variable should be
 computed from the `state`, `cache`, and `time` of the simulation. Typically,
 these are packaged within an `integrator` object (e.g., `state = integrator.u`
-or `integrator.Y`).
+or `integrator.Y`). `copy` is needed because we want to return a new area of
+memory (without `copy`, `ClimaDiagnostics` might end up modifying `state.ta`).
 
 `compute_ta!` takes another argument, `out`. `out` is an area of memory managed
 by `ClimaDiagnostics` that is used to reduce the number of allocations needed
@@ -69,9 +70,8 @@ overwritten, leading to much better performance. You should follow this pattern
 in all your diagnostics. To help you with that, `ClimaDiagnostics` provides a
 macro, `@assign`, that expands to the if switch, so that `compute_ta!` can be
 written as
-
 ```julia
-import ClimaDiagnostics: DiagnosticVariable, @assign
+import ClimaDiagnostics: @assign
 
 function compute_ta!(out, state, cache, time)
     @assign out copy(state.ta) state.ta
@@ -80,10 +80,18 @@ end
 Note that when the second and third arguments of `@assign` are the same, one of
 the two can be omitted.
 
-> Note, in the future, we hope to improve this rather clumsy way to write
-> diagnostics. Hopefully, at some point you will just have to write something like
-> `state.ta` and not worry about the `out` at all. For the time being, we recommend
-> using the `@assign` macro to help with possible future transitions.
+`ClimaDiagnostics` supports working with unevaluated expressions represented by
+`Base.Broadcast.Broadcasted` objects, such as the ones produced with
+[LazyBroadcast.jl](https://github.com/CliMA/LazyBroadcast.jl). Using
+`LazyBroadcast.jl`, the previous snippet can be rewritten as
+```julia
+import LazyBroadcast: @lazy
+
+function compute_ta!(out, state, cache, time)
+    @lazy @. out = state.ta
+end
+```
+Using lazy expressions can lead to improved performance and clearer code.
 
 A `DiagnosticVariable` defines what a variable is and how to compute it, but
 does not specify when to compute/output it. For that, we need

src/clima_diagnostics.jl

@@ -18,6 +18,7 @@ struct DiagnosticsHandler{
     STORAGE <: Dict,
     ACC <: Dict,
     COUNT <: Dict,
+    BROAD <: Dict,
 }
     """An iterable with the `ScheduledDiagnostic`s that are scheduled."""
     scheduled_diagnostics::SD
@@ -34,6 +35,11 @@ struct DiagnosticsHandler{
     many times the given diagnostics was computed from the last time it was output to
     disk."""
     counters::COUNT
+
+    """Dictionary that maps a given `ScheduledDiagnostic` to a Base.Broadcast.Broadcasted
+    object. This is used to allow lazy evaluation of expressions, which can lead to reduced
+    code verbosity and improved performance."""
+    broadcasted_expressions::BROAD
 end
 
 """
@@ -57,9 +63,11 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
 
     # For diagnostics that perform reductions, the storage is used for the values computed
     # at each call. Reductions also save the accumulated value in accumulators.
+    # broadcasted_expressions maps diagnostics with LazyBroadcast objects.
     storage = Dict()
     accumulators = Dict()
     counters = Dict()
+    broadcasted_expressions = Dict()
 
     unique_scheduled_diagnostics = Tuple(unique(scheduled_diagnostics))
     if length(unique_scheduled_diagnostics) != length(scheduled_diagnostics)
@@ -90,8 +98,18 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
         isa_time_reduction = !isnothing(diag.reduction_time_func)
 
         # 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
-        storage[diag] = variable.compute!(nothing, Y, p, t)
+        # (in the callbacks), we will write the result in place. ClimaDiagnostics supports
+        # LazyBroadcast.jl. In this case, the return value of `compute!` is a
+        # `Base.Broadcast.Broadcasted` and we have to manually materialize the result.
+        out_or_broadcasted_expr = variable.compute!(nothing, Y, p, t)
+        if out_or_broadcasted_expr isa Base.Broadcast.Broadcasted
+            broadcasted_expressions[diag] = out_or_broadcasted_expr
+            storage[diag] =
+                Base.Broadcast.materialize(broadcasted_expressions[diag])
+        else
+            storage[diag] = out_or_broadcasted_expr
+        end
+
         counters[diag] = 1
 
         # If it is not a reduction, call the output writer as well
@@ -115,6 +133,7 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
         storage,
         accumulators,
         counters,
+        broadcasted_expressions,
     )
 end
 
@@ -153,13 +172,42 @@ function orchestrate_diagnostics(
         active_compute[diag_index] || continue
         diag = scheduled_diagnostics[diag_index]
 
-        diag.variable.compute!(
+        diagnostic_handler.counters[diag] += 1
+
+        # ClimaDiagnostics supports LazyBroadcast.jl. When used, the return value of
+        # `compute!` is a `Base.Broadcast.Broadcasted`. We materialize the output to
+        # diagnostic_handler.storage[diag] in the next for loop. If the output is not a
+        # `Base.Broadcast.Broadcasted`, we don't have to do anything because compute! will
+        # already update diagnostic_handler.storage[diag].
+
+        out_or_broadcasted_expr = diag.variable.compute!(
             diagnostic_handler.storage[diag],
             integrator.u,
             integrator.p,
             integrator.t,
         )
-        diagnostic_handler.counters[diag] += 1
+        if out_or_broadcasted_expr isa Base.Broadcast.Broadcasted
+            diagnostic_handler.broadcasted_expressions[diag] =
+                out_or_broadcasted_expr
+        end
+    end
+
+    # Evaluate the lazy compute (aka, materialize everything)
+    for diag_index in 1:length(scheduled_diagnostics)
+        active_compute[diag_index] || continue
+        diag = scheduled_diagnostics[diag_index]
+        haskey(diagnostic_handler.broadcasted_expressions, diag) || continue
+
+        Base.Broadcast.materialize!(
+            diagnostic_handler.storage[diag],
+            diagnostic_handler.broadcasted_expressions[diag],
+        )
+    end
+
+    # Process possible time reductions (now we have evaluated storage[diag])
+    for diag_index in 1:length(scheduled_diagnostics)
+        active_compute[diag_index] || continue
+        diag = scheduled_diagnostics[diag_index]
 
         isa_time_reduction = !isnothing(diag.reduction_time_func)
         if isa_time_reduction

test/integration_test.jl

@@ -12,6 +12,8 @@ import ClimaComms
     ClimaComms.@import_required_backends
 end
 
+import LazyBroadcast: @lazy
+
 const context = ClimaComms.context()
 ClimaComms.init(context)
 
@@ -41,7 +43,15 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
     )
 
     function compute_my_var!(out, u, p, t)
-        ClimaDiagnostics.@assign out copy(u.my_var)
+        if isnothing(out)
+            return copy(u.my_var)
+        else
+            out .= u.my_var
+        end
+    end
+
+    function compute_my_var_lazy!(out, u, p, t)
+        return @lazy @. out = u.my_var
     end
 
     simple_var = ClimaDiagnostics.DiagnosticVariable(;
@@ -50,6 +60,12 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
         long_name = "YO YO",
     )
 
+    simple_var_lazy = ClimaDiagnostics.DiagnosticVariable(;
+        compute! = compute_my_var_lazy!,
+        short_name = "YO LAZY",
+        long_name = "YO YO LAZY",
+    )
+
     average_diagnostic = ClimaDiagnostics.ScheduledDiagnostic(
         variable = simple_var,
         output_writer = nc_writer,
@@ -61,6 +77,10 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
         variable = simple_var,
         output_writer = nc_writer,
     )
+    inst_diagnostic_lazy = ClimaDiagnostics.ScheduledDiagnostic(
+        variable = simple_var_lazy,
+        output_writer = nc_writer,
+    )
     inst_every3s_diagnostic = ClimaDiagnostics.ScheduledDiagnostic(
         variable = simple_var,
         output_writer = nc_writer,
@@ -76,6 +96,7 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
     scheduled_diagnostics = [
         average_diagnostic,
         inst_diagnostic,
+        inst_diagnostic_lazy,
         inst_diagnostic_h5,
         inst_every3s_diagnostic,
     ]