diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2796f7c7d4..f87975bf20 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -132,6 +132,33 @@ steps: architecture: CPU depends_on: "init_cpu" +##### +##### HydrostaticFreeSurfaceModel +##### + + - label: "💧 gpu hydrostatic free surface model tests" + env: + JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" + TEST_GROUP: "hydrostatic_free_surface" + commands: + - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" + agents: + queue: Oceananigans + architecture: GPU + depends_on: "init_gpu" + + - label: "💦 cpu hydrostatic free surface model tests" + env: + JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" + TEST_GROUP: "hydrostatic_free_surface" + CUDA_VISIBLE_DEVICES: "-1" + commands: + - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" + agents: + queue: Oceananigans + architecture: CPU + depends_on: "init_cpu" + ##### ##### ShallowWaterModel ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl new file mode 100644 index 0000000000..c4ab3f663e --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -0,0 +1,40 @@ +module HydrostaticFreeSurfaceModels + +using KernelAbstractions: @index, @kernel, Event, MultiEvent +using KernelAbstractions.Extras.LoopInfo: @unroll + +using Oceananigans.Utils: launch! + +import Oceananigans: fields + +##### +##### HydrostaticFreeSurfaceModel definition +##### + +include("compute_w_from_continuity.jl") +include("hydrostatic_free_surface_tendency_fields.jl") +include("hydrostatic_free_surface_model.jl") +include("show_hydrostatic_free_surface_model.jl") +include("set_hydrostatic_free_surface_model.jl") + +##### +##### Time-stepping HydrostaticFreeSurfaceModels +##### + +""" + fields(model::HydrostaticFreeSurfaceModel) + +Returns a flattened `NamedTuple` of the fields in `model.velocities` and `model.tracers`. +""" +fields(model::HydrostaticFreeSurfaceModel) = merge((u = model.velocities.u, + v = model.velocities.v, + η = model.free_surface.η), + model.tracers) + +include("calculate_hydrostatic_free_surface_barotropic_pressure.jl") +include("hydrostatic_free_surface_tendency_kernel_functions.jl") +include("update_hydrostatic_free_surface_model_state.jl") +include("calculate_hydrostatic_free_surface_tendencies.jl") +include("hydrostatic_free_surface_time_step.jl") + +end # module diff --git a/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_barotropic_pressure.jl b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_barotropic_pressure.jl new file mode 100644 index 0000000000..26dfa6b375 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_barotropic_pressure.jl @@ -0,0 +1,4 @@ +import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities! + +calculate_pressure_correction!(::HydrostaticFreeSurfaceModel, Δt) = nothing +pressure_correct_velocities!(::HydrostaticFreeSurfaceModel, Δt) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl new file mode 100644 index 0000000000..082ec37687 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/calculate_hydrostatic_free_surface_tendencies.jl @@ -0,0 +1,232 @@ +import Oceananigans.TimeSteppers: calculate_tendencies! + +using Oceananigans: fields +using Oceananigans.Utils: work_layout + +""" + calculate_tendencies!(model::IncompressibleModel) + +Calculate the interior and boundary contributions to tendency terms without the +contribution from non-hydrostatic pressure. +""" +function calculate_tendencies!(model::HydrostaticFreeSurfaceModel) + + # Note: + # + # "tendencies" is a NamedTuple of OffsetArrays corresponding to the tendency data for use + # in GPU computations. + # + # "model.timestepper.Gⁿ" is a NamedTuple of Fields, whose data also corresponds to + # tendency data. + + # Calculate contributions to momentum and tracer tendencies from fluxes and volume terms in the + # interior of the domain + calculate_interior_tendency_contributions!(model.timestepper.Gⁿ, + model.architecture, + model.grid, + model.advection, + model.coriolis, + model.buoyancy, + model.closure, + model.velocities, + model.free_surface.η, + model.tracers, + model.pressure.pHY′, + model.diffusivities, + model.forcing, + model.clock) + + # Calculate contributions to momentum and tracer tendencies from user-prescribed fluxes across the + # boundaries of the domain + calculate_boundary_tendency_contributions!(model.timestepper.Gⁿ, + model.architecture, + model.velocities, + model.tracers, + model.clock, + fields(model)) + + return nothing +end + +""" Store previous value of the source term and calculate current source term. """ +function calculate_interior_tendency_contributions!(tendencies, + arch, + grid, + advection, + coriolis, + buoyancy, + closure, + velocities, + free_surface_displacement, + tracers, + hydrostatic_pressure_anomaly, + diffusivities, + forcings, + clock) + + workgroup, worksize = work_layout(grid, :xyz) + xy_workgroup, xy_worksize = work_layout(grid, :xy) + + calculate_Gu_kernel! = calculate_Gu!(device(arch), workgroup, worksize) + calculate_Gv_kernel! = calculate_Gv!(device(arch), workgroup, worksize) + calculate_Gc_kernel! = calculate_Gc!(device(arch), workgroup, worksize) + + calculate_Gη_kernel! = calculate_Gη!(device(arch), xy_workgroup, xy_worksize) + + barrier = Event(device(arch)) + + Gu_event = calculate_Gu_kernel!(tendencies.u, grid, advection.momentum, coriolis, closure, + velocities, free_surface_displacement, tracers, diffusivities, + forcings, hydrostatic_pressure_anomaly, clock, dependencies=barrier) + + Gv_event = calculate_Gv_kernel!(tendencies.v, grid, advection.momentum, coriolis, closure, + velocities, free_surface_displacement, tracers, diffusivities, + forcings, hydrostatic_pressure_anomaly, clock, dependencies=barrier) + + Gη_event = calculate_Gη_kernel!(tendencies.η, grid, closure, + velocities, free_surface_displacement, tracers, diffusivities, + forcings, clock, dependencies=barrier) + + events = [Gu_event, Gv_event, Gη_event] + + for tracer_index in 1:length(tracers) + @inbounds c_tendency = tendencies[tracer_index+3] + @inbounds c_advection = advection[tracer_index+1] + @inbounds forcing = forcings[tracer_index+3] + + Gc_event = calculate_Gc_kernel!(c_tendency, grid, Val(tracer_index), c_advection, closure, buoyancy, + velocities, free_surface_displacement, tracers, diffusivities, + forcing, clock, dependencies=barrier) + + push!(events, Gc_event) + end + + wait(device(arch), MultiEvent(Tuple(events))) + + return nothing +end + +##### +##### Tendency calculators for u, v +##### + +""" Calculate the right-hand-side of the u-velocity equation. """ +@kernel function calculate_Gu!(Gu, + grid, + advection, + coriolis, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + hydrostatic_pressure_anomaly, + clock) + + i, j, k = @index(Global, NTuple) + + @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, advection, coriolis, + closure, velocities, free_surface_displacement, tracers, + diffusivities, forcings, hydrostatic_pressure_anomaly, clock) +end + +""" Calculate the right-hand-side of the v-velocity equation. """ +@kernel function calculate_Gv!(Gv, + grid, + advection, + coriolis, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + hydrostatic_pressure_anomaly, + clock) + + i, j, k = @index(Global, NTuple) + + @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, advection, coriolis, + closure, velocities, free_surface_displacement, tracers, + diffusivities, forcings, hydrostatic_pressure_anomaly, clock) +end + +""" Calculate the right-hand-side of the w-velocity equation. """ +@kernel function calculate_Gη!(Gη, + grid, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + clock) + + i, j = @index(Global, NTuple) + + @inbounds Gη[i, j, 1] = free_surface_tendency(i, j, grid, closure, + velocities, free_surface_displacement, tracers, + diffusivities, forcings, clock) +end + +##### +##### Tracer(s) +##### + +""" Calculate the right-hand-side of the tracer advection-diffusion equation. """ +@kernel function calculate_Gc!(Gc, + grid, + tracer_index, + advection, + closure, + buoyancy, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcing, + clock) + + i, j, k = @index(Global, NTuple) + + @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, tracer_index, advection, closure, + buoyancy, velocities, free_surface_displacement, tracers, + diffusivities, forcing, clock) +end + +##### +##### Boundary contributions to tendencies due to user-prescribed fluxes +##### + +""" Apply boundary conditions by adding flux divergences to the right-hand-side. """ +function calculate_boundary_tendency_contributions!(Gⁿ, arch, velocities, tracers, clock, model_fields) + + barrier = Event(device(arch)) + + events = [] + + # Velocity fields + for i in 1:3 + x_bcs_event = apply_x_bcs!(Gⁿ[i], velocities[i], arch, barrier, clock, model_fields) + y_bcs_event = apply_y_bcs!(Gⁿ[i], velocities[i], arch, barrier, clock, model_fields) + z_bcs_event = apply_z_bcs!(Gⁿ[i], velocities[i], arch, barrier, clock, model_fields) + + push!(events, x_bcs_event, y_bcs_event, z_bcs_event) + end + + # Tracer fields + for i in 4:length(Gⁿ) + x_bcs_event = apply_x_bcs!(Gⁿ[i], tracers[i-3], arch, barrier, clock, model_fields) + y_bcs_event = apply_y_bcs!(Gⁿ[i], tracers[i-3], arch, barrier, clock, model_fields) + z_bcs_event = apply_z_bcs!(Gⁿ[i], tracers[i-3], arch, barrier, clock, model_fields) + + push!(events, x_bcs_event, y_bcs_event, z_bcs_event) + end + + events = filter(e -> typeof(e) <: Event, events) + + wait(device(arch), MultiEvent(Tuple(events))) + + return nothing +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl new file mode 100644 index 0000000000..8b17ddffe9 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl @@ -0,0 +1,25 @@ +using Oceananigans.Architectures: device +using Oceananigans.Operators: div_xyᶜᶜᵃ, ΔzC + +""" +Compute the vertical velocity w by integrating the continuity equation from the bottom upwards + + `w^{n+1} = -∫ [∂/∂x (u^{n+1}) + ∂/∂y (v^{n+1})] dz` +""" +function compute_w_from_continuity!(model) + + event = launch!(model.architecture, model.grid, :xy, _compute_w_from_continuity!, model.velocities, model.grid, + dependencies=Event(device(model.architecture))) + + wait(device(model.architecture), event) + + return nothing +end + +@kernel function _compute_w_from_continuity!(U, grid) + i, j = @index(Global, NTuple) + # U.w[i, j, 1] = 0 is enforced via halo regions. + @unroll for k in 2:grid.Nz + @inbounds U.w[i, j, k] = U.w[i, j, k-1] - ΔzC(i, j, k, grid) * div_xyᶜᶜᵃ(i, j, k-1, grid, U.u, U.v) + end +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl new file mode 100644 index 0000000000..744600ad12 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -0,0 +1,164 @@ +using CUDA: has_cuda +using OrderedCollections: OrderedDict + +using Oceananigans: AbstractModel, AbstractOutputWriter, AbstractDiagnostic + +using Oceananigans.Architectures: AbstractArchitecture, GPU +using Oceananigans.Advection: AbstractAdvectionScheme, CenteredSecondOrder +using Oceananigans.Buoyancy: validate_buoyancy, SeawaterBuoyancy, g_Earth +using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions, TracerBoundaryConditions +using Oceananigans.Fields: Field, CenterField, tracernames, VelocityFields, TracerFields +using Oceananigans.Forcings: model_forcing +using Oceananigans.Grids: with_halo +using Oceananigans.Models.IncompressibleModels: extract_boundary_conditions +using Oceananigans.Solvers: PressureSolver +using Oceananigans.TimeSteppers: Clock, TimeStepper +using Oceananigans.TurbulenceClosures: ν₀, κ₀, with_tracers, DiffusivityFields, IsotropicDiffusivity +using Oceananigans.LagrangianParticleTracking: LagrangianParticles +using Oceananigans.Utils: inflate_halo_size, tupleit + +struct VectorInvariant end + +struct ExplicitFreeSurface{E, T} + η :: E + gravitational_acceleration :: T +end + +ExplicitFreeSurface(; gravitational_acceleration=g_Earth) = + ExplicitFreeSurface(nothing, gravitational_acceleration) + +function FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, arch, grid) + η = CenterField(arch, grid, TracerBoundaryConditions(grid)) + g = convert(eltype(grid), free_surface.gravitational_acceleration) + return ExplicitFreeSurface(η, g) +end + +""" Returns a default_tracer_advection, tracer_advection `tuple`. """ +validate_tracer_advection(invalid_tracer_advection, grid) = error("$invalid_tracer_advection is invalid tracer_advection!") +validate_tracer_advection(tracer_advection_tuple::NamedTuple, grid) = CenteredSecondOrder(), advection_scheme_tuple +validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, grid) = tracer_advection, NamedTuple() + +mutable struct HydrostaticFreeSurfaceModel{TS, E, A<:AbstractArchitecture, + G, T, V, B, R, S, F, P, U, C, Φ, K} <: AbstractModel{TS} + architecture :: A # Computer `Architecture` on which `Model` is run + grid :: G # Grid of physical points on which `Model` is solved + clock :: Clock{T} # Tracks iteration number and simulation time of `Model` + advection :: V # Advection scheme for tracers + buoyancy :: B # Set of parameters for buoyancy model + coriolis :: R # Set of parameters for the background rotation rate of `Model` + free_surface :: S # Free surface parameters and fields + forcing :: F # Container for forcing functions defined by the user + closure :: E # Diffusive 'turbulence closure' for all model fields + particles :: P # Particle set for Lagrangian tracking + velocities :: U # Container for velocity fields `u`, `v`, and `w` + tracers :: C # Container for tracer fields + pressure :: Φ # Container for hydrostatic pressure + diffusivities :: K # Container for turbulent diffusivities + timestepper :: TS # Object containing timestepper fields and parameters +end + +""" + HydrostaticFreeSurfaceModel(; + grid, + architecture = CPU(), + clock = Clock{eltype(grid)}(0, 0, 1), + advection = CenteredSecondOrder(), + buoyancy = SeawaterBuoyancy(eltype(grid)), + coriolis = nothing, + forcing = NamedTuple(), + closure = IsotropicDiffusivity(eltype(grid), ν=ν₀, κ=κ₀), + boundary_conditions = NamedTuple(), + tracers = (:T, :S), + particles = nothing, + velocities = nothing, + pressure = nothing, + diffusivities = nothing, + ) + +Construct an hydrostatic `Oceananigans.jl` model with a free surface on `grid`. + +Keyword arguments +================= + + - `grid`: (required) The resolution and discrete geometry on which `model` is solved. + - `architecture`: `CPU()` or `GPU()`. The computer architecture used to time-step `model`. + - `gravitational_acceleration`: The gravitational acceleration applied to the free surface + - `advection`: The scheme that advects velocities and tracers. See `Oceananigans.Advection`. + - `buoyancy`: The buoyancy model. See `Oceananigans.Buoyancy`. + - `closure`: The turbulence closure for `model`. See `Oceananigans.TurbulenceClosures`. + - `coriolis`: Parameters for the background rotation rate of the model. + - `forcing`: `NamedTuple` of user-defined forcing functions that contribute to solution tendencies. + - `boundary_conditions`: `NamedTuple` containing field boundary conditions. + - `tracers`: A tuple of symbols defining the names of the modeled tracers, or a `NamedTuple` of + preallocated `CenterField`s. +""" +function HydrostaticFreeSurfaceModel(; grid, + architecture::AbstractArchitecture = CPU(), + clock = Clock{eltype(grid)}(0, 0, 1), + momentum_advection = CenteredSecondOrder(), + tracer_advection = CenteredSecondOrder(), + buoyancy = SeawaterBuoyancy(eltype(grid)), + coriolis = nothing, + free_surface = ExplicitFreeSurface(gravitational_acceleration=g_Earth), + forcing::NamedTuple = NamedTuple(), + closure = IsotropicDiffusivity(eltype(grid), ν=ν₀, κ=κ₀), + boundary_conditions::NamedTuple = NamedTuple(), + tracers = (:T, :S), + particles::Union{Nothing, LagrangianParticles} = nothing, + velocities = nothing, + pressure = nothing, + diffusivities = nothing, + ) + + if architecture == GPU() && !has_cuda() + throw(ArgumentError("Cannot create a GPU model. No CUDA-enabled GPU was detected!")) + end + + tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) + validate_buoyancy(buoyancy, tracernames(tracers)) + + # Recursively "regularize" field-dependent boundary conditions by supplying list of tracer names. + # We also regularize boundary conditions included in velocities, tracers, pressure, and diffusivities. + # Note that we do not regularize boundary conditions contained in *tupled* diffusivity fields right now. + embedded_boundary_conditions = merge(extract_boundary_conditions(velocities), + extract_boundary_conditions(tracers), + extract_boundary_conditions(pressure), + extract_boundary_conditions(diffusivities)) + + boundary_conditions = merge(embedded_boundary_conditions, boundary_conditions) + + boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, tracernames(tracers), nothing) + + # Either check grid-correctness, or construct tuples of fields + velocities = VelocityFields(velocities, architecture, grid, boundary_conditions) + tracers = TracerFields(tracers, architecture, grid, boundary_conditions) + pressure = (pHY′ = CenterField(architecture, grid, TracerBoundaryConditions(grid)),) + diffusivities = DiffusivityFields(diffusivities, architecture, grid, + tracernames(tracers), boundary_conditions, closure) + + # Instantiate timestepper if not already instantiated + timestepper = TimeStepper(:QuasiAdamsBashforth2, architecture, grid, tracernames(tracers); + Gⁿ = HydrostaticFreeSurfaceTendencyFields(architecture, grid, tracernames(tracers)), + G⁻ = HydrostaticFreeSurfaceTendencyFields(architecture, grid, tracernames(tracers))) + + free_surface = FreeSurface(free_surface, architecture, grid) + + # Regularize forcing and closure for model tracer and velocity fields. + model_fields = merge((u=velocities.u, v=velocities.v, η=free_surface.η), tracers) + forcing = model_forcing(model_fields; forcing...) + closure = with_tracers(tracernames(tracers), closure) + + default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid) + + # Advection schemes + tracer_advection_tuple = with_tracers(tracernames(tracers), + tracer_advection, + (name, tracer_advection) -> default_tracer_advection, + with_velocities=false) + + advection = merge((momentum=momentum_advection,), tracer_advection_tuple) + + return HydrostaticFreeSurfaceModel(architecture, grid, clock, advection, buoyancy, coriolis, + free_surface, forcing, closure, particles, velocities, tracers, + pressure, diffusivities, timestepper) +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl new file mode 100644 index 0000000000..a434190716 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_fields.jl @@ -0,0 +1,14 @@ +using Oceananigans.BoundaryConditions: UVelocityBoundaryConditions, VVelocityBoundaryConditions, TracerBoundaryConditions +using Oceananigans.Fields: XFaceField, YFaceField, CenterField, TracerFields + +function HydrostaticFreeSurfaceTendencyFields(arch, grid, tracer_names) + + u = XFaceField(arch, grid, UVelocityBoundaryConditions(grid)) + v = YFaceField(arch, grid, VVelocityBoundaryConditions(grid)) + η = CenterField(arch, grid, TracerBoundaryConditions(grid)) + tracers = TracerFields(tracer_names, arch, grid) + + return merge((u=u, v=v, η=η), tracers) +end + + diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl new file mode 100644 index 0000000000..b644b0d51e --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -0,0 +1,173 @@ +using Oceananigans.Advection +using Oceananigans.Buoyancy +using Oceananigans.Coriolis +using Oceananigans.Operators +using Oceananigans.SurfaceWaves +using Oceananigans.TurbulenceClosures: ∂ⱼ_2ν_Σ₁ⱼ, ∂ⱼ_2ν_Σ₂ⱼ, ∂ⱼ_2ν_Σ₃ⱼ, ∇_κ_∇c + +# Temporary implementation of momentum advection for hydrostatic free surface model +U_dot_∇u(i, j, k, grid, advection::VectorInvariant, velocities) = 0 +U_dot_∇v(i, j, k, grid, advection::VectorInvariant, velocities) = 0 + +U_dot_∇u(i, j, k, grid, advection::AbstractAdvectionScheme, velocities) = div_Uu(i, j, k, grid, advection, velocities, velocities.u) +U_dot_∇v(i, j, k, grid, advection::AbstractAdvectionScheme, velocities) = div_Uv(i, j, k, grid, advection, velocities, velocities.v) + +""" + hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, + advection, + coriolis, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + hydrostatic_pressure_anomaly, + clock) + +Return the tendency for the horizontal velocity in the x-direction, or the east-west +direction, ``u``, at grid point `i, j, k` for a HydrostaticFreeSurfaceModel. + +The tendency for ``u`` is called ``G_u`` and defined via + + ``∂_t u = G_u - ∂_x ϕ_n`` + +where ∂_x ϕ_n is the barotropic pressure gradient associated with +free surface displacement in the x-direction. +""" +@inline function hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, + advection, + coriolis, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + hydrostatic_pressure_anomaly, + clock) + + return ( - U_dot_∇u(i, j, k, grid, advection, velocities) + - ∂xᶠᵃᵃ(i, j, k, grid, hydrostatic_pressure_anomaly) + + ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, clock, closure, velocities, diffusivities) + + forcings.u(i, j, k, grid, clock, merge(velocities, (η=free_surface_displacement,), tracers))) +end + +""" + hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, + advection, + coriolis, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + hydrostatic_pressure_anomaly, + clock) + +Return the tendency for the horizontal velocity in the y-direction, or the east-west +direction, ``v``, at grid point `i, j, k` for a HydrostaticFreeSurfaceModel. + +The tendency for ``v`` is called ``G_v`` and defined via + + ``∂_t v = G_v - ∂_y ϕ_n`` + +where ∂_y ϕ_n is the barotropic pressure gradient associated with +free surface displacement in the y-direction. +""" +@inline function hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, + advection, + coriolis, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + hydrostatic_pressure_anomaly, + clock) + + return ( - U_dot_∇v(i, j, k, grid, advection, velocities) + - y_f_cross_U(i, j, k, grid, coriolis, velocities) + - ∂yᵃᶠᵃ(i, j, k, grid, hydrostatic_pressure_anomaly) + + ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, clock, closure, velocities, diffusivities) + + forcings.v(i, j, k, grid, clock, merge(velocities, (η=free_surface_displacement,), tracers))) +end + +""" + hydrostatic_free_surface_tracer_tendency(i, j, k, grid, + val_tracer_index::Val{tracer_index}, + advection, + closure, + buoyancy, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcing, + clock) + +Return the tendency for a tracer field with index `tracer_index` +at grid point `i, j, k`. + +The tendency is called ``G_c`` and defined via + + ``∂_t c = G_c`` + +where `c = C[tracer_index]`. +""" +@inline function hydrostatic_free_surface_tracer_tendency(i, j, k, grid, + val_tracer_index::Val{tracer_index}, + advection, + closure, + buoyancy, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcing, + clock) where tracer_index + + @inbounds c = tracers[tracer_index] + + return ( - div_Uc(i, j, k, grid, advection, velocities, c) + + ∇_κ_∇c(i, j, k, grid, clock, closure, c, val_tracer_index, diffusivities, tracers, buoyancy) + + forcing(i, j, k, grid, clock, merge(velocities, (η=free_surface_displacement,), tracers))) +end + +@inline ∇_ν_∇η(args...) = 0 + +""" + free_surface_tendency(i, j, k, grid, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcing, + clock) + +Return the tendency for a tracer field with index `tracer_index` +at grid point `i, j, k`. + +The tendency is called ``G_c`` and defined via + + ``∂_t c = G_c`` + +where `c = C[tracer_index]`. +""" +@inline function free_surface_tendency(i, j, grid, + closure, + velocities, + free_surface_displacement, + tracers, + diffusivities, + forcings, + clock) + + k_surface = grid.Nz + 1 + + return @inbounds ( - grid.Lz * velocities.w[i, j, k_surface] + + ∇_ν_∇η(i, j, grid, clock, closure, free_surface_displacement, diffusivities) + + forcings.η(i, j, k_surface, grid, clock, merge(velocities, (η=free_surface_displacement,), tracers))) +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_time_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_time_step.jl new file mode 100644 index 0000000000..7f3d7aa63a --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_time_step.jl @@ -0,0 +1,83 @@ +using Oceananigans.TimeSteppers: store_tendencies!, update_particle_properties! + +import Oceananigans.TimeSteppers: time_step!, tick!, ab2_step! + +""" + time_step!(model::HydrostaticFreeSurfaceModel, Δt; euler=false) + +Step forward `model` one time step `Δt` with a 2nd-order Adams-Bashforth method and +pressure-correction substep. Setting `euler=true` will take a forward Euler time step. +""" +function time_step!(model::HydrostaticFreeSurfaceModel, Δt; euler=false) + + Δt == 0 && @warn "Δt == 0 may cause model blowup!" + + χ = ifelse(euler, convert(eltype(model.grid), -0.5), model.timestepper.χ) + + # Be paranoid and update state at iteration 0, in case run! is not used: + model.clock.iteration == 0 && update_state!(model) + + calculate_tendencies!(model) + + #ab2_step!(model, Δt, χ) # full step for tracers, fractional step for velocities. + + calculate_pressure_correction!(model, Δt) + pressure_correct_velocities!(model, Δt) + + update_state!(model) + store_tendencies!(model) + update_particle_properties!(model, Δt) + + tick!(model.clock, Δt) + + return nothing +end + +#= +##### +##### Time stepping in each step +##### + +function ab2_step!(model, Δt, χ) + + workgroup, worksize = work_layout(model.grid, :xyz) + + barrier = Event(device(model.architecture)) + + step_field_kernel! = ab2_step_field!(device(model.architecture), workgroup, worksize) + + model_fields = fields(model) + + events = [] + + for (i, field) in enumerate(model_fields) + + field_event = step_field_kernel!(field, Δt, χ, + model.timestepper.Gⁿ[i], + model.timestepper.G⁻[i], + dependencies=Event(device(model.architecture))) + + push!(events, field_event) + end + + wait(device(model.architecture), MultiEvent(Tuple(events))) + + return nothing +end + +""" +Time step via + + `U^{n+1} = U^n + Δt ( (3/2 + χ) * G^{n} - (1/2 + χ) G^{n-1} )` + +""" + +@kernel function ab2_step_field!(U, Δt, χ::FT, Gⁿ, G⁻) where FT + i, j, k = @index(Global, NTuple) + + @inbounds begin + U[i, j, k] += Δt * ( (FT(1.5) + χ) * Gⁿ[i, j, k] - (FT(0.5) + χ) * G⁻[i, j, k] ) + + end +end +=# diff --git a/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl new file mode 100644 index 0000000000..9e545f89f8 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl @@ -0,0 +1,49 @@ +using Oceananigans.TimeSteppers: update_state! + +import Oceananigans.Fields: set! + +""" + set!(model; kwargs...) + +Set velocity and tracer fields of `model`. The keyword arguments +`kwargs...` take the form `name=data`, where `name` refers to one of the +fields of `model.velocities` or `model.tracers`, and the `data` may be an array, +a function with arguments `(x, y, z)`, or any data type for which a +`set!(ϕ::AbstractField, data)` function exists. + +Example +======= +```julia +model = HydrostaticFreeSurfaceModel(grid=RegularCartesianGrid(size=(32, 32, 32), length=(1, 1, 1)) + +# Set u to a parabolic function of z, v to random numbers damped +# at top and bottom, and T to some silly array of half zeros, +# half random numbers. + +u₀(x, y, z) = z/model.grid.Lz * (1 + z/model.grid.Lz) +v₀(x, y, z) = 1e-3 * rand() * u₀(x, y, z) + +T₀ = rand(size(model.grid)...) +T₀[T₀ .< 0.5] .= 0 + +set!(model, u=u₀, v=v₀, T=T₀) +``` +""" +function set!(model::HydrostaticFreeSurfaceModel; kwargs...) + for (fldname, value) in kwargs + if fldname ∈ propertynames(model.velocities) + ϕ = getproperty(model.velocities, fldname) + elseif fldname ∈ propertynames(model.tracers) + ϕ = getproperty(model.tracers, fldname) + elseif fldname ∈ propertynames(model.free_surface) + ϕ = getproperty(model.free_surface, fldname) + else + throw(ArgumentError("name $fldname not found in model.velocities, model.tracers, or model.free_surface.")) + end + set!(ϕ, value) + end + + update_state!(model) + + return nothing +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl new file mode 100644 index 0000000000..e170937bb6 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/show_hydrostatic_free_surface_model.jl @@ -0,0 +1,21 @@ +using Oceananigans.Utils: prettytime, ordered_dict_show +using Oceananigans: short_show + +"""Show the innards of a `Model` in the REPL.""" +function Base.show(io::IO, model::HydrostaticFreeSurfaceModel{TS, C, A}) where {TS, C, A} + print(io, "HydrostaticFreeSurfaceModel{$A, $(eltype(model.grid))}", + "(time = $(prettytime(model.clock.time)), iteration = $(model.clock.iteration)) \n", + "├── grid: $(short_show(model.grid))\n", + "├── tracers: $(tracernames(model.tracers))\n", + "├── closure: $(typeof(model.closure))\n", + "├── buoyancy: $(typeof(model.buoyancy))\n") + + if isnothing(model.particles) + print(io, "└── coriolis: $(typeof(model.coriolis))") + else + particles = model.particles.properties + properties = propertynames(particles) + print(io, "├── coriolis: $(typeof(model.coriolis))\n") + print(io, "└── particles: $(length(particles)) Lagrangian particles with $(length(properties)) properties: $properties") + end +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl new file mode 100644 index 0000000000..52f26e728e --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -0,0 +1,40 @@ +using Oceananigans.Architectures +using Oceananigans.BoundaryConditions +using Oceananigans.TurbulenceClosures: calculate_diffusivities! +using Oceananigans.Models.IncompressibleModels: update_hydrostatic_pressure! + +import Oceananigans.TimeSteppers: update_state! + +""" + update_state!(model::HydrostaticFreeSurfaceModel) + +Update peripheral aspects of the model (halo regions, diffusivities, hydrostatic pressure) to the current model state. +""" +function update_state!(model::HydrostaticFreeSurfaceModel) + + # Fill halos for velocities and tracers + fill_halo_regions!(fields(model), model.architecture, model.clock, fields(model)) + + compute_w_from_continuity!(model) + + fill_halo_regions!(model.velocities.w, model.architecture, model.clock, fields(model)) + + # Calculate diffusivities + calculate_diffusivities!(model.diffusivities, model.architecture, model.grid, model.closure, + model.buoyancy, model.velocities, model.tracers) + + fill_halo_regions!(model.diffusivities, model.architecture, model.clock, fields(model)) + + # Calculate hydrostatic pressure + pressure_calculation = launch!(model.architecture, model.grid, :xy, update_hydrostatic_pressure!, + model.pressure.pHY′, model.grid, model.buoyancy, model.tracers, + dependencies=Event(device(model.architecture))) + + # Fill halo regions for pressure + wait(device(model.architecture), pressure_calculation) + + fill_halo_regions!(model.pressure.pHY′, model.architecture) + + return nothing +end + diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 8250066145..1f0e410793 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -2,10 +2,16 @@ module Models export IncompressibleModel, NonDimensionalModel +using Oceananigans: AbstractModel + +abstract type AbstractIncompressibleModel{TS} <: AbstractModel{TS} end + include("IncompressibleModels/IncompressibleModels.jl") +include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl") include("ShallowWaterModels/ShallowWaterModels.jl") using .IncompressibleModels: IncompressibleModel, NonDimensionalModel +using .HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel using .ShallowWaterModels: ShallowWaterModel end diff --git a/test/runtests.jl b/test/runtests.jl index 90b5077701..60b6500571 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -136,4 +136,8 @@ group = get(ENV, "TEST_GROUP", :all) |> Symbol if group == :shallow_water || group == :all include("test_shallow_water_models.jl") end + + if group == :hydrostatic_free_surface || group == :all + include("test_hydrostatic_free_surface_models.jl") + end end diff --git a/test/test_hydrostatic_free_surface_models.jl b/test/test_hydrostatic_free_surface_models.jl new file mode 100644 index 0000000000..737f558707 --- /dev/null +++ b/test/test_hydrostatic_free_surface_models.jl @@ -0,0 +1,117 @@ +using Oceananigans: CPU, GPU +using Oceananigans.Models: HydrostaticFreeSurfaceModel +using Oceananigans.Grids: Periodic, Bounded + +function time_stepping_hydrostatic_free_surface_model_works(arch, topo, coriolis) + grid = RegularCartesianGrid(size=(1, 1, 1), extent=(2π, 2π, 2π), topology=topo) + model = HydrostaticFreeSurfaceModel(grid=grid, architecture=arch, coriolis=coriolis) + simulation = Simulation(model, Δt=1.0, stop_iteration=1) + run!(simulation) + + return model.clock.iteration == 1 +end + +function hydrostatic_free_surface_model_tracers_and_forcings_work(arch) + grid = RegularCartesianGrid(size=(1, 1, 1), extent=(2π, 2π, 2π)) + model = HydrostaticFreeSurfaceModel(grid=grid, architecture=arch, tracers=(:T, :S, :c, :d)) + + @test model.tracers.T isa Field + @test model.tracers.S isa Field + @test model.tracers.c isa Field + @test model.tracers.d isa Field + + @test haskey(model.forcing, :u) + @test haskey(model.forcing, :v) + @test haskey(model.forcing, :η) + @test haskey(model.forcing, :T) + @test haskey(model.forcing, :S) + @test haskey(model.forcing, :c) + @test haskey(model.forcing, :d) + + simulation = Simulation(model, Δt=1.0, stop_iteration=1) + run!(simulation) + + @test model.clock.iteration == 1 + + return nothing +end + +@testset "Hydrostatic free surface Models" begin + @info "Testing hydrostatic free surface models..." + + @testset "Model constructor errors" begin + grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1)) + @test_throws TypeError HydrostaticFreeSurfaceModel(architecture=CPU, grid=grid) + @test_throws TypeError HydrostaticFreeSurfaceModel(architecture=GPU, grid=grid) + end + + topos = ( + (Periodic, Periodic, Bounded), + (Periodic, Bounded, Bounded), + (Bounded, Bounded, Bounded), + ) + + for topo in topos + @testset "$topo model construction" begin + @info " Testing $topo model construction..." + for arch in archs, FT in float_types + grid = RegularCartesianGrid(FT, topology=topo, size=(1, 1, 1), extent=(1, 2, 3)) + model = HydrostaticFreeSurfaceModel(grid=grid, architecture=arch) + + # Just testing that the model was constructed with no errors/crashes. + @test model isa HydrostaticFreeSurfaceModel + + # Test that the grid didn't get mangled (sort of) + @test size(grid) === size(model.grid) + end + end + end + + @testset "Setting HydrostaticFreeSurfaceModel fields" begin + @info " Testing setting hydrostatic free surface model fields..." + for arch in archs, FT in float_types + N = (4, 4, 1) + L = (2π, 3π, 5π) + + grid = RegularCartesianGrid(FT, size=N, extent=L) + model = HydrostaticFreeSurfaceModel(grid=grid, architecture=arch) + + x, y, z = nodes((Face, Center, Center), model.grid, reshape=true) + + u₀(x, y, z) = x * y^2 + u_answer = @. x * y^2 + + η₀ = rand(size(grid)...) + η_answer = deepcopy(η₀) + + set!(model, u=u₀, η=η₀) + + u, v, w = model.velocities + η = model.free_surface.η + + @test all(interior(u) .≈ u_answer) + @test all(interior(η) .≈ η_answer) + end + end + + for arch in archs + for topo in topos + @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $topo]" begin + @info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $topo]..." + @test time_stepping_hydrostatic_free_surface_model_works(arch, topo, nothing) + end + end + + for coriolis in (nothing, FPlane(f=1), BetaPlane(f₀=1, β=0.1)) + @testset "Time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]" begin + @info " Testing time-stepping HydrostaticFreeSurfaceModels [$arch, $(typeof(coriolis))]..." + @test time_stepping_hydrostatic_free_surface_model_works(arch, topos[1], coriolis) + end + end + + @testset "HydrostaticFreeSurfaceModel with tracers and forcings [$arch]" begin + @info " Testing HydrostaticFreeSurfaceModel with tracers and forcings [$arch]..." + hydrostatic_free_surface_model_tracers_and_forcings_work(arch) + end + end +end