From b3ecc39c26807a22278a8e76c7383c914a32bcd2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Rueda-Ram=C3=ADrez?= <aruedara@uni-koeln.de>
Date: Mon, 2 Dec 2024 17:55:17 +0100
Subject: [PATCH 1/3] Added routines to plot vorticity for SWE 3D

---
 Project.toml                                  |   2 +
 ...wwater_quad_icosahedron_shell_advection.jl |   2 +-
 src/TrixiAtmo.jl                              |   2 +
 src/callbacks_step/callbacks_step.jl          |   1 +
 ...ve_solution_2d_manifold_in_3d_cartesian.jl | 197 ++++++++++++++++++
 5 files changed, 203 insertions(+), 1 deletion(-)
 create mode 100644 src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl

diff --git a/Project.toml b/Project.toml
index 0a4f134..7acc0d0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,7 @@ 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"
@@ -17,6 +18,7 @@ StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
 Trixi = "a7f1ee26-1774-49b1-8366-f1abc58fbfcb"
 
 [compat]
+HDF5 = "0.16.10, 0.17"
 LinearAlgebra = "1"
 LoopVectorization = "0.12.118"
 MuladdMacro = "0.2.2"
diff --git a/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl b/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl
index a9ff9dc..55fd161 100644
--- a/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl
+++ b/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl
@@ -70,7 +70,7 @@ analysis_callback = AnalysisCallback(semi, interval = 100,
 
 # The SaveSolutionCallback allows to save the solution to a file in regular intervals
 save_solution = SaveSolutionCallback(interval = 100,
-                                     solution_variables = cons2prim)
+                                     solution_variables = cons2prim_plus_vorticity)
 
 # 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
index 6803370..7903000 100644
--- a/src/TrixiAtmo.jl
+++ b/src/TrixiAtmo.jl
@@ -18,6 +18,7 @@ using StaticArrayInterface: static_size
 using LinearAlgebra: norm, dot
 using Reexport: @reexport
 using LoopVectorization: @turbo
+using HDF5: HDF5, h5open, attributes, create_dataset, datatype, dataspace
 
 @reexport using StaticArrays: SVector, SMatrix
 
@@ -37,6 +38,7 @@ export flux_chandrashekar, flux_LMARS
 export velocity, waterheight, pressure, energy_total, energy_kinetic, energy_internal,
        lake_at_rest_error, source_terms_lagrange_multiplier,
        clean_solution_lagrange_multiplier!
+export cons2prim_plus_vorticity
 export P4estMeshCubedSphere2D, P4estMeshQuadIcosahedron2D, MetricTermsCrossProduct,
        MetricTermsInvariantCurl
 export EARTH_RADIUS, EARTH_GRAVITATIONAL_ACCELERATION,
diff --git a/src/callbacks_step/callbacks_step.jl b/src/callbacks_step/callbacks_step.jl
index 3f8aec6..10df9bb 100644
--- a/src/callbacks_step/callbacks_step.jl
+++ b/src/callbacks_step/callbacks_step.jl
@@ -1,3 +1,4 @@
 include("analysis_covariant.jl")
 include("save_solution_covariant.jl")
 include("stepsize_dg2d.jl")
+include("save_solution_2d_manifold_in_3d_cartesian.jl")
diff --git a/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
new file mode 100644
index 0000000..d65aaec
--- /dev/null
+++ b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
@@ -0,0 +1,197 @@
+@muladd begin
+#! format: noindent
+
+# 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::AbstractEquations{3}, dg, cache)
+    (; contravariant_vectors) = cache.elements
+    # Extract the number of solution variables to be output 
+    # (may be different than the number of conservative variables) 
+    n_vars = length(Trixi.varnames(solution_variables, 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)
+            
+            normal_vector_node = Trixi.get_contravariant_vector(3, contravariant_vectors, i, j, element)
+
+            data_node = solution_variables(u, normal_vector_node, equations, dg, cache, i, j, element)
+
+            for v in 1:n_vars
+                data[v, i, j, element] = data_node[v]
+            end
+        end
+    end
+    return data
+end
+
+# If no auxiliary variables are passed into the conversion to spherical coordinates, do not 
+# do any conversion.
+@inline cons2prim_plus_vorticity(u, equations) = u
+
+# # 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{2},
+#                                                                      <:AbstractEquations{3}},
+#                                   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)
+
+#     solution_callback2 = copy(solution_callback)
+#     solution_callback2
+#     # 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
+
+
+# Calculate the relative vorticity at a given node using the collocation derivative
+@inline function cons2prim_plus_vorticity(u, normal_vector, equations::ShallowWaterEquations3D,
+                                          dg::DGSEM, cache, i, j, element)
+    (; derivative_matrix) = dg.basis
+    (; contravariant_vectors, inverse_jacobian) = cache.elements
+
+    # Compute gradients in reference space
+    dv1dxi1 = dv1dxi2 = zero(eltype(u))
+    dv2dxi1 = dv2dxi2 = zero(eltype(u))
+    dv3dxi1 = dv3dxi2 = zero(eltype(u))
+    for ii in eachnode(dg)
+        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, ii, j, element)
+        dv1dxi1 += derivative_matrix[i, ii] * hv_1 / h
+        dv2dxi1 += derivative_matrix[i, ii] * hv_2 / h
+        dv3dxi1 += derivative_matrix[i, ii] * hv_3 / h
+    end
+
+    for jj in eachnode(dg)
+        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, i, jj, element)
+        dv1dxi2 += derivative_matrix[j, jj] * hv_1 / h
+        dv2dxi2 += derivative_matrix[j, jj] * hv_2 / h
+        dv3dxi2 += derivative_matrix[j, jj] * hv_3 / h
+    end
+
+    # Transform gradients to Cartesian space
+    Ja11, Ja12, Ja13 = Trixi.get_contravariant_vector(1, contravariant_vectors, i, j, element)
+    Ja21, Ja22, Ja23 = Trixi.get_contravariant_vector(2, contravariant_vectors, i, j, element)
+    
+    dv1dy = (Ja12 * dv1dxi1 + Ja22 * dv1dxi2) * inverse_jacobian[i, j, element]
+    dv1dz = (Ja13 * dv1dxi1 + Ja23 * dv1dxi2) * inverse_jacobian[i, j, element]
+    dv2dx = (Ja11 * dv2dxi1 + Ja21 * dv2dxi2) * inverse_jacobian[i, j, element]
+    dv2dz = (Ja13 * dv2dxi1 + Ja23 * dv2dxi2) * inverse_jacobian[i, j, element]
+    dv3dx = (Ja11 * dv3dxi1 + Ja21 * dv3dxi2) * inverse_jacobian[i, j, element]
+    dv3dy = (Ja12 * dv3dxi1 + Ja22 * dv3dxi2) * inverse_jacobian[i, j, element]
+
+    # compute the vorticity and project onto normal vector
+    vorticity = ((dv3dy - dv2dz) * normal_vector[1] + 
+                 (dv1dz - dv3dx) * normal_vector[2] + 
+                 (dv2dx - dv1dy) * normal_vector[3]) / norm(normal_vector)
+
+    u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
+
+    return SVector(cons2prim(u_node, equations)..., vorticity)
+end
+
+Trixi.varnames(::typeof(cons2prim_plus_vorticity), equations::ShallowWaterEquations3D) = (varnames(cons2prim, equations)..., "vorticity")
+
+function Trixi.save_solution_file(u, time, dt, timestep,
+                            mesh::P4estMesh{2},
+                            equations::AbstractEquations{3}, dg::DG, cache,
+                            solution_callback,
+                            element_variables = Dict{Symbol, Any}(),
+                            node_variables = Dict{Symbol, Any}();
+                            system = "")
+    @unpack 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
+        n_vars = nvariables(equations)
+    else
+        data = convert_variables(u, solution_variables, mesh, equations, dg, cache)
+        n_vars = size(data, 1)
+    # else
+    #     # Reinterpret the solution array as an array of conservative variables,
+    #     # compute the solution variables via broadcasting, and reinterpret the
+    #     # result as a plain array of floating point numbers
+    #     data = Array(reinterpret(eltype(u),
+    #                              solution_variables.(reinterpret(SVector{nvariables(equations),
+    #                                                                      eltype(u)}, u),
+    #                                                  Ref(equations))))
+
+    #     # Find out variable count by looking at output from `solution_variables` function
+    #     n_vars = size(data, 1)
+    end
+
+    # Open file (clobber existing content)
+    # TODO: Create a function to do this in Trixi.jl to avoid duplicated code.
+    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"] = 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
+end
\ No newline at end of file

From 123e198645ff715b7054a140ea8f5f9315950557 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Rueda-Ram=C3=ADrez?= <aruedara@uni-koeln.de>
Date: Tue, 3 Dec 2024 11:11:43 +0100
Subject: [PATCH 2/3] Reorganized code and made convert_variables compatible
 with functions with the signature solution_variables(u_node, equations)

---
 ...wwater_cubed_sphere_shell_EC_projection.jl |   2 +-
 ...wwater_quad_icosahedron_shell_advection.jl |   2 +-
 src/TrixiAtmo.jl                              |   2 +-
 ...ve_solution_2d_manifold_in_3d_cartesian.jl | 179 ++++++++----------
 4 files changed, 78 insertions(+), 107 deletions(-)

diff --git a/examples/elixir_shallowwater_cubed_sphere_shell_EC_projection.jl b/examples/elixir_shallowwater_cubed_sphere_shell_EC_projection.jl
index 21934c5..eace0c8 100644
--- a/examples/elixir_shallowwater_cubed_sphere_shell_EC_projection.jl
+++ b/examples/elixir_shallowwater_cubed_sphere_shell_EC_projection.jl
@@ -99,7 +99,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 = cons2prim)
+                                     solution_variables = cons2prim_and_vorticity)
 
 # The StepsizeCallback handles the re-calculation of the maximum Δt after each time step
 stepsize_callback = StepsizeCallback(cfl = 0.7)
diff --git a/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl b/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl
index 55fd161..e7a7593 100644
--- a/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl
+++ b/examples/elixir_shallowwater_quad_icosahedron_shell_advection.jl
@@ -70,7 +70,7 @@ analysis_callback = AnalysisCallback(semi, interval = 100,
 
 # The SaveSolutionCallback allows to save the solution to a file in regular intervals
 save_solution = SaveSolutionCallback(interval = 100,
-                                     solution_variables = cons2prim_plus_vorticity)
+                                     solution_variables = cons2prim_and_vorticity)
 
 # 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
index 7903000..ba96306 100644
--- a/src/TrixiAtmo.jl
+++ b/src/TrixiAtmo.jl
@@ -38,7 +38,7 @@ export flux_chandrashekar, flux_LMARS
 export velocity, waterheight, pressure, energy_total, energy_kinetic, energy_internal,
        lake_at_rest_error, source_terms_lagrange_multiplier,
        clean_solution_lagrange_multiplier!
-export cons2prim_plus_vorticity
+export cons2prim_and_vorticity
 export P4estMeshCubedSphere2D, P4estMeshQuadIcosahedron2D, MetricTermsCrossProduct,
        MetricTermsInvariantCurl
 export EARTH_RADIUS, EARTH_GRAVITATIONAL_ACCELERATION,
diff --git a/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
index d65aaec..441ee91 100644
--- a/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
+++ b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
@@ -1,8 +1,7 @@
 @muladd begin
 #! format: noindent
 
-# Convert to another set of variables using a solution_variables function that depends on 
-# auxiliary variables
+# Convert to another set of variables using a solution_variables function
 function convert_variables(u, solution_variables, mesh::P4estMesh{2},
                            equations::AbstractEquations{3}, dg, cache)
     (; contravariant_vectors) = cache.elements
@@ -16,10 +15,20 @@ function convert_variables(u, solution_variables, mesh::P4estMesh{2},
     # 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)
-            
-            normal_vector_node = Trixi.get_contravariant_vector(3, contravariant_vectors, i, j, element)
-
-            data_node = solution_variables(u, normal_vector_node, equations, dg, cache, i, j, element)
+            normal_vector_node = Trixi.get_contravariant_vector(3,
+                                                                contravariant_vectors,
+                                                                i, j, element)
+
+            if applicable(solution_variables, u, normal_vector_node, equations, dg,
+                          cache, i, j, element)
+                # The solution_variables function depends on the solution on other nodes, the normal_vector_node, etc.
+                data_node = solution_variables(u, normal_vector_node, equations, dg,
+                                               cache, i, j, element)
+            else
+                # The solution_variables function depends on u_node and equations
+                u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
+                data_node = solution_variables(u_node, equations)
+            end
 
             for v in 1:n_vars
                 data[v, i, j, element] = data_node[v]
@@ -29,93 +38,13 @@ function convert_variables(u, solution_variables, mesh::P4estMesh{2},
     return data
 end
 
-# If no auxiliary variables are passed into the conversion to spherical coordinates, do not 
-# do any conversion.
-@inline cons2prim_plus_vorticity(u, equations) = u
-
-# # 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{2},
-#                                                                      <:AbstractEquations{3}},
-#                                   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)
-
-#     solution_callback2 = copy(solution_callback)
-#     solution_callback2
-#     # 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
-
-
-# Calculate the relative vorticity at a given node using the collocation derivative
-@inline function cons2prim_plus_vorticity(u, normal_vector, equations::ShallowWaterEquations3D,
-                                          dg::DGSEM, cache, i, j, element)
-    (; derivative_matrix) = dg.basis
-    (; contravariant_vectors, inverse_jacobian) = cache.elements
-
-    # Compute gradients in reference space
-    dv1dxi1 = dv1dxi2 = zero(eltype(u))
-    dv2dxi1 = dv2dxi2 = zero(eltype(u))
-    dv3dxi1 = dv3dxi2 = zero(eltype(u))
-    for ii in eachnode(dg)
-        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, ii, j, element)
-        dv1dxi1 += derivative_matrix[i, ii] * hv_1 / h
-        dv2dxi1 += derivative_matrix[i, ii] * hv_2 / h
-        dv3dxi1 += derivative_matrix[i, ii] * hv_3 / h
-    end
-
-    for jj in eachnode(dg)
-        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, i, jj, element)
-        dv1dxi2 += derivative_matrix[j, jj] * hv_1 / h
-        dv2dxi2 += derivative_matrix[j, jj] * hv_2 / h
-        dv3dxi2 += derivative_matrix[j, jj] * hv_3 / h
-    end
-
-    # Transform gradients to Cartesian space
-    Ja11, Ja12, Ja13 = Trixi.get_contravariant_vector(1, contravariant_vectors, i, j, element)
-    Ja21, Ja22, Ja23 = Trixi.get_contravariant_vector(2, contravariant_vectors, i, j, element)
-    
-    dv1dy = (Ja12 * dv1dxi1 + Ja22 * dv1dxi2) * inverse_jacobian[i, j, element]
-    dv1dz = (Ja13 * dv1dxi1 + Ja23 * dv1dxi2) * inverse_jacobian[i, j, element]
-    dv2dx = (Ja11 * dv2dxi1 + Ja21 * dv2dxi2) * inverse_jacobian[i, j, element]
-    dv2dz = (Ja13 * dv2dxi1 + Ja23 * dv2dxi2) * inverse_jacobian[i, j, element]
-    dv3dx = (Ja11 * dv3dxi1 + Ja21 * dv3dxi2) * inverse_jacobian[i, j, element]
-    dv3dy = (Ja12 * dv3dxi1 + Ja22 * dv3dxi2) * inverse_jacobian[i, j, element]
-
-    # compute the vorticity and project onto normal vector
-    vorticity = ((dv3dy - dv2dz) * normal_vector[1] + 
-                 (dv1dz - dv3dx) * normal_vector[2] + 
-                 (dv2dx - dv1dy) * normal_vector[3]) / norm(normal_vector)
-
-    u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
-
-    return SVector(cons2prim(u_node, equations)..., vorticity)
-end
-
-Trixi.varnames(::typeof(cons2prim_plus_vorticity), equations::ShallowWaterEquations3D) = (varnames(cons2prim, equations)..., "vorticity")
-
 function Trixi.save_solution_file(u, time, dt, timestep,
-                            mesh::P4estMesh{2},
-                            equations::AbstractEquations{3}, dg::DG, cache,
-                            solution_callback,
-                            element_variables = Dict{Symbol, Any}(),
-                            node_variables = Dict{Symbol, Any}();
-                            system = "")
+                                  mesh::P4estMesh{2},
+                                  equations::AbstractEquations{3}, dg::DG, cache,
+                                  solution_callback,
+                                  element_variables = Dict{Symbol, Any}(),
+                                  node_variables = Dict{Symbol, Any}();
+                                  system = "")
     @unpack output_directory, solution_variables = solution_callback
 
     # Filename based on current time step
@@ -133,17 +62,6 @@ function Trixi.save_solution_file(u, time, dt, timestep,
     else
         data = convert_variables(u, solution_variables, mesh, equations, dg, cache)
         n_vars = size(data, 1)
-    # else
-    #     # Reinterpret the solution array as an array of conservative variables,
-    #     # compute the solution variables via broadcasting, and reinterpret the
-    #     # result as a plain array of floating point numbers
-    #     data = Array(reinterpret(eltype(u),
-    #                              solution_variables.(reinterpret(SVector{nvariables(equations),
-    #                                                                      eltype(u)}, u),
-    #                                                  Ref(equations))))
-
-    #     # Find out variable count by looking at output from `solution_variables` function
-    #     n_vars = size(data, 1)
     end
 
     # Open file (clobber existing content)
@@ -194,4 +112,57 @@ function Trixi.save_solution_file(u, time, dt, timestep,
 
     return filename
 end
-end
\ No newline at end of file
+
+# Calculate the primitive variables and the relative vorticity at a given node
+@inline function cons2prim_and_vorticity(u, normal_vector,
+                                         equations::ShallowWaterEquations3D,
+                                         dg::DGSEM, cache, i, j, element)
+    (; derivative_matrix) = dg.basis
+    (; contravariant_vectors, inverse_jacobian, node_coordinates) = cache.elements
+
+    # Compute gradients in reference space
+    dv1dxi1 = dv1dxi2 = zero(eltype(u))
+    dv2dxi1 = dv2dxi2 = zero(eltype(u))
+    dv3dxi1 = dv3dxi2 = zero(eltype(u))
+    for ii in eachnode(dg)
+        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, ii, j, element)
+        dv1dxi1 += derivative_matrix[i, ii] * hv_1 / h
+        dv2dxi1 += derivative_matrix[i, ii] * hv_2 / h
+        dv3dxi1 += derivative_matrix[i, ii] * hv_3 / h
+    end
+
+    for jj in eachnode(dg)
+        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, i, jj, element)
+        dv1dxi2 += derivative_matrix[j, jj] * hv_1 / h
+        dv2dxi2 += derivative_matrix[j, jj] * hv_2 / h
+        dv3dxi2 += derivative_matrix[j, jj] * hv_3 / h
+    end
+
+    # Transform gradients to Cartesian space
+    Ja11, Ja12, Ja13 = Trixi.get_contravariant_vector(1, contravariant_vectors, i, j,
+                                                      element)
+    Ja21, Ja22, Ja23 = Trixi.get_contravariant_vector(2, contravariant_vectors, i, j,
+                                                      element)
+
+    dv1dy = (Ja12 * dv1dxi1 + Ja22 * dv1dxi2) * inverse_jacobian[i, j, element]
+    dv1dz = (Ja13 * dv1dxi1 + Ja23 * dv1dxi2) * inverse_jacobian[i, j, element]
+    dv2dx = (Ja11 * dv2dxi1 + Ja21 * dv2dxi2) * inverse_jacobian[i, j, element]
+    dv2dz = (Ja13 * dv2dxi1 + Ja23 * dv2dxi2) * inverse_jacobian[i, j, element]
+    dv3dx = (Ja11 * dv3dxi1 + Ja21 * dv3dxi2) * inverse_jacobian[i, j, element]
+    dv3dy = (Ja12 * dv3dxi1 + Ja22 * dv3dxi2) * inverse_jacobian[i, j, element]
+
+    # compute the vorticity and project onto normal vector
+    vorticity = ((dv3dy - dv2dz) * normal_vector[1] +
+                 (dv1dz - dv3dx) * normal_vector[2] +
+                 (dv2dx - dv1dy) * normal_vector[3]) / norm(normal_vector)
+
+    u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
+
+    return SVector(cons2prim(u_node, equations)..., vorticity)
+end
+
+# Variable names for cons2prim_and_vorticity
+Trixi.varnames(::typeof(cons2prim_and_vorticity), equations::ShallowWaterEquations3D) = (varnames(cons2prim,
+                                                                                                  equations)...,
+                                                                                         "vorticity")
+end

From fcef9a2c0a072880d5af4be8e235a529da56ced6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Rueda-Ram=C3=ADrez?= <aruedara@uni-koeln.de>
Date: Tue, 3 Dec 2024 14:08:28 +0100
Subject: [PATCH 3/3] Added function for vorticity

---
 ...ve_solution_2d_manifold_in_3d_cartesian.jl | 53 ++++++++++++-------
 1 file changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
index 441ee91..4982812 100644
--- a/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
+++ b/src/callbacks_step/save_solution_2d_manifold_in_3d_cartesian.jl
@@ -19,10 +19,12 @@ function convert_variables(u, solution_variables, mesh::P4estMesh{2},
                                                                 contravariant_vectors,
                                                                 i, j, element)
 
-            if applicable(solution_variables, u, normal_vector_node, equations, dg,
+            if applicable(solution_variables, u, normal_vector_node, mesh, equations,
+                          dg,
                           cache, i, j, element)
                 # The solution_variables function depends on the solution on other nodes, the normal_vector_node, etc.
-                data_node = solution_variables(u, normal_vector_node, equations, dg,
+                data_node = solution_variables(u, normal_vector_node, mesh, equations,
+                                               dg,
                                                cache, i, j, element)
             else
                 # The solution_variables function depends on u_node and equations
@@ -115,27 +117,44 @@ end
 
 # Calculate the primitive variables and the relative vorticity at a given node
 @inline function cons2prim_and_vorticity(u, normal_vector,
-                                         equations::ShallowWaterEquations3D,
+                                         mesh::P4estMesh{2},
+                                         equations::AbstractEquations{3},
                                          dg::DGSEM, cache, i, j, element)
+
+    # compute the vorticity and project onto normal vector
+    vorticity = calc_vorticity_node(u, mesh, equations, dg, cache, i, j, element)
+    relative_vorticity = dot(vorticity, normal_vector) / norm(normal_vector)
+
+    u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
+
+    # Return th solution variables
+    return SVector(cons2prim(u_node, equations)..., relative_vorticity)
+end
+
+@inline function calc_vorticity_node(u, mesh::P4estMesh{2},
+                                     equations::AbstractEquations{3}, dg::DGSEM, cache,
+                                     i, j, element)
     (; derivative_matrix) = dg.basis
-    (; contravariant_vectors, inverse_jacobian, node_coordinates) = cache.elements
+    (; contravariant_vectors, inverse_jacobian) = cache.elements
 
     # Compute gradients in reference space
     dv1dxi1 = dv1dxi2 = zero(eltype(u))
     dv2dxi1 = dv2dxi2 = zero(eltype(u))
     dv3dxi1 = dv3dxi2 = zero(eltype(u))
     for ii in eachnode(dg)
-        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, ii, j, element)
-        dv1dxi1 += derivative_matrix[i, ii] * hv_1 / h
-        dv2dxi1 += derivative_matrix[i, ii] * hv_2 / h
-        dv3dxi1 += derivative_matrix[i, ii] * hv_3 / h
+        u_node = Trixi.get_node_vars(u, equations, dg, ii, j, element)
+        v1, v2, v3 = velocity(u_node, equations)
+        dv1dxi1 += derivative_matrix[i, ii] * v1
+        dv2dxi1 += derivative_matrix[i, ii] * v2
+        dv3dxi1 += derivative_matrix[i, ii] * v3
     end
 
     for jj in eachnode(dg)
-        h, hv_1, hv_2, hv_3, _ = Trixi.get_node_vars(u, equations, dg, i, jj, element)
-        dv1dxi2 += derivative_matrix[j, jj] * hv_1 / h
-        dv2dxi2 += derivative_matrix[j, jj] * hv_2 / h
-        dv3dxi2 += derivative_matrix[j, jj] * hv_3 / h
+        u_node = Trixi.get_node_vars(u, equations, dg, i, jj, element)
+        v1, v2, v3 = velocity(u_node, equations)
+        dv1dxi2 += derivative_matrix[j, jj] * v1
+        dv2dxi2 += derivative_matrix[j, jj] * v2
+        dv3dxi2 += derivative_matrix[j, jj] * v3
     end
 
     # Transform gradients to Cartesian space
@@ -151,14 +170,8 @@ end
     dv3dx = (Ja11 * dv3dxi1 + Ja21 * dv3dxi2) * inverse_jacobian[i, j, element]
     dv3dy = (Ja12 * dv3dxi1 + Ja22 * dv3dxi2) * inverse_jacobian[i, j, element]
 
-    # compute the vorticity and project onto normal vector
-    vorticity = ((dv3dy - dv2dz) * normal_vector[1] +
-                 (dv1dz - dv3dx) * normal_vector[2] +
-                 (dv2dx - dv1dy) * normal_vector[3]) / norm(normal_vector)
-
-    u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
-
-    return SVector(cons2prim(u_node, equations)..., vorticity)
+    # compute the vorticity
+    return SVector(dv3dy - dv2dz, dv1dz - dv3dx, dv2dx - dv1dy)
 end
 
 # Variable names for cons2prim_and_vorticity