added extra methods to allow for in-place ifft without modification

The-ReSolver · Dec 17, 2021 · 0741e73 · 0741e73
1 parent 2dbfdf3
commit 0741e73
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/src/fft.jl b/src/fft.jl
@@ -87,6 +87,25 @@ function (f::IFFTPlan!{Ny, Nz, Nt})(u::VectorField{N, P}, û::VectorField{N, S}
     for i in 1:N
         f(u[i], û[i])
     end
+end
+
+function (f::IFFTPlan!{Ny, Nz, Nt})(u::PhysicalField{Ny, Nz, Nt},
+                                    û::SpectralField{Ny, Nz, Nt},
+                                    û_tmp::SpectralField{Ny, Nz, Nt}) where {Ny, Nz, Nt}
+    # copy spectral contents to temporary field
+    û_tmp .= û
+
+    # perform transform
+    FFTW.unsafe_execute!(f.plan, parent(û_tmp), parent(u))
+
+    return u
+end
+
+function (f::IFFTPlan!{Ny, Nz, Nt})(u::VectorField{N, P}, û::VectorField{N, S}, û_tmp::VectorField{N, S}) where
+            {Ny, Nz, Nt, N, P<:PhysicalField{Ny, Nz, Nt}, S<:SpectralField{Ny, Nz, Nt}}
+    for i in 1:N
+        f(u[i], û[i], û_tmp[i])
+    end
 
     return u
 end
diff --git a/test/test_fft.jl b/test/test_fft.jl
@@ -1,8 +1,9 @@
 @testset "FFT Transforms                " begin
     # randon signal
-    Ny = rand(3:50)
-    Nz = rand(3:50)
-    Nt = rand(3:50)
+    # Ny = rand(3:50)
+    # Nz = rand(3:50)
+    # Nt = rand(3:50)
+    Ny = 8; Nz = 8; Nt = 8
     ω = abs(randn())
     β = abs(randn())
     grid = Grid(rand(Float64, Ny), Nz, Nt,
@@ -26,9 +27,9 @@
 
     # is the transform invertible correctly
     FFT(Â1, A1)
-    IFFT(B1, Â1)
+    IFFT(B1, Â1, copy(Â1))
     FFT(𝐀̂, 𝐀)
-    IFFT(𝐁, 𝐀̂)
+    IFFT(𝐁, 𝐀̂, VectorField(grid; N = 2))
     @test A1 ≈ B1
     @test 𝐀 ≈ 𝐁
 end