add thin plan

MikaelSlevinsky · MikaelSlevinsky · commit ad21867a34c4 · 2017-05-10T15:54:41.000-05:00
diff --git a/src/FastTransforms.jl b/src/FastTransforms.jl
@@ -19,6 +19,8 @@ export gaunt
 export paduatransform, ipaduatransform, paduatransform!, ipaduatransform!, paduapoints
 export plan_paduatransform!, plan_ipaduatransform!
 
+export SlowSphericalHarmonicPlan, FastSphericalHarmonicPlan, ThinSphericalHarmonicPlan
+
 # Other module methods and constants:
 #export ChebyshevJacobiPlan, jac2cheb, cheb2jac
 #export sqrtpi, pochhammer, stirlingseries, stirlingremainder, Aratio, Cratio, Anαβ
diff --git a/src/SphericalHarmonics/SphericalHarmonics.jl b/src/SphericalHarmonics/SphericalHarmonics.jl
@@ -1,3 +1,16 @@
+function zero_spurious_modes!(A::AbstractMatrix)
+    M, N = size(A)
+    n = N÷2
+    @inbounds for j = 1:n
+        @simd for i = M-j+1:M
+            A[i,2j] = 0
+            A[i,2j+1] = 0
+        end
+    end
+    A
+end
+
 include("slowplan.jl")
 include("Butterfly.jl")
 include("fastplan.jl")
+include("thinplan.jl")
diff --git a/src/SphericalHarmonics/fastplan.jl b/src/SphericalHarmonics/fastplan.jl
@@ -10,6 +10,8 @@ end
 
 function FastSphericalHarmonicPlan{T}(A::Matrix{T}, L::Int)
     M, N = size(A)
+    @assert ispow2(M)
+    @assert N == 2M-1
     n = (N+1)÷2
     RP = RotationPlan(T, n-1)
     a1 = A[:,1]
@@ -24,16 +26,17 @@ function FastSphericalHarmonicPlan{T}(A::Matrix{T}, L::Int)
     for j = 1:2:n-2
         A_mul_B!(Ce, RP.layers[j])
         BF[j] = orthogonalButterfly(Ce, L)
-        #println("Level: ",j)
+        println("Layer: ",j)
     end
     for j = 2:2:n-2
         A_mul_B!(Co, RP.layers[j])
         BF[j] = orthogonalButterfly(Co, L)
-        #println("Level: ",j)
+        println("Layer: ",j)
     end
     FastSphericalHarmonicPlan(RP, BF, p1, p2, p1inv, p2inv, B)
 end
 
+FastSphericalHarmonicPlan(A::Matrix) = FastSphericalHarmonicPlan(A, round(Int, log2(size(A, 1)+1)-6))
 
 function A_mul_B!(Y::Matrix, FP::FastSphericalHarmonicPlan, X::Matrix)
     RP, BF, p1, p2, B = FP.RP, FP.BF, FP.p1, FP.p2, FP.B
@@ -76,7 +79,7 @@ function At_mul_B!(Y::Matrix, FP::FastSphericalHarmonicPlan, X::Matrix)
         At_mul_B_col_J!(Y, BF[J-1], B, 2J)
         At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
     end
-    Y
+    zero_spurious_modes!(Y)
 end
 
 Ac_mul_B!(Y::Matrix, FP::FastSphericalHarmonicPlan, X::Matrix) = At_mul_B!(Y, FP, X)
diff --git a/src/SphericalHarmonics/slowplan.jl b/src/SphericalHarmonics/slowplan.jl
@@ -176,7 +176,7 @@ function At_mul_B!(Y::Matrix, SP::SlowSphericalHarmonicPlan, X::Matrix)
         A_mul_B_col_J!!(Y, p1inv, B, J)
         A_mul_B_col_J!!(Y, p1inv, B, J+1)
     end
-    At_mul_B!(RP, Y)
+    zero_spurious_modes!(At_mul_B!(RP, Y))
 end
 
 Ac_mul_B!(Y::Matrix, SP::SlowSphericalHarmonicPlan, X::Matrix) = At_mul_B!(Y, SP, X)
diff --git a/src/SphericalHarmonics/thinplan.jl b/src/SphericalHarmonics/thinplan.jl
@@ -0,0 +1,185 @@
+const LAYERSKELETON = 64
+
+checklayer(j::Int) = j÷LAYERSKELETON == j/LAYERSKELETON
+
+immutable ThinSphericalHarmonicPlan{T}
+    RP::RotationPlan{T}
+    BF::Vector{Butterfly{T}}
+    p1::NormalizedLegendreToChebyshevPlan{T}
+    p2::NormalizedLegendre1ToChebyshev2Plan{T}
+    p1inv::ChebyshevToNormalizedLegendrePlan{T}
+    p2inv::Chebyshev2ToNormalizedLegendre1Plan{T}
+    B::Matrix{T}
+end
+
+function ThinSphericalHarmonicPlan{T}(A::Matrix{T}, L::Int)
+    M, N = size(A)
+    @assert ispow2(M)
+    @assert N == 2M-1
+    n = (N+1)÷2
+    RP = RotationPlan(T, n-1)
+    a1 = A[:,1]
+    p1 = plan_normleg2cheb(a1)
+    p2 = plan_normleg12cheb2(a1)
+    p1inv = plan_cheb2normleg(a1)
+    p2inv = plan_cheb22normleg1(a1)
+    B = zeros(A)
+    Ce = eye(T, M)
+    Co = eye(T, M)
+    BF = Vector{Butterfly{T}}(n-2)
+    for j = 1:2:n-2
+        A_mul_B!(Ce, RP.layers[j])
+        checklayer(j+1) && (BF[j] = orthogonalButterfly(Ce, L);println("Layer: ",j))
+    end
+    for j = 2:2:n-2
+        A_mul_B!(Co, RP.layers[j])
+        checklayer(j) && (BF[j] = orthogonalButterfly(Co, L);println("Layer: ",j))
+    end
+    ThinSphericalHarmonicPlan(RP, BF, p1, p2, p1inv, p2inv, B)
+end
+
+ThinSphericalHarmonicPlan(A::Matrix) = ThinSphericalHarmonicPlan(A, round(Int, log2(size(A, 1)+1)-6))
+
+function A_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
+    RP, BF, p1, p2, B = TP.RP, TP.BF, TP.p1, TP.p2, TP.B
+    copy!(B, X)
+    M, N = size(X)
+
+    for j = 3:2:N÷2
+        if checklayer(j-1)
+            A_mul_B_col_J!(Y, BF[j-1], B, 2j)
+            A_mul_B_col_J!(Y, BF[j-1], B, 2j+1)
+        else
+            ℓ = round(Int, (j-1)÷LAYERSKELETON)*LAYERSKELETON
+            A_mul_B_col_J!(RP, B, 2j, ℓ+1, j-1)
+            A_mul_B_col_J!(RP, B, 2j+1, ℓ+1, j-1)
+            if ℓ > LAYERSKELETON-2
+                A_mul_B_col_J!(Y, BF[ℓ], B, 2j)
+                A_mul_B_col_J!(Y, BF[ℓ], B, 2j+1)
+            else
+                copy!(Y, 1+M*(2j-1), B, 1+M*(2j-1), 2M)
+            end
+        end
+    end
+
+    for j = 2:2:N÷2
+        if checklayer(j)
+            A_mul_B_col_J!(Y, BF[j-1], B, 2j)
+            A_mul_B_col_J!(Y, BF[j-1], B, 2j+1)
+        else
+            ℓ = round(Int, j÷LAYERSKELETON)*LAYERSKELETON
+            A_mul_B_col_J!(RP, B, 2j, ℓ, j-1)
+            A_mul_B_col_J!(RP, B, 2j+1, ℓ, j-1)
+            if ℓ > LAYERSKELETON-2
+                A_mul_B_col_J!(Y, BF[ℓ-1], B, 2j)
+                A_mul_B_col_J!(Y, BF[ℓ-1], B, 2j+1)
+            else
+                copy!(Y, 1+M*(2j-1), B, 1+M*(2j-1), 2M)
+            end
+        end
+    end
+
+    copy!(Y, 1, X, 1, 3M)
+    copy!(B, Y)
+    fill!(Y, zero(eltype(Y)))
+
+    A_mul_B_col_J!!(Y, p1, B, 1)
+    for J = 2:4:N
+        A_mul_B_col_J!!(Y, p2, B, J)
+        A_mul_B_col_J!!(Y, p2, B, J+1)
+    end
+    for J = 4:4:N
+        A_mul_B_col_J!!(Y, p1, B, J)
+        A_mul_B_col_J!!(Y, p1, B, J+1)
+    end
+    Y
+end
+
+
+function At_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
+    RP, BF, p1inv, p2inv, B = TP.RP, TP.BF, TP.p1inv, TP.p2inv, TP.B
+    copy!(B, X)
+    M, N = size(X)
+    A_mul_B_col_J!!(Y, p1inv, B, 1)
+    for J = 2:4:N
+        A_mul_B_col_J!!(Y, p2inv, B, J)
+        A_mul_B_col_J!!(Y, p2inv, B, J+1)
+    end
+    for J = 4:4:N
+        A_mul_B_col_J!!(Y, p1inv, B, J)
+        A_mul_B_col_J!!(Y, p1inv, B, J+1)
+    end
+
+    copy!(B, Y)
+    fill!(Y, zero(eltype(Y)))
+    copy!(Y, 1, B, 1, 3M)
+
+    for j = 3:2:N÷2
+        if checklayer(j-1)
+            At_mul_B_col_J!(Y, BF[j-1], B, 2j)
+            At_mul_B_col_J!(Y, BF[j-1], B, 2j+1)
+        else
+            ℓ = round(Int, (j-1)÷LAYERSKELETON)*LAYERSKELETON
+            if ℓ > LAYERSKELETON-2
+                At_mul_B_col_J!(Y, BF[ℓ], B, 2j)
+                At_mul_B_col_J!(Y, BF[ℓ], B, 2j+1)
+            else
+                copy!(Y, 1+M*(2j-1), B, 1+M*(2j-1), 2M)
+            end
+            At_mul_B_col_J!(RP, Y, 2j, ℓ+1, j-1)
+            At_mul_B_col_J!(RP, Y, 2j+1, ℓ+1, j-1)
+        end
+    end
+
+    for j = 2:2:N÷2
+        if checklayer(j)
+            At_mul_B_col_J!(Y, BF[j-1], B, 2j)
+            At_mul_B_col_J!(Y, BF[j-1], B, 2j+1)
+        else
+            ℓ = round(Int, j÷LAYERSKELETON)*LAYERSKELETON
+            if ℓ > LAYERSKELETON-2
+                At_mul_B_col_J!(Y, BF[ℓ-1], B, 2j)
+                At_mul_B_col_J!(Y, BF[ℓ-1], B, 2j+1)
+            else
+                copy!(Y, 1+M*(2j-1), B, 1+M*(2j-1), 2M)
+            end
+            At_mul_B_col_J!(RP, Y, 2j, ℓ, j-1)
+            At_mul_B_col_J!(RP, Y, 2j+1, ℓ, j-1)
+        end
+    end
+
+    zero_spurious_modes!(Y)
+end
+
+Ac_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix) = At_mul_B!(Y, TP, X)
+
+allranks(TP::ThinSphericalHarmonicPlan) = mapreduce(i->allranks(TP.BF[i]),vcat,sort!([LAYERSKELETON-1:LAYERSKELETON:length(TP.BF);LAYERSKELETON:LAYERSKELETON:length(TP.BF)]))
+
+
+function A_mul_B_col_J!(P::RotationPlan, A::AbstractMatrix, J::Int, L1::Int, L2::Int)
+    M, N = size(A)
+    @inbounds for m = L2-1:-2:L1
+        layer = P.layers[m+1]
+        @simd for i = 1:length(layer)
+            G = layer[i]
+            a1, a2 = A[G.i1,J], A[G.i2,J]
+            A[G.i1,J] = G.c*a1 + G.s*a2
+            A[G.i2,J] = G.c*a2 - G.s*a1
+        end
+    end
+    A
+end
+
+function At_mul_B_col_J!(P::RotationPlan, A::AbstractMatrix, J::Int, L1::Int, L2::Int)
+    M, N = size(A)
+    @inbounds for m = L1:2:L2-1
+        layer = P.layers[m+1]
+        @simd for i = length(layer):-1:1
+            G = layer[i]
+            a1, a2 = A[G.i1,J], A[G.i2,J]
+            A[G.i1,J] = G.c*a1 - G.s*a2
+            A[G.i2,J] = G.c*a2 + G.s*a1
+        end
+    end
+    A
+end
diff --git a/test/butterflytests.jl b/test/butterflytests.jl
@@ -79,7 +79,7 @@ for n in 7:N
 end
 
 
-N = 14
+N = 12
 A = Vector{Matrix{Complex{Float64}}}(N)
 B = Vector{Butterfly{Complex{Float64}}}(N)
 for n in 7:N
diff --git a/test/sphericalharmonictests.jl b/test/sphericalharmonictests.jl
@@ -0,0 +1,64 @@
+using FastTransforms, Base.Test
+
+function sphrand{T}(::Type{T}, m, n)
+    A = zeros(T, m, 2n-1)
+    for i = 1:m
+        A[i,1] = rand(T)
+    end
+    for j = 1:n
+        for i = 1:m-j
+            A[i,2j] = rand(T)
+            A[i,2j+1] = rand(T)
+        end
+    end
+    A
+end
+
+function sphrandn{T}(::Type{T}, m, n)
+    A = zeros(T, m, 2n-1)
+    for i = 1:m
+        A[i,1] = randn(T)
+    end
+    for j = 1:n
+        for i = 1:m-j
+            A[i,2j] = randn(T)
+            A[i,2j+1] = randn(T)
+        end
+    end
+    A
+end
+
+function normalizecolumns!(A::AbstractMatrix)
+    m, n = size(A)
+    @inbounds for j = 1:n
+        nrm = zero(eltype(A))
+        for i = 1:m
+            nrm += abs2(A[i,j])
+        end
+        nrm = sqrt(nrm)
+        for i = 1:m
+            A[i,j] /= nrm
+        end
+    end
+    A
+end
+
+function maxcolnorm(A::AbstractMatrix)
+    m, n = size(A)
+    nrm = zeros(n)
+    @inbounds for j = 1:n
+        nrm[j] = 0
+        for i = 1:m
+            nrm[j] += abs2(A[i,j])
+        end
+        nrm[j] = sqrt(nrm[j])
+    end
+    norm(nrm, Inf)
+end
+
+println("Testing slow plan")
+include("test_slowplan.jl")
+println("Testing fast plan")
+include("test_fastplan.jl")
+println("Testing thin plan")
+include("test_thinplan.jl")
diff --git a/test/test_fastplan.jl b/test/test_fastplan.jl
diff --git a/test/test_slowplan.jl b/test/test_slowplan.jl
diff --git a/test/test_thinplan.jl b/test/test_thinplan.jl
