Merge #1294

1294: Add `reduce_[nodal|element]_stack()` in `SingleStackUtils` r=kpamnany a=kpamnany

# Description

And a couple of tests. Closes #1175.



Co-authored-by: K Pamnany <kpamnany@gmail.com>

docs/src/APIs/Utilities/SingleStackUtils.md

@@ -11,4 +11,6 @@ get_vars_from_nodal_stack
 get_vars_from_element_stack
 get_horizontal_mean
 get_horizontal_variance
+reduce_nodal_stack
+reduce_element_stack
 ```

src/Utilities/SingleStackUtils/SingleStackUtils.jl

@@ -1,7 +1,11 @@
 module SingleStackUtils
 
-export get_vars_from_nodal_stack, get_vars_from_element_stack
-export get_horizontal_variance, get_horizontal_mean
+export get_vars_from_nodal_stack,
+    get_vars_from_element_stack,
+    get_horizontal_variance,
+    get_horizontal_mean,
+    reduce_nodal_stack,
+    reduce_element_stack
 
 using OrderedCollections
 using StaticArrays
@@ -204,4 +208,93 @@ function get_horizontal_variance(
     return vars_var
 end
 
+"""
+    reduce_nodal_stack(
+        op::Function,
+        grid::DiscontinuousSpectralElementGrid{T, dim, N},
+        Q::MPIStateArray,
+        vars::NamedTuple,
+        var::String;
+        vrange::UnitRange = 1:size(Q, 3),
+    ) where {T, dim, N}
+
+Reduce `var` from `vars` within `Q` over all nodal points in the specified
+`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is
+the final value returned by `op` and `z` is the index within `vrange` where the
+`result` was determined.
+"""
+function reduce_nodal_stack(
+    op::Function,
+    grid::DiscontinuousSpectralElementGrid{T, dim, N},
+    Q::MPIStateArray,
+    vars::Type,
+    var::String;
+    vrange::UnitRange = 1:size(Q, 3),
+    i::Int = 1,
+    j::Int = 1,
+) where {T, dim, N}
+    Nq = N + 1
+    Nqk = dimensionality(grid) == 2 ? 1 : Nq
+
+    var_names = flattenednames(vars)
+    var_ind = findfirst(s -> s == var, var_names)
+    if var_ind === nothing
+        return
+    end
+
+    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)
+    z = vrange.start
+    result = state_data[1, var_ind, z]
+    for ev in vrange
+        for k in 1:Nqk
+            ijk = i + Nq * ((j - 1) + Nq * (k - 1))
+            new_result = op(result, state_data[ijk, var_ind, ev])
+            if !isequal(new_result, result)
+                result = new_result
+                z = ev
+            end
+        end
+    end
+
+    return (result, z)
+end
+
+"""
+    reduce_element_stack(
+        op::Function,
+        grid::DiscontinuousSpectralElementGrid{T, dim, N},
+        Q::MPIStateArray,
+        vars::NamedTuple,
+        var::String;
+        vrange::UnitRange = 1:size(Q, 3),
+    ) where {T, dim, N}
+
+Reduce `var` from `vars` within `Q` over all nodal points in the specified
+`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is
+the final value returned by `op` and `z` is the index within `vrange` where the
+`result` was determined.
+"""
+function reduce_element_stack(
+    op::Function,
+    grid::DiscontinuousSpectralElementGrid{T, dim, N},
+    Q::MPIStateArray,
+    vars::Type,
+    var::String;
+    vrange::UnitRange = 1:size(Q, 3),
+) where {T, dim, N}
+    Nq = N + 1
+    return [
+        reduce_nodal_stack(
+            op,
+            grid,
+            Q,
+            vars,
+            var,
+            vrange = vrange,
+            i = i,
+            j = j,
+        ) for i in 1:Nq, j in 1:Nq
+    ]
+end
+
 end # module

test/Utilities/SingleStackUtils/runtests.jl

@@ -2,5 +2,5 @@ using MPI, Test
 include(joinpath("..", "..", "testhelpers.jl"))
 
 @testset "SingleStackUtils" begin
-    runmpi(joinpath(@__DIR__, "horizontal_stats_test.jl"))
+    runmpi(joinpath(@__DIR__, "ssu_tests.jl"))
 end

test/Utilities/SingleStackUtils/horizontal_stats_test.jl → test/Utilities/SingleStackUtils/ssu_tests.jl

@@ -127,7 +127,7 @@ function main()
     dt = 0.1
     # Establish a `ClimateMachine` single stack configuration
     driver_config = ClimateMachine.SingleStackConfiguration(
-        "HstatsTest",
+        "SingleStackUtilsTest",
         N_poly,
         nelem_vert,
         zmax,
@@ -141,17 +141,49 @@ function main()
         ode_dt = dt,
     )
 
-    @testset "horizontal_stats" begin
-        test_hmean(
-            driver_config.grid,
-            solver_config.Q,
-            vars_state_conservative(m, FT),
-        )
-        test_hvar(
-            driver_config.grid,
-            solver_config.Q,
-            vars_state_conservative(m, FT),
-        )
+    # tests
+    test_hmean(
+        driver_config.grid,
+        solver_config.Q,
+        vars_state_conservative(m, FT),
+    )
+    test_hvar(
+        driver_config.grid,
+        solver_config.Q,
+        vars_state_conservative(m, FT),
+    )
+
+    r1, z1 = reduce_nodal_stack(
+        max,
+        solver_config.dg.grid,
+        solver_config.Q,
+        vars_state_conservative(m, FT),
+        "ρ",
+        i = 6,
+        j = 6,
+    )
+    @test r1 ≈ 8.880558532968455e-16 && z1 == 10
+    r2, z2 = reduce_nodal_stack(
+        +,
+        solver_config.dg.grid,
+        solver_config.Q,
+        vars_state_conservative(m, FT),
+        "ρ",
+        i = 3,
+        j = 3,
+    )
+    @test r2 ≈ 102.73283921735293 && z2 == 20
+    ns = reduce_element_stack(
+        +,
+        solver_config.dg.grid,
+        solver_config.Q,
+        vars_state_conservative(m, FT),
+        "ρ",
+    )
+    (r3, z3) = let
+        f(a, b) = (a[1] + b[1], b[2])
+        reduce(f, ns)
     end
+    @test r3 ≈ FT(2877.6) && z3 == 20
 end
 main()