Putting together a near-global prototype OMIP setup (#172)

* Fix import of melting_temperature from ClimaSeaIce

* Update compat to ClimaSeaIce

* Update ClimaSeaIce compat

* Get rid of ocean only model for now

* Cosmetic changes

* SlabSeaIce -> SeaIce

* Delete ocena only model file

* pop callbacks from the ocean simulation

* Export LatitudeDependentAlbedo

* Add newline

* Add lat lon near global script

* Improve show method for OceanSeaIceModel

prototype_omip_simulation/lat_lon_near_global_simulation.jl

@@ -0,0 +1,171 @@
+using Oceananigans
+using Oceananigans.Units
+
+using ClimaOcean
+using ClimaOcean.ECCO: ECCO_restoring_forcing, ECCO4Monthly, ECCO2Daily, ECCOMetadata
+
+using OrthogonalSphericalShellGrids
+using Printf
+
+using CFTime
+using Dates
+
+#####
+##### Global ocean simulation
+#####
+
+Nx = 4 * 360
+Ny = 4 * 160
+Nz = 60
+arch = GPU()
+z_faces = exponential_z_faces(depth=5000; Nz)
+prefix = "near_global_omip_Nz$(Nz)_"
+
+#arch = Distributed(GPU(), partition = Partition(2))
+
+grid = LatitudeLongitudeGrid(arch; 
+                             size = (Nx, Ny, Nz), 
+                             halo = (5, 5, 5), 
+                             longitude = (0, 360),
+                             latitude = (-80, 80),
+                             z = z_faces)
+
+@info "Put together a grid:"
+@info grid
+
+bottom_height = retrieve_bathymetry(grid;
+                                    minimum_depth = 10,
+                                    interpolation_passes = 10,
+                                    major_basins = 1)
+
+grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) 
+
+#####
+##### Add restoring to ECCO fields for temperature and salinity in the artic and antarctic
+#####
+
+@inline function restoring_mask(λ, φ, z, t=0)
+    ϵN = (φ - 75) / 5
+    ϵN = clamp(ϵN, zero(ϵN), one(ϵN))
+    ϵS = - (φ + 75) / 5
+    ϵS = clamp(ϵS, zero(ϵS), one(ϵS))
+    return ϵN + ϵS
+end
+
+restoring_mask_field = CenterField(grid)
+set!(restoring_mask_field, restoring_mask)
+
+@inline sponge_layer(λ, φ, z, t, c, ω) = - restoring_mask(λ, φ, z, t) * ω * c
+Fu = Forcing(sponge_layer, field_dependencies=:u, parameters=1/1days)
+Fv = Forcing(sponge_layer, field_dependencies=:v, parameters=1/1days)
+
+dates = DateTimeProlepticGregorian(1993, 1, 1) : Month(1) : DateTimeProlepticGregorian(2003, 12, 1)
+
+temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly())
+salinity    = ECCOMetadata(:salinity,    dates, ECCO4Monthly())
+
+FT = ECCO_restoring_forcing(temperature; mask=restoring_mask_field, grid, architecture=arch, timescale=1days)
+FS = ECCO_restoring_forcing(salinity;    mask=restoring_mask_field, grid, architecture=arch, timescale=1days)
+forcing = (; T=FT, S=FS, u=Fu, v=Fv)
+
+# New advection scheme
+tracer_advection = WENO(order=5)
+momentum_advection = WENOVectorInvariant(vorticity_order=5)
+
+ocean = ocean_simulation(grid; forcing, tracer_advection, momentum_advection) 
+
+fluxes = (u = ocean.model.velocities.u.boundary_conditions.top.condition,
+          v = ocean.model.velocities.v.boundary_conditions.top.condition,
+          T = ocean.model.tracers.T.boundary_conditions.top.condition,
+          S = ocean.model.tracers.S.boundary_conditions.top.condition)
+
+ocean.output_writers[:fluxes] = JLD2OutputWriter(ocean.model, fluxes,
+                                                 schedule = TimeInterval(1days),
+                                                 with_halos = true,
+                                                 overwrite_existing = true,
+                                                 array_type = Array{Float32},
+                                                 filename = prefix * "surface_fluxes")
+
+ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, merge(ocean.model.tracers, ocean.model.velocities),
+                                                  schedule = TimeInterval(1days),
+                                                  with_halos = true,
+                                                  overwrite_existing = true,
+                                                  array_type = Array{Float32},
+                                                  filename = prefix * "surface",
+                                                  indices = (:, :, grid.Nz))
+
+ocean.output_writers[:snapshots] = JLD2OutputWriter(ocean.model, merge(ocean.model.tracers, ocean.model.velocities),
+                                                    schedule = TimeInterval(10days),
+                                                    with_halos = true,
+                                                    overwrite_existing = true,
+                                                    array_type = Array{Float32},
+                                                    filename = prefix * "snapshots")
+
+ocean.output_writers[:checkpoint] = Checkpointer(ocean.model, 
+                                                 schedule = TimeInterval(60days),
+                                                 overwrite_existing = true,
+                                                 prefix = prefix * "checkpoint")
+
+@info "Built an ocean simulation with the model:"
+@info ocean.model
+@info "\ninside the simulation:"
+@info ocean
+
+#####
+##### The atmosphere
+#####
+
+backend    = JRA55NetCDFBackend(4) 
+atmosphere = JRA55_prescribed_atmosphere(arch; backend)
+radiation  = Radiation(arch, ocean_albedo=LatitudeDependentAlbedo())
+
+sea_ice = ClimaOcean.OceanSeaIceModels.MinimumTemperatureSeaIce()
+coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
+
+@info "Set up coupled model:"
+@info coupled_model
+
+coupled_simulation = Simulation(coupled_model; Δt=30.0, stop_time=25days)
+
+# Spin up
+set!(ocean.model, 
+     T = ECCOMetadata(:temperature, dates[1], ECCO2Daily()),
+     S = ECCOMetadata(:salinity,    dates[1], ECCO2Daily()))
+
+wall_time = Ref(time_ns())
+
+function progress(sim) 
+    ocean = sim.model.ocean
+    u, v, w = ocean.model.velocities  
+    T, S = ocean.model.tracers
+
+    Tmax = maximum(interior(T))
+    Tmin = minimum(interior(T))
+    umax = maximum(abs, interior(u)), maximum(abs, interior(v)), maximum(abs, interior(w))
+    step_time = 1e-9 * (time_ns() - wall_time[])
+
+    msg = @sprintf("Time: %s, iter: %d, Δt: %s",
+                   prettytime(sim), iteration(sim), prettytime(sim.Δt))
+
+    msg *= @sprintf("max(u): (%.2e, %.2e, %.2e) m s⁻¹, extrema(T): %.2fᵒC, %.2fᵒC, wall time: %s",
+                   umax..., Tmax, Tmin, prettytime(step_time))
+
+     wall_time[] = time_ns()
+end
+
+add_callback!(coupled_simulation, progress, IterationInterval(10))
+
+@info "Running the coupled simulation:"
+@info coupled_simulation
+
+run!(coupled_simulation)
+
+# Run for real
+coupled_simulation.Δt = 5minutes
+
+# Let's reset the maximum number of iterations
+coupled_simulation.stop_time = 7200days
+coupled_simulation.stop_iteration = Inf
+
+run!(coupled_simulation)
+

src/ClimaOcean.jl

@@ -4,6 +4,7 @@ export
     OceanSeaIceModel,
     MinimumTemperatureSeaIce,
     Radiation,
+    LatitudeDependentAlbedo,
     SimilarityTheoryTurbulentFluxes,
     JRA55_prescribed_atmosphere,
     JRA55NetCDFBackend,
@@ -39,7 +40,6 @@ const SKOFTS = SomeKindOfFieldTimeSeries
 @inline function stateindex(a::Function, i, j, k, grid, time, loc)
     LX, LY, LZ = loc 
     λ, φ, z = node(i, j, k, grid, LX(), LY(), LZ())
-
     return a(λ, φ, z, time)
 end
 

src/DataWrangling/ecco_restoring.jl

@@ -273,11 +273,11 @@ function ECCO_restoring_forcing(metadata::ECCOMetadata;
     # Grab the correct Oceananigans field to restore
     variable_name = metadata.name
     field_name = oceananigans_fieldname[variable_name]
-
     ecco_restoring = ECCORestoring(ecco_fts, ecco_grid, mask, field_name, 1 / timescale)
 
     # Defining the forcing that depends on the restoring field.
     restoring_forcing = Forcing(ecco_restoring; discrete_form = true)
 
     return restoring_forcing
-end
+end
+

src/OceanSeaIceModels/CrossRealmFluxes/CrossRealmFluxes.jl

@@ -5,6 +5,7 @@ using Adapt
 
 export Radiation,
        OceanSeaIceSurfaceFluxes,
+       LatitudeDependentAlbedo,
        SimilarityTheoryTurbulentFluxes
 
 using ..OceanSeaIceModels: default_gravitational_acceleration

src/OceanSeaIceModels/OceanSeaIceModels.jl

@@ -1,6 +1,6 @@
 module OceanSeaIceModels
 
-export OceanSeaIceModel, SimilarityTheoryTurbulentFluxes, Radiation
+export OceanSeaIceModel, SimilarityTheoryTurbulentFluxes, Radiation, LatitudeDependentAlbedo
 
 using Oceananigans
 using SeawaterPolynomials
@@ -49,7 +49,6 @@ using .CrossRealmFluxes
 
 include("minimum_temperature_sea_ice.jl")
 include("ocean_sea_ice_model.jl")
-include("ocean_only_model.jl")
 include("time_step_ocean_sea_ice_model.jl")
 
 import .CrossRealmFluxes:

src/OceanSeaIceModels/ocean_sea_ice_model.jl

@@ -26,8 +26,21 @@ end
 
 const OSIM = OceanSeaIceModel
 
-Base.summary(::OSIM)                = "OceanSeaIceModel"
-Base.show(io::IO, cm::OSIM)         = print(io, summary(cm))
+function Base.summary(model::OSIM)
+    A = nameof(typeof(architecture(model.grid)))
+    G = nameof(typeof(model.grid))
+    return string("OceanSeaIceModel{$A, $G}",
+                  "(time = ", prettytime(model.clock.time), ", iteration = ", model.clock.iteration, ")")
+end
+
+function Base.show(io::IO, cm::OSIM)
+    print(io, summary(cm))
+    print(io, "├── ocean: ", summary(cm.ocean.model))
+    print(io, "├── atmosphere: ", summary(cm.atmosphere))
+    print(io, "└── sea_ice: ", summary(cm.sea_ice))
+    return nothing
+end
+
 prettytime(model::OSIM)             = prettytime(model.clock.time)
 iteration(model::OSIM)              = model.clock.iteration
 timestepper(::OSIM)                 = nothing
@@ -63,7 +76,19 @@ function OceanSeaIceModel(ocean, sea_ice=nothing;
                           ocean_reference_density = reference_density(ocean),
                           ocean_heat_capacity = heat_capacity(ocean),
                           clock = deepcopy(ocean.model.clock))
-
+
+    # Remove some potentially irksome callbacks from the ocean simulation
+    # TODO: also remove these from sea ice simulations
+    pop!(ocean.callbacks, :stop_time_exceeded)
+    pop!(ocean.callbacks, :stop_iteration_exceeded)
+    pop!(ocean.callbacks, :wall_time_limit_exceeded)
+    pop!(ocean.callbacks, :nan_checker)
+
+    # In case there was any doubt these are meaningless.
+    ocean.stop_time = Inf
+    ocean.stop_iteration = Inf
+    ocean.wall_time_limit = Inf
+
     # Contains information about flux contributions: bulk formula, prescribed fluxes, etc.
     fluxes = OceanSeaIceSurfaceFluxes(ocean, sea_ice; 
                                       atmosphere, 

src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl

@@ -28,28 +28,26 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_
         time_step!(sea_ice)
     end
 
-    ocean.Δt = Δt
-
     # TODO after ice time-step:
     #   - Adjust ocean heat flux if the ice completely melts?
 
+    ocean.Δt = Δt
     time_step!(ocean)
 
     # TODO:
     # - Store fractional ice-free / ice-covered _time_ for more
     #   accurate flux computation?
     tick!(coupled_model.clock, Δt)
-
     update_state!(coupled_model, callbacks; compute_tendencies)
 
     return nothing
 end
 
-function update_state!(coupled_model::OceanSeaIceModel, callbacks=[]; compute_tendencies=false)
+function update_state!(coupled_model::OceanSeaIceModel, callbacks=[]; compute_tendencies=true)
     time = Time(coupled_model.clock.time)
     update_model_field_time_series!(coupled_model.atmosphere, time)
     compute_atmosphere_ocean_fluxes!(coupled_model) 
-    compute_sea_ice_ocean_fluxes!(coupled_model)
+    #compute_sea_ice_ocean_fluxes!(coupled_model)
     #compute_atmosphere_sea_ice_fluxes!(coupled_model)
     return nothing
 end