test_sph_harm.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. import numpy as np
  2. import pytest
  3. from numpy.testing import assert_allclose
  4. import scipy.special as sc
  5. class TestSphHarm:
  6. @pytest.mark.slow
  7. def test_p(self):
  8. m_max = 20
  9. n_max = 10
  10. theta = np.linspace(0, np.pi)
  11. phi = np.linspace(0, 2*np.pi)
  12. theta, phi = np.meshgrid(theta, phi)
  13. y, y_jac, y_hess = sc.sph_harm_y_all(n_max, m_max, theta, phi, diff_n=2)
  14. p, p_jac, p_hess = sc.sph_legendre_p_all(n_max, m_max, theta, diff_n=2)
  15. m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)])
  16. m = np.expand_dims(m, axis=(0,)+tuple(range(2,theta.ndim+2)))
  17. assert_allclose(y, p * np.exp(1j * m * phi))
  18. assert_allclose(y_jac[..., 0], p_jac * np.exp(1j * m * phi))
  19. assert_allclose(y_jac[..., 1], 1j * m * p * np.exp(1j * m * phi))
  20. assert_allclose(y_hess[..., 0, 0], p_hess * np.exp(1j * m * phi))
  21. assert_allclose(y_hess[..., 0, 1], 1j * m * p_jac * np.exp(1j * m * phi))
  22. assert_allclose(y_hess[..., 1, 0], y_hess[..., 0, 1])
  23. assert_allclose(y_hess[..., 1, 1], -m * m * p * np.exp(1j * m * phi))
  24. @pytest.mark.parametrize("n_max", [7, 10, 50])
  25. @pytest.mark.parametrize("m_max", [1, 4, 5, 9, 14])
  26. def test_all(self, n_max, m_max):
  27. theta = np.linspace(0, np.pi)
  28. phi = np.linspace(0, 2 * np.pi)
  29. n = np.arange(n_max + 1)
  30. n = np.expand_dims(n, axis=tuple(range(1,theta.ndim+2)))
  31. m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)])
  32. m = np.expand_dims(m, axis=(0,)+tuple(range(2,theta.ndim+2)))
  33. y_actual = sc.sph_harm_y_all(n_max, m_max, theta, phi)
  34. y_desired = sc.sph_harm_y(n, m, theta, phi)
  35. np.testing.assert_allclose(y_actual, y_desired, rtol=1e-05)