Coverage for local_installation_linux/mumott/core/wigner_d_utilities.py: 91%

1import numpy as np 

2import importlib.resources 

3import h5py 

4import logging 

5from typing import List 

6from numpy.typing import NDArray 

7 

8logger = logging.getLogger(__name__) 

9 

10 

11def load_d_matrices(ell_max: int) -> List[NDArray[float]]: 

12 """ Load Wigner (small) d-matrices corresponding to a rotation by 90 degrees 

13 around the x-axis up to order :attr:`ell_max`. 

14 The sign and normalization conventions are those adopted throughout :program:`mumott`. 

15 

16 Parameters 

17 ---------- 

18 ell_max 

19 maximum order of the Wigner (small) d-matrices to return 

20 

21 Returns 

22 --------- 

23 list of Wigner (small) d-matrices as numpy arrays with shape ``(l*2+1)x(l*2+1)`` 

24 """ 

25 d_matrices = [] 

26 # Open h5 file 

27 with importlib.resources.path(__package__, 'wigner_d_matrices.h5') as p: 

28 path = p 

29 

30 with h5py.File(path, 'r') as file: 

31 for ell in range(ell_max + 1): 

32 d_matrices.append(file[str(ell)][...].T) 

33 return d_matrices 

34 

35 

36def calculate_sph_coefficients_rotated_by_90_degrees_around_positive_x( 

37 input_array: NDArray[float], 

38 ell_list: List[int], 

39 d_matrices: List[NDArray[float]], 

40 output_array: NDArray[float] = None, 

41 ) -> NDArray[float]: 

42 r""" Calculate the spherical harmonic coefficients after a rotation by 90 degrees around +x. 

43 

44 Parameters 

45 ---------- 

46 input_array 

47 spherical harmonic coefficients 

48 ell_list 

49 list of :math:`\ell` values to use. If ``None`` (default) we assume that all even orders are inluded 

50 until some :math:`\ell_\mathrm{max}`, which will be calculated from the shape of the coefficients 

51 array. 

52 d_matrices 

53 list of Wigner (small) d-matrices corresponding to a 90 degree rotation around x. 

54 If ``None`` (default) pre-calculated matrix elements are used. 

55 output_array 

56 numpy array with same shape as :attr:`input_array`. 

57 If ``None`` (default) a new array is initialized. 

58 If :attr:`output_array` is given, the calculations are carried out in place. 

59 

60 Returns 

61 --------- 

62 numpy array with the same shape as :attr:`input_array`. 

63 """ 

64 if output_array is None: 64 ↛ 65line 64 didn't jump to line 65, because the condition on line 64 was never true

65 output_array = np.zeros(input_array.shape, dtype=float) 

66 

67 start_index = 0 

68 for ell in ell_list: 

69 output_array[..., start_index:start_index+(2*ell+1)]\ 

70 = np.einsum('...m,nm->...n', input_array[..., start_index:start_index+(2*ell+1)], d_matrices[ell]) 

71 start_index += 2 * ell + 1 

72 return output_array 

73 

74 

75def calculate_sph_coefficients_rotated_by_90_degrees_around_negative_x( 

76 input_array: NDArray[float], 

77 ell_list: List[int], 

78 d_matrices: List[NDArray[float]], 

79 output_array: NDArray[float] = None, 

80 ) -> NDArray[float]: 

81 r""" Calculate the spherical harmonic coefficients after a rotation by 90 degrees around -x. 

82 

83 Parameters 

84 ---------- 

85 input_array 

86 spherical harmonic coefficients 

87 ell_list 

88 list of :math:`\ell` values to use. If ``None`` (default) we assume that all even orders are inluded 

89 until some :math:`\ell_\mathrm{max}`, which will be calculated from the shape of the coefficients 

90 array. 

91 d_matrices 

92 list of Wigner (small) d-matrices corresponding to a 90 degree rotation around x. 

93 If ``None`` (default) pre-calculated matrix elements are used. 

94 output_array 

95 numpy array with same shape as :attr:`input_array`. 

96 If ``None`` (default) a new array is initialized. 

97 If :attr:`output_array` is given, the calculations are carried out in place. 

98 

99 Returns 

100 --------- 

101 numpy array with same shape as :attr:`input_array`. 

102 """ 

103 if output_array is None: 103 ↛ 104line 103 didn't jump to line 104, because the condition on line 103 was never true

104 output_array = np.zeros(input_array.shape, dtype=float) 

105 

106 start_index = 0 

107 for ell in ell_list: 

108 output_array[..., start_index:start_index+(2*ell+1)]\ 

109 = np.einsum('...m,mn->...n', input_array[..., start_index:start_index+(2*ell+1)], d_matrices[ell]) 

110 start_index += 2 * ell + 1 

111 return output_array 

112 

113 

114def calculate_sph_coefficients_rotated_around_z( 

115 input_array: NDArray[float], 

116 angle: NDArray[float], 

117 ell_list: List[int], 

118 output_array: NDArray[float] = None, 

119 ) -> NDArray[float]: 

120 r""" Calculate the spherical harmonic coefficients after a rotation around +z. 

121 

122 Parameters 

123 ---------- 

124 input_array 

125 spherical harmonic coefficients 

126 angle 

127 either an array with the same shape as ``input_array.shape[:-1]`` or a scalar. 

128 If a scalar is given all coefficient lists are rotated by the same angle. 

129 ell_list 

130 list of :math:`\ell` values to use. If ``None`` (default) we assume that all even orders are inluded 

131 until some :math:`\ell_\mathrm{max}`, which will be calculated from the shape of the coefficients 

132 array. 

133 d_matrices 

134 list of Wigner (small) d-matrices corresponding to a 90 degree rotation around x. 

135 If ``None`` (default) pre-calculated matrix elements are used. 

136 output_array 

137 numpy array with same shape as :attr:`input_array`. 

138 If ``None`` (default) a new array is initialized. 

139 If :attr:`output_array` is given, the calculations are carried out in place. 

140 

141 Returns 

142 --------- 

143 numpy array with same shape as :attr:`input_array`. 

144 """ 

145 if output_array is None: 

146 output_array = np.zeros(input_array.shape, dtype=float) 

147 

148 start_index = 0 

149 

150 for ell in ell_list: 

151 

152 if ell == 0: 

153 output_array[..., 0] = input_array[..., 0] 

154 else: 

155 m = np.arange(-ell, ell+1) # using negative ms is a comp. trick 

156 ma = np.einsum('m,...->...m', m, angle) 

157 output_array[..., start_index:start_index+(2*ell+1)]\ 

158 = input_array[..., start_index:start_index+(2*ell+1)] * np.cos(ma)\ 

159 - input_array[..., start_index+(2*ell+1)-1:start_index-1:-1] * np.sin(ma) 

160 

161 start_index += 2 * ell + 1 

162 

163 return output_array 

164 

165 

166def calculate_sph_coefficients_rotated_around_z_derived_wrt_the_angle( 

167 input_array: NDArray[float], 

168 angle: List[float], 

169 ell_list: List[int], 

170 output_array: NDArray[float] = None, 

171 ) -> NDArray[float]: 

172 r""" Calculate the angular derivatives of spherical harmonic coefficients with respect 

173 at the specified angle, where the angle refers to the rotation around +z. 

174 

175 Parameters 

176 ---------- 

177 input_array 

178 spherical harmonic coefficients 

179 angle 

180 either an array with the same shape as :attr:`input_array.shape[:-1]` or a scalar. 

181 If a scalar is given all coefficient lists are rotated by the same angle. 

182 ell_list 

183 list of :math:`\ell` values to use. If ``None`` (default) we assume that all even orders are inluded 

184 until some :math:`\ell_\mathrm{max}`, which will be calculated from the shape of the coefficients 

185 array. 

186 d_matrices 

187 list of Wigner (small) d-matrices corresponding to a 90 degree rotation about x. 

188 If ``None`` (default) pre-calculated matrix elements are used. 

189 output_array 

190 numpy array with same shape as :attr:`input_array`. 

191 If ``None`` (default) a new array is initialized. 

192 If :attr:`output_array` is given, the calculations are carried out in place. 

193 

194 Returns 

195 --------- 

196 numpy array with same shape as :attr:`input_array`. 

197 """ 

198 if output_array is None: 

199 output_array = np.zeros(input_array.shape, dtype=float) 

200 

201 start_index = 0 

202 

203 for ell in ell_list: 

204 

205 if ell == 0: 

206 output_array[..., 0] = 0 

207 else: 

208 m = np.arange(-ell, ell+1) # using negative ms is a comp. trick 

209 ma = np.einsum('m,...->...m', m, angle) 

210 m = np.einsum('m,...->...m', m, np.ones(angle.shape)) # Hacky indexing trick 

211 

212 output_array[..., start_index:start_index+(2*ell+1)]\ 

213 = input_array[..., start_index:start_index+(2*ell+1)] * (-m) * np.sin(ma)\ 

214 - input_array[..., start_index+(2*ell+1)-1:start_index-1:-1] * m * np.cos(ma) 

215 

216 start_index += 2 * ell + 1 

217 

218 return output_array 

219 

220 

221def calculate_sph_coefficients_rotated_by_euler_angles( 

222 input_array: NDArray[float], 

223 Psi: NDArray[float], 

224 Theta: NDArray[float], 

225 Phi: NDArray[float], 

226 ell_list: List[int] = None, 

227 d_matrices: List[NDArray] = None, 

228 output_array: NDArray[float] = None, 

229 ) -> NDArray[float]: 

230 r"""Calculate the spherical harmonics coefficients after a rotation specified by a set of Euler angles. 

231 The Euler angles need to be given in 'zyz' format. 

232 A rotation with ``(0, Theta, Phi)`` will move the z-axis into the coordinates ``(Theta, Phi)``. 

233 

234 Parameters 

235 ---------- 

236 input_array 

237 array, the last dimension of which runs over the spherical harmonics coefficients. 

238 Psi 

239 First Euler angle. Initial rotation about the z-axis. Can be either a scalar 

240 or a numpy array with shape :attr:`input_array.shape[:-1]`. 

241 If ``None`` the rotation is skipped. 

242 Theta 

243 Second Euler angle. Rotation about the y-axis. Can be either a scalar 

244 or a numpy array with shape ``input_array.shape[:-1]``. 

245 If ``None`` the rotation is skipped. 

246 Phi 

247 Third Euler angle. Final rotation about the z-axis. Can be either a scalar 

248 or a numpy array with shape ``input_array.shape[:-1]``. 

249 If ``None`` the rotation is skipped. 

250 ell_list 

251 list of :math:`\ell` values to use. If ``None`` (default) we assume that all even orders are inluded 

252 until some :math:`\ell_\mathrm{max}`, which will be calculated from the shape of the coefficients 

253 array. 

254 d_matrices 

255 list of Wigner (small) d-matrices corresponding to a 90 degree rotation around x. 

256 If ``None`` (default) pre-calculated matrix elements are used. 

257 output_array 

258 numpy array with same shape as :attr:`input_array`. 

259 If ``None`` (default) a new array is initialized. 

260 If :attr:`output_array` is given, the calculations are carried out in place. 

261 

262 Returns 

263 --------- 

264 numpy array with same shape as :attr:`input_array`. 

265 """ 

266 

267 if output_array is None: 267 ↛ 270line 267 didn't jump to line 270, because the condition on line 267 was never false

268 output_array = np.zeros(input_array.shape, dtype=float) 

269 

270 if isinstance(ell_list, int): 270 ↛ 271line 270 didn't jump to line 271, because the condition on line 270 was never true

271 ell_list = np.arange(0, ell_list+1, 2) 

272 elif ell_list is None: 272 ↛ 290line 272 didn't jump to line 290, because the condition on line 272 was never false

273 # figure out what ell_max is 

274 num_coeffs = input_array.shape[-1] 

275 ell = 0 

276 cum_sum = 1 

277 while cum_sum < num_coeffs: 

278 ell += 2 

279 cum_sum += 2 * ell + 1 

280 

281 if cum_sum == num_coeffs: 281 ↛ 284line 281 didn't jump to line 284, because the condition on line 281 was never false

282 ell_list = np.arange(0, ell+1, 2) 

283 else: 

284 raise ValueError('ell_max cannot be derived from input') 

285 

286 if d_matrices is None: 

287 d_matrices = load_d_matrices(np.max(ell_list)) 

288 

289 # Do z rotation of Psi 

290 if Psi is not None: 

291 output_array = calculate_sph_coefficients_rotated_around_z(input_array, 

292 Psi, 

293 ell_list, 

294 output_array=output_array) 

295 else: 

296 output_array = np.array(input_array, dtype=float) 

297 # Do y rotation about Theta 

298 if Theta is not None: 298 ↛ 315line 298 didn't jump to line 315, because the condition on line 298 was never false

299 # Do 90 degree rotation around x 

300 calculate_sph_coefficients_rotated_by_90_degrees_around_positive_x(output_array, 

301 ell_list, 

302 d_matrices, 

303 output_array=output_array) 

304 # Do z rotation of Theta 

305 calculate_sph_coefficients_rotated_around_z(output_array, 

306 Theta, 

307 ell_list, 

308 output_array=output_array) 

309 # Do -90 degree rotation around x 

310 calculate_sph_coefficients_rotated_by_90_degrees_around_negative_x(output_array, 

311 ell_list, 

312 d_matrices, 

313 output_array=output_array) 

314 # Do z rotation of Phi 

315 if Phi is not None: 

316 # Do z rotation of Phi 

317 calculate_sph_coefficients_rotated_around_z(output_array, 

318 Phi, 

319 ell_list, 

320 output_array=output_array) 

321 

322 return output_array