Add DateTime clock timestepping test for OceanSeaIceModel

experiments/one_degree_calibration/one_degree_near_global_simulation.jl

@@ -1,12 +1,12 @@
 using ClimaOcean.NearGlobalSimulations: one_degree_near_global_simulation
 
 using Oceananigans
-using Oceananigans.Units
-using Oceananigans.Utils: WallTimeInterval
 using Oceananigans.Models: buoyancy_operation
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField
 using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities:
     MixingLength, TurbulentKineticEnergyEquation, CATKEVerticalDiffusivity
+using Oceananigans.Units
+using Oceananigans.Utils: WallTimeInterval
 
 #using ParameterEstimocean.Parameters: closure_with_parameters
 using JLD2

src/Atmospheres/prescribed_atmosphere.jl

@@ -64,6 +64,11 @@ function default_atmosphere_pressure(grid, times)
     return pa
 end
 
+"""Default clock that starts at the first provided time."""
+function default_atmosphere_clock(times)
+    isempty(times) && return Clock{Float64}(time = 0)
+    return Clock(time = first(times))
+end
 
 @inline function update_state!(atmos::PrescribedAtmosphere)
     time = Time(atmos.clock.time)
@@ -93,7 +98,7 @@ update_net_fluxes!(coupled_model, ::PrescribedAtmosphere) = nothing
 
 """
     PrescribedAtmosphere(grid, times=[zero(grid)];
-                         clock = Clock{Float64}(time = 0),
+                         clock = default_atmosphere_clock(times),
                          surface_layer_height = 10, # meters
                          boundary_layer_height = 512 # meters,
                          thermodynamics_parameters = AtmosphereThermodynamicsParameters(eltype(grid)),
@@ -108,7 +113,7 @@ Return a representation of a prescribed time-evolving atmospheric
 state with data given at `times`.
 """
 function PrescribedAtmosphere(grid, times=[zero(grid)];
-                              clock = Clock{Float64}(time = 0),
+                              clock = default_atmosphere_clock(times),
                               surface_layer_height = 10,
                               boundary_layer_height = 512,
                               thermodynamics_parameters = AtmosphereThermodynamicsParameters(eltype(grid)),

src/Diagnostics/Diagnostics.jl

@@ -4,10 +4,10 @@ export MixedLayerDepthField, MixedLayerDepthOperand
 
 using Oceananigans
 using Oceananigans.Architectures: architecture
-using Oceananigans.Models: buoyancy_operation
-using Oceananigans.Grids: new_data, inactive_cell, znode
 using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions!
 using Oceananigans.Fields: FieldStatus
+using Oceananigans.Grids: new_data, inactive_cell, znode
+using Oceananigans.Models: buoyancy_operation
 using Oceananigans.Utils: launch!
 using KernelAbstractions: @index, @kernel
 

src/OceanSeaIceModels/InterfaceComputations/InterfaceComputations.jl

@@ -1,10 +1,5 @@
 module InterfaceComputations
 
-using Oceananigans
-using Oceananigans.Fields: AbstractField
-using Oceananigans.Utils: KernelParameters
-using Adapt
-
 export
     Radiation,
     ComponentInterfaces,
@@ -21,6 +16,11 @@ export
     compute_atmosphere_sea_ice_fluxes!,
     compute_sea_ice_ocean_flu
 
+using Adapt
+using Oceananigans
+using Oceananigans.Fields: AbstractField
+using Oceananigans.Utils: KernelParameters
+
 using ..OceanSeaIceModels: default_gravitational_acceleration,
                            default_freshwater_density,
                            thermodynamics_parameters,

src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl

@@ -1,5 +1,5 @@
-using Oceananigans.Operators: intrinsic_vector
 using Oceananigans.Grids: inactive_node
+using Oceananigans.Operators: intrinsic_vector
 
 function compute_atmosphere_ocean_fluxes!(coupled_model)
     exchanger = coupled_model.interfaces.exchanger
@@ -11,7 +11,7 @@ function compute_atmosphere_ocean_fluxes!(coupled_model)
 
     # Simplify NamedTuple to reduce parameter space consumption.
     # See https://github.com/CliMA/ClimaOcean.jl/issues/116.
-    atmosphere_data = merge(atmosphere_fields, 
+    atmosphere_data = merge(atmosphere_fields,
                             (; h_bℓ = boundary_layer_height(coupled_model.atmosphere)))
 
     flux_formulation = coupled_model.interfaces.atmosphere_ocean_interface.flux_formulation
@@ -147,7 +147,7 @@ end
     ρₐ = AtmosphericThermodynamics.air_density(ℂₐ, 𝒬ₐ)
     cₚ = AtmosphericThermodynamics.cp_m(ℂₐ, 𝒬ₐ) # moist heat capacity
     ℒv = AtmosphericThermodynamics.latent_heat_vapor(ℂₐ, 𝒬ₐ)
-    
+
 
     # Store fluxes
     Qv  = interface_fluxes.latent_heat
@@ -159,7 +159,7 @@ end
 
     @inbounds begin
         # +0: cooling, -0: heating
-        Qv[i, j, 1]  = - ρₐ * ℒv * u★ * q★ 
+        Qv[i, j, 1]  = - ρₐ * ℒv * u★ * q★
         Qc[i, j, 1]  = - ρₐ * cₚ * u★ * θ★
         Fv[i, j, 1]  = - ρₐ * u★ * q★
         ρτx[i, j, 1] = + ρₐ * τx

src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl

@@ -11,12 +11,12 @@ using ..OceanSeaIceModels: reference_density,
 using ClimaSeaIce: SeaIceModel
 
 using Oceananigans: HydrostaticFreeSurfaceModel, architecture
-using Oceananigans.Units: Time
-using Oceananigans.Grids: inactive_node, node, topology
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Fields: ConstantField, interpolate, FractionalIndices
-using Oceananigans.Utils: launch!, KernelParameters
+using Oceananigans.Grids: inactive_node, node, topology
 using Oceananigans.Operators: ℑxᶜᵃᵃ, ℑyᵃᶜᵃ, ℑxᶠᵃᵃ, ℑyᵃᶠᵃ
+using Oceananigans.Units: Time
+using Oceananigans.Utils: launch!, KernelParameters
 
 using KernelAbstractions: @kernel, @index
 
@@ -80,7 +80,7 @@ atmosphere_ocean_interface(grid, ::Nothing,   ocean,    args...) = nothing
 atmosphere_ocean_interface(grid, ::Nothing,  ::Nothing, args...) = nothing
 atmosphere_ocean_interface(grid, atmosphere, ::Nothing, args...) = nothing
 
-function atmosphere_ocean_interface(grid, 
+function atmosphere_ocean_interface(grid,
                                     atmosphere,
                                     ocean,
                                     radiation,
@@ -136,7 +136,7 @@ atmosphere_sea_ice_interface(grid, atmos, ::Nothing,     args...) = nothing
 atmosphere_sea_ice_interface(grid, ::Nothing, sea_ice,   args...) = nothing
 atmosphere_sea_ice_interface(grid, ::Nothing, ::Nothing, args...) = nothing
 
-function atmosphere_sea_ice_interface(grid, 
+function atmosphere_sea_ice_interface(grid,
                                       atmosphere,
                                       sea_ice,
                                       radiation,
@@ -290,7 +290,7 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing;
 
     io_interface = sea_ice_ocean_interface(exchange_grid, sea_ice, ocean)
 
-    ai_interface = atmosphere_sea_ice_interface(exchange_grid, 
+    ai_interface = atmosphere_sea_ice_interface(exchange_grid,
                                                 atmosphere,
                                                 sea_ice,
                                                 radiation,

src/OceanSeaIceModels/InterfaceComputations/latitude_dependent_albedo.jl

@@ -52,4 +52,3 @@ Base.show(io::IO, α::LatitudeDependentAlbedo) = print(io, summary(α))
     α₁ = α.direct
     return α₀ - α₁ * hack_cosd(2φ)
 end
-

src/OceanSeaIceModels/InterfaceComputations/similarity_theory_turbulent_fluxes.jl

@@ -1,4 +1,5 @@
-using Oceananigans.Grids: AbstractGrid, prettysummary
+using Oceananigans.Grids: AbstractGrid
+using Oceananigans.Utils: prettysummary
 
 using Adapt
 using Printf

src/OceanSeaIceModels/OceanSeaIceModels.jl

@@ -15,15 +15,15 @@ export
 
 using Oceananigans
 using Oceananigans.Operators
-using Oceananigans.Utils: launch!, KernelParameters
-using Oceananigans.Units: Time
 using Oceananigans.Architectures: architecture
 using Oceananigans.BoundaryConditions: fill_halo_regions!, BoundaryCondition
-using Oceananigans.Grids: architecture
 using Oceananigans.Fields: ZeroField
-using Oceananigans.TimeSteppers: tick!
+using Oceananigans.Grids: architecture
 using Oceananigans.Models: AbstractModel
 using Oceananigans.OutputReaders: FieldTimeSeries, GPUAdaptedFieldTimeSeries
+using Oceananigans.TimeSteppers: tick!
+using Oceananigans.Units: Time
+using Oceananigans.Utils: launch!, KernelParameters
 
 using ClimaSeaIce: SeaIceModel
 using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
@@ -49,7 +49,7 @@ include("components.jl")
 
 #####
 ##### The coupled model
-##### 
+#####
 
 const default_gravitational_acceleration = Oceananigans.defaults.gravitational_acceleration
 const default_freshwater_density = 1000 # kg m⁻³
@@ -64,7 +64,7 @@ include("time_step_ocean_sea_ice_model.jl")
 #####
 #####  Fallbacks for no-interface models
 #####
-    
+
 using .InterfaceComputations: ComponentInterfaces, AtmosphereInterface, SeaIceOceanInterface
 
 const NoSeaIceInterface = ComponentInterfaces{<:AtmosphereInterface,  <:Nothing, <:Nothing}

src/OceanSeaIceModels/ocean_sea_ice_model.jl

@@ -55,6 +55,29 @@ prognostic_fields(cm::OSIM)         = nothing
 fields(::OSIM)                      = NamedTuple()
 default_clock(TT)                   = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1)
 
+"""Return a clock from a component if available."""
+function component_clock(component)
+    if component === nothing
+        return nothing
+    elseif hasproperty(component, :clock)
+        return component.clock
+    elseif hasproperty(component, :model) && hasproperty(component.model, :clock)
+        return component.model.clock
+    else
+        return nothing
+    end
+end
+
+"""Select a coupled clock based on the components, defaulting to a zero Float64 clock."""
+function default_coupled_clock(ocean, atmosphere, sea_ice)
+    for component in (ocean, atmosphere, sea_ice)
+        clock = component_clock(component)
+        !isnothing(clock) && return deepcopy(clock)
+    end
+
+    return Clock{Float64}(time=0)
+end
+
 function reset!(model::OSIM)
     reset!(model.ocean)
     return nothing
@@ -83,7 +106,7 @@ heat_capacity(unsupported) =
     OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype(ocean.model));
                      atmosphere = nothing,
                      radiation = Radiation(architecture(ocean.model)),
-                     clock = deepcopy(ocean.model.clock),
+                     clock = default_coupled_clock(ocean, atmosphere, sea_ice),
                      ocean_reference_density = reference_density(ocean),
                      ocean_heat_capacity = heat_capacity(ocean),
                      sea_ice_reference_density = reference_density(sea_ice),
@@ -119,7 +142,7 @@ using ClimaOcean
 using Oceananigans
 
 grid = RectilinearGrid(size=10, z=(-100, 0), topology=(Flat, Flat, Bounded))
-ocean = ocean_simulation(grid, timestepper = :QuasiAdamsBashforth2)
+ocean = ocean_simulation(grid)
 
 # Three choices for stability function:
 # "No stability function", which also apply to neutral boundary layers
@@ -136,6 +159,9 @@ interfaces = ClimaOcean.OceanSeaIceModels.ComponentInterfaces(nothing, ocean; at
 model = OceanSeaIceModel(ocean; interfaces)
 
 # output
+┌ Warning: Split barotropic-baroclinic time stepping with SplitRungeKutta3TimeStepper is experimental.
+│ Use at own risk, and report any issues encountered at https://github.com/CliMA/Oceananigans.jl/issues.
+└ @ Oceananigans.TimeSteppers ~/Oceananigans/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl:59
 OceanSeaIceModel{CPU}(time = 0 seconds, iteration = 0)
 ├── ocean: HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 ├── atmosphere: Nothing
@@ -153,7 +179,7 @@ The available stability function options include:
 function OceanSeaIceModel(ocean, sea_ice=default_sea_ice();
                           atmosphere = nothing,
                           radiation = Radiation(),
-                          clock = Clock{Float64}(time=0),
+                          clock = default_coupled_clock(ocean, atmosphere, sea_ice),
                           ocean_reference_density = reference_density(ocean),
                           ocean_heat_capacity = heat_capacity(ocean),
                           sea_ice_reference_density = reference_density(sea_ice),

src/Oceans/ocean_simulation.jl

@@ -98,6 +98,7 @@ end
 """
     ocean_simulation(grid;
                      Δt = estimate_maximum_Δt(grid),
+                     clock = Clock(grid),
                      closure = default_ocean_closure(),
                      tracers = (:T, :S),
                      free_surface = default_free_surface(grid),
@@ -124,7 +125,6 @@ consistent defaults for advection, closures, the equation of state, surface flux
 barotropic pressure–gradient forcing, boundary conditions, and optional biogeochemistry.
 It then wraps the model into an Oceananigans's `Simulation` with the specified timestepping options.
 
-
 ## Behaviour and automatic configuration
 
 ### Coriolis
@@ -159,6 +159,7 @@ defaults on a per-field basis.
 ## Keyword Arguments
 
 - `Δt`: Timestep used by the `Simulation`. Defaults to the maximum stable timestep estimated from the `grid`.
+- `clock`: Clock object. Defaults to `Clock(grid)`.
 - `closure`: A turbulence or mixing closure. Defaults to `default_ocean_closure()`.
 - `tracers`: Tuple of tracer names. Defaults to `(:T, :S)`.
 - `free_surface`: Free–surface solver. Defaults to `default_free_surface(grid)`.
@@ -180,6 +181,7 @@ defaults on a per-field basis.
 """
 function ocean_simulation(grid;
                           Δt = estimate_maximum_Δt(grid),
+                          clock = Clock(grid),
                           closure = default_ocean_closure(),
                           tracers = (:T, :S),
                           free_surface = default_free_surface(grid),
@@ -301,6 +303,7 @@ function ocean_simulation(grid;
     end
 
     ocean_model = HydrostaticFreeSurfaceModel(; grid,
+                                              clock,
                                               buoyancy,
                                               closure,
                                               biogeochemistry,

test/test_ocean_sea_ice_model.jl

@@ -6,6 +6,7 @@ using ClimaOcean.OceanSeaIceModels: above_freezing_ocean_temperature!
 using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
 using ClimaSeaIce.SeaIceDynamics
 using ClimaSeaIce.Rheologies
+using Oceananigans.TimeSteppers: Clock
 
 @inline kernel_melting_temperature(i, j, k, grid, liquidus, S) = @inbounds melting_temperature(liquidus, S[i, j, k])
 
@@ -105,3 +106,52 @@ using ClimaSeaIce.Rheologies
         end
     end
 end
+
+@testset "Prescribed atmosphere with DateTime clock" begin
+    start_time = Dates.DateTime(1993, 2, 1)
+    Δt_seconds = 3050.0
+    Δt_period = Dates.Second(3050)
+    steps = 5
+    stop_time = start_time + steps * Δt_period
+    times = [start_time + m * Δt_period for m in 0:steps]
+
+    for arch in test_architectures
+        grid = LatitudeLongitudeGrid(arch;
+                                     size = (1, 1, 20),
+                                     longitude = (10, 11),
+                                     latitude = (20.2, 21.2),
+                                     z = (-200, 0),
+                                     topology = (Bounded, Bounded, Bounded))
+
+        ocean_clock = Clock(time=start_time)
+        ocean = ocean_simulation(grid; clock=ocean_clock, Δt=Δt_seconds, tracers=(:T, :S), verbose=false, closure=nothing)
+        atmosphere_clock = Clock(time=start_time)
+        atmosphere = PrescribedAtmosphere(grid, times; clock=atmosphere_clock)
+        radiation = Radiation(arch)
+        coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
+
+        for _ in 1:steps
+            time_step!(coupled_model, Δt_seconds)
+        end
+
+        @test coupled_model.clock.time == stop_time
+        @test coupled_model.clock.iteration == steps
+        @test atmosphere.clock.time == coupled_model.clock.time
+        @test coupled_model.clock.time isa Dates.DateTime
+
+        ocean_clock = Clock(time=start_time)
+        ocean = ocean_simulation(grid; clock=ocean_clock, Δt=Δt_seconds, tracers=(:T, :S), verbose=false, closure=nothing)
+        atmosphere_clock = Clock(time=start_time)
+        atmosphere = PrescribedAtmosphere(grid, times; clock=atmosphere_clock)
+        radiation = Radiation(arch)
+        coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
+        coupled_simulation = Simulation(coupled_model; Δt=Δt_seconds, stop_time)
+
+        run!(coupled_simulation)
+
+        @test coupled_simulation.model.clock.time == stop_time
+        @test coupled_simulation.model.clock.iteration == steps
+        @test coupled_simulation.model.atmosphere.clock.time == coupled_model.clock.time
+        @test coupled_simulation.model.clock.time isa Dates.DateTime
+    end
+end