AbstractGridArray N-dimensional and GPU ready

Project.toml

@@ -31,9 +31,11 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 
 [weakdeps]
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 
 [extensions]
 SpeedyWeatherMakieExt = "Makie"
+SpeedyWeatherJLArraysExt = "JLArrays"
 
 [compat]
 AbstractFFTs = "1"
@@ -47,7 +49,9 @@ DocStringExtensions = "0.9"
 FFTW = "1"
 FLoops = "0.2"
 FastGaussQuadrature = "0.4, 0.5, 1"
+GPUArrays = "10"
 GenericFFT = "0.1"
+JLArrays = "0.1.4"
 JLD2 = "0.4"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1.9"
@@ -63,8 +67,8 @@ UnicodePlots = "3.3"
 julia = "1.9"
 
 [extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "JLArrays"]

docs/src/output.md

@@ -96,13 +96,13 @@ for the model integration.
 ```julia
 my_output_writer = OutputWriter(spectral_grid, ShallowWater, output_Grid=FullClenshawGrid, nlat_half=48)
 ```
-Note that by default the output is on the corresponding full of the grid used in the dynamical core
+Note that by default the output is on the corresponding full type of the grid type used in the dynamical core
 so that interpolation only happens at most in the zonal direction as they share the location of the
 latitude rings. You can check this by
 ```@example netcdf
-RingGrids.full_grid(OctahedralGaussianGrid)
+RingGrids.full_grid_type(OctahedralGaussianGrid)
 ```
-So the corresponding full grid of an `OctahedralGaussianGrid` is the `FullGaussiangrid` and the same resolution
+So the corresponding full grid of an `OctahedralGaussianGrid` is the `FullGaussianGrid` and the same resolution
 `nlat_half` is chosen by default in the output writer (which you can change though as shown above).
 Overview of the corresponding full grids
 

ext/SpeedyWeatherJLArraysExt.jl

@@ -0,0 +1,10 @@
+module SpeedyWeatherJLArraysExt
+
+using SpeedyWeather, JLArrays
+
+# for RingGrids and LowerTriangularMatrices:
+# every Array needs this method to strip away the parameters
+SpeedyWeather.RingGrids.nonparametric_type(::Type{<:JLArray}) = JLArray
+SpeedyWeather.LowerTriangularMatrices.nonparametric_type(::Type{<:JLArray}) = JLArray
+
+end # module

ext/SpeedyWeatherMakieExt.jl

@@ -12,7 +12,7 @@ function Makie.heatmap(
     title::String = "$(RingGrids.get_nlat(grid))-ring $(typeof(grid))",
     kwargs...   # pass on to Makie.heatmap
 )
-    full_grid = RingGrids.interpolate(RingGrids.full_grid(typeof(grid)), grid.nlat_half, grid)
+    full_grid = RingGrids.interpolate(RingGrids.full_grid_type(grid), grid.nlat_half, grid)
     heatmap(full_grid; title, kwargs...)
 end
 

src/RingGrids/RingGrids.jl

@@ -9,50 +9,77 @@ import Statistics: mean
 import FastGaussQuadrature
 import LinearAlgebra
 
-# GRIDS
+# GPU
+import Adapt
+import GPUArrays
+import CUDA
+
+# ABSTRACT GRIDS (2D) AND GRIDARRAYS (3D+)
+export  AbstractGridArray,
+        AbstractFullGridArray,
+        AbstractReducedGridArray
+
 export  AbstractGrid,
         AbstractFullGrid,
-        AbstractOctahedralGrid,
-        AbstractHEALPixGrid,
-        AbstractOctaHEALPixGrid
+        AbstractReducedGrid
+
+# CONCRETE GRIDS (2D) AND GRIDARRAYS (3D+)
+export  FullGaussianArray,
+        FullClenshawArray,
+        FullHEALPixArray,
+        FullOctaHEALPixArray
 
 export  FullGaussianGrid,
         FullClenshawGrid,
         FullHEALPixGrid,
         FullOctaHEALPixGrid
 
+export  OctahedralGaussianArray,
+        OctahedralClenshawArray,
+        HEALPixArray,
+        OctaHEALPixArray
+
 export  OctahedralGaussianGrid,
         OctahedralClenshawGrid,
         HEALPixGrid,
         OctaHEALPixGrid
 
-# GRID FUNCTIONS
+# SIZE
 export  grids_match,
-        get_truncation,
-        get_resolution,
         get_nlat,
         get_nlat_half,
-        get_npoints,
-        get_latdlonds,
+        get_npoints2D
+
+# COORDINATES
+export  get_latdlonds,
         get_lat,
         get_colat,
         get_latd,
         get_lond,
-        each_index_in_ring,
-        each_index_in_ring!,
-        eachgridpoint,
-        eachring,
-        whichring,
         get_nlons,
-        get_nlon_max,
-        get_quadrature_weights,
+        get_nlon_max
+
+# INTEGRATION
+export  get_quadrature_weights,
         get_solid_angles
 
+# ITERATORS
+export  eachgrid,
+        eachring,
+        whichring,
+        eachgridpoint,
+        each_index_in_ring,
+        each_index_in_ring!
+
 # SCALING
 export  scale_coslat!,
         scale_coslat²!,
         scale_coslat⁻¹!,
-        scale_coslat⁻²!
+        scale_coslat⁻²!,
+        scale_coslat,
+        scale_coslat²,
+        scale_coslat⁻¹,
+        scale_coslat⁻²
 
 # INTERPOLATION
 export  AbstractInterpolator,
@@ -67,19 +94,31 @@ export  interpolate,
         update_locator,
         update_locator!
 
-# export  plot
-
 include("utility_functions.jl")
 
-include("grids_general.jl")
+# GENERAL
+include("general.jl")
 include("full_grids.jl")
-include("octahedral.jl")
-include("healpix.jl")
-include("octahealpix.jl")
+include("reduced_grids.jl")
+include("scaling.jl")
+
+# FULL GRIDS
+include("grids/full_gaussian.jl")
+include("grids/full_clenshaw.jl")
+include("grids/full_healpix.jl")
+include("grids/full_octahealpix.jl")
+
+# REDUCED GRIDS
+include("grids/octahedral_gaussian.jl")
+include("grids/octahedral_clenshaw.jl")
+include("grids/healpix.jl")
+include("grids/octahealpix.jl")
+
+# INTEGRATION AND INTERPOLATION
 include("quadrature_weights.jl")
 include("interpolation.jl")
+
+# OUTPUT
 include("show.jl")
-include("similar.jl")
-include("scaling.jl")
 
 end