diff --git a/ext/OceananigansReactantExt/Models.jl b/ext/OceananigansReactantExt/Models.jl index 64a5fff6a2..4fe5fde947 100644 --- a/ext/OceananigansReactantExt/Models.jl +++ b/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! diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index cd8952d517..0c220d5b87 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/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)) diff --git a/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl b/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl index 00400578ef..0888026ba8 100644 --- a/src/Models/LagrangianParticleTracking/lagrangian_particle_advection.jl +++ b/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 ##### diff --git a/src/Models/Models.jl b/src/Models/Models.jl index cf3da49513..40db6c5665 100644 --- a/src/Models/Models.jl +++ b/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 ##### diff --git a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl b/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl index a1c0a8955c..c5deffbd68 100644 --- a/src/Models/NonhydrostaticModels/NonhydrostaticModels.jl +++ b/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) diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index c89fb9fdae..b435e912ba 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/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} diff --git a/src/Models/ShallowWaterModels/shallow_water_model.jl b/src/Models/ShallowWaterModels/shallow_water_model.jl index ebbbb1c1b3..0fdb4570d0 100644 --- a/src/Models/ShallowWaterModels/shallow_water_model.jl +++ b/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}} diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 1c53a1fa47..1c804b47ab 100644 --- a/src/Oceananigans.jl +++ b/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 diff --git a/src/Simulations/Simulations.jl b/src/Simulations/Simulations.jl index 83d53c841c..e0b0df057f 100644 --- a/src/Simulations/Simulations.jl +++ b/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") diff --git a/src/Solvers/Solvers.jl b/src/Solvers/Solvers.jl index 61cf0c4513..091cb9c163 100644 --- a/src/Solvers/Solvers.jl +++ b/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) diff --git a/src/TimeSteppers/TimeSteppers.jl b/src/TimeSteppers/TimeSteppers.jl index e025458f06..f90148bf1c 100644 --- a/src/TimeSteppers/TimeSteppers.jl +++ b/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