Coverage for local_installation_linux/mumott/methods/basis_sets/spherical_harmonics.py: 94%
236 statements
« prev ^ index » next coverage.py v7.3.2, created at 2024-08-11 23:08 +0000
1import logging
2from typing import Any, Dict, Iterator, List, Tuple
4import numpy as np
5from numpy.typing import NDArray
6from scipy.special import sph_harm
8from mumott import ProbedCoordinates
9from mumott.core.hashing import list_to_hash
10from mumott.methods.utilities.tensor_operations import (framewise_contraction,
11 framewise_contraction_transpose)
12from .base_basis_set import BasisSet
15logger = logging.getLogger(__name__)
18class SphericalHarmonics(BasisSet):
19 """ Basis set class for spherical harmonics, the canonical representation
20 of polynomials on the unit sphere and a simple way of representing
21 band-limited spherical functions which allows for easy computations of statistics
22 and is suitable for analyzing certain symmetries.
24 Parameters
25 ----------
26 ell_max : int
27 The bandlimit of the spherical functions that you want to be able to represent.
28 A good rule of thumb is that :attr:`ell_max` should not exceed the
29 number of detector segments minus 1.
30 probed_coordinates : ProbedCoordinates
31 Optional. A container with the coordinates on the sphere probed at each detector segment by the
32 experimental method. Its construction from the system geometry is method-dependent.
33 By default, an empty instance of
34 :class:`ProbedCoordinates <mumott.core.probed_coordinates.ProbedCoordinates>` is created.
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. This results in only even ``ell`` being used.
38 kwargs
39 Miscellaneous arguments which relate to segment integrations can be
40 passed as keyword arguments:
42 integration_mode
43 Mode to integrate line segments on the reciprocal space sphere. Possible options are
44 ``'simpson'``, ``'midpoint'``, ``'romberg'``, ``'trapezoid'``.
45 ``'simpson'``, ``'trapezoid'``, and ``'romberg'`` use adaptive
46 integration with the respective quadrature rule from ``scipy.integrate``.
47 ``'midpoint'`` uses a single mid-point approximation of the integral.
48 Default value is ``'simpson'``.
49 n_integration_starting_points
50 Number of points used in the first iteration of the adaptive integration.
51 The number increases by the rule ``N`` ← ``2 * N - 1`` for each iteration.
52 Default value is 3.
53 integration_tolerance
54 Tolerance for the maximum relative error between iterations before the integral
55 is considered converged. Default is ``1e-5``.
56 integration_maxiter
57 Maximum number of iterations. Default is ``10``.
58 """
59 def __init__(self,
60 probed_coordinates: ProbedCoordinates = None,
61 ell_max: int = 0,
62 enforce_friedel_symmetry: bool = True,
63 **kwargs):
64 super().__init__(probed_coordinates, **kwargs)
65 self._probed_coordinates_hash = hash(self.probed_coordinates)
66 self._ell_max = ell_max
67 self._ell_indices = np.zeros(1)
68 self._emm_indices = np.zeros(1)
69 # Compute initial values for indices and matrix.
70 self._enforce_friedel_symmetry = enforce_friedel_symmetry
71 self._calculate_coefficient_indices()
72 self._projection_matrix = self._get_integrated_projection_matrix()
74 def _calculate_coefficient_indices(self) -> None:
75 """
76 Computes the attributes :attr:`~.SphericalHarmonics.ell_indices` and
77 :attr:`~.SphericalHarmonics.emm_indices`. Called when :attr:`~.SphericalHarmonics.ell_max`
78 changes.
79 """
80 if self._enforce_friedel_symmetry:
81 divisor = 2
82 else:
83 divisor = 1
85 mm = np.zeros((self._ell_max + 1) * (self._ell_max // divisor + 1), dtype=int)
86 ll = np.zeros((self._ell_max + 1) * (self._ell_max // divisor + 1), dtype=int)
87 count = 0
88 for h in range(0, self._ell_max + 1, divisor):
89 for i in range(-h, h + 1):
90 ll[count] = h
91 mm[count] = i
92 count += 1
93 self._ell_indices = ll
94 self._emm_indices = mm
96 def _get_projection_matrix(self, probed_coordinates: ProbedCoordinates = None) -> None:
97 """ Computes the matrix necessary for forward and gradient calculations.
98 Called when the coordinate system has been updated or ``ell_max`` has changed."""
99 if probed_coordinates is None: 99 ↛ 100line 99 didn't jump to line 100, because the condition on line 99 was never true
100 probed_coordinates = self._probed_coordinates
101 _, probed_polar_angles, probed_azim_angles = probed_coordinates.to_spherical
102 # retrieve complex spherical harmonics with emm >= 0, shape (N, M, len(ell_indices), K)
103 complex_factors = sph_harm(abs(self._emm_indices)[np.newaxis, np.newaxis, np.newaxis, ...],
104 self._ell_indices[np.newaxis, np.newaxis, np.newaxis, ...],
105 probed_azim_angles[..., np.newaxis],
106 probed_polar_angles[..., np.newaxis])
107 # cancel Condon-Shortley phase factor in scipy.special.sph_harm
108 condon_shortley_factor = (-1.) ** self._emm_indices
109 # 4pi normalization factor and complex-to-real normalization factor for m != 0
110 norm_factor = np.sqrt(4 * np.pi) * \
111 np.sqrt(1 + (self._emm_indices != 0).astype(int)) * condon_shortley_factor
112 matrix = norm_factor[np.newaxis, np.newaxis, np.newaxis, ...] * (
113 (self._emm_indices >= 0)[np.newaxis, np.newaxis, np.newaxis, ...] * complex_factors.real +
114 (self._emm_indices < 0)[np.newaxis, np.newaxis, np.newaxis, ...] * complex_factors.imag)
115 return matrix
117 def forward(self,
118 coefficients: NDArray,
119 indices: NDArray = None) -> NDArray:
120 """ Carries out a forward computation of projections from spherical harmonic space
121 into detector space, for one or several tomographic projections.
123 Parameters
124 ----------
125 coefficients
126 An array of coefficients, of arbitrary shape so long as the last
127 axis has the same size as :attr:`~.SphericalHarmonics.ell_indices`, and if
128 :attr:`indices` is `None` or greater than one, the first axis should have the
129 same length as :attr:`indices`
130 indices
131 Optional. Indices of the tomographic projections for which the forward
132 computation is to be performed. If ``None``, the forward computation will
133 be performed for all projections.
135 Returns
136 -------
137 An array of values on the detector corresponding to the :attr:`coefficients` given.
138 If :attr:`indices` contains exactly one index, the shape is ``(coefficients.shape[:-1], J)``
139 where ``J`` is the number of detector segments. If :attr:`indices` is ``None`` or contains
140 several indices, the shape is ``(N, coefficients.shape[1:-1], J)`` where ``N``
141 is the number of tomographic projections for which the computation is performed.
143 Notes
144 -----
145 The assumption is made in this implementation that computations over several
146 indices act on sets of images from different projections. For special usage
147 where multiple projections of entire fields is desired, it may be better
148 to use :attr:`projection_matrix` directly. This also applies to
149 :meth:`gradient`.
150 """
151 assert coefficients.shape[-1] == self._ell_indices.size
152 self._update()
153 output = np.zeros(coefficients.shape[:-1] + (self._projection_matrix.shape[1],),
154 coefficients.dtype)
155 if indices is None:
156 framewise_contraction_transpose(self._projection_matrix,
157 coefficients,
158 output)
159 elif indices.size == 1: 159 ↛ 167line 159 didn't jump to line 167, because the condition on line 159 was never false
160 np.einsum('ijk, ...k -> ...j',
161 self._projection_matrix[indices],
162 coefficients,
163 out=output,
164 optimize='greedy',
165 casting='unsafe')
166 else:
167 framewise_contraction_transpose(self._projection_matrix[indices],
168 coefficients,
169 output)
170 return output
172 def gradient(self,
173 coefficients: NDArray,
174 indices: NDArray = None) -> NDArray:
175 """ Carries out a gradient computation of projections from spherical harmonic space
176 into detector space, for one or several tomographic projections.
178 Parameters
179 ----------
180 coefficients
181 An array of coefficients (or residuals) of arbitrary shape so long as the last
182 axis has the same size as the number of detector segments.
183 indices
184 Optional. Indices of the tomographic projections for which the gradient
185 computation is to be performed. If ``None``, the gradient computation will
186 be performed for all projections.
188 Returns
189 -------
190 An array of gradient values based on the `coefficients` given.
191 If :attr:`indices` contains exactly one index, the shape is ``(coefficients.shape[:-1], J)``
192 where ``J`` is the number of detector segments. If indices is ``None`` or contains
193 several indices, the shape is ``(N, coefficients.shape[1:-1], J)`` where ``N``
194 is the number of tomographic projections for which the computation is performed.
196 Notes
197 -----
198 When solving an inverse problem, one should not to attempt to optimize the
199 coefficients directly using the ``gradient`` one obtains by applying this method to the data.
200 Instead, one must either take the gradient of the residual between the
201 :meth:`~.SphericalHarmonics.forward` computation of the coefficients and the data.
202 Alternatively one can apply both the forward and the gradient computation to the
203 coefficients to be optimized, and the gradient computation to the data, and treat
204 the residual of the two as the gradient of the optimization coefficients. The approaches
205 are algebraically equivalent, but one may be more efficient than the other in some
206 circumstances.
207 """
208 self._update()
209 output = np.zeros(coefficients.shape[:-1] + (self._projection_matrix.shape[2],),
210 coefficients.dtype)
211 if indices is None:
212 framewise_contraction(self._projection_matrix,
213 coefficients,
214 output)
215 elif indices.size == 1: 215 ↛ 223line 215 didn't jump to line 223, because the condition on line 215 was never false
216 np.einsum('ikj, ...k -> ...j',
217 self._projection_matrix[indices],
218 coefficients,
219 out=output,
220 optimize='greedy',
221 casting='unsafe')
222 else:
223 framewise_contraction(self._projection_matrix[indices],
224 coefficients,
225 output)
226 return output
228 def get_inner_product(self,
229 u: NDArray,
230 v: NDArray,
231 resolve_spectrum: bool = False,
232 spectral_moments: List[int] = None) -> NDArray:
233 r""" Retrieves the inner product of two coefficient arrays.
235 Notes
236 -----
237 The canonical inner product in a spherical harmonic representation
238 is :math:`\sum_\ell N(\ell) \sum_m u_m^\ell v_m^\ell`, where :math:`N(\ell)` is
239 a normalization constant (which is unity for the :math:`4\pi` normalization).
240 This inner product is a rotational invariant. The rotational invariance also holds
241 for any partial sums over :math:`\ell`.
242 One can define a function of :math:`\ell` that returns such
243 products, namely :math:`S(\ell, u, v) = N(\ell)\sum_m u_m^\ell v_m^\ell`,
244 called the spectral power function.
245 The sum :math:`\sum_{\ell = 1}S(\ell)` is equal to the covariance of the
246 band-limited spherical functions represented by :math:`u` and :math:`v`, and each
247 :math:`S(\ell, u, v)` is the contribution to the covariance of the band :math:`\ell`.
248 See also
249 `the SHTOOLS documentation <https://shtools.github.io/SHTOOLS/real-spherical-harmonics.html>`_
250 for an excellent overview of this.
252 Parameters
253 ----------
254 u
255 The first coefficient array, of arbitrary shape and dimension,
256 except the last dimension must be the same as the length of
257 :attr:`~.SphericalHarmonics.ell_indices`.
258 v
259 The second coefficient array, of the same shape as ``u``.
260 resolve_spectrum
261 Optional. Whether to resolve the product according to each frequency band, given by the
262 coefficients of each ``ell`` in :attr:`~SphericalHarmonics.ell_indices`.
263 Defaults to ``False``, which means that the sum of every component of the spectrum is returned.
264 If ``True``, components are returned in order of ascending ``ell``.
265 The ``ell`` included in the spectrum depends on :attr:`spectral_moments`.
266 spectral_moments
267 Optional. List of particular values of ``ell`` to calculate the inner product for.
268 Defaults to ``None``, which is identical to including all values of ``ell``
269 in the calculation.
270 If :attr:`spectral_moments` contains all nonzero values of ``ell`` and
271 :attr:`resolve_spectrum` is ``False``, the covariance of
272 :attr:`v` and :attr:`u` will be calculated (the sum of the inner product over all non-zero ``ell``
273 If ``resolve_spectrum`` is ``True``, the covariance per `ell` in ``spectral_moments``, will
274 be calculated, i.e., the inner products will not be summed over.
276 Returns
277 -------
278 An array of the inner products of the spherical functions represented by ``u`` and ``v``.
279 Has the shape ``(u.shape[:-1])`` if :attr:`resolve_spectrum` is ``False``,
280 ``(u.shape[:-1] + (ell_max // 2 + 1,))`` if :attr:`resolve_spectrum` is ``True`` and
281 ``spectral_moments`` is ``None``, and finally the shape
282 ``(u.shape[:-1] + (np.unique(spectral_moments).size,))`` if :attr:`resolve_spectrum` is ``True``
283 and :attr:`spectral_moments` is a list of integers found
284 in :attr:`~.SphericalHarmonics.ell_indices`
285 """
286 assert u.shape == v.shape
287 assert u.shape[-1] == self._ell_indices.size
288 if not resolve_spectrum:
289 if spectral_moments is None:
290 return np.einsum('...i, ...i -> ...', u, v)
291 # pick out only the subset where ell matches the provided spectral moments
292 where = np.any([np.equal(self._ell_indices, ell) for ell in spectral_moments], axis=0)
293 return np.einsum('...i, ...i -> ...', u[..., where], v[..., where])
294 if spectral_moments is None: 294 ↛ 297line 294 didn't jump to line 297, because the condition on line 294 was never false
295 which_ell = np.unique(self._ell_indices)
296 else:
297 which_ell = np.unique(spectral_moments)
298 which_ell = [ell for ell in which_ell if np.any(np.equals(self._ell_indices, ell))]
299 power_spectrum = np.zeros((*u.shape[:-1], which_ell.size))
300 # power spectrum for any one ell is given by inner product over each ell
301 for i, ell in enumerate(which_ell):
302 power_spectrum[..., i] = np.einsum('...i, ...i -> ...',
303 u[..., self._ell_indices == ell],
304 v[..., self._ell_indices == ell],
305 optimize='greedy')
306 return power_spectrum
308 def get_covariances(self,
309 u: NDArray,
310 v: NDArray,
311 resolve_spectrum: bool = False) -> NDArray:
312 """ Returns the covariances of the spherical functions represented by
313 two coefficient arrays.
315 Parameters
316 ----------
317 u
318 The first coefficient array, of arbitrary shape
319 except its last dimension must be the same length as
320 the length of :attr:`~SphericalHarmonics.ell_indices`.
321 v
322 The second coefficient array, of the same shape as :attr:`u`.
323 resolve_spectrum
324 Optional. Whether to resolve the product according to each frequency band, given by the
325 coefficients of each ``ell`` in :attr:`~.SphericalHarmonics.ell_indices`.
326 Default value is ``False``.
328 Returns
329 -------
330 An array of the covariances of the spherical functions represented by ``u`` and ``v``.
331 Has the shape ``(u.shape[:-1])`` if `resolve_spectrum` is ``False``, and
332 ``(u.shape[:-1] + (ell_max // 2 + 1,))`` if `resolve_spectrum` is ``True``, where
333 ``ell_max`` is :attr:`.SphericalHarmonics.ell_max`.
335 Notes
336 -----
337 Calling this function is equivalent to calling :func:`~.SphericalHarmonics.get_inner_product`
338 with ``spectral_moments=np.unique(ell_indices[ell_indices > 0])`` where ``ell_indices`` is
339 :attr:`.SphericalHarmonics.ell_indices`. See the note to
340 :func:`~.SphericalHarmonics.get_inner_product` for mathematical details.
341 """
342 spectral_moments = np.unique(self._ell_indices[self._ell_indices > 0])
343 return self.get_inner_products(u, v, resolve_spectrum, spectral_moments)
345 def get_output(self,
346 coefficients: NDArray) -> Dict[str, Any]:
347 r""" Returns a dictionary of output data for a given array of spherical harmonic coefficients.
349 Parameters
350 ----------
351 coefficients
352 An array of coefficients of arbitrary shape and dimensions, except
353 its last dimension must be the same length as :attr:`~.SphericalHarmonics.ell_indices`.
354 Computations only operate over the last axis of :attr:`coefficients`, so derived
355 properties in the output will have the shape ``(*coefficients.shape[:-1], ...)``.
357 Returns
358 -------
359 A dictionary containing two sub-dictionaries, ``basis_set`` and ``spherical_functions``.
360 ``basis_set`` contains information particular to :class:`SphericalHarmonics`, whereas
361 ``spherical_functions`` contains information about the spherical functions
362 represented by the :attr:`coefficients` which are not specific to the chosen representation.
364 Notes
365 -----
366 In detail, the two sub-dictionaries ``basis_set`` and ``spherical_functions`` have the following
367 members:
369 basis_set
370 name
371 The name of the basis set, i.e., ``'SphericalHarmonicParameters'``
372 coefficients
373 A copy of :attr:`coefficients`.
374 ell_max
375 A copy of :attr:`~.SphericalHarmonics.ell_max`.
376 ell_indices
377 A copy of :attr:`~.SphericalHarmonics.ell_indices`.
378 emm_indices
379 A copy of :attr:`~.SphericalHarmonics.emm_indices`.
380 projection_matrix
381 A copy of :attr:`~.SphericalHarmonics.projection_matrix`.
382 spherical_functions
383 means
384 The spherical means of each function represented by :attr:`coefficients`.
385 variances
386 The spherical variances of each function represented by :attr:`coefficients`.
387 If :attr:`~.ell_max` is ``0``, all variances will equal zero.
388 r2_tensors
389 The traceless symmetric rank-2 tensor component of each function represented by
390 :attr:`coefficients`, in 6-element form, in the order ``[xx, yy, zz, yz, xz, xy]``,
391 i.e., by the Voigt convention.
392 The matrix form can be recovered as r2_tensors[..., tensor_to_matrix_indices],
393 yielding matrix elements ``[[xx, xy, xz], [xy, yy, yz], [xz, yz, zz]]``.
394 If :attr:`~.ell_max` is ``0``, all tensors have elements
395 [1, 0, -1, 0, 0, 0].
396 tensor_to_matrix_indices
397 A list of indices to help recover the matrix from the 2-element form of the
398 rank-2 tensors, equalling precisely ``[[0, 5, 4], [5, 1, 3], [4, 3, 2]]``
399 eigenvalues
400 The eigenvalues of the rank-2 tensors, sorted in ascending order in the last index.
401 If :attr:`~.ell_max` is ``0``, the eigenvalues will always be (1, 0, -1)
402 eigenvectors
403 The eigenvectors of the rank-2 tensors, sorted with their corresponding eigenvectors
404 in the last index. Thus, ``eigenvectors[..., 0]`` gives the eigenvector corresponding
405 to the smallest eigenvalue, and ``eigenvectors[..., 2]`` gives the eigenvector
406 corresponding to the largest eigenvalue. Generally, one of these two eigenvectors
407 is used to define the orientation of a function, depending on whether it is
408 characterized by a minimum (``0``) or a maximum (``2``). The middle eigenvector (``1``)
409 is typically only used for visualizations.
410 If :attr:`~.ell_max` is ``0``, the eigenvectors will be the Cartesian basis
411 vectors.
412 main_orientations
413 The estimated main orientations from the largest absolute eigenvalues.
414 If :attr:`~.ell_max` is ``0``, the main orientation will be the x-axis.
415 main_orientation_symmetries
416 The strength or definiteness of the main orientation, calculated from
417 the quotient of the absolute middle and signed largest eigenvalues of the
418 rank-2 tensor.
419 If ``0``, the orientation is totally ambiguous. The orientation is completely
420 transversal if the value is ``-1`` (orientation represents a minimum),
421 and completely longitudinal if the value is ``1`` (orientation represents a maximum).
422 If :attr:`~.ell_max` is ``0``, the main orientations are all totally ambiguous.
423 normalized_standard_deviations
424 A relative measure of the overall anisotropy of the spherical functions. Equals
425 :math:`\sqrt{\sigma^2 / \mu}`, where :math:`\sigma^2` is the variance
426 and :math:`\mu` is the mean. The places where :math:`\mu=0` have been
427 set to ``0``. If :attr:`~.ell_max` is ``0``, the normalized standard
428 deviations will always be zero.
429 power_spectra
430 The spectral powers of each ``ell`` in :attr:`~.SphericalHarmonics.ell_indices`,
431 for each spherical function, sorted in ascending ``ell``.
432 If :attr:`~.ell_max` is ``0``, each function will have only one element, equal
433 to the mean squared.
434 power_spectra_ell
435 An array containing the corresponding ``ell`` to each of the last indices
436 in :attr:`power_spectra`. Equal to ``np.unique(ell_indices)``.
437 """
438 assert coefficients.shape[-1] == self._ell_indices.size
439 # Update to ensure non-dirty output state.
440 self._update()
441 output_dictionary = {}
443 # basis set-specific information
444 basis_set = {}
445 output_dictionary['basis_set'] = basis_set
446 basis_set['name'] = type(self).__name__
447 basis_set['coefficients'] = coefficients.copy()
448 basis_set['ell_max'] = self._ell_max
449 basis_set['ell_indices'] = self._ell_indices.copy()
450 basis_set['emm_indices'] = self._emm_indices.copy()
451 basis_set['projection_matrix'] = self._projection_matrix.copy()
452 basis_set['hash'] = hex(hash(self))
454 # general properties of spherical function
455 spherical_functions = {}
456 output_dictionary['spherical_functions'] = spherical_functions
457 spherical_functions['means'] = coefficients[..., 0]
458 if coefficients.shape[-1] > 1: 458 ↛ 461line 458 didn't jump to line 461, because the condition on line 458 was never false
459 spherical_functions['variances'] = (coefficients[..., 1:] ** 2).sum(-1)
460 else:
461 spherical_functions['variances'] = np.zeros_like(coefficients[..., 0])
462 r2_tensor, eigvect, eigval = self._get_rank2_tensor_analysis(coefficients)
463 spherical_functions['eigenvectors'] = eigvect
464 spherical_functions['r2_tensors'] = np.concatenate(
465 (np.diagonal(r2_tensor, offset=0, axis1=-2, axis2=-1),
466 r2_tensor[..., 1, [2]], r2_tensor[..., 0, [2]], r2_tensor[..., 0, [1]]), axis=-1)
467 spherical_functions['tensor_to_matrix_indices'] = [[0, 5, 4], [5, 1, 3], [4, 3, 2]]
468 spherical_functions['eigenvalues'] = eigval
470 # estimate main orientation from absolute eigenvalues and select from eigenvectors
471 spherical_functions['main_orientations'] = \
472 np.take_along_axis(eigvect,
473 np.argmax(abs(eigval), axis=-1).reshape(eigval.shape[:-1] + (1, 1)), axis=-1)
474 spherical_functions['main_orientations'] = spherical_functions['main_orientations'][..., 0]
475 # estimate main orientation strength as quotient between middle and largest absolute eigenvalue.
476 eigargs = np.argmax(abs(eigval), axis=-1).reshape(eigval.shape[:-1] + (1,))
477 spherical_functions['main_orientation_symmetries'] = 2 * abs(eigval[..., 1][..., None]) / \
478 np.take_along_axis(eigval, eigargs, axis=-1)
480 # estimate overall relative anisotropy as quotient betewen standard deviation and mean
481 valid_indices = spherical_functions['means'] > 0 # mask points where mean is zero or negative
482 normalized_std = np.zeros_like(spherical_functions['means'])
483 normalized_std[valid_indices] = np.sqrt(spherical_functions['variances'][valid_indices]) / \
484 spherical_functions['means'][valid_indices]
485 spherical_functions['normalized_standard_deviations'] = normalized_std
486 spherical_functions['power_spectra'] = self.get_inner_product(coefficients,
487 coefficients,
488 resolve_spectrum=True)
489 spherical_functions['power_spectra_ell'] = np.unique(self._ell_indices)
490 return output_dictionary
492 def get_spherical_harmonic_coefficients(
493 self,
494 coefficients: NDArray[float],
495 ell_max: int = None
496 ) -> NDArray[float]:
497 """ Convert a set of spherical harmonics coefficients to a different :attr:`ell_max`
498 by either zero-padding or truncation and return the result.
500 Parameters
501 ----------
502 coefficients
503 An array of coefficients of arbitrary shape, provided that the
504 last dimension contains the coefficients for one function.
505 ell_max
506 The band limit of the spherical harmonic expansion.
507 """
509 if coefficients.shape[-1] != len(self): 509 ↛ 510line 509 didn't jump to line 510, because the condition on line 509 was never true
510 raise ValueError(f'The number of coefficients ({coefficients.shape[-1]}) does not '
511 f'match the expected value. ({len(self)})')
513 if self._enforce_friedel_symmetry:
514 num_coeff_output = (ell_max+1) * (ell_max+2) // 2
515 elif not self._enforce_friedel_symmetry: 515 ↛ 518line 515 didn't jump to line 518, because the condition on line 515 was never false
516 num_coeff_output = (ell_max+1)**2
518 output_array = np.zeros((*coefficients.shape[:-1], num_coeff_output))
519 output_array[..., :min(len(self), num_coeff_output)] = \
520 coefficients[..., :min(len(self), num_coeff_output)]
521 return output_array
523 def _get_rank2_tensor_analysis(self, coefficients: NDArray) -> Tuple[NDArray, NDArray, NDArray]:
524 """ Performs an analysis of the rank-2 tensor components of the functions represented by the
525 given coefficients.
527 Parameters
528 ----------
529 coefficients
530 An array of coefficients of arbitrary shape, as long as the last axis has the same shape
531 as :attr:`~.SphericalHarmonics.ell_indices`.
533 Returns
534 -------
535 A tuple with three memberrs.
536 First, the traceless rank-2 tensor components of all the
537 functions represented by the :attr:`coefficients`, in 3-by-3 matrix form stored
538 in the last two indices: ``[[xx, xy, xz], [xy, yy, yz], [xz, yz, zz]]``
539 Second, eigenvalues of the rank-2 tensors, sorted in ascending order.
540 Third, eigenvectors associated to each eigenvalue, sorted in the same order in the last index.
541 Thus, ``eigenvectors[..., 0]`` returns the eigenvectors corresponding to the smallest eigenvalue.
543 Notes
544 -----
545 This method handles the edge case where :attr:`~.SphericalHarmonics.ell_max` is ``0`` gracefully.
546 In this case, all elements of :attr:`r2_tensors` will be ``[[1, 0, 0], [0, 0, 0], [0, 0, -1]]``,
547 all elements of :attr:`eigenvalues` will be ``[-1, 0, 1]``, and all elements of
548 :attr:`eigenvectors` will ``[[1, 0, 0], [0, 1, 0], [0, 0, 1]]``.
550 """
551 if self._ell_max < 2:
552 logger.info('Note: ell_max < 2, so rank-2 tensors will all be diagonal matrices with diagonal'
553 ' [1, 0, -1], eigenvalues will be [1, 0, -1],'
554 ' and eigenvectors will be Cartesian basis vectors.')
555 eigenvalues = np.zeros(coefficients.shape[:-1] + (3,), dtype=coefficients.dtype)
556 eigenvectors = np.zeros(coefficients.shape[:-1] + (3, 3), dtype=coefficients.dtype)
557 eigenvalues[...] = [-1, 0, 1]
558 eigenvectors[..., 0, 0] = 1.
559 eigenvectors[..., 1, 1] = 1.
560 eigenvectors[..., 2, 2] = 1.
561 r2_tensor = np.zeros(coefficients.shape[:-1] + (3, 3), dtype=coefficients.dtype)
562 r2_tensor[..., 0, 0] = 1
563 r2_tensor[..., 1, 1] = 0
564 r2_tensor[..., 2, 2] = -1
565 return r2_tensor, eigenvectors, eigenvalues
566 A = coefficients[..., self._ell_indices == 2]
568 # Normalizing coefficients
569 c1 = np.sqrt(15)
570 c2 = np.sqrt(5)
571 c3 = np.sqrt(15)
572 r2_tensor = np.zeros(coefficients.shape[:-1] + (3, 3), dtype=coefficients.dtype)
573 r2_tensor[..., 0, 0] = c3 * A[..., 4] - c2 * A[..., 2]
574 r2_tensor[..., 1, 1] = -c3 * A[..., 4] - c2 * A[..., 2]
575 r2_tensor[..., 2, 2] = c2 * 2 * A[..., 2]
576 r2_tensor[..., 0, 1] = c1 * A[..., 0]
577 r2_tensor[..., 1, 0] = c1 * A[..., 0]
578 r2_tensor[..., 2, 1] = c1 * A[..., 1]
579 r2_tensor[..., 1, 2] = c1 * A[..., 1]
580 r2_tensor[..., 2, 0] = c1 * A[..., 3]
581 r2_tensor[..., 0, 2] = c1 * A[..., 3]
582 w, v = np.linalg.eigh(r2_tensor.reshape(-1, 3, 3))
584 # Some complicated sorting logic to sort eigenvectors per ascending eigenvalues.
585 sorting = np.argsort(w, axis=1).reshape(-1, 3, 1)
586 v = v.transpose(0, 2, 1)
587 v = np.take_along_axis(v, sorting, axis=1)
588 v = v.transpose(0, 2, 1)
589 v = v / np.sqrt(np.sum(v ** 2, axis=1).reshape(-1, 1, 3))
590 eigenvalues = w.reshape(coefficients.shape[:-1] + (3,))
591 eigenvectors = v.reshape(coefficients.shape[:-1] + (3, 3,))
592 return r2_tensor.reshape(*coefficients.shape[:-1], 3, 3), eigenvectors, eigenvalues
594 def __iter__(self) -> Iterator[Tuple[str, Any]]:
595 """ Allows class to be iterated over and in particular be cast as a dictionary.
596 """
597 yield 'name', type(self).__name__
598 yield 'ell_max', self._ell_max
599 yield 'ell_indices', self._ell_indices
600 yield 'emm_indices', self._emm_indices
601 yield 'projection_matrix', self._projection_matrix
602 yield 'hash', hex(hash(self))[2:]
604 def __len__(self) -> int:
605 return len(self._ell_indices)
607 def __hash__(self) -> int:
608 """Returns a hash reflecting the internal state of the instance.
610 Returns
611 -------
612 A hash of the internal state of the instance,
613 cast as an ``int``.
614 """
615 to_hash = [self._ell_max,
616 self._ell_indices,
617 self._emm_indices,
618 self._projection_matrix,
619 self._probed_coordinates_hash]
620 return int(list_to_hash(to_hash), 16)
622 def _update(self) -> None:
623 # We only run updates if the hashes do not match.
624 if self.is_dirty:
625 self._projection_matrix = self._get_integrated_projection_matrix()
626 self._probed_coordinates_hash = hash(self._probed_coordinates)
628 @property
629 def is_dirty(self) -> bool:
630 return hash(self._probed_coordinates) != self._probed_coordinates_hash
632 @property
633 def projection_matrix(self) -> NDArray:
634 """ The matrix used to project spherical functions from the unit sphere onto the detector.
635 If ``v`` is a vector of spherical harmonic coefficients, and ``M`` is the ``projection_matrix``,
636 then ``M @ v`` gives the corresponding values on the detector segments associated with
637 each projection. ``M[i] @ v`` gives the values on the detector segments associated with
638 projection ``i``.
640 If ``r`` is a residual between a projection from spherical to detector space and data from
641 projection ``i``, then ``M[i].T @ r`` gives the associated gradient in spherical harmonic
642 space.
643 """
644 self._update()
645 return self._projection_matrix
647 @property
648 def ell_max(self) -> int:
649 r""" The maximum ``ell`` used to represent spherical functions.
651 Notes
652 -----
653 The word ``ell`` is used to represent the cursive small L, also written
654 :math:`\ell`, often used as an index for the degree of the Legendre polynomial
655 in the definition of the spherical harmonics.
656 """
657 return self._ell_max
659 @ell_max.setter
660 def ell_max(self, val: int) -> NDArray[int]:
661 if (val % 2 != 0 and self._enforce_friedel_symmetry) or val < 0 or val != round(val):
662 raise ValueError('ell_max must be an even (if Friedel symmetry is enforced),'
663 ' non-negative integer,'
664 f' but a value of {val} was given!')
665 self._ell_max = val
666 self._calculate_coefficient_indices()
667 self._projection_matrix = self._get_integrated_projection_matrix()
669 @property
670 def ell_indices(self) -> NDArray[int]:
671 r""" The ``ell`` associated with each coefficient and its corresponding
672 spherical harmonic. Updated when :attr:`~.SphericalHarmonics.ell_max` changes.
674 Notes
675 -----
676 The word ``ell`` is used to represent the cursive small L, also written
677 :math:`\ell`, often used as an index for the degree of the Legendre polynomial
678 in the definition of the spherical harmonics.
679 """
680 return self._ell_indices
682 @property
683 def emm_indices(self) -> NDArray[int]:
684 r""" The ``emm`` associated with each coefficient and its corresponding
685 spherical harmonic. Updated when :attr:`~.SphericalHarmonics.ell_max` changes.
687 Notes
688 -----
689 For consistency with :attr:`~.SphericalHarmonics.ell_indices`, and to avoid
690 visual confusion with other letters, ``emm`` is
691 used to represent the index commonly written :math:`m` in mathematical notation,
692 the frequency of the sine-cosine parts of the spherical harmonics,
693 often called the spherical harmonic order.
694 """
695 return self._emm_indices
697 def __str__(self) -> str:
698 wdt = 74
699 s = []
700 s += ['-' * wdt]
701 s += ['SphericalHarmonics'.center(wdt)]
702 s += ['-' * wdt]
703 with np.printoptions(threshold=4, edgeitems=2, precision=5, linewidth=60):
704 s += ['{:18} : {}'.format('Maximum "ell"', self.ell_max)]
705 s += ['{:18} : {}'.format('"ell" indices', self.ell_indices)]
706 s += ['{:18} : {}'.format('"emm" indices', self.emm_indices)]
707 s += ['{:18} : {}'.format('Projection matrix', self.projection_matrix)]
708 s += ['{:18} : {}'.format('Hash', hex(hash(self))[2:8])]
709 s += ['-' * wdt]
710 return '\n'.join(s)
712 def _repr_html_(self) -> str:
713 s = []
714 s += ['<h3>SphericalHarmonics</h3>']
715 s += ['<table border="1" class="dataframe">']
716 s += ['<thead><tr><th style="text-align: left;">Field</th><th>Size</th><th>Data</th></tr></thead>']
717 s += ['<tbody>']
718 with np.printoptions(threshold=4, edgeitems=2, precision=2, linewidth=40):
719 s += ['<tr><td style="text-align: left;">Maximum "ell"</td>']
720 s += [f'<td>{1}</td><td>{self.ell_max}</td></tr>']
721 s += ['<tr><td style="text-align: left;">"ell" indices</td>']
722 s += [f'<td>{len(self.ell_indices)}</td><td>{self.ell_indices}</td></tr>']
723 s += ['<tr><td style="text-align: left;">"emm" indices</td>']
724 s += [f'<td>{len(self.emm_indices)}</td><td>{self.emm_indices}</td></tr>']
725 s += ['<tr><td style="text-align: left;">Coefficient projection matrix</td>']
726 s += [f'<td>{self.projection_matrix.shape}</td><td>{self.projection_matrix}</td></tr>']
727 s += ['<tr><td style="text-align: left;">Hash</td>']
728 s += [f'<td>{len(hex(hash(self)))}</td><td>{hex(hash(self))[2:8]}</td></tr>']
729 s += ['</tbody>']
730 s += ['</table>']
731 return '\n'.join(s)