Thermal shallow water equations on the cubed sphere manifold (intrinsic formulation)

This example solves the thermal shallow water equations. The evolution of the prognostic variables, $\widetilde{\boldsymbol{u}}$ and $\widetilde{\varphi}$ is:

\[\begin{align*} \partial_t \widetilde{\boldsymbol{u}} + \widetilde{q} \widetilde{\boldsymbol{F}}^\dagger + \nabla_{\gamma} \widetilde{\Phi} + \widetilde{b}\nabla_{\gamma}\vartheta &= 0 \quad \text{in} \quad \gamma\times[0,T], \\ \partial_t \widetilde{\varphi} + \nabla_{\gamma}\cdot \widetilde{\boldsymbol{F}} &= 0 \quad \quad \text{in} \quad \gamma\times[0,T], \\ \partial_t \widetilde{B} + \nabla_{\gamma}\cdot (\widetilde{b} \widetilde{\boldsymbol{F}}) &= 0 \quad \quad \text{in} \quad \gamma\times[0,T], \end{align*}\]

The diagnostic variables solve algebraic constraints:

\[\begin{align*} \widetilde{\boldsymbol{F}} &= \widetilde{\varphi}\widetilde{\boldsymbol{u}} , \qquad \widetilde{\Phi} = \frac{1}{2} \widetilde{\boldsymbol{u}}\cdot\widetilde{\boldsymbol{u}} + \frac{1}{2} \widetilde{B} ,\qquad \widetilde{\vartheta} = \frac{1}{2} \widetilde{\varphi} \\ \widetilde{q} &= \frac{1 }{\widetilde{\varphi}}(\nabla_{\gamma}^\dagger\cdot \widetilde{\boldsymbol{u}} + f) , \qquad \widetilde{b} = \frac{\widetilde{B}}{\widetilde{\varphi}} \end{align*}\]

where $\gamma$ is the cubed sphere manifold, $g_r$ is the acceleration due to gravitiy, and $f$ is the Coriolis force. We use compatible finite element pairings to solve in the parametric space of the cubed sphere. This problem is a differential algebraic equation (DAE), which is solved using a bespoke $DAEFEOperator$ that builds on Gridap's ODE API.

Set up

First load all required packages. In this example, we will use a distributed model and iterative solvers. So we initialise MPI.

using GridapGeosciences
using Gridap
using GridapSolvers
using GridapDistributed
using PartitionedArrays
using MPI

MPI.Init()
ranks = distribute_with_mpi(LinearIndices((prod(MPI.Comm_size(MPI.COMM_WORLD)),)))

Discrete model

To obtain a refined 2D parametric model, we pass $\ell$ levels of refinement:

ℓ = 2
radius = 1.0
coarse_mesh = CubedSphereMesh(radius)
omodel = AtlasOctreeDistributedDiscreteModel(ranks, coarse_mesh, ℓ; manifold_style=IntrinsicManifold())

Triangulation

Now we extract the triangulation, volume and skeleton, associated to each cell:

Ω = Triangulation(omodel)
Λ = SkeletonTriangulation(omodel)
n_Λ = get_normal_vector(Λ)

FE Spaces

Trial and test spaces relevant to compatible finite element pairings are defined using Gridap's high level API (refer to Cotter et al. 2012).

order = 1
R = TestFESpace(Ω, ReferenceFE(lagrangian,Float64,order+1); conformity=:H1)
H = TransientTrialFESpace(R)

V = TestFESpace(Ω, ReferenceFE(raviart_thomas,Float64,order); conformity=:HDiv)
U = TransientTrialFESpace(V)

Q = TestFESpace(Ω, ReferenceFE(lagrangian,Float64,order); conformity=:L2)
P = TransientTrialFESpace(Q)

The multifield FE spaces for the prognostic and diagnostic variables are:

X_prog = MultiFieldFESpace([U,P,P]) # u, p, B
Y_prog = MultiFieldFESpace([V,Q,Q]) # u, p, B

X_diag = MultiFieldFESpace([H,U,P,P]) # q, F, Φ, b
Y_diag = MultiFieldFESpace([R,V,Q,Q]) # q, F, Φ, b

Initial conditions

The initial condition for velocity, fluid depth, buoyancy and Coriolis are defined as per the thermogeostrophic balance test case, in Ricardo et al. 2023.

The dimensional parameters are:

aₑ = 6.37e6 # m
gₑ = 9.8 # m/s^{-2}
ωₑ = 7.29e-5 # s^{-1}
Hₑ = 2.94e4/gₑ # m
uₑ = 2*π*aₑ/(12*24*3600) # m s^{-1}

Consider a length scale of $a_{e}$ and time scale of $1/\omega$ yields the non-dimensional parameters:

gravity = gₑ*(1/ωₑ)^2/aₑ
ω = ωₑ*(1/ωₑ)
H₀ = Hₑ/aₑ
u₀ = uₑ/aₑ*(1/ωₑ)

The initial conditions are defined as functions of the ambient space. They are composed with AmbientMapCellField in the weak form section below, where we also extract the contravariant components for the velocity:

function u0(xyz)
  ζ = 0.0
  θϕr   = xyz2θϕr(xyz)
  θ,ϕ,_ = θϕr
  u     = u₀*(cos(ϕ)*cos(ζ) + cos(θ)*sin(ϕ)*sin(ζ))
  v     = - u₀*sin(θ)*sin(ζ)
  function _spherical_to_cartesian_matrix(θϕr)
    θ,ϕ,_ = θϕr
    TensorValue(-sin(θ)       , cos(θ)       ,      0,
                -sin(ϕ)*cos(θ),-sin(ϕ)*sin(θ), cos(ϕ),
                cos(ϕ)*cos(θ), cos(ϕ)*sin(θ), sin(ϕ))
  end
  _spherical_to_cartesian_matrix(θϕr)⋅VectorValue(u,v,0)
end


function h0(xyz)
  ζ = 0.0
  θϕr   = xyz2θϕr(xyz)
  θ,ϕ,_ = θϕr
  h  = -cos(θ)*cos(ϕ)*sin(ζ) + sin(ϕ)*cos(ζ)
  H₀- (ω*u₀  + 0.5*u₀*u₀)*h*h/gravity
end

function b0(xyz)
  c = 0.05
  h = h0(xyz)
  gravity*( 1 + c *H₀^2 /h^2)
end

function B0(xyz)
  h = h0(xyz)
  b = b0(xyz)
  b*h
end

function f0(xyz)
  ζ = 0.0
  θϕr   = xyz2θϕr(xyz)
  θ,ϕ,_ = θϕr
  2.0*ω*( -cos(θ)*cos(ϕ)*sin(ζ) + sin(ϕ)*cos(ζ) )
end

Weak form

To define the weak form, we require the metric and measure, as well as the the matrix that represents the perp operator. We use an increased degree of quadrature to exactly approximate the geometrical map included in the weak form.

ambient_map_cf = AmbientMapCellField(Ω)
g = MetricCellField(Ω)
ginv = InvMetricCellField(Ω)
meas = MeasureCellField(Ω)
covariant_basis_cf = transpose∘∇(ambient_map_cf)

Then converted into a cellfield, where we extract the contravariant components for the velocity:

u_cf = meas*((pinvJ∘covariant_basis_cf)⋅(u0∘ambient_map_cf))
h_cf = h0∘ambient_map_cf
B_cf = B0∘ambient_map_cf
f_cf = f0∘ambient_map_cf
Aperp = [0 -1
        1 0]
Rperp = TensorValue(Aperp)
Rperp_cf = CellField(Rperp,Ω)
dΩ = Measure(Ω,4*order)
dΛ = Measure(Λ,4*order)

Diagnostic variables

The weak forms for the diagnostic variables are:

resq(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r)) = ∫( q*p*w*meas  )dΩ - ∫( f_cf*w*meas  )dΩ - ∫( (( (Rperp_cf⋅u)⋅ginv)⋅∇(w))*meas  )dΩ
resF(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r)) = ∫( (F⋅ (g⋅v))*(1/meas) )dΩ - ∫( p*(u⋅(g⋅v))*(1/meas)   )dΩ
resΦ(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r)) = ∫( Φ*ψ*meas  )dΩ - ∫( 0.5*B*ψ*meas  )dΩ - ∫( 0.5*( u ⋅(g⋅u) )*ψ*(1/meas)  )dΩ
resb(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r)) = ∫( b*p*r*meas  )dΩ - ∫( B*r*meas  )dΩ

res_y(t,((u,p,B),(q,F,Φ,b)),(w,v,ψ,r)) = (
    resq(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r))
  + resF(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r))
  + resΦ(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r))
  + resb(((u,p,B),(q,F,Φ,b)),(w,v,ψ,r))
)
jac_y(t,((u,p,B),(q,F,Φ,b)),(dq,dF,dΦ,db),(w,v,ψ,r)) = (
    ∫( dq*p*w*meas  )dΩ
  + ∫( (dF⋅ (g⋅v))*(1/meas) )dΩ
  + ∫( dΦ*ψ*meas  )dΩ
  + ∫( db*p*r*meas  )dΩ
)

Prognostic variables

The weak forms for the prognostic variables are as follows, where we use SUPG stabilisation for the vorticity the velocity equation, and dG upwinding for the buoyancy in the velocity and weighted buoyancy equation. Note the temperature $\vartheta = 0.5 \varphi$ is used directly

mass(t,(dut,dpt,dBt),(v,r,w)) = (
    ∫( (dut⋅ (g⋅v))*(1/meas) )dΩ
  + ∫( (dpt*r)*meas )dΩ
  + ∫( (dBt*w)*meas )dΩ
)

res_p(((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0)) = ∫( r*(∇⋅F) )dΩ

function my_mean( Bu_n::Gridap.Geometry.SkeletonPair)
  plus  = ( Bu_n.plus)
  minus = ( Bu_n.minus)
  0.5*( plus - minus  )
end

Upwinding function, can be extended to have soft upwinding or excessive damping

function upwinding_sign(Fn)
  c = 0.0
  if Fn < -1e-4
    c = -0.5
  elseif Fn > 1e-4
    c = 0.5
  end
  return c
end

res_u(((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0)) = (
          ∫( ( q*( (Rperp_cf⋅ F)⋅v))  )dΩ
        + ∫( -τ*( (q-q0)/dt )*( (Rperp_cf⋅ F)⋅v)   )dΩ
        + ∫( -τ*(u⋅∇(q))*( (Rperp_cf⋅ F)⋅v)*(1/meas)   )dΩ
        - ∫( Φ*(∇⋅v) )dΩ
        + ∫( 0.5*(b*(∇(0.5*p)⋅v) ) )dΩ
        + ∫( -0.5*(b*(0.5*p))*(∇⋅v)  )dΩ
        + ∫( -0.5*((0.5*p)*(∇(b)⋅v) ) )dΩ
    )

u_s1(((u,p,B),(q,F,Φ,b)),(v,r,w)) = (
    ∫( -0.5*my_mean((v*b)⋅n_Λ)*jump(0.5*p)   )dΛ
  + ∫( 0.5*my_mean((v*(0.5*p))⋅n_Λ)*jump(b)  )dΛ
)

u_s2(((u,p,B),(q,F,Φ,b)),(v,r,w)) = ∫( -0.5*( (upwinding_sign∘((F⋅ n_Λ).plus))*(v⋅n_Λ).plus )*jump(b)*jump(0.5*p)   )dΛ


res_B(((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0)) = (
    ∫( -0.5*(b*(∇(w)⋅F) ) )dΩ
  + ∫( 0.5*(b*w)*(∇⋅F)  )dΩ
  + ∫( 0.5*(w*(∇(b)⋅F) ) )dΩ
)

B_s1(((u,p,B),(q,F,Φ,b)),(v,r,w)) = (
    ∫( 0.5*my_mean((F*b)⋅n_Λ)*jump(w)   )dΛ
  + ∫( -0.5*my_mean((F*w)⋅n_Λ)*jump(b)   )dΛ
)

B_s2(((u,p,B),(q,F,Φ,b)),(v,r,w)) = ∫( 0.5*( (upwinding_sign∘((F⋅ n_Λ).plus))*(F⋅n_Λ).plus )*jump(b)*jump(w)   )dΛ


res_x(t,((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0)) = (
    res_u(((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0))
  + u_s1(((u,p,B),(q,F,Φ,b)),(v,r,w))
  + u_s2(((u,p,B),(q,F,Φ,b)),(v,r,w))
  + res_p(((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0))
  + res_B(((u,p,B),(q,F,Φ,b)),(v,r,w),(q0,F0,Φ0,b0))
  + B_s1(((u,p,B),(q,F,Φ,b)),(v,r,w))
  + B_s2(((u,p,B),(q,F,Φ,b)),(v,r,w))
)
jac_x(t,((u,p,B),(q,F,Φ,b)),(du,dp,dB),(v,r,w),(q0,F0,Φ0,b0)) = ∫( -τ*(du⋅∇(q))*( (Rperp_cf⋅ F)⋅v)*(1/meas)   )dΩ
jac_xt(t,((u,p,B),(q,F,Φ,b)),(dut,dpt,dBt),(v,r,w),(q0,F0,Φ0,b0)) = (
    ∫( (dut⋅ (g⋅v))*(1/meas) )dΩ
  + ∫( (dpt*r)*meas )dΩ
  + ∫( (dBt*w)*meas )dΩ
)

Transient problem

The transient problem requires a linear solver for both the prognostic and diagnostic problem. In this example, we using a Conjugate Gradient method for both.

ls_diag = CGSolver(JacobiLinearSolver();rtol=1e-12,verbose=i_am_main(ranks),name="diagnostic_solver")
ls_ode = CGSolver(JacobiLinearSolver();rtol=1e-12,verbose=i_am_main(ranks),name="ode_solver")

The transient parameters are:

t0 = 0.0
tF = 2*π
nsteps = 100
dt = tF/nsteps
τ = dt/2

The transient operator, DAE operator and transient solution, integrated using SSPRK3, is defined as follows:

opT = TransientSemilinearFEOperator(mass,res_x,(jac_x,jac_xt),X_prog,Y_prog)
opFE = FEOperator(res_y,jac_y,X_diag,Y_diag)
opDAE = DAEFEOperator(opT,opFE,ls_diag)
ode_solver = RungeKutta(ls_ode,ls_ode,dt,:EXRK_SSP_3_3)
xh0 = interpolate_everywhere([u_cf,h_cf,B_cf],X_prog)
solT  = solve(ode_solver,opDAE,t0,tF,xh0)

Post processing

The diagnostic solution can be extracted from the cache of the ode solution. To do so, we utilise the iterable functionality of $solT$ as follows:

i_am_main(ranks) && mkpath("thermal_shallow_water_sol")
PartitionedArrays.barrier(ranks)

it = iterate(solT)
while !isnothing(it)
  data, state = it
  t, xh = data
  odeopcache = state[2][5][2]
  yh = odeopcache.diagnostics

  uh,ph,Bh = xh
  qh,Fh,Φh,bh = yh

  i_am_main(ranks) && println(t)

  writevtk_with_cell_geomap(LatLonMapCellField(Ω),Ω,"thermal_shallow_water_sol/solT_$t.vtu",
      cellfields=["vel"=>covariant_basis_cf⋅(1/meas*uh),"p"=>ph,"vort"=>qh],
      append=false)

  it = iterate(solT, state)
end

This page was generated using Literate.jl.