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 9cce31d92f980e5159d7f1e64717a5d9078caeb4..473960fa6f6e65b9e068978a646b1f43e8f991cc 100644 (file)
--- a/extra/math/transforms/fft/fft.factor
+++ b/extra/math/transforms/fft/fft.factor
@@ -1,46 +1,48 @@
 ! Copyright (c) 2012 John Benediktsson
 ! See http://factorcode.org/license.txt for BSD license.
-USING: arrays kernel locals math math.constants math.functions
-math.vectors sequences sequences.extras sequences.private ;
+USING: kernel math math.constants math.functions
+math.vectors sequences sequences.extras ;
 IN: math.transforms.fft
 
 <PRIVATE
 
+DEFER: (fft)
+
+! Discrete Fourier Transform
 :: (slow-fft) ( seq inverse? -- seq' )
     seq length :> N
-    inverse? 1 -1 ? 2pi * i* N / :> O
-    N iota [| k |
-        0 seq [ O k * * e^ * + ] each-index
+    inverse? 1 -1 ? 2pi * N / N <iota> n*v :> omega
+    N <iota> [| k |
+        0 seq omega [ k * cis * + ] 2each
         inverse? [ N / ] when
     ] map ; inline
 
-:: (fft) ( seq inverse? -- seq' )
+! Cooley–Tukey Algorithm
+:: (fast-fft) ( seq inverse? -- seq' )
     seq length :> N
     N 1 = [ seq ] [
-        inverse? 1 -1 ? 2pi * i* N / :> O
-        N 0 <array> :> X
-        N 2/ :> M
-        seq even-indices inverse? (fft)
-        seq odd-indices inverse? (fft)
-        [ [ [ O * e^ * + inverse? [ 2 / ] when ] [ X set-nth-unsafe ] bi ] 2each-index ]
-        [ [ [ O * e^ * - inverse? [ 2 / ] when ] [ M + X set-nth-unsafe ] bi ] 2each-index ]
-        2bi
-        X
+        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>
 
 ERROR: not-enough-data ;
 
 : fft ( seq -- seq' )
-    [ not-enough-data ] [
-        f over length even? [ (fft) ] [ (slow-fft) ] if
-    ] if-empty ;
+    [ not-enough-data ] [ f (fft) ] if-empty ;
 
 : ifft ( seq -- seq' )
-    [ not-enough-data ] [
-        t over length even? [ (fft) ] [ (slow-fft) ] if
-    ] if-empty ;
+    [ not-enough-data ] [ t (fft) ] if-empty ;
 
 : correlate ( x y -- z )
     [ fft ] [ reverse fft ] bi* v* ifft ;