Covariant formulation of the shallow water equations (#58)

* add additional aux vars * compute christoffel symbols * fix typo in christoffel calculation * covariant SWE weak form runs now * covariant SWE weak form runs now * no longer pass surface_integral into prolong2mortars * bump compat for Trixi from 0.9 to 0.9.9 * add test for spherical SWE * approach to automatically transform to desired coordinate system * added missing aux_vars to prim2cons * both spherical and cartesian works now * add EC formulation * add tests and improve elixirs/docs/cases * more detailed docs * fix typo in docs * improve docs more * add docstring for covariant SWE * sqrtG -> J * fix typo in quadrilaterals * add vorticity calculation for covariant form * improve docs and comments * fix mistake in comment * Rossby-Haurwitz is actually Case 6, not Case 5 * cleanup * reformulate flux differencing * run formatter * a few more comments * add missing h in total energy/entropy calculation * rename source terms * fix comment p -> polydeg * add comments and streamline calculation of contravariant metric * describe the Christoffel symbols * specify in comment that J = sqrtG for analysis in covariant form * fix typo in auxvarnames * comment about possibility to compute Christoffel symbols exactly * replace u_ll_spherical with u_ll_global and likewise for u_rr * run formatter * fix docs * fix another typo in docs * weak form works with central-central * separate split-covariant equations from standard covariant * update geostrophic test values to be for after one day * Apply suggestions from code review Co-authored-by: Andrés Rueda-Ramírez <[email protected]> * revise docstrings and comments --------- Co-authored-by: Andrés Rueda-Ramírez <[email protected]>
trixi-framework · Jan 29, 2025 · 2db20a7 · 2db20a7
1 parent 029c14e
commit 2db20a7
Show file tree

Hide file tree

Showing 21 changed files with 1,300 additions and 144 deletions.
diff --git a/examples/elixir_shallowwater_covariant_geostrophic_balance.jl b/examples/elixir_shallowwater_covariant_geostrophic_balance.jl
@@ -0,0 +1,75 @@
+###############################################################################
+# DGSEM for the shallow water equations in covariant form on the cubed sphere
+###############################################################################
+
+using OrdinaryDiffEq, Trixi, TrixiAtmo
+
+###############################################################################
+# Parameters
+
+initial_condition = initial_condition_geostrophic_balance
+polydeg = 3
+cells_per_dimension = 5
+n_saves = 10
+tspan = (0.0, 5.0 * SECONDS_PER_DAY)
+
+###############################################################################
+# Spatial discretization
+
+mesh = P4estMeshCubedSphere2D(cells_per_dimension, EARTH_RADIUS, polydeg = polydeg,
+                              initial_refinement_level = 0,
+                              element_local_mapping = true)
+
+equations = CovariantShallowWaterEquations2D(EARTH_GRAVITATIONAL_ACCELERATION,
+                                             EARTH_ROTATION_RATE,
+                                             global_coordinate_system = GlobalSphericalCoordinates())
+
+# Create DG solver with polynomial degree = polydeg
+solver = DGSEM(polydeg = polydeg, surface_flux = flux_lax_friedrichs,
+               volume_integral = VolumeIntegralWeakForm())
+
+# Transform the initial condition to the proper set of conservative variables
+initial_condition_transformed = transform_initial_condition(initial_condition, equations)
+
+# A semidiscretization collects data structures and functions for the spatial discretization
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    source_terms = source_terms_geometric_coriolis)
+
+###############################################################################
+# ODE solvers, callbacks etc.
+
+# Create ODE problem with time span from 0 to T
+ode = semidiscretize(semi, tspan)
+
+# At the beginning of the main loop, the SummaryCallback prints a summary of the simulation 
+# setup and resets the timers
+summary_callback = SummaryCallback()
+
+# The AnalysisCallback allows to analyse the solution in regular intervals and prints the 
+# results
+analysis_callback = AnalysisCallback(semi, interval = 200,
+                                     save_analysis = true,
+                                     extra_analysis_errors = (:conservation_error,))
+
+# The SaveSolutionCallback allows to save the solution to a file in regular intervals
+save_solution = SaveSolutionCallback(dt = (tspan[2] - tspan[1]) / n_saves,
+                                     solution_variables = cons2cons)
+
+# The StepsizeCallback handles the re-calculation of the maximum Δt after each time step
+stepsize_callback = StepsizeCallback(cfl = 0.4)
+
+# Create a CallbackSet to collect all callbacks such that they can be passed to the ODE 
+# solver
+callbacks = CallbackSet(summary_callback, analysis_callback, save_solution,
+                        stepsize_callback)
+
+###############################################################################
+# run the simulation
+
+# OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed 
+# callbacks
+sol = solve(ode, CarpenterKennedy2N54(williamson_condition = false),
+            dt = 100.0, save_everystep = false, callback = callbacks);
+
+# Print the timer summary
+summary_callback()
diff --git a/examples/elixir_shallowwater_covariant_rossby_haurwitz_EC.jl b/examples/elixir_shallowwater_covariant_rossby_haurwitz_EC.jl
@@ -0,0 +1,81 @@
+###############################################################################
+# Entropy-conservative DGSEM for the shallow water equations in covariant form 
+# on the cubed sphere
+###############################################################################
+
+using OrdinaryDiffEq, Trixi, TrixiAtmo
+
+###############################################################################
+# Parameters
+
+initial_condition = initial_condition_rossby_haurwitz
+polydeg = 3
+cells_per_dimension = 5
+n_saves = 10
+tspan = (0.0, 7.0 * SECONDS_PER_DAY)
+
+###############################################################################
+# Spatial discretization
+
+mesh = P4estMeshCubedSphere2D(cells_per_dimension, EARTH_RADIUS, polydeg = polydeg,
+                              initial_refinement_level = 0,
+                              element_local_mapping = true)
+
+equations = SplitCovariantShallowWaterEquations2D(EARTH_GRAVITATIONAL_ACCELERATION,
+                                                  EARTH_ROTATION_RATE,
+                                                  global_coordinate_system = GlobalCartesianCoordinates())
+
+# Use entropy-conservative two-point fluxes for volume and surface terms
+volume_flux = (flux_ec, flux_nonconservative_ec)
+surface_flux = (flux_ec, flux_nonconservative_ec)
+
+# Create DG solver with polynomial degree = polydeg
+solver = DGSEM(polydeg = polydeg, surface_flux = surface_flux,
+               volume_integral = VolumeIntegralFluxDifferencing(volume_flux))
+
+# Transform the initial condition to the proper set of conservative variables
+initial_condition_transformed = transform_initial_condition(initial_condition, equations)
+
+# A semidiscretization collects data structures and functions for the spatial discretization
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    source_terms = source_terms_geometric_coriolis)
+
+###############################################################################
+# ODE solvers, callbacks etc.
+
+# Create ODE problem with time span from 0 to T
+ode = semidiscretize(semi, tspan)
+
+# At the beginning of the main loop, the SummaryCallback prints a summary of the simulation 
+# setup and resets the timers
+summary_callback = SummaryCallback()
+
+# The AnalysisCallback allows to analyse the solution in regular intervals and prints the
+# results. Note that entropy should be conserved at the semi-discrete level.
+analysis_callback = AnalysisCallback(semi, interval = 200,
+                                     save_analysis = true,
+                                     extra_analysis_errors = (:conservation_error,),
+                                     extra_analysis_integrals = (entropy,))
+
+# The SaveSolutionCallback allows to save the solution to a file in regular intervals
+save_solution = SaveSolutionCallback(dt = (tspan[2] - tspan[1]) / n_saves,
+                                     solution_variables = cons2prim_and_vorticity)
+
+# The StepsizeCallback handles the re-calculation of the maximum Δt after each time step
+stepsize_callback = StepsizeCallback(cfl = 0.4)
+
+# Create a CallbackSet to collect all callbacks such that they can be passed to the ODE 
+# solver
+callbacks = CallbackSet(summary_callback, analysis_callback, save_solution,
+                        stepsize_callback)
+
+###############################################################################
+# run the simulation
+
+# OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed 
+# callbacks
+sol = solve(ode, CarpenterKennedy2N54(williamson_condition = false),
+            dt = 100.0, save_everystep = false, callback = callbacks);
+
+# Print the timer summary
+summary_callback()
diff --git a/examples/elixir_shallowwater_cubed_sphere_shell_advection.jl b/examples/elixir_shallowwater_cubed_sphere_shell_advection.jl
@@ -46,6 +46,7 @@ mesh = P4estMeshCubedSphere2D(cells_per_dimension[1], EARTH_RADIUS, polydeg = 3,
                               #initial_refinement_level = 0, # Comment to use cells_per_dimension in the convergence test
                               element_local_mapping = false)
 
+# Transform the initial condition to the proper set of conservative variables
 initial_condition_transformed = transform_initial_condition(initial_condition, equations)
 
 # A semidiscretization collects data structures and functions for the spatial discretization

diff --git a/examples/elixir_spherical_advection_covariant_cubed_sphere.jl b/examples/elixir_spherical_advection_covariant_cubed_sphere.jl
@@ -25,6 +25,7 @@ mesh = P4estMeshCubedSphere2D(cells_per_dimension[1], EARTH_RADIUS,
                               polydeg = Trixi.polydeg(solver),
                               element_local_mapping = true)
 
+# Transform the initial condition to the proper set of conservative variables
 initial_condition_transformed = transform_initial_condition(initial_condition, equations)
 
 # A semidiscretization collects data structures and functions for the spatial discretization

diff --git a/examples/elixir_spherical_advection_covariant_quad_icosahedron.jl b/examples/elixir_spherical_advection_covariant_quad_icosahedron.jl
@@ -1,5 +1,5 @@
 ###############################################################################
-# DGSEM for the linear advection equation on the cubed sphere
+# DGSEM for the linear advection equation on a quadrilateral icosahedral grid
 ###############################################################################
 # To run a convergence test, use
 # convergence_test("../examples/elixir_spherical_advection_covariant_quad_icosahedron.jl", 4, cells_per_dimension = (1,1))
@@ -18,12 +18,13 @@ equations = CovariantLinearAdvectionEquation2D(global_coordinate_system = Global
 solver = DGSEM(polydeg = 3, surface_flux = flux_lax_friedrichs,
                volume_integral = VolumeIntegralWeakForm())
 
-# Create a 2D cubed sphere mesh the size of the Earth. For the covariant form to work 
-# properly, we currently need polydeg to equal that of the solver, and
+# Create a 2D quadrilateral icosahedral mesh the size of the Earth. For the covariant form 
+# to work properly, we currently need polydeg to equal that of the solver, and
 # initial_refinement_level = 0 (default)
 mesh = P4estMeshQuadIcosahedron2D(cells_per_dimension[1], EARTH_RADIUS,
                                   polydeg = Trixi.polydeg(solver))
 
+# Transform the initial condition to the proper set of conservative variables
 initial_condition_transformed = transform_initial_condition(initial_condition, equations)
 
 # A semidiscretization collects data structures and functions for the spatial discretization
@@ -35,11 +36,12 @@ semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transform
 # Create ODE problem with time span from 0 to T
 ode = semidiscretize(semi, (0.0, 12 * SECONDS_PER_DAY))
 
-# At the beginning of the main loop, the SummaryCallback prints a summary of the simulation setup
-# and resets the timers
+# At the beginning of the main loop, the SummaryCallback prints a summary of the simulation 
+# setup and resets the timers
 summary_callback = SummaryCallback()
 
-# The AnalysisCallback allows to analyse the solution in regular intervals and prints the results
+# The AnalysisCallback allows to analyse the solution in regular intervals and prints the 
+# results
 analysis_callback = AnalysisCallback(semi, interval = 10,
                                      save_analysis = true,
                                      extra_analysis_errors = (:conservation_error,))
@@ -51,14 +53,16 @@ save_solution = SaveSolutionCallback(interval = 10,
 # The StepsizeCallback handles the re-calculation of the maximum Δt after each time step
 stepsize_callback = StepsizeCallback(cfl = 0.7)
 
-# Create a CallbackSet to collect all callbacks such that they can be passed to the ODE solver
+# Create a CallbackSet to collect all callbacks such that they can be passed to the ODE 
+# solver
 callbacks = CallbackSet(summary_callback, analysis_callback, save_solution,
                         stepsize_callback)
 
 ###############################################################################
 # run the simulation
 
-# OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed callbacks
+# OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed 
+# callbacks
 sol = solve(ode, CarpenterKennedy2N54(williamson_condition = false),
             dt = 1.0, save_everystep = false, callback = callbacks);
 

diff --git a/src/TrixiAtmo.jl b/src/TrixiAtmo.jl
@@ -31,22 +31,33 @@ include("semidiscretization/semidiscretization_hyperbolic_2d_manifold_in_3d.jl")
 include("callbacks_step/callbacks_step.jl")
 
 export CompressibleMoistEulerEquations2D, ShallowWaterEquations3D,
-       CovariantLinearAdvectionEquation2D
+       CovariantLinearAdvectionEquation2D, CovariantShallowWaterEquations2D,
+       SplitCovariantShallowWaterEquations2D
+
 export GlobalCartesianCoordinates, GlobalSphericalCoordinates
 
 export flux_chandrashekar, FluxLMARS
 
+export flux_nonconservative_zeros, flux_nonconservative_ec,
+       source_terms_geometric_coriolis
+
 export velocity, waterheight, pressure, energy_total, energy_kinetic, energy_internal,
        lake_at_rest_error, source_terms_lagrange_multiplier,
        clean_solution_lagrange_multiplier!
+
 export cons2prim_and_vorticity
+
 export P4estMeshCubedSphere2D, P4estMeshQuadIcosahedron2D, MetricTermsCrossProduct,
        MetricTermsInvariantCurl
+
 export EARTH_RADIUS, EARTH_GRAVITATIONAL_ACCELERATION,
        EARTH_ROTATION_RATE, SECONDS_PER_DAY
-export global2contravariant, contravariant2global, spherical2cartesian,
+
+export global2contravariant, contravariant2global, spherical2cartesian, cartesian2spherical,
        transform_initial_condition
-export initial_condition_gaussian, initial_condition_gaussian_cartesian
+
+export initial_condition_gaussian, initial_condition_geostrophic_balance,
+       initial_condition_rossby_haurwitz
 
 export examples_dir
 end # module TrixiAtmo
diff --git a/src/callbacks_step/analysis_covariant.jl b/src/callbacks_step/analysis_covariant.jl
@@ -1,9 +1,28 @@
 @muladd begin
 #! format: noindent
 
+# When the equations are of type AbstractCovariantEquations, the functions which we would 
+# like to integrate depend on the solution as well as the auxiliary variables
+function Trixi.integrate(func::Func, u,
+                         mesh::Union{TreeMesh{2}, StructuredMesh{2},
+                                     StructuredMeshView{2},
+                                     UnstructuredMesh2D, P4estMesh{2}, T8codeMesh{2}},
+                         equations::AbstractCovariantEquations{2}, dg::DG,
+                         cache; normalize = true) where {Func}
+    (; aux_node_vars) = cache.auxiliary_variables
+
+    Trixi.integrate_via_indices(u, mesh, equations, dg, cache;
+                                normalize = normalize) do u, i, j, element, equations,
+                                                          dg
+        u_local = Trixi.get_node_vars(u, equations, dg, i, j, element)
+        aux_local = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+        return func(u_local, aux_local, equations)
+    end
+end
+
 # For the covariant form, we want to integrate using the exact area element 
-# √G = (det(AᵀA))^(1/2), which is stored in cache.auxiliary_variables, not the approximate 
-# area element used in the Cartesian formulation, which stored in cache.elements.
+# J = √G = (det(AᵀA))^(1/2), which is stored in cache.auxiliary_variables, not the approximate 
+# area element used in the Cartesian formulation, which stored in cache.elements
 function Trixi.integrate_via_indices(func::Func, u,
                                      mesh::Union{StructuredMesh{2},
                                                  StructuredMeshView{2},
@@ -23,10 +42,10 @@ function Trixi.integrate_via_indices(func::Func, u,
     for element in eachelement(dg, cache)
         for j in eachnode(dg), i in eachnode(dg)
             aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
-            sqrtG = area_element(aux_node, equations)
-            integral += weights[i] * weights[j] * sqrtG *
+            J = area_element(aux_node, equations)
+            integral += weights[i] * weights[j] * J *
                         func(u, i, j, element, equations, dg, args...)
-            total_volume += weights[i] * weights[j] * sqrtG
+            total_volume += weights[i] * weights[j] * J
         end
     end
 
@@ -65,6 +84,7 @@ function Trixi.calc_error_norms(func, u, t, analyzer, mesh::P4estMesh{2},
     (; weights) = dg.basis
     (; node_coordinates) = cache.elements
     (; aux_node_vars) = cache.auxiliary_variables
+
     # Set up data structures
     l2_error = zero(func(Trixi.get_node_vars(u, equations, dg, 1, 1, 1), equations))
     linf_error = copy(l2_error)
@@ -88,14 +108,14 @@ function Trixi.calc_error_norms(func, u, t, analyzer, mesh::P4estMesh{2},
             diff = func(u_exact, equations) - func(u_numerical, equations)
 
             # For the L2 error, integrate with respect to area element stored in aux vars 
-            sqrtG = area_element(aux_node, equations)
-            l2_error += diff .^ 2 * (weights[i] * weights[j] * sqrtG)
+            J = area_element(aux_node, equations)
+            l2_error += diff .^ 2 * (weights[i] * weights[j] * J)
 
             # Compute Linf error as usual
             linf_error = @. max(linf_error, abs(diff))
 
             # Increment total volume according to the volume element stored in aux vars
-            total_volume += weights[i] * weights[j] * sqrtG
+            total_volume += weights[i] * weights[j] * J
         end
     end
 
@@ -104,4 +124,4 @@ function Trixi.calc_error_norms(func, u, t, analyzer, mesh::P4estMesh{2},
 
     return l2_error, linf_error
 end
-end # muladd
+end # @muladd
diff --git a/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
@@ -3,7 +3,8 @@
 
 # Convert to another set of variables using a solution_variables function
 function convert_variables(u, solution_variables, mesh::P4estMesh{2},
-                           equations::AbstractEquations{3}, dg, cache)
+                           equations::AbstractEquations{3},
+                           dg, cache)
     (; contravariant_vectors) = cache.elements
     # Extract the number of solution variables to be output 
     # (may be different than the number of conservative variables) 
@@ -129,7 +130,7 @@ end
 
     u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
 
-    # Return th solution variables
+    # Return the solution variables
     return SVector(cons2prim(u_node, equations)..., relative_vorticity)
 end
 
@@ -177,7 +178,9 @@ end
 end
 
 # Variable names for cons2prim_and_vorticity
-Trixi.varnames(::typeof(cons2prim_and_vorticity), equations::ShallowWaterEquations3D) = (varnames(cons2prim,
-                                                                                                  equations)...,
-                                                                                         "vorticity")
+function Trixi.varnames(::typeof(cons2prim_and_vorticity),
+                        equations::Union{ShallowWaterEquations3D,
+                                         AbstractCovariantEquations{2}})
+    return (varnames(cons2prim, equations)..., "vorticity")
 end
+end # @muladd