1"""Tools for using spherical harmonic models to fit diffusion data.
2
3References
4----------
5Aganj, I., et al. 2009. ODF Reconstruction in Q-Ball Imaging With Solid
6    Angle Consideration.
7Descoteaux, M., et al. 2007. Regularized, fast, and robust analytical
8    Q-ball imaging.
9Tristan-Vega, A., et al. 2010. A new methodology for estimation of fiber
10    populations in white matter of the brain with Funk-Radon transform.
11Tristan-Vega, A., et al. 2009. Estimation of fiber orientation probability
12    density functions in high angular resolution diffusion imaging.
13
14
15Note about the Transpose:
16In the literature the matrix representation of these methods is often written
17as Y = Bx where B is some design matrix and Y and x are column vectors. In our
18case the input data, a dwi stored as a nifti file for example, is stored as row
19vectors (ndarrays) of the form (x, y, z, n), where n is the number of diffusion
20directions. We could transpose and reshape the data to be (n, x*y*z), so that
21we could directly plug it into the above equation. However, I have chosen to
22keep the data as is and implement the relevant equations rewritten in the
23following form: Y.T = x.T B.T, or in python syntax data = np.dot(sh_coef, B.T)
24where data is Y.T and sh_coef is x.T.
25"""
26
27from warnings import warn
28import numpy as np
29import scipy.special as sps
30
31from numpy.random import randint
32
33from dipy.utils.deprecator import deprecate_with_version
34from dipy.reconst.odf import OdfModel, OdfFit
35from dipy.core.geometry import cart2sphere
36from dipy.core.onetime import auto_attr
37from dipy.reconst.cache import Cache
38
39
40def _copydoc(obj):
41    def bandit(f):
42        f.__doc__ = obj.__doc__
43        return f
44    return bandit
45
46
47def forward_sdeconv_mat(r_rh, n):
48    """ Build forward spherical deconvolution matrix
49
50    Parameters
51    ----------
52    r_rh : ndarray
53        Rotational harmonics coefficients for the single fiber response
54        function. Each element ``rh[i]`` is associated with spherical harmonics
55        of degree ``2*i``.
56    n : ndarray
57        The order of spherical harmonic function associated with each row of
58        the deconvolution matrix. Only even orders are allowed
59
60    Returns
61    -------
62    R : ndarray (N, N)
63        Deconvolution matrix with shape (N, N)
64
65    """
66
67    if np.any(n % 2):
68        raise ValueError("n has odd orders, expecting only even orders")
69    return np.diag(r_rh[n // 2])
70
71
72def sh_to_rh(r_sh, m, n):
73    """ Spherical harmonics (SH) to rotational harmonics (RH)
74
75    Calculate the rotational harmonic decomposition up to
76    harmonic order ``n``, degree ``m`` for an axially and antipodally
77    symmetric function. Note that all ``m != 0`` coefficients
78    will be ignored as axial symmetry is assumed. Hence, there
79    will be ``(sh_order/2 + 1)`` non-zero coefficients.
80
81    Parameters
82    ----------
83    r_sh : ndarray (N,)
84        ndarray of SH coefficients for the single fiber response function.
85        These coefficients must correspond to the real spherical harmonic
86        functions produced by `shm.real_sh_descoteaux_from_index`.
87    m : ndarray (N,)
88        The degree of the spherical harmonic function associated with each
89        coefficient.
90    n : ndarray (N,)
91        The order of the spherical harmonic function associated with each
92        coefficient.
93
94    Returns
95    -------
96    r_rh : ndarray (``(sh_order + 1)*(sh_order + 2)/2``,)
97        Rotational harmonics coefficients representing the input `r_sh`
98
99    See Also
100    --------
101    shm.real_sh_descoteaux_from_index, shm.real_sh_descoteaux
102
103    References
104    ----------
105    .. [1] Tournier, J.D., et al. NeuroImage 2007. Robust determination of the
106        fibre orientation distribution in diffusion MRI: Non-negativity
107        constrained super-resolved spherical deconvolution
108
109    """
110    mask = m == 0
111    # The delta function at theta = phi = 0 is known to have zero coefficients
112    # where m != 0, therefore we need only compute the coefficients at m=0.
113    dirac_sh = gen_dirac(0, n[mask], 0, 0)
114    r_rh = r_sh[mask] / dirac_sh
115    return r_rh
116
117
118def gen_dirac(m, n, theta, phi, legacy=True):
119    """ Generate Dirac delta function orientated in (theta, phi) on the sphere
120
121    The spherical harmonics (SH) representation of this Dirac is returned as
122    coefficients to spherical harmonic functions produced from ``descoteaux07``
123    basis.
124
125    Parameters
126    ----------
127    m : ndarray (N,)
128        The degree of the spherical harmonic function associated with each
129        coefficient.
130    n : ndarray (N,)
131        The order of the spherical harmonic function associated with each
132        coefficient.
133    theta : float [0, pi]
134        The polar (colatitudinal) coordinate.
135    phi : float [0, 2*pi]
136        The azimuthal (longitudinal) coordinate.
137
138    See Also
139    --------
140    shm.real_sh_descoteaux_from_index, shm.real_sh_descoteaux
141
142    Returns
143    -------
144    dirac : ndarray
145        SH coefficients representing the Dirac function. The shape of this is
146        `(m + 2) * (m + 1) / 2`.
147
148    """
149    return real_sh_descoteaux_from_index(m, n, theta, phi, legacy=legacy)
150
151
152def spherical_harmonics(m, n, theta, phi, use_scipy=True):
153    """Compute spherical harmonics.
154
155    This may take scalar or array arguments. The inputs will be broadcasted
156    against each other.
157
158    Parameters
159    ----------
160    m : int ``|m| <= n``
161        The degree of the harmonic.
162    n : int ``>= 0``
163        The order of the harmonic.
164    theta : float [0, 2*pi]
165        The azimuthal (longitudinal) coordinate.
166    phi : float [0, pi]
167        The polar (colatitudinal) coordinate.
168    use_scipy : bool, optional
169        If True, use scipy implementation.
170
171    Returns
172    -------
173    y_mn : complex float
174        The harmonic $Y^m_n$ sampled at ``theta`` and ``phi``.
175
176    Notes
177    -----
178    This is a faster implementation of scipy.special.sph_harm for
179    scipy version < 0.15.0. For scipy 0.15 and onwards, we use the scipy
180    implementation of the function.
181
182    The usual definitions for ``theta` and `phi`` used in DIPY are interchanged
183    in the method definition to agree with the definitions in
184    scipy.special.sph_harm, where `theta` represents the azimuthal coordinate
185    and `phi` represents the polar coordinate.
186
187    Altough scipy uses a naming convention where ``m`` is the order and ``n``
188    is the degree of the SH, the opposite of DIPY's, their definition for both
189    parameters is the same as ours, with ``n >= 0`` and ``|m| <= n``.
190    """
191    if use_scipy:
192        return sps.sph_harm(m, n, theta, phi, dtype=complex)
193
194    x = np.cos(phi)
195    val = sps.lpmv(m, n, x).astype(complex)
196    val *= np.sqrt((2 * n + 1) / 4.0 / np.pi)
197    val *= np.exp(0.5 * (sps.gammaln(n - m + 1) - sps.gammaln(n + m + 1)))
198    val = val * np.exp(1j * m * theta)
199    return val
200
201
202@deprecate_with_version('dipy.reconst.shm.real_sph_harm is deprecated, '
203                        'Please use '
204                        'dipy.reconst.shm.real_sh_descoteaux_from_index '
205                        'instead', since='1.3', until='2.0')
206def real_sph_harm(m, n, theta, phi):
207    """ Compute real spherical harmonics.
208
209    Where the real harmonic $Y^m_n$ is defined to be:
210
211        Imag($Y^m_n$) * sqrt(2)     if m > 0
212        $Y^0_n$                     if m = 0
213        Real($Y^|m|_n$) * sqrt(2)   if m < 0
214
215    This may take scalar or array arguments. The inputs will be broadcasted
216    against each other.
217
218    Parameters
219    ----------
220    m : int ``|m| <= n``
221        The degree of the harmonic.
222    n : int ``>= 0``
223        The order of the harmonic.
224    theta : float [0, pi]
225        The polar (colatitudinal) coordinate.
226    phi : float [0, 2*pi]
227        The azimuthal (longitudinal) coordinate.
228
229    Returns
230    --------
231    y_mn : real float
232        The real harmonic $Y^m_n$ sampled at `theta` and `phi`.
233
234    See Also
235    --------
236    scipy.special.sph_harm
237    """
238    return real_sh_descoteaux_from_index(m, n, theta, phi, legacy=True)
239
240
241def real_sh_tournier_from_index(m, n, theta, phi, legacy=True):
242    """ Compute real spherical harmonics as initially defined in Tournier
243    2007 [1]_ then updated in MRtrix3 [2]_, where the real harmonic $Y^m_n$
244    is defined to be:
245
246        Real($Y^m_n$) * sqrt(2)      if m > 0
247        $Y^0_n$                      if m = 0
248        Imag($Y^|m|_n$) * sqrt(2)    if m < 0
249
250    This may take scalar or array arguments. The inputs will be broadcasted
251    against each other.
252
253    Parameters
254    ----------
255    m : int ``|m| <= n``
256        The degree of the harmonic.
257    n : int ``>= 0``
258        The order of the harmonic.
259    theta : float [0, pi]
260        The polar (colatitudinal) coordinate.
261    phi : float [0, 2*pi]
262        The azimuthal (longitudinal) coordinate.
263    legacy: bool, optional
264        If true, uses MRtrix 0.2 SH basis definition, where the ``sqrt(2)``
265        factor is omitted. Else, uses the MRtrix 3 definition presented above.
266
267    Returns
268    -------
269    real_sh : real float
270        The real harmonics $Y^m_n$ sampled at ``theta`` and ``phi``.
271
272    References
273    ----------
274    .. [1] Tournier J.D., Calamante F. and Connelly A. Robust determination
275           of the fibre orientation distribution in diffusion MRI:
276           Non-negativity constrained super-resolved spherical deconvolution.
277           NeuroImage. 2007;35(4):1459-1472.
278    .. [2] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
279           Pietsch M, et al. MRtrix3: A fast, flexible and open software
280           framework for medical image processing and visualisation.
281           NeuroImage. 2019 Nov 15;202:116-137.
282    """
283    # In the m < 0 case, Tournier basis considers |m|
284    sh = spherical_harmonics(np.abs(m), n, phi, theta)
285    real_sh = np.where(m < 0, sh.imag, sh.real)
286
287    if not legacy:
288        # The Tournier basis from MRtrix3 is normalized
289        real_sh *= np.where(m == 0, 1., np.sqrt(2))
290    else:
291        warn('The legacy tournier07 basis is outdated and will be deprecated '
292             'in a future release of DIPY. Consider using the new tournier07 '
293             'basis.', category=PendingDeprecationWarning)
294
295    return real_sh
296
297
298def real_sh_descoteaux_from_index(m, n, theta, phi, legacy=True):
299    """ Compute real spherical harmonics as in Descoteaux et al. 2007 [1]_,
300    where the real harmonic $Y^m_n$ is defined to be:
301
302        Imag($Y^m_n$) * sqrt(2)      if m > 0
303        $Y^0_n$                      if m = 0
304        Real($Y^m_n$) * sqrt(2)      if m < 0
305
306    This may take scalar or array arguments. The inputs will be broadcasted
307    against each other.
308
309    Parameters
310    ----------
311    m : int ``|m| <= n``
312        The degree of the harmonic.
313    n : int ``>= 0``
314        The order of the harmonic.
315    theta : float [0, pi]
316        The polar (colatitudinal) coordinate.
317    phi : float [0, 2*pi]
318        The azimuthal (longitudinal) coordinate.
319    legacy: bool, optional
320        If true, uses DIPY's legacy descoteaux07 implementation (where |m|
321        is used for m < 0). Else, implements the basis as defined in
322        Descoteaux et al. 2007 (without the absolute value).
323
324    Returns
325    -------
326    real_sh : real float
327        The real harmonic $Y^m_n$ sampled at ``theta`` and ``phi``.
328
329    References
330    ----------
331     .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
332           Regularized, Fast, and Robust Analytical Q-ball Imaging.
333           Magn. Reson. Med. 2007;58:497-510.
334    """
335    if legacy:
336        # In the case where m < 0, legacy descoteaux basis considers |m|
337        warn('The legacy descoteaux07 SH basis is outdated and will be '
338             'deprecated in a future DIPY release. Consider using the new '
339             'descoteaux07 basis.', category=PendingDeprecationWarning)
340        sh = spherical_harmonics(np.abs(m), n, phi, theta)
341    else:
342        # In the cited paper, the basis is defined without the absolute value
343        sh = spherical_harmonics(m, n, phi, theta)
344
345    real_sh = np.where(m > 0, sh.imag, sh.real)
346    real_sh *= np.where(m == 0, 1., np.sqrt(2))
347
348    return real_sh
349
350
351def real_sh_tournier(sh_order, theta, phi,
352                     full_basis=False,
353                     legacy=True):
354    """ Compute real spherical harmonics as initially defined in Tournier
355    2007 [1]_ then updated in MRtrix3 [2]_, where the real harmonic $Y^m_n$
356    is defined to be:
357
358        Real($Y^m_n$) * sqrt(2)      if m > 0
359        $Y^0_n$                      if m = 0
360        Imag($Y^|m|_n$) * sqrt(2)    if m < 0
361
362    This may take scalar or array arguments. The inputs will be broadcasted
363    against each other.
364
365    Parameters
366    ----------
367    sh_order : int
368        The maximum degree or the spherical harmonic basis.
369    theta : float [0, pi]
370        The polar (colatitudinal) coordinate.
371    phi : float [0, 2*pi]
372        The azimuthal (longitudinal) coordinate.
373    full_basis: bool, optional
374        If true, returns a basis including odd order SH functions as well as
375        even order SH functions. Else returns only even order SH functions.
376    legacy: bool, optional
377        If true, uses MRtrix 0.2 SH basis definition, where the ``sqrt(2)``
378        factor is omitted. Else, uses MRtrix 3 definition presented above.
379
380    Returns
381    -------
382    real_sh : real float
383        The real harmonic $Y^m_n$ sampled at ``theta`` and ``phi``.
384    m : array
385        The degree of the harmonics.
386    n : array
387        The order of the harmonics.
388
389    References
390    ----------
391    .. [1] Tournier J.D., Calamante F. and Connelly A. Robust determination
392           of the fibre orientation distribution in diffusion MRI:
393           Non-negativity constrained super-resolved spherical deconvolution.
394           NeuroImage. 2007;35(4):1459-1472.
395    .. [2] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
396           Pietsch M, et al. MRtrix3: A fast, flexible and open software
397           framework for medical image processing and visualisation.
398           NeuroImage. 2019 Nov 15;202:116-137.
399    """
400    m, n = sph_harm_ind_list(sh_order, full_basis)
401
402    phi = np.reshape(phi, [-1, 1])
403    theta = np.reshape(theta, [-1, 1])
404
405    real_sh = real_sh_tournier_from_index(m, n, theta, phi, legacy)
406
407    return real_sh, m, n
408
409
410def real_sh_descoteaux(sh_order, theta, phi,
411                       full_basis=False,
412                       legacy=True):
413    """ Compute real spherical harmonics as in Descoteaux et al. 2007 [1]_,
414    where the real harmonic $Y^m_n$ is defined to be:
415
416        Imag($Y^m_n$) * sqrt(2)      if m > 0
417        $Y^0_n$                      if m = 0
418        Real($Y^m_n$) * sqrt(2)      if m < 0
419
420    This may take scalar or array arguments. The inputs will be broadcasted
421    against each other.
422
423    Parameters
424    ----------
425    sh_order : int
426        The maximum degree or the spherical harmonic basis.
427    theta : float [0, pi]
428        The polar (colatitudinal) coordinate.
429    phi : float [0, 2*pi]
430        The azimuthal (longitudinal) coordinate.
431    full_basis: bool, optional
432        If true, returns a basis including odd order SH functions as well as
433        even order SH functions. Otherwise returns only even order SH
434        functions.
435    legacy: bool, optional
436        If true, uses DIPY's legacy descoteaux07 implementation (where |m|
437        for m < 0). Else, implements the basis as defined in Descoteaux et al.
438        2007.
439
440    Returns
441    -------
442    real_sh : real float
443        The real harmonic $Y^m_n$ sampled at ``theta`` and ``phi``.
444    m : array
445        The degree of the harmonics.
446    n : array
447        The order of the harmonics.
448
449    References
450    ----------
451     .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
452           Regularized, Fast, and Robust Analytical Q-ball Imaging.
453           Magn. Reson. Med. 2007;58:497-510.
454    """
455    m, n = sph_harm_ind_list(sh_order, full_basis)
456
457    phi = np.reshape(phi, [-1, 1])
458    theta = np.reshape(theta, [-1, 1])
459
460    real_sh = real_sh_descoteaux_from_index(m, n, theta, phi, legacy)
461
462    return real_sh, m, n
463
464
465@deprecate_with_version('dipy.reconst.shm.real_sym_sh_mrtix is deprecated, '
466                        'Please use dipy.reconst.shm.real_sh_tournier instead',
467                        since='1.3', until='2.0')
468def real_sym_sh_mrtrix(sh_order, theta, phi):
469    """
470    Compute real symmetric spherical harmonics as in Tournier 2007 [2]_, where
471    the real harmonic $Y^m_n$ is defined to be::
472
473        Real($Y^m_n$)       if m > 0
474        $Y^0_n$             if m = 0
475        Imag($Y^|m|_n$)     if m < 0
476
477    This may take scalar or array arguments. The inputs will be broadcasted
478    against each other.
479
480    Parameters
481    -----------
482    sh_order : int
483        The maximum order or the spherical harmonic basis.
484    theta : float [0, pi]
485        The polar (colatitudinal) coordinate.
486    phi : float [0, 2*pi]
487        The azimuthal (longitudinal) coordinate.
488
489    Returns
490    --------
491    y_mn : real float
492        The real harmonic $Y^m_n$ sampled at ``theta`` and ``phi`` as
493        implemented in mrtrix. Warning: the basis is Tournier et al.
494        2007 [2]_; 2004 [1]_ is slightly different.
495    m : array
496        The degree of the harmonics.
497    n : array
498        The order of the harmonics.
499
500    References
501    ----------
502    .. [1] Tournier J.D., Calamante F., Gadian D.G. and Connelly A.
503           Direct estimation of the fibre orientation density function from
504           diffusion-weighted MRI data using spherical deconvolution.
505           NeuroImage. 2004;23:1176-1185.
506    .. [2] Tournier J.D., Calamante F. and Connelly A. Robust determination
507           of the fibre orientation distribution in diffusion MRI:
508           Non-negativity constrained super-resolved spherical deconvolution.
509           NeuroImage. 2007;35(4):1459-1472.
510
511    """
512    return real_sh_tournier(sh_order, theta, phi, legacy=True)
513
514
515@deprecate_with_version('dipy.reconst.shm.real_sym_sh_basis is deprecated, '
516                        'Please use dipy.reconst.shm.real_sh_descoteaux '
517                        'instead', since='1.3', until='2.0')
518def real_sym_sh_basis(sh_order, theta, phi):
519    """Samples a real symmetric spherical harmonic basis at point on the sphere
520
521    Samples the basis functions up to order `sh_order` at points on the sphere
522    given by `theta` and `phi`. The basis functions are defined here the same
523    way as in Descoteaux et al. 2007 [1]_ where the real harmonic $Y^m_n$ is
524    defined to be:
525
526        Imag($Y^m_n$) * sqrt(2)     if m > 0
527        $Y^0_n$                     if m = 0
528        Real($Y^|m|_n$) * sqrt(2)   if m < 0
529
530    This may take scalar or array arguments. The inputs will be broadcasted
531    against each other.
532
533    Parameters
534    -----------
535    sh_order : int
536        even int > 0, max spherical harmonic order
537    theta : float [0, 2*pi]
538        The azimuthal (longitudinal) coordinate.
539    phi : float [0, pi]
540        The polar (colatitudinal) coordinate.
541
542    Returns
543    --------
544    y_mn : real float
545        The real harmonic $Y^m_n$ sampled at ``theta`` and ``phi``
546    m : array
547        The degree of the harmonics.
548    n : array
549        The order of the harmonics.
550
551    References
552    ----------
553    .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
554           Regularized, Fast, and Robust Analytical Q-ball Imaging.
555           Magn. Reson. Med. 2007;58:497-510.
556
557    """
558    return real_sh_descoteaux(sh_order, theta, phi, legacy=True)
559
560
561sph_harm_lookup = {None: real_sh_descoteaux,
562                   "tournier07": real_sh_tournier,
563                   "descoteaux07": real_sh_descoteaux}
564
565
566def sph_harm_ind_list(sh_order, full_basis=False):
567    """
568    Returns the degree (``m``) and order (``n``) of all the symmetric spherical
569    harmonics of degree less then or equal to ``sh_order``. The results,
570    ``m_list`` and ``n_list`` are kx1 arrays, where k depends on ``sh_order``.
571    They can be passed to :func:`real_sh_descoteaux_from_index` and
572    :func:``real_sh_tournier_from_index``.
573
574    Parameters
575    ----------
576    sh_order : int
577        even int > 0, max order to return
578    full_basis: bool, optional
579        True for SH basis with even and odd order terms
580
581    Returns
582    -------
583    m_list : array
584        degrees of even spherical harmonics
585    n_list : array
586        orders of even spherical harmonics
587
588    See also
589    --------
590    shm.real_sh_descoteaux_from_index, shm.real_sh_tournier_from_index
591    """
592    if full_basis:
593        n_range = np.arange(0, sh_order + 1, dtype=int)
594        ncoef = int((sh_order + 1) * (sh_order + 1))
595    else:
596        if sh_order % 2 != 0:
597            raise ValueError('sh_order must be an even integer >= 0')
598        n_range = np.arange(0, sh_order + 1, 2, dtype=int)
599        ncoef = int((sh_order + 2) * (sh_order + 1) // 2)
600
601    n_list = np.repeat(n_range, n_range * 2 + 1)
602    offset = 0
603    m_list = np.empty(ncoef, 'int')
604    for ii in n_range:
605        m_list[offset:offset + 2 * ii + 1] = np.arange(-ii, ii + 1)
606        offset = offset + 2 * ii + 1
607
608    # makes the arrays ncoef by 1, allows for easy broadcasting later in code
609    return (m_list, n_list)
610
611
612def order_from_ncoef(ncoef, full_basis=False):
613    """
614    Given a number ``n`` of coefficients, calculate back the ``sh_order``
615
616    Parameters
617    ----------
618    ncoef: int
619        number of coefficients
620    full_basis: bool, optional
621        True when coefficients are for a full SH basis.
622
623    Returns
624    -------
625    sh_order: int
626        maximum order of SH basis
627    """
628    if full_basis:
629        # Solve the equation :
630        # ncoef = (sh_order + 1) * (sh_order + 1)
631        return int(np.sqrt(ncoef) - 1)
632
633    # Solve the quadratic equation derived from :
634    # ncoef = (sh_order + 2) * (sh_order + 1) / 2
635    return int(-3 + np.sqrt(9 - 4 * (2-2*ncoef)))/2
636
637
638def smooth_pinv(B, L):
639    """Regularized pseudo-inverse
640
641    Computes a regularized least square inverse of B
642
643    Parameters
644    ----------
645    B : array_like (n, m)
646        Matrix to be inverted
647    L : array_like (n,)
648
649    Returns
650    -------
651    inv : ndarray (m, n)
652        regularized least square inverse of B
653
654    Notes
655    -----
656    In the literature this inverse is often written $(B^{T}B+L^{2})^{-1}B^{T}$.
657    However here this inverse is implemented using the pseudo-inverse because
658    it is more numerically stable than the direct implementation of the matrix
659    product.
660
661    """
662    L = np.diag(L)
663    inv = np.linalg.pinv(np.concatenate((B, L)))
664    return inv[:, :len(B)]
665
666
667def lazy_index(index):
668    """Produces a lazy index
669
670    Returns a slice that can be used for indexing an array, if no slice can be
671    made index is returned as is.
672    """
673    index = np.array(index)
674    assert index.ndim == 1
675    if index.dtype.kind == 'b':
676        index = index.nonzero()[0]
677    if len(index) == 1:
678        return slice(index[0], index[0] + 1)
679    step = np.unique(np.diff(index))
680    if len(step) != 1 or step[0] == 0:
681        return index
682    else:
683        return slice(index[0], index[-1] + 1, step[0])
684
685
686def _gfa_sh(coef, sh0_index=0):
687    """The gfa of the odf, computed from the spherical harmonic coefficients
688
689    This is a private function because it only works for coefficients of
690    normalized sh bases.
691
692    Parameters
693    ----------
694    coef : array
695        The coefficients, using a normalized sh basis, that represent each odf.
696    sh0_index : int, optional
697        The index of the coefficient associated with the 0th order sh harmonic.
698
699    Returns
700    -------
701    gfa_values : array
702        The gfa of each odf.
703
704    """
705    coef_sq = coef**2
706    numer = coef_sq[..., sh0_index]
707    denom = (coef_sq).sum(-1)
708    # The sum of the square of the coefficients being zero is the same as all
709    # the coefficients being zero
710    allzero = denom == 0
711    # By adding 1 to numer and denom where both and are 0, we prevent 0/0
712    numer = numer + allzero
713    denom = denom + allzero
714    return np.sqrt(1. - (numer / denom))
715
716
717class SphHarmModel(OdfModel, Cache):
718    """To be subclassed by all models that return a SphHarmFit when fit."""
719
720    def sampling_matrix(self, sphere):
721        """The matrix needed to sample ODFs from coefficients of the model.
722
723        Parameters
724        ----------
725        sphere : Sphere
726            Points used to sample ODF.
727
728        Returns
729        -------
730        sampling_matrix : array
731            The size of the matrix will be (N, M) where N is the number of
732            vertices on sphere and M is the number of coefficients needed by
733            the model.
734        """
735        sampling_matrix = self.cache_get("sampling_matrix", sphere)
736        if sampling_matrix is None:
737            sh_order = self.sh_order
738            theta = sphere.theta
739            phi = sphere.phi
740            sampling_matrix, m, n = real_sh_descoteaux(sh_order, theta, phi)
741            self.cache_set("sampling_matrix", sphere, sampling_matrix)
742        return sampling_matrix
743
744
745class QballBaseModel(SphHarmModel):
746    """To be subclassed by Qball type models."""
747    def __init__(self, gtab, sh_order, smooth=0.006, min_signal=1e-5,
748                 assume_normed=False):
749        """Creates a model that can be used to fit or sample diffusion data
750
751        Parameters
752        ------------
753        gtab : GradientTable
754            Diffusion gradients used to acquire data
755        sh_order : even int >= 0
756            the spherical harmonic order of the model
757        smooth : float between 0 and 1, optional
758            The regularization parameter of the model
759        min_signal : float, > 0, optional
760            During fitting, all signal values less than `min_signal` are
761            clipped to `min_signal`. This is done primarily to avoid values
762            less than or equal to zero when taking logs.
763        assume_normed : bool, optional
764            If True, clipping and normalization of the data with respect to the
765            mean B0 signal are skipped during mode fitting. This is an advanced
766            feature and should be used with care.
767
768        See Also
769        --------
770        normalize_data
771
772        """
773        SphHarmModel.__init__(self, gtab)
774        self._where_b0s = lazy_index(gtab.b0s_mask)
775        self._where_dwi = lazy_index(~gtab.b0s_mask)
776        self.assume_normed = assume_normed
777        self.min_signal = min_signal
778        x, y, z = gtab.gradients[self._where_dwi].T
779        r, theta, phi = cart2sphere(x, y, z)
780        B, m, n = real_sh_descoteaux(sh_order, theta[:, None], phi[:, None])
781        L = -n * (n + 1)
782        legendre0 = sps.lpn(sh_order, 0)[0]
783        F = legendre0[n]
784        self.sh_order = sh_order
785        self.B = B
786        self.m = m
787        self.n = n
788        self._set_fit_matrix(B, L, F, smooth)
789
790    def _set_fit_matrix(self, *args):
791        """Should be set in a subclass and is called by __init__"""
792        msg = "User must implement this method in a subclass"
793        raise NotImplementedError(msg)
794
795    def fit(self, data, mask=None):
796        """Fits the model to diffusion data and returns the model fit"""
797        # Normalize the data and fit coefficients
798        if not self.assume_normed:
799            data = normalize_data(data, self._where_b0s, self.min_signal)
800
801        # Compute coefficients using abstract method
802        coef = self._get_shm_coef(data)
803
804        # Apply the mask to the coefficients
805        if mask is not None:
806            mask = np.asarray(mask, dtype=bool)
807            coef *= mask[..., None]
808        return SphHarmFit(self, coef, mask)
809
810
811class SphHarmFit(OdfFit):
812    """Diffusion data fit to a spherical harmonic model"""
813
814    def __init__(self, model, shm_coef, mask):
815        self.model = model
816        self._shm_coef = shm_coef
817        self.mask = mask
818
819    @property
820    def shape(self):
821        return self._shm_coef.shape[:-1]
822
823    def __getitem__(self, index):
824        """Allowing indexing into fit"""
825        # Index shm_coefficients
826        if isinstance(index, tuple):
827            coef_index = index + (Ellipsis,)
828        else:
829            coef_index = index
830        new_coef = self._shm_coef[coef_index]
831
832        # Index mask
833        if self.mask is not None:
834            new_mask = self.mask[index]
835            assert new_mask.shape == new_coef.shape[:-1]
836        else:
837            new_mask = None
838
839        return SphHarmFit(self.model, new_coef, new_mask)
840
841    def odf(self, sphere):
842        """Samples the odf function on the points of a sphere
843
844        Parameters
845        ----------
846        sphere : Sphere
847            The points on which to sample the odf.
848
849        Returns
850        -------
851        values : ndarray
852            The value of the odf on each point of `sphere`.
853
854        """
855        B = self.model.sampling_matrix(sphere)
856        return np.dot(self.shm_coeff, B.T)
857
858    @auto_attr
859    def gfa(self):
860        return _gfa_sh(self.shm_coeff, 0)
861
862    @property
863    def shm_coeff(self):
864        """The spherical harmonic coefficients of the odf
865
866        Make this a property for now, if there is a usecase for modifying
867        the coefficients we can add a setter or expose the coefficients more
868        directly
869        """
870        return self._shm_coef
871
872    def predict(self, gtab=None, S0=1.0):
873        """
874        Predict the diffusion signal from the model coefficients.
875
876        Parameters
877        ----------
878        gtab : a GradientTable class instance
879            The directions and bvalues on which prediction is desired
880
881        S0 : float array
882           The mean non-diffusion-weighted signal in each voxel.
883        """
884        if not hasattr(self.model, 'predict'):
885            msg = "This model does not have prediction implemented yet"
886            raise NotImplementedError(msg)
887        return self.model.predict(self._shm_coef, gtab, S0)
888
889
890class CsaOdfModel(QballBaseModel):
891    """Implementation of Constant Solid Angle reconstruction method.
892
893    References
894    ----------
895    .. [1] Aganj, I., et al. 2009. ODF Reconstruction in Q-Ball Imaging With
896           Solid Angle Consideration.
897    """
898    min = .001
899    max = .999
900    _n0_const = .5 / np.sqrt(np.pi)
901
902    def _set_fit_matrix(self, B, L, F, smooth):
903        """The fit matrix, is used by fit_coefficients to return the
904        coefficients of the odf"""
905        invB = smooth_pinv(B, np.sqrt(smooth) * L)
906        L = L[:, None]
907        F = F[:, None]
908        self._fit_matrix = (F * L) / (8 * np.pi) * invB
909
910    def _get_shm_coef(self, data, mask=None):
911        """Returns the coefficients of the model"""
912        data = data[..., self._where_dwi]
913        data = data.clip(self.min, self.max)
914        loglog_data = np.log(-np.log(data))
915        sh_coef = np.dot(loglog_data, self._fit_matrix.T)
916        sh_coef[..., 0] = self._n0_const
917        return sh_coef
918
919
920class OpdtModel(QballBaseModel):
921    """Implementation of Orientation Probability Density Transform
922    reconstruction method.
923
924    References
925    ----------
926    .. [1] Tristan-Vega, A., et al. 2010. A new methodology for estimation of
927           fiber populations in white matter of the brain with Funk-Radon
928           transform.
929    .. [2] Tristan-Vega, A., et al. 2009. Estimation of fiber orientation
930           probability density functions in high angular resolution diffusion
931           imaging.
932    """
933    def _set_fit_matrix(self, B, L, F, smooth):
934        invB = smooth_pinv(B, np.sqrt(smooth) * L)
935        L = L[:, None]
936        F = F[:, None]
937        delta_b = F * L * invB
938        delta_q = 4 * F * invB
939        self._fit_matrix = delta_b, delta_q
940
941    def _get_shm_coef(self, data, mask=None):
942        """Returns the coefficients of the model"""
943        delta_b, delta_q = self._fit_matrix
944        return _slowadc_formula(data[..., self._where_dwi], delta_b, delta_q)
945
946
947def _slowadc_formula(data, delta_b, delta_q):
948    """formula used in SlowAdcOpdfModel"""
949    logd = -np.log(data)
950    return (np.dot(logd * (1.5 - logd) * data, delta_q.T)
951            - np.dot(data, delta_b.T))
952
953
954class QballModel(QballBaseModel):
955    """Implementation of regularized Qball reconstruction method.
956
957    References
958    ----------
959    .. [1] Descoteaux, M., et al. 2007. Regularized, fast, and robust
960           analytical Q-ball imaging.
961    """
962
963    def _set_fit_matrix(self, B, L, F, smooth):
964        invB = smooth_pinv(B, np.sqrt(smooth) * L)
965        F = F[:, None]
966        self._fit_matrix = F * invB
967
968    def _get_shm_coef(self, data, mask=None):
969        """Returns the coefficients of the model"""
970        return np.dot(data[..., self._where_dwi], self._fit_matrix.T)
971
972
973def normalize_data(data, where_b0, min_signal=1e-5, out=None):
974    """Normalizes the data with respect to the mean b0
975    """
976    if out is None:
977        out = np.array(data, dtype='float32', copy=True)
978    else:
979        if out.dtype.kind != 'f':
980            raise ValueError("out must be floating point")
981        out[:] = data
982
983    out.clip(min_signal, out=out)
984    b0 = out[..., where_b0].mean(-1)
985    out /= b0[..., None]
986    return out
987
988
989def hat(B):
990    """Returns the hat matrix for the design matrix B
991    """
992
993    U, S, V = np.linalg.svd(B, False)
994    H = np.dot(U, U.T)
995    return H
996
997
998def lcr_matrix(H):
999    """Returns a matrix for computing leveraged, centered residuals from data
1000
1001    if r = (d-Hd), the leveraged centered residuals are lcr = (r/l)-mean(r/l)
1002    ruturns the matrix R, such lcr = Rd
1003
1004    """
1005    if H.ndim != 2 or H.shape[0] != H.shape[1]:
1006        raise ValueError('H should be a square matrix')
1007
1008    leverages = np.sqrt(1 - H.diagonal(), where=H.diagonal() <= 1)
1009    leverages = leverages[:, None]
1010    R = (np.eye(len(H)) - H) / leverages
1011    return R - R.mean(0)
1012
1013
1014def bootstrap_data_array(data, H, R, permute=None):
1015    """Applies the Residual Bootstraps to the data given H and R
1016
1017    data must be normalized, ie 0 < data <= 1
1018
1019    This function, and the bootstrap_data_voxel function, calculate
1020    residual-bootsrap samples given a Hat matrix and a Residual matrix. These
1021    samples can be used for non-parametric statistics or for bootstrap
1022    probabilistic tractography:
1023
1024    References
1025    ----------
1026    .. [1] J. I. Berman, et al., "Probabilistic streamline q-ball tractography
1027           using the residual bootstrap" 2008.
1028    .. [2] HA Haroon, et al., "Using the model-based residual bootstrap to
1029           quantify uncertainty in fiber orientations from Q-ball analysis"
1030           2009.
1031    .. [3] B. Jeurissen, et al., "Probabilistic Fiber Tracking Using the
1032           Residual Bootstrap with Constrained Spherical Deconvolution" 2011.
1033    """
1034
1035    if permute is None:
1036        permute = randint(data.shape[-1], size=data.shape[-1])
1037    assert R.shape == H.shape
1038    assert len(permute) == R.shape[-1]
1039    R = R[permute]
1040    data = np.dot(data, (H + R).T)
1041    return data
1042
1043
1044def bootstrap_data_voxel(data, H, R, permute=None):
1045    """Like bootstrap_data_array but faster when for a single voxel
1046
1047    data must be 1d and normalized
1048    """
1049    if permute is None:
1050        permute = randint(data.shape[-1], size=data.shape[-1])
1051    r = np.dot(data, R.T)
1052    boot_data = np.dot(data, H.T)
1053    boot_data += r[permute]
1054    return boot_data
1055
1056
1057class ResidualBootstrapWrapper(object):
1058    """Returns a residual bootstrap sample of the signal_object when indexed
1059
1060    Wraps a signal_object, this signal object can be an interpolator. When
1061    indexed, the the wrapper indexes the signal_object to get the signal.
1062    There wrapper than samples the residual boostrap distribution of signal and
1063    returns that sample.
1064    """
1065    def __init__(self, signal_object, B, where_dwi, min_signal=1e-5):
1066        """Builds a ResidualBootstrapWapper
1067
1068        Given some linear model described by B, the design matrix, and a
1069        signal_object, returns an object which can sample the residual
1070        bootstrap distribution of the signal. We assume that the signals are
1071        normalized so we clip the bootsrap samples to be between `min_signal`
1072        and 1.
1073
1074        Parameters
1075        ----------
1076        signal_object : some object that can be indexed
1077            This object should return diffusion weighted signals when indexed.
1078        B : ndarray, ndim=2
1079            The design matrix of the spherical harmonics model used to fit the
1080            data. This is the model that will be used to compute the residuals
1081            and sample the residual bootstrap distribution
1082        where_dwi :
1083            indexing object to find diffusion weighted signals from signal
1084        min_signal : float
1085            The lowest allowable signal.
1086        """
1087        self._signal_object = signal_object
1088        self._H = hat(B)
1089        self._R = lcr_matrix(self._H)
1090        self._min_signal = min_signal
1091        self._where_dwi = where_dwi
1092        self.data = signal_object.data
1093        self.voxel_size = signal_object.voxel_size
1094
1095    def __getitem__(self, index):
1096        """Indexes self._signal_object and bootstraps the result"""
1097        signal = self._signal_object[index].copy()
1098        dwi_signal = signal[self._where_dwi]
1099        boot_signal = bootstrap_data_voxel(dwi_signal, self._H, self._R)
1100        boot_signal.clip(self._min_signal, 1., out=boot_signal)
1101        signal[self._where_dwi] = boot_signal
1102        return signal
1103
1104
1105def sf_to_sh(sf, sphere, sh_order=4, basis_type=None, full_basis=False,
1106             legacy=True, smooth=0.0):
1107    """Spherical function to spherical harmonics (SH).
1108
1109    Parameters
1110    ----------
1111    sf : ndarray
1112        Values of a function on the given ``sphere``.
1113    sphere : Sphere
1114        The points on which the sf is defined.
1115    sh_order : int, optional
1116        Maximum SH order in the SH fit.  For ``sh_order``, there will be
1117        ``(sh_order + 1) * (sh_order + 2) / 2`` SH coefficients for a symmetric
1118        basis and ``(sh_order + 1) * (sh_order + 1)`` coefficients for a full
1119        SH basis.
1120    basis_type : {None, 'tournier07', 'descoteaux07'}, optional
1121        ``None`` for the default DIPY basis,
1122        ``tournier07`` for the Tournier 2007 [2]_[3]_ basis,
1123        ``descoteaux07`` for the Descoteaux 2007 [1]_ basis,
1124        (``None`` defaults to ``descoteaux07``).
1125    full_basis: bool, optional
1126        True for using a SH basis containing even and odd order SH functions.
1127        False for using a SH basis consisting only of even order SH functions.
1128    legacy: bool, optional
1129        True to use a legacy basis definition for backward compatibility
1130        with previous ``tournier07`` and ``descoteaux07`` implementations.
1131    smooth : float, optional
1132        Lambda-regularization in the SH fit.
1133
1134    Returns
1135    -------
1136    sh : ndarray
1137        SH coefficients representing the input function.
1138
1139    References
1140    ----------
1141    .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
1142           Regularized, Fast, and Robust Analytical Q-ball Imaging.
1143           Magn. Reson. Med. 2007;58:497-510.
1144    .. [2] Tournier J.D., Calamante F. and Connelly A. Robust determination
1145           of the fibre orientation distribution in diffusion MRI:
1146           Non-negativity constrained super-resolved spherical deconvolution.
1147           NeuroImage. 2007;35(4):1459-1472.
1148    .. [3] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
1149           Pietsch M, et al. MRtrix3: A fast, flexible and open software
1150           framework for medical image processing and visualisation.
1151           NeuroImage. 2019 Nov 15;202:116-137.
1152    """
1153
1154    sph_harm_basis = sph_harm_lookup.get(basis_type)
1155
1156    if sph_harm_basis is None:
1157        raise ValueError("Invalid basis name.")
1158    B, m, n = sph_harm_basis(sh_order, sphere.theta, sphere.phi,
1159                             full_basis=full_basis,
1160                             legacy=legacy)
1161
1162    L = -n * (n + 1)
1163    invB = smooth_pinv(B, np.sqrt(smooth) * L)
1164    sh = np.dot(sf, invB.T)
1165
1166    return sh
1167
1168
1169def sh_to_sf(sh, sphere, sh_order=4, basis_type=None,
1170             full_basis=False, legacy=True):
1171    """Spherical harmonics (SH) to spherical function (SF).
1172
1173    Parameters
1174    ----------
1175    sh : ndarray
1176        SH coefficients representing a spherical function.
1177    sphere : Sphere
1178        The points on which to sample the spherical function.
1179    sh_order : int, optional
1180        Maximum SH order in the SH fit.  For ``sh_order``, there will be
1181        ``(sh_order + 1) * (sh_order + 2) / 2`` SH coefficients for a symmetric
1182        basis and ``(sh_order + 1) * (sh_order + 1)`` coefficients for a full
1183        SH basis.
1184    basis_type : {None, 'tournier07', 'descoteaux07'}, optional
1185        ``None`` for the default DIPY basis,
1186        ``tournier07`` for the Tournier 2007 [2]_[3]_ basis,
1187        ``descoteaux07`` for the Descoteaux 2007 [1]_ basis,
1188        (``None`` defaults to ``descoteaux07``).
1189    full_basis: bool, optional
1190        True to use a SH basis containing even and odd order SH functions.
1191        Else, use a SH basis consisting only of even order SH functions.
1192    legacy: bool, optional
1193        True to use a legacy basis definition for backward compatibility
1194        with previous ``tournier07`` and ``descoteaux07`` implementations.
1195
1196    Returns
1197    -------
1198    sf : ndarray
1199         Spherical function values on the ``sphere``.
1200
1201    References
1202    ----------
1203    .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
1204           Regularized, Fast, and Robust Analytical Q-ball Imaging.
1205           Magn. Reson. Med. 2007;58:497-510.
1206    .. [2] Tournier J.D., Calamante F. and Connelly A. Robust determination
1207           of the fibre orientation distribution in diffusion MRI:
1208           Non-negativity constrained super-resolved spherical deconvolution.
1209           NeuroImage. 2007;35(4):1459-1472.
1210    .. [3] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
1211           Pietsch M, et al. MRtrix3: A fast, flexible and open software
1212           framework for medical image processing and visualisation.
1213           NeuroImage. 2019 Nov 15;202:116-137.
1214    """
1215
1216    sph_harm_basis = sph_harm_lookup.get(basis_type)
1217
1218    if sph_harm_basis is None:
1219        raise ValueError("Invalid basis name.")
1220    B, m, n = sph_harm_basis(sh_order, sphere.theta, sphere.phi,
1221                             full_basis=full_basis,
1222                             legacy=legacy)
1223
1224    sf = np.dot(sh, B.T)
1225
1226    return sf
1227
1228
1229def sh_to_sf_matrix(sphere, sh_order=4, basis_type=None, full_basis=False,
1230                    legacy=True, return_inv=True, smooth=0):
1231    """ Matrix that transforms Spherical harmonics (SH) to spherical
1232    function (SF).
1233
1234    Parameters
1235    ----------
1236    sphere : Sphere
1237        The points on which to sample the spherical function.
1238    sh_order : int, optional
1239        Maximum SH order in the SH fit.  For ``sh_order``, there will be
1240        ``(sh_order + 1) * (sh_order + 2) / 2`` SH coefficients for a symmetric
1241        basis and ``(sh_order + 1) * (sh_order + 1)`` coefficients for a full
1242        SH basis.
1243    basis_type : {None, 'tournier07', 'descoteaux07'}, optional
1244        ``None`` for the default DIPY basis,
1245        ``tournier07`` for the Tournier 2007 [2]_[3]_ basis,
1246        ``descoteaux07`` for the Descoteaux 2007 [1]_ basis,
1247        (``None`` defaults to ``descoteaux07``).
1248    full_basis: bool, optional
1249        If True, uses a SH basis containing even and odd order SH functions.
1250        Else, uses a SH basis consisting only of even order SH functions.
1251    legacy: bool, optional
1252        True to use a legacy basis definition for backward compatibility
1253        with previous ``tournier07`` and ``descoteaux07`` implementations.
1254    return_inv : bool, optional
1255        If True then the inverse of the matrix is also returned.
1256    smooth : float, optional
1257        Lambda-regularization in the SH fit.
1258
1259    Returns
1260    -------
1261    B : ndarray
1262        Matrix that transforms spherical harmonics to spherical function
1263        ``sf = np.dot(sh, B)``.
1264    invB : ndarray
1265        Inverse of B.
1266
1267    References
1268    ----------
1269    .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
1270           Regularized, Fast, and Robust Analytical Q-ball Imaging.
1271           Magn. Reson. Med. 2007;58:497-510.
1272    .. [2] Tournier J.D., Calamante F. and Connelly A. Robust determination
1273           of the fibre orientation distribution in diffusion MRI:
1274           Non-negativity constrained super-resolved spherical deconvolution.
1275           NeuroImage. 2007;35(4):1459-1472.
1276    .. [3] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
1277           Pietsch M, et al. MRtrix3: A fast, flexible and open software
1278           framework for medical image processing and visualisation.
1279           NeuroImage. 2019 Nov 15;202:116-137.
1280
1281    """
1282
1283    sph_harm_basis = sph_harm_lookup.get(basis_type)
1284
1285    if sph_harm_basis is None:
1286        raise ValueError("Invalid basis name.")
1287    B, m, n = sph_harm_basis(sh_order, sphere.theta, sphere.phi,
1288                             full_basis=full_basis, legacy=legacy)
1289
1290    if return_inv:
1291        L = -n * (n + 1)
1292        invB = smooth_pinv(B, np.sqrt(smooth) * L)
1293        return B.T, invB.T
1294
1295    return B.T
1296
1297
1298def calculate_max_order(n_coeffs, full_basis=False):
1299    r"""Calculate the maximal harmonic order, given that you know the
1300    number of parameters that were estimated.
1301
1302    Parameters
1303    ----------
1304    n_coeffs : int
1305        The number of SH coefficients
1306    full_basis: bool, optional
1307        True if the used SH basis contains even and odd order SH functions.
1308        False if the SH basis consists only of even order SH functions.
1309
1310    Returns
1311    -------
1312    L : int
1313        The maximal SH order, given the number of coefficients
1314
1315    Notes
1316    -----
1317    The calculation in this function for the symmetric SH basis
1318    proceeds according to the following logic:
1319    .. math::
1320        n = \frac{1}{2} (L+1) (L+2)
1321        \rarrow 2n = L^2 + 3L + 2
1322        \rarrow L^2 + 3L + 2 - 2n = 0
1323        \rarrow L^2 + 3L + 2(1-n) = 0
1324        \rarrow L_{1,2} = \frac{-3 \pm \sqrt{9 - 8 (1-n)}}{2}
1325        \rarrow L{1,2} = \frac{-3 \pm \sqrt{1 + 8n}}{2}
1326
1327    Finally, the positive value is chosen between the two options.
1328
1329    For a full SH basis, the calcultation consists in solving the equation
1330    $n = (L + 1)^2$ for $L$, which gives $L = sqrt(n) - 1$.
1331    """
1332
1333    # L2 is negative for all positive values of n_coeffs, so we don't
1334    # bother even computing it:
1335    # L2 = (-3 - np.sqrt(1 + 8 * n_coeffs)) / 2
1336    # L1 is always the larger value, so we go with that:
1337    if full_basis:
1338        L1 = np.sqrt(n_coeffs) - 1
1339        if L1.is_integer():
1340            return int(L1)
1341    else:
1342        L1 = (-3 + np.sqrt(1 + 8 * n_coeffs)) / 2.0
1343        # Check that it is a whole even number:
1344        if L1.is_integer() and not np.mod(L1, 2):
1345            return int(L1)
1346
1347    # Otherwise, the input didn't make sense:
1348    raise ValueError("The input to ``calculate_max_order`` was ",
1349                     "%s, but that is not a valid number" % n_coeffs,
1350                     "of coefficients for a spherical harmonics ",
1351                     "basis set.")
1352
1353
1354def anisotropic_power(sh_coeffs, norm_factor=0.00001, power=2,
1355                      non_negative=True):
1356    r"""Calculate anisotropic power map with a given SH coefficient matrix.
1357
1358    Parameters
1359    ----------
1360    sh_coeffs : ndarray
1361        A ndarray where the last dimension is the
1362        SH coefficients estimates for that voxel.
1363    norm_factor: float, optional
1364        The value to normalize the ap values.
1365    power : int, optional
1366        The degree to which power maps are calculated.
1367    non_negative: bool, optional
1368        Whether to rectify the resulting map to be non-negative.
1369
1370    Returns
1371    -------
1372    log_ap : ndarray
1373        The log of the resulting power image.
1374
1375    Notes
1376    -----
1377    Calculate AP image based on a IxJxKxC SH coefficient matrix based on the
1378    equation:
1379    .. math::
1380        AP = \sum_{l=2,4,6,...}{\frac{1}{2l+1} \sum_{m=-l}^l{|a_{l,m}|^n}}
1381
1382    Where the last dimension, C, is made of a flattened array of $l$x$m$
1383    coefficients, where $l$ are the SH orders, and $m = 2l+1$,
1384    So l=1 has 1 coeffecient, l=2 has 5, ... l=8 has 17 and so on.
1385    A l=2 SH coefficient matrix will then be composed of a IxJxKx6 volume.
1386    The power, $n$ is usually set to $n=2$.
1387
1388    The final AP image is then shifted by -log(norm_factor), to be strictly
1389    non-negative. Remaining values < 0 are discarded (set to 0), per default,
1390    and this option is controlled through the `non_negative` keyword argument.
1391
1392    References
1393    ----------
1394    .. [1]  Dell'Acqua, F., Lacerda, L., Catani, M., Simmons, A., 2014.
1395            Anisotropic Power Maps: A diffusion contrast to reveal low
1396            anisotropy tissues from HARDI data,
1397            in: Proceedings of International Society for Magnetic Resonance in
1398            Medicine. Milan, Italy.
1399
1400    """
1401    dim = sh_coeffs.shape[:-1]
1402    n_coeffs = sh_coeffs.shape[-1]
1403    max_order = calculate_max_order(n_coeffs)
1404    ap = np.zeros(dim)
1405    n_start = 1
1406    for L in range(2, max_order + 2, 2):
1407        n_stop = n_start + (2 * L + 1)
1408        ap_i = np.mean(np.abs(sh_coeffs[..., n_start:n_stop]) ** power, -1)
1409        ap += ap_i
1410        n_start = n_stop
1411
1412    # Shift the map to be mostly non-negative,
1413    # only applying the log operation to positive elements
1414    # to avoid getting numpy warnings on log(0).
1415    # It is impossible to get ap values smaller than 0.
1416    # Also avoids getting voxels with -inf when non_negative=False.
1417
1418    if ap.ndim < 1:
1419        # For the off chance we have a scalar on our hands
1420        ap = np.reshape(ap, (1, ))
1421    log_ap = np.zeros_like(ap)
1422    log_ap[ap > 0] = np.log(ap[ap > 0]) - np.log(norm_factor)
1423
1424    # Deal with residual negative values:
1425    if non_negative:
1426        if isinstance(log_ap, np.ndarray):
1427            # zero all values < 0
1428            log_ap[log_ap < 0] = 0
1429        else:
1430            # assume this is a singleton float (input was 1D):
1431            if log_ap < 0:
1432                return 0
1433    return log_ap
1434
1435
1436def convert_sh_to_full_basis(sh_coeffs):
1437    """Given an array of SH coeffs from a symmetric basis, returns the
1438    coefficients for the full SH basis by filling odd order SH coefficients
1439    with zeros
1440
1441    Parameters
1442    ----------
1443    sh_coeffs: ndarray
1444        A ndarray where the last dimension is the
1445        SH coefficients estimates for that voxel.
1446
1447    Returns
1448    -------
1449    full_sh_coeffs: ndarray
1450        A ndarray where the last dimension is the
1451        SH coefficients estimates for that voxel in
1452        a full SH basis.
1453    """
1454    sh_order = calculate_max_order(sh_coeffs.shape[-1])
1455    _, n = sph_harm_ind_list(sh_order, full_basis=True)
1456
1457    full_sh_coeffs =\
1458        np.zeros(np.append(sh_coeffs.shape[:-1], [n.size]).astype(int))
1459    mask = np.mod(n, 2) == 0
1460
1461    full_sh_coeffs[..., mask] = sh_coeffs
1462    return full_sh_coeffs
1463
1464
1465def convert_sh_from_legacy(sh_coeffs, sh_basis, full_basis=False):
1466    """Convert SH coefficients in legacy SH basis to SH coefficients
1467    of the new SH basis for ``descoteaux07`` [1]_ or ``tournier07`` [2]_[3]_
1468    bases.
1469
1470    Parameters
1471    ----------
1472    sh_coeffs: ndarray
1473        A ndarray where the last dimension is the
1474        SH coefficients estimates for that voxel.
1475    sh_basis: {'descoteaux07', 'tournier07'}
1476        ``tournier07`` for the Tournier 2007 [2]_[3]_ basis,
1477        ``descoteaux07`` for the Descoteaux 2007 [1]_ basis.
1478    full_basis: bool, optional
1479        True if the input SH basis includes both even and odd
1480        order SH functions, else False.
1481
1482    Returns
1483    -------
1484    out_sh_coeffs: ndarray
1485        The array of coefficients expressed in the new SH basis.
1486
1487    References
1488    ----------
1489    .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
1490           Regularized, Fast, and Robust Analytical Q-ball Imaging.
1491           Magn. Reson. Med. 2007;58:497-510.
1492    .. [2] Tournier J.D., Calamante F. and Connelly A. Robust determination
1493           of the fibre orientation distribution in diffusion MRI:
1494           Non-negativity constrained super-resolved spherical deconvolution.
1495           NeuroImage. 2007;35(4):1459-1472.
1496    .. [3] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
1497           Pietsch M, et al. MRtrix3: A fast, flexible and open software
1498           framework for medical image processing and visualisation.
1499           NeuroImage. 2019 Nov 15;202:116-137.
1500    """
1501    sh_order = calculate_max_order(sh_coeffs.shape[-1],
1502                                   full_basis=full_basis)
1503
1504    m, n = sph_harm_ind_list(sh_order, full_basis=full_basis)
1505
1506    if sh_basis == 'descoteaux07':
1507        out_sh_coeffs = sh_coeffs * np.where(m < 0, (-1.)**m, 1.)
1508    elif sh_basis == 'tournier07':
1509        out_sh_coeffs = sh_coeffs * np.where(m == 0, 1., 1./np.sqrt(2))
1510    else:
1511        raise ValueError("Invalid basis name.")
1512
1513    return out_sh_coeffs
1514
1515
1516def convert_sh_to_legacy(sh_coeffs, sh_basis, full_basis=False):
1517    """Convert SH coefficients in new SH basis to SH coefficients for
1518    the legacy SH basis for ``descoteaux07`` [1]_ or ``tournier07`` [2]_[3]_
1519    bases.
1520
1521    Parameters
1522    ----------
1523    sh_coeffs: ndarray
1524        A ndarray where the last dimension is the
1525        SH coefficients estimates for that voxel.
1526    sh_basis: {'descoteaux07', 'tournier07'}
1527        ``tournier07`` for the Tournier 2007 [2]_[3]_ basis,
1528        ``descoteaux07`` for the Descoteaux 2007 [1]_ basis.
1529    full_basis: bool, optional
1530        True if the input SH basis includes both even and odd
1531        order SH functions.
1532
1533    Returns
1534    -------
1535    out_sh_coeffs: ndarray
1536        The array of coefficients expressed in the legacy SH basis.
1537
1538    References
1539    ----------
1540    .. [1] Descoteaux, M., Angelino, E., Fitzgibbons, S. and Deriche, R.
1541           Regularized, Fast, and Robust Analytical Q-ball Imaging.
1542           Magn. Reson. Med. 2007;58:497-510.
1543    .. [2] Tournier J.D., Calamante F. and Connelly A. Robust determination
1544           of the fibre orientation distribution in diffusion MRI:
1545           Non-negativity constrained super-resolved spherical deconvolution.
1546           NeuroImage. 2007;35(4):1459-1472.
1547    .. [3] Tournier J-D, Smith R, Raffelt D, Tabbara R, Dhollander T,
1548           Pietsch M, et al. MRtrix3: A fast, flexible and open software
1549           framework for medical image processing and visualisation.
1550           NeuroImage. 2019 Nov 15;202:116-137.
1551    """
1552    sh_order = calculate_max_order(sh_coeffs.shape[-1],
1553                                   full_basis=full_basis)
1554
1555    m, n = sph_harm_ind_list(sh_order, full_basis=full_basis)
1556
1557    if sh_basis == 'descoteaux07':
1558        out_sh_coeffs = sh_coeffs * np.where(m < 0, (-1.)**m, 1.)
1559    elif sh_basis == 'tournier07':
1560        out_sh_coeffs = sh_coeffs * np.where(m == 0, 1., np.sqrt(2))
1561    else:
1562        raise ValueError("Invalid basis name.")
1563
1564    return out_sh_coeffs
1565