Coverage for local_installation_linux/mumott/methods/basis_sets/nearest_neighbor.py: 92%
206 statements
« prev ^ index » next coverage.py v7.3.2, created at 2025-05-05 21:21 +0000
1import logging
2from typing import Tuple
3from copy import deepcopy
5import numpy as np
6from numpy.typing import NDArray
8from mumott import ProbedCoordinates, DataContainer, Geometry, SphericalHarmonicMapper
9from mumott.core.hashing import list_to_hash
10from mumott.methods.utilities.tensor_operations import (framewise_contraction,
11 framewise_contraction_transpose)
12from mumott.output_handling.reconstruction_derived_quantities import\
13 (ReconstructionDerivedQuantities, get_sorted_eigenvectors)
14from .base_basis_set import BasisSet
16logger = logging.getLogger(__name__)
19class NearestNeighbor(BasisSet):
20 r""" Basis set class for nearest-neighbor interpolation. Used to construct methods similar to that
21 presented in `Schaff et al. (2015) <https://doi.org/10.1038/nature16060>`_.
22 By default this representation is sparse and maps only a single direction on the sphere
23 to each detector segment. This can be changed; see ``kwargs``.
25 Parameters
26 ----------
27 directions : NDArray[float]
28 Two-dimensional Array containing the ``N`` sensitivity directions with shape ``(N, 3)``.
29 probed_coordinates : ProbedCoordinates
30 Optional. Coordinates on the sphere probed at each detector segment by the
31 experimental method. Its construction from the system geometry is method-dependent.
32 By default, an empty instance of :class:`mumott.ProbedCoordinates` is created.
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.
36 kwargs
37 Miscellaneous arguments which relate to segment integrations can be
38 passed as keyword arguments:
40 integration_mode
41 Mode to integrate line segments on the reciprocal space sphere. Possible options are
42 ``'simpson'``, ``'midpoint'``, ``'romberg'``, ``'trapezoid'``.
43 ``'simpson'``, ``'trapezoid'``, and ``'romberg'`` use adaptive
44 integration with the respective quadrature rule from ``scipy.integrate``.
45 ``'midpoint'`` uses a single mid-point approximation of the integral.
46 Default value is ``'simpson'``.
48 n_integration_starting_points
49 Number of points used in the first iteration of the adaptive integration.
50 The number increases by the rule ``N`` ← ``2 * N - 1`` for each iteration.
51 Default value is 3.
52 integration_tolerance
53 Tolerance for the maximum relative error between iterations before the integral
54 is considered converged. Default is ``1e-3``.
56 integration_maxiter
57 Maximum number of iterations. Default is ``10``.
58 enforce_sparsity
59 If ``True``, limites the number of basis set elements
60 that can map to each detector segemnt. Default is ``False``.
61 sparsity_count
62 If ``enforce_sparsity`` is set to ``True``, the number of
63 basis set elements that can map to each detector segment.
64 Default value is ``1``.
65 """
66 def __init__(self,
67 directions: NDArray[float],
68 probed_coordinates: ProbedCoordinates = None,
69 enforce_friedel_symmetry: bool = True,
70 **kwargs):
71 # This basis set struggles with integral convergence due to sharp transitions
72 kwargs.update(dict(integration_tolerance=kwargs.get('integration_tolerance', 1e-3),
73 sparsity_count=kwargs.get('sparsity_count', 1)))
74 super().__init__(probed_coordinates, **kwargs)
75 # Handling grid of directions
76 self._number_of_coefficients = directions.shape[0]
77 if enforce_friedel_symmetry: 77 ↛ 80line 77 didn't jump to line 80, because the condition on line 77 was never false
78 self._directions_full = np.concatenate((directions, -directions), axis=0)
79 else:
80 self._directions_full = np.array(directions)
82 self._probed_coordinates_hash = hash(self.probed_coordinates)
84 self._enforce_friedel_symmetry = enforce_friedel_symmetry
85 self._projection_matrix = self._get_integrated_projection_matrix()
87 def find_nearest_neighbor_index(self, probed_directions: NDArray[float]) -> NDArray[int]:
88 """
89 Caluculate the nearest neighbor sensitivity directions for an array of x-y-z vectors.
91 Parameters
92 ----------
93 probed_directions
94 Array with length 3 along its last axis
96 Returns
97 -------
98 Array with same shape as the input except for the last dimension, which
99 contains the index of the nearest-neighbor sensitivity direction.
100 """
102 # normalize input directions
103 input_shape = probed_directions.shape
104 normed_probed_directions = probed_directions / \
105 np.linalg.norm(probed_directions, axis=-1)[..., np.newaxis]
107 # Find distance (3D euclidian) between each probed direction and sensitivity direction
108 pad_dimension = (1,) * (len(input_shape)-1)
109 distance = np.sum((normed_probed_directions[np.newaxis, ...] -
110 self._directions_full.reshape(self._directions_full.shape[0],
111 *pad_dimension, 3))**2, axis=-1)
113 # Find nearest_neighbor
114 best_dir = np.argmin(distance, axis=0)
116 if self._enforce_friedel_symmetry: 116 ↛ 118line 116 didn't jump to line 118, because the condition on line 116 was never false
117 best_dir = best_dir % self._number_of_coefficients
118 return best_dir
120 def get_function_values(self, probed_directions: NDArray) -> NDArray[float]:
121 """
122 Calculate the value of the basis functions from an array of x-y-z vectors.
124 Parameters
125 ----------
126 probed_directions
127 Array with length 3 along its last axis
129 Returns
130 -------
131 Array with same shape as input array except for the last axis, which now
132 has length ``N``, i.e., the number of sensitivity directions.
134 """
136 best_dir = self.find_nearest_neighbor_index(probed_directions)
137 input_shape = probed_directions.shape
138 output_array = np.zeros((*input_shape[:-1], self._number_of_coefficients))
139 for mode_number in range(self._number_of_coefficients):
140 output_array[best_dir == mode_number, mode_number] = 1.0
141 return output_array
143 def get_amplitudes(self, coefficients: NDArray[float],
144 probed_directions: NDArray[float]) -> NDArray[float]:
145 """
146 Calculate function values of an array of coefficients.
148 Parameters
149 ----------
150 coefficients
151 Array of coefficients with coefficient number along its last index.
152 probed_directions
153 Array with length 3 along its last axis.
155 Returns
156 -------
157 Array with function values. The shape of the array is
158 ``(*coefficients.shape[:-1], *probed_directions.shape[:-1])``.
159 """
160 final_shape = (*coefficients.shape[:-1], *probed_directions.shape[:-1])
161 nn_index = self.find_nearest_neighbor_index(probed_directions).ravel()
162 amplitudes = np.zeros((np.prod(coefficients.shape[:-1]), np.prod(probed_directions.shape[:-1])))
163 coefficients = np.reshape(coefficients, (np.prod(coefficients.shape[:-1]),
164 coefficients.shape[-1]))
165 for coeff_index in range(amplitudes.shape[0]):
166 amplitudes[coeff_index, :] = coefficients[coeff_index, nn_index]
168 return amplitudes.reshape(final_shape)
170 def get_second_moments(self, coefficients: NDArray[float]) -> NDArray[float]:
171 """
172 Calculate the second moments of the functions described by :attr:`coefficients`.
174 Parameters
175 ----------
176 coefficients
177 An array of coefficients (or residuals) of arbitrary shape so long as the last
178 axis has the same size as the number of detector channels.
180 Returns
181 -------
182 Array containing the second moments of the functions described by coefficients,
183 formatted as rank-two tensors with tensor indices in the last 2 dimensions.
184 """
186 if not self._enforce_friedel_symmetry:
187 raise NotImplementedError('NearestNeighbor.get_second_moments does not support'
188 ' cases with Friedel symmetry.')
190 second_moments_array = np.zeros((*coefficients.shape[:-1], 3, 3))
192 sumint = np.zeros(coefficients.shape[:-1])
193 sumxx = np.zeros(coefficients.shape[:-1])
194 sumxy = np.zeros(coefficients.shape[:-1])
195 sumxz = np.zeros(coefficients.shape[:-1])
196 sumyy = np.zeros(coefficients.shape[:-1])
197 sumyz = np.zeros(coefficients.shape[:-1])
198 sumzz = np.zeros(coefficients.shape[:-1])
200 for mode_number in range(len(self)):
202 sumint += coefficients[..., mode_number]
203 sumxx += coefficients[..., mode_number] * self._directions_full[mode_number, 0]**2
204 sumxy += coefficients[..., mode_number] * self._directions_full[mode_number, 0]\
205 * self._directions_full[mode_number, 1]
206 sumxz += coefficients[..., mode_number] * self._directions_full[mode_number, 0]\
207 * self._directions_full[mode_number, 2]
208 sumyy += coefficients[..., mode_number] * self._directions_full[mode_number, 1]**2
209 sumyz += coefficients[..., mode_number] * self._directions_full[mode_number, 1]\
210 * self._directions_full[mode_number, 2]
211 sumzz += coefficients[..., mode_number] * self._directions_full[mode_number, 2]**2
213 second_moments_array[..., 0, 0] = sumxx / len(self)
214 second_moments_array[..., 0, 1] = sumxy / len(self)
215 second_moments_array[..., 0, 2] = sumxz / len(self)
216 second_moments_array[..., 1, 0] = sumxy / len(self)
217 second_moments_array[..., 1, 1] = sumyy / len(self)
218 second_moments_array[..., 1, 2] = sumyz / len(self)
219 second_moments_array[..., 2, 0] = sumxz / len(self)
220 second_moments_array[..., 2, 1] = sumyz / len(self)
221 second_moments_array[..., 2, 2] = sumzz / len(self)
223 return second_moments_array
225 def get_spherical_harmonic_coefficients(
226 self,
227 coefficients: NDArray[float],
228 ell_max: int = None
229 ) -> NDArray[float]:
230 """ Computes and rturns the spherical harmonic coefficients of the spherical function
231 represented by the provided :attr:`coefficients` using a Driscoll-Healy grid.
233 For details on the Driscoll-Healy grid, see
234 `the SHTools page <https://shtools.github.io/SHTOOLS/grid-formats.html>`_ for a
235 comprehensive overview.
237 Parameters
238 ----------
239 coefficients
240 An array of coefficients of arbitrary shape, provided that the
241 last dimension contains the coefficients for one function.
242 ell_max
243 The bandlimit of the spherical harmonic expansion.
245 """
246 dh_grid_size = 2*ell_max + 1
247 mapper = SphericalHarmonicMapper(ell_max=ell_max, polar_resolution=dh_grid_size,
248 azimuthal_resolution=dh_grid_size,
249 enforce_friedel_symmetry=self._enforce_friedel_symmetry)
250 coordinates = mapper.unit_vectors
251 amplitudes = self.get_amplitudes(coefficients, coordinates)
252 spherical_harmonics_coefficients = mapper.get_harmonic_coefficients(amplitudes)
253 return spherical_harmonics_coefficients
255 def _get_projection_matrix(self, probed_coordinates: ProbedCoordinates = None) -> NDArray[float]:
256 """ Computes the matrix necessary for forward and gradient calculations.
257 Called when the coordinate system has been updated, or one of
258 :attr:`kernel_scale_parameter` or :attr:`grid_scale` has been changed."""
259 if probed_coordinates is None: 259 ↛ 260line 259 didn't jump to line 260, because the condition on line 259 was never true
260 probed_coordinates = self._probed_coordinates
261 return self.get_function_values(probed_coordinates.vector)
263 def get_sub_geometry(self,
264 direction_index: int,
265 geometry: Geometry,
266 data_container: DataContainer = None,
267 ) -> tuple[Geometry, tuple[NDArray[float], NDArray[float]]]:
268 """ Create and return a geometry object corresponding to a scalar tomography problem for
269 scattering along the sensitivity direction with index :attr:`direction_index`.
270 If optionally a :class:`mumott.DataContainer` is provided, the sinograms and weights for this
271 scalar tomography problem will alse be returned.
273 Used for an implementation of the algorithm descibed in [Schaff2015]_.
275 Parameters
276 ----------
277 direction_index
278 Index of the sensitivity direction.
279 geometry
280 :class:`mumott.Geometry` object of the full problem.
281 data_container (optional)
282 :class:`mumott.DataContainer` compatible with :attr:`Geometry` from which a scalar dataset
283 will be constructed.
285 returns
286 -------
287 sub_geometry
288 Geometry of the scalar problem.
289 data_tuple
290 :class:`Tuple` containing two numpy arrays. :attr:`data_tuple[0]` is the data of the
291 scalar problem. :attr:`data_tuple[1]` are the weights.
292 """
293 if self._integration_mode != 'midpoint':
294 logger.info("The 'Discrete Directions' reconstruction workflow has not been tested"
295 "with detector segment integration. Set :attr:`integration_mode` to ``'midpoint'``"
296 ' or proceed with caution.')
298 # Get projection weights
299 probed_coordinates = ProbedCoordinates()
300 probed_coordinates.vector = geometry.probed_coordinates.vector
301 projection_matrix = self._get_integrated_projection_matrix(probed_coordinates)[..., direction_index]
303 # Copy over certain parts of geometry
304 sub_geometry = deepcopy(geometry)
305 sub_geometry.delete_projections()
306 sub_geometry.detector_angles = np.array([0])
307 sub_geometry.detector_direction_origin = np.array([0, 0, 0])
308 sub_geometry.detector_direction_positive_90 = np.array([0, 0, 0])
310 if data_container is not None:
311 data_list = []
312 weight_list = []
314 for projection_index in range(len(geometry)):
315 if np.any(projection_matrix[projection_index, :] > 0.0):
317 # append sub geometry
318 sub_geometry.append(deepcopy(geometry[projection_index]))
320 # Load data if given
321 if data_container is not None:
323 projection_weight = projection_matrix[projection_index, :]
324 weighted_weights = data_container.projections[projection_index].weights\
325 * projection_weight[np.newaxis, np.newaxis, :]
326 weighted_data = data_container.projections[projection_index].data\
327 * weighted_weights
329 weight_list.append(np.sum(weighted_weights, axis=-1))
330 summed_data = np.sum(weighted_data, axis=-1)
331 data_list.append(
332 np.divide(summed_data,
333 weight_list[-1],
334 out=np.zeros(summed_data.shape),
335 where=weight_list[-1] != 0)
336 ) # Avoid runtime warning when weights are zero.
338 if data_container is None:
339 return sub_geometry, None
340 elif len(data_list) == 0:
341 logger.warning('No projections found for current direction.')
342 return sub_geometry, None
343 else:
344 data_array = np.stack(data_list, axis=0)
345 weight_array = np.stack(weight_list, axis=0)
346 return sub_geometry, (data_array, weight_array)
348 # TODO there could be a bit of a speedup by doing this without matrix products
349 def forward(self,
350 coefficients: NDArray[float],
351 indices: NDArray[int] = None) -> NDArray[float]:
352 """ Carries out a forward computation of projections from reciprocal space modes to
353 detector channels, for one or several tomographic projections.
355 Parameters
356 ----------
357 coefficients
358 An array of coefficients, of arbitrary shape so long as the last
359 axis has the same size as this basis set.
360 indices
361 Optional. Indices of the tomographic projections for which the forward
362 computation is to be performed. If ``None``, the forward computation will
363 be performed for all projections.
365 Returns
366 -------
367 An array of values on the detector corresponding to the :attr:`coefficients` given.
368 If :attr:`indices` contains exactly one index, the shape is ``(coefficients.shape[:-1], J)``
369 where ``J`` is the number of detector segments. If :attr:`indices` is ``None`` or contains
370 several indices, the shape is ``(N, coefficients.shape[1:-1], J)`` where ``N``
371 is the number of tomographic projections for which the computation is performed.
372 """
373 assert coefficients.shape[-1] == len(self)
374 self._update()
375 output = np.zeros(coefficients.shape[:-1] + (self._projection_matrix.shape[1],),
376 coefficients.dtype)
377 if indices is None: 377 ↛ 381line 377 didn't jump to line 381, because the condition on line 377 was never false
378 framewise_contraction_transpose(self._projection_matrix,
379 coefficients,
380 output)
381 elif indices.size == 1:
382 np.einsum('ijk, ...k -> ...j',
383 self._projection_matrix[indices],
384 coefficients,
385 out=output,
386 optimize='greedy',
387 casting='unsafe')
388 else:
389 framewise_contraction_transpose(self._projection_matrix[indices],
390 coefficients,
391 output)
392 return output
394 def gradient(self,
395 coefficients: NDArray[float],
396 indices: NDArray[int] = None) -> NDArray[float]:
397 """ Carries out a gradient computation of projections of projections from reciprocal space modes to
398 detector channels, for one or several tomographic projections.
400 Parameters
401 ----------
402 coefficients
403 An array of coefficients (or residuals) of arbitrary shape so long as the last
404 axis has the same size as the number of detector channels.
405 indices
406 Optional. Indices of the tomographic projections for which the gradient
407 computation is to be performed. If ``None``, the gradient computation will
408 be performed for all projections.
410 Returns
411 -------
412 An array of gradient values based on the :attr:`coefficients` given.
413 If :attr:`indices` contains exactly one index, the shape is ``(coefficients.shape[:-1], J)``
414 where ``J`` is the number of detector segments. If indices is ``None`` or contains
415 several indices, the shape is ``(N, coefficients.shape[1:-1], J)`` where ``N``
416 is the number of tomographic projections for which the computation is performed.
417 """
418 self._update()
419 output = np.zeros(coefficients.shape[:-1] + (self._projection_matrix.shape[2],),
420 coefficients.dtype)
421 if indices is None: 421 ↛ 425line 421 didn't jump to line 425, because the condition on line 421 was never false
422 framewise_contraction(self._projection_matrix,
423 coefficients,
424 output)
425 elif indices.size == 1:
426 np.einsum('ikj, ...k -> ...j',
427 self._projection_matrix[indices],
428 coefficients,
429 out=output,
430 optimize='greedy',
431 casting='unsafe')
432 else:
433 framewise_contraction(self._projection_matrix[indices],
434 coefficients,
435 output)
436 return output
438 def get_output(self,
439 coefficients: NDArray) -> ReconstructionDerivedQuantities:
440 r""" Returns a :class:`ReconstructionDerivedQuantities` instance of output data for
441 a given array of basis set coefficients.
443 Parameters
444 ----------
445 coefficients
446 An array of coefficients of arbitrary shape and dimensions, except
447 its last dimension must be the same length as the :attr:`len` of this instance.
448 Computations only operate over the last axis of :attr:`coefficients`, so derived
449 properties in the output will have the shape ``(*coefficients.shape[:-1], ...)``.
451 Returns
452 -------
453 :class:`ReconstructionDerivedQuantities` containing a number of quantities that
454 have been computed from the spherical functions represented by the input
455 coefficients.
456 """
458 assert coefficients.shape[-1] == len(self)
459 # Update to ensure non-dirty output state.
460 self._update()
462 mean_intensity = np.mean(coefficients, axis=-1)
463 second_moment_tensor = self.get_second_moments(coefficients)
464 eigenvalues, eigenvectors = get_sorted_eigenvectors(second_moment_tensor)
465 fractional_anisotropy = np.sqrt((eigenvalues[..., 0] - eigenvalues[..., 1])**2
466 + (eigenvalues[..., 1] - eigenvalues[..., 2])**2
467 + (eigenvalues[..., 2] - eigenvalues[..., 0])**2)
468 fractional_anisotropy = fractional_anisotropy / np.sqrt(2*np.sum(eigenvalues**2, axis=-1))
470 reconstruction_derived_quantities = ReconstructionDerivedQuantities(
471 volume_shape=tuple(coefficients.shape[:3]),
472 mean_intensity=mean_intensity,
473 fractional_anisotropy=fractional_anisotropy,
474 eigenvector_1=np.copy(eigenvectors[..., 0]),
475 eigenvector_2=np.copy(eigenvectors[..., 1]),
476 eigenvector_3=np.copy(eigenvectors[..., 2]),
477 eigenvalue_1=np.copy(eigenvalues[..., 0]),
478 eigenvalue_2=np.copy(eigenvalues[..., 1]),
479 eigenvalue_3=np.copy(eigenvalues[..., 2]),
480 second_moment_tensor=second_moment_tensor
481 )
483 return reconstruction_derived_quantities
485 def __len__(self) -> int:
486 return self._number_of_coefficients
488 def __hash__(self) -> int:
489 """Returns a hash reflecting the internal state of the instance.
491 Returns
492 -------
493 A hash of the internal state of the instance,
494 cast as an ``int``.
495 """
496 to_hash = [self.grid,
497 self._enforce_friedel_symmetry,
498 self._projection_matrix,
499 self._probed_coordinates_hash]
500 return int(list_to_hash(to_hash), 16)
502 def _update(self) -> None:
503 # We only run updates if the hashes do not match.
504 if self.is_dirty: 504 ↛ 505line 504 didn't jump to line 505, because the condition on line 504 was never true
505 self._projection_matrix = self._get_integrated_projection_matrix()
506 self._probed_coordinates_hash = hash(self._probed_coordinates)
508 @property
509 def is_dirty(self) -> bool:
510 return hash(self._probed_coordinates) != self._probed_coordinates_hash
512 @property
513 def projection_matrix(self) -> NDArray:
514 """ The matrix used to project spherical functions from the unit sphere onto the detector.
515 If ``v`` is a vector of gaussian kernel coefficients, and ``M`` is the ``projection_matrix``,
516 then ``M @ v`` gives the corresponding values on the detector segments associated with
517 each projection. ``M[i] @ v`` gives the values on the detector segments associated with
518 projection ``i``.
519 """
520 self._update()
521 return self._projection_matrix
523 @property
524 def enforce_friedel_symmetry(self) -> bool:
525 """ If ``True``, Friedel symmetry is enforced, i.e., the point
526 :math:`-r` is treated as equivalent to :math:`r`. """
527 return self._enforce_friedel_symmetry
529 @property
530 def grid(self) -> Tuple[NDArray['float'], NDArray['float']]:
531 r""" Returns the polar and azimuthal angles of the grid used by the basis.
533 Returns
534 -------
535 A ``Tuple`` with contents ``(polar_angle, azimuthal_angle)``, where the
536 polar angle is defined as :math:`\arccos(z)`.
537 """
538 return self._directions_full[:self._number_of_coefficients, :]
540 @property
541 def grid_hash(self) -> str:
542 """ Returns a hash of :attr:`grid`.
543 """
544 return list_to_hash([self.grid])
546 @property
547 def projection_matrix_hash(self) -> str:
548 """ Returns a hash of :attr:`projection_matrix`.
549 """
550 return list_to_hash([self.projection_matrix])
552 def __str__(self) -> str:
553 wdt = 74
554 s = [self.__class__.__name__]
555 s += ['-' * wdt]
556 s += [''.center(wdt)]
557 s += ['-' * wdt]
558 with np.printoptions(threshold=4, edgeitems=2, precision=5, linewidth=60):
559 s += ['{:18} : {}'.format('number of directions', len(self))]
560 s += ['{:18} : {}'.format('grid_hash', self.grid_hash[:6])]
561 s += ['{:18} : {}'.format('enforce_friedel_symmetry', self.enforce_friedel_symmetry)]
562 s += ['{:18} : {}'.format('projection_matrix_hash', self.projection_matrix_hash[2:8])]
563 s += ['{:18} : {}'.format('hash', hex(hash(self))[2:8])]
564 s += ['-' * wdt]
565 return '\n'.join(s)
567 def _repr_html_(self) -> str:
568 s = []
569 s += [f'<h3>{self.__class__.__name__}</h3>']
570 s += ['<table border="1" class="dataframe">']
571 s += ['<thead><tr><th style="text-align: left;">Field</th><th>Size</th><th>Data</th></tr></thead>']
572 s += ['<tbody>']
573 with np.printoptions(threshold=4, edgeitems=2, precision=2, linewidth=40):
574 s += ['<tr><td style="text-align: left;">grid_hash</td>']
575 s += [f'<td>{len(self.grid_hash)}</td><td>{self.grid_hash[:6]}</td></tr>']
576 s += ['<tr><td style="text-align: left;">enforce_friedel_symmetry</td>']
577 s += [f'<td>1</td>'
578 f'<td>{self.enforce_friedel_symmetry}</td></tr>']
579 s += ['<tr><td style="text-align: left;">projection_matrix</td>']
580 s += [f'<td>{len(self.projection_matrix_hash)}</td>'
581 f'<td>{self.projection_matrix_hash[:6]}</td></tr>']
582 s += ['<tr><td style="text-align: left;">hash</td>']
583 s += [f'<td>{len(hex(hash(self)))}</td><td>{hex(hash(self))[2:8]}</td></tr>']
584 s += ['</tbody>']
585 s += ['</table>']
586 return '\n'.join(s)