1import pytest
2
3from diofant import (I, Symbol, Ynm, Ynm_c, Znm, assoc_legendre, conjugate,
4                     cos, cot, diff, exp, factorial, pi, sin, sqrt)
5from diofant.abc import m, n
6from diofant.core.function import ArgumentIndexError
7
8
9__all__ = ()
10
11
12def test_Ynm():
13    # https://en.wikipedia.org/wiki/Spherical_harmonics
14    th, ph = Symbol('theta', extended_real=True), Symbol('phi', extended_real=True)
15
16    assert Ynm(0, 0, th, ph).expand(func=True) == 1/(2*sqrt(pi))
17    assert Ynm(1, -1, th, ph) == -exp(-2*I*ph)*Ynm(1, 1, th, ph)
18    assert Ynm(1, -1, th, ph).expand(func=True) == sqrt(6)*sin(th)*exp(-I*ph)/(4*sqrt(pi))
19    assert Ynm(1, -1, th, ph).expand(func=True) == sqrt(6)*sin(th)*exp(-I*ph)/(4*sqrt(pi))
20    assert Ynm(1, 0, th, ph).expand(func=True) == sqrt(3)*cos(th)/(2*sqrt(pi))
21    assert Ynm(1, 1, th, ph).expand(func=True) == -sqrt(6)*sin(th)*exp(I*ph)/(4*sqrt(pi))
22    assert Ynm(2, 0, th, ph).expand(func=True) == 3*sqrt(5)*cos(th)**2/(4*sqrt(pi)) - sqrt(5)/(4*sqrt(pi))
23    assert Ynm(2, 1, th, ph).expand(func=True) == -sqrt(30)*sin(th)*exp(I*ph)*cos(th)/(4*sqrt(pi))
24    assert Ynm(2, -2, th, ph).expand(func=True) == (-sqrt(30)*exp(-2*I*ph)*cos(th)**2/(8*sqrt(pi))
25                                                    + sqrt(30)*exp(-2*I*ph)/(8*sqrt(pi)))
26    assert Ynm(2, 2, th, ph).expand(func=True) == (-sqrt(30)*exp(2*I*ph)*cos(th)**2/(8*sqrt(pi))
27                                                   + sqrt(30)*exp(2*I*ph)/(8*sqrt(pi)))
28
29    assert diff(Ynm(n, m, th, ph), th) == (m*cot(th)*Ynm(n, m, th, ph)
30                                           + sqrt((-m + n)*(m + n + 1))*exp(-I*ph)*Ynm(n, m + 1, th, ph))
31    assert diff(Ynm(n, m, th, ph), ph) == I*m*Ynm(n, m, th, ph)
32    pytest.raises(ArgumentIndexError, lambda: Ynm(n, m, th, ph).fdiff(1))
33
34    assert conjugate(Ynm(n, m, th, ph)) == (-1)**(2*m)*exp(-2*I*m*ph)*Ynm(n, m, th, ph)
35
36    assert Ynm(n, m, -th, ph) == Ynm(n, m, th, ph)
37    assert Ynm(n, m, th, -ph) == exp(-2*I*m*ph)*Ynm(n, m, th, ph)
38    assert Ynm(n, -m, th, ph) == (-1)**m*exp(-2*I*m*ph)*Ynm(n, m, th, ph)
39
40    assert (Ynm(n, m, th, ph).rewrite(sin) ==
41            Ynm(n, m, th, ph).rewrite(cos) ==
42            exp(I*m*ph)*sqrt((2*n + 1)*factorial(-m + n)/factorial(m + n)) *
43            assoc_legendre(n, m, cos(th))/(2*sqrt(pi)))
44    assert (Ynm(n, m, th, ph).as_real_imag() ==
45            (sqrt((2*n + 1)*factorial(-m + n)/factorial(m + n))*cos(m*ph) *
46             assoc_legendre(n, m, cos(th))/(2*sqrt(pi)),
47             sqrt((2*n + 1)*factorial(-m + n)/factorial(m + n))*sin(m*ph) *
48             assoc_legendre(n, m, cos(th))/(2*sqrt(pi))))
49
50
51def test_Ynm_c():
52    th, ph = Symbol('theta', extended_real=True), Symbol('phi', extended_real=True)
53
54    assert Ynm_c(n, m, th, ph) == (-1)**(2*m)*exp(-2*I*m*ph)*Ynm(n, m, th, ph)
55
56
57def test_Znm():
58    # https://en.wikipedia.org/wiki/Solid_harmonics#List_of_lowest_functions
59    th, ph = Symbol('theta', extended_real=True), Symbol('phi', extended_real=True)
60
61    assert Znm(0, 0, th, ph) == Ynm(0, 0, th, ph)
62    assert Znm(1, -1, th, ph) == (-sqrt(2)*I*(Ynm(1, 1, th, ph)
63                                              - exp(-2*I*ph)*Ynm(1, 1, th, ph))/2)
64    assert Znm(1, 0, th, ph) == Ynm(1, 0, th, ph)
65    assert Znm(1, 1, th, ph) == (sqrt(2)*(Ynm(1, 1, th, ph)
66                                          + exp(-2*I*ph)*Ynm(1, 1, th, ph))/2)
67    assert Znm(0, 0, th, ph).expand(func=True) == 1/(2*sqrt(pi))
68    assert Znm(1, -1, th, ph).expand(func=True) == (sqrt(3)*I*sin(th)*exp(I*ph)/(4*sqrt(pi))
69                                                    - sqrt(3)*I*sin(th)*exp(-I*ph)/(4*sqrt(pi)))
70    assert Znm(1, 0, th, ph).expand(func=True) == sqrt(3)*cos(th)/(2*sqrt(pi))
71    assert Znm(1, 1, th, ph).expand(func=True) == (-sqrt(3)*sin(th)*exp(I*ph)/(4*sqrt(pi))
72                                                   - sqrt(3)*sin(th)*exp(-I*ph)/(4*sqrt(pi)))
73    assert Znm(2, -1, th, ph).expand(func=True) == (sqrt(15)*I*sin(th)*exp(I*ph)*cos(th)/(4*sqrt(pi))
74                                                    - sqrt(15)*I*sin(th)*exp(-I*ph)*cos(th)/(4*sqrt(pi)))
75    assert Znm(2, 0, th, ph).expand(func=True) == 3*sqrt(5)*cos(th)**2/(4*sqrt(pi)) - sqrt(5)/(4*sqrt(pi))
76    assert Znm(2, 1, th, ph).expand(func=True) == (-sqrt(15)*sin(th)*exp(I*ph)*cos(th)/(4*sqrt(pi))
77                                                   - sqrt(15)*sin(th)*exp(-I*ph)*cos(th)/(4*sqrt(pi)))
78