Move AbstractModel interface to top level and add docstrings

ext/OceananigansReactantExt/Models.jl

@@ -2,7 +2,7 @@ module Models
 
 import Oceananigans
 
-import Oceananigans.Models: initialization_update_state!
+import Oceananigans.TimeSteppers: initialization_update_state!
 import Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: maybe_extend_halos, FixedSubstepNumber
 import Oceananigans: initialize!
 

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -1,5 +1,6 @@
 using OrderedCollections: OrderedDict
 
+using Oceananigans: AbstractModel
 using Oceananigans.DistributedComputations
 using Oceananigans.Architectures: AbstractArchitecture
 using Oceananigans.Advection: AbstractAdvectionScheme, Centered, VectorInvariant, adapt_advection_order
@@ -10,14 +11,13 @@ using Oceananigans.Fields: Field, CenterField, tracernames, VelocityFields, Trac
 using Oceananigans.Forcings: model_forcing
 using Oceananigans.Grids: AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid, architecture, halo_size, MutableVerticalDiscretization
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
-using Oceananigans.Models: AbstractModel, validate_model_halo, validate_tracer_advection, extract_boundary_conditions, initialization_update_state!
-using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles, SplitRungeKutta3TimeStepper
+using Oceananigans.Models: validate_model_halo, validate_tracer_advection, extract_boundary_conditions
+using Oceananigans.TimeSteppers: Clock, TimeStepper, AbstractLagrangianParticles, SplitRungeKutta3TimeStepper, timestepper, update_state!, initialization_update_state!
 using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, build_diffusivity_fields, add_closure_specific_boundary_conditions
 using Oceananigans.TurbulenceClosures: time_discretization, implicit_diffusion_solver
 using Oceananigans.Utils: tupleit
 
-import Oceananigans: initialize!
-import Oceananigans.Models: total_velocities, timestepper
+import Oceananigans: initialize!, total_velocities
 
 PressureField(grid) = (; pHY′ = CenterField(grid))
 

src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl

@@ -1,5 +1,5 @@
+using Oceananigans: total_velocities
 using Oceananigans.Utils: instantiate, KernelParameters
-using Oceananigans.Models: total_velocities
 using Oceananigans.Fields: interpolator, FractionalIndices
 
 #####

src/Models/Models.jl

@@ -10,6 +10,8 @@ export
     BoundaryConditionOperation, ForcingOperation,
     seawater_density
 
+using DocStringExtensions
+
 using Oceananigans: AbstractModel, fields, prognostic_fields
 using Oceananigans.AbstractOperations: AbstractOperation
 using Oceananigans.Advection: AbstractAdvectionScheme, Centered, VectorInvariant
@@ -19,31 +21,23 @@ using Oceananigans.OutputReaders: update_field_time_series!, extract_field_time_
 using Oceananigans.TimeSteppers: AbstractTimeStepper, Clock, update_state!
 using Oceananigans.Utils: Time
 
-import Oceananigans: initialize!
+import Oceananigans: initialize!, iteration
 import Oceananigans.Architectures: architecture
-import Oceananigans.Solvers: iteration
-import Oceananigans.Simulations: timestepper
-import Oceananigans.TimeSteppers: reset!, set_clock!
+import Oceananigans.TimeSteppers: timestepper, reset!, set_clock!
 
-# A prototype interface for AbstractModel.
-#
-# TODO: decide if we like this.
-#
-# We assume that model has some properties, eg:
-#   - model.clock::Clock
-#   - model.architecture.
-#   - model.timestepper with timestepper.G⁻ and timestepper.Gⁿ :spiral_eyes:
-
-iteration(model::AbstractModel) = model.clock.iteration
-Base.time(model::AbstractModel) = model.clock.time
-Base.eltype(model::AbstractModel) = eltype(model.grid)
+"""
+    $SIGNATURES
+
+Returns the device architecture defined by the underlying model `grid`.
+"""
 architecture(model::AbstractModel) = model.grid.architecture
-initialize!(model::AbstractModel) = nothing
-total_velocities(model::AbstractModel) = nothing
-timestepper(model::AbstractModel) = model.timestepper
-initialization_update_state!(model::AbstractModel; kw...) = update_state!(model; kw...) # fallback
 
-# Fallback for any abstract model that does not contain `FieldTimeSeries`es
+"""
+    $SIGNATURES
+
+Update all `FieldTimeSeries` fields defined on the given `model.
+The default implementation assumes that the `model` that does not contain any `FieldTimeSeries` and returns `nothing`.
+"""
 update_model_field_time_series!(model::AbstractModel, clock::Clock) = nothing
 
 #####

src/Models/NonhydrostaticModels/NonhydrostaticModels.jl

@@ -20,7 +20,7 @@ using Oceananigans.Utils: sum_of_velocities
 
 import Oceananigans: fields, prognostic_fields
 import Oceananigans.Advection: cell_advection_timescale
-import Oceananigans.TimeSteppers: step_lagrangian_particles!
+import Oceananigans.TimeSteppers: timestepper, step_lagrangian_particles!
 
 function nonhydrostatic_pressure_solver(::Distributed, local_grid::XYZRegularRG)
     global_grid = reconstruct_global_grid(local_grid)

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -18,8 +18,8 @@ using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: FlavorOfCAT
 using Oceananigans.Utils: tupleit
 using Oceananigans.Grids: topology
 
+import Oceananigans: total_velocities
 import Oceananigans.Architectures: architecture
-import Oceananigans.Models: total_velocities, timestepper
 
 const ParticlesOrNothing = Union{Nothing, AbstractLagrangianParticles}
 const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry}

src/Models/ShallowWaterModels/shallow_water_model.jl

@@ -15,7 +15,7 @@ using Oceananigans.TurbulenceClosures: with_tracers, build_diffusivity_fields
 using Oceananigans.Utils: tupleit
 
 import Oceananigans.Architectures: architecture
-import Oceananigans.Simulations: timestepper
+import Oceananigans.TimeSteppers: timestepper
 
 const RectilinearGrids = Union{RectilinearGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:RectilinearGrid}}
 

src/Oceananigans.jl

@@ -84,6 +84,9 @@ export
     ConservativeFormulation, VectorInvariantFormulation,
     PressureField, fields, ZCoordinate, ZStarCoordinate,
 
+    # Model interface
+    iteration, timesteppper,
+
     # Hydrostatic free surface model stuff
     VectorInvariant, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface,
     HydrostaticSphericalCoriolis, PrescribedVelocityFields,
@@ -92,7 +95,7 @@ export
     Clock, TimeStepWizard, conjure_time_step_wizard!, time_step!,
 
     # Simulations
-    Simulation, run!, Callback, add_callback!, iteration,
+    Simulation, run!, Callback, add_callback!,
     iteration_limit_exceeded, stop_time_exceeded, wall_time_limit_exceeded,
 
     # Diagnostics
@@ -162,9 +165,52 @@ const defaults = Defaults()
     AbstractModel
 
 Abstract supertype for models.
+
+Models are required to have the following fields:
+  - `model.clock::Clock`
+  - `model.grid::AbstractGrid`
 """
 abstract type AbstractModel{TS, A} end
 
+# Method interface for AbstractModel
+
+"""
+    $SIGNATURES
+
+Return the current iteration of the `model.clock`.
+"""
+iteration(model::AbstractModel) = model.clock.iteration
+
+"""
+    $SIGNATURES
+
+Ensure consistency between internal variables in the `AbstractModel` and user-specified
+initial conditions provided by via `set(::AbstractModel; kwargs...)`.
+Default implementation does nothing.
+"""
+initialize!(model::AbstractModel) = nothing
+
+"""
+    $SIGNATURES
+
+Returns the sum of the user-prescribed background + prognostic velocities, where applicable.
+"""
+total_velocities(model::AbstractModel) = nothing
+
+"""
+    Base.time(model::AbstractModel) 
+
+Return the current `time` from the `model.clock`.
+"""
+Base.time(model::AbstractModel) = model.clock.time
+
+"""
+    Base.eltype(model::AbstractModel)
+
+Return the numeric `eltype` of the `model.grid` and all associated `Field`s.
+"""
+Base.eltype(model::AbstractModel) = eltype(model.grid)
+
 """
     AbstractDiagnostic
 

src/Simulations/Simulations.jl

@@ -18,9 +18,6 @@ using OrderedCollections: OrderedDict
 
 import Base: show
 
-# To be extended in the `Models` module
-timestepper(model) = nothing
-
 include("callback.jl")
 include("simulation.jl")
 include("run.jl")

src/Solvers/Solvers.jl

@@ -24,6 +24,8 @@ using Oceananigans.Fields
 using Oceananigans.Grids: unpack_grid, inactive_cell
 using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 
+import Oceananigans: iteration
+
 """
     ω(M, k)
 

src/TimeSteppers/TimeSteppers.jl

@@ -5,9 +5,12 @@ export
     RungeKutta3TimeStepper,
     SplitRungeKutta3TimeStepper,
     time_step!,
+    timestepper,
     Clock,
     tendencies
 
+using DocStringExtensions
+
 using KernelAbstractions
 using Oceananigans: AbstractModel, initialize!, prognostic_fields
 using Oceananigans.Architectures: device
@@ -24,6 +27,24 @@ function update_state! end
 function compute_tendencies! end
 function compute_flux_bc_tendencies! end
 
+"""
+    $SIGNATURES
+
+Define initialization routines for the first call to `update_state!` after initialization.
+Defaults to invoking `update_state!(model; kw...)`.
+"""
+function initialization_update_state!(model::AbstractModel; kw...) 
+    initialize!(model)
+    update_state!(model; kw...) # fallback
+end
+
+"""
+    $SIGNATURES
+
+Returns the `TimeStepper` used by the given `model`.
+"""
+timestepper(model::AbstractModel) = model.timestepper
+
 compute_pressure_correction!(model, Δt) = nothing
 make_pressure_correction!(model, Δt) = nothing