switch to SArrays

Author: patrickersing

src/TrixiShallowWater.jl

@@ -8,7 +8,7 @@ using Trixi
 # Import additional symbols that are not exported by Trixi.jl
 using Trixi: get_node_vars, set_node_vars!
 using MuladdMacro: @muladd
-using StaticArrays: SVector, MVector, MArray, SMatrix, @SMatrix
+using StaticArrays: SVector, MVector, MArray, SMatrix, @SMatrix, SArray
 using Static: True, False
 using LinearAlgebra: norm
 using Roots: Order2, solve, ZeroProblem

src/equations/moment_matrices.jl

@@ -49,20 +49,19 @@ function compute_B_tensor(n::Int, RealT = Float64)
         B[i, j, k] = 0.5f0 * (2 * i + 1) * res
     end # i, j, k
 
-    return Array{RealT, 3}(B)
+    return SArray{Tuple{n, n, n}, RealT, 3, n * n * n}(B)
 end
 
 # Directly compute the A tensor from B using the relation
 #  A_ijk / (2i + 1) = - B_ijk / (2i + 1) - B_kji / (2k + 1)
-function compute_A_tensor(B::Array{RealT, 3}) where {RealT}
-    N = size(B, 1) # Get number of moments
+function compute_A_tensor(B::SArray{Tuple{N, N, N}, RealT, 3}) where {RealT <: Real, N}
     A = Array{RealT, 3}(zeros(RealT, N, N, N))  # initialize the A tensor
 
     for i in 1:N, j in 1:N, k in 1:N
         A[i, j, k] = -(2 * i + 1) * (B[i, j, k] / (2 * i + 1) + B[k, j, i] / (2 * k + 1))
     end # i, j, k
 
-    return A
+    return SArray{Tuple{N, N, N}, RealT, 3, N * N * N}(A)
 end
 
 # Helper function to precompute and store every instance of the interior integral
@@ -152,7 +151,7 @@ function compute_C_matrix(n::Int, RealT = Float64)
         C[i, j] = 0.5f0 * (2 * i + 1) * res
     end # i, j
 
-    return C
+    return SArray{Tuple{n, n}, RealT, 2, n * n}(C)
 end
 
 # The tensor B and matrix C are built from the shifted and unnormalized Legendre polynomials.

src/equations/shallow_water_moments_1d.jl

@@ -41,36 +41,44 @@ For details see the paper:
   Entropy analysis and entropy stable DG methods for the shallow water moment equations
   [DOI: 10.48550/arXiv.2602.06513](https://doi.org/10.48550/arXiv.2602.06513)
 """
-struct ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT <: Real} <:
+struct ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT <: Real,
+                                     Array3 <: AbstractArray{RealT, 3},
+                                     Array2 <: AbstractArray{RealT, 2}} <:
        AbstractMomentEquations{1, NVARS, NMOMENTS}
     gravity::RealT   # gravitational acceleration
     H0::RealT        # constant "lake-at-rest" total water height
     n_moments::Integer  # number of moments
     # Moment matrices
-    A::Array{RealT, 3}
-    B::Array{RealT, 3}
-    C::Array{RealT, 2}
+    A::Array3
+    B::Array3
+    C::Array2
     # Friction related quantities
     nu::RealT       # kinematic viscosity
     lambda::RealT  # slip length
     rho::RealT    # fluid density (relevant for Manning friction)
     nman::RealT   # Manning roughness coefficient
 
-    function ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT}(gravity::RealT,
-                                                                   H0::RealT,
-                                                                   n_moments::Integer,
-                                                                   A::Array{RealT, 3},
-                                                                   B::Array{RealT, 3},
-                                                                   C::Array{RealT, 2},
-                                                                   nu::RealT,
-                                                                   lambda::RealT,
-                                                                   rho::RealT,
-                                                                   nman::RealT) where {
-                                                                                       NVARS,
-                                                                                       NMOMENTS,
-                                                                                       RealT <:
-                                                                                       Real
-                                                                                       }
+    function ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT, Array3, Array2}(gravity::RealT,
+                                                                                   H0::RealT,
+                                                                                   n_moments::Integer,
+                                                                                   A::Array3,
+                                                                                   B::Array3,
+                                                                                   C::Array2,
+                                                                                   nu::RealT,
+                                                                                   lambda::RealT,
+                                                                                   rho::RealT,
+                                                                                   nman::RealT) where {
+                                                                                                       NVARS,
+                                                                                                       NMOMENTS,
+                                                                                                       RealT <:
+                                                                                                       Real,
+                                                                                                       Array3 <:
+                                                                                                       AbstractArray{RealT,
+                                                                                                                     3},
+                                                                                                       Array2 <:
+                                                                                                       AbstractArray{RealT,
+                                                                                                                     2}
+                                                                                                       }
         new(gravity, H0, n_moments, A, B, C, nu, lambda, rho, nman)
     end
 end
@@ -93,24 +101,33 @@ function ShallowWaterMomentEquations1D(; gravity, H0 = zero(gravity), n_moments,
     A = compute_A_tensor(B)
     C = compute_C_matrix(n_moments)
 
-    return ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT}(gravity,
-                                                                 H0,
-                                                                 n_moments,
-                                                                 A,
-                                                                 B,
-                                                                 C,
-                                                                 nu,
-                                                                 lambda,
-                                                                 rho,
-                                                                 nman)
+    Array3 = promote_type(typeof(A), typeof(B))
+    Array2 = typeof(C)
+    return ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT, Array3, Array2}(gravity,
+                                                                                 H0,
+                                                                                 n_moments,
+                                                                                 A,
+                                                                                 B,
+                                                                                 C,
+                                                                                 nu,
+                                                                                 lambda,
+                                                                                 rho,
+                                                                                 nman)
 end
 
-@inline function Base.real(::ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT}) where {
-                                                                                           NVARS,
-                                                                                           NMOMENTS,
-                                                                                           RealT <:
-                                                                                           Real
-                                                                                           }
+@inline function Base.real(::ShallowWaterMomentEquations1D{NVARS, NMOMENTS, RealT,
+                                                           Array3, Array2}) where {
+                                                                                   NVARS,
+                                                                                   NMOMENTS,
+                                                                                   RealT <:
+                                                                                   Real,
+                                                                                   Array3 <:
+                                                                                   AbstractArray{RealT,
+                                                                                                 3},
+                                                                                   Array2 <:
+                                                                                   AbstractArray{RealT,
+                                                                                                 2}
+                                                                                   }
     RealT
 end