Speeding up Dynamics and removing internal heat flux when h < hc

examples/ice_advected_by_anticyclone.jl

@@ -76,8 +76,8 @@ set!(Vₐ, (x, y) -> va_time(x, y, 0))
 
 fill_halo_regions!((Uₐ, Vₐ))
 
-τₐu = - Uₐ * sqrt(Uₐ^2 + Vₐ^2) * 1.3 * 1.2e-3
-τₐv = - Vₐ * sqrt(Uₐ^2 + Vₐ^2) * 1.3 * 1.2e-3
+τₐu = Field(- Uₐ * sqrt(Uₐ^2 + Vₐ^2) * 1.3 * 1.2e-3)
+τₐv = Field(- Vₐ * sqrt(Uₐ^2 + Vₐ^2) * 1.3 * 1.2e-3)
 
 #####
 ##### Numerical details
@@ -124,6 +124,9 @@ function compute_wind_stress(sim)
 
     fill_halo_regions!((Uₐ, Vₐ))
 
+    compute!(τₐu)
+    compute!(τₐv)
+
     return nothing
 end
 
@@ -177,10 +180,10 @@ vtimeseries = FieldTimeSeries("sea_ice_advected_by_anticyclone.jld2", "v")
 Nt = length(htimeseries)
 iter = Observable(1)
 
-hi = @lift(htimeseries[$iter])
-ℵi = @lift(ℵtimeseries[$iter])
-ui = @lift(utimeseries[$iter])
-vi = @lift(vtimeseries[$iter])
+hi = @lift(interior(htimeseries[$iter], :, :, 1))
+ℵi = @lift(interior(ℵtimeseries[$iter], :, :, 1))
+ui = @lift(interior(utimeseries[$iter], :, :, 1))
+vi = @lift(interior(vtimeseries[$iter], :, :, 1))
 
 fig = Figure()
 ax = Axis(fig[1, 1], title = "sea ice thickness")

src/Rheologies/Rheologies.jl

@@ -1,16 +1,16 @@
 module Rheologies
 
 export ViscousRheology, ElastoViscoPlasticRheology
-export ∂ⱼ_σ₁ⱼ, ∂ⱼ_σ₂ⱼ, required_auxiliary_fields
+export ∂ⱼ_σ₁ⱼ, ∂ⱼ_σ₂ⱼ, required_auxiliaries
 
 using Oceananigans
 using Oceananigans.Operators
 using ClimaSeaIce: ice_mass
 
 # Nothing rheology
-required_auxiliary_fields(rheology, grid) = NamedTuple()
+required_auxiliaries(rheology, grid) = (; fields = NamedTuple())
 initialize_rheology!(model, rheology) = nothing
-compute_stresses!(model, dynamics, rheology, Δt) = nothing
+compute_stresses!(dynamics, fields, grid, rheology, Δt) = nothing
 
 # Nothing rheology or viscous rheology
 @inline compute_substep_Δtᶠᶜᶜ(i, j, grid, Δt, rheology, substeps, fields) = Δt / substeps

src/Rheologies/elasto_visco_plastic_rheology.jl

@@ -105,8 +105,21 @@ function ElastoViscoPlasticRheology(FT::DataType = Float64;
                                       pressure_formulation)
 end
 
-function required_auxiliary_fields(r::ElastoViscoPlasticRheology, grid)
-
+struct ElastoViscoPlasticAuxiliaries{F, K}
+    fields :: F
+    kernels :: K
+end
+
+Adapt.adapt_structure(to, a::ElastoViscoPlasticAuxiliaries) = Adapt.adapt(to, a.fields)
+
+function required_auxiliaries(r::ElastoViscoPlasticRheology, grid)
+
+    arch      = architecture(grid)
+    Nx, Ny, _ = size(grid)
+    Hx, Hy, _ = halo_size(grid)
+
+    parameters = KernelParameters(-Hx+2:Nx+Hx-1, -Hy+2:Ny+Hy-1)
+
     # TODO: What about boundary conditions?
     σ₁₁ = Field{Center, Center, Nothing}(grid)
     σ₂₂ = Field{Center, Center, Nothing}(grid)
@@ -122,7 +135,10 @@ function required_auxiliary_fields(r::ElastoViscoPlasticRheology, grid)
     # An initial (safe) educated guess
     fill!(α, r.max_relaxation_parameter)
 
-    return (; σ₁₁, σ₂₂, σ₁₂, ζ, Δ, α, uⁿ, vⁿ, P)
+    _viscosity_kernel!, _ = configure_kernel(arch, grid, parameters, _compute_evp_viscosities!)
+    _stresses_kernel!, _  = configure_kernel(arch, grid, parameters, _compute_evp_stresses!)
+
+    return ElastoViscoPlasticAuxiliaries((; σ₁₁, σ₂₂, σ₁₂, ζ, Δ, α, uⁿ, vⁿ, P), (; _viscosity_kernel!, _stresses_kernel!))
 end
 
 # Extend the `adapt_structure` function for the ElastoViscoPlasticRheology
@@ -150,7 +166,7 @@ function initialize_rheology!(model, rheology::ElastoViscoPlasticRheology)
     C  = rheology.ice_compaction_hardening
 
     u, v   = model.velocities
-    fields = model.dynamics.auxiliary_fields
+    fields = model.dynamics.auxiliaries.fields
 
     # compute on the whole grid including halos
     parameters = KernelParameters(size(fields.P.data)[1:2], fields.P.data.offsets[1:2])
@@ -170,25 +186,16 @@ end
 @inline ice_strength(i, j, k, grid, P★, C, h, ℵ) = @inbounds P★ * h[i, j, k] * exp(- C * (1 - ℵ[i, j, k])) 
 
 # Specific compute stresses for the EVP rheology
-function compute_stresses!(model, dynamics, rheology::ElastoViscoPlasticRheology, Δt) 
-
-    grid = model.grid
-    arch = architecture(grid)
-
-    h  = model.ice_thickness
-    ρᵢ = model.ice_density
-    ℵ  = model.ice_concentration
-
-    fields = dynamics.auxiliary_fields
-    u, v = model.velocities
-
-    Nx, Ny, _ = size(grid)
-    Hx, Hy, _ = halo_size(grid)
-
-    parameters = KernelParameters(-Hx+2:Nx+Hx-1, -Hy+2:Ny+Hy-1)
-
-    launch!(arch, grid, parameters, _compute_evp_viscosities!, fields, grid, rheology, u, v)
-    launch!(arch, grid, parameters, _compute_evp_stresses!, fields, grid, rheology, u, v, h, ℵ, ρᵢ, Δt)
+function compute_stresses!(dynamics, fields, grid, rheology::ElastoViscoPlasticRheology, Δt)
+
+    h  = fields.h
+    ρᵢ = fields.ρ
+    ℵ  = fields.ℵ
+    u  = fields.u
+    v  = fields.v
+
+    dynamics.auxiliaries.kernels._viscosity_kernel!(fields, grid, rheology, u, v)
+    dynamics.auxiliaries.kernels._stresses_kernel!(fields, grid, rheology, u, v, h, ℵ, ρᵢ, Δt)
 
     return nothing
 end

src/SeaIceDynamics/SeaIceDynamics.jl

@@ -18,7 +18,7 @@ using ClimaSeaIce.Rheologies: ∂ⱼ_σ₁ⱼ,
                               ∂ⱼ_σ₂ⱼ, 
                               immersed_∂ⱼ_σ₁ⱼ,
                               immersed_∂ⱼ_σ₂ⱼ,
-                              required_auxiliary_fields, 
+                              required_auxiliaries,
                               compute_stresses!,
                               initialize_rheology!,
                               compute_substep_Δtᶠᶜᶜ,

src/SeaIceDynamics/explicit_momentum_equations.jl

@@ -12,7 +12,7 @@ function time_step_momentum!(model, ::ExplicitMomentumEquation, Δt)
 
     dynamics = model.dynamics
 
-    model_fields = merge(dynamics.auxiliary_fields, model.velocities, 
+    model_fields = merge(dynamics.auxiliaries.fields, model.velocities, 
                       (; h = model.ice_thickness, 
                          ℵ = model.ice_concentration, 
                          ρ = model.ice_density))
@@ -77,7 +77,7 @@ function compute_momentum_tendencies!(model, ::ExplicitMomentumEquation, Δt)
     coriolis = dynamics.coriolis
     rheology = dynamics.rheology
 
-    model_fields = merge(dynamics.auxiliary_fields, model.velocities, 
+    model_fields = merge(dynamics.auxiliaries.fields, model.velocities, 
             (; h = model.ice_thickness, 
                ℵ = model.ice_concentration, 
                ρ = model.ice_density))

src/SeaIceDynamics/sea_ice_momentum_equations.jl

@@ -4,7 +4,7 @@ using Adapt
 struct SeaIceMomentumEquation{S, C, R, F, A, ES, FT}
     coriolis :: C
     rheology :: R
-    auxiliary_fields :: A
+    auxiliaries :: A
     solver :: S
     free_drift :: F
     external_momentum_stresses :: ES
@@ -19,7 +19,6 @@ struct ExplicitSolver end
     SeaIceMomentumEquation(grid; 
                            coriolis = nothing,
                            rheology = ElastoViscoPlasticRheology(eltype(grid)),
-                           auxiliary_fields = NamedTuple(),
                            top_momentum_stress    = nothing,
                            bottom_momentum_stress = nothing,
                            free_drift = nothing,
@@ -48,7 +47,6 @@ Keyword Arguments
 
 - `coriolis`: Parameters for the background rotation rate of the model.
 - `rheology`: The sea ice rheology model, default is `ElastoViscoPlasticRheology(eltype(grid))`.
-- `auxiliary_fields`: A named tuple of auxiliary fields, default is an empty `NamedTuple()`.
 - `free_drift`: The free drift velocities used to limit sea ice momentum when the mass or the concentration are
                 below a certain threshold. Default is `nothing` (indicating that the free drift velocities are zero).
 - `solver`: The momentum solver to be used.
@@ -58,28 +56,27 @@ Keyword Arguments
 function SeaIceMomentumEquation(grid; 
                                 coriolis = nothing,
                                 rheology = ElastoViscoPlasticRheology(eltype(grid)),
-                                auxiliary_fields = NamedTuple(),
                                 top_momentum_stress    = nothing,
                                 bottom_momentum_stress = nothing,
                                 free_drift = nothing,
                                 solver = SplitExplicitSolver(150),
                                 minimum_concentration = 1e-3,
                                 minimum_mass = 1.0)
 
-    auxiliary_fields = merge(auxiliary_fields, required_auxiliary_fields(rheology, grid))
+    auxiliaries = required_auxiliaries(rheology, grid)
     external_momentum_stresses = (top = top_momentum_stress,
                                   bottom = bottom_momentum_stress)
 
     FT = eltype(grid)
 
     return SeaIceMomentumEquation(coriolis, 
                                   rheology, 
-                                  auxiliary_fields, 
+                                  auxiliaries, 
                                   solver,
                                   free_drift,
                                   external_momentum_stresses,
                                   convert(FT, minimum_concentration),
                                   convert(FT, minimum_mass))
 end
 
-fields(mom::SeaIceMomentumEquation) = mom.auxiliary_fields
+fields(mom::SeaIceMomentumEquation) = mom.auxiliaries.fields

src/SeaIceDynamics/split_explicit_momentum_equations.jl

@@ -1,7 +1,8 @@
-using Oceananigans.Grids: AbstractGrid, architecture
-using Oceananigans.BoundaryConditions: fill_halo_regions!
+using Oceananigans.Grids: AbstractGrid, architecture, halo_size
+using Oceananigans.BoundaryConditions: fill_halo_regions!, fill_halo_size, fill_halo_offset
 using Oceananigans.Utils: configure_kernel
-using Oceananigans.ImmersedBoundaries: mask_immersed_field_xy!
+using Oceananigans.Fields: instantiated_location, boundary_conditions
+using Oceananigans.ImmersedBoundaries: peripheral_node
 
 struct SplitExplicitSolver 
     substeps :: Int
@@ -31,9 +32,10 @@ function time_step_momentum!(model, dynamics::SplitExplicitMomentumEquation, Δt
 
     grid = model.grid
     arch = architecture(grid)
-    rheology = dynamics.rheology
-
-    u, v = model.velocities
+    rheology   = dynamics.rheology
+    Nx, Ny, Nz = size(grid)
+    Hx, Hy, _  = halo_size(grid)
+    u, v       = model.velocities
 
     free_drift = dynamics.free_drift
     clock = model.clock
@@ -50,7 +52,7 @@ function time_step_momentum!(model, dynamics::SplitExplicitMomentumEquation, Δt
     u_immersed_bc = u.boundary_conditions.immersed
     v_immersed_bc = v.boundary_conditions.immersed
 
-    model_fields = merge(dynamics.auxiliary_fields, model.velocities, 
+    model_fields = merge(dynamics.auxiliaries.fields, model.velocities, 
                       (; h = model.ice_thickness, 
                          ℵ = model.ice_concentration, 
                          ρ = model.ice_density))
@@ -61,47 +63,52 @@ function time_step_momentum!(model, dynamics::SplitExplicitMomentumEquation, Δt
     v_velocity_kernel!, _ = configure_kernel(arch, grid, :xy, _v_velocity_step!; active_cells_map)
 
     substeps = dynamics.solver.substeps
-
-    fill_halo_regions!(model.velocities)
     initialize_rheology!(model, dynamics.rheology)
 
-    for substep in 1 : substeps
-        # Compute stresses! depending on the particular rheology implementation
-        compute_stresses!(model, dynamics, rheology, Δt)
-
-        # The momentum equations are solved using an alternating leap-frog algorithm
-        # for u and v (used for the ocean - ice stresses and the coriolis term)
-        # In even substeps we calculate uⁿ⁺¹ = f(vⁿ) and vⁿ⁺¹ = f(uⁿ⁺¹).
-        # In odd substeps we switch and calculate vⁿ⁺¹ = f(uⁿ) and uⁿ⁺¹ = f(vⁿ⁺¹).
-        if iseven(substep) 
-            u_velocity_kernel!(u, grid, Δt, substeps, rheology, model_fields, 
-                               free_drift, clock, coriolis,
-                               minimum_mass, minimum_concentration, 
-                               u_immersed_bc, top_stress, bottom_stress, u_forcing)
-
-            v_velocity_kernel!(v, grid, Δt, substeps, rheology, model_fields, 
-                               free_drift, clock, coriolis, 
-                               minimum_mass, minimum_concentration,
-                               v_immersed_bc, top_stress, bottom_stress, v_forcing)
-
-        else
-            v_velocity_kernel!(v, grid, Δt, substeps, rheology, model_fields, 
-                               free_drift, clock, coriolis, 
-                               minimum_mass, minimum_concentration,
-                               v_immersed_bc, top_stress, bottom_stress, v_forcing)
-
-
-            u_velocity_kernel!(u, grid, Δt, substeps, rheology, model_fields, 
-                               free_drift, clock, coriolis,
-                               minimum_mass, minimum_concentration, 
-                               u_immersed_bc, top_stress, bottom_stress, u_forcing)
+    u_args = (u, grid, Δt, substeps, rheology, model_fields, 
+              free_drift, clock, coriolis,
+              minimum_mass, minimum_concentration, 
+              u_immersed_bc, top_stress, bottom_stress, u_forcing)
+
+    v_args = (v, grid, Δt, substeps, rheology, model_fields, 
+              free_drift, clock, coriolis, 
+              minimum_mass, minimum_concentration,
+              v_immersed_bc, top_stress, bottom_stress, v_forcing)
+
+    u_fill_halo_args = (u.data, u.boundary_conditions, u.indices, (Face(), Center(), nothing), grid)
+    v_fill_halo_args = (v.data, v.boundary_conditions, v.indices, (Center(), Face(), nothing), grid)
+    stresses_args    = (model_fields, grid, rheology, Δt)
+
+    GC.@preserve v_args u_args u_fill_halo_args v_fill_halo_args stresses_args begin
+        # We need to perform ~150 time-steps which means
+        # launching ~300 very small kernels: we are limited by
+        # latency of argument conversion to GPU-compatible values.
+        # To alleviate this penalty we convert first and then we substep!
+        converted_u_args = Oceananigans.Architectures.convert_to_device(arch, u_args)
+        converted_v_args = Oceananigans.Architectures.convert_to_device(arch, v_args)
+        converted_u_halo = Oceananigans.Architectures.convert_to_device(arch, u_fill_halo_args)
+        converted_v_halo = Oceananigans.Architectures.convert_to_device(arch, v_fill_halo_args)
+        converted_stresses_args = Oceananigans.Architectures.convert_to_device(arch, stresses_args)
+
+        for substep in 1 : substeps
+            # Compute stresses! depending on the particular rheology implementation
+            compute_stresses!(dynamics, converted_stresses_args...)
+
+            # The momentum equations are solved using an alternating leap-frog algorithm
+            # for u and v (used for the ocean - ice stresses and the coriolis term)
+            # In even substeps we calculate uⁿ⁺¹ = f(vⁿ) and vⁿ⁺¹ = f(uⁿ⁺¹).
+            # In odd substeps we switch and calculate vⁿ⁺¹ = f(uⁿ) and uⁿ⁺¹ = f(vⁿ⁺¹).
+            if iseven(substep) 
+                u_velocity_kernel!(converted_u_args...)
+                v_velocity_kernel!(converted_v_args...)
+            else
+                v_velocity_kernel!(converted_v_args...)
+                u_velocity_kernel!(converted_u_args...)
+            end
+
+            fill_halo_regions!(converted_u_halo...)
+            fill_halo_regions!(converted_v_halo...)
         end
-
-        # TODO: This needs to be removed in some way!
-        fill_halo_regions!(model.velocities)
-
-        mask_immersed_field_xy!(model.velocities.u, k=size(grid, 3))
-        mask_immersed_field_xy!(model.velocities.v, k=size(grid, 3))
     end
 
     return nothing
@@ -134,8 +141,9 @@ end
     # If the ice mass or the ice concentration are below a certain threshold, 
     # the sea ice velocity is set to the free drift velocity
     sea_ice = (mᵢ ≥ minimum_mass) & (ℵᵢ ≥ minimum_concentration)
+    active  = !peripheral_node(i, j, kᴺ, grid, Face(), Center(), Center())
 
-    @inbounds u[i, j, 1] = ifelse(sea_ice, uᴰ, uᶠ)
+    @inbounds u[i, j, 1] = ifelse(sea_ice, uᴰ, uᶠ) * active
 end
 
 @kernel function _v_velocity_step!(v, grid, Δt, 
@@ -166,6 +174,7 @@ end
     # If the ice mass or the ice concentration are below a certain threshold, 
     # the sea ice velocity is set to the free drift velocity
     sea_ice = (mᵢ ≥ minimum_mass) & (ℵᵢ ≥ minimum_concentration)
+    active  = !peripheral_node(i, j, kᴺ, grid, Center(), Face(), Center())
 
-    @inbounds v[i, j, 1] = ifelse(sea_ice, vᴰ, vᶠ)
+    @inbounds v[i, j, 1] = ifelse(sea_ice, vᴰ, vᶠ) * active
 end

src/SeaIceThermodynamics/slab_thermodynamics_tendencies.jl

@@ -57,17 +57,17 @@ using Oceananigans
 
     # Retrieve fluxes
     Quᵢ = getflux(Qu, i, j, grid, Tuᵢ, clock, model_fields)
-    Qiᵢ = getflux(Qi, i, j, grid, Tuᵢ, clock, model_fields)
     Qbᵢ = getflux(Qb, i, j, grid, Tuᵢ, clock, model_fields)
 
+    if consolidated_ice # If the ice is consolidated, we use the internal heat flux
+        Qiᵢ = getflux(Qi, i, j, grid, Tuᵢ, clock, model_fields)
+    else # Slab is unconsolidated, there is no internal heat flux (Qi -> ∞)
+        Qiᵢ = zero(grid)
+    end
+
     # Upper (top) and bottom interface velocities
     wu = (Quᵢ - Qiᵢ) / ℰu # < 0 => melting
-    wb =      + Qiᵢ  / ℰb # < 0 => freezing
-
-    # Ice forming at the bottom.
-    # it applies to the whole area, so it need 
-    # not be multiplied by the concentration
-    wf = - Qbᵢ / ℰb
-
-    return wu + wb + wf
-end
+    wb = (Qiᵢ - Qbᵢ) / ℰb # < 0 => freezing
+
+    return wu + wb 
+end