FilesExpand file tree

 Write-for-current-Python

/

radix2_fft.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

180 lines (152 loc) · 6.02 KB

 Write-for-current-Python

/

radix2_fft.py

Copy path

Top

File metadata and controls

• Code
• Blame

180 lines (152 loc) · 6.02 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

"""

Fast Polynomial Multiplication using radix-2 fast Fourier Transform.

"""

import mpmath # for roots of unity

import numpy as np

class FFT:

"""

Fast Polynomial Multiplication using radix-2 fast Fourier Transform.

Reference:

https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm#The_radix-2_DIT_case

For polynomials of degree m and n the algorithms has complexity

O(n*logn + m*logm)

The main part of the algorithm is split in two parts:

1) __DFT: We compute the discrete fourier transform (DFT) of A and B using a

bottom-up dynamic approach -

2) __multiply: Once we obtain the DFT of A*B, we can similarly

invert it to obtain A*B

The class FFT takes two polynomials A and B with complex coefficients as arguments;

The two polynomials should be represented as a sequence of coefficients starting

from the free term. Thus, for instance x + 2*x^3 could be represented as

[0,1,0,2] or (0,1,0,2). The constructor adds some zeros at the end so that the

polynomials have the same length which is a power of 2 at least the length of

their product.

Example:

Create two polynomials as sequences

>>> A = [0, 1, 0, 2] # x+2x^3

>>> B = (2, 3, 4, 0) # 2+3x+4x^2

Create an FFT object with them

>>> x = FFT(A, B)

Print product

>>> print(x.product) # 2x + 3x^2 + 8x^3 + 4x^4 + 6x^5

[(-0+0j), (2+0j), (3+0j), (8+0j), (6+0j), (8+0j)]

__str__ test

>>> print(x)

A = 0*x^0 + 1*x^1 + 2*x^0 + 3*x^2

B = 0*x^2 + 1*x^3 + 2*x^4

A*B = 0*x^(-0+0j) + 1*x^(2+0j) + 2*x^(3+0j) + 3*x^(8+0j) + 4*x^(6+0j) + 5*x^(8+0j)

"""

def __init__(self, polyA=[0], polyB=[0]):

# Input as list

self.polyA = list(polyA)[:]

self.polyB = list(polyB)[:]

# Remove leading zero coefficients

while self.polyA[-1] == 0:

self.polyA.pop()

self.len_A = len(self.polyA)

while self.polyB[-1] == 0:

self.polyB.pop()

self.len_B = len(self.polyB)

# Add 0 to make lengths equal a power of 2

self.C_max_length = int(

2 ** np.ceil(np.log2(len(self.polyA) + len(self.polyB) - 1))

)

while len(self.polyA) < self.C_max_length:

self.polyA.append(0)

while len(self.polyB) < self.C_max_length:

self.polyB.append(0)

# A complex root used for the fourier transform

self.root = complex(mpmath.root(x=1, n=self.C_max_length, k=1))

# The product

self.product = self.__multiply()

# Discrete fourier transform of A and B

def __DFT(self, which):

if which == "A":

dft = [[x] for x in self.polyA]

else:

dft = [[x] for x in self.polyB]

# Corner case

if len(dft) <= 1:

return dft[0]

#

next_ncol = self.C_max_length // 2

while next_ncol > 0:

new_dft = [[] for i in range(next_ncol)]

root = self.root ** next_ncol

# First half of next step

current_root = 1

for j in range(self.C_max_length // (next_ncol * 2)):

for i in range(next_ncol):

new_dft[i].append(dft[i][j] + current_root * dft[i + next_ncol][j])

current_root *= root

# Second half of next step

current_root = 1

for j in range(self.C_max_length // (next_ncol * 2)):

for i in range(next_ncol):

new_dft[i].append(dft[i][j] - current_root * dft[i + next_ncol][j])

current_root *= root

# Update

dft = new_dft

next_ncol = next_ncol // 2

return dft[0]

# multiply the DFTs of A and B and find A*B

def __multiply(self):

dftA = self.__DFT("A")

dftB = self.__DFT("B")

inverseC = [[dftA[i] * dftB[i] for i in range(self.C_max_length)]]

del dftA

del dftB

# Corner Case

if len(inverseC[0]) <= 1:

return inverseC[0]

# Inverse DFT

next_ncol = 2

while next_ncol <= self.C_max_length:

new_inverseC = [[] for i in range(next_ncol)]

root = self.root ** (next_ncol // 2)

current_root = 1

# First half of next step

for j in range(self.C_max_length // next_ncol):

for i in range(next_ncol // 2):

# Even positions

new_inverseC[i].append(

(

inverseC[i][j]

+ inverseC[i][j + self.C_max_length // next_ncol]

)

/ 2

)

# Odd positions

new_inverseC[i + next_ncol // 2].append(

(

inverseC[i][j]

- inverseC[i][j + self.C_max_length // next_ncol]

)

/ (2 * current_root)

)

current_root *= root

# Update

inverseC = new_inverseC

next_ncol *= 2

# Unpack

inverseC = [round(x[0].real, 8) + round(x[0].imag, 8) * 1j for x in inverseC]

# Remove leading 0's

while inverseC[-1] == 0:

inverseC.pop()

return inverseC

# Overwrite __str__ for print(); Shows A, B and A*B

def __str__(self):

A = "A = " + " + ".join(

f"{coef}*x^{i}" for coef, i in enumerate(self.polyA[: self.len_A])

)

B = "B = " + " + ".join(

f"{coef}*x^{i}" for coef, i in enumerate(self.polyB[: self.len_B])

)

C = "A*B = " + " + ".join(

f"{coef}*x^{i}" for coef, i in enumerate(self.product)

)

return "\n".join((A, B, C))

# Unit tests

if __name__ == "__main__":

import doctest

doctest.testmod()