Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix BoundaryAdjacentMean #4189

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "Oceananigans"
uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
authors = ["Climate Modeling Alliance and contributors"]
version = "0.95.22"
version = "0.95.23"

[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
Expand All @@ -11,6 +11,7 @@ CubedSphere = "7445602f-e544-4518-8976-18f8e8ae6cdb"
Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
Expand All @@ -37,12 +38,12 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"

[weakdeps]
ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b"
Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"

[extensions]
OceananigansEnzymeExt = "Enzyme"
Expand All @@ -58,6 +59,7 @@ CubedSphere = "0.2, 0.3"
Dates = "1.9"
Distances = "0.10"
DocStringExtensions = "0.8, 0.9"
Documenter = "1.8.1"
Enzyme = "0.13.30"
FFTW = "1"
GPUArrays = "10.3, 11.2.2"
Expand Down Expand Up @@ -92,8 +94,8 @@ julia = "1.9"
CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Expand Down
115 changes: 73 additions & 42 deletions src/Models/boundary_mean.jl
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we call this boundary_adjacent_mean to match the type name?

Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
using Adapt, CUDA
using Oceananigans: instantiated_location
using Oceananigans.Fields: Center, Face
using Oceananigans.AbstractOperations: GridMetricOperation, Ax, Ay, Az
using Oceananigans.AbstractOperations: GridMetricOperation, Ax, Ay, Az, @at
using Oceananigans.BoundaryConditions: BoundaryCondition, Open, PerturbationAdvection

import Adapt: adapt_structure
Expand All @@ -11,85 +11,116 @@ import Oceananigans.BoundaryConditions: update_boundary_condition!
"""
BoundaryAdjacentMean

Stores the boundary mean `value` of a `Field`. Updated by calling
Stores the boundary mean `value` of a `Field`. The value is updated and stored
when the object is called. This automatically happens during
`update_boundary_condition!` when `BoundaryAdjacentMean` is used as the condition
value in a boundary condition.

```jldoctest
julia> using Oceananigans

julia> using Oceananigans.Models: BoundaryAdjacentMean

julia> grid = RectilinearGrid(size = (16, 16, 16), extent = (3, 4, 5))
16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── Periodic x ∈ [0.0, 3.0) regularly spaced with Δx=0.1875
julia> grid = RectilinearGrid(topology = (Bounded, Periodic, Bounded), size = (16, 16, 16), extent = (3, 4, 5))
16×16×16 RectilinearGrid{Float64, Bounded, Periodic, Bounded} on CPU with 3×3×3 halo
├── Bounded x ∈ [0.0, 3.0] regularly spaced with Δx=0.1875
├── Periodic y ∈ [0.0, 4.0) regularly spaced with Δy=0.25
└── Bounded z ∈ [-5.0, 0.0] regularly spaced with Δz=0.3125

julia> bam = BoundaryAdjacentMean(grid, :east)
BoundaryAdjacentMean: (0.0)

julia> cf = CenterField(grid)
16×16×16 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── grid: 16×16×16 RectilinearGrid{Float64, Bounded, Periodic, Bounded} on CPU with 3×3×3 halo
├── boundary conditions: FieldBoundaryConditions
│ └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
│ └── west: ZeroFlux, east: ZeroFlux, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 22×22×22 OffsetArray(::Array{Float64, 3}, -2:19, -2:19, -2:19) with eltype Float64 with indices -2:19×-2:19×-2:19
└── max=0.0, min=0.0, mean=0.0

julia> set!(cf, (x, y, z) -> sin(2π * y / 4))
16×16×16 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 16×16×16 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── grid: 16×16×16 RectilinearGrid{Float64, Bounded, Periodic, Bounded} on CPU with 3×3×3 halo
├── boundary conditions: FieldBoundaryConditions
│ └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
│ └── west: ZeroFlux, east: ZeroFlux, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 22×22×22 OffsetArray(::Array{Float64, 3}, -2:19, -2:19, -2:19) with eltype Float64 with indices -2:19×-2:19×-2:19
└── max=0.980785, min=-0.980785, mean=1.10534e-16

julia> bam = BoundaryAdjacentMean(grid, :east)
BoundaryAdjacentMean: (0.0)
└── max=0.980785, min=-0.980785, mean=1.11077e-16

julia> bam(:east, cf)
-7.806255641895632e-19

julia> ff = Field{Face, Center, Center}(grid)
17×16×16 Field{Face, Center, Center} on RectilinearGrid on CPU
├── grid: 16×16×16 RectilinearGrid{Float64, Bounded, Periodic, Bounded} on CPU with 3×3×3 halo
├── boundary conditions: FieldBoundaryConditions
│ └── west: Nothing, east: Nothing, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 23×22×22 OffsetArray(::Array{Float64, 3}, -2:20, -2:19, -2:19) with eltype Float64 with indices -2:20×-2:19×-2:19
└── max=0.0, min=0.0, mean=0.0

julia> set!(ff, (x, y, z) -> sin(2π * y / 4))
17×16×16 Field{Face, Center, Center} on RectilinearGrid on CPU
├── grid: 16×16×16 RectilinearGrid{Float64, Bounded, Periodic, Bounded} on CPU with 3×3×3 halo
├── boundary conditions: FieldBoundaryConditions
│ └── west: Nothing, east: Nothing, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 23×22×22 OffsetArray(::Array{Float64, 3}, -2:20, -2:19, -2:19) with eltype Float64 with indices -2:20×-2:19×-2:19
└── max=0.980785, min=-0.980785, mean=1.0284e-16

julia> bam(:east, ff)
-1.5612511283791264e-18

julia> using Oceananigans.BoundaryConditions: PerturbationAdvectionOpenBoundaryCondition

julia> obc = PerturbationAdvectionOpenBoundaryCondition(bam)
┌ Warning: `PerturbationAdvection` open boundaries matching scheme is experimental and un-tested/validated
└ @ Oceananigans.BoundaryConditions ~/Oceananigans.jl/src/BoundaryConditions/perturbation_advection_open_boundary_matching_scheme.jl:52
OpenBoundaryCondition{Oceananigans.BoundaryConditions.PerturbationAdvection{Val{true}, Float64}}: BoundaryAdjacentMean: (-1.5612511283791264e-18)

julia> u_bcs = FieldBoundaryConditions(east = obc)
Oceananigans.FieldBoundaryConditions, with boundary conditions
├── west: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
├── east: OpenBoundaryCondition{Oceananigans.BoundaryConditions.PerturbationAdvection{Val{true}, Float64}}: BoundaryAdjacentMean: (-1.5612511283791264e-18)
├── south: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
├── north: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
├── bottom: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
├── top: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)
└── immersed: DefaultBoundaryCondition (FluxBoundaryCondition: Nothing)

```
"""
struct BoundaryAdjacentMean{FF, BV}
flux_field :: FF
value :: BV
struct BoundaryAdjacentMean{NI, BV}
boundary_normal_integral :: NI
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so the idea is that we compute an integral first and then divide by area to get the mean?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the type doesn't say anything about "normal" and I think this works on scalars right?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a good point, I guess boundary_parallel_integral is more accurate? boundary normal integal is actually wrong

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think merely integral is enough. The name of the type is "boundary adjacent mean". If we compute the mean by first taking the integral and dividing by area, then integral combined with the name of the type give you what yo uneed to know.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The first rule is "do no harm"

value :: BV

BoundaryAdjacentMean(grid, side;
flux_field::FF = boundary_reduced_field(Val(side), grid),
value::BV = Ref(zero(grid))) where {FF, BV} =
new{FF, BV}(flux_field, value)
boundary_normal_integral::NI = boundary_reduced_field(Val(side), grid),
value::BV = Ref(zero(grid))) where {NI, BV} =
new{NI, BV}(boundary_normal_integral, value)
end

@inline (bam::BoundaryAdjacentMean)(args...) = bam.value[]

Adapt.adapt_structure(to, mo::BoundaryAdjacentMean) =
BoundaryAdjacentMean(; flux_fields = nothing, value = adapt(to, mo.value[]))
BoundaryAdjacentMean(; boundary_normal_integral = nothing, value = adapt(to, mo.value[]))

Base.show(io::IO, bam::BoundaryAdjacentMean) = print(io, summary(bam)*"\n")
Base.summary(bam::BoundaryAdjacentMean) = "BoundaryAdjacentMean: ($(bam.value[]))"

@inline boundary_reduced_field(::Union{Val{:west}, Val{:east}}, grid) = Field{Center, Nothing, Nothing}(grid)
@inline boundary_reduced_field(::Union{Val{:south}, Val{:north}}, grid) = Field{Nothing, Center, Nothing}(grid)
@inline boundary_reduced_field(::Union{Val{:bottom}, Val{:top}}, grid) = Field{Nothing, Nothing, Center}(grid)
@inline boundary_reduced_field(::Union{Val{:west}, Val{:east}}, grid) = Field{Face, Nothing, Nothing}(grid)
@inline boundary_reduced_field(::Union{Val{:south}, Val{:north}}, grid) = Field{Nothing, Face, Nothing}(grid)
@inline boundary_reduced_field(::Union{Val{:bottom}, Val{:top}}, grid) = Field{Nothing, Nothing, Face}(grid)

@inline boundary_normal_area(::Union{Val{:west}, Val{:east}}, grid) = GridMetricOperation((Face, Center, Center), Ax, grid)
@inline boundary_normal_area(::Union{Val{:south}, Val{:north}}, grid) = GridMetricOperation((Center, Face, Center), Ay, grid)
@inline boundary_normal_area(::Union{Val{:bottom}, Val{:top}}, grid) = GridMetricOperation((Center, Center, Face), Az, grid)

@inline boundary_adjacent_indices(::Val{:east}, grid, loc) = size(grid, 1), 1, 1
@inline boundary_adjacent_indices(val_side::Val{:west}, grid, loc) = first_interior_index(val_side, loc), 1, 1

@inline boundary_adjacent_indices(::Val{:north}, grid, loc) = 1, size(grid, 2), 1
@inline boundary_adjacent_indices(val_side::Val{:south}, grid, loc) = 1, first_interior_index(val_side, loc), 1

@inline boundary_adjacent_indices(::Val{:top}, grid, loc) = 1, 1, size(grid, 3)
@inline boundary_adjacent_indices(val_side::Val{:bottom}, grid, loc) = 1, 1, first_interior_index(val_side, loc)

@inline first_interior_index(::Union{Val{:west}, Val{:east}}, ::Tuple{Center, <:Any, <:Any}) = 1
@inline first_interior_index(::Union{Val{:west}, Val{:east}}, ::Tuple{Face, <:Any, <:Any}) = 2
@inline boundary_adjacent_indices(::Val{:east}, grid, loc) = size(grid, 1)+1, 1, 1
@inline boundary_adjacent_indices(val_side::Val{:west}, grid, loc) = 1, 1, 1

@inline first_interior_index(::Union{Val{:south}, Val{:north}}, ::Tuple{<:Any, Center, <:Any}) = 1
@inline first_interior_index(::Union{Val{:south}, Val{:north}}, ::Tuple{<:Any, Face, <:Any}) = 2
@inline boundary_adjacent_indices(::Val{:north}, grid, loc) = 1, size(grid, 2)+1, 1
@inline boundary_adjacent_indices(val_side::Val{:south}, grid, loc) = 1, 1, 1

@inline first_interior_index(::Union{Val{:bottom}, Val{:top}}, ::Tuple{<:Any, <:Any, Center}) = 1
@inline first_interior_index(::Union{Val{:bottom}, Val{:top}}, ::Tuple{<:Any, <:Any, Face}) = 2
@inline boundary_adjacent_indices(::Val{:top}, grid, loc) = 1, 1, size(grid, 3)+1
@inline boundary_adjacent_indices(val_side::Val{:bottom}, grid, loc) = 1, 1, 1

(bam::BoundaryAdjacentMean)(side, u) = bam(Val(side), u)

Expand All @@ -102,14 +133,14 @@ function (bam::BoundaryAdjacentMean)(val_side::Val, u)
An = boundary_normal_area(val_side, grid)

# get the total flux
sum!(bam.flux_field, u * An)
sum!(bam.boundary_normal_integral, (@at (Face, Center, Center) u) * An)

bam.value[] = CUDA.@allowscalar bam.flux_field[iB, jB, kB]
bam.value[] = CUDA.@allowscalar bam.boundary_normal_integral[iB, jB, kB]

# get the normalizing area
sum!(bam.flux_field, An)
sum!(bam.boundary_normal_integral, An)

bam.value[] /= CUDA.@allowscalar bam.flux_field[iB, jB, kB]
bam.value[] /= CUDA.@allowscalar bam.boundary_normal_integral[iB, jB, kB]

return bam.value[]
end
Expand Down