Moved the Lagrange multiplier functions to shallow_water_3d.jl and im…

…proved the description

examples/elixir_shallowwater_cubed_sphere_shell_EC_correction.jl

@@ -57,23 +57,21 @@ end
 
 # Source term function to apply Lagrange multiplier to the semi-discretization
 # The vector contravariant_normal_vector contains the normal contravariant vectors scaled with the inverse Jacobian
-function source_terms_lagrange_multiplier(u, du, x, t,
-                                          equations::ShallowWaterEquations3D,
-                                          contravariant_normal_vector)
+function source_terms_entropy_correction(u, du, x, t,
+                                         equations::ShallowWaterEquations3D,
+                                         contravariant_normal_vector)
 
     # Entropy correction term
     s2 = -contravariant_normal_vector[1] * equations.gravity * u[1]^2
     s3 = -contravariant_normal_vector[2] * equations.gravity * u[1]^2
     s4 = -contravariant_normal_vector[3] * equations.gravity * u[1]^2
 
-    # Lagrange multipliers in the x direction
-    x_dot_div_f = (x[1] * (du[2] + s2) + x[2] * (du[3] + s3) + x[3] * (du[4] + s4)) /
-                  sum(x .^ 2)
-    s2 += -x[1] * x_dot_div_f
-    s3 += -x[2] * x_dot_div_f
-    s4 += -x[3] * x_dot_div_f
+    new_du = SVector(du[1], du[2] + s2, du[3] + s3, du[4] + s4, du[5])
 
-    return SVector(0.0, s2, s3, s4, 0.0)
+    # Apply Lagrange multipliers using the exact normal vector x
+    s = source_terms_lagrange_multiplier(u, new_du, x, t, equations, x)
+
+    return SVector(s[1], s[2] + s2, s[3] + s3, s[4] + s4, s[5])
 end
 
 # Custom RHS that applies a source term that depends on du (used to convert apply Lagrange multiplier)
@@ -103,9 +101,9 @@ function rhs_semi_custom!(du_ode, u_ode, semi, t)
                                                                              element) *
                                               inverse_jacobian[i, j, element]
 
-                source = source_terms_lagrange_multiplier(u_local, du_local,
-                                                          x_local, t, equations,
-                                                          contravariant_normal_vector)
+                source = source_terms_entropy_correction(u_local, du_local,
+                                                         x_local, t, equations,
+                                                         contravariant_normal_vector)
 
                 Trixi.add_to_node_vars!(du, source, equations, solver, i, j, element)
             end

examples/elixir_shallowwater_cubed_sphere_shell_EC_projection.jl

@@ -56,36 +56,6 @@ function initial_condition_advection_sphere(x, t, equations::ShallowWaterEquatio
     return prim2cons(SVector(rho, v1, v2, v3, 0), equations)
 end
 
-# Source term function to apply Lagrange multiplier to the semi-discretization
-# The vector contravariant_normal_vector contains the normal contravariant vectors (no scaling is needed)
-function source_terms_lagrange_multiplier(u, du, x, t,
-                                          equations::ShallowWaterEquations3D,
-                                          contravariant_normal_vector)
-    x_dot_div_f = (contravariant_normal_vector[1] * du[2] +
-                   contravariant_normal_vector[2] * du[3] +
-                   contravariant_normal_vector[3] * du[4]) /
-                  sum(contravariant_normal_vector .^ 2)
-
-    s2 = -contravariant_normal_vector[1] * x_dot_div_f
-    s3 = -contravariant_normal_vector[2] * x_dot_div_f
-    s4 = -contravariant_normal_vector[3] * x_dot_div_f
-
-    return SVector(0.0, s2, s3, s4, 0.0)
-end
-
-# Function to apply Lagrange multiplier discretely to the solution
-# The vector contravariant_normal_vector contains the normal contravariant vectors (no scaling is needed)
-function clean_solution!(u, equations::ShallowWaterEquations3D, contravariant_normal_vector)
-    x_dot_div_f = (contravariant_normal_vector[1] * u[2] +
-                   contravariant_normal_vector[2] * u[3] +
-                   contravariant_normal_vector[3] * u[4]) /
-                  sum(contravariant_normal_vector .^ 2)
-
-    u[2] -= contravariant_normal_vector[1] * x_dot_div_f
-    u[3] -= contravariant_normal_vector[2] * x_dot_div_f
-    u[4] -= contravariant_normal_vector[3] * x_dot_div_f
-end
-
 # Custom RHS that applies a source term that depends on du (used to convert apply Lagrange multiplier)
 function rhs_semi_custom!(du_ode, u_ode, semi, t)
     # Compute standard Trixi RHS
@@ -145,7 +115,8 @@ for element in eachelement(solver, semi.cache)
         contravariant_normal_vector = Trixi.get_contravariant_vector(3,
                                                                      semi.cache.elements.contravariant_vectors,
                                                                      i, j, element)
-        clean_solution!(u0[:, i, j, element], equations, contravariant_normal_vector)
+        clean_solution_lagrange_multiplier!(u0[:, i, j, element], equations,
+                                            contravariant_normal_vector)
     end
 end
 

src/TrixiAtmo.jl

@@ -31,7 +31,8 @@ export CompressibleMoistEulerEquations2D, ShallowWaterEquations3D
 export flux_chandrashekar, flux_LMARS
 
 export velocity, waterheight, pressure, energy_total, energy_kinetic, energy_internal,
-       lake_at_rest_error
+       lake_at_rest_error, source_terms_lagrange_multiplier,
+       clean_solution_lagrange_multiplier!
 
 export examples_dir
 

src/equations/shallow_water_3d.jl

@@ -25,6 +25,11 @@ The equations are given by
 The unknown quantities of the SWE are the water height ``h`` and the velocities ``\mathbf{v} = (v_1, v_2, v_3)^T``.
 The gravitational constant is denoted by `g`.
 
+The 3D Shallow Water Equations (SWE) extend the 2D SWE to model shallow water flows on 2D manifolds embedded within 3D space. 
+To confine the flow to the 2D manifold, a source term incorporating a Lagrange multiplier is applied. 
+This term effectively removes momentum components that are normal to the manifold, ensuring the flow remains 
+constrained within the 2D surface.
+
 The additional quantity ``H_0`` is also available to store a reference value for the total water height that
 is useful to set initial conditions or test the "lake-at-rest" well-balancedness.
 
@@ -39,10 +44,11 @@ This affects the implementation and use of these equations in various ways:
 * [`AnalysisCallback`](https://trixi-framework.github.io/Trixi.jl/stable/reference-trixi/#Trixi.AnalysisCallback) analyzes this variable.
 * Trixi.jl's visualization tools will visualize the bottom topography by default.
 
-References for the SWE are many but a good introduction is available in Chapter 13 of the book:
-- Randall J. LeVeque (2002)
-  Finite Volume Methods for Hyperbolic Problems
-  [DOI: 10.1017/CBO9780511791253](https://doi.org/10.1017/CBO9780511791253)
+References:
+- J. Coté (1988). "A Lagrange multiplier approach for the metric terms of semi-Lagrangian models on the sphere". 
+  Quarterly Journal of the Royal Meteorological Society 114, 1347-1352. https://doi.org/10.1002/qj.49711448310
+- Giraldo (2001). "A spectral element shallow water model on spherical geodesic grids". 
+  https://doi.org/10.1002/1097-0363(20010430)35:8%3C869::AID-FLD116%3E3.0.CO;2-S
 """
 struct ShallowWaterEquations3D{RealT <: Real} <:
        Trixi.AbstractShallowWaterEquations{3, 5}
@@ -201,6 +207,53 @@ Details are available in Eq. (4.1) in the paper:
     return SVector(f1, f2, f3, f4, zero(eltype(u_ll)))
 end
 
+"""
+         source_terms_lagrange_multiplier(u, du, x, t,
+                                          equations::ShallowWaterEquations3D,
+                                          normal_direction)
+
+Source term function to apply a Lagrange multiplier to the semi-discretization
+in order to constrain the momentum to a 2D manifold.
+
+The vector normal_direction is perpendicular to the 2D manifold. By default, 
+this is the normal contravariant basis vector.
+"""
+function source_terms_lagrange_multiplier(u, du, x, t,
+                                          equations::ShallowWaterEquations3D,
+                                          normal_direction)
+    x_dot_div_f = (normal_direction[1] * du[2] +
+                   normal_direction[2] * du[3] +
+                   normal_direction[3] * du[4]) /
+                  sum(normal_direction .^ 2)
+
+    s2 = -normal_direction[1] * x_dot_div_f
+    s3 = -normal_direction[2] * x_dot_div_f
+    s4 = -normal_direction[3] * x_dot_div_f
+
+    return SVector(0, s2, s3, s4, 0)
+end
+
+"""
+         clean_solution_lagrange_multiplier!(u, equations::ShallowWaterEquations3D, normal_direction)
+
+Function to apply Lagrange multiplier discretely to the solution in order to constrain 
+the momentum to a 2D manifold.
+
+The vector normal_direction is perpendicular to the 2D manifold. By default, 
+this is the normal contravariant basis vector.
+"""
+function clean_solution_lagrange_multiplier!(u, equations::ShallowWaterEquations3D,
+                                             normal_direction)
+    x_dot_div_f = (normal_direction[1] * u[2] +
+                   normal_direction[2] * u[3] +
+                   normal_direction[3] * u[4]) /
+                  sum(normal_direction .^ 2)
+
+    u[2] -= normal_direction[1] * x_dot_div_f
+    u[3] -= normal_direction[2] * x_dot_div_f
+    u[4] -= normal_direction[3] * x_dot_div_f
+end
+
 @inline function Trixi.max_abs_speed_naive(u_ll, u_rr, normal_direction::AbstractVector,
                                            equations::ShallowWaterEquations3D)
     # Extract and compute the velocities in the normal direction