Add option to use ForwardDiff.jl to compute Christoffel symbols exactly and separate covariant form containers

Project.toml

@@ -4,6 +4,7 @@ authors = ["Benedict Geihe <bgeihe@uni-koeln.de>", "Tristan Montoya <montoya.tri
 version = "0.1.0-DEV"
 
 [deps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
@@ -19,6 +20,7 @@ StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
 Trixi = "a7f1ee26-1774-49b1-8366-f1abc58fbfcb"
 
 [compat]
+ForwardDiff = "0.10.36"
 HDF5 = "0.16.10, 0.17"
 LinearAlgebra = "1"
 LoopVectorization = "0.12.118"

examples/elixir_shallowwater_covariant_barotropic_instability.jl

@@ -39,6 +39,7 @@ initial_condition_transformed = transform_initial_condition(initial_condition, e
 
 # A semidiscretization collects data structures and functions for the spatial discretization
 semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere(),
                                     source_terms = source_terms_geometric_coriolis)
 
 ###############################################################################

examples/elixir_shallowwater_covariant_geostrophic_balance.jl

@@ -33,6 +33,7 @@ initial_condition_transformed = transform_initial_condition(initial_condition, e
 
 # A semidiscretization collects data structures and functions for the spatial discretization
 semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere(),
                                     source_terms = source_terms_geometric_coriolis)
 
 ###############################################################################

examples/elixir_shallowwater_covariant_isolated_mountain.jl

@@ -43,6 +43,7 @@ initial_condition_transformed = transform_initial_condition(initial_condition, e
 # specify the bottom topography.
 semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
                                     source_terms = source_terms_geometric_coriolis,
+                                    metric_terms = MetricTermsCovariantSphere(),
                                     auxiliary_field = bottom_topography_isolated_mountain)
 
 ###############################################################################

examples/elixir_shallowwater_covariant_rossby_haurwitz.jl

@@ -39,6 +39,7 @@ initial_condition_transformed = transform_initial_condition(initial_condition, e
 
 # A semidiscretization collects data structures and functions for the spatial discretization
 semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere(),
                                     source_terms = source_terms_geometric_coriolis)
 
 ###############################################################################

examples/elixir_shallowwater_covariant_unsteady_solid_body_rotation_EC.jl

@@ -41,6 +41,7 @@ initial_condition_transformed = transform_initial_condition(initial_condition, e
 # discretization. Here, we pass in the additional keyword argument "auxiliary_field" to 
 # specify the bottom topography.
 semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere(),
                                     source_terms = source_terms_geometric_coriolis,
                                     auxiliary_field = bottom_topography_unsteady_solid_body_rotation)
 

examples/elixir_shallowwater_covariant_well_balanced.jl

@@ -44,6 +44,7 @@ initial_condition_transformed = transform_initial_condition(initial_condition, e
 # specify the bottom topography, which is the same as for the standard isolated mountain 
 # case.
 semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere(),
                                     source_terms = source_terms_geometric_coriolis,
                                     auxiliary_field = bottom_topography_isolated_mountain)
 

examples/elixir_spherical_advection_covariant_cubed_sphere.jl

@@ -29,7 +29,8 @@ mesh = P4estMeshCubedSphere2D(cells_per_dimension[1], EARTH_RADIUS,
 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)
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere())
 
 ###############################################################################
 # ODE solvers, callbacks etc.

examples/elixir_spherical_advection_covariant_quad_icosahedron.jl

@@ -28,7 +28,8 @@ mesh = P4estMeshQuadIcosahedron2D(cells_per_dimension[1], EARTH_RADIUS,
 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)
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
+                                    metric_terms = MetricTermsCovariantSphere())
 
 ###############################################################################
 # ODE solvers, callbacks etc.

src/TrixiAtmo.jl

@@ -19,6 +19,7 @@ using LinearAlgebra: cross, norm, dot, det
 using Reexport: @reexport
 using LoopVectorization: @turbo
 using QuadGK: quadgk
+using ForwardDiff: derivative
 using HDF5: HDF5, h5open, attributes, create_dataset, datatype, dataspace
 
 @reexport using StaticArrays: SVector, SMatrix
@@ -29,7 +30,6 @@ include("equations/equations.jl")
 include("meshes/meshes.jl")
 include("semidiscretization/semidiscretization.jl")
 include("solvers/solvers.jl")
-include("semidiscretization/semidiscretization_hyperbolic_2d_manifold_in_3d.jl")
 include("callbacks_step/callbacks_step.jl")
 
 export CompressibleMoistEulerEquations2D, ShallowWaterEquations3D,
@@ -48,7 +48,8 @@ export source_terms_lagrange_multiplier, clean_solution_lagrange_multiplier!
 export cons2prim_and_vorticity, contravariant2global
 
 export P4estMeshCubedSphere2D, P4estMeshQuadIcosahedron2D, MetricTermsCrossProduct,
-       MetricTermsInvariantCurl
+       MetricTermsInvariantCurl, MetricTermsCovariantSphere, ChristoffelSymbolsAutodiff,
+       ChristoffelSymbolsCollocationDerivative
 
 export EARTH_RADIUS, EARTH_GRAVITATIONAL_ACCELERATION,
        EARTH_ROTATION_RATE, SECONDS_PER_DAY

src/semidiscretization/metric_terms.jl

@@ -0,0 +1,73 @@
+"""
+    MetricTermsCrossProduct()
+
+Struct used for multiple dispatch on functions that compute the metric terms.
+When the argument `metric_terms` is of type `MetricTermsCrossProduct`, the 
+contravariant vectors are computed using the cross-product form.
+"""
+struct MetricTermsCrossProduct end
+
+"""
+    MetricTermsInvariantCurl()
+
+Struct used for multiple dispatch on functions that compute the metric terms.
+When the argument `metric_terms` is of type `MetricTermsInvariantCurl`, the 
+contravariant vectors are computed using the invariant curl form.
+## References
+- Kopriva, D. A. (2006). Metric identities and the discontinuous spectral element method on 
+  curvilinear meshes. Journal of Scientific Computing 26, 301-327. 
+  [DOI: 10.1007/s10915-005-9070-8](https://doi.org/10.1007/s10915-005-9070-8)
+- Vinokur, M. and Yee, H. C. (2001). Extension of efficient low dissipation high order
+  schemes for 3-D curvilinear moving grids. In Caughey, D. A., and Hafez, M. M. (eds.),
+  Frontiers of Computational Fluid Dynamics 2002, World Scientific, Singapore, pp. 129–164.
+  [DOI: 10.1142/9789812810793_0008](https://doi.org/10.1142/9789812810793_0008)
+"""
+struct MetricTermsInvariantCurl end
+
+"""
+    MetricTermsCovariantSphere(christoffel_symbols = ChristoffelSymbolsAutodiff())
+
+Struct specifying options for computing geometric information for discretizations in 
+covariant form based on an exact representation of the spherical geometry.  Currently, the 
+only field is `christoffel_symbols`, specifying the approach used to compute the 
+Christoffel symbols, for which the options are [`ChristoffelSymbolsAutodiff`](@ref) or 
+[`ChristoffelSymbolsCollocationDerivative`](@ref). 
+"""
+struct MetricTermsCovariantSphere{ChristoffelSymbols}
+    christoffel_symbols::ChristoffelSymbols
+    function MetricTermsCovariantSphere(;
+                                        christoffel_symbols::ChristoffelSymbols = ChristoffelSymbolsAutodiff()) where {ChristoffelSymbols}
+        return new{ChristoffelSymbols}(christoffel_symbols)
+    end
+end
+
+@doc raw"""
+    ChristoffelSymbolsAutodiff()
+
+Struct used for multiple dispatch on functions that compute the Christoffel symbols. This 
+option uses [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) to compute
+```math
+\Gamma_{jk}^i =
+\frac{1}{2}G^{il}\big(\partial_j G_{kl} + \partial_k G_{jl} - \partial_l G_{jk}\big)
+```
+using forward-mode automatic differentiation. 
+!!! warning
+    Using this option with [`GlobalSphericalCoordinates`](@ref) is prone to `NaN` values as 
+    a result of the polar singularity. This is remedied through the use of 
+    [`GlobalCartesianCoordinates`](@ref).
+"""
+struct ChristoffelSymbolsAutodiff end
+
+@doc raw"""
+    ChristoffelSymbolsCollocationDerivative()
+
+Struct used for multiple dispatch on functions that compute the Christoffel symbols. 
+Letting $I^N$ denote the degree $N$ polynomial interpolation operator collocated with the 
+scheme's quadrature nodes, this option computes the Christoffel symbols at each quadrature 
+node using the approximation
+```math
+\Gamma_{jk}^i \approx
+\frac{1}{2}G^{il}\big(\partial_j I^N G_{kl} + \partial_k I^N G_{jl} - \partial_l I^N G_{jk}\big).
+```
+"""
+struct ChristoffelSymbolsCollocationDerivative end

src/semidiscretization/semidiscretization.jl

@@ -1,25 +1,2 @@
-"""
-    MetricTermsCrossProduct()
-
-Struct used for multiple dispatch on functions that compute the metric terms.
-When the argument `metric_terms` is of type `MetricTermsCrossProduct`, the 
-contravariant vectors are computed using the cross-product form.
-"""
-struct MetricTermsCrossProduct end
-
-"""
-    MetricTermsInvariantCurl()
-
-Struct used for multiple dispatch on functions that compute the metric terms.
-When the argument `metric_terms` is of type `MetricTermsInvariantCurl`, the 
-contravariant vectors are computed using the invariant curl form.
-## References
-- Kopriva, D. A. (2006). Metric identities and the discontinuous spectral element method on 
-  curvilinear meshes. Journal of Scientific Computing 26, 301-327. 
-  [DOI: 10.1007/s10915-005-9070-8](https://doi.org/10.1007/s10915-005-9070-8)
-- Vinokur, M. and Yee, H. C. (2001). Extension of efficient low dissipation high order
-  schemes for 3-D curvilinear moving grids. In Caughey, D. A., and Hafez, M. M. (eds.),
-  Frontiers of Computational Fluid Dynamics 2002, World Scientific, Singapore, pp. 129–164.
-  [DOI: 10.1142/9789812810793_0008](https://doi.org/10.1142/9789812810793_0008)
-"""
-struct MetricTermsInvariantCurl end
+include("metric_terms.jl")
+include("semidiscretization_hyperbolic_2d_manifold_in_3d.jl")

src/semidiscretization/semidiscretization_hyperbolic_2d_manifold_in_3d.jl

@@ -1,17 +1,14 @@
 # The SemidiscretizationHyperbolic constructor has been modified to remove the assertion
 # that checks the compatibility between the mesh dimensionality and the equations' 
 # dimensionality. Instead, we now directly dispatch using the specific mesh type 
-# (P4estMesh{2}) for 2D meshes and AbstractEquations{3} for 3D equations or 
-# AbstractCovariantEquations{2, 3} for 2D covariant equations in three-dimensional space. 
+# (P4estMesh{2}) for 2D meshes and AbstractEquations{3} for 3D equations.
 # This change is necessary to support the implementation of a 2D manifold embedded in a 3D 
 # space. We also pass in the keyword arguments "metric_terms" and "auxiliary_field", which 
 # specify the approach used to compute the metric terms as well as any auxiliary fields 
 # (e.g. bottom topography or a background state) that may be needed by the solver but are 
 # not stored in the solution variables or elsewhere in the cache.
 function Trixi.SemidiscretizationHyperbolic(mesh::P4estMesh{2},
-                                            equations::Union{AbstractEquations{3},
-                                                             AbstractCovariantEquations{2,
-                                                                                        3}},
+                                            equations::AbstractEquations{3},
                                             initial_condition,
                                             solver;
                                             source_terms = nothing,
@@ -38,3 +35,37 @@ function Trixi.SemidiscretizationHyperbolic(mesh::P4estMesh{2},
                                                                 source_terms, solver,
                                                                 cache)
 end
+
+# Constructor for SemidiscretizationHyperbolic for the covariant form. Requires 
+# compatibility between the mesh and equations (i.e. the same `NDIMS` and `NDIMS_AMBIENT`)
+# and sets the default metric terms to MetricTermsCovariantSphere.
+function Trixi.SemidiscretizationHyperbolic(mesh::P4estMesh{NDIMS, NDIMS_AMBIENT},
+                                            equations::AbstractCovariantEquations{NDIMS,
+                                                                                  NDIMS_AMBIENT},
+                                            initial_condition,
+                                            solver;
+                                            source_terms = nothing,
+                                            boundary_conditions = boundary_condition_periodic,
+                                            # `RealT` is used as real type for node locations etc.
+                                            # while `uEltype` is used as element type of solutions etc.
+                                            RealT = real(solver), uEltype = RealT,
+                                            initial_cache = NamedTuple(),
+                                            metric_terms = MetricTermsCovariantSphere(),
+                                            auxiliary_field = nothing) where {NDIMS,
+                                                                              NDIMS_AMBIENT}
+    cache = (;
+             Trixi.create_cache(mesh, equations, solver, RealT, metric_terms,
+                                auxiliary_field, uEltype)..., initial_cache...)
+    _boundary_conditions = Trixi.digest_boundary_conditions(boundary_conditions, mesh,
+                                                            solver,
+                                                            cache)
+
+    SemidiscretizationHyperbolic{typeof(mesh), typeof(equations),
+                                 typeof(initial_condition),
+                                 typeof(_boundary_conditions), typeof(source_terms),
+                                 typeof(solver), typeof(cache)}(mesh, equations,
+                                                                initial_condition,
+                                                                _boundary_conditions,
+                                                                source_terms, solver,
+                                                                cache)
+end

src/solvers/dgsem_p4est/containers_2d_manifold_in_3d_cartesian.jl

@@ -65,8 +65,7 @@ end
 # This function dispatches on the dimensions of the mesh and the equation (AbstractEquations{3})
 function Trixi.init_elements(mesh::Union{P4estMesh{2, 3, RealT},
                                          T8codeMesh{2}},
-                             equations::Union{AbstractEquations{3},
-                                              AbstractCovariantEquations{2, 3}},
+                             equations::AbstractEquations{3},
                              basis,
                              metric_terms,
                              ::Type{uEltype}) where {RealT <: Real, uEltype <: Real}
@@ -140,7 +139,7 @@ function init_elements_2d_manifold_in_3d!(elements,
     contravariant_vectors, inverse_jacobian) = elements
 
     # The standard calc_node_coordinates! can be used, since Trixi.jl now dispatches on
-    # P4estMesh{NDIMS, NDIMS_AMBIENT}, so it can be used here.
+    # P4estMesh{NDIMS, NDIMS_AMBIENT}.
     Trixi.calc_node_coordinates!(node_coordinates, mesh, basis)
 
     for element in 1:Trixi.ncells(mesh)