Set a default Clock depending on the grid

ext/OceananigansReactantExt/Architectures.jl

@@ -3,12 +3,15 @@ module Architectures
 using Reactant
 using Oceananigans
 
-import Oceananigans.Architectures: device, architecture, array_type, on_architecture, unified_array, ReactantState, device_copy_to!
+import Oceananigans.Architectures: device, architecture, array_type, on_architecture
+import Oceananigans.Architectures: unified_array, ReactantState, device_copy_to!
 
 const ReactantKernelAbstractionsExt = Base.get_extension(
     Reactant, :ReactantKernelAbstractionsExt
 )
+
 const ReactantBackend = ReactantKernelAbstractionsExt.ReactantBackend
+
 device(::ReactantState) = ReactantBackend()
 
 architecture(::Reactant.AnyConcreteRArray) = ReactantState
@@ -24,6 +27,7 @@ on_architecture(::ReactantState, a::SubArray{<:Any, <:Any, <:Array}) = ConcreteR
 
 unified_array(::ReactantState, a) = a
 
-@inline device_copy_to!(dst::Reactant.AnyConcreteRArray, src::Reactant.AnyConcreteRArray; kw...) = Base.copyto!(dst, src)
+@inline device_copy_to!(dst::Reactant.AnyConcreteRArray, src::Reactant.AnyConcreteRArray; kw...) =
+    Base.copyto!(dst, src)
 
 end # module

ext/OceananigansReactantExt/OceananigansReactantExt.jl

@@ -5,6 +5,7 @@ using Reactant
 include("Architectures.jl")
 using .Architectures
 
+
 # These are additional modules that may need to be Reactantified in the future:
 #
 # include("Utils.jl")

ext/OceananigansReactantExt/TimeSteppers.jl

@@ -0,0 +1,24 @@
+module TimeSteppers
+
+using Reactant
+using Oceananigans
+
+import Oceananigans.Grids: AbstractGrid
+import Oceananigans.TimeSteppers: Clock
+
+using OceananigansReactantExt: ReactantState
+
+const ReactantGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:ReactantState} where {FT, TX, TY, TZ}
+
+function Clock(grid::ReactantGrid)
+    FT = Float64 # may change in the future
+    t = ConcreteRNumber(zero(FT))
+    iter = ConcreteRNumber(0)
+    stage = ConcreteRNumber(0)
+    last_Δt = ConcreteRNumber(zero(FT))
+    last_stage_Δt = ConcreteRNumber(zero(FT))
+    return Clock(; time=t, iteration=iter, stage, last_Δt, last_stage_Δt)
+end
+
+end # module
+

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -108,7 +108,7 @@ Keyword arguments
   - `vertical_coordinate`: Rulesets that define the time-evolution of the grid (ZStar/ZCoordinate). Default: `ZCoordinate()`.
 """
 function HydrostaticFreeSurfaceModel(; grid,
-                                     clock = Clock{Float64}(time = 0),
+                                     clock = Clock(grid),
                                      momentum_advection = VectorInvariant(),
                                      tracer_advection = Centered(),
                                      buoyancy = nothing,

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -112,7 +112,7 @@ Keyword arguments
   - `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing`         
 """
 function NonhydrostaticModel(; grid,
-                             clock = Clock{Float64}(time = 0),
+                             clock = Clock(grid),
                              advection = Centered(),
                              buoyancy = nothing,
                              coriolis = nothing,

src/Models/ShallowWaterModels/shallow_water_model.jl

@@ -112,7 +112,7 @@ Keyword arguments
 function ShallowWaterModel(;
                            grid,
                            gravitational_acceleration,
-                               clock = Clock{eltype(grid)}(time=0),
+                               clock = Clock(grid),
                   momentum_advection = UpwindBiased(order=5),
                     tracer_advection = WENO(),
                       mass_advection = WENO(),

src/TimeSteppers/clock.jl

@@ -1,6 +1,7 @@
 using Adapt
 using Dates: AbstractTime, DateTime, Nanosecond, Millisecond
 using Oceananigans.Utils: prettytime
+using Oceananigans.Grids: AbstractGrid
 
 import Base: show
 import Oceananigans.Units: Time
@@ -54,6 +55,9 @@ function Clock{TT}(; time,
     return Clock{TT, DT}(time, last_Δt, last_stage_Δt, iteration, stage)
 end
 
+# helpful default
+Clock(grid::AbstractGrid) = Clock{Float64}(time=0)
+
 function Base.summary(clock::Clock)
     TT = typeof(clock.time)
     DT = typeof(clock.last_Δt)