Specialize save_solution_file for AbstractCovariantEquations to support solution_variables which depend on auxiliary variables

Project.toml

@@ -4,10 +4,12 @@ authors = ["Benedict Geihe <bgeihe@uni-koeln.de>", "Tristan Montoya <montoya.tri
 version = "0.1.0-DEV"
 
 [deps]
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 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"
@@ -16,10 +18,12 @@ StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
 Trixi = "a7f1ee26-1774-49b1-8366-f1abc58fbfcb"
 
 [compat]
+HDF5 = "0.17"
 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"

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)

src/TrixiAtmo.jl

@@ -10,7 +10,9 @@ module TrixiAtmo
 
 using Reexport: @reexport
 using Trixi
+using HDF5: HDF5, h5open, attributes
 using MuladdMacro: @muladd
+using Printf: @sprintf
 using Static: True, False
 using StrideArrays: PtrArray
 using StaticArrayInterface: static_size

src/callbacks_step/callbacks_step.jl

@@ -1,2 +1,3 @@
 include("analysis_covariant.jl")
+include("save_solution_covariant.jl")
 include("stepsize_dg2d.jl")

src/callbacks_step/save_solution_covariant.jl

@@ -0,0 +1,106 @@
+# 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)
+            a_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
+            u_node_converted = solution_variables(u_node, a_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 
+# depending on auxiliary variables
+function Trixi.save_solution_file(u, time, dt, timestep,
+                                  mesh::Union{Trixi.SerialTreeMesh, StructuredMesh,
+                                              StructuredMeshView,
+                                              UnstructuredMesh2D, Trixi.SerialP4estMesh,
+                                              Trixi.SerialT8codeMesh},
+                                  equations::AbstractCovariantEquations, dg::DG,
+                                  cache,
+                                  solution_callback,
+                                  element_variables = Dict{Symbol, Any}(),
+                                  node_variables = Dict{Symbol, Any}();
+                                  system = "")
+    (; output_directory, solution_variables) = solution_callback
+
+    # Filename based on current time step
+    if isempty(system)
+        filename = joinpath(output_directory, @sprintf("solution_%09d.h5", timestep))
+    else
+        filename = joinpath(output_directory,
+                            @sprintf("solution_%s_%09d.h5", system, timestep))
+    end
+
+    # Convert to different set of variables if requested
+    if solution_variables === cons2cons
+        data = u
+    else
+        data = convert_variables(u, solution_variables, mesh, equations, dg, cache)
+    end
+    n_vars = size(data, 1)
+
+    # Open file (clobber existing content)
+    h5open(filename, "w") do file
+        # Add context information as attributes
+        attributes(file)["ndims"] = ndims(mesh)
+        attributes(file)["equations"] = Trixi.get_name(equations)
+        attributes(file)["polydeg"] = Trixi.polydeg(dg)
+        attributes(file)["n_vars"] = n_vars
+        attributes(file)["n_elements"] = nelements(dg, cache)
+        attributes(file)["mesh_type"] = Trixi.get_name(mesh)
+        attributes(file)["mesh_file"] = splitdir(mesh.current_filename)[2]
+        attributes(file)["time"] = convert(Float64, time) # Ensure that `time` is written as a double precision scalar
+        attributes(file)["dt"] = convert(Float64, dt) # Ensure that `dt` is written as a double precision scalar
+        attributes(file)["timestep"] = timestep
+
+        # Store each variable of the solution data
+        for v in 1:n_vars
+            # Convert to 1D array
+            file["variables_$v"] = vec(data[v, .., :])
+
+            # Add variable name as attribute
+            var = file["variables_$v"]
+            attributes(var)["name"] = Trixi.varnames(solution_variables, equations)[v]
+        end
+
+        # Store element variables
+        for (v, (key, element_variable)) in enumerate(element_variables)
+            # Add to file
+            file["element_variables_$v"] = element_variable
+
+            # Add variable name as attribute
+            var = file["element_variables_$v"]
+            attributes(var)["name"] = string(key)
+        end
+
+        # Store node variables
+        for (v, (key, node_variable)) in enumerate(node_variables)
+            # Add to file
+            file["node_variables_$v"] = node_variable
+
+            # Add variable name as attribute
+            var = file["node_variables_$v"]
+            attributes(var)["name"] = string(key)
+        end
+    end
+
+    return filename
+end

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

src/equations/equations.jl

@@ -48,6 +48,9 @@
 """
 @inline have_aux_node_vars(::AbstractEquations) = False()
 
+# cons2cons method which takes in unused aux_vars variable
+@inline Trixi.cons2cons(u, aux_vars, 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()