Specialize save_solution_file for AbstractCovariantEquations to s…

…upport `solution_variables` which depend on auxiliary variables (#51) * new save_solution_file method for covariant form * save using conversion that depends on auxiliary variables * improve comment * move back spherical2cartesian * a_node -> aux_node * remove duplicated code and HDF5 dependency * remove using HDF5...
trixi-framework · Nov 30, 2024 · 8ce924e · 8ce924e
1 parent 40800a6
commit 8ce924e
Show file tree

Hide file tree

Showing 10 changed files with 106 additions and 35 deletions.
diff --git a/Project.toml b/Project.toml
@@ -8,6 +8,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
@@ -20,6 +21,7 @@ LinearAlgebra = "1"
 LoopVectorization = "0.12.118"
 MuladdMacro = "0.2.2"
 NLsolve = "4.5.1"
+Printf = "1"
 Reexport = "1.0"
 Static = "0.8, 1"
 StaticArrayInterface = "1.5.1"

diff --git a/examples/elixir_spherical_advection_covariant.jl b/examples/elixir_spherical_advection_covariant.jl
@@ -48,7 +48,7 @@ analysis_callback = AnalysisCallback(semi, interval = 10,
 
 # The SaveSolutionCallback allows to save the solution to a file in regular intervals
 save_solution = SaveSolutionCallback(interval = 10,
-                                     solution_variables = cons2cons)
+                                     solution_variables = contravariant2spherical)
 
 # The StepsizeCallback handles the re-calculation of the maximum Δt after each time step
 stepsize_callback = StepsizeCallback(cfl = 0.7)

diff --git a/src/TrixiAtmo.jl b/src/TrixiAtmo.jl
@@ -11,6 +11,7 @@ module TrixiAtmo
 using Reexport: @reexport
 using Trixi
 using MuladdMacro: @muladd
+using Printf: @sprintf
 using Static: True, False
 using StrideArrays: PtrArray
 using StaticArrayInterface: static_size

diff --git a/src/callbacks_step/analysis_covariant.jl b/src/callbacks_step/analysis_covariant.jl
@@ -11,13 +11,13 @@ function Trixi.analyze(::typeof(Trixi.entropy_timederivative), du, u, t,
     Trixi.integrate_via_indices(u, mesh, equations, dg, cache,
                                 du) do u, i, j, element, equations, dg, du_node
         # Get auxiliary variables, solution variables, and time derivative at given node
-        a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+        aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
         u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
         du_node = Trixi.get_node_vars(du, equations, dg, i, j, element)
 
         # compute ∂S/∂u ⋅ ∂u/∂t, where the entropy variables ∂S/∂u depend on the solution 
         # and auxiliary variables
-        dot(cons2entropy(u_node, a_node, equations), du_node)
+        dot(cons2entropy(u_node, aux_node, equations), du_node)
     end
 end
 
@@ -44,15 +44,15 @@ function Trixi.calc_error_norms(func, u, t, analyzer, mesh::P4estMesh{2},
 
             # Convert exact solution into contravariant components using geometric
             # information stored in aux vars
-            a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
-            u_exact = initial_condition(x_node, t, a_node, equations)
+            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+            u_exact = initial_condition(x_node, t, aux_node, equations)
 
             # Compute the difference as usual
             u_numerical = Trixi.get_node_vars(u, equations, dg, i, j, element)
             diff = func(u_exact, equations) - func(u_numerical, equations)
 
             # For the L2 error, integrate with respect to volume element stored in aux vars 
-            sqrtG = area_element(a_node, equations)
+            sqrtG = area_element(aux_node, equations)
             l2_error += diff .^ 2 * (weights[i] * weights[j] * sqrtG)
 
             # Compute Linf error as usual

diff --git a/src/callbacks_step/callbacks_step.jl b/src/callbacks_step/callbacks_step.jl
@@ -1,2 +1,3 @@
 include("analysis_covariant.jl")
+include("save_solution_covariant.jl")
 include("stepsize_dg2d.jl")
diff --git a/src/callbacks_step/save_solution_covariant.jl b/src/callbacks_step/save_solution_covariant.jl
@@ -0,0 +1,53 @@
+# Convert to another set of variables using a solution_variables function that depends on 
+# auxiliary variables
+function convert_variables(u, solution_variables, mesh::P4estMesh{2},
+                           equations, 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))
+    # 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)
+            for v in 1:n_vars
+                data[v, i, j, element] = u_node_converted[v]
+            end
+        end
+    end
+    return data
+end
+
+# Specialized save_solution_file method that supports a solution_variables function which 
+# depends on auxiliary variables. The conversion must be defined as solution_variables(u, 
+# aux_vars, equations), and an additional method must be defined as solution_variables(u,
+# equations) = u, such that no conversion is done when auxiliary variables are not provided.
+function Trixi.save_solution_file(u_ode, t, dt, iter,
+                                  semi::SemidiscretizationHyperbolic{<:Trixi.AbstractMesh,
+                                                                     <:AbstractCovariantEquations},
+                                  solution_callback,
+                                  element_variables = Dict{Symbol, Any}(),
+                                  node_variables = Dict{Symbol, Any}();
+                                  system = "")
+    mesh, equations, solver, cache = Trixi.mesh_equations_solver_cache(semi)
+    u = Trixi.wrap_array_native(u_ode, semi)
+
+    # Perform the variable conversion at each node
+    data = convert_variables(u, solution_callback.solution_variables, mesh, equations,
+                             solver, cache)
+
+    # Call the existing Trixi.save_solution_file, which will use solution_variables(u, 
+    # equations). Since the variables are already converted above, we therefore require 
+    # solution_variables(u, equations) = u.
+    Trixi.save_solution_file(data, t, dt, iter, mesh, equations, solver, cache,
+                             solution_callback, element_variables,
+                             node_variables, system = system)
+end
diff --git a/src/callbacks_step/stepsize_dg2d.jl b/src/callbacks_step/stepsize_dg2d.jl
@@ -59,8 +59,8 @@ function Trixi.max_dt(u, t, mesh::P4estMesh{2}, constant_speed::False,
         max_lambda1 = max_lambda2 = zero(max_scaled_speed)
         for j in eachnode(dg), i in eachnode(dg)
             u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
-            a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
-            lambda1, lambda2 = Trixi.max_abs_speeds(u_node, a_node, equations)
+            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+            lambda1, lambda2 = Trixi.max_abs_speeds(u_node, aux_node, equations)
             max_lambda1 = max(max_lambda1, lambda1)
             max_lambda2 = max(max_lambda2, lambda2)
         end

diff --git a/src/equations/covariant_advection.jl b/src/equations/covariant_advection.jl
@@ -47,6 +47,26 @@ function velocity(u, ::CovariantLinearAdvectionEquation2D)
     return SVector(u[2], u[3])
 end
 
+# Convert contravariant velocity/momentum components to zonal and meridional components
+@inline function contravariant2spherical(u, aux_vars,
+                                         equations::CovariantLinearAdvectionEquation2D)
+    vlon, vlat = basis_covariant(aux_vars, equations) * velocity(u, equations)
+    return SVector(u[1], vlon, vlat)
+end
+
+# Convert zonal and meridional velocity/momentum components to contravariant components
+@inline function spherical2contravariant(u_spherical, aux_vars,
+                                         equations::CovariantLinearAdvectionEquation2D)
+    vcon1, vcon2 = basis_covariant(aux_vars, equations) \
+                   velocity(u_spherical, equations)
+    return SVector(u_spherical[1], vcon1, vcon2)
+end
+
+function Trixi.varnames(::typeof(contravariant2spherical),
+                        ::CovariantLinearAdvectionEquation2D)
+    return ("scalar", "vlon", "vlat")
+end
+
 # 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, aux_vars,
@@ -91,19 +111,4 @@ end
                                       equations::CovariantLinearAdvectionEquation2D)
     return abs.(velocity(u, equations))
 end
-
-# Convert contravariant velocity/momentum components to zonal and meridional components
-@inline function contravariant2spherical(u, aux_vars,
-                                         equations::CovariantLinearAdvectionEquation2D)
-    vlon, vlat = basis_covariant(aux_vars, equations) * velocity(u, equations)
-    return SVector(u[1], vlon, vlat)
-end
-
-# Convert zonal and meridional velocity/momentum components to contravariant components
-@inline function spherical2contravariant(u_spherical, aux_vars,
-                                         equations::CovariantLinearAdvectionEquation2D)
-    vcon1, vcon2 = basis_covariant(aux_vars, equations) \
-                   velocity(u_spherical, equations)
-    return SVector(u_spherical[1], vcon1, vcon2)
-end
 end # @muladd
diff --git a/src/equations/equations.jl b/src/equations/equations.jl
@@ -48,6 +48,13 @@ dispatching on the return type.
 """
 @inline have_aux_node_vars(::AbstractEquations) = False()
 
+# cons2cons method which takes in unused aux_vars variable
+@inline Trixi.cons2cons(u, aux_vars, equations) = u
+
+# If no auxiliary variables are passed into the conversion to spherical coordinates, do not 
+# do any conversion.
+@inline contravariant2spherical(u, equations) = u
+
 # For the covariant form, the auxiliary variables are the the NDIMS^2 entries of the 
 # covariant basis matrix
 @inline have_aux_node_vars(::AbstractCovariantEquations) = True()

diff --git a/src/solvers/dgsem_p4est/dg_2d_manifold_in_3d_covariant.jl b/src/solvers/dgsem_p4est/dg_2d_manifold_in_3d_covariant.jl
@@ -15,10 +15,10 @@ function Trixi.compute_coefficients!(u, func, t, mesh::P4estMesh{2},
                                            element)
 
             # Get auxiliary variables at node (i, j)
-            a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
 
             # Compute nodal values of the provided function
-            u_node = func(x_node, t, a_node, equations)
+            u_node = func(x_node, t, aux_node, equations)
 
             # Set nodal variables in cache
             Trixi.set_node_vars!(u, u_node, equations, dg, i, j, element)
@@ -38,11 +38,11 @@ end
     for j in eachnode(dg), i in eachnode(dg)
         # Get solution variables and auxiliary variables at node (i, j)
         u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
-        a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+        aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
 
         # Evaluate the contravariant flux components
-        flux1 = flux(u_node, a_node, 1, equations)
-        flux2 = flux(u_node, a_node, 2, equations)
+        flux1 = flux(u_node, aux_node, 1, equations)
+        flux2 = flux(u_node, aux_node, 2, equations)
 
         # Apply weak form derivative with respect to ξ¹ 
         for ii in eachnode(dg)
@@ -73,16 +73,17 @@ end
     for j in eachnode(dg), i in eachnode(dg)
         # Get solution variables and auxiliary variables at node (i, j)
         u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
-        a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+        aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
 
         # ξ¹ direction
         for ii in (i + 1):nnodes(dg)
             # Get solution variables and auxiliary variables at node (ii, j)
             u_node_ii = Trixi.get_node_vars(u, equations, dg, ii, j, element)
-            a_node_ii = get_node_aux_vars(aux_node_vars, equations, dg, ii, j, element)
+            aux_node_ii = get_node_aux_vars(aux_node_vars, equations, dg, ii, j,
+                                            element)
 
             # Evaluate contravariant component 1 of the variable-coefficient two-point flux
-            flux1 = volume_flux(u_node, u_node_ii, a_node, a_node_ii, 1, equations)
+            flux1 = volume_flux(u_node, u_node_ii, aux_node, aux_node_ii, 1, equations)
 
             # Multiply by entry of split derivative matrix and add to right-hand side
             Trixi.multiply_add_to_node_vars!(du, alpha * derivative_split[i, ii], flux1,
@@ -95,10 +96,11 @@ end
         for jj in (j + 1):nnodes(dg)
             # Get solution variables and auxiliary variables at node (i, jj)
             u_node_jj = Trixi.get_node_vars(u, equations, dg, i, jj, element)
-            a_node_jj = get_node_aux_vars(aux_node_vars, equations, dg, i, jj, element)
+            aux_node_jj = get_node_aux_vars(aux_node_vars, equations, dg, i, jj,
+                                            element)
 
             # Evaluate contravariant component 2 of the variable-coefficient two-point flux
-            flux2 = volume_flux(u_node, u_node_jj, a_node, a_node_jj, 2, equations)
+            flux2 = volume_flux(u_node, u_node_jj, aux_node, aux_node_jj, 2, equations)
 
             # Multiply by entry of split derivative matrix and add to right-hand side
             Trixi.multiply_add_to_node_vars!(du, alpha * derivative_split[j, jj], flux2,
@@ -181,8 +183,8 @@ function Trixi.apply_jacobian!(du, mesh::P4estMesh{2},
 
     Trixi.@threaded for element in eachelement(dg, cache)
         for j in eachnode(dg), i in eachnode(dg)
-            a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
-            factor = -1 / area_element(a_node, equations)
+            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+            factor = -1 / area_element(aux_node, equations)
 
             for v in eachvariable(equations)
                 du[v, i, j, element] *= factor