streamline initial condition definition and improve docs

examples/elixir_shallowwater_cubed_sphere_shell_EC_correction.jl

@@ -45,12 +45,12 @@ function initial_condition_advection_sphere(x, t, equations::ShallowWaterEquatio
     v0 = 1.0 # Velocity at the "equator"
     alpha = 0.0 #pi / 4
     v_long = v0 * (cos(lat) * cos(alpha) + sin(lat) * cos(phi) * sin(alpha))
-    v_lat = -v0 * sin(phi) * sin(alpha)
+    vlat = -v0 * sin(phi) * sin(alpha)
 
     # Transform to Cartesian coordinate system
-    v1 = -cos(colat) * cos(phi) * v_lat - sin(phi) * v_long
-    v2 = -cos(colat) * sin(phi) * v_lat + cos(phi) * v_long
-    v3 = sin(colat) * v_lat
+    v1 = -cos(colat) * cos(phi) * vlat - sin(phi) * v_long
+    v2 = -cos(colat) * sin(phi) * vlat + cos(phi) * v_long
+    v3 = sin(colat) * vlat
 
     return prim2cons(SVector(rho, v1, v2, v3, 0), equations)
 end

examples/elixir_shallowwater_cubed_sphere_shell_EC_projection.jl

@@ -47,12 +47,12 @@ function initial_condition_advection_sphere(x, t, equations::ShallowWaterEquatio
     v0 = 1.0 # Velocity at the "equator"
     alpha = 0.0 #pi / 4
     v_long = v0 * (cos(lat) * cos(alpha) + sin(lat) * cos(phi) * sin(alpha))
-    v_lat = -v0 * sin(phi) * sin(alpha)
+    vlat = -v0 * sin(phi) * sin(alpha)
 
     # Transform to Cartesian coordinate system
-    v1 = -cos(colat) * cos(phi) * v_lat - sin(phi) * v_long
-    v2 = -cos(colat) * sin(phi) * v_lat + cos(phi) * v_long
-    v3 = sin(colat) * v_lat
+    v1 = -cos(colat) * cos(phi) * vlat - sin(phi) * v_long
+    v2 = -cos(colat) * sin(phi) * vlat + cos(phi) * v_long
+    v3 = sin(colat) * vlat
 
     return prim2cons(SVector(rho, v1, v2, v3, 0), equations)
 end

examples/elixir_shallowwater_cubed_sphere_shell_advection.jl

@@ -18,7 +18,7 @@ equations = ShallowWaterEquations3D(gravity_constant = 0.0)
 # Create DG solver with polynomial degree = 3 and (local) Lax-Friedrichs/Rusanov flux as surface flux
 solver = DGSEM(polydeg = polydeg, surface_flux = flux_lax_friedrichs)
 
-# Source term function to transform the Euler equations into the linear advection equations with variable advection velocity
+# Source term function to transform the Euler equations into a linear advection equation with variable advection velocity
 function source_terms_convert_to_linear_advection(u, du, x, t,
                                                   equations::ShallowWaterEquations3D,
                                                   normal_direction)
@@ -37,8 +37,12 @@ mesh = P4estMeshCubedSphere2D(cells_per_dimension, EARTH_RADIUS, polydeg = polyd
                               initial_refinement_level = 0,
                               element_local_mapping = true)
 
+# Convert initial condition given in terms of zonal and meridional velocity components to 
+# one given in terms of Cartesian momentum components
+initial_condition_transformed = spherical2cartesian(initial_condition, equations)
+
 # A semidiscretization collects data structures and functions for the spatial discretization
-semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
                                     source_terms = source_terms_convert_to_linear_advection)
 
 ###############################################################################

examples/elixir_shallowwater_cubed_sphere_shell_standard.jl

@@ -41,12 +41,12 @@ function initial_condition_advection_sphere(x, t, equations::ShallowWaterEquatio
     v0 = 1.0 # Velocity at the "equator"
     alpha = 0.0 #pi / 4
     v_long = v0 * (cos(lat) * cos(alpha) + sin(lat) * cos(phi) * sin(alpha))
-    v_lat = -v0 * sin(phi) * sin(alpha)
+    vlat = -v0 * sin(phi) * sin(alpha)
 
     # Transform to Cartesian coordinate system
-    v1 = -cos(colat) * cos(phi) * v_lat - sin(phi) * v_long
-    v2 = -cos(colat) * sin(phi) * v_lat + cos(phi) * v_long
-    v3 = sin(colat) * v_lat
+    v1 = -cos(colat) * cos(phi) * vlat - sin(phi) * v_long
+    v2 = -cos(colat) * sin(phi) * vlat + cos(phi) * v_long
+    v3 = sin(colat) * vlat
 
     return prim2cons(SVector(rho, v1, v2, v3, 0), equations)
 end

examples/elixir_spherical_advection_covariant.jl

@@ -28,8 +28,12 @@ volume_integral = VolumeIntegralWeakForm()
 solver = DGSEM(polydeg = polydeg, surface_flux = surface_flux,
                volume_integral = volume_integral)
 
+# Convert initial condition given in terms of zonal and meridional velocity components to 
+# one given in terms of contravariant velocity components
+initial_condition_transformed = spherical2contravariant(initial_condition, equations)
+
 # A semidiscretization collects data structures and functions for the spatial discretization
-semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver)
 
 ###############################################################################
 # ODE solvers, callbacks etc.

src/TrixiAtmo.jl

@@ -37,7 +37,8 @@ export velocity, waterheight, pressure, energy_total, energy_kinetic, energy_int
        clean_solution_lagrange_multiplier!
 export P4estCubedSphere2D, MetricTermsCrossProduct, MetricTermsInvariantCurl
 export EARTH_RADIUS, EARTH_GRAVITATIONAL_ACCELERATION,
-       EARTH_ROTATION_RATE, SECONDS_PER_DAY, spherical2cartesian
+       EARTH_ROTATION_RATE, SECONDS_PER_DAY
+export spherical2contravariant, contravariant2spherical, spherical2cartesian
 export initial_condition_gaussian
 
 export examples_dir

src/callbacks_step/analysis_covariant.jl

@@ -31,8 +31,7 @@ function Trixi.calc_error_norms(func, u, t, analyzer, mesh::P4estMesh{2},
         for j in eachnode(dg), i in eachnode(dg)
             x = Trixi.get_node_coords(node_coordinates, equations, dg, i, j, element)
 
-            u_exact = spherical2contravariant(initial_condition(x, t, equations),
-                                              equations, cache, (i, j), element)
+            u_exact = initial_condition(x, t, equations, cache, (i, j), element)
 
             u_numerical = Trixi.get_node_vars(u, equations, dg, i, j, element)
 

src/equations/covariant_advection.jl

@@ -25,19 +25,15 @@ J \frac{\partial}{\partial t}\left[\begin{array}{c} h \\ v^1 \\ v^2 \end{array}\
 = \left[\begin{array}{c} 0 \\ 0 \\ 0 \end{array}\right],
 ```
 where $J$ is the area element (see the documentation for [`P4estElementContainerCovariant`](@ref)). 
-!!! note
-    The initial condition should be prescribed as $[h, u, v]^{\mathrm{T}}$ in terms of the 
-    global velocity components $u$ and $v$ (i.e. zonal and meridional components in the 
-    case of a spherical shell). The transformation to local contravariant components $v^1$ 
-    and $v^2$ is handled internally within `Trixi.compute_coefficients!`.
 """
 struct CovariantLinearAdvectionEquation2D <: AbstractCovariantEquations{2, 3, 3} end
 
 function Trixi.varnames(::typeof(cons2cons), ::CovariantLinearAdvectionEquation2D)
-    return ("scalar", "v_con_1", "v_con_2")
+    return ("scalar", "vcon1", "vcon2")
 end
 
-# Compute the entropy variables (requires element container and indices)
+# We will define the "entropy variables" here to just be the scalar variable in the first 
+# slot, with zeros in the second and third positions
 @inline function Trixi.cons2entropy(u, equations::CovariantLinearAdvectionEquation2D,
                                     cache, node, element)
     z = zero(eltype(u))
@@ -82,19 +78,17 @@ end
 end
 
 # Convert contravariant velocity/momentum components to zonal and meridional components
-@inline function contravariant2spherical(u::SVector{3},
-                                         ::CovariantLinearAdvectionEquation2D,
+@inline function contravariant2spherical(u, ::CovariantLinearAdvectionEquation2D,
                                          cache, node, element)
-    v_lon, v_lat = contravariant2spherical(u[2], u[3], cache.elements, node, element)
-    return SVector(u[1], v_lon, v_lat)
+    vlon, vlat = contravariant2spherical(u[2], u[3], cache.elements, node, element)
+    return SVector(u[1], vlon, vlat)
 end
 
 # Convert zonal and meridional velocity/momentum components to contravariant components
-@inline function spherical2contravariant(u::SVector{3},
-                                         ::CovariantLinearAdvectionEquation2D,
+@inline function spherical2contravariant(u, ::CovariantLinearAdvectionEquation2D,
                                          cache, node, element)
-    v_con_1, v_con_2 = spherical2contravariant(u[2], u[3], cache.elements, node,
-                                               element)
-    return SVector(u[1], v_con_1, v_con_2)
+    vcon1, vcon2 = spherical2contravariant(u[2], u[3], cache.elements, node,
+                                           element)
+    return SVector(u[1], vcon1, vcon2)
 end
 end # @muladd

src/equations/equations.jl

@@ -28,23 +28,40 @@ Some references on discontinuous Galerkin methods in covariant flux form are lis
   Journal of Computational Physics 271:224-243. 
   [DOI: 10.1016/j.jcp.2013.11.033](https://doi.org/10.1016/j.jcp.2013.11.033)
 
-!!! note 
-    Components of vector-valued fields should be prescibed within the global coordinate 
-    system (i.e. zonal and meridional components in the case of a spherical shell). 
-    When dispatched on this type, the function `Trixi.compute_coefficients!` will 
-    internally use the `covariant_basis` field of the container type 
-    [`P4estElementContainerCovariant`](@ref) to obtain the local contravariant components
-    used in the solver. 
+When using this equation type, functions which are evaluated pointwise, such as fluxes, 
+source terms, and initial conditions take in the extra arguments `cache`, `node`, and 
+`element`, corresponding to the `cache` field of a `SemidiscretizationHyperbolic`, the node 
+index (for tensor-product elements, this should be a tuple of length `NDIMS`), and the 
+element index, respectively. To convert an initial condition given in terms of global 
+spherical velocity or momentum components to one given in terms of local contravariant 
+components, see [`spherical2contravariant`](@ref).
 """
 abstract type AbstractCovariantEquations{NDIMS,
                                          NDIMS_AMBIENT,
                                          NVARS} <: AbstractEquations{NDIMS, NVARS} end
 
+@doc raw"""
+    spherical2contravariant(initial_condition, ::AbstractCovariantEquations)
+
+Takes in a function with the signature `initial_condition(x, t)` which returns an initial 
+condition given in terms of zonal and meridional velocity or momentum components, and 
+returns another function with the signature  
+`initial_condition_transformed(x, t, equations, cache, node, element)` which returns
+the same initial condition with the velocity or momentum vector given in terms of 
+contravariant components.
+"""
+function spherical2contravariant(initial_condition, ::AbstractCovariantEquations)
+    function initial_condition_transformed(x, t, equations, cache, node, element)
+        return spherical2contravariant(initial_condition(x, t), equations, cache,
+                                       node, element)
+    end
+    return initial_condition_transformed
+end
 # Numerical flux plus dissipation which passes node/element indices and cache. 
 # We assume that u_ll and u_rr have been transformed into the same local coordinate system.
 @inline function (numflux::Trixi.FluxPlusDissipation)(u_ll, u_rr,
                                                       orientation_or_normal_direction,
-                                                      equations::AbstractCovariantEquations{2},
+                                                      equations::AbstractCovariantEquations,
                                                       cache, node_ll, node_rr, element)
 
     # The flux and dissipation need to be defined for the specific equation set
@@ -59,7 +76,7 @@ end
 # We assume that u_ll and u_rr have been transformed into the same local coordinate system.
 @inline function Trixi.flux_central(u_ll, u_rr,
                                     orientation_or_normal_direction,
-                                    equations::AbstractCovariantEquations{2},
+                                    equations::AbstractCovariantEquations,
                                     cache, node_ll, node_rr, element)
     flux_ll = Trixi.flux(u_ll, orientation_or_normal_direction, equations,
                          cache, node_ll, element)

src/equations/reference_data.jl

@@ -2,38 +2,56 @@
 #! format: noindent
 
 # Physical constants
-const EARTH_RADIUS = 6.37122  # m
+const EARTH_RADIUS = 6.37122e6   # m
 const EARTH_GRAVITATIONAL_ACCELERATION = 9.80616  # m/s²
 const EARTH_ROTATION_RATE = 7.292e-5  # rad/s
 const SECONDS_PER_DAY = 8.64e4
 
-function spherical2cartesian(vlon, vlat, x)
-    # Co-latitude
-    colat = acos(x[3] / sqrt(x[1]^2 + x[2]^2 + x[3]^2))
-
-    # Longitude
-    if sign(x[2]) == 0.0
-        signy = 1.0
-    else
-        signy = sign(x[2])
-    end
-    r_xy = sqrt(x[1]^2 + x[2]^2)
-    if r_xy == 0.0
-        lon = pi / 2
-    else
-        lon = signy * acos(x[1] / r_xy)
-    end
-
-    v1 = -cos(colat) * cos(lon) * vlat - sin(lon) * vlon
-    v2 = -cos(colat) * sin(lon) * vlat + cos(lon) * vlon
-    v3 = sin(colat) * vlat
-    return SVector(v1, v2, v3)
-end
+@doc raw"""
+    initial_condition_gaussian(x, t)
+
+This Gaussian bell case is a smooth initial condition suitable for testing the convergence 
+of discretizations of the linear advection equation on a spherical domain of radius $a = 6.
+37122 \times 10^3\ \mathrm{m}$, representing the surface of the Earth. The height field is 
+given in terms of the Cartesian coordinates $x$, $y$, $z$ as 
+```math
+h(x,y,z) = h_0 \exp\Big(-b_0 \big((x-x_0)^2 + (y-y_0)^2 + (z-z_0)^2\big) \Big),
+```
+where $h_0 = 1 \times 10^3\ \mathrm{m}$ is the height of the bell, $b_0 = 5 / a$ is the 
+width parameter, and the Cartesian coordinates of the bell's centre at longitude 
+$\lambda_0 = 3\pi/2$ and latitude $\theta_0 = 0$ are given by 
+```math
+\begin{aligned}
+x_0 &= a\cos\theta_0\cos\lambda_0,\\
+y_0 &= a\cos\theta_0\sin\lambda_0,\\
+z_0 &= a\sin\theta_0.
+\end{aligned}
+```
+The velocity field corresponds to a solid body rotation with a period of 12 days, at an 
+angle of $\alpha = \pi/4$ from the polar axis. The zonal and meridional components of the 
+velocity field are then given in terms of the longitude $\lambda$ and latitude $\theta$ as 
+```math
+\begin{aligned}
+u(\lambda,\theta) &= V (\cos\theta\cos\alpha + \sin\theta\cos\lambda\sin\alpha),\\
+v(\lambda,\theta)  &= -V \sin\lambda\sin\alpha,
+\end{aligned}
+```
+where we take $V = 2\pi a / 12\ \mathrm{days}$. This test case is adapted from Case 1 of 
+the test suite described in the following paper:
+- D. L. Williamson, J. B. Drake, J. J. Hack, R. Jakob, and P. N. Swarztrauber (1992). A  
+  standard test set for numerical approximations to the shallow water equations in
+  spherical geometry. Journal of Computational Physics, 102(1):211-224. 
+  [DOI: 10.1016/S0021-9991(05)80016-6](https://doi.org/10.1016/S0021-9991(05)80016-6)
 
-function initial_condition_gaussian(x, t)
+This function returns `SVector(h, vlon, vlat, b)`, where the first three entries are the 
+height, zonal velocity, and meridional velocity. The fourth entry, representing variable 
+bottom topography, is set to zero. The functions [`spherical2cartesian`](@ref) and 
+[`spherical2contravariant`](@ref) are available for converting to the prognostic variables 
+for Cartesian and covariant formulations.
+"""
+@inline function initial_condition_gaussian(x, t)
     RealT = eltype(x)
 
-    # set parameters
     a = EARTH_RADIUS  # radius of the sphere in metres
     V = convert(RealT, 2π) * a / (12 * SECONDS_PER_DAY)  # speed of rotation
     alpha = convert(RealT, π / 4)  # angle of rotation
@@ -42,7 +60,8 @@ function initial_condition_gaussian(x, t)
     lon_0, lat_0 = convert(RealT, 3π / 2), 0.0f0  # initial bump location
 
     # convert initial position to Cartesian coordinates
-    x_0 = SVector(a * cos(lat_0) * cos(lon_0), a * cos(lat_0) * sin(lon_0),
+    x_0 = SVector(a * cos(lat_0) * cos(lon_0),
+                  a * cos(lat_0) * sin(lon_0),
                   a * sin(lat_0))
 
     # compute Gaussian bump profile
@@ -53,10 +72,7 @@ function initial_condition_gaussian(x, t)
     vlon = V * (cos(lat) * cos(alpha) + sin(lat) * cos(lon) * sin(alpha))
     vlat = -V * sin(lon) * sin(alpha)
 
-    return SVector(h, vlon, vlat)
-end
-
-function initial_condition_gaussian(x, t, equations)
-    return initial_condition_gaussian(x, t)
+    # the last variable is the bottom topography, which we set to zero
+    return SVector(h, vlon, vlat, 0.0)
 end
 end # muladd

src/equations/shallow_water_3d.jl

@@ -204,7 +204,7 @@ Details are available in Eq. (4.1) in the paper:
 end
 
 """
-         source_terms_lagrange_multiplier(u, du, x, t,
+    source_terms_lagrange_multiplier(u, du, x, t,
                                           equations::ShallowWaterEquations3D,
                                           normal_direction)
 
@@ -381,9 +381,46 @@ end
     return energy_total(cons, equations) - energy_kinetic(cons, equations)
 end
 
-function initial_condition_gaussian(x, t, equations::ShallowWaterEquations3D)
-    h, vlon, vlat = initial_condition_gaussian(x, t)
-    v1, v2, v3 = spherical2cartesian(vlon, vlat, x)
-    return Trixi.prim2cons(SVector(h, v1, v2, v3, 0.0), equations)
+# Transform zonal and meridional velocity/momentum components to Cartesian components
+function spherical2cartesian(vlon, vlat, x)
+    # Co-latitude
+    colat = acos(x[3] / sqrt(x[1]^2 + x[2]^2 + x[3]^2))
+
+    # Longitude
+    if sign(x[2]) == 0.0
+        signy = 1.0
+    else
+        signy = sign(x[2])
+    end
+    r_xy = sqrt(x[1]^2 + x[2]^2)
+    if r_xy == 0.0
+        lon = pi / 2
+    else
+        lon = signy * acos(x[1] / r_xy)
+    end
+
+    v1 = -cos(colat) * cos(lon) * vlat - sin(lon) * vlon
+    v2 = -cos(colat) * sin(lon) * vlat + cos(lon) * vlon
+    v3 = sin(colat) * vlat
+
+    return SVector(v1, v2, v3)
+end
+
+@doc raw"""
+    spherical2cartesian(initial_condition, equations)
+
+Takes in a function with the signature `initial_condition(x, t)` which returns an initial 
+condition given in terms of zonal and meridional velocity or momentum components, and 
+returns another function with the signature 
+`initial_condition_transformed(x, t, equations)` which returns the same initial condition 
+with the velocity or momentum vector given in terms of Cartesian components.
+"""
+function spherical2cartesian(initial_condition, ::ShallowWaterEquations3D)
+    function initial_condition_transformed(x, t, equations)
+        h, vlon, vlat, b = initial_condition(x, t)
+        v1, v2, v3 = spherical2cartesian(vlon, vlat, x)
+        return Trixi.prim2cons(SVector(h, v1, v2, v3, b), equations)
+    end
+    return initial_condition_transformed
 end
 end # @muladd