Merge pull request #136 from milankl/mk/spectrans

Float32 spectral transform

.github/workflows/Documenter.yml

@@ -14,7 +14,7 @@ jobs:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@latest
         with:
-          version: '1.6'
+          version: '1.8'
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy

Project.toml

@@ -14,6 +14,7 @@ CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+GenericFFT = "a8297547-1b15-4a5a-a998-a2ac5f1cef28"
 Healpix = "9f4e344d-96bc-545a-84a3-ae6b9e1b672b"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -38,6 +39,7 @@ CUDAKernels = "0.3, 0.4"
 CodecZlib = "0.7"
 FFTW = "1"
 FastGaussQuadrature = "0.4"
+GenericFFT = "0.1"
 Healpix = "4"
 JLD2 = "0.4"
 KernelAbstractions = "0.7, 0.8"

src/SpeedyWeather.jl

@@ -9,7 +9,9 @@ module SpeedyWeather
     import FastGaussQuadrature
     import AssociatedLegendrePolynomials as Legendre
     import Healpix
+    import AbstractFFTs
     import FFTW
+    import GenericFFT
     import Primes
     import LinearAlgebra
 

src/grids.jl

@@ -227,6 +227,11 @@ function get_nlon_per_ring(::Type{<:AbstractHEALPixGrid},nside::Integer,j::Integ
     return nlon_healpix(nside,j)
 end
 
+function get_nlons(Grid::Type{<:AbstractGrid},nresolution::Integer;both_hemispheres::Bool=false)
+    n = both_hemispheres ? get_nlat(Grid,nresolution) : get_nlat_half(Grid,nresolution)
+    return [get_nlon_per_ring(Grid,nresolution,j) for j in 1:n]
+end
+
 # total number of grid points per grid type
 get_npoints(::Type{<:FullGaussianGrid},nlat_half::Integer) = npoints_gaussian(nlat_half)
 get_npoints(::Type{<:FullClenshawGrid},nlat_half::Integer) = npoints_clenshaw(nlat_half)
@@ -370,7 +375,6 @@ eachgridpoint(grid::G) where {G<:AbstractGrid} = Base.OneTo(get_npoints(G,get_nr
 # QUADRATURE WEIGHTS
 # gaussian_weights are exact for Gaussian latitudes when nlat > (2T+1)/2
 # clenshaw_curtis_weights are exact for equi-angle latitudes when nlat > 2T+1
-# riemann_weights not exact but used for HEALPix
 gaussian_weights(nlat_half::Integer) = FastGaussQuadrature.gausslegendre(2nlat_half)[2][1:nlat_half]
 
 function clenshaw_curtis_weights(nlat_half::Integer)
@@ -383,5 +387,8 @@ get_quadrature_weights(::Type{<:FullClenshawGrid},nlat_half::Integer) = clenshaw
 get_quadrature_weights(::Type{<:FullGaussianGrid},nlat_half::Integer) = gaussian_weights(nlat_half)
 get_quadrature_weights(::Type{<:OctahedralGaussianGrid},nlat_half::Integer) = gaussian_weights(nlat_half)
 
-# HEALPix's solid angle is always 4π/npoints
-get_quadrature_weights(G::Type{<:HEALPixGrid},nside::Integer) = 4π/get_npoints(G,nside)*ones(get_nlat_half(G,nside))
+# SOLID ANGLES ΔΩ = sinθ Δθ Δϕ
+get_solid_angles(Grid::Type{<:AbstractGrid},nlat_half::Integer) = 
+    get_quadrature_weights(Grid,nlat_half) .* (2π./get_nlons(Grid,nlat_half))
+get_solid_angles(Grid::Type{<:HEALPixGrid},nside::Integer) =
+    4π/get_npoints(Grid,nside)*ones(get_nlat_half(Grid,nside))

src/lower_triangular_matrix.jl

@@ -168,6 +168,42 @@ function Base.copyto!(  L1::LowerTriangularMatrix{T},   # copy to L1
     L1
 end
 
+function Base.copyto!(  L::LowerTriangularMatrix{T},    # copy to L
+                        M::AbstractMatrix) where T      # copy from M
+
+    @boundscheck size(L) == size(M) || throw(BoundsError)
+    lmax,mmax = size(L)
+
+    lm = 0
+    @inbounds for m in 1:mmax
+        for l in m:lmax
+            lm += 1
+            L[lm] = convert(T,M[l,m])
+        end
+    end
+    L
+end
+
+function Base.copyto!(  M::AbstractMatrix{T},               # copy to L
+                        L::LowerTriangularMatrix) where T   # copy from M
+
+    @boundscheck size(L) == size(M) || throw(BoundsError)
+    lmax,mmax = size(L)
+
+    lm = 0
+    @inbounds for m in 1:mmax
+        for l in 1:m-1          # zero for upper triangle (excl diagonal)
+            M[l,m] = zero(T)
+        end
+
+        for l in m:lmax         # convert and copy for lower triangle
+            lm += 1
+            M[l,m] = convert(T,M[lm])
+        end
+    end
+    L
+end
+
 function LowerTriangularMatrix{T}(M::LowerTriangularMatrix) where T
     L = LowerTriangularMatrix{T}(undef,size(M)...)
     copyto!(L,M)

src/utility_functions.jl

@@ -45,24 +45,36 @@ function roundup_fft(n::Integer;small_primes::Vector{T}=[2,3,5]) where {T<:Integ
     return n-2      # subtract unnecessary last += 2 addition
 end
 
-"""Set all negative entries in an array to zero."""
+"""
+    clip_negatives!(A::AbstractArray)
+
+Set all negative entries `a` in `A` to zero."""
 function clip_negatives!(A::AbstractArray{T}) where T
     @inbounds for i in eachindex(A)
         A[i] = max(A[i],zero(T))
     end
 end
 
-function underflow_small!(A::AbstractArray{T},ϵ::Real) where T
+"""
+    underflow!(A::AbstractArray,ϵ::Real)
+
+Underflows element `a` in `A` to zero if `abs(a) < ϵ`."""
+function underflow!(A::AbstractArray{T},ϵ::Real) where T
     ϵT = convert(T,abs(ϵ))
     @inbounds for i in eachindex(A)
         A[i] = abs(A[i]) < ϵT ? zero(T) : A[i]
     end
 end
 
+"""
+    flipgsign!(A::AbstractArray)
+
+Like `-A` but in-place."""
 function flipsign!(A::AbstractArray)
     @inbounds for i in eachindex(A)
         A[i] = -A[i]
     end
+    A
 end
 
 """
@@ -91,4 +103,5 @@ function readable_secs(secs::Real)
 end
 
 # define as will only become availale in Julia 1.9
-pkgversion(m::Module) = VersionNumber(TOML.parsefile(joinpath(dirname(string(first(methods(m.eval)).file)), "..", "Project.toml"))["version"])
+pkgversion(m::Module) = VersionNumber(TOML.parsefile(joinpath(
+    dirname(string(first(methods(m.eval)).file)), "..", "Project.toml"))["version"])

test/spectral_transform.jl

@@ -91,12 +91,13 @@ end
 @testset "Transform: Geopotential" begin
 
     # Test for variable resolution
-    @testset for trunc in spectral_resolutions[1:2]
-        for NF in (Float64,Float32)
-            for Grid in (   FullGaussianGrid,
-                            FullClenshawGrid,
-                            OctahedralGaussianGrid,)
-                            # HEALPixGrid)
+    @testset for trunc in [31,42]
+        @testset for NF in (Float64,Float32)
+            @testset for Grid in (  FullGaussianGrid,
+                                    FullClenshawGrid,
+                                    OctahedralGaussianGrid,)
+                                    # HEALPixGrid)
+
                 P = Parameters(;NF,trunc,model=:shallowwater)
                 S = SpectralTransform(P)
                 B = Boundaries(P,S)
@@ -117,65 +118,4 @@ end
             end
         end
     end
-end
-
-# @testset "Transform: with one more l" begin
-#     @testset for NF in (Float32,Float64)
-#         p,d,m = initialize_speedy(  NF,
-#                                     model=:shallowwater,
-#                                     initial_conditions=:rest,
-#                                     layer_thickness=0)
-
-#         # make sure vorticity and divergence are 0
-#         fill!(p.vor,0)
-#         fill!(p.div,0)
-
-#         # make sure vorticity and divergence are 0
-#         fill!(d.tendencies.vor_tend,0)                  
-#         fill!(d.tendencies.div_tend,0)
-
-#         # create initial conditions
-#         vor0 = zero(p.vor[:,:,1,1])
-#         div0 = zero(p.div[:,:,1,1])
-#         vor0[:,:] .= randn(Complex{NF},size(vor0)...)
-#         div0[:,:] .= randn(Complex{NF},size(div0)...)
-
-#         # remove non-zero entries in upper triangle
-#         SpeedyWeather.spectral_truncation!(vor0)
-#         SpeedyWeather.spectral_truncation!(div0)
-
-#         p.vor[:,:,1,1] .= vor0
-#         p.div[:,:,1,1] .= div0
-
-#         # get corresponding irrotational u_grid, v_grid (incl *coslat scaling)
-#         SpeedyWeather.gridded!(d,p,m)   
-
-#         # check we've actually created non-zero U,V
-#         @test all(d.grid_variables.U_grid .!= 0)
-#         @test all(d.grid_variables.V_grid .!= 0)
-
-#         lmax,mmax = size(vor0)      # 1-based maximum l,m
-
-#         @testset for lmax in (lmax,lmax+1)
-
-#             U = zeros(Complex{NF},lmax,mmax,1)
-#             U_grid = zero(d.grid_variables.U_grid)
-#             U_grid2 = zero(d.grid_variables.U_grid)
-
-#             # transform back and forth first as the original U_grid
-#             # is not exactly representable for lmax=lmax,lmax+1 (spectral truncation)
-#             S = m.geospectral.spectral_transform
-#             SpeedyWeather.spectral!(U,d.grid_variables.U_grid,S)
-#             SpeedyWeather.gridded!(U_grid,U,S)
-#             SpeedyWeather.spectral!(U,U_grid,S)
-#             SpeedyWeather.gridded!(U_grid2,U,S)
-
-#             testall = .≈(U_grid,U_grid2,rtol=sqrt(eps(NF)))
-#             @test sum(testall) > 0.99*length(testall)
-
-#             for i in eachindex(U_grid, U_grid2)
-#                 @test U_grid[i] ≈ U_grid2[i] rtol=5*sqrt(sqrt(eps(NF)))
-#             end
-#         end
-#     end
-# end
+end