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
« 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
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.
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`).
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]
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
56 # Load d matrices
57 d_matrices = load_d_matrices(ell_max)
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))
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, :]
84 return optimal_zonal_coeffs, optimal_theta, optimal_phi
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.
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`).
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]
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
131 # Load d matrices
132 d_matrices = load_d_matrices(ell_max)
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))
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
152 theta, phi = np.meshgrid(theta_points, phi_points, indexing='ij')
153 dos = dos / np.sum(truncated_coefficients**2)
155 return dos, theta, phi
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.
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.
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]
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
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