Add a `biogeochemistry kwarg to ocean_simulation

experiments/three_degree_simulation/three_degree_simulation.jl

@@ -0,0 +1,72 @@
+using ClimaOcean
+using OrthogonalSphericalShellGrids
+using Oceananigans
+using Oceananigans.Units
+using CFTime
+using Dates
+using Printf
+
+arch = CPU()
+Nx = 120
+Ny = 60
+Nz = 50
+
+z_faces = exponential_z_faces(; Nz, depth=6000, h=34)
+
+underlying_grid = TripolarGrid(arch; size=(Nx, Ny, Nz), z=z_faces)
+bottom_height = regrid_bathymetry(underlying_grid)
+grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
+
+gm = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=4000, κ_symmetric=4000)
+catke = ClimaOcean.OceanSimulations.default_ocean_closure()
+viscous_closure = Oceananigans.TurbulenceClosures.HorizontalScalarDiffusivity(ν=4000)
+
+dates = DateTimeProlepticGregorian(1993, 1, 1) : Month(1) : DateTimeProlepticGregorian(1993, 12, 1)
+temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly())
+salinity = ECCOMetadata(:salinity, dates, ECCO4Monthly())
+
+restoring_rate  = 1/2days
+mask = LinearlyTaperedPolarMask(southern=(-80, -70), northern=(70, 90))
+FT = ECCORestoring(temperature, grid; mask, rate=restoring_rate)
+FS = ECCORestoring(salinity, grid; mask, rate=restoring_rate)
+
+ocean = ocean_simulation(grid;
+                         momentum_advection =  VectorInvariant(),
+                         tracer_advection = Centered(order=2),
+                         closure = (gm, catke, viscous_closure),
+                         forcing = (T=FT, S=FT),
+                         tracers = (:T, :S, :e))
+
+set!(ocean.model, T=ECCOMetadata(:temperature; dates=first(dates)),
+                  S=ECCOMetadata(:salinity; dates=first(dates)))
+
+radiation  = Radiation(arch)
+atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(20))
+coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) 
+simulation = Simulation(coupled_model; Δt=20minutes, stop_time=30days)
+
+wall_time = Ref(time_ns())
+
+function progress(sim)
+    ocean = sim.model.ocean
+    u, v, w = ocean.model.velocities
+    T = ocean.model.tracers.T
+    Tmax = maximum(interior(T))
+    Tmin = minimum(interior(T))
+    umax = (maximum(abs, interior(u)), maximum(abs, interior(v)), maximum(abs, interior(w)))
+    step_time = 1e-9 * (time_ns() - wall_time[])
+
+    @info @sprintf("Time: %s, n: %d, Δt: %s, max|u|: (%.2e, %.2e, %.2e) m s⁻¹, \
+                   extrema(T): (%.2f, %.2f) ᵒC, wall time: %s \n",
+                   prettytime(sim), iteration(sim), prettytime(sim.Δt),
+                   umax..., Tmax, Tmin, prettytime(step_time))
+
+    wall_time[] = time_ns()
+
+    return nothing
+end
+
+add_callback!(simulation, progress, IterationInterval(10))
+
+run!(simulation)
+

src/OceanSimulations/OceanSimulations.jl

@@ -73,6 +73,12 @@ default_tracer_advection() = FluxFormAdvection(WENO(order=7),
 @inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, fields) 
 @inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, fields) 
 
+function add_required_boundary_conditions(user_boundary_conditions, grid, bottom_drag_coefficient)
+
+    return boundary_conditions
+end
+
+
 # TODO: Specify the grid to a grid on the sphere; otherwise we can provide a different
 # function that requires latitude and longitude etc for computing coriolis=FPlane...
 function ocean_simulation(grid;
@@ -85,10 +91,12 @@ function ocean_simulation(grid;
                           gravitational_acceleration = g_Earth,
                           bottom_drag_coefficient = Default(0.003),
                           forcing = NamedTuple(),
+                          biogeochemistry = nothing,
                           timestepper = :QuasiAdamsBashforth2,
                           coriolis = Default(HydrostaticSphericalCoriolis(; rotation_rate)),
                           momentum_advection = default_momentum_advection(),
                           equation_of_state = TEOS10EquationOfState(; reference_density),
+                          boundary_conditions::NamedTuple = NamedTuple(),
                           tracer_advection = default_tracer_advection(),
                           verbose = false)
 
@@ -136,7 +144,7 @@ function ocean_simulation(grid;
     end
 
     bottom_drag_coefficient = convert(FT, bottom_drag_coefficient)
-    
+
     # Set up boundary conditions using Field
     top_zonal_momentum_flux      = τx = Field{Face, Center, Nothing}(grid)
     top_meridional_momentum_flux = τy = Field{Center, Face, Nothing}(grid)
@@ -148,18 +156,19 @@ function ocean_simulation(grid;
     v_top_bc = FluxBoundaryCondition(τy)
     T_top_bc = FluxBoundaryCondition(Jᵀ)
     S_top_bc = FluxBoundaryCondition(Jˢ)
-
+        
     u_bot_bc = FluxBoundaryCondition(u_quadratic_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
     v_bot_bc = FluxBoundaryCondition(v_quadratic_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient)
 
-    ocean_boundary_conditions = (u = FieldBoundaryConditions(top=u_top_bc, bottom=u_bot_bc, immersed=u_immersed_bc),
-                                 v = FieldBoundaryConditions(top=v_top_bc, bottom=v_bot_bc, immersed=v_immersed_bc),
-                                 T = FieldBoundaryConditions(top=T_top_bc),
-                                 S = FieldBoundaryConditions(top=S_top_bc))
+    default_boundary_conditions = (u = FieldBoundaryConditions(top=u_top_bc, bottom=u_bot_bc, immersed=u_immersed_bc),
+                                   v = FieldBoundaryConditions(top=v_top_bc, bottom=v_bot_bc, immersed=v_immersed_bc),
+                                   T = FieldBoundaryConditions(top=T_top_bc),
+                                   S = FieldBoundaryConditions(top=S_top_bc))
 
-    # Use the TEOS10 equation of state
-    teos10 = TEOS10EquationOfState(; reference_density)
-    buoyancy = SeawaterBuoyancy(; gravitational_acceleration, equation_of_state=teos10)
+    # Merge boundary conditions with preference to user
+    # TODO: support users specifying only _part_ of the bcs for u, v, T, S (ie adding the top and immersed
+    # conditions even when a user-bc is supplied).
+    boundary_conditions = merge(default_boundary_conditions, boundary_conditions)
 
     buoyancy = SeawaterBuoyancy(; gravitational_acceleration, equation_of_state)
 
@@ -183,14 +192,15 @@ function ocean_simulation(grid;
     ocean_model = HydrostaticFreeSurfaceModel(; grid,
                                               buoyancy,
                                               closure,
+                                              biogeochemistry,
                                               tracer_advection,
                                               momentum_advection,
                                               tracers,
                                               timestepper,
                                               free_surface,
                                               coriolis,
                                               forcing,
-                                              boundary_conditions = ocean_boundary_conditions)
+                                              boundary_conditions)
 
     ocean = Simulation(ocean_model; Δt, verbose)