From 8bdb5c725f1a79f181a1ad4fcb09e149a8053b5f Mon Sep 17 00:00:00 2001 From: tomaschor Date: Mon, 7 Jul 2025 10:34:01 -0700 Subject: [PATCH 01/11] add conditional_length method for NoConditionCO --- src/AbstractOperations/conditional_operations.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl index e3ad796c9ee..665e1a5ff11 100644 --- a/src/AbstractOperations/conditional_operations.jl +++ b/src/AbstractOperations/conditional_operations.jl @@ -198,6 +198,7 @@ end @inline conditional_length(c::ConditionalOperation) = sum(conditional_one(c, 0)) @inline conditional_length(c::ConditionalOperation, dims) = sum(conditional_one(c, 0); dims) +@inline conditional_length(c::NoConditionCO, args...) = conditional_length(c.operand, args...) compute_at!(c::ConditionalOperation, time) = compute_at!(c.operand, time) indices(c::ConditionalOperation) = indices(c.operand) From 4077f38dbafc529d31503db07e6438125da6cac7 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Mon, 7 Jul 2025 11:02:23 -0700 Subject: [PATCH 02/11] add method for reducing ConditionalOperations with condition==nothing --- src/AbstractOperations/conditional_operations.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl index 665e1a5ff11..36dc45e8d68 100644 --- a/src/AbstractOperations/conditional_operations.jl +++ b/src/AbstractOperations/conditional_operations.jl @@ -151,6 +151,11 @@ end return ifelse(conditioned, value, co.mask) end +# Fallbacks for reductions +for reduction! in (:sum!, :maximum!, :minimum!, :all!, :any!, :prod!) + @eval Base.$(reduction!)(id, int_r, co::NoConditionCO, kwargs...) = Base.$(reduction!)(id, int_r, co.operand; kwargs...) +end + # Conditions: general, nothing, array @inline evaluate_condition(condition, i, j, k, grid, args...) = condition(i, j, k, grid, args...) @inline evaluate_condition(::Nothing, i, j, k, grid, args...) = true From c54963663a413c269d8421f7ae6111a99faaf8fc Mon Sep 17 00:00:00 2001 From: tomaschor Date: Tue, 8 Jul 2025 10:10:25 -0700 Subject: [PATCH 03/11] calculate mean by integration --- .../metric_field_reductions.jl | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 5ddc0c8a520..47b72aebf41 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -1,4 +1,5 @@ using Statistics: mean!, sum! +using CUDA: @allowscalar using Oceananigans.Utils: tupleit using Oceananigans.Grids: regular_dimensions @@ -43,24 +44,17 @@ for information and examples using `condition` and `mask` kwargs. """ function Average(field::AbstractField; dims=:, condition=nothing, mask=0) dims = dims isa Colon ? (1, 2, 3) : tupleit(dims) + dx = reduction_grid_metric(dims) - if all(d in regular_dimensions(field.grid) for d in dims) - # Dimensions being reduced are regular; just use mean! - operand = condition_operand(field, condition, mask) - return Scan(Averaging(), mean!, operand, dims) - else - # Compute "size" (length, area, or volume) of averaging region - dx = reduction_grid_metric(dims) - metric = GridMetricOperation(location(field), dx, field.grid) - L = sum(metric; condition, mask, dims) - - # Construct summand of the Average - L⁻¹_field_dx = field * dx / L + # Compute "size" (length, area, or volume) of averaging region + metric = GridMetricOperation(location(field), dx, field.grid) + L = @allowscalar sum(metric; condition, mask, dims)[] # Pluck out L as a Number - operand = condition_operand(L⁻¹_field_dx, condition, mask) + # Construct summand of the Average + L⁻¹_field_dx = field * dx / L - return Scan(Averaging(), sum!, operand, dims) - end + operand = condition_operand(L⁻¹_field_dx, condition, mask) + return Scan(Averaging(), sum!, operand, dims) end struct Integrating <: AbstractReducing end From 71b429b50e2981c8b64e417acc8a6f634884a471 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Tue, 8 Jul 2025 10:51:00 -0700 Subject: [PATCH 04/11] simplify --- src/AbstractOperations/metric_field_reductions.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 47b72aebf41..c62d5c37393 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -49,11 +49,9 @@ function Average(field::AbstractField; dims=:, condition=nothing, mask=0) # Compute "size" (length, area, or volume) of averaging region metric = GridMetricOperation(location(field), dx, field.grid) L = @allowscalar sum(metric; condition, mask, dims)[] # Pluck out L as a Number + L = convert(eltype(field.grid), L) - # Construct summand of the Average - L⁻¹_field_dx = field * dx / L - - operand = condition_operand(L⁻¹_field_dx, condition, mask) + operand = condition_operand(field * dx / L, condition, mask) return Scan(Averaging(), sum!, operand, dims) end From f301b85b85d0aad6e5e99faaebd48c7f1395f275 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Tue, 8 Jul 2025 10:57:23 -0700 Subject: [PATCH 05/11] better names --- src/AbstractOperations/metric_field_reductions.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index c62d5c37393..07210aab993 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -48,10 +48,10 @@ function Average(field::AbstractField; dims=:, condition=nothing, mask=0) # Compute "size" (length, area, or volume) of averaging region metric = GridMetricOperation(location(field), dx, field.grid) - L = @allowscalar sum(metric; condition, mask, dims)[] # Pluck out L as a Number - L = convert(eltype(field.grid), L) + ∫dx = @allowscalar sum(metric; condition, mask, dims)[] # Pluck out L as a Number + ∫dx = convert(eltype(field.grid), ∫dx) - operand = condition_operand(field * dx / L, condition, mask) + operand = condition_operand(field * dx / ∫dx, condition, mask) return Scan(Averaging(), sum!, operand, dims) end From 673cbd7415ee6591410253515021f0ba54d4117f Mon Sep 17 00:00:00 2001 From: tomaschor Date: Tue, 8 Jul 2025 10:58:41 -0700 Subject: [PATCH 06/11] suggestion --- src/AbstractOperations/conditional_operations.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl index 36dc45e8d68..43a49ab09f5 100644 --- a/src/AbstractOperations/conditional_operations.jl +++ b/src/AbstractOperations/conditional_operations.jl @@ -153,7 +153,7 @@ end # Fallbacks for reductions for reduction! in (:sum!, :maximum!, :minimum!, :all!, :any!, :prod!) - @eval Base.$(reduction!)(id, int_r, co::NoConditionCO, kwargs...) = Base.$(reduction!)(id, int_r, co.operand; kwargs...) + @eval Base.$(reduction!)(f, int_r, co::NoConditionCO, kwargs...) = Base.$(reduction!)(f, int_r, co.operand; kwargs...) end # Conditions: general, nothing, array From 99160de5d330a9864fcb63226938ad32d432f0bc Mon Sep 17 00:00:00 2001 From: tomaschor Date: Thu, 10 Jul 2025 08:32:48 -0700 Subject: [PATCH 07/11] use integral of unity field to define area in Average --- src/AbstractOperations/metric_field_reductions.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 07210aab993..22b314e4846 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -46,10 +46,8 @@ function Average(field::AbstractField; dims=:, condition=nothing, mask=0) dims = dims isa Colon ? (1, 2, 3) : tupleit(dims) dx = reduction_grid_metric(dims) - # Compute "size" (length, area, or volume) of averaging region - metric = GridMetricOperation(location(field), dx, field.grid) - ∫dx = @allowscalar sum(metric; condition, mask, dims)[] # Pluck out L as a Number - ∫dx = convert(eltype(field.grid), ∫dx) + field_ones = Field(location(field), field.grid); set!(field_ones, 1) + ∫dx = Integral(field_ones; dims) |> Field operand = condition_operand(field * dx / ∫dx, condition, mask) return Scan(Averaging(), sum!, operand, dims) From e490abad8c4a88d5ef0b52e61e84dfbca3ff3387 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Mon, 14 Jul 2025 20:09:54 -0700 Subject: [PATCH 08/11] fix some tests --- test/test_field_scans.jl | 34 +++++++++++++++------------------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/test/test_field_scans.jl b/test/test_field_scans.jl index 576852b8d2b..881c6d65ac4 100644 --- a/test/test_field_scans.jl +++ b/test/test_field_scans.jl @@ -2,7 +2,7 @@ include("dependencies_for_runtests.jl") using Statistics using Oceananigans.Architectures: on_architecture -using Oceananigans.AbstractOperations: BinaryOperation +using Oceananigans.AbstractOperations: BinaryOperation, GridMetricOperation using Oceananigans.Fields: ReducedField, CenterField, ZFaceField, compute_at!, @compute, reverse_cumsum! using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Grids: halo_size @@ -90,15 +90,24 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) @compute ζrz = Field(CumulativeIntegral(ζ, dims=3, reverse=true)) for T′ in (Tx, Txy) - @test T′.operand.operand === T + @test T′.operand.operand.a.a === T + @test T′.operand.operand.a.b isa GridMetricOperation + @test T′.operand.operand.b.operand isa Integral + @test T′.operand.operand.b.operand.operand isa BinaryOperation end for w′ in (wx, wxy) - @test w′.operand.operand === w + @test w′.operand.operand.a.a === w + @test w′.operand.operand.a.b isa GridMetricOperation + @test w′.operand.operand.b.operand isa Integral + @test w′.operand.operand.b.operand.operand isa BinaryOperation end for ζ′ in (ζx, ζxy) - @test ζ′.operand.operand === ζ + @test ζ′.operand.operand.a.a === ζ + @test ζ′.operand.operand.a.b isa GridMetricOperation + @test ζ′.operand.operand.b.operand isa Integral + @test ζ′.operand.operand.b.operand.operand isa BinaryOperation end for f in (wx, wxy, Tx, Txy, ζx, ζxy, Wx, Wxy, Θx, Θxy, Zx, Zxy) @@ -112,11 +121,11 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) end for f in (wx, wxy, Tx, Txy, ζx, ζxy) - @test f.operand.scan! === mean! + @test f.operand.scan! === sum! end for f in (wx, wxy, Tx, Txy, ζx, ζxy) - @test f.operand.scan! === mean! + @test f.operand.scan! === sum! end for f in (Tcx, Tcy, Tcz, wcx, wcy, wcz, ζcx, ζcy, ζcz) @@ -131,19 +140,6 @@ interior_array(a, i, j, k) = Array(interior(a, i, j, k)) @test wxyz.operand isa Reduction @test ζxyz.operand isa Reduction - # Different behavior for regular grid z vs not. - if grid === regular_grid - @test Txyz.operand.scan! === mean! - @test wxyz.operand.scan! === mean! - @test Txyz.operand.operand === T - @test wxyz.operand.operand === w - else - @test Txyz.operand.scan! === sum! - @test wxyz.operand.scan! === sum! - @test Txyz.operand.operand isa BinaryOperation - @test wxyz.operand.operand isa BinaryOperation - end - @test Tx.operand.dims === tuple(1) @test wx.operand.dims === tuple(1) @test Txy.operand.dims === (1, 2) From 0485ce8bc00d92620f5225e42d83d7d21c69222c Mon Sep 17 00:00:00 2001 From: tomaschor Date: Mon, 14 Jul 2025 20:23:57 -0700 Subject: [PATCH 09/11] fix bug from main merge --- src/AbstractOperations/metric_field_reductions.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 22b314e4846..9c3ed4af0ef 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -1,5 +1,4 @@ -using Statistics: mean!, sum! -using CUDA: @allowscalar +using Statistics: sum! using Oceananigans.Utils: tupleit using Oceananigans.Grids: regular_dimensions From c731b19ca0ef6cf3c697a9af2aff478c5096ab66 Mon Sep 17 00:00:00 2001 From: tomaschor Date: Mon, 14 Jul 2025 20:31:57 -0700 Subject: [PATCH 10/11] fix call to integral in Average --- src/AbstractOperations/metric_field_reductions.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 9c3ed4af0ef..07f17d4d3db 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -46,7 +46,7 @@ function Average(field::AbstractField; dims=:, condition=nothing, mask=0) dx = reduction_grid_metric(dims) field_ones = Field(location(field), field.grid); set!(field_ones, 1) - ∫dx = Integral(field_ones; dims) |> Field + ∫dx = Integral(field_ones; dims, condition, mask) |> Field operand = condition_operand(field * dx / ∫dx, condition, mask) return Scan(Averaging(), sum!, operand, dims) From a30a48a04bafd0ffc4c4aa3cf79e61b97944c94b Mon Sep 17 00:00:00 2001 From: tomaschor Date: Mon, 28 Jul 2025 10:16:25 -0700 Subject: [PATCH 11/11] instantiate location when building the field --- src/AbstractOperations/metric_field_reductions.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AbstractOperations/metric_field_reductions.jl b/src/AbstractOperations/metric_field_reductions.jl index 07f17d4d3db..3899c0473c4 100644 --- a/src/AbstractOperations/metric_field_reductions.jl +++ b/src/AbstractOperations/metric_field_reductions.jl @@ -45,7 +45,7 @@ function Average(field::AbstractField; dims=:, condition=nothing, mask=0) dims = dims isa Colon ? (1, 2, 3) : tupleit(dims) dx = reduction_grid_metric(dims) - field_ones = Field(location(field), field.grid); set!(field_ones, 1) + field_ones = Field(instantiated_location(field), field.grid); set!(field_ones, 1) ∫dx = Integral(field_ones; dims, condition, mask) |> Field operand = condition_operand(field * dx / ∫dx, condition, mask)