From 45f0a6f1910f4626b6fb966bc0375bbbdbde0005 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 26 Sep 2024 12:19:20 -0600 Subject: [PATCH 1/5] Massive clean up for conditional operations --- .../conditional_operations.jl | 77 ++++++++++++------- src/Fields/field.jl | 59 +++++++------- src/ImmersedBoundaries/immersed_reductions.jl | 46 +++++++---- 3 files changed, 110 insertions(+), 72 deletions(-) diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl index 7eba3699f7..5e31d2bbb0 100644 --- a/src/AbstractOperations/conditional_operations.jl +++ b/src/AbstractOperations/conditional_operations.jl @@ -5,23 +5,27 @@ import Oceananigans.Architectures: on_architecture import Oceananigans.Fields: condition_operand, conditional_length, set!, compute_at!, indices # For conditional reductions such as mean(u * v, condition = u .> 0)) -struct ConditionalOperation{LX, LY, LZ, O, F, G, C, M, T} <: AbstractOperation{LX, LY, LZ, G, T} +struct ConditionalOperation{LX, LY, LZ, F, C, O, G, M, T} <: AbstractOperation{LX, LY, LZ, G, T} operand :: O func :: F grid :: G condition :: C mask :: M - function ConditionalOperation{LX, LY, LZ}(operand::O, func::F, grid::G, - condition::C, mask::M) where {LX, LY, LZ, O, F, G, C, M} + function ConditionalOperation{LX, LY, LZ}(operand::O, func, grid::G, + condition::C, mask::M) where {LX, LY, LZ, O, G, C, M} + if func === Base.identity + func = nothing + end T = eltype(operand) - return new{LX, LY, LZ, O, F, G, C, M, T}(operand, func, grid, condition, mask) + F = typeof(func) + return new{LX, LY, LZ, F, C, O, G, M, T}(operand, func, grid, condition, mask) end end """ ConditionalOperation(operand::AbstractField; - func = identity, + func = nothing, condition = nothing, mask = 0) @@ -31,19 +35,19 @@ described by `func(operand)`. Positional arguments ==================== -- `operand`: The `AbstractField` to be masked (it must have a `grid` property!) +- `operand`: The `AbstractField` to be masked. Keyword arguments ================= - `func`: A unary transformation applied element-wise to the field `operand` at locations where - `condition == true`. Default is `identity`. + `condition == true`. Default is `nothing` which applies no transformation. - `condition`: either a function of `(i, j, k, grid, operand)` returning a Boolean, or a 3-dimensional Boolean `AbstractArray`. At locations where `condition == false`, - operand will be masked by `mask` + operand will be masked by `mask`. -- `mask`: the scalar mask +- `mask`: the scalar mask. Default: 0. `condition_operand` is a convenience function used to construct a `ConditionalOperation` @@ -78,10 +82,9 @@ julia> d[2, 1, 1] ``` """ function ConditionalOperation(operand::AbstractField; - func = identity, + func = nothing, condition = nothing, mask = zero(eltype(operand))) - LX, LY, LZ = location(operand) return ConditionalOperation{LX, LY, LZ}(operand, func, operand.grid, condition, mask) end @@ -95,29 +98,43 @@ function ConditionalOperation(c::ConditionalOperation; return ConditionalOperation{LX, LY, LZ}(c.operand, func, c.grid, condition, mask) end -struct TrueCondition end +@inline function Base.getindex(co::ConditionalOperation, i, j, k) + conditioned = evaluate_condition(co.condition, i, j, k, c.grid, c) + value = getindex(co.operand, i, j, k) + func_value = co.func(value) + return ifelse(conditioned, value, c.mask) +end + +# Some special cases +const NoFuncCO = ConditionalOperation{<:Any, <:Any, <:Any, Nothing} +const NoConditionCO = ConditionalOperation{<:Any, <:Any, <:Any, <:Any, Nothing} +const NoFuncNoConditionCO = ConditionalOperation{<:Any, <:Any, <:Any, Nothing, Nothing} + +@inline function Base.getindex(co::NoConditionCO, i, j, k) + value = getindex(co.operand, i, j, k) + return co.func(value) +end + +@inline Base.getindex(co::NoFuncNoConditionCO, i, j, k) = getindex(co.operand, i, j, k) -@inline function Base.getindex(c::ConditionalOperation, i, j, k) - return ifelse(evaluate_condition(c.condition, i, j, k, c.grid, c), - c.func(getindex(c.operand, i, j, k)), - c.mask) +@inline function Base.getindex(co::NoFuncCO, i, j, k) + conditioned = evaluate_condition(co.condition, i, j, k, co.grid, co) + value = getindex(co.operand, i, j, k) + return ifelse(conditioned, value, co.mask) end +# Conditions: general, nothing, array @inline evaluate_condition(condition, i, j, k, grid, args...) = condition(i, j, k, grid, args...) -@inline evaluate_condition(::TrueCondition, i, j, k, grid, args...) = true +@inline evaluate_condition(::Nothing, i, j, k, grid, args...) = true @inline evaluate_condition(condition::AbstractArray, i, j, k, grid, args...) = @inbounds condition[i, j, k] -@inline condition_operand(func::Function, op::AbstractField, condition, mask) = ConditionalOperation(op; func, condition, mask) -@inline condition_operand(func::Function, op::AbstractField, ::Nothing, mask) = ConditionalOperation(op; func, condition=TrueCondition(), mask) +@inline condition_operand(func, op, condition, mask) = ConditionalOperation(op; func, condition, mask) -@inline function condition_operand(func::Function, operand::AbstractField, condition::AbstractArray, mask) +@inline function condition_operand(func, operand, condition::AbstractArray, mask) condition = on_architecture(architecture(operand.grid), condition) return ConditionalOperation(operand; func, condition, mask) end -@inline condition_operand(func::typeof(identity), c::ConditionalOperation, ::Nothing, mask) = ConditionalOperation(c; mask) -@inline condition_operand(func::Function, c::ConditionalOperation, ::Nothing, mask) = ConditionalOperation(c; func, mask) - @inline materialize_condition!(c::ConditionalOperation) = set!(c.operand, c) function materialize_condition(c::ConditionalOperation) @@ -126,14 +143,17 @@ function materialize_condition(c::ConditionalOperation) return f end -@inline condition_onefield(c::ConditionalOperation{LX, LY, LZ}, mask) where {LX, LY, LZ} = - ConditionalOperation{LX, LY, LZ}(OneField(Int), identity, c.grid, c.condition, mask) +@inline function conditional_one(c::ConditionalOperation, mask) + LX, LY, LZ = location(c) + one_field = OneField(Int) + return ConditionalOperation{LX, LY, LZ}(one_field, nothing, c.grid, c.condition, mask) +end -@inline conditional_length(c::ConditionalOperation) = sum(condition_onefield(c, 0)) -@inline conditional_length(c::ConditionalOperation, dims) = sum(condition_onefield(c, 0); dims = dims) +@inline conditional_length(c::ConditionalOperation) = sum(conditional_one(c, 0)) +@inline conditional_length(c::ConditionalOperation, dims) = sum(conditional_one(c, 0); dims = dims) Adapt.adapt_structure(to, c::ConditionalOperation{LX, LY, LZ}) where {LX, LY, LZ} = - ConditionalOperation{LX, LY, LZ}(adapt(to, c.operand), + ConditionalOperation{LX, LY, LZ}(adapt(to, c.operand), adapt(to, c.func), adapt(to, c.grid), adapt(to, c.condition), @@ -159,3 +179,4 @@ Base.show(io::IO, operation::ConditionalOperation) = "├── func: ", summary(operation.func), "\n", "├── condition: ", summary(operation.condition), "\n", "└── mask: ", operation.mask) + diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 0e8c87ffcb..3e78ce1587 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -642,15 +642,12 @@ function reduced_dimension(loc) return dims end -## Allow support for ConditionalOperation - get_neutral_mask(::Union{AllReduction, AnyReduction}) = true -get_neutral_mask(::Union{SumReduction, MeanReduction}) = 0 -get_neutral_mask(::MinimumReduction) = Inf -get_neutral_mask(::MaximumReduction) = - Inf -get_neutral_mask(::ProdReduction) = 1 +get_neutral_mask(::Union{SumReduction, MeanReduction}) = 0 +get_neutral_mask(::MinimumReduction) = +Inf +get_neutral_mask(::MaximumReduction) = -Inf +get_neutral_mask(::ProdReduction) = 1 -# If func = identity and condition = nothing, nothing happens """ condition_operand(f::Function, op::AbstractField, condition, mask) @@ -660,8 +657,11 @@ If `f isa identity` and `isnothing(condition)` then `op` is returned without wra Otherwise return `ConditionedOperand`, even when `isnothing(condition)` but `!(f isa identity)`. """ -@inline condition_operand(op::AbstractField, condition, mask) = condition_operand(identity, op, condition, mask) -@inline condition_operand(::typeof(identity), operand::AbstractField, ::Nothing, mask) = operand +@inline condition_operand(op::AbstractField, condition, mask) = condition_operand(nothing, op, condition, mask) + +# Do NOT condition if condition=nothing. +# All non-trivial conditioning is found in AbstractOperations/conditional_operations.jl +@inline condition_operand(::Nothing, operand, ::Nothing, mask) = operand @inline conditional_length(c::AbstractField) = length(c) @inline conditional_length(c::AbstractField, dims) = mapreduce(i -> size(c, i), *, unique(dims); init=1) @@ -674,24 +674,24 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) @eval begin # In-place - function Base.$(reduction!)(f::Function, - r::ReducedAbstractField, - a::AbstractField; - condition = nothing, - mask = get_neutral_mask(Base.$(reduction!)), - kwargs...) + @inline function Base.$(reduction!)(f::Function, + r::ReducedAbstractField, + a::AbstractField; + condition = nothing, + mask = get_neutral_mask(Base.$(reduction!)), + kwargs...) return Base.$(reduction!)(identity, interior(r), - condition_operand(f, a, condition, mask); + a.operand; #condition_operand(f, a, condition, mask); kwargs...) end - function Base.$(reduction!)(r::ReducedAbstractField, - a::AbstractField; - condition = nothing, - mask = get_neutral_mask(Base.$(reduction!)), - kwargs...) + @inline function Base.$(reduction!)(r::ReducedAbstractField, + a::AbstractField; + condition = nothing, + mask = get_neutral_mask(Base.$(reduction!)), + kwargs...) return Base.$(reduction!)(identity, interior(r), @@ -700,17 +700,20 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) end # Allocating - function Base.$(reduction)(f::Function, - c::AbstractField; - condition = nothing, - mask = get_neutral_mask(Base.$(reduction!)), - dims = :) + @inline function Base.$(reduction)(f::Function, + c::AbstractField; + condition = nothing, + mask = get_neutral_mask(Base.$(reduction!)), + dims = :) + + conditioned_c = condition_operand(f, c, condition, mask) T = filltype(Base.$(reduction!), c) loc = reduced_location(location(c); dims) r = Field(loc, c.grid, T; indices=indices(c)) - conditioned_c = condition_operand(f, c, condition, mask) initialize_reduced_field!(Base.$(reduction!), identity, r, conditioned_c) + + @show interior(r) conditioned_c Base.$(reduction!)(identity, r, conditioned_c, init=false) if dims isa Colon @@ -720,7 +723,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) end end - Base.$(reduction)(c::AbstractField; kwargs...) = Base.$(reduction)(identity, c; kwargs...) + @inline Base.$(reduction)(c::AbstractField; kwargs...) = Base.$(reduction)(identity, c; kwargs...) end end diff --git a/src/ImmersedBoundaries/immersed_reductions.jl b/src/ImmersedBoundaries/immersed_reductions.jl index a9969363c6..6fe5fb52c6 100644 --- a/src/ImmersedBoundaries/immersed_reductions.jl +++ b/src/ImmersedBoundaries/immersed_reductions.jl @@ -11,30 +11,44 @@ import Oceananigans.Fields: condition_operand, conditional_length @inline truefunc(args...) = true struct NotImmersed{F} <: Function - func :: F + condition :: F end +NotImmersed() = NotImmersed(nothing) +Base.summary(::NotImmersed{Nothing}) = "NotImmersed()" +Base.summary(::NotImmersed) = string("NotImmersed(", summary(condition), ")") + # ImmersedField const IF = AbstractField{<:Any, <:Any, <:Any, <:ImmersedBoundaryGrid} -@inline condition_operand(func::Function, op::IF, cond, mask) = ConditionalOperation(op; func, condition=NotImmersed(cond), mask) -@inline condition_operand(func::Function, op::IF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersed(truefunc), mask) -@inline condition_operand(func::typeof(identity), op::IF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersed(truefunc), mask) +function ConditionalOperation(operand::IF; + func = nothing, + condition = nothing, + mask = zero(eltype(operand))) -@inline function condition_operand(func::Function, op::IF, cond::AbstractArray, mask) - arch = architecture(op.grid) - arch_condition = on_architecture(arch, cond) - ni_condition = NotImmersed(arch_condition) - return ConditionalOperation(op; func, condition=ni_condition, mask) + immersed_condition = NotImmersed(condition) + LX, LY, LZ = location(operand) + return ConditionalOperation{LX, LY, LZ}(operand, func, operand.grid, immersed_condition, mask) end -@inline conditional_length(c::IF) = conditional_length(condition_operand(identity, c, nothing, 0)) -@inline conditional_length(c::IF, dims) = conditional_length(condition_operand(identity, c, nothing, 0), dims) +@inline conditional_length(c::IF) = conditional_length(condition_operand(c, nothing, 0)) +@inline conditional_length(c::IF, dims) = conditional_length(condition_operand(c, nothing, 0), dims) + +@inline function evaluate_condition(::NotImmersed{Nothing}, i, j, k, + grid::ImmersedBoundaryGrid, + co::ConditionalOperation, args...) + + ℓx, ℓy, ℓz = map(instantiate, location(co)) + return peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) +end + +@inline function evaluate_condition(condition::NotImmersed, i, j, k, + grid::ImmersedBoundaryGrid, + co::ConditionalOperation, args...) -@inline function evaluate_condition(condition::NotImmersed, i, j, k, ibg, co::ConditionalOperation, args...) ℓx, ℓy, ℓz = map(instantiate, location(co)) - immersed = immersed_peripheral_node(i, j, k, ibg, ℓx, ℓy, ℓz) | inactive_node(i, j, k, ibg, ℓx, ℓy, ℓz) - return !immersed & evaluate_condition(condition.func, i, j, k, ibg, args...) + immersed = peripheral_node(i, j, k, ibg, ℓx, ℓy, ℓz) + return !immersed & evaluate_condition(ni.condition, i, j, k, ibg, args...) end ##### @@ -78,9 +92,9 @@ end @inline center_to_nothing(::Type{Center}) = Center @inline center_to_nothing(::Type{Nothing}) = Center -@inline function evaluate_condition(condition::NotImmersedColumn, i, j, k, ibg, co::ConditionalOperation, args...) +@inline function evaluate_condition(nic::NotImmersedColumn, i, j, k, ibg, co::ConditionalOperation, args...) LX, LY, LZ = location(co) - return evaluate_condition(condition.func, i, j, k, ibg, args...) & !(is_immersed_column(i, j, k, condition.immersed_column)) + return evaluate_condition(ni.condition, i, j, k, ibg, args...) & !(is_immersed_column(i, j, k, nic.immersed_column)) end @inline is_immersed_column(i, j, k, column) = @inbounds column[i, j, k] == 0 From 4252b288f0c12e5f1c90c5397748b56951fa833e Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 26 Sep 2024 12:25:19 -0600 Subject: [PATCH 2/5] Get rid of inlines --- src/Fields/field.jl | 36 +++++++++++++++++------------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 3e78ce1587..dadd808e0d 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -674,12 +674,12 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) @eval begin # In-place - @inline function Base.$(reduction!)(f::Function, - r::ReducedAbstractField, - a::AbstractField; - condition = nothing, - mask = get_neutral_mask(Base.$(reduction!)), - kwargs...) + function Base.$(reduction!)(f::Function, + r::ReducedAbstractField, + a::AbstractField; + condition = nothing, + mask = get_neutral_mask(Base.$(reduction!)), + kwargs...) return Base.$(reduction!)(identity, interior(r), @@ -687,11 +687,11 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) kwargs...) end - @inline function Base.$(reduction!)(r::ReducedAbstractField, - a::AbstractField; - condition = nothing, - mask = get_neutral_mask(Base.$(reduction!)), - kwargs...) + function Base.$(reduction!)(r::ReducedAbstractField, + a::AbstractField; + condition = nothing, + mask = get_neutral_mask(Base.$(reduction!)), + kwargs...) return Base.$(reduction!)(identity, interior(r), @@ -700,11 +700,11 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) end # Allocating - @inline function Base.$(reduction)(f::Function, - c::AbstractField; - condition = nothing, - mask = get_neutral_mask(Base.$(reduction!)), - dims = :) + function Base.$(reduction)(f::Function, + c::AbstractField; + condition = nothing, + mask = get_neutral_mask(Base.$(reduction!)), + dims = :) conditioned_c = condition_operand(f, c, condition, mask) @@ -712,8 +712,6 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) loc = reduced_location(location(c); dims) r = Field(loc, c.grid, T; indices=indices(c)) initialize_reduced_field!(Base.$(reduction!), identity, r, conditioned_c) - - @show interior(r) conditioned_c Base.$(reduction!)(identity, r, conditioned_c, init=false) if dims isa Colon @@ -723,7 +721,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) end end - @inline Base.$(reduction)(c::AbstractField; kwargs...) = Base.$(reduction)(identity, c; kwargs...) + Base.$(reduction)(c::AbstractField; kwargs...) = Base.$(reduction)(identity, c; kwargs...) end end From 445ca6a8d424e5a85ec9a6988c96e22a83e79ca6 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 26 Sep 2024 12:37:36 -0600 Subject: [PATCH 3/5] Bugfix --- src/AbstractOperations/conditional_operations.jl | 2 +- src/Fields/field.jl | 2 +- src/ImmersedBoundaries/immersed_reductions.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl index 5e31d2bbb0..252a80da85 100644 --- a/src/AbstractOperations/conditional_operations.jl +++ b/src/AbstractOperations/conditional_operations.jl @@ -149,7 +149,7 @@ end return ConditionalOperation{LX, LY, LZ}(one_field, nothing, c.grid, c.condition, mask) end -@inline conditional_length(c::ConditionalOperation) = sum(conditional_one(c, 0)) +@inline conditional_length(c::ConditionalOperation) = sum(conditional_one(c, 0)) @inline conditional_length(c::ConditionalOperation, dims) = sum(conditional_one(c, 0); dims = dims) Adapt.adapt_structure(to, c::ConditionalOperation{LX, LY, LZ}) where {LX, LY, LZ} = diff --git a/src/Fields/field.jl b/src/Fields/field.jl index dadd808e0d..7eed699d2c 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -683,7 +683,7 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) return Base.$(reduction!)(identity, interior(r), - a.operand; #condition_operand(f, a, condition, mask); + condition_operand(f, a, condition, mask); kwargs...) end diff --git a/src/ImmersedBoundaries/immersed_reductions.jl b/src/ImmersedBoundaries/immersed_reductions.jl index 6fe5fb52c6..2299e6b798 100644 --- a/src/ImmersedBoundaries/immersed_reductions.jl +++ b/src/ImmersedBoundaries/immersed_reductions.jl @@ -31,7 +31,7 @@ function ConditionalOperation(operand::IF; return ConditionalOperation{LX, LY, LZ}(operand, func, operand.grid, immersed_condition, mask) end -@inline conditional_length(c::IF) = conditional_length(condition_operand(c, nothing, 0)) +@inline conditional_length(c::IF) = conditional_length(condition_operand(c, nothing, 0)) @inline conditional_length(c::IF, dims) = conditional_length(condition_operand(c, nothing, 0), dims) @inline function evaluate_condition(::NotImmersed{Nothing}, i, j, k, From 419c718842cc7096d4d346c3d056fced33503cda Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 26 Sep 2024 14:21:41 -0600 Subject: [PATCH 4/5] Ouchies --- src/AbstractOperations/AbstractOperations.jl | 9 ---- .../conditional_operations.jl | 26 ++++++------ src/Fields/abstract_field.jl | 33 +++++++++++++++ src/Fields/field.jl | 29 ++----------- src/ImmersedBoundaries/immersed_reductions.jl | 42 ++++++++++++------- 5 files changed, 78 insertions(+), 61 deletions(-) diff --git a/src/AbstractOperations/AbstractOperations.jl b/src/AbstractOperations/AbstractOperations.jl index 390597cc6c..576f633179 100644 --- a/src/AbstractOperations/AbstractOperations.jl +++ b/src/AbstractOperations/AbstractOperations.jl @@ -34,15 +34,6 @@ abstract type AbstractOperation{LX, LY, LZ, G, T} <: AbstractField{LX, LY, LZ, G const AF = AbstractField # used in unary_operations.jl, binary_operations.jl, etc -function Base.axes(f::AbstractOperation) - idx = indices(f) - if idx === (:, : ,:) - return Base.OneTo.(size(f)) - else - return Tuple(idx[i] isa Colon ? Base.OneTo(size(f, i)) : idx[i] for i = 1:3) - end -end - # We have no halos to fill @inline fill_halo_regions!(::AbstractOperation, args...; kwargs...) = nothing diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl index 252a80da85..7f9c7251cb 100644 --- a/src/AbstractOperations/conditional_operations.jl +++ b/src/AbstractOperations/conditional_operations.jl @@ -110,23 +110,24 @@ const NoFuncCO = ConditionalOperation{<:Any, <:Any, <:Any, Nothing} const NoConditionCO = ConditionalOperation{<:Any, <:Any, <:Any, <:Any, Nothing} const NoFuncNoConditionCO = ConditionalOperation{<:Any, <:Any, <:Any, Nothing, Nothing} -@inline function Base.getindex(co::NoConditionCO, i, j, k) +using Base: @propagate_inbounds +@propagate_inbounds function Base.getindex(co::NoConditionCO, i, j, k) value = getindex(co.operand, i, j, k) return co.func(value) end -@inline Base.getindex(co::NoFuncNoConditionCO, i, j, k) = getindex(co.operand, i, j, k) +@propagate_inbounds Base.getindex(co::NoFuncNoConditionCO, i, j, k) = getindex(co.operand, i, j, k) -@inline function Base.getindex(co::NoFuncCO, i, j, k) +@propagate_inbounds function Base.getindex(co::NoFuncCO, i, j, k) conditioned = evaluate_condition(co.condition, i, j, k, co.grid, co) value = getindex(co.operand, i, j, k) return ifelse(conditioned, value, co.mask) 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 -@inline evaluate_condition(condition::AbstractArray, i, j, k, grid, args...) = @inbounds condition[i, j, k] +@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 +@propagate_inbounds evaluate_condition(condition::AbstractArray, i, j, k, grid, args...) = condition[i, j, k] @inline condition_operand(func, op, condition, mask) = ConditionalOperation(op; func, condition, mask) @@ -172,11 +173,10 @@ compute_at!(c::ConditionalOperation, time) = compute_at!(c.operand, time) indices(c::ConditionalOperation) = indices(c.operand) Base.show(io::IO, operation::ConditionalOperation) = - print(io, - "ConditionalOperation at $(location(operation))", "\n", - "├── operand: ", summary(operation.operand), "\n", - "├── grid: ", summary(operation.grid), "\n", - "├── func: ", summary(operation.func), "\n", - "├── condition: ", summary(operation.condition), "\n", - "└── mask: ", operation.mask) + print(io, "ConditionalOperation at $(location(operation))", '\n', + "├── operand: ", summary(operation.operand), '\n', + "├── grid: ", summary(operation.grid), '\n', + "├── func: ", summary(operation.func), '\n', + "├── condition: ", summary(operation.condition), '\n', + "└── mask: ", operation.mask) diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index ab4475796c..1b1d06cc5e 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -58,6 +58,20 @@ Base.size(f::AbstractField) = size(f.grid, location(f)) Base.length(f::AbstractField) = prod(size(f)) Base.parent(f::AbstractField) = f +@inline axis(::Colon, N) = Base.OneTo(N) +@inline axis(index::UnitRange, N) = index + +@inline function Base.axes(f::AbstractField) + Nx, Ny, Nz = size(f) + ix, iy, iz = indices(f) + + ax = axis(ix, Nx) + ay = axis(iy, Ny) + az = axis(iz, Nz) + + return (ax, ay, az) +end + """ total_size(field::AbstractField) @@ -92,3 +106,22 @@ for f in (:+, :-) @eval Base.$f(ϕ::AbstractField, ψ::AbstractArray) = $f(interior(ϕ), ψ) end +const XReducedAF = AbstractField{Nothing} +const YReducedAF = AbstractField{<:Any, Nothing} +const ZReducedAF = AbstractField{<:Any, <:Any, Nothing} + +const YZReducedAF = AbstractField{<:Any, Nothing, Nothing} +const XZReducedAF = AbstractField{Nothing, <:Any, Nothing} +const XYReducedAF = AbstractField{Nothing, Nothing, <:Any} + +const XYZReducedAF = AbstractField{Nothing, Nothing, Nothing} + +reduced_dimensions(field::AbstractField) = () +reduced_dimensions(field::XReducedAF) = tuple(1) +reduced_dimensions(field::YReducedAF) = tuple(2) +reduced_dimensions(field::ZReducedAF) = tuple(3) +reduced_dimensions(field::YZReducedAF) = (2, 3) +reduced_dimensions(field::XZReducedAF) = (1, 3) +reduced_dimensions(field::XYReducedAF) = (1, 2) +reduced_dimensions(field::XYZReducedAF) = (1, 2, 3) + diff --git a/src/Fields/field.jl b/src/Fields/field.jl index 7eed699d2c..98760de274 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -398,20 +398,6 @@ interior(f::Field, I...) = view(interior(f), I...) # Don't use axes(f) to checkbounds; use axes(f.data) Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...) -@inline axis(::Colon, N) = Base.OneTo(N) -@inline axis(index::UnitRange, N) = index - -@inline function Base.axes(f::Field) - Nx, Ny, Nz = size(f) - ix, iy, iz = f.indices - - ax = axis(ix, Nx) - ay = axis(iy, Ny) - az = axis(iz, Nz) - - return (ax, ay, az) -end - @propagate_inbounds Base.getindex(f::Field, inds...) = getindex(f.data, inds...) @propagate_inbounds Base.getindex(f::Field, i::Int) = parent(f)[i] @propagate_inbounds Base.setindex!(f::Field, val, i, j, k) = setindex!(f.data, val, i, j, k) @@ -527,15 +513,6 @@ const ReducedField = Union{XReducedField, XYReducedField, XYZReducedField} -reduced_dimensions(field::Field) = () -reduced_dimensions(field::XReducedField) = tuple(1) -reduced_dimensions(field::YReducedField) = tuple(2) -reduced_dimensions(field::ZReducedField) = tuple(3) -reduced_dimensions(field::YZReducedField) = (2, 3) -reduced_dimensions(field::XZReducedField) = (1, 3) -reduced_dimensions(field::XYReducedField) = (1, 2) -reduced_dimensions(field::XYZReducedField) = (1, 2, 3) - @propagate_inbounds Base.getindex(r::XReducedField, i, j, k) = getindex(r.data, 1, j, k) @propagate_inbounds Base.getindex(r::YReducedField, i, j, k) = getindex(r.data, i, 1, k) @propagate_inbounds Base.getindex(r::ZReducedField, i, j, k) = getindex(r.data, i, j, 1) @@ -644,9 +621,11 @@ end get_neutral_mask(::Union{AllReduction, AnyReduction}) = true get_neutral_mask(::Union{SumReduction, MeanReduction}) = 0 +get_neutral_mask(::ProdReduction) = 1 + +# TODO make this Float32 friendly get_neutral_mask(::MinimumReduction) = +Inf get_neutral_mask(::MaximumReduction) = -Inf -get_neutral_mask(::ProdReduction) = 1 """ condition_operand(f::Function, op::AbstractField, condition, mask) @@ -707,7 +686,6 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod) dims = :) conditioned_c = condition_operand(f, c, condition, mask) - T = filltype(Base.$(reduction!), c) loc = reduced_location(location(c); dims) r = Field(loc, c.grid, T; indices=indices(c)) @@ -782,3 +760,4 @@ function fill_halo_regions!(field::Field, args...; kwargs...) return nothing end + diff --git a/src/ImmersedBoundaries/immersed_reductions.jl b/src/ImmersedBoundaries/immersed_reductions.jl index 2299e6b798..ba0d2c8bd0 100644 --- a/src/ImmersedBoundaries/immersed_reductions.jl +++ b/src/ImmersedBoundaries/immersed_reductions.jl @@ -28,38 +28,46 @@ function ConditionalOperation(operand::IF; immersed_condition = NotImmersed(condition) LX, LY, LZ = location(operand) - return ConditionalOperation{LX, LY, LZ}(operand, func, operand.grid, immersed_condition, mask) + grid = operand.grid + # if operand isa Field + # operand = interior(operand) + # end + return ConditionalOperation{LX, LY, LZ}(operand, func, grid, immersed_condition, mask) end @inline conditional_length(c::IF) = conditional_length(condition_operand(c, nothing, 0)) @inline conditional_length(c::IF, dims) = conditional_length(condition_operand(c, nothing, 0), dims) -@inline function evaluate_condition(::NotImmersed{Nothing}, i, j, k, +@inline function evaluate_condition(::NotImmersed{Nothing}, + i, j, k, grid::ImmersedBoundaryGrid, - co::ConditionalOperation, args...) + co::ConditionalOperation) #, args...) ℓx, ℓy, ℓz = map(instantiate, location(co)) return peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) end -@inline function evaluate_condition(condition::NotImmersed, i, j, k, +@inline function evaluate_condition(ni::NotImmersed, + i, j, k, grid::ImmersedBoundaryGrid, co::ConditionalOperation, args...) ℓx, ℓy, ℓz = map(instantiate, location(co)) - immersed = peripheral_node(i, j, k, ibg, ℓx, ℓy, ℓz) - return !immersed & evaluate_condition(ni.condition, i, j, k, ibg, args...) + immersed = peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) + return !immersed & evaluate_condition(ni.condition, i, j, k, grid, co, args...) end ##### ##### Reduction operations on Reduced Fields test the immersed condition on the entirety of the immersed direction ##### -struct NotImmersedColumn{IC, F} <:Function +struct NotImmersedColumn{F, IC} <:Function immersed_column :: IC - func :: F + condition :: F end +NotImmersedColumn(immersed_column) = NotImmersedColumn(immersed_column, nothing) + using Oceananigans.Fields: reduced_dimensions, OneField using Oceananigans.AbstractOperations: ConditionalOperation @@ -76,15 +84,16 @@ const XYZIRF = AbstractField{Nothing, Nothing, Nothing, <:ImmersedBoundaryGrid} const IRF = Union{XIRF, YIRF, ZIRF, YZIRF, XZIRF, XYIRF, XYZIRF} -@inline condition_operand(func::Function, op::IRF, cond, mask) = ConditionalOperation(op; func, condition=NotImmersedColumn(immersed_column(op), cond ), mask) -@inline condition_operand(func::Function, op::IRF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersedColumn(immersed_column(op), truefunc), mask) -@inline condition_operand(func::typeof(identity), op::IRF, ::Nothing, mask) = ConditionalOperation(op; func, condition=NotImmersedColumn(immersed_column(op), truefunc), mask) +@inline function condition_operand(func, op::IRF, condition, mask) + immersed_condition = NotImmersedColumn(immersed_column(op), condition) + return ConditionalOperation(op; func, condition, mask) +end @inline function immersed_column(field::IRF) grid = field.grid reduced_dims = reduced_dimensions(field) LX, LY, LZ = map(center_to_nothing, location(field)) - one_field = ConditionalOperation{LX, LY, LZ}(OneField(Int), identity, grid, NotImmersed(truefunc), zero(grid)) + one_field = ConditionalOperation{LX, LY, LZ}(OneField(Int), identity, grid, NotImmersed(), zero(grid)) return sum(one_field, dims=reduced_dims) end @@ -92,9 +101,14 @@ end @inline center_to_nothing(::Type{Center}) = Center @inline center_to_nothing(::Type{Nothing}) = Center -@inline function evaluate_condition(nic::NotImmersedColumn, i, j, k, ibg, co::ConditionalOperation, args...) +@inline function evaluate_condition(nic::NotImmersedColumn, + i, j, k, + grid::ImmersedBoundaryGrid, + co::ConditionalOperation, args...) LX, LY, LZ = location(co) - return evaluate_condition(ni.condition, i, j, k, ibg, args...) & !(is_immersed_column(i, j, k, nic.immersed_column)) + immersed = is_immersed_column(i, j, k, nic.immersed_column) + return !immersed & evaluate_condition(nic.condition, i, j, k, grid, args...) end @inline is_immersed_column(i, j, k, column) = @inbounds column[i, j, k] == 0 + From 714c1cac2b30580dac1cf82b2025e4c1a8312257 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Sun, 29 Sep 2024 20:53:37 -0400 Subject: [PATCH 5/5] Update immersed_reductions.jl --- src/ImmersedBoundaries/immersed_reductions.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/ImmersedBoundaries/immersed_reductions.jl b/src/ImmersedBoundaries/immersed_reductions.jl index ba0d2c8bd0..a5a5224cb1 100644 --- a/src/ImmersedBoundaries/immersed_reductions.jl +++ b/src/ImmersedBoundaries/immersed_reductions.jl @@ -29,9 +29,6 @@ function ConditionalOperation(operand::IF; immersed_condition = NotImmersed(condition) LX, LY, LZ = location(operand) grid = operand.grid - # if operand isa Field - # operand = interior(operand) - # end return ConditionalOperation{LX, LY, LZ}(operand, func, grid, immersed_condition, mask) end @@ -44,7 +41,8 @@ end co::ConditionalOperation) #, args...) ℓx, ℓy, ℓz = map(instantiate, location(co)) - return peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) + immersed = immersed_peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) | inactive_node(i, j, k, grid, ℓx, ℓy, ℓz) + return !immersed end @inline function evaluate_condition(ni::NotImmersed, @@ -53,7 +51,7 @@ end co::ConditionalOperation, args...) ℓx, ℓy, ℓz = map(instantiate, location(co)) - immersed = peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) + immersed = immersed_peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) | inactive_node(i, j, k, grid, ℓx, ℓy, ℓz) return !immersed & evaluate_condition(ni.condition, i, j, k, grid, co, args...) end