426 lines
11 KiB
Python
426 lines
11 KiB
Python
"""
|
|
Discrete Fourier Transform, Number Theoretic Transform,
|
|
Walsh Hadamard Transform, Mobius Transform
|
|
"""
|
|
|
|
from sympy.core import S, Symbol, sympify
|
|
from sympy.core.function import expand_mul
|
|
from sympy.core.numbers import pi, I
|
|
from sympy.functions.elementary.trigonometric import sin, cos
|
|
from sympy.ntheory import isprime, primitive_root
|
|
from sympy.utilities.iterables import ibin, iterable
|
|
from sympy.utilities.misc import as_int
|
|
|
|
|
|
#----------------------------------------------------------------------------#
|
|
# #
|
|
# Discrete Fourier Transform #
|
|
# #
|
|
#----------------------------------------------------------------------------#
|
|
|
|
def _fourier_transform(seq, dps, inverse=False):
|
|
"""Utility function for the Discrete Fourier Transform"""
|
|
|
|
if not iterable(seq):
|
|
raise TypeError("Expected a sequence of numeric coefficients "
|
|
"for Fourier Transform")
|
|
|
|
a = [sympify(arg) for arg in seq]
|
|
if any(x.has(Symbol) for x in a):
|
|
raise ValueError("Expected non-symbolic coefficients")
|
|
|
|
n = len(a)
|
|
if n < 2:
|
|
return a
|
|
|
|
b = n.bit_length() - 1
|
|
if n&(n - 1): # not a power of 2
|
|
b += 1
|
|
n = 2**b
|
|
|
|
a += [S.Zero]*(n - len(a))
|
|
for i in range(1, n):
|
|
j = int(ibin(i, b, str=True)[::-1], 2)
|
|
if i < j:
|
|
a[i], a[j] = a[j], a[i]
|
|
|
|
ang = -2*pi/n if inverse else 2*pi/n
|
|
|
|
if dps is not None:
|
|
ang = ang.evalf(dps + 2)
|
|
|
|
w = [cos(ang*i) + I*sin(ang*i) for i in range(n // 2)]
|
|
|
|
h = 2
|
|
while h <= n:
|
|
hf, ut = h // 2, n // h
|
|
for i in range(0, n, h):
|
|
for j in range(hf):
|
|
u, v = a[i + j], expand_mul(a[i + j + hf]*w[ut * j])
|
|
a[i + j], a[i + j + hf] = u + v, u - v
|
|
h *= 2
|
|
|
|
if inverse:
|
|
a = [(x/n).evalf(dps) for x in a] if dps is not None \
|
|
else [x/n for x in a]
|
|
|
|
return a
|
|
|
|
|
|
def fft(seq, dps=None):
|
|
r"""
|
|
Performs the Discrete Fourier Transform (**DFT**) in the complex domain.
|
|
|
|
The sequence is automatically padded to the right with zeros, as the
|
|
*radix-2 FFT* requires the number of sample points to be a power of 2.
|
|
|
|
This method should be used with default arguments only for short sequences
|
|
as the complexity of expressions increases with the size of the sequence.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
seq : iterable
|
|
The sequence on which **DFT** is to be applied.
|
|
dps : Integer
|
|
Specifies the number of decimal digits for precision.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy import fft, ifft
|
|
|
|
>>> fft([1, 2, 3, 4])
|
|
[10, -2 - 2*I, -2, -2 + 2*I]
|
|
>>> ifft(_)
|
|
[1, 2, 3, 4]
|
|
|
|
>>> ifft([1, 2, 3, 4])
|
|
[5/2, -1/2 + I/2, -1/2, -1/2 - I/2]
|
|
>>> fft(_)
|
|
[1, 2, 3, 4]
|
|
|
|
>>> ifft([1, 7, 3, 4], dps=15)
|
|
[3.75, -0.5 - 0.75*I, -1.75, -0.5 + 0.75*I]
|
|
>>> fft(_)
|
|
[1.0, 7.0, 3.0, 4.0]
|
|
|
|
References
|
|
==========
|
|
|
|
.. [1] https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm
|
|
.. [2] https://mathworld.wolfram.com/FastFourierTransform.html
|
|
|
|
"""
|
|
|
|
return _fourier_transform(seq, dps=dps)
|
|
|
|
|
|
def ifft(seq, dps=None):
|
|
return _fourier_transform(seq, dps=dps, inverse=True)
|
|
|
|
ifft.__doc__ = fft.__doc__
|
|
|
|
|
|
#----------------------------------------------------------------------------#
|
|
# #
|
|
# Number Theoretic Transform #
|
|
# #
|
|
#----------------------------------------------------------------------------#
|
|
|
|
def _number_theoretic_transform(seq, prime, inverse=False):
|
|
"""Utility function for the Number Theoretic Transform"""
|
|
|
|
if not iterable(seq):
|
|
raise TypeError("Expected a sequence of integer coefficients "
|
|
"for Number Theoretic Transform")
|
|
|
|
p = as_int(prime)
|
|
if not isprime(p):
|
|
raise ValueError("Expected prime modulus for "
|
|
"Number Theoretic Transform")
|
|
|
|
a = [as_int(x) % p for x in seq]
|
|
|
|
n = len(a)
|
|
if n < 1:
|
|
return a
|
|
|
|
b = n.bit_length() - 1
|
|
if n&(n - 1):
|
|
b += 1
|
|
n = 2**b
|
|
|
|
if (p - 1) % n:
|
|
raise ValueError("Expected prime modulus of the form (m*2**k + 1)")
|
|
|
|
a += [0]*(n - len(a))
|
|
for i in range(1, n):
|
|
j = int(ibin(i, b, str=True)[::-1], 2)
|
|
if i < j:
|
|
a[i], a[j] = a[j], a[i]
|
|
|
|
pr = primitive_root(p)
|
|
|
|
rt = pow(pr, (p - 1) // n, p)
|
|
if inverse:
|
|
rt = pow(rt, p - 2, p)
|
|
|
|
w = [1]*(n // 2)
|
|
for i in range(1, n // 2):
|
|
w[i] = w[i - 1]*rt % p
|
|
|
|
h = 2
|
|
while h <= n:
|
|
hf, ut = h // 2, n // h
|
|
for i in range(0, n, h):
|
|
for j in range(hf):
|
|
u, v = a[i + j], a[i + j + hf]*w[ut * j]
|
|
a[i + j], a[i + j + hf] = (u + v) % p, (u - v) % p
|
|
h *= 2
|
|
|
|
if inverse:
|
|
rv = pow(n, p - 2, p)
|
|
a = [x*rv % p for x in a]
|
|
|
|
return a
|
|
|
|
|
|
def ntt(seq, prime):
|
|
r"""
|
|
Performs the Number Theoretic Transform (**NTT**), which specializes the
|
|
Discrete Fourier Transform (**DFT**) over quotient ring `Z/pZ` for prime
|
|
`p` instead of complex numbers `C`.
|
|
|
|
The sequence is automatically padded to the right with zeros, as the
|
|
*radix-2 NTT* requires the number of sample points to be a power of 2.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
seq : iterable
|
|
The sequence on which **DFT** is to be applied.
|
|
prime : Integer
|
|
Prime modulus of the form `(m 2^k + 1)` to be used for performing
|
|
**NTT** on the sequence.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy import ntt, intt
|
|
>>> ntt([1, 2, 3, 4], prime=3*2**8 + 1)
|
|
[10, 643, 767, 122]
|
|
>>> intt(_, 3*2**8 + 1)
|
|
[1, 2, 3, 4]
|
|
>>> intt([1, 2, 3, 4], prime=3*2**8 + 1)
|
|
[387, 415, 384, 353]
|
|
>>> ntt(_, prime=3*2**8 + 1)
|
|
[1, 2, 3, 4]
|
|
|
|
References
|
|
==========
|
|
|
|
.. [1] http://www.apfloat.org/ntt.html
|
|
.. [2] https://mathworld.wolfram.com/NumberTheoreticTransform.html
|
|
.. [3] https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general%29
|
|
|
|
"""
|
|
|
|
return _number_theoretic_transform(seq, prime=prime)
|
|
|
|
|
|
def intt(seq, prime):
|
|
return _number_theoretic_transform(seq, prime=prime, inverse=True)
|
|
|
|
intt.__doc__ = ntt.__doc__
|
|
|
|
|
|
#----------------------------------------------------------------------------#
|
|
# #
|
|
# Walsh Hadamard Transform #
|
|
# #
|
|
#----------------------------------------------------------------------------#
|
|
|
|
def _walsh_hadamard_transform(seq, inverse=False):
|
|
"""Utility function for the Walsh Hadamard Transform"""
|
|
|
|
if not iterable(seq):
|
|
raise TypeError("Expected a sequence of coefficients "
|
|
"for Walsh Hadamard Transform")
|
|
|
|
a = [sympify(arg) for arg in seq]
|
|
n = len(a)
|
|
if n < 2:
|
|
return a
|
|
|
|
if n&(n - 1):
|
|
n = 2**n.bit_length()
|
|
|
|
a += [S.Zero]*(n - len(a))
|
|
h = 2
|
|
while h <= n:
|
|
hf = h // 2
|
|
for i in range(0, n, h):
|
|
for j in range(hf):
|
|
u, v = a[i + j], a[i + j + hf]
|
|
a[i + j], a[i + j + hf] = u + v, u - v
|
|
h *= 2
|
|
|
|
if inverse:
|
|
a = [x/n for x in a]
|
|
|
|
return a
|
|
|
|
|
|
def fwht(seq):
|
|
r"""
|
|
Performs the Walsh Hadamard Transform (**WHT**), and uses Hadamard
|
|
ordering for the sequence.
|
|
|
|
The sequence is automatically padded to the right with zeros, as the
|
|
*radix-2 FWHT* requires the number of sample points to be a power of 2.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
seq : iterable
|
|
The sequence on which WHT is to be applied.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy import fwht, ifwht
|
|
>>> fwht([4, 2, 2, 0, 0, 2, -2, 0])
|
|
[8, 0, 8, 0, 8, 8, 0, 0]
|
|
>>> ifwht(_)
|
|
[4, 2, 2, 0, 0, 2, -2, 0]
|
|
|
|
>>> ifwht([19, -1, 11, -9, -7, 13, -15, 5])
|
|
[2, 0, 4, 0, 3, 10, 0, 0]
|
|
>>> fwht(_)
|
|
[19, -1, 11, -9, -7, 13, -15, 5]
|
|
|
|
References
|
|
==========
|
|
|
|
.. [1] https://en.wikipedia.org/wiki/Hadamard_transform
|
|
.. [2] https://en.wikipedia.org/wiki/Fast_Walsh%E2%80%93Hadamard_transform
|
|
|
|
"""
|
|
|
|
return _walsh_hadamard_transform(seq)
|
|
|
|
|
|
def ifwht(seq):
|
|
return _walsh_hadamard_transform(seq, inverse=True)
|
|
|
|
ifwht.__doc__ = fwht.__doc__
|
|
|
|
|
|
#----------------------------------------------------------------------------#
|
|
# #
|
|
# Mobius Transform for Subset Lattice #
|
|
# #
|
|
#----------------------------------------------------------------------------#
|
|
|
|
def _mobius_transform(seq, sgn, subset):
|
|
r"""Utility function for performing Mobius Transform using
|
|
Yate's Dynamic Programming method"""
|
|
|
|
if not iterable(seq):
|
|
raise TypeError("Expected a sequence of coefficients")
|
|
|
|
a = [sympify(arg) for arg in seq]
|
|
|
|
n = len(a)
|
|
if n < 2:
|
|
return a
|
|
|
|
if n&(n - 1):
|
|
n = 2**n.bit_length()
|
|
|
|
a += [S.Zero]*(n - len(a))
|
|
|
|
if subset:
|
|
i = 1
|
|
while i < n:
|
|
for j in range(n):
|
|
if j & i:
|
|
a[j] += sgn*a[j ^ i]
|
|
i *= 2
|
|
|
|
else:
|
|
i = 1
|
|
while i < n:
|
|
for j in range(n):
|
|
if j & i:
|
|
continue
|
|
a[j] += sgn*a[j ^ i]
|
|
i *= 2
|
|
|
|
return a
|
|
|
|
|
|
def mobius_transform(seq, subset=True):
|
|
r"""
|
|
Performs the Mobius Transform for subset lattice with indices of
|
|
sequence as bitmasks.
|
|
|
|
The indices of each argument, considered as bit strings, correspond
|
|
to subsets of a finite set.
|
|
|
|
The sequence is automatically padded to the right with zeros, as the
|
|
definition of subset/superset based on bitmasks (indices) requires
|
|
the size of sequence to be a power of 2.
|
|
|
|
Parameters
|
|
==========
|
|
|
|
seq : iterable
|
|
The sequence on which Mobius Transform is to be applied.
|
|
subset : bool
|
|
Specifies if Mobius Transform is applied by enumerating subsets
|
|
or supersets of the given set.
|
|
|
|
Examples
|
|
========
|
|
|
|
>>> from sympy import symbols
|
|
>>> from sympy import mobius_transform, inverse_mobius_transform
|
|
>>> x, y, z = symbols('x y z')
|
|
|
|
>>> mobius_transform([x, y, z])
|
|
[x, x + y, x + z, x + y + z]
|
|
>>> inverse_mobius_transform(_)
|
|
[x, y, z, 0]
|
|
|
|
>>> mobius_transform([x, y, z], subset=False)
|
|
[x + y + z, y, z, 0]
|
|
>>> inverse_mobius_transform(_, subset=False)
|
|
[x, y, z, 0]
|
|
|
|
>>> mobius_transform([1, 2, 3, 4])
|
|
[1, 3, 4, 10]
|
|
>>> inverse_mobius_transform(_)
|
|
[1, 2, 3, 4]
|
|
>>> mobius_transform([1, 2, 3, 4], subset=False)
|
|
[10, 6, 7, 4]
|
|
>>> inverse_mobius_transform(_, subset=False)
|
|
[1, 2, 3, 4]
|
|
|
|
References
|
|
==========
|
|
|
|
.. [1] https://en.wikipedia.org/wiki/M%C3%B6bius_inversion_formula
|
|
.. [2] https://people.csail.mit.edu/rrw/presentations/subset-conv.pdf
|
|
.. [3] https://arxiv.org/pdf/1211.0189.pdf
|
|
|
|
"""
|
|
|
|
return _mobius_transform(seq, sgn=+1, subset=subset)
|
|
|
|
def inverse_mobius_transform(seq, subset=True):
|
|
return _mobius_transform(seq, sgn=-1, subset=subset)
|
|
|
|
inverse_mobius_transform.__doc__ = mobius_transform.__doc__
|