Merge pull request #503 from SpeedyWeather/mg/ltm-gpu-ndim

LowerTriangularMatrices: N-dimensional and GPU ready

Project.toml

@@ -15,6 +15,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GenericFFT = "a8297547-1b15-4a5a-a998-a2ac5f1cef28"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -63,6 +64,7 @@ julia = "1.9"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 
 [targets]
-test = ["Test"]
+test = ["Test", "JLArrays"]

docs/src/lowertriangularmatrices.md

@@ -4,34 +4,37 @@ LowerTriangularMatrices is a submodule that has been developed for SpeedyWeather
 technically independent (SpeedyWeather.jl however imports it and so does SpeedyTransforms)
 and can also be used without running simulations. It is just not put into its own respective repository.
 
-This module defines `LowerTriangularMatrix`, a lower triangular matrix, which in contrast to
+This module defines `LowerTriangularArray`, a lower triangular matrix format, which in contrast to
 [`LinearAlgebra.LowerTriangular`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.LowerTriangular) does not store the entries above the diagonal. SpeedyWeather.jl
-uses `LowerTriangularMatrix` which is defined as a subtype of `AbstractMatrix` to store
-the spherical harmonic coefficients (see [Spectral packing](@ref)). 
+uses `LowerTriangularArray` which is defined as a subtype of `AbstractArray` to store
+the spherical harmonic coefficients (see [Spectral packing](@ref)). For 2D `LowerTriangularArray` the alias `LowerTriangularMatrix` exists. Higher dimensional `LowerTriangularArray` are 'batches' of 2D `LowerTriangularMatrix`. So, for example a ``(10\times 10\times 10)`` `LowerTriangularArray` holds 10 `LowerTriangularMatrix` of size ``(10\times 10)`` in one array.  
 
-## Creation of `LowerTriangularMatrix`
+## Creation of `LowerTriangularArray` 
 
-A `LowerTriangularMatrix` can be created using `zeros`, `ones`, `rand`, or `randn`
+A `LowerTriangularMatrix` and `LowerTriangularArray` can be created using `zeros`, `ones`, `rand`, or `randn`
 ```@repl LowerTriangularMatrices
 using SpeedyWeather.LowerTriangularMatrices
 
 L = rand(LowerTriangularMatrix{Float32}, 5, 5)
+L2 = rand(LowerTriangularArray{Float32}, 5, 5, 5)
 ```
 or the `undef` initializor `LowerTriangularMatrix{Float32}(undef, 3, 3)`.
 The element type is arbitrary though, you can use any type `T` too.
 
-Alternatively, it can be created through conversion from `Matrix`, which drops the upper triangle
+Alternatively, it can be created through conversion from `Array`, which drops the upper triangle
 entries and sets them to zero
 ```@repl LowerTriangularMatrices
 M = rand(Float16, 3, 3)
-
 L = LowerTriangularMatrix(M)
+
+M2 = rand(Float16, 3, 3, 5)
+L2 = LowerTriangularArray(M)
 ```
 
-## Indexing `LowerTriangularMatrix`
+## Indexing `LowerTriangularArray`
 
-`LowerTriangularMatrix` supports two types of indexing: 1) by denoting two indices, column and row `[l, m]`
-or 2) by denoting a single index `[lm]`. The double index works as expected
+`LowerTriangularArray` supports two types of indexing: 1) by denoting two indices, column and row `[l, m, ..]`
+or 2) by denoting a single index `[lm, ..]`. The double index works as expected
 ```@repl LowerTriangularMatrices
 L
 
@@ -60,7 +63,7 @@ for j in 1:m, i in j:n
     L[ij] = i+j
 end
 ```
-which loops over all lower triangle entries of `L::LowerTriangularMatrix` and the single
+which loops over all lower triangle entries of `L::LowerTriangularArray` and the single
 index `ij` is simply counted up. However, one could also use `[i, j]` as indices in the
 loop body or to perform any calculation (`i+j` here).
 An iterator over all entries in the lower triangle can be created by
@@ -70,14 +73,26 @@ for ij in eachindex(L)
 end
 ```
 The `setindex!` functionality of matrixes will throw a `BoundsError` when trying to write
-into the upper triangle of a `LowerTriangularMatrix`, for example
+into the upper triangle of a `LowerTriangularArray`, for example
 ```@repl LowerTriangularMatrices
 L[2, 1] = 0    # valid index
 
 L[1, 2] = 0    # invalid index in the upper triangle
 ```
 
-## Linear algebra with `LowerTriangularMatrix`
+Higher dimensional `LowerTriangularArray` can be indexed with multidimensional array indices 
+like most other arrays types. Both the single index and the double index for the lower 
+triangle work as shown here
+```@repl LowerTriangularMatrices
+L = rand(LowerTriangularMatrix{Float32}, 3, 3, 5)
+
+L[2, 1] # second lower triangle element of the first lower triangle matrix 
+
+L[2, 1, 1] # (2,1) element of the first lower triangle matrix 
+```
+The `setindex!` functionality follows accordingly. 
+
+## Linear algebra with `LowerTriangularArray`
 
 The [LowerTriangularMatrices](@ref lowertriangularmatrices) module's main purpose is not linear algebra, and it's
 implementation may not be efficient, however, many operations work as expected
@@ -93,6 +108,15 @@ operations, including `inv` or `\`. Hence when trying to do more sophisticated l
 algebra with `LowerTriangularMatrix` we quickly leave lower triangular-land and go
 back to normal matrix-land.
 
+## GPU 
+
+`LowerTriangularArray{T,N,ArrayType}` wraps around an array of type `ArrayType`. If this array is a GPU array (e.g. `CuArray`), all operations are performed on GPU as well. The implementation was written so that scalar indexing is avoided in almost all cases, so that GPU operation should be performant. To use `LowerTriangularArray` on GPU you can e.g. just `adapt` an exisiting `LowerTriangularArray`.
+```julia 
+using Adapt
+L = rand(LowerTriangularArray{Float32}, 5, 5, 5)
+L_gpu = adapt(CuArray, L)
+```
+
 ## Function and type index
 
 ```@autodocs

src/LowerTriangularMatrices/LowerTriangularMatrices.jl

@@ -1,12 +1,22 @@
 module LowerTriangularMatrices
 
+# STRUCTURE
 using DocStringExtensions
+
+# GPU
 import Adapt
-import UnicodePlots
+import GPUArrays
+
+# NUMERICS
+import LinearAlgebra: tril!
 
-export LowerTriangularMatrix, eachharmonic
+# VISUALISATION
+import UnicodePlots
 # export plot
 
+export LowerTriangularMatrix, LowerTriangularArray
+export eachharmonic
+
 include("lower_triangular_matrix.jl")
 include("plot.jl")