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

[tristanmontoya]

Previously, the collocation derivative was applied to the metric tensor components to evaluate the Christoffel symbols of the second kind approximately, using $\partial_k G_{ij} \approx \partial_k (I^N G_{ij})$ to compute

$$\Gamma_{jk}^i := \frac{1}{2}G^{il}\big(\partial_j G_{kl} + \partial_k G_{jl} - \partial_l G_{jk} \big).$$

Since the rest of the geometric information for the covariant solver is computed analytically for the spherical shell, I would like the Christoffel symbols to also be exact. Because the expressions get quite messy when differentiating by hand, an easy way to compute them is to use automatic differentiation. Since reverse-mode automatic differentiation is not needed here and Trixi.jl already has ForwardDiff.jl as a dependency, I've added the option to use ForwardDiff.jl to compute the Christoffel symbols with forward-mode automatic differentiation.

If the Christoffel symbols needed to be recomputed frequently (e.g. in an adaptive simulation, or for a time-varying metric), perhaps there would be an argument for coding them by hand for efficiency purposes, but, since for now they are only calculated once per simulation and stored at each node, using automatic differentiation makes the most sense to me.

As with the Cartesian solver, options for computing metric terms are specified using the keyword argument metric_terms in the SemidiscretizationHyperbolic constructor. Currently, the options are MetricTermsCovariantSphere(christoffel_symbols=ChristoffelSymbolsAutodiff()) (the new default) and MetricTermsCovariantSphere(christoffel_symbols=ChristoffelSymbolsCollocationDerivative()) (the old way). Other approaches could be added similarly if we so choose.

Adding this as one of several options for computing the metric terms in the covariant solver led to some restructuring of the code, particularly the containers. Previously, the same P4estElementContainerPtrArray was used for both the covariant and Cartesian formulations' cache.elements container, despite most of that information not being used for the covariant form (the covariant solver stores the metric information in cache.auxiliary_node_variables so that the nodal values can be passed into the flux, which cannot access cache.elements). Now, we have a new container type P4estElementContainerCovariant, which stores only what is needed for the covariant solver and specializes the metric terms separately from the Cartesian solver. This made Trixi.calc_mortar_flux! no longer work with covariant-form containers, but since we haven't implemented mortars yet for the covariant solver anyway, I could safely remove the Trixi.calc_mortar_flux! from the specialized Trixi.rhs!. That call will be added back if/when an adaptive covariant solver is implemented, in which case we can go back to using the standard Trixi.rhs! method for P4estMesh{2}.

To make these changes, several methods which dispatched on types like Union{AbstractEquations{3}, AbstractCovariantEquations{2,3}} were separated into the two cases to handle the different containers used by the different solvers.

The structs for specifying metric terms were also moved from semidiscretization/semidiscretization.jl to the more aptly named semidiscretization/metric_terms.jl.

When computing the Christoffel symbols exactly, the sensitivity of the covariant basis vector components described in #49 can be seen from the fact that ForwardDiff returns NaNs when using GlobalSphericalCoordinates and an even number of cells per dimension on each face of the cubed sphere, presumably due to this pole problem. The default option GlobalCartesianCoordinates avoids this problem, and should therefore always be used instead of GlobalSphericalCoordinates.

[codecov]

Codecov Report

Attention: Patch coverage is 99.17355% with 1 line in your changes missing coverage. Please review.

Project coverage is 90.99%. Comparing base (dcc9c9a) to head (bb9e921).

Files with missing lines	Patch %	Lines
...vers/dgsem_p4est/dg_2d_manifold_in_3d_covariant.jl	97.82%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #78      +/-   ##
==========================================
+ Coverage   90.68%   90.99%   +0.31%     
==========================================
  Files          22       22              
  Lines        2189     2277      +88     
==========================================
+ Hits         1985     2072      +87     
- Misses        204      205       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[amrueda]

Great work, @tristanmontoya!
Some minor comments below.

[amrueda]

I understand that this line can be removed because MetricTermsCovariantSphere() is the default value of metric_terms for this constructor. Isn't that the case? Any reason why we would like to keep it here and in all other elixirs?

[tristanmontoya]

Is it possible to pass metric_terms as a parameter to trixi_include otherwise?

[amrueda]

I don't think it is. However, for the testing, it is only needed to specify metric_terms explicitly in elixir_shallowwater_covariant_rossby_haurwitz.jl. I added a second review with code suggestions to improve the readability of the elixirs. What do you think of this variant?

[amrueda]

I'd suggest using more verbose names for these functions

Suggested change

	# Use ForwardDiff.jl to automatically differentiate the covariant metric tensor components
	function calc_metric_derivatives(v1, v2, v3, v4, xi1, xi2, radius, equations)
	dGdxi1 = derivative(x -> calc_metric_covariant(v1, v2, v3, v4, x, xi2, radius,
	equations), xi1)
	dGdxi2 = derivative(x -> calc_metric_covariant(v1, v2, v3, v4, xi1, x, radius,
	equations), xi2)
	return dGdxi1, dGdxi2
	end
	# Use the collocation derivative operator to numerically differentiate the covariant
	# metric tensor components
	function calc_metric_derivatives(aux_node_vars, equations, dg, i, j, element)
	# Use ForwardDiff.jl to automatically differentiate the covariant metric tensor components
	function calc_metric_derivatives_ad(v1, v2, v3, v4, xi1, xi2, radius, equations)
	dGdxi1 = derivative(x -> calc_metric_covariant(v1, v2, v3, v4, x, xi2, radius,
	equations), xi1)
	dGdxi2 = derivative(x -> calc_metric_covariant(v1, v2, v3, v4, xi1, x, radius,
	equations), xi2)
	return dGdxi1, dGdxi2
	end
	# Use the collocation derivative operator to numerically differentiate the covariant
	# metric tensor components
	function calc_metric_derivatives_collocation(aux_node_vars, equations, dg, i, j, element)

[amrueda]

See above

Suggested change

	dGdxi1, dGdxi2 = calc_metric_derivatives(v1, v2, v3, v4,
	dg.basis.nodes[i], dg.basis.nodes[j],
	radius, equations)
	# Compute Christoffel symbols of the first kind
	dGdxi1, dGdxi2 = calc_metric_derivatives_ad(v1, v2, v3, v4,
	dg.basis.nodes[i], dg.basis.nodes[j],
	radius, equations)
	# Compute Christoffel symbols of the first kind

[amrueda]

Suggested change

	aux_node_vars[21:26, i, j, element] = calc_christoffel_symbols(dGdxi1, dGdxi2,
	Gcon)
	aux_node_vars[21:26, i, j, element] = calc_christoffel_symbols_collocation(dGdxi1,
	dGdxi2,
	Gcon)

[amrueda]

Following the discussion in #78 (comment), I add these suggestions to improve the readability of the elixirs.

[amrueda]

Suggested change

metric_terms = MetricTermsCovariantSphere(),

[amrueda]

Suggested change

metric_terms = MetricTermsCovariantSphere(),

[amrueda]

Suggested change

metric_terms = MetricTermsCovariantSphere(),

[amrueda]

Suggested change

	# 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)
	# A semidiscretization collects data structures and functions for the spatial discretization.
	# Even though `metric_terms = MetricTermsCovariantSphere()` is default, we pass it here
	# explicitly, such that `metric_terms` can be adjusted from the `trixi_include()` call in the tests
	semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
	metric_terms = MetricTermsCovariantSphere(),
	source_terms = source_terms_geometric_coriolis)

[amrueda]

Suggested change

metric_terms = MetricTermsCovariantSphere(),

[amrueda]

Suggested change

metric_terms = MetricTermsCovariantSphere(),

[amrueda]

Suggested change

	semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
	metric_terms = MetricTermsCovariantSphere())
	semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver)

[amrueda]

Suggested change

	semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver,
	metric_terms = MetricTermsCovariantSphere())
	semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition_transformed, solver)