WIP: p4est solver on a spherical shell

examples/p4est_2d_dgsem/elixir_euler_cubed_sphere_shell.jl

@@ -0,0 +1,114 @@
+
+using OrdinaryDiffEq
+using Trixi
+
+###############################################################################
+# semidiscretization of the linear advection equation
+
+equations = CompressibleEulerEquations3D(1.4)
+
+# Create DG solver with polynomial degree = 3 and (local) Lax-Friedrichs/Rusanov flux as surface flux
+polydeg = 3
+solver = DGSEM(polydeg = polydeg, surface_flux = flux_lax_friedrichs)
+
+# Initial condition for a Gaussian density profile with constant pressure
+# and the velocity of a rotating solid body
+function initial_condition_advection_sphere(x, t, equations::CompressibleEulerEquations3D)
+    # Gaussian density
+    rho = 1.0 + exp(-20 * (x[1]^2 + x[3]^2))
+    # Constant pressure
+    p = 1.0
+
+    # Spherical coordinates for the point x
+    if sign(x[2]) == 0.0
+        signy = 1.0
+    else
+        signy = sign(x[2])
+    end
+    # Co-latitude
+    colat = acos(x[3] / sqrt(x[1]^2 + x[2]^2 + x[3]^2))
+    # Latitude (auxiliary variable)
+    lat = -colat + 0.5 * pi
+    # Longitude
+    r_xy = sqrt(x[1]^2 + x[2]^2)
+    if r_xy == 0.0
+        phi = pi / 2
+    else
+        phi = signy * acos(x[1] / r_xy)
+    end
+
+    # Compute the velocity of a rotating solid body
+    # (alpha is the angle between the rotation axis and the polar axis of the spherical coordinate system)
+    v0 = 1.0 # Velocity at the "equator"
+    alpha = 0.0 #pi / 4
+    v_long = v0 * (cos(lat) * cos(alpha) + sin(lat) * cos(phi) * sin(alpha))
+    v_lat = -v0 * sin(phi) * sin(alpha)
+
+    # Transform to Cartesian coordinate system
+    v1 = -cos(colat) * cos(phi) * v_lat - sin(phi) * v_long
+    v2 = -cos(colat) * sin(phi) * v_lat + cos(phi) * v_long
+    v3 = sin(colat) * v_lat
+
+    return prim2cons(SVector(rho, v1, v2, v3, p), equations)
+end
+
+# Source term function to transform the Euler equations into the linear advection equations with variable advection velocity
+function source_terms_convert_to_linear_advection(u, du, x, t,
+                                                  equations::CompressibleEulerEquations3D)
+    v1 = u[2] / u[1]
+    v2 = u[3] / u[1]
+    v3 = u[4] / u[1]
+
+    s2 = du[1] * v1 - du[2]
+    s3 = du[1] * v2 - du[3]
+    s4 = du[1] * v3 - du[4]
+    s5 = 0.5 * (s2 * v1 + s3 * v2 + s4 * v3) - du[5]
+
+    return SVector(0.0, s2, s3, s4, s5)
+end
+
+initial_condition = initial_condition_advection_sphere
+
+mesh = Trixi.P4estMeshCubedSphere2D(5, 1.0, polydeg = polydeg, initial_refinement_level = 0)
+
+# A semidiscretization collects data structures and functions for the spatial discretization
+semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
+                                    source_terms = source_terms_convert_to_linear_advection)
+
+###############################################################################
+# ODE solvers, callbacks etc.
+
+# Create ODE problem with time span from 0.0 to π
+tspan = (0.0, pi)
+ode = semidiscretize(semi, tspan)
+
+# At the beginning of the main loop, the SummaryCallback prints a summary of the simulation setup
+# and resets the timers
+summary_callback = SummaryCallback()
+
+# The AnalysisCallback allows to analyse the solution in regular intervals and prints the results
+analysis_callback = AnalysisCallback(semi, interval = 10,
+                                     save_analysis = true,
+                                     extra_analysis_integrals = (Trixi.density,))
+
+# The SaveSolutionCallback allows to save the solution to a file in regular intervals
+save_solution = SaveSolutionCallback(interval = 10,
+                                     solution_variables = cons2prim)
+
+# The StepsizeCallback handles the re-calculation of the maximum Δt after each time step
+stepsize_callback = StepsizeCallback(cfl = 0.7)
+
+# Create a CallbackSet to collect all callbacks such that they can be passed to the ODE solver
+callbacks = CallbackSet(summary_callback, analysis_callback, save_solution,
+                        stepsize_callback)
+
+###############################################################################
+# run the simulation
+
+# OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed callbacks
+sol = solve(ode, CarpenterKennedy2N54(williamson_condition = false),
+            dt = 1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
+            save_everystep = false, callback = callbacks);
+
+# Print the timer summary
+summary_callback()

src/semidiscretization/semidiscretization_hyperbolic.jl

@@ -40,7 +40,7 @@ struct SemidiscretizationHyperbolic{Mesh, Equations, InitialCondition,
                                                                  BoundaryConditions,
                                                                  SourceTerms, Solver,
                                                                  Cache}
-        @assert ndims(mesh) == ndims(equations)
+        #@assert ndims(mesh) == ndims(equations)
 
         performance_counter = PerformanceCounter()
 

src/solvers/dgsem_p4est/containers.jl

@@ -87,19 +87,29 @@ function init_elements(mesh::Union{P4estMesh{NDIMS, RealT}, T8codeMesh{NDIMS, Re
                        ::Type{uEltype}) where {NDIMS, RealT <: Real, uEltype <: Real}
     nelements = ncells(mesh)
 
-    _node_coordinates = Vector{RealT}(undef, NDIMS * nnodes(basis)^NDIMS * nelements)
+    ndims_spa = size(mesh.tree_node_coordinates, 1)
+
+    _node_coordinates = Vector{RealT}(undef,
+                                      ndims_spa * nnodes(basis)^NDIMS * nelements)
     node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
-                                   (NDIMS, ntuple(_ -> nnodes(basis), NDIMS)...,
+                                   (ndims_spa, ntuple(_ -> nnodes(basis), NDIMS)...,
                                     nelements))
 
-    _jacobian_matrix = Vector{RealT}(undef, NDIMS^2 * nnodes(basis)^NDIMS * nelements)
+    _jacobian_matrix = Vector{RealT}(undef,
+                                     ndims_spa * NDIMS * nnodes(basis)^NDIMS *
+                                     nelements)
     jacobian_matrix = unsafe_wrap(Array, pointer(_jacobian_matrix),
-                                  (NDIMS, NDIMS, ntuple(_ -> nnodes(basis), NDIMS)...,
+                                  (ndims_spa, NDIMS,
+                                   ntuple(_ -> nnodes(basis), NDIMS)...,
                                    nelements))
 
-    _contravariant_vectors = similar(_jacobian_matrix)
+    _contravariant_vectors = Vector{RealT}(undef,
+                                           ndims_spa^2 * nnodes(basis)^NDIMS *
+                                           nelements)
     contravariant_vectors = unsafe_wrap(Array, pointer(_contravariant_vectors),
-                                        size(jacobian_matrix))
+                                        (ndims_spa, ndims_spa,
+                                         ntuple(_ -> nnodes(basis), NDIMS)...,
+                                         nelements))
 
     _inverse_jacobian = Vector{RealT}(undef, nnodes(basis)^NDIMS * nelements)
     inverse_jacobian = unsafe_wrap(Array, pointer(_inverse_jacobian),

src/solvers/dgsem_p4est/containers_2d.jl

@@ -13,12 +13,32 @@ function init_elements!(elements, mesh::Union{P4estMesh{2}, T8codeMesh{2}},
 
     calc_node_coordinates!(node_coordinates, mesh, basis)
 
-    for element in 1:ncells(mesh)
-        calc_jacobian_matrix!(jacobian_matrix, element, node_coordinates, basis)
+    if size(node_coordinates, 1) == 3
+        # The mesh is a spherical shell
+        for element in 1:ncells(mesh)
+            # Compute Jacobian matrix as usual
+            calc_jacobian_matrix!(jacobian_matrix, element, node_coordinates, basis)
+            # Compute contravariant vectors with Giraldo's formula (cross-product form)
+            calc_contravariant_vectors_cubed_sphere!(contravariant_vectors, element,
+                                                     jacobian_matrix, node_coordinates,
+                                                     basis)
+            # Compute the inverse Jacobian as the norm of the cross product of the covariant vectors
+            for j in eachnode(basis), i in eachnode(basis)
+                inverse_jacobian[i, j, element] = 1 /
+                                                  norm(cross(jacobian_matrix[:, 1, i, j,
+                                                                             element],
+                                                             jacobian_matrix[:, 2, i, j,
+                                                                             element]))
+            end
+        end
+    else
+        for element in 1:ncells(mesh)
+            calc_jacobian_matrix!(jacobian_matrix, element, node_coordinates, basis)
 
-        calc_contravariant_vectors!(contravariant_vectors, element, jacobian_matrix)
+            calc_contravariant_vectors!(contravariant_vectors, element, jacobian_matrix)
 
-        calc_inverse_jacobian!(inverse_jacobian, element, jacobian_matrix)
+            calc_inverse_jacobian!(inverse_jacobian, element, jacobian_matrix)
+        end
     end
 
     return nothing
@@ -43,7 +63,8 @@ function calc_node_coordinates!(node_coordinates,
     # places and the additional information passed to the compiler makes them faster
     # than native `Array`s.
     tmp1 = StrideArray(undef, real(mesh),
-                       StaticInt(2), static_length(nodes), static_length(mesh.nodes))
+                       StaticInt(size(mesh.tree_node_coordinates, 1)),
+                       static_length(nodes), static_length(mesh.nodes))
     matrix1 = StrideArray(undef, real(mesh),
                           static_length(nodes), static_length(mesh.nodes))
     matrix2 = similar(matrix1)
@@ -84,6 +105,63 @@ function calc_node_coordinates!(node_coordinates,
     return node_coordinates
 end
 
+# Calculate contravariant vectors, multiplied by the Jacobian determinant J of the transformation mapping,
+# using eq (12) of :
+#   Giraldo, F. X. (2001). A spectral element shallow water model on spherical geodesic grids. 
+#   International Journal for Numerical Methods in Fluids, 35(8), 869-901. 
+#   https://doi.org/10.1002/1097-0363(20010430)35:8<869::AID-FLD116>3.0.CO;2-S
+# This is nothing but the cross-product form, but we end up with three contravariant vectors
+# because there are three space dimensions.
+function calc_contravariant_vectors_cubed_sphere!(contravariant_vectors::AbstractArray{
+                                                                                       <:Any,
+                                                                                       5
+                                                                                       },
+                                                  element,
+                                                  jacobian_matrix, node_coordinates,
+                                                  basis::LobattoLegendreBasis)
+    @unpack derivative_matrix = basis
+
+    for j in eachnode(basis), i in eachnode(basis)
+        for n in 1:3
+            # (n, m, l) cyclic
+            m = (n % 3) + 1
+            l = ((n + 1) % 3) + 1
+
+            contravariant_vectors[n, 1, i, j, element] = (jacobian_matrix[m, 2, i, j,
+                                                                          element] *
+                                                          node_coordinates[l, i, j,
+                                                                           element]
+                                                          -
+                                                          jacobian_matrix[l, 2, i, j,
+                                                                          element] *
+                                                          node_coordinates[m, i, j,
+                                                                           element])
+
+            contravariant_vectors[n, 2, i, j, element] = (jacobian_matrix[l, 1, i, j,
+                                                                          element] *
+                                                          node_coordinates[m, i, j,
+                                                                           element]
+                                                          -
+                                                          jacobian_matrix[m, 1, i, j,
+                                                                          element] *
+                                                          node_coordinates[l, i, j,
+                                                                           element])
+
+            contravariant_vectors[n, 3, i, j, element] = (jacobian_matrix[m, 1, i, j,
+                                                                          element] *
+                                                          jacobian_matrix[l, 2, i, j,
+                                                                          element]
+                                                          -
+                                                          jacobian_matrix[m, 2, i, j,
+                                                                          element] *
+                                                          jacobian_matrix[l, 1, i, j,
+                                                                          element])
+        end
+    end
+
+    return contravariant_vectors
+end
+
 # Initialize node_indices of interface container
 @inline function init_interface_node_indices!(interfaces::P4estInterfaceContainer{2},
                                               faces, orientation, interface_id)

src/solvers/dgsem_structured/dg.jl

@@ -22,9 +22,10 @@ function create_cache(mesh::StructuredMesh, equations::AbstractEquations, dg::DG
 end
 
 # Extract contravariant vector Ja^i (i = index) as SVector
+# TODO: Here, size() causes a lot of allocations! Find an alternative to improve performance
 @inline function get_contravariant_vector(index, contravariant_vectors, indices...)
     SVector(ntuple(@inline(dim->contravariant_vectors[dim, index, indices...]),
-                   Val(ndims(contravariant_vectors) - 3)))
+                   Val(size(contravariant_vectors, 1))))
 end
 
 @inline function calc_boundary_flux_by_direction!(surface_flux_values, u, t,

src/solvers/dgsem_structured/dg_2d.jl

@@ -67,30 +67,78 @@ See also https://github.com/trixi-framework/Trixi.jl/issues/1671#issuecomment-17
     @unpack derivative_dhat = dg.basis
     @unpack contravariant_vectors = cache.elements
 
-    for j in eachnode(dg), i in eachnode(dg)
-        u_node = get_node_vars(u, equations, dg, i, j, element)
+    if size(contravariant_vectors, 1) == 2
+        for j in eachnode(dg), i in eachnode(dg)
+            u_node = get_node_vars(u, equations, dg, i, j, element)
 
-        flux1 = flux(u_node, 1, equations)
-        flux2 = flux(u_node, 2, equations)
+            flux1 = flux(u_node, 1, equations)
+            flux2 = flux(u_node, 2, equations)
+            # Compute the contravariant flux by taking the scalar product of the
+            # first contravariant vector Ja^1 and the flux vector
+            #Ja11, Ja12 = get_contravariant_vector(1, contravariant_vectors, i, j, element)
+            Ja11 = contravariant_vectors[1, 1, i, j, element]
+            Ja12 = contravariant_vectors[2, 1, i, j, element]
+            contravariant_flux1 = Ja11 * flux1 + Ja12 * flux2
+            for ii in eachnode(dg)
+                multiply_add_to_node_vars!(du, alpha * derivative_dhat[ii, i],
+                                           contravariant_flux1, equations, dg, ii, j,
+                                           element)
+            end
 
-        # Compute the contravariant flux by taking the scalar product of the
-        # first contravariant vector Ja^1 and the flux vector
-        Ja11, Ja12 = get_contravariant_vector(1, contravariant_vectors, i, j, element)
-        contravariant_flux1 = Ja11 * flux1 + Ja12 * flux2
-        for ii in eachnode(dg)
-            multiply_add_to_node_vars!(du, alpha * derivative_dhat[ii, i],
-                                       contravariant_flux1, equations, dg, ii, j,
-                                       element)
+            # Compute the contravariant flux by taking the scalar product of the
+            # second contravariant vector Ja^2 and the flux vector
+            #Ja21, Ja22 = get_contravariant_vector(2, contravariant_vectors, i, j, element)
+            Ja21 = contravariant_vectors[1, 2, i, j, element]
+            Ja22 = contravariant_vectors[2, 2, i, j, element]
+            contravariant_flux2 = Ja21 * flux1 + Ja22 * flux2
+            for jj in eachnode(dg)
+                multiply_add_to_node_vars!(du, alpha * derivative_dhat[jj, j],
+                                           contravariant_flux2, equations, dg, i, jj,
+                                           element)
+            end
         end
+    else #size(contravariant_vectors,1) == 3
+        for j in eachnode(dg), i in eachnode(dg)
+            u_node = get_node_vars(u, equations, dg, i, j, element)
 
-        # Compute the contravariant flux by taking the scalar product of the
-        # second contravariant vector Ja^2 and the flux vector
-        Ja21, Ja22 = get_contravariant_vector(2, contravariant_vectors, i, j, element)
-        contravariant_flux2 = Ja21 * flux1 + Ja22 * flux2
-        for jj in eachnode(dg)
-            multiply_add_to_node_vars!(du, alpha * derivative_dhat[jj, j],
-                                       contravariant_flux2, equations, dg, i, jj,
-                                       element)
+            flux1 = flux(u_node, 1, equations)
+            flux2 = flux(u_node, 2, equations)
+            flux3 = flux(u_node, 3, equations)
+
+            # Compute the contravariant flux by taking the scalar product of the
+            # first contravariant vector Ja^1 and the flux vector
+            #Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors, i, j, element)
+            Ja11 = contravariant_vectors[1, 1, i, j, element]
+            Ja12 = contravariant_vectors[2, 1, i, j, element]
+            Ja13 = contravariant_vectors[3, 1, i, j, element]
+            contravariant_flux1 = Ja11 * flux1 + Ja12 * flux2 + Ja13 * flux3
+            for ii in eachnode(dg)
+                multiply_add_to_node_vars!(du, alpha * derivative_dhat[ii, i],
+                                           contravariant_flux1, equations, dg, ii, j,
+                                           element)
+            end
+
+            # Compute the contravariant flux by taking the scalar product of the
+            # second contravariant vector Ja^2 and the flux vector
+            #Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors, i, j, element)
+            Ja21 = contravariant_vectors[1, 2, i, j, element]
+            Ja22 = contravariant_vectors[2, 2, i, j, element]
+            Ja23 = contravariant_vectors[3, 2, i, j, element]
+            contravariant_flux2 = Ja21 * flux1 + Ja22 * flux2 + Ja23 * flux3
+            for jj in eachnode(dg)
+                multiply_add_to_node_vars!(du, alpha * derivative_dhat[jj, j],
+                                           contravariant_flux2, equations, dg, i, jj,
+                                           element)
+            end
+
+            #Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors, i, j, element)
+            Ja31 = contravariant_vectors[1, 3, i, j, element]
+            Ja32 = contravariant_vectors[2, 3, i, j, element]
+            Ja33 = contravariant_vectors[3, 3, i, j, element]
+            for v in eachvariable(equations)
+                du[v, i, j, element] += (Ja31 * flux1[v] + Ja32 * flux2[v] +
+                                         Ja33 * flux3[v])
+            end
         end
     end