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
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_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 # y_r schedule depends on size
50 for i
in range(0, n
, size
):
53 if SVSHAPE
.invxyz
[1]: y_r
.reverse()
54 for i
in y_r
: # loop over 2nd order dimension
58 for j
in range(i
, i
+halfsize
):
63 if SVSHAPE
.invxyz
[2]: k_r
.reverse()
64 if SVSHAPE
.invxyz
[2]: j_r
.reverse()
65 for j
, k
in zip(j_r
, k_r
): # loop over 1st order dimension
66 # skip the first entries up to offset
67 if skip
< SVSHAPE
.offset
:
70 # now depending on MODE return the index
71 if SVSHAPE
.skip
== 0b00:
72 result
= j
# for vec[j]
73 elif SVSHAPE
.skip
== 0b01:
74 result
= j
+ halfsize
# for vec[j+halfsize]
75 elif SVSHAPE
.skip
== 0b10:
76 result
= k
# for exptable[k]
81 # set the dimension sizes here
84 zdim
= 0 # again, not needed
86 # set total. err don't know how to calculate how many there are...
87 # do it manually for now
94 for i
in range(0, n
, size
):
95 for j
in range(i
, i
+ halfsize
):
104 SVSHAPE0
.lims
= [xdim
, ydim
, zdim
]
105 SVSHAPE0
.order
= [0,1,2] # experiment with different permutations, here
108 SVSHAPE0
.offset
= 0 # experiment with different offset, here
109 SVSHAPE0
.invxyz
= [0,0,0] # inversion if desired
110 # j+halfstep schedule
112 SVSHAPE1
.lims
= [xdim
, ydim
, zdim
]
113 SVSHAPE1
.order
= [0,1,2] # experiment with different permutations, here
116 SVSHAPE1
.offset
= 0 # experiment with different offset, here
117 SVSHAPE1
.invxyz
= [0,0,0] # inversion if desired
120 SVSHAPE2
.lims
= [xdim
, ydim
, zdim
]
121 SVSHAPE2
.order
= [0,1,2] # experiment with different permutations, here
124 SVSHAPE2
.offset
= 0 # experiment with different offset, here
125 SVSHAPE2
.invxyz
= [0,0,0] # inversion if desired
127 # enumerate over the iterator function, getting new indices
129 for idx
, (jl
, jh
, k
) in enumerate(zip(iterate_indices(SVSHAPE0
),
130 iterate_indices(SVSHAPE1
),
131 iterate_indices(SVSHAPE2
))):
134 schedule
.append((jl
, jh
, k
))
136 # ok now pretty-print the results, with some debug output
141 tablestep
= n
// size
142 print ("size %d halfsize %d tablestep %d" % \
143 (size
, halfsize
, tablestep
))
144 for i
in range(0, n
, size
):
145 prefix
= "i %d\t" % i
147 for j
in range(i
, i
+ halfsize
):
148 jl
, jh
, ks
= schedule
[idx
]
149 print (" %-3d\t%s j=%-2d jh=%-2d k=%-2d -> "
150 "j[jl=%-2d] j[jh=%-2d] exptable[k=%d]" % \
151 (idx
, prefix
, j
, j
+halfsize
, k
,
158 if __name__
== '__main__':