A new grid for SpeedyWeather.jl? about speedyweather.jl

[milankl]

With #127 the spectral transform in T31 is timed with

for grid in (FullGaussianGrid,FullClenshawGrid,OctahedralGaussianGrid,HEALPixGrid)
   map = gridded(alms;grid)
   S = SpeedyWeather.SpectralTransform(map,recompute_legendre=false)
   @btime SpeedyWeather.gridded!($map,$alms,$S)
   @btime SpeedyWeather.spectral!($alms,$map,$S)
end

as (all Float64)

Grid	truncation	gridded!	spectral!
FullGaussianGrid	F32 cubic	38.398μs	34.898 μs
FullClenshawGrid	C32 cubic	41.082 μs	35.644 μs
OctahedralGaussianGrid	O32 cubic	45.065 μs	31.473 μs
HEALPixGrid	H16 cubic-lat	27.383 μs	26.257 μs

(EDIT: Now with shortened loops over the zonal wavenumber $m$)

Note that the new grids are compared to F32 (so the current non-default cubic version of linear F24 we normally use with T31). HEALPix is only cubic in latitude, but linear in longitude. But sanity check: A single spectral transform can be fast for a range of grids ✅

[milankl]

[milankl]

Daisuke Hotta @hottad was proposing HEALPix4x1 grid, that uses 4 base pixels, no equatorial zone, from Gorski's paper the pixel positions then become
$$z = 1 - \frac{i^2}{N_{side}^2}$$
with ring index $1 \leq i \leq N_{side}$ (northern hemisphere + equator)
$$\phi = \frac{\pi}{2i}(j-\frac{1}{2})$$
with pixel-in-ring index $1 \leq j \leq 4i$. So just the $1/3$ is dropped from Gorski eq. (4) and (5). Similarly the boundaries (eq. (19) and (20)) drop the $1/3$ factor and extend all the way to the equator, $z > 0$. This yields the HEALPix4 grid which looks like

It has at most $4N_{side}$ longitude points around the equator, $2N_{side}-1$ latitude rings and $4N_{side}^2$ grid points in total. It may be more attractive for our purposes as it retains a 2:1 ratio of maximum points around the equator over latitude rings. Therefore cubic in latitude is also cubic in longitude. A cubic truncation with T31 could be H32 (HEALPix4 with $N_{side}=32$ which has 63 latitude rings, at most 128 longitudes and 4096 points in total, which is exactly half of what a full Gaussian grid would have with cubic truncation (8192 points) and less than an octahedral grid (5248).

[milankl]

Summing this up, I think if we want cubic truncation with T31 then these would be good alternative candidates, F32 (128x64 grid points) the current default but higher resolution for cubic, O32 (64 latitude rings, 144 longitude points around the equator), H32 (63 latitude rings, 128 points around the equator)

Cool thing is, H32 would have even fewer grid points than F24 what we currently use for T31. In terms of the FFT complexity, which is $n\log(n)$, we currently have with F24 $mn \log(n) = 48*96\log{96} = 21033$ whereas H32 would have $2\sum_{i=1:32} 4i \log(4i) = 18541$ so about 15% lower. In essence, we may get cubic truncation, speed and equal area all at once!

[milankl]

We could also, without any interpolation needed, store the data on a HEALPix4 grid directly ncview-able into a square matrix like so

@samhatfield I think that would produce some funky videos, with storm tracks running diagonally across the chess board ♟️ ↗️

[milankl]

just a little overview about the characteristics of the latitudinal resolution of the different grids / quadratures with 1024 latitude rings (Gaussian) or 1023 (Clenshaw-Curtis, HEALPix).

• Gaussian latitudes / Clenshaw-Curtis are practically / exactly equi-spaced
• The standard HEALPix grid (4x3 = 12 basepixels) has a higher meridional resolution around the Equator, and lower at mid-latitudes
• Vice versa for the HEALPix4 grid (4x1 = 4 basepixels)
• Both HEALPix grids compensate that with a respectively higher/lower resolution in the zonal direction because they are equal-area.

[milankl]

Also useful https://arxiv.org/abs/astro-ph/0409513 / https://iopscience.iop.org/article/10.1086/427976/pdf
To summarize, some of the HEALPix grid's properties

• 12 base pixels (4 diamonds centred on the equator + 4 diamonds merging at north/south pole, picture from ref above
  
• Choose $N_{side}$ as the parameter that effectively controls the resolution
• These base pixels are subdivided into $N_{side}^2$ pixels, such there is $12N_{side}^2$ pixels/grid points in total
• The area of each base pixel is $\Omega = 4\pi r^2/12$, the area of each grid point is $\Omega = 4\pi r^2/(12N_{side}^2)$
• Polar and equatorial zone are divided by $\cos(\theta) = 2/3$, i.e. $\pm 48.18968510422141$
• In the equatorial zone there are $4N_{side}$ grid points/pixels per latitude ring
• In the polar zones there are $4,8,12,16,20,...,4N_{side}$ pixels per latitude ring

We can create a dependency on Healpix.jl and use their functions and the data structure

Nside = 64
H = HealpixMap{Float32,RingOrder}(Nside)

to hold the data in grid point space. We could use the following resolutions

[milankl]

After playing around with creating different grids, I believe we'll always have the Fourier constraints that at the equator $N_{lon} \geq 2T+1$, but with the regular Gaussian grid we currently use $N_{lon} \geq 3T+1$ and in the meridional the quadrature requires $N_{lat} \geq (2T+1)/2$. Whether the quadrature has to be exact is not clear though. For equi-latitude grids, Clenshaw-Curtis quadrature would have $N_{lat} \geq 2T+1$ instead, but I believe we could fall back to the less contraining condition in the zonal. So I believe a fair comparison would be

$N_{side}$	$N_{lon}$	$N_{lat}$	$N_{pix}^H = 12N_{side}^2$	spectral	reg. Gaussian grid $N_{lon}$x$N_{lat}$	$N_{pix}^G = N_{lon}N_{lat}$
16	64	63	3072	T31	96x48	4608
32	128	127	12288	T63	192x96	18432
64	256	255	49152	T127	384x192	73728
128	512	511	196608	T255	768x384	294912
256	1024	1023	786432	T511	1536x768	1179648
512	2048	2047	3145728	T1023	3072x1536	4718592

Then the number of grid points for HEALPix is always 2/3 of the regular Gaussian grid. The resolution of a HEALPix grid could be estimated via $\Delta x = 4\pi R^2 / 12 N_{side}^2$, roughly

$N_{side}$	$\Delta x$
8	800km
16	400km
32	200km
64	100km
128	50km
256	25km
512	13km
1024	6km
2048	3km

[milankl]

Following Hotta 2018 one should be able to swap out the Legendre weights we currently use with the Clenshaw-Curtis quadrature that uses the following weights

with J being the number of latitudes $N_{lat}$ and $\theta_j$ the equi-distantly spaced colatitudes following

so that's as simple as

julia> nlat = 32
julia> θjs = [j/(nlat+1)*π for j in 1:nlat]
32-element Vector{Float64}:
 0.09519977738150888
 0.19039955476301776
 0.28559933214452665
 ⋮
 2.855993321445266
 2.9511930988267756
 3.046392876208284

julia> wj = [4sin(θj)/(nlat+1)*sum([sin(p*θj)/p for p in 1:2:nlat]) for θj in θjs]
32-element Vector{Float64}:
 0.010663416953504472
 0.01628798008667748
 0.028546546759225664
 ⋮
 0.02854654675922572
 0.016287980086677478
 0.010663416953504501

julia> sum(wj)
2.0000000000000004

[milankl]

Curtis Clenshaw weights can only be applied for equi-latitude grids, or grids with imposed equi-latitude, e.g. the octahedral grid with Gaussian latitudes could be swapped for equi-distant latitudes. In any case, one can always fall back to a Riemann sum, which doesn't guarantee exactness as the Gaussian/Clenshaw-Curtis quadratures, but this might not be needed. An equi-latitude HEALPix would look like

which compensates for the higher meridional density of points in the transition zone betwenn the equatorial belt and the polar caps

Riemann weights for the normal HEALPix grid would be

julia> using Healpix
julia> Nside = 16
julia> res = Healpix.Resolution(Nside)
julia> pixels_per_ring = [getringinfo(res,i).numOfPixels for i in 1:4Nside-1]
63-element Vector{Int64}:
  4
  8
 12
 16
 20
 24
  ⋮
 24
 20
 16
 12
  8
  4

julia> weights = 2pixels_per_ring/12Nside^2
63-element Vector{Float64}:
 0.0026041666666666665
 0.005208333333333333
 0.0078125
 0.010416666666666666
 0.013020833333333334
 0.015625
 ⋮
 0.015625
 0.013020833333333334
 0.010416666666666666
 0.0078125
 0.005208333333333333
 0.0026041666666666665

julia> sum(weights)
2.0

[milankl]

Started the implementation in #122

[milankl]

If one were to use a HEALPixGrid but move the latitude rings to Gaussian latitudes, this would look like

[milankl]

#124 implements further grid abstractions towards a single spectral transform that can support all of the grids discussed above

[hottad]

Great!
I suppose you implemented the regular HEALPix (with 12 base pixels).
With #127, has it become mature enough to allow experimenting with various other grid possibilities? If so, I'd be happy to try HEALPix4 and other things.

[milankl]

Looping over rings was changed with #132