Coverage for local_installation_linux / mumott / core / spherical_harmonic_utilities.py: 73%

1import numpy as np 

2from scipy.special import factorial 

3 

4try: 

5 from scipy.special import assoc_legendre_p 

6 

7 def legendre(n, m, z): 

8 # Get rid of extra dimension 

9 return assoc_legendre_p(n, m, z)[0] 

10except ImportError: 

11 from scipy.special import lpmv 

12 

13 def legendre(n, m, z): 

14 # Call signature compatibility 

15 return lpmv(m, n, z) 

16 

17 

18def sph_harm(emms: np.ndarray[int], ells: np.ndarray[int], 

19 azimuthal_angles: np.ndarray[float], 

20 z: np.ndarray[float]): 

21 """ 

22 Utility for computing spherical harmonics using the z-coordinates from 

23 :class:`ProbedCoordinates <mumott.core.probed_coordinates.ProbedCoordinates>` 

24 irrespective of SciPy version. 

25 """ 

26 legendre_part = legendre(ells, emms, z) 

27 norm = np.sqrt(((2 * ells + 1) / (4 * np.pi)) * (factorial(ells - emms) / factorial(ells + emms))) 

28 phase = np.exp(1j * emms * azimuthal_angles) 

29 return legendre_part * norm * phase