diff --git a/docs/src/APIs/Utilities/SingleStackUtils.md b/docs/src/APIs/Utilities/SingleStackUtils.md index e6998a3be92..6932b760255 100644 --- a/docs/src/APIs/Utilities/SingleStackUtils.md +++ b/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 ``` diff --git a/src/Utilities/SingleStackUtils/SingleStackUtils.jl b/src/Utilities/SingleStackUtils/SingleStackUtils.jl index 161e23d3622..c5d83ebed0f 100644 --- a/src/Utilities/SingleStackUtils/SingleStackUtils.jl +++ b/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 diff --git a/test/Utilities/SingleStackUtils/runtests.jl b/test/Utilities/SingleStackUtils/runtests.jl index ed04fe7909c..b6292a028c6 100644 --- a/test/Utilities/SingleStackUtils/runtests.jl +++ b/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 diff --git a/test/Utilities/SingleStackUtils/horizontal_stats_test.jl b/test/Utilities/SingleStackUtils/ssu_tests.jl similarity index 77% rename from test/Utilities/SingleStackUtils/horizontal_stats_test.jl rename to test/Utilities/SingleStackUtils/ssu_tests.jl index 8b4b85fa888..4d29250869f 100644 --- a/test/Utilities/SingleStackUtils/horizontal_stats_test.jl +++ b/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()