Coverage for local_installation_linux/mumott/core/spherical_harmonic_mapper.py: 96%
124 statements
« prev ^ index » next coverage.py v7.3.2, created at 2025-05-05 21:21 +0000
1""" Container for class SphericalHarmonicMapper. """
2import numpy as np
3from numpy.typing import NDArray
4from scipy.spatial.transform import Rotation as rot
5from scipy.special import sph_harm, factorial2
6from typing import Tuple
9class SphericalHarmonicMapper:
10 """
11 Helper class for visualizing and analyzing spherical harmonics.
12 Using this class, one can obtain the amplitudes for a given set
13 of even-ordered harmonics and apply the Funk-Radon transform.
14 These can then be plotted, analyzed, and so forth. In addition, the class
15 allows one to represent functions in terms of spherical harmonics,
16 if they are given as functions of the azimuthal and polar angles
17 of the class instance.
19 Parameters
20 ----------
21 ell_max : int, optional
22 Maximum order of the spherical harmonics. Default is ``2``.
23 polar_resolution : int, optional
24 Number of samples in the polar direction. Default is ``16``.
25 azimuthal_resolution : int, optional
26 Number of samples in the azimuthal direction. Default is ``32``.
27 polar_zero : float, optional
28 The polar angle of the spherical harmonic coordinate system's
29 pole, relative to the reference coordinate system. Default is ``0``.
30 azimuthal_zero : float, optional
31 The azimuthal angle of the spherical harmonic coordinate system's
32 pole, relative to a reference coordinate system. Default is ``0``.
33 enforce_friedel_symmetry : bool
34 If set to ``True``, Friedel symmetry will be enforced, using the assumption that points
35 on opposite sides of the sphere are equivalent.
37 Example
38 -------
39 >>> import numpy as np
40 >>> import matplotlib.pyplot as plt
41 >>> S = SphericalHarmonicMapper(ell_max=4, polar_resolution=16, azimuthal_resolution=8)
42 >>> S.ell_indices
43 array([0, 2, 2, 2, 2, 2, 4, 4, ...])
44 >>> S.emm_indices
45 array([ 0, -2, -1, 0, 1, 2, -4, -3, ...])
46 >>> a = S.get_amplitudes(np.random.rand(S.ell_indices.size) - 0.5)
47 >>> plt.pcolormesh(S.phi * np.sqrt(1 / 2), # basic cylindrical equal-area projection
48 np.cos(S.theta) / np.sqrt(1 / 2),
49 a[:-1, :-1])
50 ...
51 """
53 def __init__(self,
54 ell_max: int = 2,
55 polar_resolution: int = 16,
56 azimuthal_resolution: int = 32,
57 polar_zero: float = 0,
58 azimuthal_zero: float = 0,
59 enforce_friedel_symmetry: bool = True):
60 self._polar_zero = polar_zero
61 self._azimuthal_zero = azimuthal_zero
62 self._ell_max = ell_max
63 polar_coordinates = np.linspace(0, np.pi, polar_resolution)
64 azimuthal_coordinates = np.linspace(-np.pi, np.pi, azimuthal_resolution + 1)[:-1]
65 self._theta, self._phi = np.meshgrid(
66 polar_coordinates,
67 azimuthal_coordinates,
68 indexing='ij')
69 self._enforce_friedel_symmetry = enforce_friedel_symmetry
70 self._update_l_and_m()
72 self._map = np.zeros(self._theta.shape + self._ell_indices.shape)
73 self._complex_map = np.zeros(self._theta.shape + self._ell_indices.shape).astype(complex)
74 self._update_map()
76 self._rotated_system = True
77 self._get_coordinates()
78 self._update_funk_coefficients()
80 self._xyz = np.concatenate((self._X[..., None], self._Y[..., None], self._Z[..., None]), axis=2)
81 self.update_zeros(self._polar_zero, self._azimuthal_zero)
83 def get_amplitudes(self,
84 coefficients: NDArray[float],
85 apply_funk_transform: bool = False) -> NDArray[float]:
86 """
87 Returns the amplitudes of a set of spherical harmonics. For sorting
88 of the coefficients, see the :attr:`orders` and :attr:`degrees` attributes.
90 Parameters
91 ----------
92 coefficients
93 The coefficients for which the amplitudes are to be calculated.
94 Size must be equal to ``(ell_max + 1) * (ell_max / 2 + 1)``.
95 apply_funk_fransform
96 Whether to apply the Funk-Radon transform to the coefficients.
97 This is useful for orientation analysis in some cases.
98 Default is ``False``.
99 """
100 if apply_funk_transform:
101 coefficients = (self._funk_coefficients * coefficients.ravel()).reshape(1, 1, -1)
102 else:
103 coefficients = coefficients.reshape(1, 1, -1)
104 return np.einsum('...i, ...i', coefficients, self._map)
106 def get_harmonic_coefficients(self, amplitude: NDArray[float]) -> NDArray[float]:
107 """ Returns the spherical harmonic coefficients for the given amplitudes.
108 Can be used with amplitudes from another instance of
109 :class:`SphericalHarmonicParameters` with a different orientation
110 in order to solve for the rotated coefficients. The accuracy of the
111 representation depends on the maximum order and the polar and azimuthal
112 resolution.
114 Parameters
115 ----------
116 amplitude
117 The amplitude of the spherical function to be
118 represented, as a function of ``theta`` and ``phi``.
119 """
120 assert np.allclose(amplitude.shape[-1:-3:-1], self.theta.shape)
121 area_normer = np.sin(self.theta)
122 area_normer /= np.sum(area_normer)
123 scaled_amp = np.einsum('...ij, ij -> ...ij',
124 amplitude,
125 area_normer, optimize='greedy')
126 coeffs = np.einsum('ijk, ...ij -> ...k', self.map, scaled_amp, optimize='greedy')
127 return coeffs
129 def _get_coordinates(self) -> Tuple[NDArray[float], NDArray[float], NDArray[float]]:
130 """ Gets the X, Y, and Z-coordinates. Updates them only if the system has
131 been rotated since the last call. """
132 if self._rotated_system:
133 self._X, self._Y, self._Z = \
134 (np.multiply(0.5, np.einsum('..., ...', np.sin(self._theta), np.cos(self._phi))),
135 np.multiply(0.5, np.einsum('..., ...', np.sin(self._theta), np.sin(self._phi))),
136 np.multiply(0.5, np.cos(self._theta)))
137 self._rotated_system = False
138 return self._X, self._Y, self._Z
140 def _update_funk_coefficients(self) -> None:
141 """ Updates the Funk coefficients used for the Funk transform. """
142 funk_coefficients = []
143 for i in range(self._ell_max+1):
144 if i % 2 == 0:
145 funk_coefficients.append(((-1) ** (i // 2)) * factorial2(i - 1) / factorial2(i))
146 funk_coefficients = np.array(funk_coefficients)
147 self._funk_coefficients = funk_coefficients[self._ell_indices // 2]
149 def _update_l_and_m(self) -> None:
150 """ Updates the order and degree vectors of the system. """
151 if self._enforce_friedel_symmetry:
152 ell_step = 2
153 else:
154 ell_step = 1
155 self._ell_indices = []
156 self._emm_indices = []
157 for ll in range(0, self._ell_max+1, ell_step):
158 for mm in range(-ll, ll+1):
159 self._ell_indices.append(ll)
160 self._emm_indices.append(mm)
161 self._ell_indices = np.array(self._ell_indices)
162 self._emm_indices = np.array(self._emm_indices)
164 def _update_map(self) -> None:
165 """ Updates the map of the system. """
166 self._update_complex_map()
167 self._update_real_map()
169 def _update_complex_map(self) -> None:
170 """ Retrieves the complex map used for determining
171 the real map. """
172 self._complex_map[...] = sph_harm(
173 abs(self._emm_indices.reshape(1, 1, -1)),
174 self._ell_indices.reshape(1, 1, -1),
175 self._phi[:, :, None],
176 self._theta[:, :, None])
178 def _update_real_map(self) -> None:
179 """ Calculates the real spherical harmonic map based on
180 the complex map. """
181 self._map[:, :, self._emm_indices == 0] = \
182 np.sqrt(4 * np.pi) * self._complex_map[:, :, self._emm_indices == 0].real
183 self._map[:, :, self._emm_indices > 0] = \
184 ((-1.) ** (self._emm_indices[self._emm_indices > 0])) * np.sqrt(2) * \
185 np.sqrt(4 * np.pi) * self._complex_map[:, :, self._emm_indices > 0].real
186 self._map[:, :, self._emm_indices < 0] = \
187 ((-1.) ** (self._emm_indices[self._emm_indices < 0])) * np.sqrt(2) * \
188 np.sqrt(4 * np.pi) * self._complex_map[:, :, self._emm_indices < 0].imag
190 def update_zeros(self, polar_zero: float, azimuthal_zero: float) -> None:
191 """Changes the orientation of the coordinate system.
193 Parameters
194 ----------
195 polar_zero
196 The new polar angle at which the pole should be,
197 relative to a fixed reference system.
198 azimuthal_zero
199 The new azimuthal angle at which the pole should be,
200 relative to a fixed reference system.
201 """
202 if polar_zero == 0 and azimuthal_zero == 0:
203 self._theta = np.arccos(self._xyz[..., 2] / np.sqrt((self._xyz ** 2).sum(-1)))
204 self._phi = np.arctan2(self._xyz[..., 1], self._xyz[..., 0])
205 return
206 xyz = np.copy(self._xyz)
207 xyz = xyz.reshape(-1, 3)
208 rotvec = rot.from_euler('Z', (-azimuthal_zero))
209 rotvec = rot.from_euler('Y', (-polar_zero)) * rotvec
210 xyz = rotvec.apply(xyz).reshape(self._theta.shape + (3,))
211 self._theta = np.arccos(xyz[..., 2] / np.sqrt((xyz ** 2).sum(-1)))
212 self._phi = np.arctan2(xyz[..., 1], xyz[..., 0])
213 self._polar_zero = polar_zero
214 self._azimuthal_zero = azimuthal_zero
215 self._rotated_system = True
216 self._update_map()
218 @property
219 def unit_vectors(self):
220 """ Probed coordinates object used by BasisSets to calculate function maps.
221 """
222 X, Y, Z = self.coordinates
223 vectors = 2 * np.stack((X, Y, Z), axis=-1) # The mapper uses a sphere of radius 0.5
224 return vectors
226 @property
227 def ell_indices(self) -> NDArray[int]:
228 """ The orders of the harmonics calculated
229 by the class instance. """
230 return self._ell_indices
232 @property
233 def emm_indices(self) -> NDArray[int]:
234 """ The degrees of the harmonics calculated
235 by the class instance. """
236 return self._emm_indices
238 @property
239 def polar_zero(self) -> NDArray[float]:
240 """ The polar angle of the spherical harmonic pole,
241 relative to a fixed reference system. """
242 return self._polar_zero
244 @property
245 def azimuthal_zero(self) -> NDArray[float]:
246 """ The azimuthal angle of the spherical harmonic pole,
247 relative to a fixed reference system. """
248 return self._azimuthal_zero
250 @property
251 def theta(self) -> NDArray[float]:
252 """ The polar angle to which the amplitude is mapped. """
253 return self._theta
255 @property
256 def phi(self) -> NDArray[float]:
257 """ The azimuthal angle to which the amplitude is mapped. """
258 return self._phi
260 @property
261 def ell_max(self) -> int:
262 """ Maximum order of the spherical harmonics. """
263 return self._ell_max
265 @property
266 def map(self) -> NDArray[float]:
267 """ Map between amplitude and harmonics. """
268 return self._map
270 @property
271 def coordinates(self) -> Tuple[NDArray[float], NDArray[float], NDArray[float]]:
272 """ The X, Y, Z coordinates that the amplitudes
273 are mapped to. """
274 return self._get_coordinates()
276 @ell_max.setter
277 def ell_max(self, new_ell_max: int):
278 self._ell_max = new_ell_max
279 self._update_l_and_m()
280 self._map = np.zeros(self._theta.shape + self._ell_indices.shape)
281 self._complex_map = np.zeros(self._theta.shape + self._ell_indices.shape).astype(complex)
282 self._update_map()
283 self._update_funk_coefficients()