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.

Sign up for GitHub

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

[WIP] Extend MultiRegion to NonhydrostaticModel #2523

Closed
wants to merge 3 commits into from
Closed

[WIP] Extend MultiRegion to NonhydrostaticModel #2523

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cc25f86
first commit
simone-silvestri May 5, 2022
3fba2f0
task number 1 complete
simone-silvestri May 6, 2022
038d5be
allow with_halo on nonhydrostatic
simone-silvestri May 6, 2022
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
2 changes: 1 addition & 1 deletion src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model,
arch = model.architecture

@apply_regionally prognostic_field_events = wait_velocity_event(arch, prognostic_field_events)
fill_halo_regions!(model.velocities)
fill_halo_regions!(model.velocities, model.clock, fields(model))

# Compute right hand side of implicit free surface equation
@apply_regionally local_compute_integrated_volume_flux!(∫ᶻQ, model.velocities, arch)
Expand Down
1 change: 1 addition & 0 deletions src/Models/NonhydrostaticModels/NonhydrostaticModels.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ using DocStringExtensions
using KernelAbstractions: @index, @kernel, Event, MultiEvent
using KernelAbstractions.Extras.LoopInfo: @unroll

using Oceananigans.Utils
using Oceananigans.Utils: launch!
using Oceananigans.Grids
using Oceananigans.Solvers
Expand Down
14 changes: 13 additions & 1 deletion src/Models/NonhydrostaticModels/nonhydrostatic_model.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ function NonhydrostaticModel(; grid,

grid = with_halo((Hx, Hy, Hz), grid)
end

# Collect boundary conditions for all model prognostic fields and, if specified, some model
# auxiliary fields. Boundary conditions are "regularized" based on the _name_ of the field:
# boundary conditions on u, v, w are regularized assuming they represent momentum at appropriate
Expand Down Expand Up @@ -225,3 +225,15 @@ function extract_boundary_conditions(field_tuple::NamedTuple)
end

extract_boundary_conditions(field::Field) = field.boundary_conditions


function validate_model_halo(grid, advection, closure)
user_halo = halo_size(grid)
required_halo = inflate_halo_size(user_halo..., topology(grid),
advection,
closure)

any(user_halo .< required_halo) &&
throw(ArgumentError("The grid halo $user_halo must be larger than $required_halo."))
end

20 changes: 13 additions & 7 deletions src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,25 +13,31 @@ Update peripheral aspects of the model (halo regions, diffusivities, hydrostatic
"""
function update_state!(model::NonhydrostaticModel)

@apply_regionally local_calculations!(model)

# Fill halos for velocities and tracers
fill_halo_regions!(merge(model.velocities, model.tracers), model.clock, fields(model))
fill_halo_regions!(model.diffusivity_fields, model.clock, fields(model))
fill_halo_regions!(model.pressures.pHY′)

return nothing
end

function local_calculations!(model)
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
function local_calculations!(model)
function local_nonhydrostatic_state_update!(model)

even though this function appears right next to where it's used, it's possible that we encounter it in external contexts (ie in an error). So it's helpful if we give it a descriptive name (since it's not much extra effort to do that, we imght as well...)


# Mask immersed tracers
tracer_masking_events = Tuple(mask_immersed_field!(c) for c in model.tracers)

wait(device(model.architecture), MultiEvent(tracer_masking_events))

# Fill halos for velocities and tracers
fill_halo_regions!(merge(model.velocities, model.tracers), model.clock, fields(model))

# Compute auxiliary fields
for aux_field in model.auxiliary_fields
compute!(aux_field)
end

# Calculate diffusivities
calculate_diffusivities!(model.diffusivity_fields, model.closure, model)
fill_halo_regions!(model.diffusivity_fields, model.clock, fields(model))

update_hydrostatic_pressure!(model)
fill_halo_regions!(model.pressures.pHY′)

return nothing
end
end
2 changes: 1 addition & 1 deletion src/MultiRegion/multi_region_abstract_operations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ Base.size(f::MultiRegionAbstractOperation) = size(getregion(f.grid, 1))

@inline isregional(f::MultiRegionAbstractOperation) = true
@inline devices(f::MultiRegionAbstractOperation) = devices(f.grid)
sync_all_devices!(f::MultiRegionAbstractOperation) = sync_all_devices!(devices(f.grid))
@inline sync_all_devices!(f::MultiRegionAbstractOperation) = sync_all_devices!(devices(f.grid))

@inline switch_device!(f::MultiRegionAbstractOperation, d) = switch_device!(f.grid, d)
@inline getdevice(f::MultiRegionAbstractOperation, d) = getdevice(f.grid, d)
Expand Down
6 changes: 5 additions & 1 deletion src/MultiRegion/multi_region_models.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@ import Oceananigans.Diagnostics: accurate_cell_advection_timescale
import Oceananigans.Advection: WENO5
import Oceananigans.Models.HydrostaticFreeSurfaceModels: build_implicit_step_solver, validate_tracer_advection
import Oceananigans.TurbulenceClosures: implicit_diffusion_solver
import Oceananigans.Models.NonhydrostaticModels: PressureSolver

const MultiRegionModel = HydrostaticFreeSurfaceModel{<:Any, <:Any, <:AbstractArchitecture, <:Any, <:MultiRegionGrid}
const MultiRegionModel = Union{HydrostaticFreeSurfaceModel{<:Any, <:Any, <:AbstractArchitecture, <:Any, <:MultiRegionGrid}, NonhydrostaticModel{<:Any, <:Any, <:AbstractArchitecture, <:MultiRegionGrid}}

# Utility to generate the inputs to complex `getregion`s
function getregionalproperties(T, inner=true)
Expand All @@ -26,6 +27,7 @@ function getregionalproperties(T, inner=true)
end

Types = (:HydrostaticFreeSurfaceModel,
:NonhydrostaticModel,
:ImplicitFreeSurface,
:ExplicitFreeSurface,
:QuasiAdamsBashforth2TimeStepper,
Expand All @@ -45,6 +47,8 @@ end

validate_tracer_advection(tracer_advection::MultiRegionObject, grid::MultiRegionGrid) = tracer_advection, NamedTuple()

PressureSolver(arch, ::MultiRegionGrid) = nothing

@inline isregional(mrm::MultiRegionModel) = true
@inline devices(mrm::MultiRegionModel) = devices(mrm.grid)
@inline getdevice(mrm::MultiRegionModel, d) = getdevice(mrm.grid, d)
Expand Down