p4est solver on a spherical mesh

Project.toml

@@ -1,6 +1,6 @@
 name = "TrixiAtmo"
 uuid = "c9ed1054-d170-44a9-8ee2-d5566f5d1389"
-authors = ["Benedict Geihe <bgeihe@uni-koeln.de>", "Tristan Montoya <montoya.tristan@gmail.com>", "Hendrik Ranocha <hendrik.ranocha@uni-mainz.de>", "Michael Schlottke-Lakemper <michael@sloede.com>"]
+authors = ["Benedict Geihe <bgeihe@uni-koeln.de>", "Tristan Montoya <montoya.tristan@gmail.com>", "Hendrik Ranocha <hendrik.ranocha@uni-mainz.de>", "Andrés Rueda-Ramírez <aruedara@uni-koeln.de>", "Michael Schlottke-Lakemper <michael@sloede.com>"]
 version = "0.1.0-DEV"
 
 [deps]
@@ -9,6 +9,8 @@ MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
+StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
 Trixi = "a7f1ee26-1774-49b1-8366-f1abc58fbfcb"
 
 [compat]
@@ -17,5 +19,7 @@ MuladdMacro = "0.2.2"
 Reexport = "1.0"
 Static = "0.8, 1"
 StaticArrays = "1"
+StaticArrayInterface = "1.5.1"
+StrideArrays = "0.1.28"
 Trixi = "0.7, 0.8"
 julia = "1.9"

src/TrixiAtmo.jl

@@ -12,12 +12,16 @@ using Trixi
 using MuladdMacro: @muladd
 using StaticArrays: SVector
 using Static: True, False
+using StrideArrays: PtrArray
+using StaticArrayInterface: static_size
 using LinearAlgebra: norm
 using Reexport: @reexport
 @reexport using StaticArrays: SVector
 
 include("auxiliary/auxiliary.jl")
 include("equations/equations.jl")
+include("meshes/meshes.jl")
+include("solvers/solvers.jl")
 
 export CompressibleMoistEulerEquations2D
 

src/meshes/meshes.jl

@@ -0,0 +1,2 @@
+export P4estMeshCubedSphere2D
+include("p4est_cubed_sphere.jl")

src/solvers/dgsem_p4est/containers_2d_manifold_in_3d.jl

@@ -0,0 +1,258 @@
+@muladd begin
+#! format: noindent
+
+# New p4est element container that allows the use of a PtrArray for the contravariant_vectors
+mutable struct P4estElementContainerPtrArray{NDIMS, RealT <: Real, uEltype <: Real,
+                                             NDIMSP1,
+                                             NDIMSP2, NDIMSP3,
+                                             ContravariantVectors <:
+                                             AbstractArray{RealT, NDIMSP3}} <:
+               Trixi.AbstractContainer
+    # Physical coordinates at each node
+    node_coordinates::Array{RealT, NDIMSP2}   # [orientation, node_i, node_j, node_k, element]
+    # Jacobian matrix of the transformation
+    # [jacobian_i, jacobian_j, node_i, node_j, node_k, element] where jacobian_i is the first index of the Jacobian matrix,...
+    jacobian_matrix::Array{RealT, NDIMSP3}
+    # Contravariant vectors, scaled by J, in Kopriva's blue book called Ja^i_n (i index, n dimension)
+    contravariant_vectors::ContravariantVectors  # [dimension, index, node_i, node_j, node_k, element]
+    # 1/J where J is the Jacobian determinant (determinant of Jacobian matrix)
+    inverse_jacobian::Array{RealT, NDIMSP1}   # [node_i, node_j, node_k, element]
+    # Buffer for calculated surface flux
+    surface_flux_values::Array{uEltype, NDIMSP2} # [variable, i, j, direction, element]
+
+    # internal `resize!`able storage
+    _node_coordinates::Vector{RealT}
+    _jacobian_matrix::Vector{RealT}
+    _contravariant_vectors::Vector{RealT}
+    _inverse_jacobian::Vector{RealT}
+    _surface_flux_values::Vector{uEltype}
+end
+
+@inline function nelements(elements::P4estElementContainerPtrArray)
+    size(elements.node_coordinates, ndims(elements) + 2)
+end
+@inline Base.ndims(::P4estElementContainerPtrArray{NDIMS}) where {NDIMS} = NDIMS
+@inline function Base.eltype(::P4estElementContainerPtrArray{NDIMS, RealT, uEltype}) where {
+                                                                                            NDIMS,
+                                                                                            RealT,
+                                                                                            uEltype
+                                                                                            }
+    uEltype
+end
+
+# Create element container and initialize element data.
+# This function dispatches on the dimensions of the mesh and the equation (AbstractEquations{3})
+function Trixi.init_elements(mesh::Union{P4estMesh{2, RealT},
+                                         T8codeMesh{2, RealT}},
+                             equations::AbstractEquations{3},
+                             basis,
+                             ::Type{uEltype}) where {RealT <: Real,
+                                                     uEltype <: Real}
+    nelements = Trixi.ncells(mesh)
+
+    NDIMS = 2 #Dimension of the manifold
+    ndims_spa = size(mesh.tree_node_coordinates, 1)
+
+    _node_coordinates = Vector{RealT}(undef,
+                                      ndims_spa * nnodes(basis)^NDIMS * nelements)
+    node_coordinates = Trixi.unsafe_wrap(Array, pointer(_node_coordinates),
+                                         (ndims_spa,
+                                          ntuple(_ -> nnodes(basis), NDIMS)...,
+                                          nelements))
+
+    _jacobian_matrix = Vector{RealT}(undef,
+                                     ndims_spa * NDIMS * nnodes(basis)^NDIMS *
+                                     nelements)
+    jacobian_matrix = Trixi.unsafe_wrap(Array, pointer(_jacobian_matrix),
+                                        (ndims_spa, NDIMS,
+                                         ntuple(_ -> nnodes(basis), NDIMS)...,
+                                         nelements))
+
+    _contravariant_vectors = Vector{RealT}(undef,
+                                           ndims_spa^2 * nnodes(basis)^NDIMS *
+                                           nelements)
+    contravariant_vectors = PtrArray(pointer(_contravariant_vectors),
+                                     (Trixi.StaticInt(ndims_spa), ndims_spa,
+                                      ntuple(_ -> nnodes(basis), NDIMS)...,
+                                      nelements))
+
+    _inverse_jacobian = Vector{RealT}(undef, nnodes(basis)^NDIMS * nelements)
+    inverse_jacobian = Trixi.unsafe_wrap(Array, pointer(_inverse_jacobian),
+                                         (ntuple(_ -> nnodes(basis), NDIMS)...,
+                                          nelements))
+
+    _surface_flux_values = Vector{uEltype}(undef,
+                                           nvariables(equations) *
+                                           nnodes(basis)^(NDIMS - 1) * (NDIMS * 2) *
+                                           nelements)
+    surface_flux_values = Trixi.unsafe_wrap(Array, pointer(_surface_flux_values),
+                                            (nvariables(equations),
+                                             ntuple(_ -> nnodes(basis), NDIMS - 1)...,
+                                             NDIMS * 2, nelements))
+
+    elements = P4estElementContainerPtrArray{NDIMS, RealT, uEltype, NDIMS + 1,
+                                             NDIMS + 2,
+                                             NDIMS + 3, typeof(contravariant_vectors)}(node_coordinates,
+                                                                                       jacobian_matrix,
+                                                                                       contravariant_vectors,
+                                                                                       inverse_jacobian,
+                                                                                       surface_flux_values,
+                                                                                       _node_coordinates,
+                                                                                       _jacobian_matrix,
+                                                                                       _contravariant_vectors,
+                                                                                       _inverse_jacobian,
+                                                                                       _surface_flux_values)
+
+    init_elements_2d_manifold_in_3d!(elements, mesh, basis)
+
+    return elements
+end
+
+function init_elements_2d_manifold_in_3d!(elements,
+                                          mesh::Union{P4estMesh{2}, T8codeMesh{2}},
+                                          basis::LobattoLegendreBasis)
+    (; node_coordinates, jacobian_matrix,
+    contravariant_vectors, inverse_jacobian) = elements
+
+    calc_node_coordinates_2d_shell!(node_coordinates, mesh, basis)
+
+    for element in 1:Trixi.ncells(mesh)
+
+        # Compute Jacobian matrix as usual
+        Trixi.calc_jacobian_matrix!(jacobian_matrix, element, node_coordinates, basis)
+
+        # Compute contravariant vectors
+        calc_contravariant_vectors_2d_shell!(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(Trixi.cross(jacobian_matrix[:, 1, i,
+                                                                               j,
+                                                                               element],
+                                                               jacobian_matrix[:, 2, i,
+                                                                               j,
+                                                                               element]))
+        end
+    end
+
+    return nothing
+end
+
+# Interpolate tree_node_coordinates to each quadrant at the nodes of the specified basis
+function calc_node_coordinates_2d_shell!(node_coordinates,
+                                         mesh::Union{P4estMesh{2}, T8codeMesh{2}},
+                                         basis::LobattoLegendreBasis)
+    # Hanging nodes will cause holes in the mesh if its polydeg is higher
+    # than the polydeg of the solver.
+    @assert length(basis.nodes)>=length(mesh.nodes) "The solver can't have a lower polydeg than the mesh"
+
+    calc_node_coordinates_2d_shell!(node_coordinates, mesh, basis.nodes)
+end
+
+# Interpolate tree_node_coordinates to each quadrant at the specified nodes
+function calc_node_coordinates_2d_shell!(node_coordinates,
+                                         mesh::P4estMesh{2},
+                                         nodes::AbstractVector)
+    # We use `StrideArray`s here since these buffers are used in performance-critical
+    # places and the additional information passed to the compiler makes them faster
+    # than native `Array`s.
+    tmp1 = Trixi.StrideArray(undef, real(mesh),
+                             Trixi.StaticInt(size(mesh.tree_node_coordinates, 1)),
+                             Trixi.static_length(nodes),
+                             Trixi.static_length(mesh.nodes))
+    matrix1 = Trixi.StrideArray(undef, real(mesh),
+                                Trixi.static_length(nodes),
+                                Trixi.static_length(mesh.nodes))
+    matrix2 = similar(matrix1)
+    baryweights_in = Trixi.barycentric_weights(mesh.nodes)
+
+    # Macros from `p4est`
+    p4est_root_len = 1 << Trixi.P4EST_MAXLEVEL
+    p4est_root_len = 1 << Trixi.P4EST_MAXLEVEL
+    p4est_quadrant_len(l) = 1 << (Trixi.P4EST_MAXLEVEL - l)
+
+    trees = Trixi.unsafe_wrap_sc(Trixi.p4est_tree_t, mesh.p4est.trees)
+
+    for tree in eachindex(trees)
+        offset = trees[tree].quadrants_offset
+        quadrants = Trixi.unsafe_wrap_sc(Trixi.p4est_quadrant_t, trees[tree].quadrants)
+
+        for i in eachindex(quadrants)
+            element = offset + i
+            quad = quadrants[i]
+
+            quad_length = p4est_quadrant_len(quad.level) / p4est_root_len
+
+            nodes_out_x = 2 * (quad_length * 1 / 2 * (nodes .+ 1) .+
+                           quad.x / p4est_root_len) .- 1
+            nodes_out_y = 2 * (quad_length * 1 / 2 * (nodes .+ 1) .+
+                           quad.y / p4est_root_len) .- 1
+            Trixi.polynomial_interpolation_matrix!(matrix1, mesh.nodes, nodes_out_x,
+                                                   baryweights_in)
+            Trixi.polynomial_interpolation_matrix!(matrix2, mesh.nodes, nodes_out_y,
+                                                   baryweights_in)
+
+            Trixi.multiply_dimensionwise!(view(node_coordinates, :, :, :, element),
+                                          matrix1, matrix2,
+                                          view(mesh.tree_node_coordinates, :, :, :,
+                                               tree),
+                                          tmp1)
+        end
+    end
+
+    return node_coordinates
+end
+
+function calc_contravariant_vectors_2d_shell!(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
+end # @muladd

src/solvers/dgsem_p4est/dg.jl

@@ -0,0 +1,8 @@
+# Extract contravariant vector Ja^i (i = index) as SVector
+# This function dispatches on the type of contravariant_vectors
+static2val(::Trixi.StaticInt{N}) where {N} = Val{N}()
+@inline function get_contravariant_vector(index, contravariant_vectors::PtrArray,
+                                          indices...)
+    SVector(ntuple(@inline(dim->contravariant_vectors[dim, index, indices...]),
+                   static2val(static_size(contravariant_vectors, Trixi.StaticInt(1)))))
+end

src/solvers/solvers.jl

@@ -0,0 +1,2 @@
+include("dgsem_p4est/dg.jl")
+include("dgsem_p4est/containers_2d_manifold_in_3d.jl")