1 # a "yield" version of the REMAP algorithm, for FFT Tukey-Cooley schedules
2 # original code for the FFT Tukey-Cooley schedul:
3 # https://www.nayuki.io/res/free-small-fft-in-multiple-languages/fft.py
5 # Radix-2 decimation-in-time FFT (real, not complex)
10 for i in range(0, n, size):
12 for j in range(i, i + halfsize):
15 temp1 = vec[jh] * exptable[k]
17 vec[jh] = temp2 - temp1
18 vec[jl] = temp2 + temp1
23 # python "yield" can be iterated. use this to make it clear how
24 # the indices are generated by using natural-looking nested loops
25 def iterate_butterfly_indices(SVSHAPE
):
26 # get indices to iterate over, in the required order
28 # createing lists of indices to iterate over in each dimension
29 # has to be done dynamically, because it depends on the size
30 # first, the size-based loop (which can be done statically)
36 # invert order if requested
37 if SVSHAPE
.invxyz
[0]: x_r
.reverse()
42 # start an infinite (wrapping) loop
45 for size
in x_r
: # loop over 3rd order dimension (size)
46 x_end
= size
== x_r
[-1]
47 # y_r schedule depends on size
51 for i
in range(0, n
, size
):
54 if SVSHAPE
.invxyz
[1]: y_r
.reverse()
55 for i
in y_r
: # loop over 2nd order dimension
60 for j
in range(i
, i
+halfsize
):
65 if SVSHAPE
.invxyz
[2]: k_r
.reverse()
66 if SVSHAPE
.invxyz
[2]: j_r
.reverse()
67 for j
, k
in zip(j_r
, k_r
): # loop over 1st order dimension
69 # now depending on MODE return the index
70 if SVSHAPE
.skip
== 0b00:
71 result
= j
# for vec[j]
72 elif SVSHAPE
.skip
== 0b01:
73 result
= j
+ halfsize
# for vec[j+halfsize]
74 elif SVSHAPE
.skip
== 0b10:
75 result
= k
# for exptable[k]
78 ((y_end
and z_end
)<<1) |
79 ((y_end
and x_end
and z_end
)<<2))
81 yield result
+ SVSHAPE
.offset
, loopends
84 # set the dimension sizes here
87 zdim
= 0 # again, not needed
89 # set total. err don't know how to calculate how many there are...
90 # do it manually for now
97 for i
in range(0, n
, size
):
98 for j
in range(i
, i
+ halfsize
):
107 SVSHAPE0
.lims
= [xdim
, ydim
, zdim
]
108 SVSHAPE0
.order
= [0,1,2] # experiment with different permutations, here
111 SVSHAPE0
.offset
= 0 # experiment with different offset, here
112 SVSHAPE0
.invxyz
= [0,0,0] # inversion if desired
113 # j+halfstep schedule
115 SVSHAPE1
.lims
= [xdim
, ydim
, zdim
]
116 SVSHAPE1
.order
= [0,1,2] # experiment with different permutations, here
119 SVSHAPE1
.offset
= 0 # experiment with different offset, here
120 SVSHAPE1
.invxyz
= [0,0,0] # inversion if desired
123 SVSHAPE2
.lims
= [xdim
, ydim
, zdim
]
124 SVSHAPE2
.order
= [0,1,2] # experiment with different permutations, here
127 SVSHAPE2
.offset
= 0 # experiment with different offset, here
128 SVSHAPE2
.invxyz
= [0,0,0] # inversion if desired
130 # enumerate over the iterator function, getting new indices
132 for idx
, (jl
, jh
, k
) in enumerate(zip(iterate_butterfly_indices(SVSHAPE0
),
133 iterate_butterfly_indices(SVSHAPE1
),
134 iterate_butterfly_indices(SVSHAPE2
))):
137 schedule
.append((jl
, jh
, k
))
139 # ok now pretty-print the results, with some debug output
144 tablestep
= n
// size
145 print ("size %d halfsize %d tablestep %d" % \
146 (size
, halfsize
, tablestep
))
147 for i
in range(0, n
, size
):
148 prefix
= "i %d\t" % i
150 for j
in range(i
, i
+ halfsize
):
151 (jl
, je
), (jh
, he
), (ks
, ke
) = schedule
[idx
]
152 print (" %-3d\t%s j=%-2d jh=%-2d k=%-2d -> "
153 "j[jl=%-2d] j[jh=%-2d] ex[k=%d]" % \
154 (idx
, prefix
, j
, j
+halfsize
, k
,
157 "end", bin(je
)[2:], bin(je
)[2:], bin(ke
)[2:])
163 if __name__
== '__main__':