(0.4.7) Update for latest oceananigans (#390)

* Update for latest oceananigans

* More cleanup

* Bump ClimaSeaIce compat

Author: glwagner

Project.toml

@@ -2,7 +2,7 @@ name = "ClimaOcean"
 uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 license = "MIT"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.4.6"
+version = "0.4.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -33,7 +33,7 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 Adapt = "4"
 CFTime = "0.1"
 CUDA = "4, 5"
-ClimaSeaIce = "0.2.3"
+ClimaSeaIce = "0.2.4"
 CubicSplines = "0.2"
 DataDeps = "0.7"
 Downloads = "1.6"
@@ -42,7 +42,7 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.95.18 - 0.99"
+Oceananigans = "0.95.21 - 0.99"
 OffsetArrays = "1.14"
 OrthogonalSphericalShellGrids = "0.2.2"
 Scratch = "1"

src/OceanSeaIceModels/ocean_sea_ice_model.jl

@@ -19,7 +19,8 @@ import Oceananigans.TimeSteppers: time_step!, update_state!, time
 import Oceananigans.Utils: prettytime
 import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker
 
-struct OceanSeaIceModel{I, A, O, F, C} <: AbstractModel{Nothing}
+struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch}
+    architecture :: Arch
     clock :: C
     atmosphere :: A
     sea_ice :: I
@@ -57,7 +58,6 @@ Base.eltype(model::OSIM)            = Base.eltype(model.ocean.model)
 prettytime(model::OSIM)             = prettytime(model.clock.time)
 iteration(model::OSIM)              = model.clock.iteration
 timestepper(::OSIM)                 = nothing
-initialize!(::OSIM)                 = nothing
 default_included_properties(::OSIM) = tuple()
 prognostic_fields(cm::OSIM)         = nothing
 fields(::OSIM)                      = NamedTuple()
@@ -68,6 +68,11 @@ function reset!(model::OSIM)
     return nothing
 end
 
+function initialize!(model::OSIM)
+    initialize!(model.ocean)
+    return nothing
+end
+
 reference_density(unsupported) =
     throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))"))
 
@@ -102,26 +107,21 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype(
                           sea_ice_heat_capacity = heat_capacity(sea_ice),
                           interfaces = nothing)
 
-    # Remove some potentially irksome callbacks from the ocean simulation
-    pop!(ocean.callbacks, :stop_time_exceeded, nothing)
-    pop!(ocean.callbacks, :stop_iteration_exceeded, nothing)
-    pop!(ocean.callbacks, :wall_time_limit_exceeded, nothing)
-    pop!(ocean.callbacks, :nan_checker, nothing)
-
-    # In case there was any doubt these are meaningless.
-    ocean.stop_time = Inf
-    ocean.stop_iteration = Inf
-    ocean.wall_time_limit = Inf
-
+    if !isnothing(ocean.callbacks)
+        # Remove some potentially irksome callbacks from the ocean simulation
+        pop!(ocean.callbacks, :stop_time_exceeded, nothing)
+        pop!(ocean.callbacks, :stop_iteration_exceeded, nothing)
+        pop!(ocean.callbacks, :wall_time_limit_exceeded, nothing)
+        pop!(ocean.callbacks, :nan_checker, nothing)
+    end
+    
     if sea_ice isa SeaIceSimulation
-        pop!(sea_ice.callbacks, :stop_time_exceeded, nothing)
-        pop!(sea_ice.callbacks, :stop_iteration_exceeded, nothing)
-        pop!(sea_ice.callbacks, :wall_time_limit_exceeded, nothing)
-        pop!(sea_ice.callbacks, :nan_checker, nothing)
-
-        sea_ice.stop_time = Inf
-        sea_ice.stop_iteration = Inf
-        sea_ice.wall_time_limit = Inf
+        if !isnothing(sea_ice.callbacks)
+            pop!(sea_ice.callbacks, :stop_time_exceeded, nothing)
+            pop!(sea_ice.callbacks, :stop_iteration_exceeded, nothing)
+            pop!(sea_ice.callbacks, :wall_time_limit_exceeded, nothing)
+            pop!(sea_ice.callbacks, :nan_checker, nothing)
+        end
     end
 
     # Contains information about flux contributions: bulk formula, prescribed fluxes, etc.
@@ -134,7 +134,10 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype(
                                          radiation)
     end
 
-    ocean_sea_ice_model = OceanSeaIceModel(clock,
+    arch = architecture(ocean.model.grid)
+
+    ocean_sea_ice_model = OceanSeaIceModel(arch,
+                                           clock,
                                            atmosphere,
                                            sea_ice,
                                            ocean,

src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl

@@ -10,14 +10,14 @@ using Oceananigans.Grids: φnode
 
 using Printf
 
-function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_tendencies=true)
+function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_tendencies=true, kw...)
     ocean = coupled_model.ocean
     sea_ice = coupled_model.sea_ice
     atmosphere = coupled_model.atmosphere
     clock = coupled_model.clock
 
     # Be paranoid and update state at iteration 0
-    coupled_model.clock.iteration == 0 && update_state!(coupled_model, callbacks)
+    #coupled_model.clock.iteration == 0 && update_state!(coupled_model, callbacks)
 
     # Eventually, split out into OceanOnlyModel
     if sea_ice isa SeaIceSimulation
@@ -37,7 +37,7 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_
 
     # TODO after ice time-step:
     #  - Adjust ocean heat flux if the ice completely melts?
-    time_step!(ocean, Δt)
+    time_step!(ocean.model, Δt; kw...)
 
     # Time step the atmosphere
     time_step!(atmosphere, Δt)