Refactor SeaIceModel for simpler combined timestepping

src/sea_ice_advection.jl → src/Advection.jl

@@ -1,3 +1,7 @@
+module Advection
+
+export div_Uℵh, horizontal_div_Uc
+
 using Oceananigans.Operators
 using Oceananigans.ImmersedBoundaries
 using Oceananigans.Advection: FluxFormAdvection, 
@@ -16,7 +20,6 @@ using Oceananigans.Advection: FluxFormAdvection,
 # A = ∇ ⋅ (uh)
 
 _advective_thickness_flux_x(args...) = advective_thickness_flux_x(args...)
-
 _advective_thickness_flux_y(args...) = advective_thickness_flux_y(args...)
 
 _advective_thickness_flux_x(i, j, k, ibg::ImmersedBoundaryGrid, args...) = 
@@ -44,7 +47,9 @@ end
     ∇Uℵh = 1 / Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, _advective_thickness_flux_x, advection, U.u, ℵ, h) +
                                       δyᵃᶜᵃ(i, j, k, grid, _advective_thickness_flux_y, advection, U.v, ℵ, h))
 
-    @inbounds ℵ⁻¹ = ifelse(ℵ[i, j, k] != 0, 1 / ℵ[i, j, k], zero(grid))
+    ℵᵢ = @inbounds ℵ[i, j, k]
+    non_zero_area = @inbounds ℵᵢ != 0
+    ℵ⁻¹ = ifelse(non_zero_area, 1 / ℵᵢ, zero(grid))
 
     return ℵ⁻¹ * ∇Uℵh
 end
@@ -54,7 +59,9 @@ end
     ∇Uℵh = 1 / Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, _advective_thickness_flux_x, advection.x, U.u, ℵ, h) +
                                       δyᵃᶜᵃ(i, j, k, grid, _advective_thickness_flux_y, advection.y, U.v, ℵ, h))
 
-    @inbounds ℵ⁻¹ = ifelse(ℵ[i, j, k] != 0, 1 / ℵ[i, j, k], zero(grid))
+    ℵᵢ = @inbounds ℵ[i, j, k]
+    non_zero_area = @inbounds ℵᵢ != 0
+    ℵ⁻¹ = ifelse(non_zero_area, 1 / ℵᵢ, zero(grid))
 
     return ℵ⁻¹ * ∇Uℵh
 end
@@ -67,4 +74,5 @@ end
 @inline horizontal_div_Uc(i, j, k, grid, advection::FluxFormAdvection, U, c) = 
     1 / Vᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, _advective_tracer_flux_x, advection.x, U.u, c) +
                                δyᵃᶜᵃ(i, j, k, grid, _advective_tracer_flux_y, advection.y, U.v, c))
-
+
+end # Advection

src/ClimaSeaIce.jl

@@ -1,4 +1,4 @@
-""" Ocean 🌊 Sea ice component of CliMa's Earth system model. """
+""" Ocean 🌊 Sea ice component of CliMA's Earth system model. """
 module ClimaSeaIce
 
 using Oceananigans
@@ -40,15 +40,13 @@ mask_immersed_field!(::ConstantField) = nothing
 mask_immersed_field!(::ZeroField)     = nothing
 
 include("SeaIceThermodynamics/SeaIceThermodynamics.jl")
-include("SeaIceDynamics/SeaIceDynamics.jl")
-include("sea_ice_model.jl")
-include("sea_ice_advection.jl")
-include("tracer_tendency_kernel_functions.jl")
-include("sea_ice_time_stepping.jl")
-include("sea_ice_ab2_time_stepping.jl")
-include("sea_ice_rk3_time_stepping.jl")
+include("Rheologies.jl")
+include("Advection.jl")
+include("SeaIceModels/SeaIceModels.jl")
 include("EnthalpyMethodSeaIceModel.jl")
 
 using .SeaIceThermodynamics
+using .SeaIceModels
 
 end # module
+

src/SeaIceDynamics/SeaIceDynamics.jl → src/Rheologies.jl

@@ -1,4 +1,4 @@
-module SeaIceDynamics
+module Rheologies
 
 using ClimaSeaIce
 
@@ -23,6 +23,7 @@ using Oceananigans.Grids: architecture
 ## - ice-atmosphere boundary stress (provided as an external flux)
 ## - ocean dynamic surface
 
-include("nothing_dynamics.jl")
+compute_momentum_tendencies!(model) = nothing
+
+end
 
-end

src/SeaIceModels/SeaIceModels.jl

@@ -0,0 +1,7 @@
+module SeaIceModels
+
+include("sea_ice_model.jl")
+include("compute_tendencies.jl")
+include("sea_ice_time_stepping.jl")
+
+end # SeaIceModels

src/tracer_tendency_kernel_functions.jl → src/SeaIceModels/compute_tendencies.jl

@@ -1,6 +1,12 @@
-using Oceananigans.Advection
+using ClimaSeaIce.Advection: div_Uℵh, horizontal_div_Uc
 using ClimaSeaIce.SeaIceThermodynamics: thickness_thermodynamic_tendency
 
+function compute_tendencies!(model::SIM)
+    compute_tracer_tendencies!(model)
+    # compute_momentum_tendencies!(model)
+    return nothing
+end
+
 function compute_tracer_tendencies!(model::SIM)
     grid = model.grid
     arch = architecture(grid)
@@ -17,7 +23,6 @@ function compute_tracer_tendencies!(model::SIM)
             model.ice_thermodynamics,
             model.external_heat_fluxes.top,
             model.external_heat_fluxes.bottom,
-            nothing, #model.forcing.h,
             fields(model))
 
     return nothing
@@ -32,36 +37,9 @@ end
                                              thermodynamics,
                                              top_external_heat_flux,
                                              bottom_external_heat_flux,
-                                             h_forcing,
                                              model_fields)
 
     i, j, k = @index(Global, NTuple)
-
-    @inbounds Gⁿ.h[i, j, k] = ice_thickness_tendency(i, j, k, grid, clock,
-                                                     velocities,
-                                                     advection,
-                                                     ice_thickness,
-                                                     concentration,
-                                                     thermodynamics,
-                                                     top_external_heat_flux,
-                                                     bottom_external_heat_flux,
-                                                     h_forcing,
-                                                     model_fields)
-
-    @inbounds Gⁿ.ℵ[i, j, k] = - horizontal_div_Uc(i, j, k, grid, advection, velocities, concentration)
-end
-
-# Thickness change due to accretion and melting, restricted by minimum allowable value
-function ice_thickness_tendency(i, j, k, grid, clock,
-                                velocities,
-                                advection,
-                                ice_thickness,
-                                ice_concentration,
-                                thermodynamics,
-                                top_external_heat_flux,
-                                bottom_external_heat_flux,
-                                h_forcing,
-                                model_fields)
 
     Gh_advection = - div_Uℵh(i, j, k, grid, advection, velocities, ice_concentration, ice_thickness)
 
@@ -72,10 +50,8 @@ function ice_thickness_tendency(i, j, k, grid, clock,
                                                          top_external_heat_flux,
                                                          bottom_external_heat_flux,
                                                          clock, model_fields)
-
 
-    # Compute forcing
-    Fh = zero(grid) #h_forcing(i, j, grid, clock, model_fields)
-
-    return Gh_advection + Gh_thermodynamics + Fh
+    @inbounds Gⁿ.h[i, j, k] = Gh_advection + Gh_thermodynamics
+    @inbounds Gⁿ.ℵ[i, j, k] = - horizontal_div_Uc(i, j, k, grid, advection, velocities, concentration)
 end
+

src/sea_ice_model.jl → src/SeaIceModels/sea_ice_model.jl

@@ -15,7 +15,7 @@ struct SeaIceModel{GR, TD, D, TS, CL, U, T, IT, IC, STF, SMS, A} <: AbstractMode
     ice_concentration :: IC
     # Thermodynamics
     ice_thermodynamics :: TD
-    ice_dynamics :: D
+    rheology :: D
     # External boundary conditions
     external_heat_fluxes :: STF
     external_momentum_stresses :: SMS
@@ -32,13 +32,11 @@ function SeaIceModel(grid;
                      top_heat_flux       = nothing,
                      bottom_heat_flux    = 0,
                      velocities          = nothing,
-                     timestepper         = :RungeKutta3,
                      advection           = nothing,
-                     top_momentum_stress = nothing, # Fix when introducing dynamics
                      tracers             = (),
                      boundary_conditions = NamedTuple(),
                      ice_thermodynamics  = SlabSeaIceThermodynamics(grid),
-                     ice_dynamics        = nothing)
+                     rheology            = nothing)
 
     if isnothing(velocities)
         velocities = (u = ZeroField(), v=ZeroField(), w=ZeroField())
@@ -62,7 +60,7 @@ function SeaIceModel(grid;
     # TODO: should we have ice thickness and concentration as part of the tracers or
     # just additional fields of the sea ice model?
     tracers = merge(tracers, (; S = ice_salinity))
-    timestepper = TimeStepper(timestepper, grid, prognostic_tracers)
+    timestepper = TimeStepper(:RungeKutta3, grid, prognostic_tracers)
 
     top_heat_flux = external_top_heat_flux(ice_thermodynamics, top_heat_flux)
 
@@ -77,9 +75,9 @@ function SeaIceModel(grid;
                        ice_thickness,
                        ice_concentration,
                        ice_thermodynamics,
-                       ice_dynamics,
+                       rheology,
                        external_heat_fluxes,
-                       top_momentum_stress,
+                       nothing,
                        timestepper,
                        advection)
 end
@@ -88,7 +86,7 @@ const SIM = SeaIceModel
 
 function set!(model::SIM; h=nothing, ℵ=nothing)
     !isnothing(h) && set!(model.ice_thickness, h)
-    !isnothing(ℵ) && set!(model.ice_conentration, ℵ)
+    !isnothing(ℵ) && set!(model.ice_concentration, ℵ)
     return nothing
 end
 

src/sea_ice_time_stepping.jl → src/SeaIceModels/sea_ice_time_stepping.jl

@@ -1,4 +1,4 @@
-function step_tracers!(model::SIM, Δt, substep)
+function rk3_step!(model::SIM, Δt, γ, ζ)
     grid = model.grid
     arch = architecture(grid)
 
@@ -9,9 +9,8 @@ function step_tracers!(model::SIM, Δt, substep)
     Gⁿ = model.timestepper.Gⁿ
     G⁻ = model.timestepper.G⁻
 
-    α, β = timestepping_coefficients(model.timestepper, substep)
-
-    launch!(arch, grid, :xyz, _step_tracers!, h, ℵ, tracers, Gⁿ, G⁻, Δt, α, β)
+    launch!(arch, grid, :xyz, _step_tracers!, h, ℵ, tracers, Gⁿ, G⁻, Δt, γ, ζ)
+    #launch!(arch, grid, :xyz, _rk3_step_fields!, h, ℵ, tracers, Gⁿ, G⁻, Δt, γ, ζ)
 
     return nothing
 end
@@ -25,7 +24,6 @@ end
 
     Ghⁿ = Gⁿ.h
     Gℵⁿ = Gⁿ.ℵ
-
     Gh⁻ = G⁻.h
     Gℵ⁻ = G⁻.ℵ
 
@@ -45,16 +43,15 @@ end
 end
 
 function store_tendencies!(model::SIM) 
-
     grid = model.grid
     arch = architecture(grid)
     Nx, Ny, _ = size(grid)
 
     Gⁿ = model.timestepper.Gⁿ
     G⁻ = model.timestepper.G⁻
-    Nt = length(Gⁿ)
+    NG = length(Gⁿ)
 
-    params = KernelParameters((Nx, Ny, Nt), (0, 0, 0))
+    params = KernelParameters((Nx, Ny, NG), (0, 0, 0))
     launch!(arch, model.grid, params, _store_tendencies!, G⁻, Gⁿ)
 
     return nothing
@@ -65,8 +62,58 @@ end
     @inbounds G⁻[n][i, j, 1] = Gⁿ[n][i, j, 1]
 end
 
+const RK3SeaIceModel = SeaIceModel{<:Any, <:Any, <:Any, <:RungeKutta3TimeStepper}
+
+function time_step!(model::RK3SeaIceModel, Δt; callbacks = [])
+
+    # Be paranoid and update state at iteration 0, in case run! is not used:
+    model.clock.iteration == 0 && update_state!(model)
+
+    γ¹ = model.timestepper.γ¹
+    γ² = model.timestepper.γ²
+    γ³ = model.timestepper.γ³
+
+    ζ² = model.timestepper.ζ²
+    ζ³ = model.timestepper.ζ³
+
+    first_stage_Δt  = γ¹ * Δt
+    second_stage_Δt = (γ² + ζ²) * Δt
+    third_stage_Δt  = (γ³ + ζ³) * Δt
+
+    #
+    # First stage
+    #
+
+    compute_tendencies!(model)
+    rk3_step!(model, Δt, γ¹, zero(γ¹))
+    store_tendencies!(model)
+    tick!(model.clock, first_stage_Δt)
+    update_state!(model)
+
+    #
+    # Second stage
+    #
+
+    compute_tendencies!(model)
+    rk3_step!(model, Δt, γ², ζ²)
+    store_tendencies!(model)
+    tick!(model.clock, second_stage_Δt)
+    update_state!(model)
+
+    #
+    # Third stage
+    #
+
+    compute_tendencies!(model)
+    rk3_step!(model, Δt, γ³, ζ³)
+    store_tendencies!(model)
+    tick!(model.clock, third_stage_Δt)
+    update_state!(model)
+
+    return nothing
+end
+
 function update_state!(model::SIM)
-
     foreach(prognostic_fields(model)) do field
         mask_immersed_field!(field)
     end
@@ -75,3 +122,5 @@ function update_state!(model::SIM)
 
     return nothing
 end
+
+