Skip to content

Commit 3b8a1b2

Browse files
dannys4dannys4
authored
Merge pull request #40 from dannys4/dannys4/issue22
Add support for AbstractFFTs Interface
2 parents b213f97 + 8ba9289 commit 3b8a1b2

File tree

8 files changed

+196
-141
lines changed

8 files changed

+196
-141
lines changed

‎Project.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,14 @@ authors = ["Danny Sharp <[email protected]> and contributors"]
44
version = "0.2.2"
55

66
[deps]
7+
AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
78
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
89
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
10+
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
911
MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
1012
Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
1113
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
14+
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
1215

1316
[compat]
1417
DocStringExtensions = "0.8"

‎src/FFTA.jl

Lines changed: 9 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -1,72 +1,15 @@
11
module FFTA
22

3-
using Primes, DocStringExtensions, MuladdMacro
4-
export fft, bfft
3+
using Primes, DocStringExtensions, Reexport, MuladdMacro, LinearAlgebra
4+
@reexport using AbstractFFTs
5+
6+
import AbstractFFTs: Plan
57

68
include("callgraph.jl")
79
include("algos.jl")
10+
include("plan.jl")
811

9-
"""
10-
$(TYPEDSIGNATURES)
11-
Perform a fast Fourier transform of a vector. Preserves types given by the
12-
user.
13-
14-
# Arguments
15-
x::AbstractVector: The vector to transform.
16-
17-
# Examples
18-
```julia
19-
julia> x = rand(ComplexF64, 10)
20-
21-
julia> y = fft(x)
22-
```
23-
"""
24-
function fft(x::AbstractVector{T}) where {T}
25-
y = similar(x)
26-
g = CallGraph{T}(length(x))
27-
fft!(y, x, 1, 1, FFT_FORWARD, g[1].type, g, 1)
28-
y
29-
end
30-
31-
function fft(x::AbstractVector{T}) where {T <: Real}
32-
y = similar(x, Complex{T})
33-
g = CallGraph{Complex{T}}(length(x))
34-
fft!(y, x, 1, 1, FFT_FORWARD, g[1].type, g, 1)
35-
y
36-
end
37-
38-
"""
39-
$(TYPEDSIGNATURES)
40-
Perform a fast Fourier transform of a matrix. Preserves types given by the
41-
user.
42-
43-
# Arguments
44-
x::AbstractMatrix: The matrix to transform (columnwise then rowwise).
45-
46-
# Examples
47-
```julia
48-
julia> x = rand(ComplexF64, 10, 10)
49-
50-
julia> y = fft(x)
51-
```
52-
"""
53-
function fft(x::AbstractMatrix{T}) where {T}
54-
M,N = size(x)
55-
y1 = similar(x)
56-
y2 = similar(x)
57-
g1 = CallGraph{T}(size(x,1))
58-
g2 = CallGraph{T}(size(x,2))
59-
60-
for k in 1:N
61-
@views fft!(y1[:,k], x[:,k], 1, 1, FFT_FORWARD, g1[1].type, g1, 1)
62-
end
63-
64-
for k in 1:M
65-
@views fft!(y2[k,:], y1[k,:], 1, 1, FFT_FORWARD, g2[1].type, g2, 1)
66-
end
67-
y2
68-
end
69-
12+
#=
7013
function fft(x::AbstractMatrix{T}) where {T <: Real}
7114
M,N = size(x)
7215
y1 = similar(x, Complex{T})
@@ -84,79 +27,13 @@ function fft(x::AbstractMatrix{T}) where {T <: Real}
8427
y2
8528
end
8629
87-
"""
88-
$(TYPEDSIGNATURES)
89-
Perform a backward fast Fourier transform of a vector, where "backward"
90-
indicates the same output signal down to a constant factor. Preserves types
91-
given by the user.
92-
93-
# Arguments
94-
x::AbstractVector: The vector to transform
95-
96-
# Examples
97-
```julia
98-
julia> x = rand(ComplexF64, 10)
99-
100-
julia> y = bfft(x)
101-
102-
julia> z = fft(y)
103-
104-
julia> x ≈ z/10
105-
true
106-
```
107-
"""
108-
function bfft(x::AbstractVector{T}) where {T}
109-
y = similar(x)
110-
g = CallGraph{T}(length(x))
111-
fft!(y, x, 1, 1, FFT_BACKWARD, g[1].type, g, 1)
112-
y
113-
end
114-
11530
function bfft(x::AbstractVector{T}) where {T <: Real}
11631
y = similar(x, Complex{T})
11732
g = CallGraph{Complex{T}}(length(x))
11833
fft!(y, x, 1, 1, FFT_BACKWARD, g[1].type, g, 1)
11934
y
12035
end
12136
122-
"""
123-
$(TYPEDSIGNATURES)
124-
Perform a backward fast Fourier transform of a matrix, where "backward"
125-
indicates the same output signal down to a constant factor. Preserves types
126-
given by the user.
127-
128-
# Arguments
129-
x::AbstractMatrix: The matrix to transform
130-
131-
# Examples
132-
```julia
133-
julia> x = rand(ComplexF64, 10, 10)
134-
135-
julia> y = bfft(x)
136-
137-
julia> z = fft(y)
138-
139-
julia> x ≈ z/100
140-
true
141-
```
142-
"""
143-
function bfft(x::AbstractMatrix{T}) where {T}
144-
M,N = size(x)
145-
y1 = similar(x)
146-
y2 = similar(x)
147-
g1 = CallGraph{T}(size(x,1))
148-
g2 = CallGraph{T}(size(x,2))
149-
150-
for k in 1:N
151-
@views fft!(y1[:,k], x[:,k], 1, 1, FFT_BACKWARD, g1[1].type, g1, 1)
152-
end
153-
154-
for k in 1:M
155-
@views fft!(y2[k,:], y1[k,:], 1, 1, FFT_BACKWARD, g2[1].type, g2, 1)
156-
end
157-
y2
158-
end
159-
16037
function bfft(x::AbstractMatrix{T}) where {T <: Real}
16138
M,N = size(x)
16239
y1 = similar(x, Complex{T})
@@ -172,6 +49,8 @@ function bfft(x::AbstractMatrix{T}) where {T <: Real}
17249
@views fft!(y2[k,:], y1[k,:], 1, 1, FFT_BACKWARD, g2[1].type, g2, 1)
17350
end
17451
y2
175-
end
52+
end =#
53+
54+
17655

17756
end

‎src/plan.jl

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
import Base: *
2+
import LinearAlgebra: mul!
3+
4+
abstract type FFTAPlan{T,N} <: Plan{T} end
5+
6+
struct FFTAInvPlan{T,N} <: FFTAPlan{T,N} end
7+
8+
struct FFTAPlan_cx{T,N} <: FFTAPlan{T,N}
9+
callgraph::NTuple{N, CallGraph{T}}
10+
region::Union{Int,AbstractVector{<:Int}}
11+
dir::Direction
12+
pinv::FFTAInvPlan{T}
13+
end
14+
15+
struct FFTAPlan_re{T,N} <: FFTAPlan{T,N}
16+
callgraph::NTuple{N, CallGraph{T}}
17+
region::Union{Int,AbstractVector{<:Int}}
18+
dir::Direction
19+
pinv::FFTAInvPlan{T}
20+
flen::Int
21+
end
22+
23+
function AbstractFFTs.plan_fft(x::AbstractArray{T}, region; kwargs...)::FFTAPlan_cx{T} where {T <: Complex}
24+
N = length(region)
25+
@assert N <= 2 "Only supports vectors and matrices"
26+
if N == 1
27+
g = CallGraph{T}(size(x,region[]))
28+
pinv = FFTAInvPlan{T,N}()
29+
return FFTAPlan_cx{T,N}((g,), region, FFT_FORWARD, pinv)
30+
else
31+
sort!(region)
32+
g1 = CallGraph{T}(size(x,region[1]))
33+
g2 = CallGraph{T}(size(x,region[2]))
34+
pinv = FFTAInvPlan{T,N}()
35+
return FFTAPlan_cx{T,N}((g1,g2), region, FFT_FORWARD, pinv)
36+
end
37+
end
38+
39+
function AbstractFFTs.plan_bfft(x::AbstractArray{T}, region; kwargs...)::FFTAPlan_cx{T} where {T <: Complex}
40+
N = length(region)
41+
@assert N <= 2 "Only supports vectors and matrices"
42+
if N == 1
43+
g = CallGraph{T}(size(x,region[]))
44+
pinv = FFTAInvPlan{T,N}()
45+
return FFTAPlan_cx{T,N}((g,), region, FFT_BACKWARD, pinv)
46+
else
47+
sort!(region)
48+
g1 = CallGraph{T}(size(x,region[1]))
49+
g2 = CallGraph{T}(size(x,region[2]))
50+
pinv = FFTAInvPlan{T,N}()
51+
return FFTAPlan_cx{T,N}((g1,g2), region, FFT_BACKWARD, pinv)
52+
end
53+
end
54+
55+
function AbstractFFTs.plan_rfft(x::AbstractArray{T}, region; kwargs...)::FFTAPlan_re{Complex{T}} where {T <: Real}
56+
N = length(region)
57+
@assert N <= 2 "Only supports vectors and matrices"
58+
if N == 1
59+
g = CallGraph{Complex{T}}(size(x,region[]))
60+
pinv = FFTAInvPlan{Complex{T},N}()
61+
return FFTAPlan_re{Complex{T},N}(tuple(g), region, FFT_FORWARD, pinv, size(x,region[]))
62+
else
63+
sort!(region)
64+
g1 = CallGraph{Complex{T}}(size(x,region[1]))
65+
g2 = CallGraph{Complex{T}}(size(x,region[2]))
66+
pinv = FFTAInvPlan{Complex{T},N}()
67+
return FFTAPlan_re{Complex{T},N}(tuple(g1,g2), region, FFT_FORWARD, pinv, size(x,region[1]))
68+
end
69+
end
70+
71+
function AbstractFFTs.plan_brfft(x::AbstractArray{T}, len, region; kwargs...)::FFTAPlan_re{T} where {T}
72+
N = length(region)
73+
@assert N <= 2 "Only supports vectors and matrices"
74+
if N == 1
75+
g = CallGraph{T}(len)
76+
pinv = FFTAInvPlan{T,N}()
77+
return FFTAPlan_re{T,N}((g,), region, FFT_BACKWARD, pinv, len)
78+
else
79+
sort!(region)
80+
g1 = CallGraph{T}(len)
81+
g2 = CallGraph{T}(size(x,region[2]))
82+
pinv = FFTAInvPlan{T,N}()
83+
return FFTAPlan_re{T,N}((g1,g2), region, FFT_BACKWARD, pinv, len)
84+
end
85+
end
86+
87+
function AbstractFFTs.plan_bfft(p::FFTAPlan_cx{T,N}) where {T,N}
88+
return FFTAPlan_cx{T,N}(p.callgraph, p.region, -p.dir, p.pinv)
89+
end
90+
91+
function AbstractFFTs.plan_brfft(p::FFTAPlan_re{T,N}) where {T,N}
92+
return FFTAPlan_re{T,N}(p.callgraph, p.region, -p.dir, p.pinv, p.flen)
93+
end
94+
95+
function LinearAlgebra.mul!(y::AbstractVector{U}, p::FFTAPlan{T,1}, x::AbstractVector{T}) where {T,U}
96+
fft!(y, x, 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
97+
end
98+
99+
function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan{T,1}, x::AbstractArray{T,N}) where {T,U,N}
100+
Rpre = CartesianIndices(size(x)[1:p.region-1])
101+
Rpost = CartesianIndices(size(x)[p.region+1:end])
102+
for Ipre in Rpre
103+
for Ipost in Rpost
104+
@views fft!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
105+
end
106+
end
107+
end
108+
109+
function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan{T,2}, x::AbstractArray{T,N}) where {T,U,N}
110+
R1 = CartesianIndices(size(x)[1:p.region[1]-1])
111+
R2 = CartesianIndices(size(x)[p.region[1]+1:p.region[2]-1])
112+
R3 = CartesianIndices(size(x)[p.region[2]+1:end])
113+
y_tmp = similar(y, axes(y)[p.region])
114+
rows,cols = size(x)[p.region]
115+
for I1 in R1
116+
for I2 in R2
117+
for I3 in R3
118+
for k in 1:cols
119+
@views fft!(y_tmp[:,k], x[I1,:,I2,k,I3], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
120+
end
121+
122+
for k in 1:rows
123+
@views fft!(y[I1,k,I2,:,I3], y_tmp[k,:], 1, 1, p.dir, p.callgraph[2][1].type, p.callgraph[2], 1)
124+
end
125+
end
126+
end
127+
end
128+
end
129+
130+
function *(p::FFTAPlan{T,1}, x::AbstractVector{T}) where {T<:Union{Real,Complex}}
131+
y = similar(x, T <: Real ? Complex{T} : T)
132+
LinearAlgebra.mul!(y, p, x)
133+
y
134+
end
135+
136+
function *(p::FFTAPlan{T,N1}, x::AbstractArray{T,N2}) where {T<:Union{Real, Complex}, N1, N2}
137+
y = similar(x, T <: Real ? Complex{T} : T)
138+
LinearAlgebra.mul!(y, p, x)
139+
y
140+
end
141+
142+
function *(p::FFTAPlan_re{T,1}, x::AbstractVector{T}) where {T<:Union{Real, Complex}}
143+
if p.dir == FFT_FORWARD
144+
y = similar(x, T <: Real ? Complex{T} : T)
145+
LinearAlgebra.mul!(y, p, x)
146+
return y[1:end÷2 + 1]
147+
else
148+
x_tmp = similar(x, p.flen)
149+
x_tmp[1:end÷2 + 1] .= x
150+
x_tmp[end÷2 + 2:end] .= iseven(p.flen) ? conj.(x[end-1:-1:2]) : conj.(x[end:-1:2])
151+
y = similar(x_tmp)
152+
LinearAlgebra.mul!(y, p, x_tmp)
153+
return y
154+
end
155+
end
156+
157+
function *(p::FFTAPlan_re{T,N}, x::AbstractArray{T,2}) where {T<:Union{Real, Complex}, N}
158+
if p.dir == FFT_FORWARD
159+
y = similar(x, T <: Real ? Complex{T} : T)
160+
LinearAlgebra.mul!(y, p, x)
161+
return y[1:end÷2 + 1,:]
162+
else
163+
x_tmp = similar(x, p.flen, size(x)[2])
164+
x_tmp[1:end÷2 + 1,:] .= x
165+
x_tmp[end÷2 + 2:end,:] .= iseven(p.flen) ? conj.(x[end-1:-1:2,:]) : conj.(x[end:-1:2,:])
166+
y = similar(x_tmp)
167+
LinearAlgebra.mul!(y, p, x_tmp)
168+
return y
169+
end
170+
end

‎test/onedim/real_backward.jl

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
1-
using FFTA, Test
1+
using FFTA, Test, LinearAlgebra
22
test_nums = [8, 11, 15, 16, 27, 100]
33
@testset "backward" begin
44
for N in test_nums
55
x = ones(Float64, N)
6-
y = bfft(x)
6+
y = brfft(x, 2*(N-1))
77
y_ref = 0*y
8-
y_ref[1] = N
8+
y_ref[1] = 2*(N-1)
9+
if !isapprox(y_ref, y, atol=1e-12)
10+
println(norm(y_ref - y))
11+
end
912
@test y_ref ≈ y atol=1e-12
1013
end
1114
end

‎test/onedim/real_forward.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ test_nums = [8, 11, 15, 16, 27, 100]
33
@testset verbose = true " forward" begin
44
for N in test_nums
55
x = ones(Float64, N)
6-
y = fft(x)
6+
y = rfft(x)
77
y_ref = 0*y
88
y_ref[1] = N
99
@test y ≈ y_ref atol=1e-12

0 commit comments

Comments
 (0)