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