fix N-dimensional transform when dims[1] != 1

Author: ioannisPApapadopoulos

src/arrays.jl

@@ -31,7 +31,7 @@ function *(P::ArrayPlan, f::AbstractArray)
     perm = [d; setdiff(1:ndims(f), d)]
     fp = permutedims(f, perm)
 
-    fr = reshape(fp, size(fp,1), prod(size(fp)[2:end]))
+    fr = reshape(fp, size(fp,1), :)
 
     permutedims(reshape(F*fr, size(fp)...), inv_perm(perm))
 end
@@ -45,33 +45,39 @@ function \(P::ArrayPlan, f::AbstractArray)
     perm = [d; setdiff(1:ndims(f), d)]
     fp = permutedims(f, perm)
 
-    fr = reshape(fp, size(fp,1), prod(size(fp)[2:end]))
+    fr = reshape(fp, size(fp,1), :)
 
     permutedims(reshape(F\fr, size(fp)...), inv_perm(perm))
 end
 
-struct NDimsPlan{T, FF<:ArrayPlan{<:T}, Dims<:Tuple} <: Plan{T}
+struct NDimsPlan{T, FF<:ArrayPlan{<:T}, Szs<:Tuple, Dims<:Tuple} <: Plan{T}
     F::FF
+    szs::Szs
     dims::Dims
-    function NDimsPlan(F, dims)
-        if length(Set(size(F, dims...))) > 1
+    function NDimsPlan(F, szs, dims)
+        if length(Set(szs[[dims...]])) > 1
             error("Different size in dims axes not yet implemented in N-dimensional transform.")
         end
-        new{eltype(F), typeof(F), typeof(dims)}(F, dims)
+        new{eltype(F), typeof(F), typeof(szs), typeof(dims)}(F, szs, dims)
     end
 end
 
-size(P::NDimsPlan) = size(P.F)
-size(P::NDimsPlan, k::Int) = size(P.F, k)
-size(P::NDimsPlan, k...) = size(P.F, k...)
+size(P::NDimsPlan) = P.szs
+size(P::NDimsPlan, k::Int) = P.szs[k]
+size(P::NDimsPlan, k...) = P.szs[[k...]]
+
+function NDimsPlan(F::FTPlan, szs::Tuple, dims::Tuple)
+    NDimsPlan(ArrayPlan(F, szs, (first(dims),)), szs, dims)
+end
 
 function *(P::NDimsPlan, f::AbstractArray)
     F, dims = P.F, P.dims
-    @assert size(F) == size(f)
+    @assert size(P) == size(f)
     g = copy(f)
     t = 1:ndims(g)
+    d1 = dims[1]
     for d in dims
-        perm = ntuple(k -> k == 1 ? t[d] : k == d ? t[1] : t[k], ndims(g))
+        perm = ntuple(k -> k == d1 ? t[d] : k == d ? t[d1] : t[k], ndims(g))
         gp = permutedims(g, perm)
         g = permutedims(F*gp, inv_perm(perm))
     end
@@ -80,11 +86,12 @@ end
 
 function \(P::NDimsPlan, f::AbstractArray)
     F, dims = P.F, P.dims
-    @assert size(F) == size(f)
+    @assert size(P) == size(f)
     g = copy(f)
     t = 1:ndims(g)
+    d1 = dims[1]
     for d in dims
-        perm = ntuple(k -> k == 1 ? t[d] : k == d ? t[1] : t[k], ndims(g))
+        perm = ntuple(k -> k == d1 ? t[d] : k == d ? t[d1] : t[k], ndims(g))
         gp = permutedims(g, perm)
         g = permutedims(F\gp, inv_perm(perm))
     end

test/arraystests.jl

@@ -29,18 +29,18 @@ import FastTransforms: ArrayPlan, NDimsPlan
 
     @testset "NDimsPlan" begin
         c = randn(20,10,20)
-        @test_throws ErrorException("Different size in dims axes not yet implemented in N-dimensional transform.") NDimsPlan(ArrayPlan(plan_cheb2leg(c), c), (1,2))        
+        @test_throws ErrorException("Different size in dims axes not yet implemented in N-dimensional transform.") NDimsPlan(ArrayPlan(plan_cheb2leg(c), c), size(c), (1,2))        
 
         c = randn(5,20)
         F = plan_cheb2leg(c)
         FT = ArrayPlan(F, c)
-        P = NDimsPlan(FT, (1,))
+        P = NDimsPlan(F, size(c), (1,))
         @test F*c ≈ FT*c ≈ P*c
 
         c = randn(20,20,5);
         F = plan_cheb2leg(c)
         FT = ArrayPlan(F, c)
-        P = NDimsPlan(FT, (1,2))
+        P = NDimsPlan(FT, size(c), (1,2))
 
         @test size(P) == size(c)
         @test size(P,1) == size(c,1)
@@ -53,12 +53,12 @@ import FastTransforms: ArrayPlan, NDimsPlan
         @test f ≈ P*c
         @test c ≈ P\f
 
-        c = randn(10,5,10,60)
-        F = plan_cheb2leg(c)
-        P = NDimsPlan(ArrayPlan(F, c), (1,3))
+        c = randn(5,10,10,60)
+        F = plan_cheb2leg(randn(10))
+        P = NDimsPlan(F, size(c), (2,3))
         f = similar(c)
-        for i in axes(f,2), j in axes(f,4)
-            f[:,i,:,j] = (F*(F*c[:,i,:,j])')'
+        for i in axes(f,1), j in axes(f,4)
+            f[i,:,:,j] = (F*(F*c[i,:,:,j])')'
         end
         @test f ≈ P*c
         @test c ≈ P\f