Merge pull request #124 from CliMA/glw/latest-oceananigans

Updates to Oceananigans 0.68

Author: navidcy

Manifest.toml

@@ -25,7 +25,8 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-Oceananigans = "^0.58"
+Oceananigans = "^0.68.2"
+Oceanostics = "^0.6.3"
 julia = "^1.6"
 
 [extras]

Project.toml

@@ -3,18 +3,14 @@
 # This script runs a simulation of convection driven by cooling at the
 # surface of an idealized, stratified, rotating ocean surface boundary layer.
 
-using LESbrary, Printf, Statistics
+using LESbrary, Printf, Statistics, Oceananigans, Oceananigans.Units
 
 # Domain
 
-using Oceananigans.Grids
-
-grid = RegularRectilinearGrid(size=(32, 32, 32), x=(0, 128), y=(0, 128), z=(-64, 0))
+grid = RectilinearGrid(CPU(), size=(32, 32, 32), x=(0, 128), y=(0, 128), z=(-64, 0))
 
 # Buoyancy and boundary conditions
 
-using Oceananigans.BuoyancyModels, Oceananigans.BoundaryConditions
-
 Qᵇ = 1e-7
 N² = 1e-5
 
@@ -27,50 +23,48 @@ g = buoyancy.gravitational_acceleration
 Qᵀ = Qᵇ / (α * g)
 dTdz = N² / (α * g)
 
-T_bcs = TracerBoundaryConditions(grid, top = BoundaryCondition(Flux, Qᵀ),
-                                       bottom = BoundaryCondition(Gradient, dTdz))
+T_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵀ),
+                                bottom = GradientBoundaryCondition(dTdz))
 
 # LES Model
 
 # Wall-aware AMD model constant which is 'enhanced' near the upper boundary.
 # Necessary to obtain a smooth temperature distribution.
 using LESbrary.NearSurfaceTurbulenceModels: SurfaceEnhancedModelConstant
 
-Cᴬᴹᴰ = SurfaceEnhancedModelConstant(grid.Δz, C₀ = 1/12, enhancement = 7, decay_scale = 4 * grid.Δz)
+Δz = grid.Δzᵃᵃᶜ
+Cᴬᴹᴰ = SurfaceEnhancedModelConstant(Δz, C₀ = 1/12, enhancement = 7, decay_scale = 4Δz)
 
-# Instantiate Oceananigans.IncompressibleModel
+# Instantiate Oceananigans.NonhydrostaticModel
 
 using Oceananigans
 using Oceananigans.Advection: WENO5
 
-model = IncompressibleModel(architecture = CPU(),
-                             timestepper = :RungeKutta3,
-                               advection = WENO5(),
-                                    grid = grid,
-                                 tracers = :T,
-                                buoyancy = buoyancy,
-                                coriolis = FPlane(f=1e-4),
-                                 closure = AnisotropicMinimumDissipation(C=Cᴬᴹᴰ),
-                     boundary_conditions = (T=T_bcs,))
+model = NonhydrostaticModel(; grid, buoyancy,
+                            timestepper = :RungeKutta3,
+                            advection = WENO5(),
+                            tracers = :T,
+                            coriolis = FPlane(f=1e-4),
+                            closure = AnisotropicMinimumDissipation(C=Cᴬᴹᴰ),
+                            boundary_conditions = (; T=T_bcs))
 
 # # Initial condition
 
 Ξ(z) = rand() * exp(z / 8)
-
 Tᵢ(x, y, z) = dTdz * z + 1e-6 * Ξ(z) * dTdz * grid.Lz
-
 set!(model, T=Tᵢ)
 
 # # Prepare the simulation
 
-using Oceananigans.Utils: hour, minute
 using LESbrary.Utils: SimulationProgressMessenger
 
 # Adaptive time-stepping
-wizard = TimeStepWizard(cfl=1.5, Δt=2.0, max_change=1.1, max_Δt=30.0)
 
-simulation = Simulation(model, Δt=wizard, stop_time=8hour, iteration_interval=100,
-                        progress=SimulationProgressMessenger(wizard))
+simulation = Simulation(model, Δt=2.0, stop_time=8hour)
+
+wizard = TimeStepWizard(cfl=1.5, max_change=1.1, max_Δt=30.0)
+simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(100))
+simulation.callbacks[:progress] = Callback(SimulationProgressMessenger(wizard), IterationInterval(100))
 
 # Prepare Output
 
@@ -86,20 +80,19 @@ mkpath(data_directory)
 cp(@__FILE__, joinpath(data_directory, basename(@__FILE__)), force=true)
 
 simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers);
-                                                           schedule = TimeInterval(4hour),
-                                                             prefix = prefix * "_fields",
-                                                                dir = data_directory,
-                                                       max_filesize = 2GiB,
-                                                              force = true)
+                                                      schedule = TimeInterval(4hour),
+                                                      prefix = prefix * "_fields",
+                                                      dir = data_directory,
+                                                      max_filesize = 2GiB,
+                                                      force = true)
 
 simulation.output_writers[:slices] = JLD2OutputWriter(model, merge(model.velocities, model.tracers),
-                                                           schedule = AveragedTimeInterval(1hour, window=15minute),
-                                                             prefix = prefix * "_slices",
-                                                       field_slicer = FieldSlicer(j=floor(Int, grid.Ny/2)),
-                                                                dir = data_directory,
-                                                       max_filesize = 2GiB,
-                                                              force = true)
-
+                                                      schedule = AveragedTimeInterval(1hour, window=15minute),
+                                                      prefix = prefix * "_slices",
+                                                      field_slicer = FieldSlicer(j=floor(Int, grid.Ny/2)),
+                                                      dir = data_directory,
+                                                      max_filesize = 2GiB,
+                                                      force = true)
 
 # Horizontally-averaged turbulence statistics
 turbulence_statistics = LESbrary.TurbulenceStatistics.first_through_second_order(model)
@@ -108,9 +101,9 @@ tke_budget_statistics = LESbrary.TurbulenceStatistics.turbulent_kinetic_energy_b
 simulation.output_writers[:statistics] =
     JLD2OutputWriter(model, merge(turbulence_statistics, tke_budget_statistics),
                      schedule = AveragedTimeInterval(1hour, window=15minute),
-                       prefix = prefix * "_statistics",
-                          dir = data_directory,
-                        force = true)
+                     prefix = prefix * "_statistics",
+                     dir = data_directory,
+                     force = true)
 
 # # Run
 
@@ -151,7 +144,7 @@ advective_flux = - file["timeseries/tke_advective_flux/$iter"][1, 1, :]
 
 transport = zeros(grid.Nz)
 transport = (pressure_flux[2:end] .+ advective_flux[2:end]
-             .- pressure_flux[1:end-1] .- advective_flux[1:end-1]) / grid.Δz
+             .- pressure_flux[1:end-1] .- advective_flux[1:end-1]) / Δz
 
 ## For mixing length calculation
 wT = file["timeseries/wT/$iter"][1, 1, 2:end-1]
@@ -161,7 +154,7 @@ close(file)
 ## Post-process the data to determine the mixing length
 
 ## Mixing length, computed at cell interfaces and omitting boundaries
-Tz = @. (T[2:end] - T[1:end-1]) / grid.Δz
+Tz = @. (T[2:end] - T[1:end-1]) / Δz
 bz = @. α * g * Tz
 eᶠ = @. (e[1:end-1] + e[2:end]) / 2
 
@@ -204,3 +197,4 @@ mixing_length = plot([ℓ_measured ℓ_estimated], zF[2:end-1], size = plot_size
                                                            label = ["measured" "estimated"])
 
 plot(temperature, variances, budget, mixing_length, layout=(1, 4))
+

examples/free_convection.jl

@@ -19,24 +19,13 @@ using Plots
 using ArgParse
 
 using Oceananigans
-using Oceananigans.Advection
-using Oceananigans.Grids
-using Oceananigans.Fields
-using Oceananigans.BuoyancyModels
-using Oceananigans.BoundaryConditions
-using Oceananigans.Diagnostics
-using Oceananigans.OutputWriters
-using Oceananigans.Utils
-
-using LESbrary
-
+using Oceananigans.Units
 using Oceananigans.BuoyancyModels: g_Earth
-using Oceananigans.Fields: PressureField
-using Oceananigans.BuoyancyModels: BuoyancyTracer
-using Oceananigans.SurfaceWaves: UniformStokesDrift
 
-using LESbrary.TurbulenceStatistics: TurbulentKineticEnergy, ShearProduction, ViscousDissipation
+using LESbrary
+using LESbrary.TurbulenceStatistics: TurbulentKineticEnergy, ViscousDissipation
 using LESbrary.TurbulenceStatistics: first_through_second_order, turbulent_kinetic_energy_budget
+using Oceanostics.TKEBudgetTerms: ZShearProduction
 
 "Returns a dictionary of command line arguments."
 function parse_command_line_arguments()
@@ -121,16 +110,16 @@ uˢ(z) = Uˢ * exp(2wavenumber * z)
 ##### Grid
 #####
 
-grid = RegularRectilinearGrid(size=(Nh, Nh, Nz), extent=(Lh, Lh, Lz))
+grid = RectilinearGrid(CPU(), size=(Nh, Nh, Nz), extent=(Lh, Lh, Lz))
 
 #####
 ##### Boundary conditions
 #####
 
-u_bcs = UVelocityBoundaryConditions(grid, top = BoundaryCondition(Flux, Qᵘ))
+u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))
 
-b_bcs = TracerBoundaryConditions(grid, top = BoundaryCondition(Flux, Qᵇ),
-                                       bottom = BoundaryCondition(Gradient, N²))
+b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ),
+                                bottom = GradientBoundaryCondition(N²))
 
 #####
 ##### Sponge layer
@@ -150,19 +139,15 @@ b_sponge = Relaxation(rate = 1/hour,
 
 advection = eval(args["advection-scheme"])()
 
-model = IncompressibleModel(
-           architecture = GPU(),
-            timestepper = :RungeKutta3,
-              advection = advection,
-                   grid = grid,
-                tracers = :b,
-               buoyancy = BuoyancyTracer(),
-               coriolis = FPlane(f=f),
-                closure = AnisotropicMinimumDissipation(),
-          surface_waves = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
-    boundary_conditions = (u=u_bcs, b=b_bcs),
-                forcing = (u=u_sponge, v=v_sponge, w=w_sponge, b=b_sponge),
-)
+model = NonhydrostaticModel(; advection, grid,
+                            timestepper = :RungeKutta3,
+                            tracers = :b,
+                            buoyancy = BuoyancyTracer(),
+                            coriolis = FPlane(f=f),
+                            closure = AnisotropicMinimumDissipation(),
+                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
+                            boundary_conditions = (u=u_bcs, b=b_bcs),
+                            forcing = (u=u_sponge, v=v_sponge, w=w_sponge, b=b_sponge))
 
 #####
 ##### Initial conditions: Stokes drift + stratification + noise
@@ -182,23 +167,26 @@ set!(model, u=uᵢ, w=wᵢ, b=bᵢ)
 ##### Simulation setup
 #####
 
-wizard = TimeStepWizard(cfl=1.0, Δt=1.0, max_change=1.1, max_Δt=30.0)
+simulation = Simulation(model; Δt=1.0, stop_time)
 
-umax = FieldMaximum(abs, model.velocities.u)
-vmax = FieldMaximum(abs, model.velocities.v)
-wmax = FieldMaximum(abs, model.velocities.w)
+wizard = TimeStepWizard(cfl=1.0, max_change=1.1, max_Δt=30.0)
+simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
 
 wall_clock = [time_ns()]
 
 function print_progress(simulation)
     model = simulation.model
 
+    umax = maximum(abs, model.velocities.u)
+    vmax = maximum(abs, model.velocities.v)
+    wmax = maximum(abs, model.velocities.w)
+
     ## Print a progress message
     msg = @sprintf("i: %04d, t: %s, Δt: %s, umax = (%.1e, %.1e, %.1e) ms⁻¹, wall time: %s\n",
-                   model.clock.iteration,
-                   prettytime(model.clock.time),
-                   prettytime(wizard.Δt),
-                   umax(), vmax(), wmax(),
+                   iteration(simulation),
+                   prettytime(simulation),
+                   prettytime(simulation.Δt),
+                   umax, vmax, wmax,
                    prettytime(1e-9 * (time_ns() - wall_clock[1]))
                   )
 
@@ -209,10 +197,7 @@ function print_progress(simulation)
     return nothing
 end
 
-simulation = Simulation(model, iteration_interval = 10,
-                                               Δt = wizard,
-                                        stop_time = stop_time,
-                                         progress = print_progress)
+simulation.callbacks[:progress] = Callback(print_progress, IterationInterval(10))
 
 #####
 ##### Output setup
@@ -224,9 +209,9 @@ data_directory = joinpath(@__DIR__, "..", "data", prefix) # save data in /data/p
 # "Primitive" statistics
 
 b = BuoyancyField(model)
-p = PressureField(model)
-w_scratch = ZFaceField(model.architecture, model.grid)
-c_scratch = CenterField(model.architecture, model.grid)
+p = model.pressures.pHY′ + model.pressures.pNHS
+w_scratch = ZFaceField(model.grid)
+c_scratch = CenterField(model.grid)
 
 primitive_statistics = first_through_second_order(model, b=b, p=p, w_scratch=w_scratch, c_scratch=c_scratch)
 
@@ -236,23 +221,26 @@ V = primitive_statistics[:v]
 # Turbulent kinetic energy budget terms
 
 e = TurbulentKineticEnergy(model, U=U, V=V)
-shear_production = ShearProduction(model, data=c_scratch.data, U=U, V=V)
-dissipation = ViscousDissipation(model, data=c_scratch.data)
+shear_production = ZShearProduction(model, U=U, V=V)
+dissipation = ViscousDissipation(model)
 
 tke_budget_statistics = turbulent_kinetic_energy_budget(model, b=b, p=p, U=U, V=V, e=e,
                                                         shear_production=shear_production, dissipation=dissipation)
 
 statistics_to_output = merge(primitive_statistics, tke_budget_statistics)
 
 fields_to_output = merge(model.velocities, model.tracers,
-                         (νₑ=model.diffusivities.νₑ, e=e, sp=shear_production, ϵ=dissipation))
+                         (νₑ = model.diffusivity_fields.νₑ,
+                           e = Field(e),
+                          sp = Field(shear_production),
+                           ϵ = Field(dissipation)))
 
 # Output configured for pickup
 
 pickup = args["pickup"]
 force = pickup ? false : true
 
-k_xy_slice = searchsortedfirst(grid.zF[:], -slice_depth)
+k_xy_slice = searchsortedfirst(znodes(Face, grid), -slice_depth)
 
 simulation.output_writers[:checkpointer] =
     Checkpointer(model, schedule = TimeInterval(26hour), prefix = prefix * "_checkpointer", dir = data_directory)
@@ -418,8 +406,9 @@ pressure_flux = statistics_file["timeseries/tke_pressure_flux/$last_iteration"][
 
 close(statistics_file)
 
-pressure_flux_divergence = @. (pressure_flux[2:end] - pressure_flux[1:end-1]) ./ grid.Δz
-advective_flux_divergence = @. (advective_flux[2:end] - advective_flux[1:end-1]) ./ grid.Δz
+Δz = grid.Δzᵃᵃᶜ
+pressure_flux_divergence = @. (pressure_flux[2:end] - pressure_flux[1:end-1]) ./ Δz
+advective_flux_divergence = @. (advective_flux[2:end] - advective_flux[1:end-1]) ./ Δz
 
 tendency = @. - pressure_flux_divergence - advective_flux_divergence + shear_production + buoyancy_flux - dissipation