Coverage for /usr/lib/python3/dist-packages/sympy/discrete/transforms.py: 16%

1"""

2Discrete Fourier Transform, Number Theoretic Transform,

3Walsh Hadamard Transform, Mobius Transform

4"""

6from sympy.core import S, Symbol, sympify

7from sympy.core.function import expand_mul

8from sympy.core.numbers import pi, I

9from sympy.functions.elementary.trigonometric import sin, cos

10from sympy.ntheory import isprime, primitive_root

11from sympy.utilities.iterables import ibin, iterable

12from sympy.utilities.misc import as_int

15#----------------------------------------------------------------------------#

16# #

17# Discrete Fourier Transform #

18# #

19#----------------------------------------------------------------------------#

21def _fourier_transform(seq, dps, inverse=False):

22 """Utility function for the Discrete Fourier Transform"""

24 if not iterable(seq):

25 raise TypeError("Expected a sequence of numeric coefficients "

26 "for Fourier Transform")

28 a = [sympify(arg) for arg in seq]

29 if any(x.has(Symbol) for x in a):

30 raise ValueError("Expected non-symbolic coefficients")

32 n = len(a)

33 if n < 2:

34 return a

36 b = n.bit_length() - 1

37 if n&(n - 1): # not a power of 2

38 b += 1

39 n = 2**b

41 a += [S.Zero]*(n - len(a))

42 for i in range(1, n):

43 j = int(ibin(i, b, str=True)[::-1], 2)

44 if i < j:

45 a[i], a[j] = a[j], a[i]

47 ang = -2*pi/n if inverse else 2*pi/n

49 if dps is not None:

50 ang = ang.evalf(dps + 2)

52 w = [cos(ang*i) + I*sin(ang*i) for i in range(n // 2)]

54 h = 2

55 while h <= n:

56 hf, ut = h // 2, n // h

57 for i in range(0, n, h):

58 for j in range(hf):

59 u, v = a[i + j], expand_mul(a[i + j + hf]*w[ut * j])

60 a[i + j], a[i + j + hf] = u + v, u - v

61 h *= 2

63 if inverse:

64 a = [(x/n).evalf(dps) for x in a] if dps is not None \

65 else [x/n for x in a]

67 return a

70def fft(seq, dps=None):

71 r"""

72 Performs the Discrete Fourier Transform (**DFT**) in the complex domain.

74 The sequence is automatically padded to the right with zeros, as the

75 *radix-2 FFT* requires the number of sample points to be a power of 2.

77 This method should be used with default arguments only for short sequences

78 as the complexity of expressions increases with the size of the sequence.

80 Parameters

81 ==========

83 seq : iterable

84 The sequence on which **DFT** is to be applied.

85 dps : Integer

86 Specifies the number of decimal digits for precision.

88 Examples

89 ========

91 >>> from sympy import fft, ifft

93 >>> fft([1, 2, 3, 4])

94 [10, -2 - 2*I, -2, -2 + 2*I]

95 >>> ifft(_)

96 [1, 2, 3, 4]

98 >>> ifft([1, 2, 3, 4])

99 [5/2, -1/2 + I/2, -1/2, -1/2 - I/2]

100 >>> fft(_)

101 [1, 2, 3, 4]

102

103 >>> ifft([1, 7, 3, 4], dps=15)

104 [3.75, -0.5 - 0.75*I, -1.75, -0.5 + 0.75*I]

105 >>> fft(_)

106 [1.0, 7.0, 3.0, 4.0]

107

108 References

109 ==========

110

111 .. [1] https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm

112 .. [2] https://mathworld.wolfram.com/FastFourierTransform.html

113

114 """

115

116 return _fourier_transform(seq, dps=dps)

117

118

119def ifft(seq, dps=None):

120 return _fourier_transform(seq, dps=dps, inverse=True)

121

122ifft.__doc__ = fft.__doc__

123

124

125#----------------------------------------------------------------------------#

126# #

127# Number Theoretic Transform #

128# #

129#----------------------------------------------------------------------------#

130

131def _number_theoretic_transform(seq, prime, inverse=False):

132 """Utility function for the Number Theoretic Transform"""

133

134 if not iterable(seq):

135 raise TypeError("Expected a sequence of integer coefficients "

136 "for Number Theoretic Transform")

137

138 p = as_int(prime)

139 if not isprime(p):

140 raise ValueError("Expected prime modulus for "

141 "Number Theoretic Transform")

142

143 a = [as_int(x) % p for x in seq]

144

145 n = len(a)

146 if n < 1:

147 return a

148

149 b = n.bit_length() - 1

150 if n&(n - 1):

151 b += 1

152 n = 2**b

153

154 if (p - 1) % n:

155 raise ValueError("Expected prime modulus of the form (m*2**k + 1)")

156

157 a += [0]*(n - len(a))

158 for i in range(1, n):

159 j = int(ibin(i, b, str=True)[::-1], 2)

160 if i < j:

161 a[i], a[j] = a[j], a[i]

162

163 pr = primitive_root(p)

164

165 rt = pow(pr, (p - 1) // n, p)

166 if inverse:

167 rt = pow(rt, p - 2, p)

168

169 w = [1]*(n // 2)

170 for i in range(1, n // 2):

171 w[i] = w[i - 1]*rt % p

172

173 h = 2

174 while h <= n:

175 hf, ut = h // 2, n // h

176 for i in range(0, n, h):

177 for j in range(hf):

178 u, v = a[i + j], a[i + j + hf]*w[ut * j]

179 a[i + j], a[i + j + hf] = (u + v) % p, (u - v) % p

180 h *= 2

181

182 if inverse:

183 rv = pow(n, p - 2, p)

184 a = [x*rv % p for x in a]

185

186 return a

187

188

189def ntt(seq, prime):

190 r"""

191 Performs the Number Theoretic Transform (**NTT**), which specializes the

192 Discrete Fourier Transform (**DFT**) over quotient ring `Z/pZ` for prime

193 `p` instead of complex numbers `C`.

194

195 The sequence is automatically padded to the right with zeros, as the

196 *radix-2 NTT* requires the number of sample points to be a power of 2.

197

198 Parameters

199 ==========

200

201 seq : iterable

202 The sequence on which **DFT** is to be applied.

203 prime : Integer

204 Prime modulus of the form `(m 2^k + 1)` to be used for performing

205 **NTT** on the sequence.

206

207 Examples

208 ========

209

210 >>> from sympy import ntt, intt

211 >>> ntt([1, 2, 3, 4], prime=3*2**8 + 1)

212 [10, 643, 767, 122]

213 >>> intt(_, 3*2**8 + 1)

214 [1, 2, 3, 4]

215 >>> intt([1, 2, 3, 4], prime=3*2**8 + 1)

216 [387, 415, 384, 353]

217 >>> ntt(_, prime=3*2**8 + 1)

218 [1, 2, 3, 4]

219

220 References

221 ==========

222

223 .. [1] http://www.apfloat.org/ntt.html

224 .. [2] https://mathworld.wolfram.com/NumberTheoreticTransform.html

225 .. [3] https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general%29

226

227 """

228

229 return _number_theoretic_transform(seq, prime=prime)

230

231

232def intt(seq, prime):

233 return _number_theoretic_transform(seq, prime=prime, inverse=True)

234

235intt.__doc__ = ntt.__doc__

236

237

238#----------------------------------------------------------------------------#

239# #

240# Walsh Hadamard Transform #

241# #

242#----------------------------------------------------------------------------#

243

244def _walsh_hadamard_transform(seq, inverse=False):

245 """Utility function for the Walsh Hadamard Transform"""

246

247 if not iterable(seq):

248 raise TypeError("Expected a sequence of coefficients "

249 "for Walsh Hadamard Transform")

250

251 a = [sympify(arg) for arg in seq]

252 n = len(a)

253 if n < 2:

254 return a

255

256 if n&(n - 1):

257 n = 2**n.bit_length()

258

259 a += [S.Zero]*(n - len(a))

260 h = 2

261 while h <= n:

262 hf = h // 2

263 for i in range(0, n, h):

264 for j in range(hf):

265 u, v = a[i + j], a[i + j + hf]

266 a[i + j], a[i + j + hf] = u + v, u - v

267 h *= 2

268

269 if inverse:

270 a = [x/n for x in a]

271

272 return a

273

274

275def fwht(seq):

276 r"""

277 Performs the Walsh Hadamard Transform (**WHT**), and uses Hadamard

278 ordering for the sequence.

279

280 The sequence is automatically padded to the right with zeros, as the

281 *radix-2 FWHT* requires the number of sample points to be a power of 2.

282

283 Parameters

284 ==========

285

286 seq : iterable

287 The sequence on which WHT is to be applied.

288

289 Examples

290 ========

291

292 >>> from sympy import fwht, ifwht

293 >>> fwht([4, 2, 2, 0, 0, 2, -2, 0])

294 [8, 0, 8, 0, 8, 8, 0, 0]

295 >>> ifwht(_)

296 [4, 2, 2, 0, 0, 2, -2, 0]

297

298 >>> ifwht([19, -1, 11, -9, -7, 13, -15, 5])

299 [2, 0, 4, 0, 3, 10, 0, 0]

300 >>> fwht(_)

301 [19, -1, 11, -9, -7, 13, -15, 5]

302

303 References

304 ==========

305

306 .. [1] https://en.wikipedia.org/wiki/Hadamard_transform

307 .. [2] https://en.wikipedia.org/wiki/Fast_Walsh%E2%80%93Hadamard_transform

308

309 """

310

311 return _walsh_hadamard_transform(seq)

312

313

314def ifwht(seq):

315 return _walsh_hadamard_transform(seq, inverse=True)

316

317ifwht.__doc__ = fwht.__doc__

318

319

320#----------------------------------------------------------------------------#

321# #

322# Mobius Transform for Subset Lattice #

323# #

324#----------------------------------------------------------------------------#

325

326def _mobius_transform(seq, sgn, subset):

327 r"""Utility function for performing Mobius Transform using

328 Yate's Dynamic Programming method"""

329

330 if not iterable(seq):

331 raise TypeError("Expected a sequence of coefficients")

332

333 a = [sympify(arg) for arg in seq]

334

335 n = len(a)

336 if n < 2:

337 return a

338

339 if n&(n - 1):

340 n = 2**n.bit_length()

341

342 a += [S.Zero]*(n - len(a))

343

344 if subset:

345 i = 1

346 while i < n:

347 for j in range(n):

348 if j & i:

349 a[j] += sgn*a[j ^ i]

350 i *= 2

351

352 else:

353 i = 1

354 while i < n:

355 for j in range(n):

356 if j & i:

357 continue

358 a[j] += sgn*a[j ^ i]

359 i *= 2

360

361 return a

362

363

364def mobius_transform(seq, subset=True):

365 r"""

366 Performs the Mobius Transform for subset lattice with indices of

367 sequence as bitmasks.

368

369 The indices of each argument, considered as bit strings, correspond

370 to subsets of a finite set.

371

372 The sequence is automatically padded to the right with zeros, as the

373 definition of subset/superset based on bitmasks (indices) requires

374 the size of sequence to be a power of 2.

375

376 Parameters

377 ==========

378

379 seq : iterable

380 The sequence on which Mobius Transform is to be applied.

381 subset : bool

382 Specifies if Mobius Transform is applied by enumerating subsets

383 or supersets of the given set.

384

385 Examples

386 ========

387

388 >>> from sympy import symbols

389 >>> from sympy import mobius_transform, inverse_mobius_transform

390 >>> x, y, z = symbols('x y z')

391

392 >>> mobius_transform([x, y, z])

393 [x, x + y, x + z, x + y + z]

394 >>> inverse_mobius_transform(_)

395 [x, y, z, 0]

396

397 >>> mobius_transform([x, y, z], subset=False)

398 [x + y + z, y, z, 0]

399 >>> inverse_mobius_transform(_, subset=False)

400 [x, y, z, 0]

401

402 >>> mobius_transform([1, 2, 3, 4])

403 [1, 3, 4, 10]

404 >>> inverse_mobius_transform(_)

405 [1, 2, 3, 4]

406 >>> mobius_transform([1, 2, 3, 4], subset=False)

407 [10, 6, 7, 4]

408 >>> inverse_mobius_transform(_, subset=False)

409 [1, 2, 3, 4]

410

411 References

412 ==========

413

414 .. [1] https://en.wikipedia.org/wiki/M%C3%B6bius_inversion_formula

415 .. [2] https://people.csail.mit.edu/rrw/presentations/subset-conv.pdf

416 .. [3] https://arxiv.org/pdf/1211.0189.pdf

417

418 """

419

420 return _mobius_transform(seq, sgn=+1, subset=subset)

421

422def inverse_mobius_transform(seq, subset=True):

423 return _mobius_transform(seq, sgn=-1, subset=subset)

424

425inverse_mobius_transform.__doc__ = mobius_transform.__doc__