Merge branch 'main' into TBendall/update_ref_profile

gusto/solvers/linear_solvers.py

@@ -9,8 +9,8 @@
     split, LinearVariationalProblem, Constant, LinearVariationalSolver,
     TestFunctions, TrialFunctions, TestFunction, TrialFunction, lhs,
     rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b,
-    ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis,
-    BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC
+    ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, cross,
+    BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector
 )
 from firedrake.fml import Term, drop
 from firedrake.petsc import flatten_parameters
@@ -319,12 +319,15 @@ def L_tr(f):
             + dl('+')*jump(u, n=n)*(dS_vp + dS_hp)
             + dl*dot(u, n)*(ds_tbp + ds_vp)
         )
-
         # TODO: can we get this term using FML?
         # contribution of the sponge term
         if hasattr(self.equations, "mu"):
             eqn += dt*self.equations.mu*inner(w, k)*inner(u, k)*dx
 
+        if equations.parameters.Omega is not None:
+            Omega = as_vector([0, 0, equations.parameters.Omega])
+            eqn += inner(w, cross(2*Omega, u))*dx
+
         aeqn = lhs(eqn)
         Leqn = rhs(eqn)
 
@@ -512,6 +515,11 @@ def V(u):
 
         if hasattr(self.equations, "mu"):
             eqn += dt*self.equations.mu*inner(w, k)*inner(u, k)*dx
+
+        if equation.parameters.Omega is not None:
+            Omega = as_vector((0, 0, equation.parameter.Omega))
+            eqn += inner(w, cross(2*Omega, u))*dx
+
         aeqn = lhs(eqn)
         Leqn = rhs(eqn)
 
@@ -661,6 +669,10 @@ def _setup_solver(self):
             + beta_d * phi * Dbar * div(u) * dx
         )
 
+        if 'coriolis' in equation.prescribed_fields._field_names:
+            f = equation.prescribed_fields('coriolis')
+            eqn += beta_u_ * f * inner(w, equation.domain.perp(u)) * dx
+
         aeqn = lhs(eqn)
         Leqn = rhs(eqn)
 
@@ -873,6 +885,10 @@ def _setup_solver(self):
             + beta_d * phi * Dbar * div(u) * dx
         )
 
+        if 'coriolis' in equation.prescribed_fields._field_names:
+            f = equation.prescribed_fields('coriolis')
+            eqn += beta_u_ * f * inner(w, equation.domain.perp(u)) * dx
+
         aeqn = lhs(eqn)
         Leqn = rhs(eqn)