You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
154 lines
5.4 KiB
154 lines
5.4 KiB
from sympy.functions.elementary.miscellaneous import sqrt
|
|
from sympy.core import S, Symbol, symbols, I, Rational
|
|
from sympy.discrete import (fft, ifft, ntt, intt, fwht, ifwht,
|
|
mobius_transform, inverse_mobius_transform)
|
|
from sympy.testing.pytest import raises
|
|
|
|
|
|
def test_fft_ifft():
|
|
assert all(tf(ls) == ls for tf in (fft, ifft)
|
|
for ls in ([], [Rational(5, 3)]))
|
|
|
|
ls = list(range(6))
|
|
fls = [15, -7*sqrt(2)/2 - 4 - sqrt(2)*I/2 + 2*I, 2 + 3*I,
|
|
-4 + 7*sqrt(2)/2 - 2*I - sqrt(2)*I/2, -3,
|
|
-4 + 7*sqrt(2)/2 + sqrt(2)*I/2 + 2*I,
|
|
2 - 3*I, -7*sqrt(2)/2 - 4 - 2*I + sqrt(2)*I/2]
|
|
|
|
assert fft(ls) == fls
|
|
assert ifft(fls) == ls + [S.Zero]*2
|
|
|
|
ls = [1 + 2*I, 3 + 4*I, 5 + 6*I]
|
|
ifls = [Rational(9, 4) + 3*I, I*Rational(-7, 4), Rational(3, 4) + I, -2 - I/4]
|
|
|
|
assert ifft(ls) == ifls
|
|
assert fft(ifls) == ls + [S.Zero]
|
|
|
|
x = Symbol('x', real=True)
|
|
raises(TypeError, lambda: fft(x))
|
|
raises(ValueError, lambda: ifft([x, 2*x, 3*x**2, 4*x**3]))
|
|
|
|
|
|
def test_ntt_intt():
|
|
# prime moduli of the form (m*2**k + 1), sequence length
|
|
# should be a divisor of 2**k
|
|
p = 7*17*2**23 + 1
|
|
q = 2*500000003 + 1 # only for sequences of length 1 or 2
|
|
r = 2*3*5*7 # composite modulus
|
|
|
|
assert all(tf(ls, p) == ls for tf in (ntt, intt)
|
|
for ls in ([], [5]))
|
|
|
|
ls = list(range(6))
|
|
nls = [15, 801133602, 738493201, 334102277, 998244350, 849020224,
|
|
259751156, 12232587]
|
|
|
|
assert ntt(ls, p) == nls
|
|
assert intt(nls, p) == ls + [0]*2
|
|
|
|
ls = [1 + 2*I, 3 + 4*I, 5 + 6*I]
|
|
x = Symbol('x', integer=True)
|
|
|
|
raises(TypeError, lambda: ntt(x, p))
|
|
raises(ValueError, lambda: intt([x, 2*x, 3*x**2, 4*x**3], p))
|
|
raises(ValueError, lambda: intt(ls, p))
|
|
raises(ValueError, lambda: ntt([1.2, 2.1, 3.5], p))
|
|
raises(ValueError, lambda: ntt([3, 5, 6], q))
|
|
raises(ValueError, lambda: ntt([4, 5, 7], r))
|
|
raises(ValueError, lambda: ntt([1.0, 2.0, 3.0], p))
|
|
|
|
|
|
def test_fwht_ifwht():
|
|
assert all(tf(ls) == ls for tf in (fwht, ifwht) \
|
|
for ls in ([], [Rational(7, 4)]))
|
|
|
|
ls = [213, 321, 43235, 5325, 312, 53]
|
|
fls = [49459, 38061, -47661, -37759, 48729, 37543, -48391, -38277]
|
|
|
|
assert fwht(ls) == fls
|
|
assert ifwht(fls) == ls + [S.Zero]*2
|
|
|
|
ls = [S.Half + 2*I, Rational(3, 7) + 4*I, Rational(5, 6) + 6*I, Rational(7, 3), Rational(9, 4)]
|
|
ifls = [Rational(533, 672) + I*3/2, Rational(23, 224) + I/2, Rational(1, 672), Rational(107, 224) - I,
|
|
Rational(155, 672) + I*3/2, Rational(-103, 224) + I/2, Rational(-377, 672), Rational(-19, 224) - I]
|
|
|
|
assert ifwht(ls) == ifls
|
|
assert fwht(ifls) == ls + [S.Zero]*3
|
|
|
|
x, y = symbols('x y')
|
|
|
|
raises(TypeError, lambda: fwht(x))
|
|
|
|
ls = [x, 2*x, 3*x**2, 4*x**3]
|
|
ifls = [x**3 + 3*x**2/4 + x*Rational(3, 4),
|
|
-x**3 + 3*x**2/4 - x/4,
|
|
-x**3 - 3*x**2/4 + x*Rational(3, 4),
|
|
x**3 - 3*x**2/4 - x/4]
|
|
|
|
assert ifwht(ls) == ifls
|
|
assert fwht(ifls) == ls
|
|
|
|
ls = [x, y, x**2, y**2, x*y]
|
|
fls = [x**2 + x*y + x + y**2 + y,
|
|
x**2 + x*y + x - y**2 - y,
|
|
-x**2 + x*y + x - y**2 + y,
|
|
-x**2 + x*y + x + y**2 - y,
|
|
x**2 - x*y + x + y**2 + y,
|
|
x**2 - x*y + x - y**2 - y,
|
|
-x**2 - x*y + x - y**2 + y,
|
|
-x**2 - x*y + x + y**2 - y]
|
|
|
|
assert fwht(ls) == fls
|
|
assert ifwht(fls) == ls + [S.Zero]*3
|
|
|
|
ls = list(range(6))
|
|
|
|
assert fwht(ls) == [x*8 for x in ifwht(ls)]
|
|
|
|
|
|
def test_mobius_transform():
|
|
assert all(tf(ls, subset=subset) == ls
|
|
for ls in ([], [Rational(7, 4)]) for subset in (True, False)
|
|
for tf in (mobius_transform, inverse_mobius_transform))
|
|
|
|
w, x, y, z = symbols('w x y z')
|
|
|
|
assert mobius_transform([x, y]) == [x, x + y]
|
|
assert inverse_mobius_transform([x, x + y]) == [x, y]
|
|
assert mobius_transform([x, y], subset=False) == [x + y, y]
|
|
assert inverse_mobius_transform([x + y, y], subset=False) == [x, y]
|
|
|
|
assert mobius_transform([w, x, y, z]) == [w, w + x, w + y, w + x + y + z]
|
|
assert inverse_mobius_transform([w, w + x, w + y, w + x + y + z]) == \
|
|
[w, x, y, z]
|
|
assert mobius_transform([w, x, y, z], subset=False) == \
|
|
[w + x + y + z, x + z, y + z, z]
|
|
assert inverse_mobius_transform([w + x + y + z, x + z, y + z, z], subset=False) == \
|
|
[w, x, y, z]
|
|
|
|
ls = [Rational(2, 3), Rational(6, 7), Rational(5, 8), 9, Rational(5, 3) + 7*I]
|
|
mls = [Rational(2, 3), Rational(32, 21), Rational(31, 24), Rational(1873, 168),
|
|
Rational(7, 3) + 7*I, Rational(67, 21) + 7*I, Rational(71, 24) + 7*I,
|
|
Rational(2153, 168) + 7*I]
|
|
|
|
assert mobius_transform(ls) == mls
|
|
assert inverse_mobius_transform(mls) == ls + [S.Zero]*3
|
|
|
|
mls = [Rational(2153, 168) + 7*I, Rational(69, 7), Rational(77, 8), 9, Rational(5, 3) + 7*I, 0, 0, 0]
|
|
|
|
assert mobius_transform(ls, subset=False) == mls
|
|
assert inverse_mobius_transform(mls, subset=False) == ls + [S.Zero]*3
|
|
|
|
ls = ls[:-1]
|
|
mls = [Rational(2, 3), Rational(32, 21), Rational(31, 24), Rational(1873, 168)]
|
|
|
|
assert mobius_transform(ls) == mls
|
|
assert inverse_mobius_transform(mls) == ls
|
|
|
|
mls = [Rational(1873, 168), Rational(69, 7), Rational(77, 8), 9]
|
|
|
|
assert mobius_transform(ls, subset=False) == mls
|
|
assert inverse_mobius_transform(mls, subset=False) == ls
|
|
|
|
raises(TypeError, lambda: mobius_transform(x, subset=True))
|
|
raises(TypeError, lambda: inverse_mobius_transform(y, subset=False))
|