Added element container with PtrArray for performance and modified th…

…e structure
trixi-framework · Jul 26, 2024 · eb768d9 · eb768d9
1 parent 24a3793
commit eb768d9
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 25 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "TrixiAtmo"
 uuid = "c9ed1054-d170-44a9-8ee2-d5566f5d1389"
-authors = ["Benedict Geihe <[email protected]>", "Tristan Montoya <[email protected]>", "Hendrik Ranocha <[email protected]>", "Michael Schlottke-Lakemper <[email protected]>"]
+authors = ["Benedict Geihe <[email protected]>", "Tristan Montoya <[email protected]>", "Hendrik Ranocha <[email protected]>", "Andrés Rueda-Ramírez <[email protected]>", "Michael Schlottke-Lakemper <[email protected]>"]
 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"
diff --git a/src/TrixiAtmo.jl b/src/TrixiAtmo.jl
@@ -12,6 +12,8 @@ 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

diff --git a/...lvers/dgsem_p4est_covariant/containers.jl → ...sem_p4est/containers_2d_manifold_in_3d.jl b/...lvers/dgsem_p4est_covariant/containers.jl → ...sem_p4est/containers_2d_manifold_in_3d.jl
@@ -1,10 +1,38 @@
 @muladd begin
 #! format: noindent
 
-# Create element container and initialize element data
-function Trixi.init_elements(mesh::Union{P4estMesh{NDIMS, RealT},
-                                         T8codeMesh{NDIMS, RealT}},
-                             equations::AbstractCovariantEquations2D,
+# 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
+
+# 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 {NDIMS, RealT <: Real,
                                                      uEltype <: Real}
@@ -30,10 +58,10 @@ function Trixi.init_elements(mesh::Union{P4estMesh{NDIMS, RealT},
     _contravariant_vectors = Vector{RealT}(undef,
                                            ndims_spa^2 * nnodes(basis)^NDIMS *
                                            nelements)
-    contravariant_vectors = Trixi.unsafe_wrap(Array, pointer(_contravariant_vectors),
-                                              (ndims_spa, ndims_spa,
-                                               ntuple(_ -> 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),
@@ -49,24 +77,27 @@ function Trixi.init_elements(mesh::Union{P4estMesh{NDIMS, RealT},
                                              ntuple(_ -> nnodes(basis), NDIMS - 1)...,
                                              NDIMS * 2, nelements))
 
-    elements = Trixi.P4estElementContainer{NDIMS, RealT, uEltype, NDIMS + 1, NDIMS + 2,
-                                           NDIMS + 3}(node_coordinates, jacobian_matrix,
-                                                      contravariant_vectors,
-                                                      inverse_jacobian,
-                                                      surface_flux_values,
-                                                      _node_coordinates,
-                                                      _jacobian_matrix,
-                                                      _contravariant_vectors,
-                                                      _inverse_jacobian,
-                                                      _surface_flux_values)
-
-    init_elements_covariant!(elements, mesh, basis)
+    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_covariant!(elements, mesh::Union{P4estMesh{2}, T8codeMesh{2}},
-                                  basis::LobattoLegendreBasis)
+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
 
@@ -118,7 +149,7 @@ function calc_node_coordinates_2d_shell!(node_coordinates,
     # places and the additional information passed to the compiler makes them faster
     # than native `Array`s.
     tmp1 = StrideArray(undef, real(mesh),
-                       StaticInt(size(mesh.tree_node_coordinates, 1)),
+                       Trixi.StaticInt(size(mesh.tree_node_coordinates, 1)),
                        Trixi.static_length(nodes), Trixi.static_length(mesh.nodes))
     matrix1 = StrideArray(undef, real(mesh),
                           Trixi.static_length(nodes), Trixi.static_length(mesh.nodes))

diff --git a/src/solvers/dgsem_p4est/dg.jl b/src/solvers/dgsem_p4est/dg.jl
@@ -0,0 +1,7 @@
+# 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
diff --git a/src/solvers/solvers.jl b/src/solvers/solvers.jl
@@ -1 +1,2 @@
-include("dgsem_p4est_covariant/containers.jl")
+include("dgsem_p4est/dg.jl")
+include("dgsem_p4est/containers_2d_manifold_in_3d.jl")