1from diofant import (DiracDelta, Function, Heaviside, Rational, cos, exp, pi,
2                     sin, symbols)
3from diofant.integrals.deltafunctions import change_mul, deltaintegrate
4
5
6__all__ = ()
7
8f = Function('f')
9x_1, x_2, x, y, z = symbols('x_1 x_2 x y z')
10
11
12def test_change_mul():
13    assert change_mul(x, x) == x
14    assert change_mul(x*y, x) == (None, None)
15    assert change_mul(x*y*DiracDelta(x), x) == (DiracDelta(x), x*y)
16    assert change_mul(x*y*DiracDelta(x)*DiracDelta(y), x) == \
17        (DiracDelta(x), x*y*DiracDelta(y))
18    assert change_mul(DiracDelta(x)**2, x) == \
19        (DiracDelta(x), DiracDelta(x))
20    assert change_mul(y*DiracDelta(x)**2, x) == \
21        (DiracDelta(x), y*DiracDelta(x))
22
23
24def test_deltaintegrate():
25    assert deltaintegrate(x, x) is None
26    assert deltaintegrate(x + DiracDelta(x), x) is None
27    assert deltaintegrate(DiracDelta(x, 0), x) == Heaviside(x)
28    for n in range(10):
29        assert deltaintegrate(DiracDelta(x, n + 1), x) == DiracDelta(x, n)
30    assert deltaintegrate(DiracDelta(x), x) == Heaviside(x)
31    assert deltaintegrate(DiracDelta(-x), x) == Heaviside(x)
32    assert deltaintegrate(DiracDelta(x - y), x) == Heaviside(x - y)
33    assert deltaintegrate(DiracDelta(y - x), x) == Heaviside(x - y)
34
35    assert deltaintegrate(x*DiracDelta(x), x) == 0
36    assert deltaintegrate((x - y)*DiracDelta(x - y), x) == 0
37
38    assert deltaintegrate(DiracDelta(x)**2, x) == DiracDelta(0)*Heaviside(x)
39    assert deltaintegrate(y*DiracDelta(x)**2, x) == \
40        y*DiracDelta(0)*Heaviside(x)
41    assert deltaintegrate(DiracDelta(x, 1)**2, x) is None
42    assert deltaintegrate(y*DiracDelta(x, 1)**2, x) is None
43
44    assert deltaintegrate(DiracDelta(x) * f(x), x) == f(0) * Heaviside(x)
45    assert deltaintegrate(DiracDelta(-x) * f(x), x) == f(0) * Heaviside(x)
46    assert deltaintegrate(DiracDelta(x - 1) * f(x), x) == f(1) * Heaviside(x - 1)
47    assert deltaintegrate(DiracDelta(1 - x) * f(x), x) == f(1) * Heaviside(x - 1)
48    assert deltaintegrate(DiracDelta(x**2 + x - 2), x) == \
49        Heaviside(x - 1)/3 + Heaviside(x + 2)/3
50
51    p = cos(x)*(DiracDelta(x) + DiracDelta(x**2 - 1))*sin(x)*(x - pi)
52    assert deltaintegrate(p, x) - (-pi*(cos(1)*Heaviside(-1 + x)*sin(1)/2 -
53                                        cos(1)*Heaviside(1 + x)*sin(1)/2) +
54                                   cos(1)*Heaviside(1 + x)*sin(1)/2 +
55                                   cos(1)*Heaviside(-1 + x)*sin(1)/2) == 0
56
57    p = x_2*DiracDelta(x - x_2)*DiracDelta(x_2 - x_1)
58    assert deltaintegrate(p, x_2) == x*DiracDelta(x - x_1)*Heaviside(x_2 - x)
59
60    p = x*y**2*z*DiracDelta(y - x)*DiracDelta(y - z)*DiracDelta(x - z)
61    assert deltaintegrate(p, y) == x**3*z*DiracDelta(x - z)**2*Heaviside(y - x)
62    assert deltaintegrate((x + 1)*DiracDelta(2*x), x) == Rational(1, 2) * Heaviside(x)
63    assert deltaintegrate((x + 1)*DiracDelta(2*x/3 + Rational(4, 9)), x) == \
64        Rational(1, 2) * Heaviside(x + Rational(2, 3))
65
66    a, b, c = symbols('a b c', commutative=False)
67    assert deltaintegrate(DiracDelta(x - y)*f(x - b)*f(x - a), x) == \
68        f(y - b)*f(y - a)*Heaviside(x - y)
69
70    p = f(x - a)*DiracDelta(x - y)*f(x - c)*f(x - b)
71    assert deltaintegrate(p, x) == f(y - a)*f(y - c)*f(y - b)*Heaviside(x - y)
72
73    p = DiracDelta(x - z)*f(x - b)*f(x - a)*DiracDelta(x - y)
74    assert deltaintegrate(p, x) == DiracDelta(y - z)*f(y - b)*f(y - a) * \
75        Heaviside(x - y)
76
77    # for test coverage
78    assert deltaintegrate(DiracDelta(exp(x) - 1), x) is None
79    assert deltaintegrate((x - 1)*DiracDelta(exp(x) - 1), x) is None
80