Spaces:
Running
Running
File size: 5,546 Bytes
6a86ad5 |
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 |
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))
|