synthesis and analysis of Fourier series on S^2 with grids that match spherefun

CC @Hadrien-Montanelli

Author: MikaelSlevinsky

src/SphericalHarmonics/synthesisanalysis.jl

@@ -15,6 +15,16 @@ function plan_synthesis(A::Matrix{T}) where T<:fftwNumber
     SynthesisPlan(planθ, planφ, C, zeros(T, n))
 end
 
+function plan_synthesis2(A::Matrix{T}) where T<:fftwNumber
+    m, n = size(A)
+    x = FFTW.FakeArray(T, m)
+    y = FFTW.FakeArray(T, n)
+    planθ = FFTW.plan_r2r!(x, FFTW.REDFT00), FFTW.plan_r2r!(FFTW.FakeArray(T, m-2), FFTW.RODFT00)
+    planφ = FFTW.plan_r2r!(y, FFTW.HC2R)
+    C = ColumnPermutation(vcat(1:2:n, 2:2:n))
+    SynthesisPlan(planθ, planφ, C, zeros(T, n))
+end
+
 struct AnalysisPlan{T, P1, P2}
     planθ::P1
     planφ::P2
@@ -32,7 +42,17 @@ function plan_analysis(A::Matrix{T}) where T<:fftwNumber
     AnalysisPlan(planθ, planφ, C, zeros(T, n))
 end
 
-function Base.A_mul_B!(Y::Matrix{T}, P::SynthesisPlan{T}, X::Matrix{T}) where T
+function plan_analysis2(A::Matrix{T}) where T<:fftwNumber
+    m, n = size(A)
+    x = FFTW.FakeArray(T, m)
+    y = FFTW.FakeArray(T, n)
+    planθ = FFTW.plan_r2r!(x, FFTW.REDFT00), FFTW.plan_r2r!(FFTW.FakeArray(T, m-2), FFTW.RODFT00)
+    planφ = FFTW.plan_r2r!(y, FFTW.R2HC)
+    C = ColumnPermutation(vcat(1:2:n, 2:2:n))
+    AnalysisPlan(planθ, planφ, C, zeros(T, n))
+end
+
+function Base.A_mul_B!(Y::Matrix{T}, P::SynthesisPlan{T, Tuple{r2rFFTWPlan{T,(REDFT01,),true,1}, r2rFFTWPlan{T,(RODFT01,),true,1}}}, X::Matrix{T}) where T
     M, N = size(X)
 
     # Column synthesis
@@ -73,7 +93,52 @@ function Base.A_mul_B!(Y::Matrix{T}, P::SynthesisPlan{T}, X::Matrix{T}) where T
     Y
 end
 
-function Base.A_mul_B!(Y::Matrix{T}, P::AnalysisPlan{T}, X::Matrix{T}) where T
+function Base.A_mul_B!(Y::Matrix{T}, P::SynthesisPlan{T, Tuple{r2rFFTWPlan{T,(REDFT00,),true,1}, r2rFFTWPlan{T,(RODFT00,),true,1}}}, X::Matrix{T}) where T
+    M, N = size(X)
+
+    # Column synthesis
+    PCe = P.planθ[1]
+    PCo = P.planθ[2]
+
+    X[1] *= two(T)
+    X[M,1] *= two(T)
+    A_mul_B_col_J!(Y, PCe, X, 1)
+    X[1] *= half(T)
+    X[M,1] *= half(T)
+
+    for J = 2:4:N
+        A_mul_B_col_J!(Y, PCo, X, J, false)
+        J < N && A_mul_B_col_J!(Y, PCo, X, J+1, false)
+    end
+    for J = 4:4:N
+        X[1,J] *= two(T)
+        X[M,J] *= two(T)
+        J < N && (X[1,J+1] *= two(T); X[M,J+1] *= two(T))
+        A_mul_B_col_J!(Y, PCe, X, J)
+        J < N && A_mul_B_col_J!(Y, PCe, X, J+1)
+        X[1,J] *= half(T)
+        X[M,J] *= half(T)
+        J < N && (X[1,J+1] *= half(T); X[M,J+1] *= half(T))
+    end
+    scale!(half(T), Y)
+
+    # Row synthesis
+    scale!(inv(sqrt(π)), Y)
+    invsqrttwo = inv(sqrt(2))
+    @inbounds for i = 1:M Y[i] *= invsqrttwo end
+
+    temp = P.temp
+    planφ = P.planφ
+    C = P.C
+    for I = 1:M
+        copy_row_I!(temp, Y, I)
+        row_synthesis!(planφ, C, temp)
+        copy_row_I!(Y, temp, I)
+    end
+    Y
+end
+
+function Base.A_mul_B!(Y::Matrix{T}, P::AnalysisPlan{T, Tuple{r2rFFTWPlan{T,(REDFT10,),true,1}, r2rFFTWPlan{T,(RODFT10,),true,1}}}, X::Matrix{T}) where T
     M, N = size(X)
 
     # Row analysis
@@ -109,6 +174,45 @@ function Base.A_mul_B!(Y::Matrix{T}, P::AnalysisPlan{T}, X::Matrix{T}) where T
     Y
 end
 
+function Base.A_mul_B!(Y::Matrix{T}, P::AnalysisPlan{T, Tuple{r2rFFTWPlan{T,(REDFT00,),true,1}, r2rFFTWPlan{T,(RODFT00,),true,1}}}, X::Matrix{T}) where T
+    M, N = size(X)
+
+    # Row analysis
+    temp = P.temp
+    planφ = P.planφ
+    C = P.C
+    for I = 1:M
+        copy_row_I!(temp, X, I)
+        row_analysis!(planφ, C, temp)
+        copy_row_I!(Y, temp, I)
+    end
+
+    # Column analysis
+    PCe = P.planθ[1]
+    PCo = P.planθ[2]
+
+    A_mul_B_col_J!(Y, PCe, Y, 1)
+    Y[1] *= half(T)
+    Y[M, 1] *= half(T)
+    for J = 2:4:N
+        A_mul_B_col_J!(Y, PCo, Y, J, true)
+        J < N && A_mul_B_col_J!(Y, PCo, Y, J+1, true)
+        Y[M-1,J] = zero(T)
+        J < N && (Y[M-1,J+1] = zero(T))
+    end
+    for J = 4:4:N
+        A_mul_B_col_J!(Y, PCe, Y, J)
+        J < N && A_mul_B_col_J!(Y, PCe, Y, J+1)
+        Y[1,J] *= half(T)
+        Y[M,J] *= half(T)
+        J < N && (Y[1,J+1] *= half(T); Y[M,J+1] *= half(T))
+    end
+    scale!(sqrt(π)*inv(T(M-1)), Y)
+    sqrttwo = sqrt(2)
+    @inbounds for i = 1:M Y[i] *= sqrttwo end
+
+    Y
+end
 
 
 
@@ -185,3 +289,26 @@ function unsafe_execute_col_J!(plan::r2rFFTWPlan{T}, X::Matrix{T}, Y::Matrix{T},
     M = size(X, 1)
     ccall((:fftwf_execute_r2r, libfftwf), Void, (PlanPtr, Ptr{T}, Ptr{T}), plan, pointer(X, M*(J-1)+1), pointer(Y, M*(J-1)+1))
 end
+
+function A_mul_B_col_J!(Y::Matrix{T}, P::r2rFFTWPlan{T}, X::Matrix{T}, J::Int, TF::Bool) where T
+    unsafe_execute_col_J!(P, X, Y, J, TF)
+    return Y
+end
+
+function unsafe_execute_col_J!(plan::r2rFFTWPlan{T}, X::Matrix{T}, Y::Matrix{T}, J::Int, TF::Bool) where T<:fftwDouble
+    M = size(X, 1)
+    if TF
+        ccall((:fftw_execute_r2r, libfftw), Void, (PlanPtr, Ptr{T}, Ptr{T}), plan, pointer(X, M*(J-1)+2), pointer(Y, M*(J-1)+1))
+    else
+        ccall((:fftw_execute_r2r, libfftw), Void, (PlanPtr, Ptr{T}, Ptr{T}), plan, pointer(X, M*(J-1)+1), pointer(Y, M*(J-1)+2))
+    end
+end
+
+function unsafe_execute_col_J!(plan::r2rFFTWPlan{T}, X::Matrix{T}, Y::Matrix{T}, J::Int, TF::Bool) where T<:fftwSingle
+    M = size(X, 1)
+    if TF
+        ccall((:fftwf_execute_r2r, libfftwf), Void, (PlanPtr, Ptr{T}, Ptr{T}), plan, pointer(X, M*(J-1)+2), pointer(Y, M*(J-1)+1))
+    else
+        ccall((:fftwf_execute_r2r, libfftwf), Void, (PlanPtr, Ptr{T}, Ptr{T}), plan, pointer(X, M*(J-1)+1), pointer(Y, M*(J-1)+2))
+    end
+end

test/sphericalharmonics/synthesisanalysistests.jl

@@ -78,3 +78,55 @@ end
         @test vecnorm(UO[:,1:end-1] - UE) < n*vecnorm(UO)*eps()
     end
 end
+
+@testset "Test for sampling through the poles" begin
+    for f in ((θ,φ)->cos(50*cospi(φ)*sinpi(θ)*sinpi(φ)*sinpi(θ)),
+              (θ,φ)->cos(50*cospi(φ)*sinpi(θ)+80*sinpi(φ)*sinpi(θ)),
+              (θ,φ)->sqrt(5+cospi(φ)*sinpi(θ)+exp(sinpi(φ)*sinpi(θ))+sin(cospi(θ))))
+        n = 200
+
+        θ = (0.5:n-0.5)/n
+        φ = (0:2n-2)*2/(2n-1)
+        F = [f(θ,φ) for θ in θ, φ in φ]
+        V = zero(F)
+        A_mul_B!(V, FastTransforms.plan_analysis(F), F)
+        G = zero(V)
+        A_mul_B!(G, FastTransforms.plan_synthesis(V), V)
+
+        θ2 = (0.0:n-1)/(n-1)
+        F2 = [f(θ,φ) for θ in θ2, φ in φ]
+        V2 = zero(F2)
+        A_mul_B!(V2, FastTransforms.plan_analysis2(F2), F2)
+        G2 = zero(V2)
+        A_mul_B!(G2, FastTransforms.plan_synthesis2(V2), V2)
+
+        @test vecnorm(V-V2) < n*vecnorm(V)*eps()
+        @test vecnorm(F-G) < n*vecnorm(F)*eps()
+        @test vecnorm(F2-G2) < n*vecnorm(F)*eps()
+    end
+end
+
+# This test confirms numerically that [P_4(z⋅y) - P_4(x⋅y)]/(z⋅y - x⋅y) is actually a degree-3 polynomial on 𝕊²
+x = [0,0,1]
+y = normalize!([.123,.456,.789])
+
+z = (θ,φ) -> [sinpi(θ)*cospi(φ), sinpi(θ)*sinpi(φ), cospi(θ)]
+
+P4 = x -> (35*x^4-30*x^2+3)/8
+
+n = 5
+θ = (0.5:n-0.5)/n
+φ = (0:2n-2)*2/(2n-1)
+F = [(P4(z(θ,φ)⋅y) - P4(x⋅y))/(z(θ,φ)⋅y - x⋅y) for θ in θ, φ in φ]
+V = zero(F)
+A_mul_B!(V, FastTransforms.plan_analysis(F), F)
+U3 = fourier2sph(V)
+
+# U3 is degree-3
+
+F = [P4(z(θ,φ)⋅y) for θ in θ, φ in φ]
+V = zero(F)
+A_mul_B!(V, FastTransforms.plan_analysis(F), F)
+U4 = fourier2sph(V)
+
+# U4 is degree-4