factorize -> plan_transform for Fourier (#117)

* factorize -> plan_transform for Fourier

* transforms for annuli

* Laplace's eqn

* Add more examples, move MulPlan

Author: dlfivefifty

Project.toml

@@ -29,11 +29,11 @@ ArrayLayouts = "0.8"
 BandedMatrices = "0.17"
 BlockArrays = "0.16.9"
 BlockBandedMatrices = "0.11.6"
-ContinuumArrays = "0.12.1"
+ContinuumArrays = "0.12.3"
 DomainSets = "0.5.6, 0.6"
 FFTW = "1.1"
 FastGaussQuadrature = "0.4.3, 0.5"
-FastTransforms = "0.14.4"
+FastTransforms = "0.14.9"
 FillArrays = "0.13"
 HypergeometricFunctions = "0.3.4"
 InfiniteArrays = "0.12.3"

examples/annulipdes.jl

@@ -1,44 +1,89 @@
 using ClassicalOrthogonalPolynomials, Plots
+plotly()
 
-ρ = 0.5; T = chebyshevt(ρ..1); U = chebyshevu(T); C = ultraspherical(2, ρ..1); r = axes(T,1); D = Derivative(r);
+###
+# Laplace's equation
+###
+
+ρ = 0.5
+T,C,F = chebyshevt(ρ..1),ultraspherical(2, ρ..1),Fourier()
+r = axes(T,1)
+D = Derivative(r)
 
 L = C \ (r.^2 .* (D^2 * T)) + C \ (r .* (D * T)) # r^2 * ∂^2 + r*∂
 M = C\T # Identity
+R = C \ (r .* C) # mult by r
+
+uᵨ = transform(F, θ -> exp(cos(θ)+sin(θ-1)))
+u₁ = transform(F, θ -> cos(100cos(θ)-sin(θ+1)-1))
+
+n = 300
+X = zeros(n,n)
+
+for j = 1:n
+    m = j ÷ 2
+    Δₘ = L - m^2*M # r^2 * Laplacian for exp(im*m*θ)*u(r), i.e. (r^2 * ∂^2 + r*∂ - m^2*I)*u
+    d = [T[[begin,end],:]; Δₘ] \ [uᵨ[j]; u₁[j]; zeros(∞)]
+    X[:,j] .= d[1:n]
+end
 
-m = 2
-Δₘ = L - m^2*M # r^2 * Laplacian for exp(im*m*θ)*u(r), i.e. (r^2 * ∂^2 + r*∂ - m^2*I)*u
+𝐫,𝛉 = ClassicalOrthogonalPolynomials.grid(T, n),ClassicalOrthogonalPolynomials.grid(F, n)
+PT,PF = plan_transform(T, (n,n), 1),plan_transform(F, (n,n), 2)
+U = PT \ (PF \ X); U = [U U[:,1]]
 
+𝐱 = 𝐫 .* cos.([𝛉; 2π]')
+𝐲 = 𝐫 .* sin.([𝛉; 2π]')
+surface(𝐱, 𝐲, U; zlims=(-2,2))
 
 
-# Poisson solve
-c = C \ exp.(r)
-R = C \ (r .* C)
+###
+# Helmholtz equation
+###
 
+uᵨ = transform(F, θ -> exp(cos(θ)+sin(θ-1)))
+u₁ = transform(F, θ -> cos(cos(θ)-sin(θ+1)-1))
 
-d = [T[[begin,end],:];
-            Δₘ] \ [1; 2; R^2 * c]
 
-plot(T*d)
+k = 100
 
+for j = 1:n
+    m = j ÷ 2
+    Δₘ = L - m^2*M + k^2 * R^2*M
+    d = [T[[begin,end],:]; Δₘ] \ [uᵨ[j]; u₁[j]; zeros(∞)]
+    X[:,j] .= d[1:n]
+end
 
-# Helmholtz
+U = PT \ (PF \ X); U = [U U[:,1]]
+surface(𝐱, 𝐲, U; zlims=(-10,10))
 
 
-Q = R^2 * M # r^2, needed for Helmholtz (Δ + k^2)*u = f
+###
+# Poisson
+###
 
-k = 5 # f
+n = 1000
+𝐫,𝛉 = ClassicalOrthogonalPolynomials.grid(T, n),ClassicalOrthogonalPolynomials.grid(F, n)
+PT,PF = plan_transform(T, (n,n), 1),plan_transform(F, (n,n), 2)
+f = (x,y) -> exp(-4000((x-0.8)^2 + (y-0.1)^2))
+𝐱 = 𝐫 .* cos.(𝛉')
+𝐲 = 𝐫 .* sin.(𝛉')
 
-d = [T[[begin,end],:];
-            Δₘ+k^2*Q] \ [1; 2; R^2 * c]
+F = PT * (PF * f.(𝐱, 𝐲))
 
 
-# transform
 
-f = (r,θ) -> exp(r*cos(θ))
+X = zeros(n+2, n+2)
 
-n = 10 # create a 10 x 10 transform
+# multiply RHS by r^2 and convert to C
+S = (R^2*M)[1:n,1:n]
 
-F = Fourier()
-𝐫,𝛉 = grid(T, n),grid(F, n)
+for j = 1:n
+    m = j ÷ 2
+    Δₘ = L - m^2*M
+    X[:,j] = [T[[begin,end],:]; Δₘ][1:n+2,1:n+2] \ [0; 0; S*F[:,j]]
+end
 
-transform(T, transform(F, f.(𝐫, 𝛉'); dims=2); dims=1)
+U = PT \ (PF \ X[1:n,1:n]); U = [U U[:,1]]
+𝐱 = 𝐫 .* cos.([𝛉; 2π]')
+𝐲 = 𝐫 .* sin.([𝛉; 2π]')
+surface(𝐱, 𝐲, U)

src/ClassicalOrthogonalPolynomials.jl

@@ -38,9 +38,9 @@ import ContinuumArrays: Basis, Weight, basis, @simplify, Identity, AbstractAffin
     inbounds_getindex, grid, plotgrid, _plotgrid, _grid, transform_ldiv, TransformFactorization, QInfAxes, broadcastbasis, ExpansionLayout, basismap,
     AffineQuasiVector, AffineMap, WeightLayout, AbstractWeightedBasisLayout, WeightedBasisLayout, WeightedBasisLayouts, demap, AbstractBasisLayout, BasisLayout,
     checkpoints, weight, unweighted, MappedBasisLayouts, __sum, invmap, plan_ldiv, layout_broadcasted, MappedBasisLayout, SubBasisLayout, _broadcastbasis,
-    plan_transform, plan_grid_transform, MAX_PLOT_POINTS
+    plan_transform, plan_grid_transform, MAX_PLOT_POINTS, MulPlan
 import FastTransforms: Λ, forwardrecurrence, forwardrecurrence!, _forwardrecurrence!, clenshaw, clenshaw!,
-                        _forwardrecurrence_next, _clenshaw_next, check_clenshaw_recurrences, ChebyshevGrid, chebyshevpoints, Plan
+                        _forwardrecurrence_next, _clenshaw_next, check_clenshaw_recurrences, ChebyshevGrid, chebyshevpoints, Plan, ScaledPlan
 
 import FastGaussQuadrature: jacobimoment
 
@@ -262,68 +262,14 @@ function golubwelsch(V::SubQuasiArray)
     x,w
 end
 
-"""
-    MulPlan(matrix, dims)
-
-Takes a matrix and supports it applied to different dimensions.
-"""
-struct MulPlan{T, Fact, Dims} # <: Plan{T} We don't depend on AbstractFFTs
-    matrix::Fact
-    dims::Dims
-end
-
-MulPlan(fact, dims) = MulPlan{eltype(fact), typeof(fact), typeof(dims)}(fact, dims)
-
-function *(P::MulPlan{<:Any,<:Any,Int}, x::AbstractVector)
-    @assert P.dims == 1
-    P.matrix * x
-end
-
-function *(P::MulPlan{<:Any,<:Any,Int}, X::AbstractMatrix)
-    if P.dims == 1
-        P.matrix * X
-    else
-        @assert P.dims == 2
-        permutedims(P.matrix * permutedims(X))
-    end
-end
-
-function *(P::MulPlan{<:Any,<:Any,Int}, X::AbstractArray{<:Any,3})
-    Y = similar(X)
-    if P.dims == 1
-        for j in axes(X,3)
-            Y[:,:,j] = P.matrix * X[:,:,j]
-        end
-    elseif P.dims == 2
-        for k in axes(X,1)
-            Y[k,:,:] = P.matrix * X[k,:,:]
-        end
-    else
-        @assert P.dims == 3
-        for k in axes(X,1), j in axes(X,2)
-            Y[k,j,:] = P.matrix * X[k,j,:]
-        end
-    end
-    Y
-end
-
-function *(P::MulPlan, X::AbstractArray)
-    for d in P.dims
-        X = MulPlan(P.matrix, d) * X
-    end
-    X
-end
-
-*(A::AbstractMatrix, P::MulPlan) = MulPlan(A*P.matrix, P.dims)
-
 
 function plan_grid_transform(Q::Normalized, szs::NTuple{N,Int}, dims=1:N) where N
     L = Q[:,OneTo(szs[1])]
     x,w = golubwelsch(L)
     x, MulPlan(L[x,:]'*Diagonal(w), dims)
 end
 
-function plan_grid_transform(P::OrthogonalPolynomial, szs::NTuple{N,Int}, dims=1:N) where N
+function plan_grid_transform(::AbstractOPLayout, P, szs::NTuple{N,Int}, dims=1:N) where N
     Q = Normalized(P)
     x, A = plan_grid_transform(Q, szs, dims...)
     n = szs[1]
@@ -332,6 +278,9 @@ function plan_grid_transform(P::OrthogonalPolynomial, szs::NTuple{N,Int}, dims=1
 end
 
 
+plan_grid_transform(::MappedOPLayout, L, szs::NTuple{N,Int}, dims=1:N) where N =
+    plan_grid_transform(MappedBasisLayout(), L, szs, dims)
+
 function \(A::SubQuasiArray{<:Any,2,<:OrthogonalPolynomial}, B::SubQuasiArray{<:Any,2,<:OrthogonalPolynomial})
     axes(A,1) == axes(B,1) || throw(DimensionMismatch())
     _,jA = parentindices(A)

src/classical/fourier.jl

@@ -37,6 +37,14 @@ struct ShuffledR2HC{T,Pl<:Plan} <: AbstractShuffledPlan{T}
     plan::Pl
 end
 
+"""
+Gives a shuffled version of the real IFFT, with order
+1,sin(θ),cos(θ),sin(2θ)…
+"""
+struct ShuffledIR2HC{T,Pl<:Plan} <: AbstractShuffledPlan{T}
+    plan::Pl
+end
+
 """
 Gives a shuffled version of the FFT, with order
 1,sin(θ),cos(θ),sin(2θ)…
@@ -45,6 +53,14 @@ struct ShuffledFFT{T,Pl<:Plan} <: AbstractShuffledPlan{T}
     plan::Pl
 end
 
+"""
+Gives a shuffled version of the IFFT, with order
+1,sin(θ),cos(θ),sin(2θ)…
+"""
+struct ShuffledIFFT{T,Pl<:Plan} <: AbstractShuffledPlan{T}
+    plan::Pl
+end
+
 size(F::AbstractShuffledPlan, k) = size(F.plan,k)
 size(F::AbstractShuffledPlan) = size(F.plan)
 
@@ -54,84 +70,70 @@ ShuffledR2HC{T}(n, d...) where T = ShuffledR2HC{T}(FFTW.plan_r2r(Array{T}(undef,
 ShuffledFFT{T}(p::Pl) where {T,Pl<:Plan} = ShuffledFFT{T,Pl}(p)
 ShuffledFFT{T}(n, d...) where T = ShuffledFFT{T}(FFTW.plan_fft(Array{T}(undef, n), d...))
 
+ShuffledIFFT{T}(p::Pl) where {T,Pl<:Plan} = ShuffledIFFT{T,Pl}(p)
+ShuffledIR2HC{T}(p::Pl) where {T,Pl<:Plan} = ShuffledIR2HC{T,Pl}(p)
 
-function _shuffledR2HC_postscale!(_, ret::AbstractVector{T}) where T
+inv(P::ShuffledR2HC{T}) where T = ShuffledIR2HC{T}(inv(P.plan))
+inv(P::ShuffledFFT{T}) where T = ShuffledIFFT{T}(inv(P.plan))
+
+
+_shuffled_prescale!(_, ret::AbstractVector) = ret
+_shuffled_postscale!(_, ret::AbstractVector) = ret
+
+function _shuffled_postscale!(::ShuffledR2HC, ret::AbstractVector{T}) where T
     n = length(ret)
     lmul!(convert(T,2)/n, ret)
     ret[1] /= 2
     iseven(n) && (ret[n÷2+1] /= 2)
     negateeven!(reverseeven!(interlace!(ret,1)))
 end
 
-function _shuffledR2HC_postscale!(d::Number, ret::AbstractMatrix{T}) where T
-    if isone(d)
-        n = size(ret,1)
-        lmul!(convert(T,2)/n, ret)
-        ldiv!(2, view(ret,1,:))
-        iseven(n) && ldiv!(2, view(ret,n÷2+1,:))
-        for j in axes(ret,2)
-            negateeven!(reverseeven!(interlace!(view(ret,:,j),1)))
-        end
-    else
-        n = size(ret,2)
-        lmul!(convert(T,2)/n, ret)
-        ldiv!(2, view(ret,:,1))
-        iseven(n) && ldiv!(2, view(ret,:,n÷2+1))
-        for k in axes(ret,1)
-            negateeven!(reverseeven!(interlace!(view(ret,k,:),1)))
-        end
-    end
-    ret
+function _shuffled_prescale!(::ShuffledIR2HC, ret::AbstractVector{T}) where T
+    n = length(ret)
+    reverseeven!(negateeven!(ret))
+    ret .= [ret[1:2:end]; ret[2:2:end]] # todo: non-allocating
+    iseven(n) && (ret[n÷2+1] *= 2)
+    ret[1] *= 2
+    lmul!(convert(T,n)/2, ret)
 end
 
-function _shuffledFFT_postscale!(_, ret::AbstractVector{T}) where T
+function _shuffled_postscale!(::ShuffledFFT, ret::AbstractVector{T}) where T
     n = length(ret)
     cfs = lmul!(inv(convert(T,n)), ret)
     reverseeven!(interlace!(cfs,1))
 end
 
-function _shuffledFFT_postscale!(d::Number, ret::AbstractMatrix{T}) where T
-    if isone(d)
-        n = size(ret,1)
-        lmul!(inv(convert(T,n)), ret)
-        for j in axes(ret,2)
-            reverseeven!(interlace!(view(ret,:,j),1))
-        end
-    else
-        n = size(ret,2)
-        lmul!(inv(convert(T,n)), ret)
-        for k in axes(ret,1)
-            reverseeven!(interlace!(view(ret,k,:),1))
+_region(F::Plan) = F.region
+_region(F::ScaledPlan) = F.p.region
+
+for func in (:_shuffled_postscale!, :_shuffled_prescale!)
+    @eval function $func(F::AbstractShuffledPlan, ret::AbstractMatrix{T}) where T
+        d = _region(F.plan)
+        if isone(d)
+            for j in axes(ret,2)
+                $func(F, view(ret,:,j))
+            end
+        else
+            @assert d == 2
+            for k in axes(ret,1)
+                $func(F, view(ret,k,:))
+            end
         end
+        ret
     end
-    ret
 end
 
-function mul!(ret::AbstractArray{T}, F::ShuffledR2HC{T}, b::AbstractArray) where T
-    mul!(ret, F.plan, convert(Array{T}, b))
-    _shuffledR2HC_postscale!(F.plan.region, ret)
+function mul!(ret::AbstractArray{T}, F::AbstractShuffledPlan{T}, bin::AbstractArray) where T
+    b = _shuffled_prescale!(F, Array{T}(bin))
+    mul!(ret, F.plan, b)
+    _shuffled_postscale!(F, ret)
 end
 
-function mul!(ret::AbstractArray{T}, F::ShuffledFFT{T}, b::AbstractArray) where T
-    mul!(ret, F.plan, convert(Array{T}, b))
-    _shuffledFFT_postscale!(F.plan.region, ret)
-end
 
 *(F::AbstractShuffledPlan{T}, b::AbstractVecOrMat) where T = mul!(similar(b, T), F, b)
 
-factorize(L::SubQuasiArray{T,2,<:Fourier,<:Tuple{<:Inclusion,<:OneTo}}) where T =
-    TransformFactorization(grid(L), ShuffledR2HC{T}(size(L,2)))
-factorize(L::SubQuasiArray{T,2,<:Fourier,<:Tuple{<:Inclusion,<:OneTo}}, d) where T =
-    TransformFactorization(grid(L), ShuffledR2HC{T}((size(L,2),d),1))
-
-factorize(L::SubQuasiArray{T,2,<:Laurent,<:Tuple{<:Inclusion,<:OneTo}}) where T =
-    TransformFactorization(grid(L), ShuffledFFT{T}(size(L,2)))
-factorize(L::SubQuasiArray{T,2,<:Laurent,<:Tuple{<:Inclusion,<:OneTo}}, d) where T =
-    TransformFactorization(grid(L), ShuffledFFT{T}((size(L,2),d),1))
-
-
-factorize(L::SubQuasiArray{T,2,<:AbstractFourier,<:Tuple{<:Inclusion,<:BlockSlice}},d...) where T =
-    ProjectionFactorization(factorize(parent(L)[:,OneTo(size(L,2))],d...),parentindices(L)[2])
+plan_grid_transform(F::Fourier{T}, szs::NTuple{N,Int}, dims...) where {T,N} = grid(F, szs[1]), ShuffledR2HC{T}(szs, dims...)
+plan_grid_transform(F::Laurent{T}, szs::NTuple{N,Int}, dims...) where {T,N} = grid(F, szs[1]), ShuffledFFT{T}(szs, dims...)
 
 import BlockBandedMatrices: _BlockSkylineMatrix