Coverage for local_installation_linux/mumott/optimization/regularizers/laplacian.py: 92%

1import numpy as np 

2from numpy.typing import NDArray 

3 

4from mumott.optimization.regularizers.base_regularizer import Regularizer 

5import logging 

6logger = logging.getLogger(__name__) 

7 

8 

9class Laplacian(Regularizer): 

10 

11 """Regularizes using the Laplacian of the coefficients. Suitable for orthonormal representations, 

12 e.g., spherical harmonics. 

13 """ 

14 

15 def __init__(self): 

16 super().__init__() 

17 

18 def get_regularization_norm(self, 

19 coefficients: NDArray[float], 

20 get_gradient: bool = False, 

21 gradient_part: str = None) -> dict[str, NDArray[float]]: 

22 """Retrieves the regularization norm and possibly the gradient based on the provided coefficients. 

23 The norm is the 2-norm of the discrete nearest-neighbour Laplacian. 

24 

25 This is in effect a smoothing kernel that enforces continuity between the tensors of 

26 neighbouring voxels. The calculation is most suitable for orthonormal representations. 

27 If the representation is in spherical harmonics, the norm corresponds to maximimzing 

28 the covariance between neighbours. 

29 

30 Parameters 

31 ---------- 

32 coefficients 

33 An ``np.ndarray`` of values, with shape ``(X, Y, Z, W)``, where 

34 the last channel contains, e.g., tensor components. 

35 get_gradient 

36 If ``True``, returns a ``'gradient'`` of the same shape as :attr:`coefficients`. 

37 Otherwise, the entry ``'gradient'`` will be ``None``. Defaults to ``False``. 

38 gradient_part 

39 Used for the zonal harmonics (ZH) reconstructions to determine what part of the gradient is 

40 being calculated. Default is ``None``. 

41 If a flag is passed in (``'full'``, ``'angles'``, ``'coefficients'``), 

42 we assume that the ZH workflow is used and that the last two coefficients are Euler angles, 

43 which should not be regularized by this regularizer. 

44 

45 Returns 

46 ------- 

47 A dictionary with two entries, ``regularization_norm`` and ``gradient``. 

48 """ 

49 

50 r = 6 * coefficients 

51 slices_r = [np.s_[1:, :, :], np.s_[:-1, :, :], 

52 np.s_[:, 1:, :], np.s_[:, :-1, :], 

53 np.s_[:, :, 1:], np.s_[:, :, :-1]] 

54 slices_coeffs = [np.s_[:-1, :, :], np.s_[1:, :, :], 

55 np.s_[:, :-1, :], np.s_[:, 1:, :], 

56 np.s_[:, :, :-1], np.s_[:, :, 1:]] 

57 

58 for s, v in zip(slices_r, slices_coeffs): 

59 r[s] -= coefficients[v] 

60 

61 r *= 0.5 

62 

63 result = dict(regularization_norm=None, gradient=None) 

64 if gradient_part is None: 

65 result['regularization_norm'] = np.dot(r.ravel(), r.ravel()) 

66 elif gradient_part in ('full', 'coefficients', 'angles'): 66 ↛ 70line 66 didn't jump to line 70, because the condition on line 66 was never false

67 r = r[..., :-2] 

68 result['regularization_norm'] = np.dot(r.ravel(), r.ravel()) 

69 else: 

70 raise ValueError('Unexpected argument given for gradient part.') 

71 

72 if get_gradient: 

73 if gradient_part is None: 

74 result['gradient'] = r 

75 elif gradient_part in ('full', 'coefficients'): 

76 result['gradient'] = np.zeros(coefficients.shape) 

77 result['gradient'][..., :-2] = r 

78 elif gradient_part in ('angles'): 78 ↛ 81line 78 didn't jump to line 81, because the condition on line 78 was never false

79 result['gradient'] = np.zeros(coefficients.shape) 

80 else: 

81 raise ValueError('Unexpected argument given for gradient part.') 

82 

83 return result 

84 

85 @property 

86 def _function_as_str(self) -> str: 

87 return 'R(x) = 0.5 * lambda * (2 * x[i] - x[i + 1] - x[i - 1]) ** 2' 

88 

89 @property 

90 def _function_as_tex(self) -> str: 

91 return r'$R(\vec{x}) = \lambda \frac{\Vert \nabla^2 \vec{x} \Vert^2}{2}$'