Skip to content

Commit

Permalink
Merge branch 'main' into ss/comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@simone-silvestri
simone-silvestri authored Jan 24, 2025
2 parents 67bdd9c + 12ee14e commit 24054b3
Show file tree
Hide file tree
Showing 6 changed files with 162 additions and 133 deletions.
80 changes: 54 additions & 26 deletions src/Advection/topologically_conditional_interpolation.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,39 +10,67 @@
##### close to the boundary, or a second-order interpolation if i is close to a boundary.
#####

using Oceananigans.Grids: AbstractUnderlyingGrid, Bounded
using Oceananigans.Grids: AbstractUnderlyingGrid,
Bounded,
RightConnected,
LeftConnected,
topology,
architecture

const AUG = AbstractUnderlyingGrid

# topologies bounded at least on one side
const BT = Union{Bounded, RightConnected, LeftConnected}

# Bounded underlying Grids
const AUGX = AUG{<:Any, <:Bounded}
const AUGY = AUG{<:Any, <:Any, <:Bounded}
const AUGZ = AUG{<:Any, <:Any, <:Any, <:Bounded}
const AUGXY = AUG{<:Any, <:Bounded, <:Bounded}
const AUGXZ = AUG{<:Any, <:Bounded, <:Any, <:Bounded}
const AUGYZ = AUG{<:Any, <:Any, <:Bounded, <:Bounded}
const AUGXYZ = AUG{<:Any, <:Bounded, <:Bounded, <:Bounded}
const AUGX = AUG{<:Any, <:BT}
const AUGY = AUG{<:Any, <:Any, <:BT}
const AUGZ = AUG{<:Any, <:Any, <:Any, <:BT}
const AUGXY = AUG{<:Any, <:BT, <:BT}
const AUGXZ = AUG{<:Any, <:BT, <:Any, <:BT}
const AUGYZ = AUG{<:Any, <:Any, <:BT, <:BT}
const AUGXYZ = AUG{<:Any, <:BT, <:BT, <:BT}

# Left-biased buffers are smaller by one grid point on the right side; vice versa for right-biased buffers
# Center interpolation stencil look at i + 1 (i.e., require one less point on the left)

for dir in (:x, :y, :z)
outside_symmetric_haloᶠ = Symbol(:outside_symmetric_halo_, dir, :ᶠ)
outside_symmetric_haloᶜ = Symbol(:outside_symmetric_halo_, dir, :ᶜ)
outside_biased_haloᶠ = Symbol(:outside_biased_halo_, dir, :ᶠ)
outside_biased_haloᶜ = Symbol(:outside_biased_halo_, dir, :ᶜ)
required_halo_size = Symbol(:required_halo_size_, dir)
outside_biased_haloᶠ = Symbol(:outside_biased_halo_, dir, :ᶠ)
outside_biased_haloᶜ = Symbol(:outside_biased_halo_, dir, :ᶜ)
required_halo_size = Symbol(:required_halo_size_, dir)

@eval begin
@inline $outside_symmetric_haloᶠ(i, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv))
@inline $outside_symmetric_haloᶜ(i, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv))

@inline $outside_biased_haloᶠ(i, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias
(i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv)) # Right bias
@inline $outside_biased_haloᶜ(i, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias
(i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv)) # Right bias
# Bounded topologies
@inline $outside_symmetric_haloᶠ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv))
@inline $outside_symmetric_haloᶜ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv))

@inline $outside_biased_haloᶠ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias
(i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv)) # Right bias
@inline $outside_biased_haloᶜ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias
(i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv)) # Right bias

# Right connected topologies (only test the left side, i.e. the bounded side)
@inline $outside_symmetric_haloᶠ(i, ::Type{RightConnected}, N, adv) = i >= $required_halo_size(adv) + 1
@inline $outside_symmetric_haloᶜ(i, ::Type{RightConnected}, N, adv) = i >= $required_halo_size(adv)

@inline $outside_biased_haloᶠ(i, ::Type{RightConnected}, N, adv) = (i >= $required_halo_size(adv) + 1) & # Left bias
(i >= $required_halo_size(adv)) # Right bias
@inline $outside_biased_haloᶜ(i, ::Type{RightConnected}, N, adv) = (i >= $required_halo_size(adv)) & # Left bias
(i >= $required_halo_size(adv) - 1) # Right bias

# Left bounded topologies (only test the right side, i.e. the bounded side)
@inline $outside_symmetric_haloᶠ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - $required_halo_size(adv))
@inline $outside_symmetric_haloᶜ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - $required_halo_size(adv))

@inline $outside_biased_haloᶠ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias
(i <= N + 1 - $required_halo_size(adv)) # Right bias
@inline $outside_biased_haloᶜ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias
(i <= N + 1 - $required_halo_size(adv)) # Right bias
end
end

# Separate High order advection from low order advection
const HOADV = Union{WENO,
Tuple(Centered{N} for N in advection_buffers[2:end])...,
Expand Down Expand Up @@ -75,22 +103,22 @@ for bias in (:symmetric, :biased)
# Conditional high-order interpolation in Bounded directions
if ξ == :x
@eval begin
@inline $alt1_interp(i, j, k, grid::AUGX, scheme::HOADV, args...) =
ifelse($outside_buffer(i, grid.Nx, scheme),
$interp(i, j, k, grid, scheme, args...),
$alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...))
@inline $alt1_interp(i, j, k, grid::AUGX, scheme::HOADV, args...) =
ifelse($outside_buffer(i, topology(grid, 1), grid.Nx, scheme),
$interp(i, j, k, grid, scheme, args...),
$alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...))
end
elseif ξ == :y
@eval begin
@inline $alt1_interp(i, j, k, grid::AUGY, scheme::HOADV, args...) =
ifelse($outside_buffer(j, grid.Ny, scheme),
ifelse($outside_buffer(j, topology(grid, 2), grid.Ny, scheme),
$interp(i, j, k, grid, scheme, args...),
$alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...))
end
elseif ξ == :z
@eval begin
@inline $alt1_interp(i, j, k, grid::AUGZ, scheme::HOADV, args...) =
ifelse($outside_buffer(k, grid.Nz, scheme),
ifelse($outside_buffer(k, topology(grid, 3), grid.Nz, scheme),
$interp(i, j, k, grid, scheme, args...),
$alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...))
end
Expand All @@ -100,11 +128,11 @@ for bias in (:symmetric, :biased)
end

@inline _multi_dimensional_reconstruction_x(i, j, k, grid::AUGX, scheme, interp, args...) =
ifelse(outside_symmetric_bufferᶜ(i, grid.Nx, scheme),
ifelse(outside_symmetric_bufferᶜ(i, topology(grid, 1), grid.Nx, scheme),
multi_dimensional_reconstruction_x(i, j, k, grid::AUGX, scheme, interp, args...),
interp(i, j, k, grid, scheme, args...))

@inline _multi_dimensional_reconstruction_y(i, j, k, grid::AUGY, scheme, interp, args...) =
ifelse(outside_symmetric_bufferᶜ(j, grid.Ny, scheme),
ifelse(outside_symmetric_bufferᶜ(j, topology(grid, 2), grid.Ny, scheme),
multi_dimensional_reconstruction_y(i, j, k, grid::AUGY, scheme, interp, args...),
interp(i, j, k, grid, scheme, args...))
23 changes: 13 additions & 10 deletions src/DistributedComputations/distributed_immersed_boundaries.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,22 +121,25 @@ function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any
Nx, Ny, Nz = size(ibg)
Hx, Hy, _ = halo_size(ibg)

x_boundary = (Hx, Ny, Nz)
y_boundary = (Nx, Hy, Nz)

left_offsets = (0, 0, 0)
right_x_offsets = (Nx-Hx, 0, 0)
right_y_offsets = (0, Ny-Hy, 0)
west_boundary = (1:Hx, 1:Ny, 1:Nz)
east_boundary = (Nx-Hx+1:Nx, 1:Ny, 1:Nz)
south_boundary = (1:Nx, 1:Hy, 1:Nz)
north_boundary = (1:Nx, Ny-Hy+1:Ny, 1:Nz)

include_west = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected)
include_east = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected)
include_south = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected)
include_north = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected)

west_halo_dependent_cells = include_west ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, left_offsets)) : nothing
east_halo_dependent_cells = include_east ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, right_x_offsets)) : nothing
south_halo_dependent_cells = include_south ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, left_offsets)) : nothing
north_halo_dependent_cells = include_north ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, right_y_offsets)) : nothing
west_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(west_boundary...))
east_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(east_boundary...))
south_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(south_boundary...))
north_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(north_boundary...))

west_halo_dependent_cells = ifelse(include_west, west_halo_dependent_cells, nothing)
east_halo_dependent_cells = ifelse(include_east, east_halo_dependent_cells, nothing)
south_halo_dependent_cells = ifelse(include_south, south_halo_dependent_cells, nothing)
north_halo_dependent_cells = ifelse(include_north, north_halo_dependent_cells, nothing)

nx = Rx == 1 ? Nx : (Tx == RightConnected || Tx == LeftConnected ? Nx - Hx : Nx - 2Hx)
ny = Ry == 1 ? Ny : (Ty == RightConnected || Ty == LeftConnected ? Ny - Hy : Ny - 2Hy)
Expand Down
24 changes: 10 additions & 14 deletions src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,17 +35,16 @@ end
function compute_buffer_tendency_contributions!(grid::DistributedActiveCellsIBG, arch, model)
maps = grid.interior_active_cells

for (name, map) in zip(keys(maps), maps)

# If there exists a buffer map, then we compute the buffer contributions. If not, the
# buffer contributions have already been calculated. We exclude the interior because it has
# already been calculated
compute_buffer = (name != :interior) && !isnothing(map)

if compute_buffer
active_cells_map = retrieve_interior_active_cells_map(grid, Val(name))
compute_hydrostatic_free_surface_tendency_contributions!(model, :xyz; active_cells_map)
end
for name in (:west_halo_dependent_cells,
:east_halo_dependent_cells,
:south_halo_dependent_cells,
:north_halo_dependent_cells)

active_cells_map = @inbounds maps[name]

# If the map == nothing, we don't need to compute the buffer because
# the buffer is not adjacent to a processor boundary
!isnothing(map) && compute_hydrostatic_free_surface_tendency_contributions!(model, :xyz; active_cells_map)
end

return nothing
Expand All @@ -56,9 +55,6 @@ function buffer_w_kernel_parameters(grid, arch)
Nx, Ny, _ = size(grid)
Hx, Hy, _ = halo_size(grid)

Sx = (Hx, Ny+2)
Sy = (Nx+2, Hy)

# Offsets in tangential direction are == -1 to
# cover the required corners
param_west = (-Hx+2:1, 0:Ny+1)
Expand Down
3 changes: 2 additions & 1 deletion src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,8 @@ function update_state!(model::NonhydrostaticModel, callbacks=[]; compute_tendenc

# Calculate diffusivities and hydrostatic pressure
@apply_regionally compute_auxiliaries!(model)
fill_halo_regions!(model.diffusivity_fields; only_local_halos = true)

fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)

for callback in callbacks
callback.callsite isa UpdateStateCallsite && callback(model)
Expand Down
Loading

0 comments on commit 24054b3

Please sign in to comment.