Coverage for yasfpy/functions/misc.py: 78%

1import numpy as np 

2 

3from scipy.special import spherical_jn 

4from scipy.special import hankel1, lpmv 

5 

6from yasfpy.functions.legendre_normalized_trigon import legendre_normalized_trigon 

7 

8 

9def jmult_max(num_part, lmax): 

10 return 2 * num_part * lmax * (lmax + 2) 

11 

12 

13def multi2single_index(j_s, tau, l, m, lmax): 

14 return ( 

15 j_s * 2 * lmax * (lmax + 2) 

16 + (tau - 1) * lmax * (lmax + 2) 

17 + (l - 1) * (l + 1) 

18 + m 

19 + l 

20 ) 

21 

22 

23def single_index2multi(idx, lmax): 

24 j_s = idx // (2 * lmax * (lmax + 2)) 

25 idx_new = idx % (2 * lmax * (lmax + 2)) 

26 tau = idx_new // (lmax * (lmax + 2)) + 1 

27 idx_new = idx_new % (lmax * (lmax + 2)) 

28 l = np.floor(np.sqrt(idx_new + 1)) 

29 m = idx_new - (l * l + l - 1) 

30 return j_s, tau, l, m 

31 

32 

33def transformation_coefficients(pilm, taulm, tau, l, m, pol, dagger: bool = False): 

34 ifac = 1j 

35 if dagger: 

36 ifac *= -1 

37 

38 # Polarized light 

39 if np.any(np.equal(pol, [1, 2])): 

40 if tau == pol: 

41 spher_fun = taulm[l, np.abs(m)] 

42 else: 

43 spher_fun = m * pilm[l, np.abs(m)] 

44 

45 return ( 

46 -1 

47 / np.power(ifac, l + 1) 

48 / np.sqrt(2 * l * (l + 1)) 

49 * (ifac * (pol == 1) + (pol == 2)) 

50 * spher_fun 

51 ) 

52 

53 # Unpolarized light 

54 return ( 

55 -1 

56 / np.power(ifac, l + 1) 

57 / np.sqrt(2 * l * (l + 1)) 

58 * (ifac + 1) 

59 * (taulm[l, np.abs(m)] + m * pilm[l, np.abs(m)]) 

60 / 2 

61 ) 

62 

63 

64def mutual_lookup( 

65 lmax: int, 

66 positions_1: np.ndarray, 

67 positions_2: np.ndarray, 

68 refractive_index: np.ndarray, 

69 derivatives: bool = False, 

70 parallel: bool = False, 

71): 

72 differences = positions_1[:, np.newaxis, :] - positions_2[np.newaxis, :, :] 

73 distances = np.sqrt(np.sum(differences**2, axis=2)) 

74 distances = distances[:, :, np.newaxis] 

75 e_r = np.divide( 

76 differences, distances, out=np.zeros_like(differences), where=distances != 0 

77 ) 

78 cosine_theta = e_r[:, :, 2] 

79 # cosine_theta = np.divide( 

80 # differences[:, :, 2], 

81 # distances, 

82 # out = np.zeros_like(distances), 

83 # where = distances != 0 

84 # ) 

85 # correct possible rounding errors 

86 cosine_theta[cosine_theta < -1] = -1 

87 cosine_theta[cosine_theta > 1] = 1 

88 sine_theta = np.sqrt(1 - cosine_theta**2) 

89 phi = np.arctan2(differences[:, :, 1], differences[:, :, 0]) 

90 e_theta = np.stack( 

91 [cosine_theta * np.cos(phi), cosine_theta * np.sin(phi), -sine_theta], axis=-1 

92 ) 

93 e_phi = np.stack([-np.sin(phi), np.cos(phi), np.zeros_like(phi)], axis=-1) 

94 

95 size_parameter = distances * refractive_index[np.newaxis, np.newaxis, :] 

96 

97 if parallel: 

98 from yasfpy.functions.cpu_numba import compute_lookup_tables 

99 

100 spherical_bessel, spherical_hankel, e_j_dm_phi, p_lm = compute_lookup_tables( 

101 lmax, size_parameter, phi, cosine_theta 

102 ) 

103 else: 

104 p_range = np.arange(2 * lmax + 1) 

105 p_range = p_range[:, np.newaxis, np.newaxis, np.newaxis] 

106 size_parameter_extended = size_parameter[np.newaxis, :, :, :] 

107 spherical_hankel = np.sqrt( 

108 np.divide( 

109 np.pi / 2, 

110 size_parameter_extended, 

111 out=np.zeros_like(size_parameter_extended), 

112 where=size_parameter_extended != 0, 

113 ) 

114 ) * hankel1(p_range + 1 / 2, size_parameter_extended) 

115 spherical_bessel = spherical_jn(p_range, size_parameter_extended) 

116 

117 if derivatives: 

118 spherical_hankel_lower = np.sqrt( 

119 np.divide( 

120 np.pi / 2, 

121 size_parameter_extended, 

122 out=np.zeros_like(size_parameter_extended), 

123 where=size_parameter_extended != 0, 

124 ) 

125 ) * hankel1(-1 / 2, size_parameter_extended) 

126 spherical_hankel_lower = np.vstack( 

127 (spherical_hankel_lower, spherical_hankel[:-1, :, :, :]) 

128 ) 

129 spherical_hankel_derivative = ( 

130 size_parameter_extended * spherical_hankel_lower 

131 - p_range * spherical_hankel 

132 ) 

133 

134 # p_range = np.arange(2 * lmax + 2) - 1 

135 # p_range = p_range[:, np.newaxis, np.newaxis, np.newaxis] 

136 # spherical_hankel = np.sqrt(np.divide(np.pi / 2, size_parameter_extended, out = np.zeros_like(size_parameter_extended), where = size_parameter_extended != 0)) * hankel1(p_range + 1/2, size_parameter_extended) 

137 # spherical_hankel_derivative = size_parameter_extended * spherical_hankel[:-1, :, :, :] - p_range[1:, :, :, :] * spherical_hankel[1:, :, :, :] 

138 

139 p_lm = legendre_normalized_trigon(lmax, cosine_theta, sine_theta) 

140 else: 

141 spherical_hankel_derivative = None 

142 

143 p_lm = np.zeros( 

144 (lmax + 1) * (2 * lmax + 1) * np.prod(size_parameter.shape[:2]) 

145 ).reshape(((lmax + 1) * (2 * lmax + 1),) + size_parameter.shape[:2]) 

146 for p in range(2 * lmax + 1): 

147 for absdm in range(p + 1): 

148 cml = np.sqrt( 

149 (2 * p + 1) 

150 / 2 

151 * np.prod(1 / np.arange(p - absdm + 1, p + absdm + 1)) 

152 ) 

153 # if np.isnan(cml): 

154 # print(p) 

155 # print(absdm) 

156 p_lm[p * (p + 1) // 2 + absdm, :, :] = ( 

157 cml * np.power(-1.0, absdm) * lpmv(absdm, p, cosine_theta) 

158 ) 

159 

160 phi = phi[np.newaxis, :, :] 

161 p_range = np.arange(-2 * lmax, 2 * lmax + 1) 

162 p_range = p_range[:, np.newaxis, np.newaxis] 

163 e_j_dm_phi = np.exp(1j * p_range * phi) 

164 

165 return ( 

166 spherical_bessel, 

167 spherical_hankel, 

168 e_j_dm_phi, 

169 p_lm, 

170 e_r, 

171 e_theta, 

172 e_phi, 

173 cosine_theta, 

174 sine_theta, 

175 size_parameter, 

176 spherical_hankel_derivative, 

177 )