Do not interweave transform

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- De-interweave SpectralTransform [#587](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/587)
 - Rossby-Haurwitz wave initial conditions [#591](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/591)
 - Haversine formula and AbstractSphericalDistance [#588](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/588)
 - Array-agnostic SpectralTransform [#583](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/583)

docs/src/speedytransforms.md

@@ -129,9 +129,7 @@ S = SpectralTransform(spectral_grid)
 Note that because we chose `dealiasing=3` (cubic truncation) we now match a T5 spectral field
 with a 12-ring octahedral Gaussian grid, instead of the 8 rings as above. So going from
 `dealiasing=2` (default) to `dealiasing=3` increased our resolution on the grid while the
-spectral resolution remains the same. The `SpectralTransform` also has options for the
-recomputation or pre-computation of the Legendre polynomials, fore more information see
-[(P)recompute Legendre polynomials](@ref).
+spectral resolution remains the same.
 
 Passing on `S` the `SpectralTransform` now allows us to transform directly on the grid
 defined therein. Note that we recreate `alms` to be of size 7x6 instead of 6x6 for T5
@@ -173,6 +171,43 @@ As you can see the `7x6 LowerTriangularMatrix` in the description above dropped
 `6x6 LowerTriangularMatrix`, this is the size of the input that is expected (otherwise
 you will get a `BoundsError`).
 
+## `SpectralTransform` generators
+
+While you can always create a `SpectralTransform` from a `SpectralGrid` (which defines
+both spectral and grid space) there are other constructors/generators available:
+
+```@example speedytransforms
+SpectralTransform(alms)
+```
+
+Now we have defined the resolution of the spectral space through `alms` but create
+a `SpectralTransform` by making assumption about the grid space. E.g. `Grid=FullGaussianGrid`
+by default, `dealiasing=2` and `nlat_half` correspondingly. But you can also pass them
+on as keyword arguments, for example
+
+```@example speedytransforms
+SpectralTransform(alms, Grid=OctahedralClenshawGrid, nlat_half=24)
+```
+
+Only note that you don't want to specify both `nlat_half` and `dealiasing` as you would
+otherwise overspecify the grid resolution (`dealiasing` will be ignored in that case).
+This also works starting from the grid space
+
+```@example speedytransforms
+grid = rand(FullClenshawGrid, 12)
+SpectralTransform(grid)
+```
+
+where you can also provide spectral resolution `trunc` or `dealiasing`. You can also
+provide both a grid and a lower triangular matrix to describe both spaces
+
+```@example speedytransforms
+SpectralTransform(grid, alms)
+```
+
+and you will precompute the transform between them. For more generators see the
+docstrings at `?SpectralTransform`.
+
 ## Power spectrum
 
 How to take some data and compute a power spectrum with SpeedyTransforms you may ask.
@@ -227,7 +262,6 @@ to be used in grid-point space!
 
 ## Example: Creating k^n-distributed noise
 
-
 How would we construct random noise in spectral space that follows a certain
 power law and transform it back into grid-point space? Define the wavenumber ``k``
 for T31, the spectral resolution we are interested in.
@@ -260,37 +294,35 @@ nothing # hide
 You can always access the underlying data in `map` via `map.data` in case you
 need to get rid of the wrapping into a grid again!
 
-## (P)recompute Legendre polynomials
+## Precomputed polynomials and allocated memory
+
+!!! info "Reuse `SpectralTransform`s wherever possible"
+    Depending on horizontal and vertical resolution of spectral and grid space,
+    a `SpectralTransform` can be become very large in memory. Also the recomputation
+    of the polynomials and the planning of the FFTs are expensive compared to the
+    actual transform itself. Therefore reuse a `SpectralTransform` wherever possible.
 
 The spectral transform uses a Legendre transform in meridional direction. For this the
 Legendre polynomials are required, at each latitude ring this is a ``l_{max} \times m_{max}``
 lower triangular matrix. Storing precomputed Legendre polynomials therefore quickly
-increase in size with resolution. One can recompute them to save memory, but that 
-uses more arithmetic operations. There is therefore a memory-compute tradeoff.
+increase in size with resolution. It is therefore advised to reuse a precomputed
+`SpectralTransform` object wherever possible. This prevents transforms to allocate
+large memory which would otherwise be garbage collected again after the transform.
 
-For a single transform, there is no need to precompute the polynomials as the
-`SpectralTransform` object will be garbage collected again anyway. For low resolution
-simulations with many repeated small transforms it makes sense to precompute the
-polynomials and SpeedyWeather.jl does that automatically anyway. At very high resolution
-the polynomials may, however, become prohibitively large. An example at T127,
-about 100km resolution
+You get information about the size of that memory (both polynomials and required scratch memory)
+in the terminal "show" of a `SpectralTransform` object, e.g. at T127 resolution
+with 8 layers these are
 
 ```@example speedytransforms
-spectral_grid = SpectralGrid(trunc=127)
-SpectralTransform(spectral_grid, recompute_legendre=false)
-```
-the polynomials are about 3MB in size. Easy that is not much. But at T1023 on the
-O1536 octahedral Gaussian grid, this is already 1.5GB, cubically increasing with the
-spectral truncation T. `recompute_legendre=true` (default `false` when
-constructing a `SpectralTransform` object which may be reused) would lower this
-to kilobytes
-```@example speedytransforms
-SpectralTransform(spectral_grid, recompute_legendre=true)
+spectral_grid = SpectralGrid(trunc=127, nlayers=8)
+SpectralTransform(spectral_grid)
 ```
 
 ## Batched Transforms 
 
-SpeedyTransforms also supports batched transforms. With batched input data the `transform` is performed along the leading dimension, and all further dimensions are interpreted as batch dimensions. Take for example 
+SpeedyTransforms also supports batched transforms. With batched input data the `transform`
+is performed along the leading dimension, and all further dimensions are interpreted as
+batch dimensions. Take for example 
 
 ```@example speedytransforms 
 alms = randn(LowerTriangularMatrix{Complex{Float32}}, 32, 32, 5) 

src/LowerTriangularMatrices/lower_triangular_array.jl

@@ -55,19 +55,32 @@ abstract type OneBased <: IndexBasis end
 # get matrix size of LTA from its data array and m, n (number of rows and columns)
 matrix_size(data::AbstractArray, m::Integer, n::Integer) = (m, n, size(data)[2:end]...)
 
+# extend to get the size of the i-th dimension, with 1 returned for any additional dimension
+# as it is also defined for Array
+function matrix_size(data::AbstractArray, m::Integer, n::Integer, i::Integer)
+    i == 1 && return m      # first dimension is the number of rows m
+    i == 2 && return n      # second dimension is the number of columns n
+    return size(data, i-1)  # -1 as m, n are collapsed into a vector in the data array
+end 
+
 """$(TYPEDSIGNATURES)
 Size of a `LowerTriangularArray` defined as size of the flattened array if `as <: AbstractVector`
 and as if it were a full matrix when `as <: AbstractMatrix`` ."""
 Base.size(L::LowerTriangularArray, base::Type{<:IndexBasis}=OneBased; as=Vector) = size(L, base, as)
-Base.size(L::LowerTriangularArray, i::Integer, base::Type{<:IndexBasis}=OneBased; as=Vector) = size(L, base; as=as)[i]
+Base.size(L::LowerTriangularArray, i::Integer, base::Type{<:IndexBasis}=OneBased; as=Vector) = size(L, i, base, as)
 
 # use multiple dispatch to chose the right options of basis and vector/flat vs matrix indexing
 # matrix indexing can be zero based (natural for spherical harmonics) or one-based,
-# vector/flat indexing has only one based indexing
+# vector/flat indexing has only one-based indexing
 Base.size(L::LowerTriangularArray, base::Type{OneBased}, as::Type{Matrix}) = matrix_size(L.data, L.m, L.n)
 Base.size(L::LowerTriangularArray, base::Type{ZeroBased}, as::Type{Matrix}) = matrix_size(L.data, L.m-1, L.n-1)
 Base.size(L::LowerTriangularArray, base::Type{OneBased}, as::Type{Vector}) = size(L.data)
 
+# size(L, i, ...) to get the size of the i-th dimension, with 1 returned for any additional dimension 
+Base.size(L::LowerTriangularArray, i::Integer, base::Type{OneBased}, as::Type{Matrix}) = matrix_size(L.data, L.m, L.n, i)
+Base.size(L::LowerTriangularArray, i::Integer, base::Type{ZeroBased}, as::Type{Matrix}) = matrix_size(L.data, L.m-1, L.n-1, i)
+Base.size(L::LowerTriangularArray, i::Integer, base::Type{OneBased}, as::Type{Vector}) = size(L.data, i)
+
 # sizeof the underlying data vector
 Base.sizeof(L::LowerTriangularArray) = sizeof(L.data)
 
@@ -107,7 +120,6 @@ end
 
 function Base.array_summary(io::IO, L::LowerTriangularMatrix{T}, inds::Tuple{Vararg{Base.OneTo}}) where T
     mn = size(L; as=Matrix)
-    @info Base.dims2string(length.(inds))
     print(io, Base.dims2string(length.(inds)), ", $(mn[1])x$(mn[2]) LowerTriangularMatrix{$T}")
 end
 

src/SpeedyTransforms/SpeedyTransforms.jl

@@ -14,8 +14,6 @@ import Primes
 using ..LowerTriangularMatrices
 using ..RingGrids
 
-const DEFAULT_GRID = FullGaussianGrid
-
 # TRANSFORM
 export  SpectralTransform,
         transform!,
@@ -42,7 +40,10 @@ export  spectral_truncation,
 
 include("aliasing.jl")
 include("legendrepolarray.jl")
+include("legendre_shortcuts.jl")
 include("spectral_transform.jl")
+include("fourier.jl")
+include("legendre.jl")
 include("spectral_gradients.jl")
 include("spectral_truncation.jl")
 include("spectrum.jl")