Merge pull request #174 from FourierFlows/glw/extend-plan-flows

Extends plan flows to take arbitrary args and kwargs but filter flags…

Author: navidcy

src/CuFourierFlows.jl

@@ -1,8 +1,8 @@
 using .CuArrays
 
 # Discard `effort` argument for CuArrays
-plan_flows_fft(a::CuArray, effort) = plan_fft(a)
-plan_flows_rfft(a::CuArray, effort) = plan_rfft(a)
+plan_flows_fft(a::CuArray, args...; flags=nothing, kwargs...) = plan_fft(a, args...; kwargs...)
+plan_flows_rfft(a::CuArray, args...; flags=nothing, kwargs...) = plan_rfft(a, args...; kwargs...)
 
 OneDGrid(dev::GPU, args...; kwargs...) = OneDGrid(args...; ArrayType=CuArray, kwargs...)
 TwoDGrid(dev::GPU, args...; kwargs...) = TwoDGrid(args...; ArrayType=CuArray, kwargs...)

src/domains.jl

@@ -1,5 +1,5 @@
-plan_flows_fft(a::Array, effort) = plan_fft(a; flags=effort)
-plan_flows_rfft(a::Array, effort) = plan_rfft(a; flags=effort)
+plan_flows_fft(a::Array, args...; kwargs...) = plan_fft(a, args...; kwargs...)
+plan_flows_rfft(a::Array, args...; kwargs...) = plan_rfft(a, args...; kwargs...)
 
 """
     ZeroDGrid()
@@ -75,8 +75,8 @@ function OneDGrid(nx, Lx; x0=-Lx/2, nthreads=Sys.CPU_THREADS, effort=FFTW.MEASUR
   invkrsq[1] = 0
 
   FFTW.set_num_threads(nthreads)
-   fftplan = plan_flows_fft(ArrayType{Complex{T}, 1}(undef, nx), effort)
-  rfftplan = plan_flows_rfft(ArrayType{T, 1}(undef, nx), effort)
+   fftplan = plan_flows_fft(ArrayType{Complex{T}, 1}(undef, nx), flags=effort)
+  rfftplan = plan_flows_rfft(ArrayType{T, 1}(undef, nx), flags=effort)
 
   kalias, kralias = getaliasedwavenumbers(nk, nkr, dealias)
 
@@ -155,8 +155,8 @@ function TwoDGrid(nx, Lx, ny=nx, Ly=Lx; x0=-Lx/2, y0=-Ly/2, nthreads=Sys.CPU_THR
 
   # FFT plans
   FFTW.set_num_threads(nthreads)
-  fftplan = plan_flows_fft(ArrayType{Complex{T}, 2}(undef, nx, ny), effort)
-  rfftplan = plan_flows_rfft(ArrayType{T, 2}(undef, nx, ny), effort)
+  fftplan = plan_flows_fft(ArrayType{Complex{T}, 2}(undef, nx, ny), flags=effort)
+  rfftplan = plan_flows_rfft(ArrayType{T, 2}(undef, nx, ny), flags=effort)
 
   # Index endpoints for aliasfrac i, j wavenumbers
   kalias, kralias = getaliasedwavenumbers(nk, nkr, dealias)
@@ -248,8 +248,8 @@ function ThreeDGrid(nx, Lx, ny=nx, Ly=Lx, nz=nx, Lz=Lx; x0=-Lx/2, y0=-Ly/2, z0=-
 
   # FFT plans
   FFTW.set_num_threads(nthreads)
-  fftplan = plan_flows_fft(ArrayType{Complex{T}, 3}(undef, nx, ny, nz), effort)
-  rfftplan = plan_flows_rfft(ArrayType{T, 3}(undef, nx, ny, nz), effort)
+  fftplan = plan_flows_fft(ArrayType{Complex{T}, 3}(undef, nx, ny, nz), flags=effort)
+  rfftplan = plan_flows_rfft(ArrayType{T, 3}(undef, nx, ny, nz), flags=effort)
 
   # Index endpoints for aliasfrac i, j wavenumbers
   kalias, kralias = getaliasedwavenumbers(nk, nkr, dealias)

test/runtests.jl

@@ -95,9 +95,11 @@ for dev in devices
     @test testgridpoints(dev, g₃)
     @test testdealias(g₃)
     @test testmakefilter(dev, g₃)
+    @test test_plan_flows_fftrfft(dev)
 
     # Test typed grids
     T = Float32
+    @test test_plan_flows_fftrfft(dev, T=T)
     @test testtypedonedgrid(dev, nx, Lx; T=T)
     @test testtypedtwodgrid(dev, nx, Lx, ny, Ly; T=T)
     @test testtypedthreedgrid(dev, nx, Lx, ny, Ly, nz, Lz; T=T)

test/test_grid.jl

@@ -123,10 +123,35 @@ function testtypedthreedgrid(dev::Device, nx, Lx, ny=nx, Ly=Lx, nz=nx, Lz=Lx; T=
   typeof(gr.dx)==T && typeof(gr.dy)==T && typeof(gr.dz)==T && typeof(gr.x[1])==T && typeof(gr.y[1])==T && typeof(gr.z[1])==T && typeof(gr.Lx)==T && typeof(gr.Ly)==T && typeof(gr.Lz)==T && eltype(gr) == T
 end
 
-function testmakefilter(dev::Device, g::AbstractGrid{T}) where T
+function testmakefilter(dev::Device, g::AbstractGrid)
+  T = eltype(g)
   filter = FourierFlows.makefilter(g)
   G = typeof(g)
   nofilter = G<:OneDGrid ? filter[@. g.kr*g.dx/π < 0.65] : ( G<:TwoDGrid ? filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2) < 0.65] : filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2 + (g.m*g.dz/π)^2) < 0.65] )
   fullfilter = G<:OneDGrid ? filter[@. g.kr*g.dx/π > 0.999] : ( G<:TwoDGrid ? filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2) > 0.999] : filter[@. sqrt((g.kr*g.dx/π)^2 + (g.l*g.dy/π)^2 + (g.m*g.dz/π)^2) > 0.999] )
   nofilter== zeros(dev, T, size(nofilter)) .+ 1 && isapprox(fullfilter, zeros(dev, T, size(fullfilter)), atol=1e-12)
 end
+
+function test_plan_flows_fftrfft(::CPU; T=Float64)
+  A = ArrayType(CPU())
+  return ( typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4,))))) == FFTW.cFFTWPlan{Complex{T},-1,false,1} &&
+  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6))))) == FFTW.cFFTWPlan{Complex{T},-1,false,2} &&
+  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))))) == FFTW.cFFTWPlan{Complex{T},-1,false,3} &&
+  FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))), [1, 2]).region == [1, 2] &&
+  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4,))))) == FFTW.rFFTWPlan{T,-1,false,1} &&
+  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6))))) == FFTW.rFFTWPlan{T,-1,false,2} &&
+  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))))) == FFTW.rFFTWPlan{T,-1,false,3} &&
+  FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))), [1, 2]).region == [1, 2] )
+end
+
+function test_plan_flows_fftrfft(::GPU; T=Float64)
+  A = ArrayType(GPU())
+  return ( typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4,))))) == CuArrays.CUFFT.cCuFFTPlan{Complex{T},-1,false,1} &&
+  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6))))) == CuArrays.CUFFT.cCuFFTPlan{Complex{T},-1,false,2} &&
+  typeof(FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))))) == CuArrays.CUFFT.cCuFFTPlan{Complex{T},-1,false,3} &&
+  FourierFlows.plan_flows_fft(A(rand(Complex{T}, (4, 6, 8))), [1, 2]).region == [1, 2] &&
+  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4,))))) == CuArrays.CUFFT.rCuFFTPlan{T,-1,false,1} &&
+  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6))))) == CuArrays.CUFFT.rCuFFTPlan{T,-1,false,2} &&
+  typeof(FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))))) == CuArrays.CUFFT.rCuFFTPlan{T,-1,false,3} &&
+  FourierFlows.plan_flows_rfft(A(rand(T, (4, 6, 8))), [1, 2]).region == [1, 2] )
+end