Covariant formulation of the shallow water equations

examples/elixir_shallowwater_covariant_geostrophic_balance.jl

@@ -0,0 +1,72 @@
+###############################################################################
+# DGSEM for the shallow water equations 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, 1.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 = p
+solver = DGSEM(polydeg = polydeg,
+               surface_flux = (flux_lax_friedrichs, flux_nonconservative_weak_form),
+               volume_integral = VolumeIntegralWeakForm())
+
+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_weak_form)
+
+###############################################################################
+# 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 = 50,
+                                     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()

examples/elixir_shallowwater_covariant_rossby_haurwitz_EC.jl

@@ -0,0 +1,75 @@
+###############################################################################
+# DGSEM for the shallow water equations 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, 1.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 = GlobalCartesianCoordinates())
+
+# Create DG solver with polynomial degree = p
+volume_flux = (flux_split_covariant, flux_nonconservative_split_covariant)
+surface_flux = (flux_split_covariant, flux_nonconservative_split_covariant)
+
+solver = DGSEM(polydeg = polydeg, surface_flux = surface_flux,
+               volume_integral = VolumeIntegralFluxDifferencing(volume_flux))
+
+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_split_covariant)
+
+###############################################################################
+# 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 = 50,
+                                     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()

src/TrixiAtmo.jl

@@ -31,22 +31,31 @@ include("semidiscretization/semidiscretization_hyperbolic_2d_manifold_in_3d.jl")
 include("callbacks_step/callbacks_step.jl")
 
 export CompressibleMoistEulerEquations2D, ShallowWaterEquations3D,
-       CovariantLinearAdvectionEquation2D
+       CovariantLinearAdvectionEquation2D, CovariantShallowWaterEquations2D
+
 export GlobalCartesianCoordinates, GlobalSphericalCoordinates
 
-export flux_chandrashekar, flux_LMARS
+export flux_chandrashekar, flux_LMARS, flux_split_covariant, flux_nonconservative_weak_form,
+       flux_nonconservative_split_covariant, flux_split_covariant_lax_friedrichs,
+       source_terms_split_covariant
 
 export velocity, waterheight, pressure, energy_total, energy_kinetic, energy_internal,
-       lake_at_rest_error, source_terms_lagrange_multiplier,
+       lake_at_rest_error, source_terms_lagrange_multiplier, source_terms_weak_form,
        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

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 = (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
 

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

src/callbacks_step/save_solution_covariant.jl

@@ -1,32 +1,71 @@
 @muladd begin
 #! format: noindent
 
-# Convert to another set of variables using a solution_variables function that depends on 
-# auxiliary variables
+# Convert to another set of variables using a solution_variables function
 function convert_variables(u, solution_variables, mesh::P4estMesh{2},
                            equations::AbstractCovariantEquations{2}, dg, cache)
     (; aux_node_vars) = cache.auxiliary_variables
-
     # Extract the number of solution variables to be output 
     # (may be different than the number of conservative variables) 
-    n_vars = length(solution_variables(Trixi.get_node_vars(u, equations, dg, 1, 1, 1),
-                                       get_node_aux_vars(aux_node_vars, equations, dg,
-                                                         1, 1,
-                                                         1), equations))
+    n_vars = length(Trixi.varnames(solution_variables, equations))
+
     # Allocate storage for output
     data = Array{eltype(u)}(undef, n_vars, nnodes(dg), nnodes(dg), nelements(dg, cache))
 
     # Loop over all nodes and convert variables, passing in auxiliary variables
     Trixi.@threaded for element in eachelement(dg, cache)
         for j in eachnode(dg), i in eachnode(dg)
-            u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
-            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
-            u_node_converted = solution_variables(u_node, aux_node, equations)
+            if applicable(solution_variables, u, equations, dg, cache, i, j, element)
+                # The solution_variables function depends on the solution on other nodes
+                data_node = solution_variables(u, equations, dg, cache, i, j, element)
+            else
+                u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
+                aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j,
+                                             element)
+                data_node = solution_variables(u_node, aux_node, equations)
+            end
+
             for v in 1:n_vars
-                data[v, i, j, element] = u_node_converted[v]
+                data[v, i, j, element] = data_node[v]
             end
         end
     end
     return data
 end
-end # muladd
+
+# Calculate the primitive variables and the relative vorticity at a given node
+@inline function cons2prim_and_vorticity(u, equations::AbstractCovariantEquations{2},
+                                         dg::DGSEM, cache, i, j, element)
+    (; aux_node_vars) = cache.auxiliary_variables
+    u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
+    aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+    relative_vorticity = calc_vorticity_node(u, equations, dg, cache, i, j, element)
+    return SVector(cons2prim(u_node, aux_node, equations)..., relative_vorticity)
+end
+
+# Calculate relative vorticity ζ = ∂₁v₂ - ∂₂v₁ for equations in covariant form
+@inline function calc_vorticity_node(u, equations::AbstractCovariantEquations{2},
+                                     dg::DGSEM, cache, i, j, element)
+    (; derivative_matrix) = dg.basis
+    (; aux_node_vars) = cache.auxiliary_variables
+
+    dv2dxi1 = dv1dxi2 = zero(eltype(u))
+    for ii in eachnode(dg)
+        u_node_ii = Trixi.get_node_vars(u, equations, dg, ii, j, element)
+        aux_node_ii = get_node_aux_vars(aux_node_vars, equations, dg, ii, j, element)
+        vcov = metric_covariant(aux_node_ii, equations) * velocity(u_node_ii, equations)
+        dv2dxi1 = dv2dxi1 + derivative_matrix[i, ii] * vcov[2]
+    end
+
+    for jj in eachnode(dg)
+        u_node_jj = Trixi.get_node_vars(u, equations, dg, i, jj, element)
+        aux_node_jj = get_node_aux_vars(aux_node_vars, equations, dg, i, jj, element)
+        vcov = metric_covariant(aux_node_jj, equations) * velocity(u_node_jj, equations)
+        dv1dxi2 = dv1dxi2 + derivative_matrix[j, jj] * vcov[1]
+    end
+
+    # compute the relative vorticity
+    aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+    return (dv2dxi1 - dv1dxi2) / area_element(aux_node, equations)
+end
+end # @muladd