/
nonhydrostatic_model.jl
225 lines (192 loc) 路 12.9 KB
/
nonhydrostatic_model.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
using CUDA: has_cuda
using OrderedCollections: OrderedDict
using Oceananigans.Architectures: AbstractArchitecture
using Oceananigans.DistributedComputations: Distributed
using Oceananigans.Advection: CenteredSecondOrder
using Oceananigans.BuoyancyModels: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy
using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields
using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
using Oceananigans.Fields: BackgroundFields, Field, tracernames, VelocityFields, TracerFields, PressureFields
using Oceananigans.Forcings: model_forcing
using Oceananigans.Grids: inflate_halo_size, with_halo, architecture
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
using Oceananigans.Models: AbstractModel, NaNChecker, extract_boundary_conditions
using Oceananigans.Solvers: FFTBasedPoissonSolver
using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles
using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, DiffusivityFields, time_discretization, implicit_diffusion_solver
using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: FlavorOfCATKE
using Oceananigans.Utils: tupleit
using Oceananigans.Grids: topology
import Oceananigans.Architectures: architecture
import Oceananigans.Models: total_velocities, default_nan_checker, timestepper
const ParticlesOrNothing = Union{Nothing, AbstractLagrangianParticles}
const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry}
mutable struct NonhydrostaticModel{TS, E, A<:AbstractArchitecture, G, T, B, R, SD, U, C, 桅, F,
V, S, K, BG, P, BGC, I, AF} <: 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 velocities _and_ tracers
buoyancy :: B # Set of parameters for buoyancy model
coriolis :: R # Set of parameters for the background rotation rate of `Model`
stokes_drift :: SD # Set of parameters for surfaces waves via the Craik-Leibovich approximation
forcing :: F # Container for forcing functions defined by the user
closure :: E # Diffusive 'turbulence closure' for all model fields
background_fields :: BG # Background velocity and tracer fields
particles :: P # Particle set for Lagrangian tracking
biogeochemistry :: BGC # Biogeochemistry for Oceananigans tracers
velocities :: U # Container for velocity fields `u`, `v`, and `w`
tracers :: C # Container for tracer fields
pressures :: 桅 # Container for hydrostatic and nonhydrostatic pressure
diffusivity_fields :: K # Container for turbulent diffusivities
timestepper :: TS # Object containing timestepper fields and parameters
pressure_solver :: S # Pressure/Poisson solver
immersed_boundary :: I # Models the physics of immersed boundaries within the grid
auxiliary_fields :: AF # User-specified auxiliary fields for forcing functions and boundary conditions
end
"""
NonhydrostaticModel(; grid,
clock = Clock{eltype(grid)}(time = 0),
advection = CenteredSecondOrder(),
buoyancy = nothing,
coriolis = nothing,
stokes_drift = nothing,
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
tracers = (),
timestepper = :QuasiAdamsBashforth2,
background_fields::NamedTuple = NamedTuple(),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
velocities = nothing,
pressures = nothing,
diffusivity_fields = nothing,
pressure_solver = nothing,
immersed_boundary = nothing,
auxiliary_fields = NamedTuple())
Construct a model for a non-hydrostatic, incompressible fluid on `grid`, using the Boussinesq
approximation when `buoyancy != nothing`. By default, all Bounded directions are rigid and impenetrable.
Keyword arguments
=================
- `grid`: (required) The resolution and discrete geometry on which the `model` is solved. The
architecture (CPU/GPU) that the model is solved on is inferred from the architecture
of the `grid`. Note that the grid needs to be regularly spaced in the horizontal
dimensions, ``x`` and ``y``.
- `advection`: The scheme that advects velocities and tracers. See `Oceananigans.Advection`.
- `buoyancy`: The buoyancy model. See `Oceananigans.BuoyancyModels`.
- `coriolis`: Parameters for the background rotation rate of the model.
- `stokes_drift`: Parameters for Stokes drift fields associated with surface waves. Default: `nothing`.
- `forcing`: `NamedTuple` of user-defined forcing functions that contribute to solution tendencies.
- `closure`: The turbulence closure for `model`. See `Oceananigans.TurbulenceClosures`.
- `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.
- `timestepper`: A symbol that specifies the time-stepping method. Either `:QuasiAdamsBashforth2` or
`:RungeKutta3`.
- `background_fields`: `NamedTuple` with background fields (e.g., background flow). Default: `nothing`.
- `particles`: Lagrangian particles to be advected with the flow. Default: `nothing`.
- `biogeochemistry`: Biogeochemical model for `tracers`.
- `velocities`: The model velocities. Default: `nothing`.
- `pressures`: Hydrostatic and non-hydrostatic pressure fields. Default: `nothing`.
- `diffusivity_fields`: Diffusivity fields. Default: `nothing`.
- `pressure_solver`: Pressure solver to be used in the model. If `nothing` (default), the model constructor
chooses the default based on the `grid` provide.
- `immersed_boundary`: The immersed boundary. Default: `nothing`.
- `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing`
"""
function NonhydrostaticModel(; grid,
clock = Clock{eltype(grid)}(time = 0),
advection = CenteredSecondOrder(),
buoyancy = nothing,
coriolis = nothing,
stokes_drift = nothing,
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
tracers = (),
timestepper = :QuasiAdamsBashforth2,
background_fields::NamedTuple = NamedTuple(),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
velocities = nothing,
pressures = nothing,
diffusivity_fields = nothing,
pressure_solver = nothing,
immersed_boundary = nothing,
auxiliary_fields = NamedTuple())
arch = architecture(grid)
tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)
# We don't support CAKTE for NonhydrostaticModel yet.
closure = validate_closure(closure)
first_closure = closure isa Tuple ? first(closure) : closure
first_closure isa FlavorOfCATKE &&
error("CATKEVerticalDiffusivity is not supported for " *
"NonhydrostaticModel --- yet!")
tracers, auxiliary_fields = validate_biogeochemistry(tracers, merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)), biogeochemistry, grid, clock)
validate_buoyancy(buoyancy, tracernames(tracers))
buoyancy = regularize_buoyancy(buoyancy)
# Adjust halos when the advection scheme or turbulence closure requires it.
# Note that halos are isotropic by default; however we respect user-input here
# by adjusting each (x, y, z) halo individually.
grid = inflate_grid_halo_size(grid, advection, closure)
# 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
# staggered locations. All other fields are regularized assuming they are tracers.
# Note that we do not regularize boundary conditions contained in *tupled* diffusivity fields right now.
# First, we extract boundary conditions that are embedded within any _user-specified_ field tuples:
embedded_boundary_conditions = merge(extract_boundary_conditions(velocities),
extract_boundary_conditions(tracers),
extract_boundary_conditions(pressures),
extract_boundary_conditions(diffusivity_fields))
# Next, we form a list of default boundary conditions:
prognostic_field_names = (:u, :v, :w, tracernames(tracers)..., keys(auxiliary_fields)...)
default_boundary_conditions = NamedTuple{prognostic_field_names}(FieldBoundaryConditions() for name in prognostic_field_names)
# Finally, we merge specified, embedded, and default boundary conditions. Specified boundary conditions
# have precedence, followed by embedded, followed by default.
boundary_conditions = merge(default_boundary_conditions, embedded_boundary_conditions, boundary_conditions)
boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, prognostic_field_names)
# Ensure `closure` describes all tracers
closure = with_tracers(tracernames(tracers), closure)
# Either check grid-correctness, or construct tuples of fields
velocities = VelocityFields(velocities, grid, boundary_conditions)
tracers = TracerFields(tracers, grid, boundary_conditions)
pressures = PressureFields(pressures, grid, boundary_conditions)
diffusivity_fields = DiffusivityFields(diffusivity_fields, grid, tracernames(tracers), boundary_conditions, closure)
if isnothing(pressure_solver)
pressure_solver = PressureSolver(arch, grid)
end
# Materialize background fields
background_fields = BackgroundFields(background_fields, tracernames(tracers), grid, clock)
# Instantiate timestepper if not already instantiated
implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid)
timestepper = TimeStepper(timestepper, grid, tracernames(tracers), implicit_solver=implicit_solver)
# Regularize forcing for model tracer and velocity fields.
model_fields = merge(velocities, tracers, auxiliary_fields)
forcing = model_forcing(model_fields; forcing...)
model = NonhydrostaticModel(arch, grid, clock, advection, buoyancy, coriolis, stokes_drift,
forcing, closure, background_fields, particles, biogeochemistry, velocities, tracers,
pressures, diffusivity_fields, timestepper, pressure_solver, immersed_boundary,
auxiliary_fields)
update_state!(model)
return model
end
architecture(model::NonhydrostaticModel) = model.architecture
function inflate_grid_halo_size(grid, tendency_terms...)
user_halo = grid.Hx, grid.Hy, grid.Hz
required_halo = Hx, Hy, Hz = inflate_halo_size(user_halo..., grid, tendency_terms...)
if any(user_halo .< required_halo) # Replace grid
@warn "Inflating model grid halo size to ($Hx, $Hy, $Hz) and recreating grid. " *
"Note that an ImmersedBoundaryGrid requires an extra halo point in all non-flat directions compared to a non-immersed boundary grid."
"The model grid will be different from the input grid. To avoid this warning, " *
"pass halo=($Hx, $Hy, $Hz) when constructing the grid."
grid = with_halo((Hx, Hy, Hz), grid)
end
return grid
end
# return the total advective velocities
@inline total_velocities(m::NonhydrostaticModel) =
(u = SumOfArrays{2}(m.velocities.u, m.background_fields.velocities.u),
v = SumOfArrays{2}(m.velocities.v, m.background_fields.velocities.v),
w = SumOfArrays{2}(m.velocities.w, m.background_fields.velocities.w))