Coverage for local_installation_linux/mumott/methods/utilities/fiber_fit.py: 93%

73 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-08-11 23:08 +0000

1import sys 

2from tqdm import tqdm 

3from typing import Tuple 

4import numpy as np 

5from numpy.typing import NDArray 

6from mumott.core.wigner_d_utilities import load_d_matrices, calculate_sph_coefficients_rotated_by_euler_angles 

7 

8 

9def find_approximate_symmetry_axis(coefficients: NDArray[float], 

10 ell_max: int, 

11 resolution: int = 10, 

12 filter: str = None) -> Tuple[NDArray[float]]: 

13 r""" Find the axis of highest apparent symmetry voxel-by-voxel for a voxel map of sperical harmonics 

14 coefficients. As a default, the measure of degree of symmetry is a power of the function rotationally 

15 averaged around the given axis. 

16 

17 Parameters 

18 ---------- 

19 coefficients 

20 Voxel map of spherical harmonic coefficients. 

21 ell_max 

22 Largest order :math:`\ell` used in the fitting. If :attr:`coefficients` contains higher orders, 

23 it will be truncated. 

24 resolution 

25 Number or angular steps along a half-circle, used in the search for the optimal axis. 

26 filter 

27 Weighing of different orders used to calculate the degree of symmetry. 

28 Possible values are "ramp" and "square". By default no filter is applied (`None`). 

29 

30 Returns 

31 --------- 

32 optimal_zonal_coeffs 

33 Zonal harmonics coefficients in the frame-of-reference corresponding 

34 to the axis of highest degree of symmetry. 

35 optimal_theta 

36 Voxel-by-voxel polar angles of the axis with highest degree of symmetry. 

37 optimal_phi 

38 Voxel-by-voxel azimuthal angles of the axis with highest degree of symmetry. 

39 """ 

40 # Use only the coefficients up until ell_max 

41 truncated_coefficients = coefficients[..., :(ell_max+1)*(ell_max+2)//2] 

42 volume_shape = coefficients.shape[:-1] 

43 

44 # Find the zonal coefficients 

45 zonal_indexes = np.zeros((ell_max+1)*(ell_max+2)//2, dtype=bool) 

46 inc = 0 

47 f = np.ones(ell_max//2+1) 

48 for ell in range(0, ell_max+1, 2): 

49 zonal_indexes[inc + ell] = True 

50 if filter == 'ramp': 50 ↛ 51line 50 didn't jump to line 51, because the condition on line 50 was never true

51 f[ell//2] = ell 

52 elif filter == 'square': 52 ↛ 53line 52 didn't jump to line 53, because the condition on line 52 was never true

53 f[ell//2] = ell**2 

54 inc += 2*ell + 1 

55 

56 # Load d matrices 

57 d_matrices = load_d_matrices(ell_max) 

58 

59 optimal_theta = np.zeros(volume_shape) 

60 optimal_phi = np.zeros(volume_shape) 

61 maximum_degree_of_symmetry = np.zeros(volume_shape) 

62 optimal_zonal_coeffs = np.zeros((*volume_shape, ell_max//2 + 1)) 

63 

64 # Loop through grid of directions 

65 theta_points = np.linspace(0, np.pi/2, resolution//2, endpoint=True) 

66 phi_points = np.linspace(0, 2*np.pi, 2*resolution, endpoint=False) 

67 for theta in tqdm(theta_points, file=sys.stdout): 

68 for phi in phi_points: 

69 rotated_coefficients = calculate_sph_coefficients_rotated_by_euler_angles( 

70 truncated_coefficients, 

71 Psi=-phi, 

72 Theta=-theta, 

73 Phi=None, 

74 d_matrices=d_matrices) 

75 # Check if degree of symmetry (DOS) is higher than maximum so far 

76 zonal_coeffs = rotated_coefficients[..., zonal_indexes] 

77 degree_of_symmetry = np.sum(zonal_coeffs**2*f[np.newaxis, np.newaxis, np.newaxis, :], axis=-1) 

78 indices = degree_of_symmetry > maximum_degree_of_symmetry 

79 maximum_degree_of_symmetry[indices] = degree_of_symmetry[indices] 

80 optimal_theta[indices] = theta 

81 optimal_phi[indices] = phi 

82 optimal_zonal_coeffs[indices, :] = zonal_coeffs[indices, :] 

83 

84 return optimal_zonal_coeffs, optimal_theta, optimal_phi 

85 

86 

87def degree_of_symmetry_map(coefficients: NDArray[float], 

88 ell_max: int, 

89 resolution: int = 10, 

90 filter: str = None) -> Tuple[NDArray[float]]: 

91 r""" Make a longitude-latitude map of the degree of symmetry. Can be used to make illustrating 

92 plots and to decide which filter is more appropriate. 

93 

94 Parameters 

95 ---------- 

96 coefficients 

97 Spherical harmonic coefficients of a single voxel. 

98 ell_max 

99 Largest order :math:`\ell` used in the fitting. If :attr:`coefficients` contains higher 

100 orders, it will be truncated. 

101 resolution 

102 Number or angular steps along a half-circle used in the search for the optimal axis. 

103 filter 

104 Weighing of different orders used to calculate the degree of symmetry. 

105 Possible values are "ramp" and "square". By default no filter is applied (`None`). 

106 

107 Returns 

108 --------- 

109 dos 

110 Map of the calculated degree of symmetry. 

111 theta 

112 Polar angle coordinates of the map. 

113 phi 

114 Azimuthal angle coordinates of the map. 

115 """ 

116 # Use only the coefficients up until ell_max 

117 truncated_coefficients = coefficients[..., :(ell_max+1)*(ell_max+2)//2] 

118 

119 # Find the zonal coefficients 

120 zonal_indexes = np.zeros((ell_max+1)*(ell_max+2)//2, dtype=bool) 

121 inc = 0 

122 f = np.ones(ell_max//2+1) 

123 for ell in range(0, ell_max+1, 2): 

124 zonal_indexes[inc + ell] = True 

125 if filter == 'ramp': 125 ↛ 126line 125 didn't jump to line 126, because the condition on line 125 was never true

126 f[ell//2] = ell 

127 elif filter == 'square': 127 ↛ 129line 127 didn't jump to line 129, because the condition on line 127 was never false

128 f[ell//2] = ell**2 

129 inc += 2*ell + 1 

130 

131 # Load d matrices 

132 d_matrices = load_d_matrices(ell_max) 

133 

134 # Loop through grid of directions 

135 theta_points = np.linspace(0, np.pi, resolution, endpoint=True) 

136 phi_points = np.linspace(0, 2*np.pi, 2*resolution, endpoint=True) 

137 dos = np.zeros((resolution, 2*resolution)) 

138 

139 for ii, theta in enumerate(theta_points): 

140 for jj, phi in enumerate(phi_points): 

141 rotated_coefficients = calculate_sph_coefficients_rotated_by_euler_angles( 

142 truncated_coefficients, 

143 Psi=-phi, 

144 Theta=-theta, 

145 Phi=None, 

146 d_matrices=d_matrices) 

147 # Check if degree of symmetry (DOS) is higher than maximum so far 

148 zonal_coeffs = rotated_coefficients[..., zonal_indexes] 

149 degree_of_symmetry = np.sum(zonal_coeffs**2*f, axis=-1) 

150 dos[ii, jj] = degree_of_symmetry 

151 

152 theta, phi = np.meshgrid(theta_points, phi_points, indexing='ij') 

153 dos = dos / np.sum(truncated_coefficients**2) 

154 

155 return dos, theta, phi 

156 

157 

158def symmetric_part_along_given_direction(coefficients: NDArray[float], 

159 theta: NDArray[float], 

160 phi: NDArray[float], 

161 ell_max: int) -> NDArray[float]: 

162 r""" Find the zonal harmonic coefficients along the given input directions. This can be 

163 used if eigenvector analysis is used to generate the symmetry directions used for a 

164 further zonal-harmonics refinement step. 

165 

166 Parameters 

167 ---------- 

168 coefficients 

169 Voxel map of spherical harmonic coefficients. 

170 theta 

171 Voxel map of polar angles. 

172 phi 

173 Voxel map of azimuthal angles. 

174 ell_max 

175 Largest order :math:`\ell` used in the fitting. If :attr:`coefficients` contains higher orders, 

176 it will be truncated. 

177 

178 Returns 

179 --------- 

180 Zonal harmonics coefficients in the frame of reference corresponding 

181 to the axis of highest degree of symmetry. 

182 """ 

183 # Use only the coefficients up until ell_max 

184 truncated_coefficients = coefficients[..., :(ell_max+1)*(ell_max+2)//2] 

185 

186 # Find the zonal coefficients 

187 zonal_indexes = np.zeros((ell_max+1)*(ell_max+2)//2, dtype=bool) 

188 inc = 0 

189 for ell in range(0, ell_max+1, 2): 

190 zonal_indexes[inc + ell] = True 

191 inc += 2*ell + 1 

192 

193 # Load d matrices 

194 d_matrices = load_d_matrices(ell_max) 

195 rotated_coefficients = calculate_sph_coefficients_rotated_by_euler_angles( 

196 truncated_coefficients, 

197 Psi=-phi, 

198 Theta=-theta, 

199 Phi=None, 

200 d_matrices=d_matrices) 

201 # Pick out symmetric part 

202 zonal_coeffs = rotated_coefficients[..., zonal_indexes] 

203 return zonal_coeffs