Move Trixi imports to TrixiAtmo.jl for all the equations (#115)

MarcoArtiano · ranocha · web-flow · commit 9dd3b48450a4 · 2025-10-02T14:22:10.000+02:00
* first commit

* add import AbstractShallowWaterEquations in equations.jl

* fix typo

* increase atol to make CI pass for MacOS

* update atol

* increase coverage

* Apply suggestions from code review

Co-authored-by: Hendrik Ranocha &lt;ranocha@users.noreply.github.com&gt;

* format

---------

Co-authored-by: Hendrik Ranocha &lt;ranocha@users.noreply.github.com&gt;
diff --git a/examples/elixir_shallowwater_cartesian_geostrophic_balance.jl b/examples/elixir_shallowwater_cartesian_geostrophic_balance.jl
@@ -66,7 +66,9 @@ summary_callback = SummaryCallback()
 # results
 analysis_callback = AnalysisCallback(semi, interval = 200,
                                      save_analysis = true,
-                                     extra_analysis_errors = (:conservation_error,))
+                                     extra_analysis_errors = (:conservation_error,),
+                                     extra_analysis_integrals = (energy_internal,
+                                                                 energy_kinetic, pressure))
 
 # The SaveSolutionCallback allows to save the solution to a file in regular intervals
 save_solution = SaveSolutionCallback(dt = (tspan[2] - tspan[1]) / n_saves,
diff --git a/src/TrixiAtmo.jl b/src/TrixiAtmo.jl
@@ -23,11 +23,15 @@ using ForwardDiff: derivative
 using HDF5: HDF5, h5open, attributes, create_dataset, datatype, dataspace
 
 @reexport using StaticArrays: SVector, SMatrix
-@reexport import Trixi: waterheight, varnames, cons2cons, cons2prim, prim2cons,
-                        cons2entropy, boundary_condition_slip_wall, flux, flux_ec,
-                        flux_chandrashekar, FluxLMARS, energy_total,
-                        max_abs_speeds, max_abs_speed_naive, max_abs_speed, energy_kinetic,
-                        energy_total, entropy, pressure, have_nonconservative_terms
+@reexport import Trixi: waterheight, varnames, cons2cons, cons2prim,
+                        prim2cons, cons2entropy, entropy2cons, velocity,
+                        max_abs_speeds, max_abs_speed_naive, max_abs_speed,
+                        have_nonconservative_terms, boundary_condition_slip_wall,
+                        energy_kinetic, energy_internal, energy_total, entropy, pressure,
+                        flux, flux_ec, flux_chandrashekar, flux_wintermeyer_etal,
+                        flux_fjordholm_etal, flux_nonconservative_wintermeyer_etal,
+                        flux_nonconservative_fjordholm_etal, FluxLMARS
+
 using Trixi: ln_mean, stolarsky_mean, inv_ln_mean
 
 include("auxiliary/auxiliary.jl")
diff --git a/src/equations/covariant_advection.jl b/src/equations/covariant_advection.jl
@@ -48,7 +48,7 @@ end
 
 # The conservative variables are the scalar conserved quantity and two contravariant 
 # velocity components.
-function Trixi.varnames(::typeof(cons2cons), ::CovariantLinearAdvectionEquation2D)
+function varnames(::typeof(cons2cons), ::CovariantLinearAdvectionEquation2D)
     return ("h", "vcon1", "vcon2")
 end
 
@@ -73,23 +73,23 @@ end
 end
 
 # Scalar conserved quantity and three global velocity components
-function Trixi.varnames(::typeof(contravariant2global),
-                        ::CovariantLinearAdvectionEquation2D)
+function varnames(::typeof(contravariant2global),
+                  ::CovariantLinearAdvectionEquation2D)
     return ("h", "v1", "v2", "v3")
 end
 
 # We will define the "entropy variables" here to just be the scalar variable in the first 
 # slot, with zeros in all other positions
-@inline function Trixi.cons2entropy(u, aux_vars,
-                                    equations::CovariantLinearAdvectionEquation2D)
+@inline function cons2entropy(u, aux_vars,
+                              equations::CovariantLinearAdvectionEquation2D)
     z = zero(eltype(u))
     return SVector(u[1], z, z)
 end
 
 # Flux as a function of the state vector u, as well as the auxiliary variables aux_vars, 
 # which contain the geometric information required for the covariant form
-@inline function Trixi.flux(u, aux_vars, orientation::Integer,
-                            equations::CovariantLinearAdvectionEquation2D)
+@inline function flux(u, aux_vars, orientation::Integer,
+                      equations::CovariantLinearAdvectionEquation2D)
     z = zero(eltype(u))
     J = area_element(aux_vars, equations)
     vcon = velocity_contravariant(u, equations)
@@ -110,17 +110,17 @@ end
 end
 
 # Maximum contravariant wave speed with respect to specific basis vector
-@inline function Trixi.max_abs_speed(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
-                                     orientation::Integer,
-                                     equations::CovariantLinearAdvectionEquation2D)
+@inline function max_abs_speed(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
+                               orientation::Integer,
+                               equations::CovariantLinearAdvectionEquation2D)
     vcon_ll = velocity_contravariant(u_ll, equations)  # Contravariant components on left side
     vcon_rr = velocity_contravariant(u_rr, equations)  # Contravariant components on right side
     return max(abs(vcon_ll[orientation]), abs(vcon_rr[orientation]))
 end
 
 # Maximum wave speeds in each direction for CFL calculation
-@inline function Trixi.max_abs_speeds(u, aux_vars,
-                                      equations::CovariantLinearAdvectionEquation2D)
+@inline function max_abs_speeds(u, aux_vars,
+                                equations::CovariantLinearAdvectionEquation2D)
     return abs.(velocity_contravariant(u, equations))
 end
 
diff --git a/src/equations/covariant_shallow_water.jl b/src/equations/covariant_shallow_water.jl
@@ -77,15 +77,15 @@ struct CovariantShallowWaterEquations2D{GlobalCoordinateSystem, RealT <: Real} <
 end
 
 # Until we implement bottom topography, there are no nonconservative terms
-Trixi.have_nonconservative_terms(::CovariantShallowWaterEquations2D) = False()
+have_nonconservative_terms(::CovariantShallowWaterEquations2D) = False()
 
 # The conservative variables are the height and contravariant momentum components
-function Trixi.varnames(::typeof(cons2cons), ::AbstractCovariantShallowWaterEquations2D)
+function varnames(::typeof(cons2cons), ::AbstractCovariantShallowWaterEquations2D)
     return ("h", "h_vcon1", "h_vcon2")
 end
 
 # The primitive variables are the height and contravariant velocity components
-function Trixi.varnames(::typeof(cons2prim), ::AbstractCovariantShallowWaterEquations2D)
+function varnames(::typeof(cons2prim), ::AbstractCovariantShallowWaterEquations2D)
     return ("H", "vcon1", "vcon2")
 end
 
@@ -94,13 +94,13 @@ end
 # equations.global_coordinate_system (e.g. spherical or Cartesian). This transformation 
 # works for both primitive and conservative variables, although varnames refers 
 # specifically to transformations from conservative variables.
-function Trixi.varnames(::typeof(contravariant2global),
-                        ::AbstractCovariantShallowWaterEquations2D)
+function varnames(::typeof(contravariant2global),
+                  ::AbstractCovariantShallowWaterEquations2D)
     return ("h", "h_v1", "h_v2", "h_v3")
 end
 
 # Convenience functions to extract physical variables from state vector
-@inline Trixi.waterheight(u, ::AbstractCovariantShallowWaterEquations2D) = u[1]
+@inline waterheight(u, ::AbstractCovariantShallowWaterEquations2D) = u[1]
 @inline velocity_contravariant(u,
 ::AbstractCovariantShallowWaterEquations2D) = SVector(u[2] /
                                                       u[1],
@@ -110,24 +110,24 @@ end
 ::AbstractCovariantShallowWaterEquations2D) = SVector(u[2],
                                                       u[3])
 
-@inline function Trixi.cons2prim(u, aux_vars,
-                                 equations::AbstractCovariantShallowWaterEquations2D)
+@inline function cons2prim(u, aux_vars,
+                           equations::AbstractCovariantShallowWaterEquations2D)
     h, h_vcon1, h_vcon2 = u
     h_s = bottom_topography(aux_vars, equations)
     return SVector(h + h_s, h_vcon1 / h, h_vcon2 / h)
 end
 
-@inline function Trixi.prim2cons(u, aux_vars,
-                                 equations::AbstractCovariantShallowWaterEquations2D)
+@inline function prim2cons(u, aux_vars,
+                           equations::AbstractCovariantShallowWaterEquations2D)
     H, vcon1, vcon2 = u
     h_s = bottom_topography(aux_vars, equations)
     h = H - h_s
     return SVector(h, h * vcon1, h * vcon2)
 end
 
 # Entropy variables are w = (g(h+hₛ) - (v₁v¹ + v₂v²)/2, v₁, v₂)ᵀ
-@inline function Trixi.cons2entropy(u, aux_vars,
-                                    equations::AbstractCovariantShallowWaterEquations2D)
+@inline function cons2entropy(u, aux_vars,
+                              equations::AbstractCovariantShallowWaterEquations2D)
     h = waterheight(u, equations)
     h_s = bottom_topography(aux_vars, equations)
     vcon = velocity_contravariant(u, equations)
@@ -153,8 +153,8 @@ end
 end
 
 # Entropy function (total energy) given by S = (h(v₁v¹ + v₂v²) + gh² + ghhₛ)/2
-@inline function Trixi.entropy(u, aux_vars,
-                               equations::AbstractCovariantShallowWaterEquations2D)
+@inline function entropy(u, aux_vars,
+                         equations::AbstractCovariantShallowWaterEquations2D)
     h = waterheight(u, equations)
     h_s = bottom_topography(aux_vars, equations)
     vcon = velocity_contravariant(u, equations)
@@ -165,8 +165,8 @@ end
 
 # Flux as a function of the state vector u, as well as the auxiliary variables aux_vars, 
 # which contain the geometric information required for the covariant form
-@inline function Trixi.flux(u, aux_vars, orientation::Integer,
-                            equations::CovariantShallowWaterEquations2D)
+@inline function flux(u, aux_vars, orientation::Integer,
+                      equations::CovariantShallowWaterEquations2D)
     # Geometric variables
     Gcon = metric_contravariant(aux_vars, equations)
     J = area_element(aux_vars, equations)
@@ -217,9 +217,9 @@ end
 end
 
 # Maximum wave speed along the normal direction in reference space
-@inline function Trixi.max_abs_speed(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
-                                     orientation,
-                                     equations::AbstractCovariantShallowWaterEquations2D)
+@inline function max_abs_speed(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
+                               orientation,
+                               equations::AbstractCovariantShallowWaterEquations2D)
     # Geometric variables
     Gcon_ll = metric_contravariant(aux_vars_ll, equations)
     Gcon_rr = metric_contravariant(aux_vars_rr, equations)
@@ -237,8 +237,8 @@ end
 end
 
 # Maximum wave speeds with respect to the covariant basis
-@inline function Trixi.max_abs_speeds(u, aux_vars,
-                                      equations::AbstractCovariantShallowWaterEquations2D)
+@inline function max_abs_speeds(u, aux_vars,
+                                equations::AbstractCovariantShallowWaterEquations2D)
     vcon = velocity_contravariant(u, equations)
     h = waterheight(u, equations)
     Gcon = metric_contravariant(aux_vars, equations)
diff --git a/src/equations/covariant_shallow_water_split.jl b/src/equations/covariant_shallow_water_split.jl
@@ -58,13 +58,13 @@ end
 
 # This alternative flux formulation has non-conservative terms even in the absence of 
 # variable bottom topography
-Trixi.have_nonconservative_terms(::SplitCovariantShallowWaterEquations2D) = True()
+have_nonconservative_terms(::SplitCovariantShallowWaterEquations2D) = True()
 
 # Flux as a function of the state vector u, as well as the auxiliary variables aux_vars, 
 # which contain the geometric information required for the covariant form.
 # Note that this flux does not include the pressure term.
-@inline function Trixi.flux(u, aux_vars, orientation::Integer,
-                            equations::SplitCovariantShallowWaterEquations2D)
+@inline function flux(u, aux_vars, orientation::Integer,
+                      equations::SplitCovariantShallowWaterEquations2D)
     # Geometric variables
     J = area_element(aux_vars, equations)
 
@@ -80,7 +80,7 @@ Trixi.have_nonconservative_terms(::SplitCovariantShallowWaterEquations2D) = True
 end
 
 @doc raw"""
-    Trixi.flux_ec(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
+    flux_ec(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
                                orientation::Integer,
                                equations::SplitCovariantShallowWaterEquations2D)
 
@@ -97,9 +97,9 @@ paper for the special case of the Euclidean metric $G_{ab} = \delta_{ab}$:
     The use of entropy-stable split-form/flux-differencing formulations for covariant 
     equations is an experimental feature and may change in future releases.
 """
-@inline function Trixi.flux_ec(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
-                               orientation::Integer,
-                               equations::SplitCovariantShallowWaterEquations2D)
+@inline function flux_ec(u_ll, u_rr, aux_vars_ll, aux_vars_rr,
+                         orientation::Integer,
+                         equations::SplitCovariantShallowWaterEquations2D)
     # Geometric variables
     J_ll = area_element(aux_vars_ll, equations)
     J_rr = area_element(aux_vars_rr, equations)
diff --git a/src/equations/equations.jl b/src/equations/equations.jl
@@ -1,7 +1,8 @@
 @muladd begin
 #! format: noindent
 
-using Trixi: AbstractEquations, AbstractCompressibleEulerEquations
+using Trixi: AbstractEquations, AbstractCompressibleEulerEquations,
+             AbstractShallowWaterEquations
 
 @doc raw"""
     AbstractCovariantEquations{NDIMS, 
diff --git a/src/equations/shallow_water_3d.jl b/src/equations/shallow_water_3d.jl
diff --git a/test/test_3d_potential_temperature.jl b/test/test_3d_potential_temperature.jl
