factor: trim using lists

[factor.git] / extra / math / transforms / fft / fft.factor

diff --git a/extra/math/transforms/fft/fft.factor b/extra/math/transforms/fft/fft.factor
index 0688c00468ed268e16235c3a576944bc22fd6fbc..473960fa6f6e65b9e068978a646b1f43e8f991cc 100644 (file)
--- a/extra/math/transforms/fft/fft.factor
+++ b/extra/math/transforms/fft/fft.factor
@@ -1,38 +1,48 @@
-! Copyright (c) 2007 Hans Schmid.
+! Copyright (c) 2012 John Benediktsson
 ! See http://factorcode.org/license.txt for BSD license.
-USING: columns grouping kernel math math.constants math.functions math.vectors
-    sequences ;
+USING: kernel math math.constants math.functions
+math.vectors sequences sequences.extras ;
 IN: math.transforms.fft
 
-! Fast Fourier Transform
-
 <PRIVATE
 
-: n^v ( n v -- w ) [ ^ ] with map ;
-
-: omega ( n -- n' )
-    recip -2 pi i* * * exp ;
-
-: twiddle ( seq -- seq )
-    dup length [ omega ] [ n^v ] bi v* ;
+DEFER: (fft)
+
+! Discrete Fourier Transform
+:: (slow-fft) ( seq inverse? -- seq' )
+    seq length :> N
+    inverse? 1 -1 ? 2pi * N / N <iota> n*v :> omega
+    N <iota> [| k |
+        0 seq omega [ k * cis * + ] 2each
+        inverse? [ N / ] when
+    ] map ; inline
+
+! Cooley–Tukey Algorithm
+:: (fast-fft) ( seq inverse? -- seq' )
+    seq length :> N
+    N 1 = [ seq ] [
+        seq even-indices inverse? (fast-fft)
+        seq odd-indices inverse? (fast-fft)
+        inverse? 1 -1 ? 2pi * N /
+        [ * cis * ] curry map-index!
+        [ [ + inverse? [ 2 / ] when ] 2map ]
+        [ [ - inverse? [ 2 / ] when ] 2map ]
+        2bi append
+    ] if ; inline recursive
+
+: (fft) ( seq inverse? -- seq' )
+    over length power-of-2?
+    [ (fast-fft) ] [ (slow-fft) ] if ; inline
 
 PRIVATE>
 
-DEFER: fft
+ERROR: not-enough-data ;
 
-: two ( seq -- seq )
-    fft 2 v/n dup append ;
-
-<PRIVATE
-
-: even ( seq -- seq ) 2 group 0 <column> ;
-: odd ( seq -- seq ) 2 group 1 <column> ;
-
-: (fft) ( seq -- seq )
-    [ odd two twiddle ] [ even two ] bi v+ ;
-
-PRIVATE>
+: fft ( seq -- seq' )
+    [ not-enough-data ] [ f (fft) ] if-empty ;
 
-: fft ( seq -- seq )
-    dup length 1 = [ (fft) ] unless ;
+: ifft ( seq -- seq' )
+    [ not-enough-data ] [ t (fft) ] if-empty ;
 
+: correlate ( x y -- z )
+    [ fft ] [ reverse fft ] bi* v* ifft ;