Merge pull request #1349 from CliMA/glw/hydrostatic-free-surface-model

Begins the implementation of Oceananigans.HydrostaticFreeSurfaceModel for general circulation modeling

.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
 #####

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

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

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

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