Skip to content

Instantly share code, notes, and snippets.

@yuyichao
Created July 9, 2015 22:38
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save yuyichao/6c737467f90020325bfa to your computer and use it in GitHub Desktop.
Save yuyichao/6c737467f90020325bfa to your computer and use it in GitHub Desktop.
diff --git a/base/fft/FFTW.jl b/base/fft/FFTW.jl
index a8e7108..e20edc0 100644
--- a/base/fft/FFTW.jl
+++ b/base/fft/FFTW.jl
@@ -73,13 +73,12 @@ typealias fftwTypeSingle Union(Type{Float32},Type{Complex64})
# since it is not written to. Hence, it is convenient to create an
# array-like type that carries a size and a stride like a "real" array
# but which is converted to C_NULL as a pointer.
-immutable FakeArray{T, N}
+immutable FakeArray{T, N} <: AbstractArray{T, N}
sz::NTuple{N, Int}
st::NTuple{N, Int}
end
size(a::FakeArray) = a.sz
strides(a::FakeArray) = a.st
-ndims(a::FakeArray) = length(a.sz)
unsafe_convert{T}(::Type{Ptr{T}}, a::FakeArray{T}) = convert(Ptr{T}, C_NULL)
pointer{T}(a::FakeArray{T}) = convert(Ptr{T}, C_NULL)
FakeArray{N}(T, sz::NTuple{N, Int}) = FakeArray{T, N}(sz, colmajorstrides(sz))
@@ -220,7 +219,7 @@ for P in (:cFFTWPlan, :rFFTWPlan, :r2rFFTWPlan) # complex, r2c/c2r, and r2r
return p
end
end
- $P{T<:fftwNumber}(plan::PlanPtr, flags::Integer, R::Any, X::StridedArrayish{T}, Y::StridedArrayish, K) = $P{T,K,pointer(X) == pointer(Y)}(plan, size(X), size(Y), strides(X), strides(Y), alignment_of(X), alignment_of(Y), flags, R)
+ $P{T<:fftwNumber}(plan::PlanPtr, flags::Integer, R::Any, X::StridedArrayish{T}, Y::StridedArrayish, K) = $P{T,K,pointer(X) == pointer(Y),ndims(X)}(plan, size(X), size(Y), strides(X), strides(Y), alignment_of(X), alignment_of(Y), flags, R)
end
end
@@ -574,28 +573,28 @@ for (f,direction) in ((:fft,FORWARD), (:bfft,BACKWARD))
plan_f! = symbol("plan_",f,"!")
idirection = -direction
@eval begin
- function $plan_f{T<:fftwComplex}(X::StridedArray{T}, region;
- flags::Integer=ESTIMATE,
- timelimit::Real=NO_TIMELIMIT)
+ function $plan_f{T<:fftwComplex, N}(X::StridedArray{T, N}, region;
+ flags::Integer=ESTIMATE,
+ timelimit::Real=NO_TIMELIMIT)
cFFTWPlan(X, fakesimilar(flags, X, T), region,
- $direction, flags,timelimit)::cFFTWPlan{T,$direction,false}
+ $direction, flags,timelimit)::cFFTWPlan{T,$direction,false,N}
end
- function $plan_f!{T<:fftwComplex}(X::StridedArray{T}, region;
- flags::Integer=ESTIMATE,
- timelimit::Real=NO_TIMELIMIT)
- cFFTWPlan(X, X, region, $direction, flags, timelimit)::cFFTWPlan{T,$direction,true}
+ function $plan_f!{T<:fftwComplex, N}(X::StridedArray{T, N}, region;
+ flags::Integer=ESTIMATE,
+ timelimit::Real=NO_TIMELIMIT)
+ cFFTWPlan(X, X, region, $direction, flags, timelimit)::cFFTWPlan{T,$direction,true,N}
end
$plan_f{T<:fftwComplex}(X::StridedArray{T}; kws...) =
$plan_f(X, 1:ndims(X); kws...)
$plan_f!{T<:fftwComplex}(X::StridedArray{T}; kws...) =
$plan_f!(X, 1:ndims(X); kws...)
- function plan_inv{T<:fftwComplex,inplace}(p::cFFTWPlan{T,$direction,inplace})
+ function plan_inv{T<:fftwComplex,inplace,N}(p::cFFTWPlan{T,$direction,inplace,N})
X = Array(T, p.sz)
Y = inplace ? X : fakesimilar(p.flags, X, T)
ScaledPlan(cFFTWPlan(X, Y, p.region, $idirection,
- p.flags, NO_TIMELIMIT)::cFFTWPlan{T,$idirection,inplace},
+ p.flags, NO_TIMELIMIT)::cFFTWPlan{T,$idirection,inplace,N},
normalization(X, p.region))
end
end
@@ -625,15 +624,16 @@ end
for (Tr,Tc) in ((:Float32,:Complex64),(:Float64,:Complex128))
# Note: use $FORWARD and $BACKWARD below because of issue #9775
@eval begin
- function plan_rfft(X::StridedArray{$Tr}, region;
- flags::Integer=ESTIMATE, timelimit::Real=NO_TIMELIMIT)
+ function plan_rfft{N}(X::StridedArray{$Tr,N}, region;
+ flags::Integer=ESTIMATE, timelimit::Real=NO_TIMELIMIT)
osize = rfft_output_size(X, region)
Y = flags&ESTIMATE != 0 ? FakeArray($Tc,osize...) : Array($Tc,osize...)
- rFFTWPlan(X, Y, region, flags, timelimit)::rFFTWPlan{$Tr,$FORWARD,false}
+ rFFTWPlan(X, Y, region, flags, timelimit)::rFFTWPlan{$Tr,$FORWARD,false,N}
end
- function plan_brfft(X::StridedArray{$Tc}, d::Integer, region;
- flags::Integer=ESTIMATE, timelimit::Real=NO_TIMELIMIT)
+ function plan_brfft{N}(X::StridedArray{$Tc,N}, d::Integer, region;
+ flags::Integer=ESTIMATE,
+ timelimit::Real=NO_TIMELIMIT)
osize = brfft_output_size(X, d, region)
Y = flags&ESTIMATE != 0 ? FakeArray($Tr,osize...) : Array($Tr,osize...)
@@ -645,25 +645,25 @@ for (Tr,Tc) in ((:Float32,:Complex64),(:Float64,:Complex128))
else
Xc = copy(X)
rFFTWPlan(X, Y, region, flags, timelimit)
- end::rFFTWPlan{$Tc,$BACKWARD,false}
+ end::rFFTWPlan{$Tc,$BACKWARD,false,N}
end
plan_rfft(X::StridedArray{$Tr};kws...)=plan_rfft(X,1:ndims(X);kws...)
plan_brfft(X::StridedArray{$Tr};kws...)=plan_brfft(X,1:ndims(X);kws...)
- function plan_inv(p::rFFTWPlan{$Tr,$FORWARD,false})
+ function plan_inv{N}(p::rFFTWPlan{$Tr,$FORWARD,false,N})
X = Array($Tr, p.sz)
Y = p.flags&ESTIMATE != 0 ? FakeArray($Tc,p.osz) : Array($Tc,p.osz)
ScaledPlan(rFFTWPlan(Y, X, p.region,
length(p.region)<=1 ? p.flags | PRESERVE_INPUT
- : p.flags, NO_TIMELIMIT)::rFFTWPlan{$Tc,$BACKWARD,false},
+ : p.flags, NO_TIMELIMIT)::rFFTWPlan{$Tc,$BACKWARD,false,N},
normalization(X, p.region))
end
- function plan_inv(p::rFFTWPlan{$Tc,$BACKWARD,false})
+ function plan_inv{N}(p::rFFTWPlan{$Tc,$BACKWARD,false,N})
X = Array($Tc, p.sz)
Y = p.flags&ESTIMATE != 0 ? FakeArray($Tr,p.osz) : Array($Tr,p.osz)
- ScaledPlan(rFFTWPlan(Y, X, p.region, p.flags, NO_TIMELIMIT)::rFFTWPlan{$Tr,$FORWARD,false},
+ ScaledPlan(rFFTWPlan(Y, X, p.region, p.flags, NO_TIMELIMIT)::rFFTWPlan{$Tr,$FORWARD,false,N},
normalization(Y, p.region))
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment