lib/numo/pocketfft.rb in numo-pocketfft-0.2.2 vs lib/numo/pocketfft.rb in numo-pocketfft-0.3.0

@@ -62,12 +62,15 @@
     # @return [Numo::DComplex] Transformed data.
     def fftn(a)
       raise ArgumentError, 'Expect class of input array to be Numo::NArray.' unless a.is_a?(Numo::NArray)
       raise ArgumentError, 'Expect input array to be non-empty.' if a.empty?
 
-      b = a.dup
-      (0...b.ndim).to_a.reverse.each { |ax_id| b = raw_fft(b, ax_id, inverse: false, real: false) }
+      return raw_fft(a, 0, inverse: false, real: false) if a.ndim == 1
+
+      last_axis_id = a.ndim - 1
+      b = raw_fft(a, last_axis_id, inverse: false, real: false)
+      (last_axis_id - 1).downto(0) { |ax_id| b = raw_fft(b, ax_id, inverse: false, real: false) }
       b
     end
 
     # Compute the N-dimensional inverse discrete Fourier Transform.
     # @param a [Numo::DComplex] Complex input array with any-dimension.
@@ -75,12 +78,15 @@
     # @return [Numo::DComplex] Inversed transformed data.
     def ifftn(a)
       raise ArgumentError, 'Expect class of input array to be Numo::NArray.' unless a.is_a?(Numo::NArray)
       raise ArgumentError, 'Expect input array to be non-empty.' if a.empty?
 
-      b = a.dup
-      (0...b.ndim).to_a.each { |ax_id| b = raw_fft(b, ax_id, inverse: true, real: false) }
+      return raw_fft(a, 0, inverse: true, real: false) if a.ndim == 1
+
+      last_axis_id = a.ndim - 1
+      b = raw_fft(a, 0, inverse: true, real: false)
+      1.upto(last_axis_id) { |ax_id| b = raw_fft(b, ax_id, inverse: true, real: false) }
       b
     end
 
     # Compute the 1-dimensional discrete Fourier Transform for real input.
     # @param a [Numo::DFloat] Real 1-dimensional input array.
@@ -136,13 +142,15 @@
     # @return [Numo::DComplex] Transformed data.
     def rfftn(a)
       raise ArgumentError, 'Expect class of input array to be Numo::NArray.' unless a.is_a?(Numo::NArray)
       raise ArgumentError, 'Expect input array to be non-empty.' if a.empty?
 
+      return raw_fft(a, 0, inverse: false, real: true) if a.ndim == 1
+
       last_axis_id = a.ndim - 1
       b = raw_fft(a, last_axis_id, inverse: false, real: true)
-      (0...last_axis_id).to_a.reverse.each { |ax_id| b = raw_fft(b, ax_id, inverse: false, real: false) }
+      (last_axis_id - 1).downto(0) { |ax_id| b = raw_fft(b, ax_id, inverse: false, real: false) }
       b
     end
 
     # Compute the inverse of the N-dimensional discrete Fourier Transform of real input.
     # @param a [Numo::DComplex] Complex input array with any-dimension.
@@ -150,13 +158,15 @@
     # @return [Numo::DFloat] Inverse transformed data.
     def irfftn(a)
       raise ArgumentError, 'Expect class of input array to be Numo::NArray.' unless a.is_a?(Numo::NArray)
       raise ArgumentError, 'Expect input array to be non-empty.' if a.empty?
 
+      return raw_fft(a, 0, inverse: true, real: true) if a.ndim == 1
+
       last_axis_id = a.ndim - 1
-      b = a.dup
-      (0...last_axis_id).to_a.each { |ax_id| b = raw_fft(b, ax_id, inverse: true, real: false) }
+      b = raw_fft(a, 0, inverse: true, real: false)
+      1.upto(last_axis_id - 1) { |ax_id| b = raw_fft(b, ax_id, inverse: true, real: false) }
       raw_fft(b, last_axis_id, inverse: true, real: true)
     end
 
     # Convolve two N-dimensinal arrays using dscrete Fourier Transform.
     # @example
@@ -207,11 +217,10 @@
             # zero padding
             n = (a.shape[-1] - 1) * 2
             b_shape = a.shape
             b_shape[-1] = n
             b = Numo::DComplex.zeros(*b_shape)
-            b_range = [true] * b.ndim
-            b_range[-1] = 0...a.shape[-1]
+            b_range = Array.new(b.ndim) { |idx| idx < b.ndim - 1 ? true : 0...a.shape[-1] }
             b[*b_range] = a
             # inverse of dft for real data
             ext_irfft(b)
           else
             ext_rfft(a)