Trying to solve Primitve Equations in Firedrake

[adominious]

Hello! I have decided to try and solve the following system of equations using firedrake:

• Continuity Equation (representing the Conservation of Mass):
  $$\frac{D\rho}{Dt} + \rho\nabla\cdot V = 0$$
  Solving for density, $\rho$; given velocity field $V$.

• Momentum Equations (A form of the Navier-Stokes (NS) equations):
  $$\frac{D V}{Dt} = -\frac{1}{\rho} \nabla p + F$$
  Solving for velocity field, $V$; given density, $\rho$, external force vector, $F$, pressure $p$.

• Thermal Energy Equation (Conservation of Energy):
  $$\frac{D \Theta}{Dt} = \frac{1}{c_p}\frac{\Theta}{T}Q$$
  Solving for potential temperature $\Theta$; given heat flux, $Q$, temperature $T$.

Here are the initial conditions I am using:

• Density ($\rho$):
  $$\rho(x, y, z, t=0) = \rho_0 + \rho_{\text{perturbation}} \sin\left(\frac{2\pi x}{L_x}\right) \cos\left(\frac{2\pi y}{L_y}\right).$$

• Velocity ($V = (u,v,w)$):
  $$u(x, y, z, t=0) = U_0 \cos\left(\frac{\pi z}{H}\right)$$
  $$v(x, y, z, t=0) = V_0 \sin\left(\frac{\pi x}{L_x}\right)$$
  $$w(x, y, z, t=0) = 0.$$

• Pressure ($p$):
  $$p(x, y, z, t=0) = p_0 - \rho_0 g z + p_{\text{perturbation}} \cos\left(\frac{2\pi x}{L_x}\right) \sin\left(\frac{2\pi y}{L_y}\right).$$

• Temperature ($T$):
  $$T(x, y, z, t=0) = T_0 + T_{\text{perturbation}} \sin\left(\frac{2\pi x}{L_x}\right) \cos\left(\frac{2\pi y}{L_y}\right).$$

Along with periodic boundary conditions.

I have mad an attempt to code this in firedrake, however with my lack of experience in this library, FEM, and continuum mechanics, I fear I might have made quite a mess!

I am having issues with trying to make use of some form of nonlinear solver due to the nonlinearity of the momentum equation, any help in resolving the error would be much appreciated (along with some general guidance with getting better at firedrake).
Thank you!

from firedrake import *

# Define the spatial domain covering the entire Earth's surface
mesh = UnitCubeMesh(50, 25, 12)  # Adjust the mesh resolution as needed

# Define function spaces
P1 = FiniteElement('CG', mesh.ufl_cell(), 1)
element = MixedElement([P1, P1, P1, P1, P1])  # For rho, u, v, w, T
V = FunctionSpace(mesh, element)

# Define spatial coordinates
x, y, z = SpatialCoordinate(mesh)

# Define trial and test functions
rho, u, v, w, T = TrialFunctions(V)
phi_rho, phi_u, phi_v, phi_w, phi_T = TestFunctions(V)

# Define parameters (you can adjust these as needed)
rho_0 = Constant(1.0)  # Initial density
u_0 = Constant(1.0)  # Initial x-velocity component
v_0 = Constant(1.0)  # Initial y-velocity component
w_0 = Constant(0.0) # Initial z-velocity component
p_0 = Constant(101325.0) #Initia; pressure (in pascals)
T_0 = Constant(273.0)  # Initial temperature
c_p = Constant(1000.0) # Specific heat capacity in air
Q = Constant(10.0) # Heat flux

rho_perturbation=Constant(0.1) # Perturbation value for density
T_perturbation=Constant(0.1) # Perturbation value for temperature
p_perturbation=Constant(0.1) # Perturbation value for pressure

Lx=Constant(1.0)
Ly=Constant(1.0)
H=Constant(1.0)

# Define the Coriolis parameter (for demonstration purposes)
f = Constant(1e-4)  # Example value, adjust as needed

# Define the Coriolis force in the x and y direction
F_x = -2*f*u
F_y = -2*f*v

# Define the initial conditions
rho_init = rho_0 + rho_perturbation * sin(2 * pi * x / Lx) * cos(2 * pi * y / Ly)
u_init = u_0 * cos(pi * z / H)
v_init = v_0 * sin(pi * x / Lx)
w_init = w_0
p_init = p_0 + p_perturbation * sin(2 * pi * x / Lx) * cos(2 * pi * y / Ly)  
T_init = T_0 + T_perturbation * sin(2 * pi * x / Lx) * cos(2 * pi * y / Ly)

# Define the equations
velocity=as_vector([u, v, w])
continuity_eq = (dot(rho, phi_rho) + rho * div(velocity) * phi_rho) * dx

momentum_eq_x = (dot(u, dot(grad(phi_u),velocity)) 
                + u * dot(u, dot(grad(u),velocity)) 
                + 1/rho * grad(p_init)[0] - F_x) * phi_u * dx

momentum_eq_y = (dot(v, dot(grad(phi_v),velocity)) 
                + v * dot(v, dot(grad(v),velocity)) 
                + 1/rho * grad(p_init)[1] - F_y) * phi_v * dx

momentum_eq_z = (dot(w, dot(grad(phi_w),velocity)) 
                + w * dot(w, dot(grad(w),velocity)) 
                + 1/rho * grad(p_init)[2] - F_y) * phi_w * dx

'''
# Define linearized momentum equations
momentum_eq_x = (dot(u, dot(grad(phi_u), velocity)) + 1/rho * phi_rho.dx(0) - F_x) * dx
momentum_eq_y = (dot(v, dot(grad(phi_v), velocity)) + 1/rho * phi_rho.dx(1) - F_y) * dx
'''

thermal_eq = (dot(T, phi_T) + 1/c_p * T * div(velocity) * Q) * dx

# Define boundary conditions (periodic boundary conditions)
bc_rho = [DirichletBC(V.sub(0), rho_init, (1, 2, 3, 4, 5, 6))]  # Adjust boundary IDs as needed
bc_u = [DirichletBC(V.sub(1), u_init, (1, 2, 3, 4, 5, 6))]  # Adjust boundary IDs as needed
bc_v = [DirichletBC(V.sub(2), v_init, (1, 2, 3, 4, 5, 6))]  # Adjust boundary IDs as needed
bc_w = [DirichletBC(V.sub(3), w_init, (1, 2, 3, 4, 5, 6))]  # Adjust boundary IDs as needed
bc_T = [DirichletBC(V.sub(4), T_init, (1, 2, 3, 4, 5, 6))]  # Adjust boundary IDs as needed

# No-slip boundary conditions for velocity components (assuming zero velocity)
bc_noslip_u = [DirichletBC(V.sub(1), 0.0, "on_boundary")]
bc_noslip_v = [DirichletBC(V.sub(2), 0.0, "on_boundary")]

bcs = bc_rho + bc_u + bc_v + bc_T + bc_noslip_u + bc_noslip_v

# Combine the equations and boundary conditions
equations = continuity_eq + momentum_eq_x + momentum_eq_y + momentum_eq_z + thermal_eq

# Define the variational problem
F = equations
solution = Function(V)

# Define the Jacobian matrix
J = derivative(F, solution)

'''
# Define the nonlinear solver parameters
params = {"snes_type": "newtonls",
          "snes_linesearch_type": "basic",
          "snes_max_it": 50,
          "snes_monitor": None,
          "ksp_type": "preonly",
          "pc_type": "lu"}

# Solve the equations using Newton's method
solve(F == 0, solution, bcs=bcs, J=J, solver_parameters=params)
'''

# Define the nonlinear variational problem
problem = NonlinearVariationalProblem(F, solution, bcs=bcs, J=J)

# Define the nonlinear solver
solver = NonlinearVariationalSolver(problem, solver_parameters={"snes_type": "newtonls",
                                                                "snes_linesearch_type": "basic",
                                                                "snes_max_it": 50,
                                                                "snes_monitor": None,
                                                                "ksp_type": "preonly",
                                                                "pc_type": "lu"})

# Solve the equations
solver.solve()

[colinjcotter]

There are a lot of things going on here. I think it is important to view Firedrake as the framework to implement the discretisation and solver that you have selected/designed. It’s a good idea to first check out what other people have done with FEM for primitive equations and why, and then pick one of those approaches. There are a whole load of related issues: stable element choices for the pressure waves, stabilisation of advection schemes, splitting methods to enable efficient solution etc. Unfortunately it is not as simple as taking the equations and trying the first thing.

I wrote a survey article on one approach:
https://www.cambridge.org/core/journals/acta-numerica/article/compatible-finite-element-methods-for-geophysical-fluid-dynamics/D2A87A43005D82261B9C7C4BB7CD000E
We have a library Gusto that is implementing some of these choices.

For FEM the other main approach that has been shown to work is spectral elements as implemented in the HOMME model.