Inverse Laplace ∇⁻²! via kernel in ∇²!

src/spectral_gradients.jl

@@ -206,81 +206,68 @@ function UV_from_vordiv!(   U::LowerTriangularMatrix{Complex{NF}},
     end
 end
 
-# In-place update of spectral coeffs alms with their (inverse) Laplace operator-ed version 
-# ∇⁻²!(alms::AbstractMatrix{Complex{NF}},S::SpectralTransform{NF}) where NF = alms .* S.eigenvalues⁻¹
-# ∇²!( alms::AbstractMatrix{Complex{NF}},S::SpectralTransform{NF}) where NF = alms .* S.eigenvalues
-
-# """
-#     ∇⁻²!(   ∇⁻²alms::AbstractMatrix{Complex},
-#             alms::AbstractMatrix{Complex},
-#             S::SpectralTransform)
-
-# Inverse Laplace operator ∇⁻² applied to the spectral coefficients `alms` in spherical
-# coordinates. The radius `R` is omitted in the eigenvalues which are precomputed in `S`.
-# ∇⁻²! is the in-place version which directly stores the output in the first argument `∇⁻²alms`.
-# The integration constant for Legendre polynomial `l=m=0` is zero. The inverse spherical
-# Laplace operator is generally
-
-#     ∇⁻²alms = alms*R²/(-l(l+1))
-
-# with the degree `l` (0-based) of the Legendre polynomial."""
-# function ∇⁻²!(  ∇⁻²alms::AbstractMatrix{Complex{NF}},   # Output: inverse Laplacian of alms
-#                 alms::AbstractMatrix{Complex{NF}},      # spectral coefficients
-#                 S::SpectralTransform{NF}                # precomputed arrays for spectral space
-#                 ) where {NF<:AbstractFloat}
-
-#     @boundscheck size(alms) == size(∇⁻²alms) || throw(BoundsError)
-#     lmax,mmax = size(alms) .- 1     # degree l, order m of the Legendre polynomials
-
-#     @unpack eigenvalues⁻¹ = S
-#     @boundscheck length(eigenvalues⁻¹) >= lmax+1 || throw(BoundsError)
-
-#     @inbounds for m in 1:mmax+1     # order m = 0:mmax but 1-based
-#         for l in m:lmax+1           # degree l = m:lmax but 1-based
-#             ∇⁻²alms[l,m] = alms[l,m]*eigenvalues⁻¹[l]
-#         end
-#     end
-
-#     return ∇⁻²alms
-# end
-
 """
-    ∇²!(    ∇⁻²alms::LowerTriangularMatrix,
+    ∇²!(    ∇²alms::LowerTriangularMatrix,
             alms::LowerTriangularMatrix,
-            S::SpectralTransform)
+            S::SpectralTransform;
+            add::Bool=false,
+            flipsign::Bool=false,
+            inverse::Bool=false)
 
 Laplace operator ∇² applied to the spectral coefficients `alms` in spherical
 coordinates. The radius `R` is omitted in the eigenvalues which are precomputed in `S`.
 ∇²! is the in-place version which directly stores the output in the first argument `∇²alms`.
-The spherical Laplace operator is generally
 
-    ∇⁻²alms = alms*(-l(l+1))/R²
+Options:
+    - `add=true` adds the ∇²(alms) to the output
+    - `flipsign=true` computes -∇²(alms) instead
+    - `inverse=true` computes ∇⁻²(alms) instead
 
-with the degree `l` (0-based) of the Legendre polynomial."""
-function ∇²!(   ∇²alms::LowerTriangularMatrix{Complex{NF}}, # Output: Laplacian of alms
+Default is `add=false`, `flipsign=false`, `inverse=false`. These options can be combined."""
+function ∇²!(   ∇²alms::LowerTriangularMatrix{Complex{NF}}, # Output: (inverse) Laplacian of alms
                 alms::LowerTriangularMatrix{Complex{NF}},   # Input: spectral coefficients
                 S::SpectralTransform{NF};                   # precomputed eigenvalues
                 add::Bool=false,                            # add to output array or overwrite
                 flipsign::Bool=false,                       # -∇² or ∇²
+                inverse::Bool=false,                        # ∇⁻² or ∇²
                 ) where {NF<:AbstractFloat}
 
     @boundscheck size(alms) == size(∇²alms) || throw(BoundsError)
-    lmax,mmax = size(alms) .- (2,1)     # degree l, order m of the Legendre polynomials
+    lmax,mmax = size(alms) .- (1,1)     # 0-based degree l, order m of the Legendre polynomials
 
-    @unpack eigenvalues = S
+    # use eigenvalues⁻¹/eigenvalues for ∇⁻²/∇² based but name both eigenvalues
+    eigenvalues = inverse ? S.eigenvalues⁻¹ : S.eigenvalues
     @boundscheck length(eigenvalues) >= lmax+1 || throw(BoundsError)
 
-    kernel(o,a) = flipsign ? (add ? (o-a) : -a) :
-                             (add ? (o+a) :  a)    
+    @inline kernel(o,a) = flipsign ? (add ? (o-a) : -a)  :
+                                     (add ? (o+a) :  a)
 
     lm = 0
     @inbounds for m in 1:mmax+1     # order m = 0:mmax but 1-based
         for l in m:lmax+1           # degree l = m:lmax but 1-based
             lm += 1
             ∇²alms[lm] = kernel(∇²alms[lm],alms[lm]*eigenvalues[l])
         end
-        lm += 1
     end
 
     return ∇²alms
-end
+end
+
+"""
+    ∇⁻²!(   ∇⁻²alms::LowerTriangularMatrix,
+            alms::LowerTriangularMatrix,
+            S::SpectralTransform;
+            add::Bool=false,
+            flipsign::Bool=false)
+
+Calls ∇²!(∇⁻²alms,alms,S;add,flipsign,inverse=true)."""
+function ∇⁻²!(  ∇⁻²alms::LowerTriangularMatrix{Complex{NF}},# Output: inverse Laplacian of alms
+                alms::LowerTriangularMatrix{Complex{NF}},   # Input: spectral coefficients
+                S::SpectralTransform{NF};                   # precomputed eigenvalues
+                add::Bool=false,                            # add to output array or overwrite
+                flipsign::Bool=false,                       # -∇⁻² or ∇⁻²
+                ) where {NF<:AbstractFloat}
+
+    inverse = true
+    return ∇²!(∇⁻²alms,alms,S;add,flipsign,inverse)
+end

src/spectral_transform.jl

@@ -184,7 +184,7 @@ function SpectralTransform( ::Type{NF},                     # Number format NF
     eigenvalues = [-l*(l+1) for l in 0:lmax+1]
     eigenvalues⁻¹ = inv.(eigenvalues)
     eigenvalues⁻¹[1] = 0                    # set the integration constant to 0
-        
+
     # conversion to NF happens here
     SpectralTransform{NF}(  Grid,nresolution,order,
                             lmax,mmax,nfreq_max,

test/spectral_gradients.jl

@@ -236,3 +236,52 @@ end
         end
     end
 end
+
+@testset "(Inverse) Laplace operator" begin
+
+    for NF in (Float32,Float64)
+        alms = LowerTriangularMatrix(randn(Complex{NF},33,32))
+        alms2 = copy(alms)
+        alms3 = copy(alms)
+
+        S = SpectralTransform(alms,recompute_legendre=false)
+
+        # ∇⁻²! same as inverse=true
+        SpeedyWeather.∇²!(alms2,alms,S,inverse=true);
+        SpeedyWeather.∇⁻²!(alms3,alms,S);
+        @test alms2 == alms3
+
+        # test add=true
+        fill!(alms2,0)
+        SpeedyWeather.∇²!(alms2,alms,S,add=true);
+        SpeedyWeather.∇²!(alms3,alms,S);
+        @test alms2 == alms3
+
+        # also for inverse
+        fill!(alms2,0)
+        SpeedyWeather.∇⁻²!(alms2,alms,S,add=true);
+        SpeedyWeather.∇⁻²!(alms3,alms,S);
+        @test alms2 == alms3
+
+        # test flipsign
+        SpeedyWeather.∇²!(alms2,alms,S,flipsign=true);
+        SpeedyWeather.∇²!(alms3,alms,S);
+        @test alms2 == -alms3
+
+        # also for inverse
+        SpeedyWeather.∇⁻²!(alms2,alms,S,flipsign=true);
+        SpeedyWeather.∇⁻²!(alms3,alms,S);
+        @test alms2 == -alms3
+
+        # test ∇²(∇⁻²) = 1
+        alms[1] = 0     # remove 0-mode which is set to zero
+        SpeedyWeather.∇²!(alms2,alms,S);
+        SpeedyWeather.∇⁻²!(alms3,alms2,S);
+        @test alms ≈ alms3
+
+        # and ∇⁻²(∇²) = 1
+        SpeedyWeather.∇⁻²!(alms2,alms,S);
+        SpeedyWeather.∇²!(alms3,alms2,S);
+        @test alms ≈ alms3
+    end
+end