Speeding up Dynamics and removing internal heat flux when h < hc (#89)

* this should speed up tremendously my computation

* maybe this works?

* go

* ok this should work?

* try it out

* should work

* this works!

* let's go

* no need to fill the halos here

* fix the periodic BC

* fix the periodic BC

* go like this

* converted args

* typo

* remove this

* try computing stresses

* go ahead

* this should work?

* correct

* try it out

* fix

* send it

* bump patch release

* Apply suggestions from code review

* Update README.md

* Update Project.toml

* RheologyAuxiliaries

* using Adapt

* tidy up

* more tidying up

* rename to auxiliaries

* extend

* Update elasto_visco_plastic_rheology.jl

* removed required auxiliaires

* change

---------

Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>

Author: simone-silvestri

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaSeaIce"
 uuid = "6ba0ff68-24e6-4315-936c-2e99227c95a4"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.3.5"
+version = "0.3.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -14,7 +14,7 @@ SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 [compat]
 Adapt = "3, 4"
 KernelAbstractions = "0.9"
-Oceananigans = "0.97.6"
+Oceananigans = "0.98.0"
 Roots = "2"
 RootSolvers = "0.3, 0.4"
 SeawaterPolynomials = "0.3.4"

README.md

@@ -36,4 +36,4 @@ When things have progressed, we'll put an outline here.
 If you use ClimaSeaIcea for your research, teaching, or fun 🤩, everyone in our community will be grateful
 if you give credit by citing the corresponding Zenodo record, e.g.,
 
-> Silvestri, S. et al. (2025). CliMA/ClimaSeaIce.jl: v0.3.3 (v0.3.3). Zenodo. https://doi.org/10.5281/zenodo.16143708
+> Silvestri, S. et al. (2025). CliMA/ClimaSeaIce.jl: v0.3.6 (v0.3.6). Zenodo. https://doi.org/10.5281/zenodo.16143708

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,35 @@
 module Rheologies
 
 export ViscousRheology, ElastoViscoPlasticRheology
-export ∂ⱼ_σ₁ⱼ, ∂ⱼ_σ₂ⱼ, required_auxiliary_fields
+export ∂ⱼ_σ₁ⱼ, ∂ⱼ_σ₂ⱼ, Auxiliaries
 
 using Oceananigans
 using Oceananigans.Operators
+using Oceananigans.Grids: AbstractGrid
 using ClimaSeaIce: ice_mass
+using Adapt 
+
+struct Auxiliaries{F, K}
+    fields :: F
+    kernels :: K
+end
+
+# When adapted, only the fields need to be passed to the GPU.
+# kernels operate only on the CPU.
+Adapt.adapt_structure(to, a::Auxiliaries) = 
+    Auxiliaries(Adapt.adapt(to, a.fields), nothing)
+
+""" 
+    Auxiliaries(rheology, grid)
+
+A struct holding any auxiliary fields and kernels needed for the computation of 
+sea ice stresses.
+"""
+Auxiliaries(rheology, grid::AbstractGrid) = Auxiliaries(NamedTuple(), nothing)
 
 # Nothing rheology
-required_auxiliary_fields(rheology, grid) = 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,15 @@ function ElastoViscoPlasticRheology(FT::DataType = Float64;
                                       pressure_formulation)
 end
 
-function required_auxiliary_fields(r::ElastoViscoPlasticRheology, grid)
-    
+# Extend Auxiliaries to hold auxiliaries for the ElastoViscoPlasticRheology
+function Auxiliaries(r::ElastoViscoPlasticRheology, grid::AbstractGrid)
+
+    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 +129,16 @@ 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!)[1]
+    _stresses_kernel!  = configure_kernel(arch, grid, parameters, _compute_evp_stresses!)[1]
+
+    parameters = KernelParameters(size(P.data)[1:2], P.data.offsets[1:2])
+    _initialize_rhology! = configure_kernel(arch, grid, parameters, _initialize_evp_rhology!)[1]
+    
+    fields  = (; σ₁₁, σ₂₂, σ₁₂, ζ, Δ, α, uⁿ, vⁿ, P)
+    kernels = (; _viscosity_kernel!, _stresses_kernel!, _initialize_rhology!)
+
+    return Auxiliaries(fields, kernels)
 end
 
 # Extend the `adapt_structure` function for the ElastoViscoPlasticRheology
@@ -149,13 +165,11 @@ function initialize_rheology!(model, rheology::ElastoViscoPlasticRheology)
     P★ = rheology.ice_compressive_strength
     C  = rheology.ice_compaction_hardening
 
-    u, v   = model.velocities
-    fields = model.dynamics.auxiliary_fields
+    u, v    = model.velocities
+    fields  = model.dynamics.auxiliaries.fields
+    kernels = model.dynamics.auxiliaries.kernels
+    kernels._initialize_rhology!(fields, model.grid, P★, C, h, ℵ, u, v)
 
-    # compute on the whole grid including halos
-    parameters = KernelParameters(size(fields.P.data)[1:2], fields.P.data.offsets[1:2])
-    launch!(architecture(model.grid), model.grid, parameters, _initialize_evp_rhology!, fields, model.grid, P★, C, h, ℵ, u, v)
-    
     return nothing
 end
 
@@ -170,25 +184,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,14 +18,14 @@ using ClimaSeaIce.Rheologies: ∂ⱼ_σ₁ⱼ,
                               ∂ⱼ_σ₂ⱼ, 
                               immersed_∂ⱼ_σ₁ⱼ,
                               immersed_∂ⱼ_σ₂ⱼ,
-                              required_auxiliary_fields, 
+                              Auxiliaries,
                               compute_stresses!,
                               initialize_rheology!,
                               compute_substep_Δtᶠᶜᶜ,
                               compute_substep_Δtᶜᶠᶜ,
                               sum_of_forcing_u,
                               sum_of_forcing_v
-
+                              
 import Oceananigans: fields
 
 ## A Framework to solve for the ice momentum equation, in the form:

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 = 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