Coverage for yasfpy/functions/spherical_functions_trigon.py: 93%

1import numpy as np 

2import math 

3 

4 

5def spherical_functions_trigon(lmax, theta, st=None): 

6 size = np.array([1]) 

7 if isinstance(theta, np.ndarray): 

8 size = theta.shape 

9 

10 if st is None: 

11 ct = np.cos(theta) 

12 st = np.sin(theta) 

13 else: 

14 ct = np.array(theta) 

15 st = np.array(st) 

16 

17 ct = ct.ravel() 

18 st = st.ravel() 

19 

20 plm = np.zeros((lmax + 1, lmax + 1, ct.shape[0])) * np.nan 

21 pilm = np.zeros(plm.shape) * np.nan 

22 taulm = np.zeros(plm.shape) * np.nan 

23 pprimel0 = np.zeros((lmax + 1, ct.shape[0])) * np.nan 

24 

25 plm[0, 0, :] = np.sqrt(1 / 2) * np.ones_like(ct) 

26 plm[1, 0, :] = np.sqrt(3 / 2) * ct 

27 

28 pilm[0, 0, :] = np.zeros_like(ct) 

29 pilm[1, 0, :] = np.zeros_like(ct) 

30 

31 pprimel0[0, :] = np.zeros_like(ct) 

32 pprimel0[1, :] = np.sqrt(3) * plm[0, 0, :] 

33 

34 taulm[0, 0, :] = -st * pprimel0[0, :] 

35 taulm[1, 0, :] = -st * pprimel0[1, :] 

36 

37 for l in range(1, lmax): 

38 plm[l + 1, 0, :] = ( 

39 1 / (l + 1) * np.sqrt((2 * l + 1) * (2 * l + 3)) * ct * plm[l, 0, :] 

40 - l / (l + 1) * np.sqrt((2 * l + 3) / (2 * l - 1)) * plm[l - 1, 0, :] 

41 ) 

42 pilm[l + 1, 0, :] = np.zeros_like(ct) 

43 pprimel0[l + 1, :] = np.sqrt((2 * l + 3) / (2 * l + 1)) * ( 

44 (l + 1) * plm[l, 0, :] + ct * pprimel0[l, :] 

45 ) 

46 taulm[l + 1, 0, :] = -st * pprimel0[l + 1, :] 

47 

48 for m in range(1, lmax + 1): 

49 plm[m - 1, m, :] = np.zeros_like(ct) 

50 pilm[m - 1, m, :] = np.zeros_like(ct) 

51 coeff = np.sqrt((2 * m + 1) / 2 / math.factorial(2 * m)) * np.prod( 

52 np.arange(1, 2 * m, 2) 

53 ) 

54 plm[m, m, :] = coeff * np.power(st, m) 

55 pilm[m, m, :] = coeff * np.power(st, m - 1) 

56 taulm[m, m, :] = m * ct * pilm[m, m, :] 

57 for l in range(m, lmax): 

58 coeff1 = np.sqrt((2 * l + 1) * (2 * l + 3) / (l + 1 - m) / (l + 1 + m)) * ct 

59 coeff2 = np.sqrt( 

60 (2 * l + 3) 

61 * (l - m) 

62 * (l + m) 

63 / (2 * l - 1) 

64 / (l + 1 - m) 

65 / (l + 1 + m) 

66 ) 

67 plm[l + 1, m, :] = coeff1 * plm[l, m, :] - coeff2 * plm[l - 1, m, :] 

68 pilm[l + 1, m, :] = coeff1 * pilm[l, m, :] - coeff2 * pilm[l - 1, m, :] 

69 taulm[l + 1, m, :] = (l + 1) * ct * pilm[l + 1, m, :] - ( 

70 l + 1 + m 

71 ) * np.sqrt((2 * l + 3) * (l + 1 - m) / (2 * l + 1) / (l + 1 + m)) * pilm[ 

72 l, m, : 

73 ] 

74 

75 pilm = np.reshape(pilm, np.concatenate(([lmax + 1, lmax + 1], size))) 

76 taulm = np.reshape(taulm, np.concatenate(([lmax + 1, lmax + 1], size))) 

77 return pilm, taulm