1#!/usr/bin/python
2""" Classes and functions for fitting tensors """
3
4import warnings
5
6import functools
7
8import numpy as np
9
10import scipy.optimize as opt
11
12from dipy.utils.arrfuncs import pinv
13from dipy.data import get_sphere
14from dipy.core.gradients import gradient_table
15from dipy.core.geometry import vector_norm
16from dipy.reconst.vec_val_sum import vec_val_vect
17from dipy.core.onetime import auto_attr
18from dipy.reconst.base import ReconstModel
19
20
21MIN_POSITIVE_SIGNAL = 0.0001
22
23
24def _roll_evals(evals, axis=-1):
25    """Check evals shape.
26
27    Helper function to check that the evals provided to functions calculating
28    tensor statistics have the right shape
29
30    Parameters
31    ----------
32    evals : array-like
33        Eigenvalues of a diffusion tensor. shape should be (...,3).
34
35    axis : int
36        The axis of the array which contains the 3 eigenvals. Default: -1
37
38    Returns
39    -------
40    evals : array-like
41        Eigenvalues of a diffusion tensor, rolled so that the 3 eigenvals are
42        the last axis.
43
44    """
45    if evals.shape[-1] != 3:
46        msg = "Expecting 3 eigenvalues, got {}".format(evals.shape[-1])
47        raise ValueError(msg)
48
49    evals = np.rollaxis(evals, axis)
50
51    return evals
52
53
54def fractional_anisotropy(evals, axis=-1):
55    r"""Return Fractional anisotropy (FA) of a diffusion tensor.
56
57    Parameters
58    ----------
59    evals : array-like
60        Eigenvalues of a diffusion tensor.
61    axis : int
62        Axis of `evals` which contains 3 eigenvalues.
63
64    Returns
65    -------
66    fa : array
67        Calculated FA. Range is 0 <= FA <= 1.
68
69    Notes
70    -----
71    FA is calculated using the following equation:
72
73    .. math::
74
75        FA = \sqrt{\frac{1}{2}\frac{(\lambda_1-\lambda_2)^2+(\lambda_1-
76                    \lambda_3)^2+(\lambda_2-\lambda_3)^2}{\lambda_1^2+
77                    \lambda_2^2+\lambda_3^2}}
78
79    """
80    evals = _roll_evals(evals, axis)
81    # Make sure not to get nans
82    all_zero = (evals == 0).all(axis=0)
83    ev1, ev2, ev3 = evals
84    fa = np.sqrt(0.5 * ((ev1 - ev2) ** 2 +
85                        (ev2 - ev3) ** 2 +
86                        (ev3 - ev1) ** 2) /
87                 ((evals * evals).sum(0) + all_zero))
88
89    return fa
90
91
92def geodesic_anisotropy(evals, axis=-1):
93    r"""
94    Geodesic anisotropy (GA) of a diffusion tensor.
95
96    Parameters
97    ----------
98    evals : array-like
99        Eigenvalues of a diffusion tensor.
100    axis : int
101        Axis of `evals` which contains 3 eigenvalues.
102
103    Returns
104    -------
105    ga : array
106        Calculated GA. In the range 0 to +infinity
107
108    Notes
109    -----
110    GA is calculated using the following equation given in [1]_:
111
112    .. math::
113
114        GA = \sqrt{\sum_{i=1}^3
115        \log^2{\left ( \lambda_i/<\mathbf{D}> \right )}},
116        \quad \textrm{where} \quad <\mathbf{D}> =
117        (\lambda_1\lambda_2\lambda_3)^{1/3}
118
119    Note that the notation, $<D>$, is often used as the mean diffusivity (MD)
120    of the diffusion tensor and can lead to confusions in the literature
121    (see [1]_ versus [2]_ versus [3]_ for example). Reference [2]_ defines
122    geodesic anisotropy (GA) with $<D>$ as the MD in the denominator of the
123    sum. This is wrong. The original paper [1]_ defines GA with
124    $<D> = det(D)^{1/3}$, as the isotropic part of the distance. This might be
125    an explanation for the confusion. The isotropic part of the diffusion
126    tensor in Euclidean space is the MD whereas the isotropic part of the
127    tensor in log-Euclidean space is $det(D)^{1/3}$. The Appendix of [1]_ and
128    log-Euclidean derivations from [3]_ are clear on this. Hence, all that to
129    say that $<D> = det(D)^{1/3}$ here for the GA definition and not MD.
130
131    References
132    ----------
133
134    .. [1] P. G. Batchelor, M. Moakher, D. Atkinson, F. Calamante,
135        A. Connelly, "A rigorous framework for diffusion tensor calculus",
136        Magnetic Resonance in Medicine, vol. 53, pp. 221-225, 2005.
137
138    .. [2] M. M. Correia, V. F. Newcombe, G.B. Williams.
139        "Contrast-to-noise ratios for indices of anisotropy obtained from
140        diffusion MRI: a study with standard clinical b-values at 3T".
141        NeuroImage, vol. 57, pp. 1103-1115, 2011.
142
143    .. [3] A. D. Lee, etal, P. M. Thompson.
144        "Comparison of fractional and geodesic anisotropy in diffusion tensor
145        images of 90 monozygotic and dizygotic twins". 5th IEEE International
146        Symposium on Biomedical Imaging (ISBI), pp. 943-946, May 2008.
147
148    .. [4] V. Arsigny, P. Fillard, X. Pennec, N. Ayache.
149        "Log-Euclidean metrics for fast and simple calculus on diffusion
150        tensors." Magnetic Resonance in Medecine, vol 56, pp. 411-421, 2006.
151
152    """
153
154    evals = _roll_evals(evals, axis)
155    ev1, ev2, ev3 = evals
156
157    log1 = np.zeros(ev1.shape)
158    log2 = np.zeros(ev1.shape)
159    log3 = np.zeros(ev1.shape)
160    idx = np.nonzero(ev1)
161
162    # this is the definition in [1]_
163    detD = np.power(ev1 * ev2 * ev3, 1 / 3.)
164    log1[idx] = np.log(ev1[idx] / detD[idx])
165    log2[idx] = np.log(ev2[idx] / detD[idx])
166    log3[idx] = np.log(ev3[idx] / detD[idx])
167
168    ga = np.sqrt(log1 ** 2 + log2 ** 2 + log3 ** 2)
169
170    return ga
171
172
173def mean_diffusivity(evals, axis=-1):
174    r"""
175    Mean Diffusivity (MD) of a diffusion tensor.
176
177    Parameters
178    ----------
179    evals : array-like
180        Eigenvalues of a diffusion tensor.
181    axis : int
182        Axis of `evals` which contains 3 eigenvalues.
183
184    Returns
185    -------
186    md : array
187        Calculated MD.
188
189    Notes
190    -----
191    MD is calculated with the following equation:
192
193    .. math::
194
195        MD = \frac{\lambda_1 + \lambda_2 + \lambda_3}{3}
196
197    """
198    evals = _roll_evals(evals, axis)
199    return evals.mean(0)
200
201
202def axial_diffusivity(evals, axis=-1):
203    r"""
204    Axial Diffusivity (AD) of a diffusion tensor.
205    Also called parallel diffusivity.
206
207    Parameters
208    ----------
209    evals : array-like
210        Eigenvalues of a diffusion tensor, must be sorted in descending order
211        along `axis`.
212    axis : int
213        Axis of `evals` which contains 3 eigenvalues.
214
215    Returns
216    -------
217    ad : array
218        Calculated AD.
219
220    Notes
221    -----
222    AD is calculated with the following equation:
223
224    .. math::
225
226        AD = \lambda_1
227
228    """
229    evals = _roll_evals(evals, axis)
230    ev1, ev2, ev3 = evals
231    return ev1
232
233
234def radial_diffusivity(evals, axis=-1):
235    r"""
236    Radial Diffusivity (RD) of a diffusion tensor.
237    Also called perpendicular diffusivity.
238
239    Parameters
240    ----------
241    evals : array-like
242        Eigenvalues of a diffusion tensor, must be sorted in descending order
243        along `axis`.
244    axis : int
245        Axis of `evals` which contains 3 eigenvalues.
246
247    Returns
248    -------
249    rd : array
250        Calculated RD.
251
252    Notes
253    -----
254    RD is calculated with the following equation:
255
256    .. math::
257
258        RD = \frac{\lambda_2 + \lambda_3}{2}
259
260    """
261    evals = _roll_evals(evals, axis)
262    return evals[1:].mean(0)
263
264
265def trace(evals, axis=-1):
266    r"""
267    Trace of a diffusion tensor.
268
269    Parameters
270    ----------
271    evals : array-like
272        Eigenvalues of a diffusion tensor.
273    axis : int
274        Axis of `evals` which contains 3 eigenvalues.
275
276    Returns
277    -------
278    trace : array
279        Calculated trace of the diffusion tensor.
280
281    Notes
282    -----
283    Trace is calculated with the following equation:
284
285    .. math::
286
287        Trace = \lambda_1 + \lambda_2 + \lambda_3
288
289    """
290    evals = _roll_evals(evals, axis)
291    return evals.sum(0)
292
293
294def color_fa(fa, evecs):
295    r""" Color fractional anisotropy of diffusion tensor
296
297    Parameters
298    ----------
299    fa : array-like
300        Array of the fractional anisotropy (can be 1D, 2D or 3D)
301
302    evecs : array-like
303        eigen vectors from the tensor model
304
305    Returns
306    -------
307    rgb : Array with 3 channels for each color as the last dimension.
308        Colormap of the FA with red for the x value, y for the green
309        value and z for the blue value.
310
311    Notes
312    -----
313
314    It is computed from the clipped FA between 0 and 1 using the following
315    formula
316
317    .. math::
318
319        rgb = abs(max(\vec{e})) \times fa
320    """
321
322    if (fa.shape != evecs[..., 0, 0].shape) or ((3, 3) != evecs.shape[-2:]):
323        raise ValueError("Wrong number of dimensions for evecs")
324
325    return np.abs(evecs[..., 0]) * np.clip(fa, 0, 1)[..., None]
326
327
328# The following are used to calculate the tensor mode:
329def determinant(q_form):
330    """
331    The determinant of a tensor, given in quadratic form
332
333    Parameters
334    ----------
335    q_form : ndarray
336        The quadratic form of a tensor, or an array with quadratic forms of
337        tensors. Should be of shape (x, y, z, 3, 3) or (n, 3, 3) or (3, 3).
338
339    Returns
340    -------
341    det : array
342        The determinant of the tensor in each spatial coordinate
343    """
344
345    # Following the conventions used here:
346    # http://en.wikipedia.org/wiki/Determinant
347    aei = q_form[..., 0, 0] * q_form[..., 1, 1] * q_form[..., 2, 2]
348    bfg = q_form[..., 0, 1] * q_form[..., 1, 2] * q_form[..., 2, 0]
349    cdh = q_form[..., 0, 2] * q_form[..., 1, 0] * q_form[..., 2, 1]
350    ceg = q_form[..., 0, 2] * q_form[..., 1, 1] * q_form[..., 2, 0]
351    bdi = q_form[..., 0, 1] * q_form[..., 1, 0] * q_form[..., 2, 2]
352    afh = q_form[..., 0, 0] * q_form[..., 1, 2] * q_form[..., 2, 1]
353    return aei + bfg + cdh - ceg - bdi - afh
354
355
356def isotropic(q_form):
357    r"""
358    Calculate the isotropic part of the tensor [1]_.
359
360    Parameters
361    ----------
362    q_form : ndarray
363        The quadratic form of a tensor, or an array with quadratic forms of
364        tensors. Should be of shape (x,y,z,3,3) or (n, 3, 3) or (3,3).
365
366    Returns
367    -------
368    A_hat: ndarray
369        The isotropic part of the tensor in each spatial coordinate
370
371    Notes
372    -----
373    The isotropic part of a tensor is defined as (equations 3-5 of [1]_):
374
375    .. math ::
376        \bar{A} = \frac{1}{2} tr(A) I
377
378    References
379    ----------
380    .. [1] Daniel B. Ennis and G. Kindlmann, "Orthogonal Tensor
381        Invariants and the Analysis of Diffusion Tensor Magnetic Resonance
382        Images", Magnetic Resonance in Medicine, vol. 55, no. 1, pp. 136-146,
383        2006.
384    """
385    tr_A = q_form[..., 0, 0] + q_form[..., 1, 1] + q_form[..., 2, 2]
386    my_I = np.eye(3)
387    tr_AI = (tr_A.reshape(tr_A.shape + (1, 1)) * my_I)
388    return (1 / 3.0) * tr_AI
389
390
391def deviatoric(q_form):
392    r"""
393    Calculate the deviatoric (anisotropic) part of the tensor [1]_.
394
395    Parameters
396    ----------
397    q_form : ndarray
398        The quadratic form of a tensor, or an array with quadratic forms of
399        tensors. Should be of shape (x,y,z,3,3) or (n, 3, 3) or (3,3).
400
401    Returns
402    -------
403    A_squiggle : ndarray
404        The deviatoric part of the tensor in each spatial coordinate.
405
406    Notes
407    -----
408    The deviatoric part of the tensor is defined as (equations 3-5 in [1]_):
409
410    .. math ::
411         \widetilde{A} = A - \bar{A}
412
413    Where $A$ is the tensor quadratic form and $\bar{A}$ is the anisotropic
414    part of the tensor.
415
416    References
417    ----------
418    .. [1] Daniel B. Ennis and G. Kindlmann, "Orthogonal Tensor
419        Invariants and the Analysis of Diffusion Tensor Magnetic Resonance
420        Images", Magnetic Resonance in Medicine, vol. 55, no. 1, pp. 136-146,
421        2006.
422    """
423    A_squiggle = q_form - isotropic(q_form)
424    return A_squiggle
425
426
427def norm(q_form):
428    r"""
429    Calculate the Frobenius norm of a tensor quadratic form
430
431    Parameters
432    ----------
433    q_form: ndarray
434        The quadratic form of a tensor, or an array with quadratic forms of
435        tensors. Should be of shape (x,y,z,3,3) or (n, 3, 3) or (3,3).
436
437    Returns
438    -------
439    norm : ndarray
440        The Frobenius norm of the 3,3 tensor q_form in each spatial
441        coordinate.
442
443    Notes
444    -----
445    The Frobenius norm is defined as:
446
447    :math:
448        ||A||_F = [\sum_{i,j} abs(a_{i,j})^2]^{1/2}
449
450    See also
451    --------
452    np.linalg.norm
453    """
454    return np.sqrt(np.sum(np.sum(np.abs(q_form ** 2), -1), -1))
455
456
457def mode(q_form):
458    r"""
459    Mode (MO) of a diffusion tensor [1]_.
460
461    Parameters
462    ----------
463    q_form : ndarray
464        The quadratic form of a tensor, or an array with quadratic forms of
465        tensors. Should be of shape (x, y, z, 3, 3) or (n, 3, 3) or (3, 3).
466
467    Returns
468    -------
469    mode : array
470        Calculated tensor mode in each spatial coordinate.
471
472    Notes
473    -----
474    Mode ranges between -1 (planar anisotropy) and +1 (linear anisotropy)
475    with 0 representing orthotropy. Mode is calculated with the
476    following equation (equation 9 in [1]_):
477
478    .. math::
479
480        Mode = 3*\sqrt{6}*det(\widetilde{A}/norm(\widetilde{A}))
481
482    Where $\widetilde{A}$ is the deviatoric part of the tensor quadratic form.
483
484    References
485    ----------
486
487    .. [1] Daniel B. Ennis and G. Kindlmann, "Orthogonal Tensor
488        Invariants and the Analysis of Diffusion Tensor Magnetic Resonance
489        Images", Magnetic Resonance in Medicine, vol. 55, no. 1, pp. 136-146,
490        2006.
491    """
492
493    A_squiggle = deviatoric(q_form)
494    A_s_norm = norm(A_squiggle)
495    # Add two dims for the (3,3), so that it can broadcast on A_squiggle:
496    A_s_norm = A_s_norm.reshape(A_s_norm.shape + (1, 1))
497
498    return 3 * np.sqrt(6) * determinant((A_squiggle / A_s_norm))
499
500
501def linearity(evals, axis=-1):
502    r"""
503    The linearity of the tensor [1]_
504
505    Parameters
506    ----------
507    evals : array-like
508        Eigenvalues of a diffusion tensor.
509    axis : int
510        Axis of `evals` which contains 3 eigenvalues.
511
512    Returns
513    -------
514    linearity : array
515        Calculated linearity of the diffusion tensor.
516
517    Notes
518    -----
519    Linearity is calculated with the following equation:
520
521    .. math::
522
523        Linearity = \frac{\lambda_1-\lambda_2}{\lambda_1+\lambda_2+\lambda_3}
524
525    References
526    ----------
527    .. [1] Westin C.-F., Peled S., Gubjartsson H., Kikinis R., Jolesz F.,
528        "Geometrical diffusion measures for MRI from tensor basis analysis" in
529        Proc. 5th Annual ISMRM, 1997.
530    """
531    evals = _roll_evals(evals, axis)
532    ev1, ev2, ev3 = evals
533    return (ev1 - ev2) / evals.sum(0)
534
535
536def planarity(evals, axis=-1):
537    r"""
538    The planarity of the tensor [1]_
539
540    Parameters
541    ----------
542    evals : array-like
543        Eigenvalues of a diffusion tensor.
544    axis : int
545        Axis of `evals` which contains 3 eigenvalues.
546
547    Returns
548    -------
549    linearity : array
550        Calculated linearity of the diffusion tensor.
551
552    Notes
553    -----
554    Planarity is calculated with the following equation:
555
556    .. math::
557
558        Planarity =
559        \frac{2 (\lambda_2-\lambda_3)}{\lambda_1+\lambda_2+\lambda_3}
560
561    References
562    ----------
563    .. [1] Westin C.-F., Peled S., Gubjartsson H., Kikinis R., Jolesz F.,
564        "Geometrical diffusion measures for MRI from tensor basis analysis" in
565        Proc. 5th Annual ISMRM, 1997.
566    """
567    evals = _roll_evals(evals, axis)
568    ev1, ev2, ev3 = evals
569    return (2 * (ev2 - ev3) / evals.sum(0))
570
571
572def sphericity(evals, axis=-1):
573    r"""
574    The sphericity of the tensor [1]_
575
576    Parameters
577    ----------
578    evals : array-like
579        Eigenvalues of a diffusion tensor.
580    axis : int
581        Axis of `evals` which contains 3 eigenvalues.
582
583    Returns
584    -------
585    sphericity : array
586        Calculated sphericity of the diffusion tensor.
587
588    Notes
589    -----
590    Sphericity is calculated with the following equation:
591
592    .. math::
593
594        Sphericity = \frac{3 \lambda_3)}{\lambda_1+\lambda_2+\lambda_3}
595
596    References
597    ----------
598    .. [1] Westin C.-F., Peled S., Gubjartsson H., Kikinis R., Jolesz F.,
599        "Geometrical diffusion measures for MRI from tensor basis analysis" in
600        Proc. 5th Annual ISMRM, 1997.
601    """
602    evals = _roll_evals(evals, axis)
603    ev1, ev2, ev3 = evals
604    return (3 * ev3) / evals.sum(0)
605
606
607def apparent_diffusion_coef(q_form, sphere):
608    r"""
609    Calculate the apparent diffusion coefficient (ADC) in each direction of a
610    sphere.
611
612    Parameters
613    ----------
614    q_form : ndarray
615        The quadratic form of a tensor, or an array with quadratic forms of
616        tensors. Should be of shape (..., 3, 3)
617
618    sphere : a Sphere class instance
619        The ADC will be calculated for each of the vertices in the sphere
620
621    Notes
622    -----
623    The calculation of ADC, relies on the following relationship:
624
625    .. math ::
626            ADC = \vec{b} Q \vec{b}^T
627
628    Where Q is the quadratic form of the tensor.
629
630    """
631    bvecs = sphere.vertices
632    bvals = np.ones(bvecs.shape[0])
633    gtab = gradient_table(bvals, bvecs)
634    D = design_matrix(gtab)[:, :6]
635    return -np.dot(lower_triangular(q_form), D.T)
636
637
638def tensor_prediction(dti_params, gtab, S0):
639    """
640    Predict a signal given tensor parameters.
641
642    Parameters
643    ----------
644    dti_params : ndarray
645        Tensor parameters. The last dimension should have 12 tensor
646        parameters: 3 eigenvalues, followed by the 3 corresponding
647        eigenvectors.
648
649    gtab : a GradientTable class instance
650        The gradient table for this prediction
651
652    S0 : float or ndarray
653        The non diffusion-weighted signal in every voxel, or across all
654        voxels. Default: 1
655
656    Notes
657    -----
658    The predicted signal is given by: $S(\theta, b) = S_0 * e^{-b ADC}$, where
659    $ADC = \theta Q \theta^T$, $\theta$ is a unit vector pointing at any
660    direction on the sphere for which a signal is to be predicted, $b$ is the b
661    value provided in the GradientTable input for that direction, $Q$ is the
662    quadratic form of the tensor determined by the input parameters.
663    """
664    evals = dti_params[..., :3]
665    evecs = dti_params[..., 3:].reshape(dti_params.shape[:-1] + (3, 3))
666    qform = vec_val_vect(evecs, evals)
667    del evals, evecs
668    lower_tri = lower_triangular(qform, S0)
669    del qform
670
671    D = design_matrix(gtab)
672    return np.exp(np.dot(lower_tri, D.T))
673
674
675class TensorModel(ReconstModel):
676    """ Diffusion Tensor
677    """
678
679    def __init__(self, gtab, fit_method="WLS", return_S0_hat=False, *args,
680                 **kwargs):
681        """ A Diffusion Tensor Model [1]_, [2]_.
682
683        Parameters
684        ----------
685        gtab : GradientTable class instance
686
687        fit_method : str or callable
688            str can be one of the following:
689
690            'WLS' for weighted least squares
691                :func:`dti.wls_fit_tensor`
692            'LS' or 'OLS' for ordinary least squares
693                :func:`dti.ols_fit_tensor`
694            'NLLS' for non-linear least-squares
695                :func:`dti.nlls_fit_tensor`
696            'RT' or 'restore' or 'RESTORE' for RESTORE robust tensor
697                fitting [3]_
698                :func:`dti.restore_fit_tensor`
699
700            callable has to have the signature:
701              fit_method(design_matrix, data, *args, **kwargs)
702
703        return_S0_hat : bool
704            Boolean to return (True) or not (False) the S0 values for the fit.
705
706        args, kwargs : arguments and key-word arguments passed to the
707           fit_method. See dti.wls_fit_tensor, dti.ols_fit_tensor for details
708
709        min_signal : float
710            The minimum signal value. Needs to be a strictly positive
711            number. Default: minimal signal in the data provided to `fit`.
712
713        Notes
714        -----
715        In order to increase speed of processing, tensor fitting is done
716        simultaneously over many voxels. Many fit_methods use the 'step'
717        parameter to set the number of voxels that will be fit at once in each
718        iteration. This is the chunk size as a number of voxels. A larger step
719        value should speed things up, but it will also take up more memory. It
720        is advisable to keep an eye on memory consumption as this value is
721        increased.
722
723        E.g., in :func:`iter_fit_tensor` we have a default step value of
724        1e4
725
726        References
727        ----------
728        .. [1] Basser, P.J., Mattiello, J., LeBihan, D., 1994. Estimation of
729           the effective self-diffusion tensor from the NMR spin echo. J Magn
730           Reson B 103, 247-254.
731        .. [2] Basser, P., Pierpaoli, C., 1996. Microstructural and
732           physiological features of tissues elucidated by quantitative
733           diffusion-tensor MRI.  Journal of Magnetic Resonance 111, 209-219.
734        .. [3] Lin-Ching C., Jones D.K., Pierpaoli, C. 2005. RESTORE: Robust
735           estimation of tensors by outlier rejection. MRM 53: 1088-1095
736
737        """
738        ReconstModel.__init__(self, gtab)
739
740        if not callable(fit_method):
741            try:
742                fit_method = common_fit_methods[fit_method]
743            except KeyError:
744                e_s = '"' + str(fit_method) + '" is not a known fit '
745                e_s += 'method, the fit method should either be a '
746                e_s += 'function or one of the common fit methods'
747                raise ValueError(e_s)
748        self.fit_method = fit_method
749        self.return_S0_hat = return_S0_hat
750        self.design_matrix = design_matrix(self.gtab)
751        self.args = args
752        self.kwargs = kwargs
753        self.min_signal = self.kwargs.pop('min_signal', None)
754        if self.min_signal is not None and self.min_signal <= 0:
755            e_s = "The `min_signal` key-word argument needs to be strictly"
756            e_s += " positive."
757            raise ValueError(e_s)
758
759    def fit(self, data, mask=None):
760        """ Fit method of the DTI model class
761
762        Parameters
763        ----------
764        data : array
765            The measured signal from one voxel.
766
767        mask : array
768            A boolean array used to mark the coordinates in the data that
769            should be analyzed that has the shape data.shape[:-1]
770
771        """
772        S0_params = None
773
774        if mask is not None:
775            # Check for valid shape of the mask
776            if mask.shape != data.shape[:-1]:
777                raise ValueError("Mask is not the same shape as data.")
778            mask = np.array(mask, dtype=bool, copy=False)
779        data_in_mask = np.reshape(data[mask], (-1, data.shape[-1]))
780
781        if self.min_signal is None:
782            min_signal = MIN_POSITIVE_SIGNAL
783        else:
784            min_signal = self.min_signal
785
786        data_in_mask = np.maximum(data_in_mask, min_signal)
787
788        params_in_mask = self.fit_method(
789                self.design_matrix,
790                data_in_mask,
791                return_S0_hat=self.return_S0_hat,
792                *self.args,
793                **self.kwargs)
794        if self.return_S0_hat:
795            params_in_mask, model_S0 = params_in_mask
796
797        if mask is None:
798            out_shape = data.shape[:-1] + (-1, )
799            dti_params = params_in_mask.reshape(out_shape)
800            if self.return_S0_hat:
801                S0_params = model_S0.reshape(out_shape[:-1])
802        else:
803            dti_params = np.zeros(data.shape[:-1] + (12,))
804            dti_params[mask, :] = params_in_mask
805            if self.return_S0_hat:
806                S0_params = np.zeros(data.shape[:-1])
807                S0_params[mask] = model_S0
808
809        return TensorFit(self, dti_params, model_S0=S0_params)
810
811    def predict(self, dti_params, S0=1.):
812        """
813        Predict a signal for this TensorModel class instance given parameters.
814
815        Parameters
816        ----------
817        dti_params : ndarray
818            The last dimension should have 12 tensor parameters: 3
819            eigenvalues, followed by the 3 eigenvectors
820
821        S0 : float or ndarray
822            The non diffusion-weighted signal in every voxel, or across all
823            voxels. Default: 1
824        """
825        return tensor_prediction(dti_params, self.gtab, S0)
826
827
828class TensorFit(object):
829
830    def __init__(self, model, model_params, model_S0=None):
831        """ Initialize a TensorFit class instance.
832        """
833        self.model = model
834        self.model_params = model_params
835        self.model_S0 = model_S0
836
837    def __getitem__(self, index):
838        model_params = self.model_params
839        model_S0 = self.model_S0
840        N = model_params.ndim
841        if type(index) is not tuple:
842            index = (index,)
843        elif len(index) >= model_params.ndim:
844            raise IndexError("IndexError: invalid index")
845        index = index + (slice(None),) * (N - len(index))
846        if model_S0 is not None:
847            model_S0 = model_S0[index[:-1]]
848        return type(self)(self.model, model_params[index], model_S0=model_S0)
849
850    @property
851    def S0_hat(self):
852        return self.model_S0
853
854    @property
855    def shape(self):
856        return self.model_params.shape[:-1]
857
858    @property
859    def directions(self):
860        """
861        For tracking - return the primary direction in each voxel
862        """
863        return self.evecs[..., None, :, 0]
864
865    @property
866    def evals(self):
867        """
868        Returns the eigenvalues of the tensor as an array
869        """
870        return self.model_params[..., :3]
871
872    @property
873    def evecs(self):
874        """
875        Returns the eigenvectors of the tensor as an array, columnwise
876        """
877        evecs = self.model_params[..., 3:12]
878        return evecs.reshape(self.shape + (3, 3))
879
880    @property
881    def quadratic_form(self):
882        """Calculates the 3x3 diffusion tensor for each voxel"""
883        # do `evecs * evals * evecs.T` where * is matrix multiply
884        # einsum does this with:
885        # np.einsum('...ij,...j,...kj->...ik', evecs, evals, evecs)
886        return vec_val_vect(self.evecs, self.evals)
887
888    def lower_triangular(self, b0=None):
889        return lower_triangular(self.quadratic_form, b0)
890
891    @auto_attr
892    def fa(self):
893        """Fractional anisotropy (FA) calculated from cached eigenvalues."""
894        return fractional_anisotropy(self.evals)
895
896    @auto_attr
897    def color_fa(self):
898        """Color fractional anisotropy of diffusion tensor"""
899        return color_fa(self.fa, self.evecs)
900
901    @auto_attr
902    def ga(self):
903        """Geodesic anisotropy (GA) calculated from cached eigenvalues."""
904        return geodesic_anisotropy(self.evals)
905
906    @auto_attr
907    def mode(self):
908        """
909        Tensor mode calculated from cached eigenvalues.
910        """
911        return mode(self.quadratic_form)
912
913    @auto_attr
914    def md(self):
915        r"""
916        Mean diffusivity (MD) calculated from cached eigenvalues.
917
918        Returns
919        -------
920        md : array (V, 1)
921            Calculated MD.
922
923        Notes
924        -----
925        MD is calculated with the following equation:
926
927        .. math::
928
929            MD = \frac{\lambda_1+\lambda_2+\lambda_3}{3}
930
931        """
932        return self.trace / 3.0
933
934    @auto_attr
935    def rd(self):
936        r"""
937        Radial diffusivity (RD) calculated from cached eigenvalues.
938
939        Returns
940        -------
941        rd : array (V, 1)
942            Calculated RD.
943
944        Notes
945        -----
946        RD is calculated with the following equation:
947
948        .. math::
949
950          RD = \frac{\lambda_2 + \lambda_3}{2}
951
952
953        """
954        return radial_diffusivity(self.evals)
955
956    @auto_attr
957    def ad(self):
958        r"""
959        Axial diffusivity (AD) calculated from cached eigenvalues.
960
961        Returns
962        -------
963        ad : array (V, 1)
964            Calculated AD.
965
966        Notes
967        -----
968        RD is calculated with the following equation:
969
970        .. math::
971
972          AD = \lambda_1
973
974
975        """
976        return axial_diffusivity(self.evals)
977
978    @auto_attr
979    def trace(self):
980        r"""
981        Trace of the tensor calculated from cached eigenvalues.
982
983        Returns
984        -------
985        trace : array (V, 1)
986            Calculated trace.
987
988        Notes
989        -----
990        The trace is calculated with the following equation:
991
992        .. math::
993
994          trace = \lambda_1 + \lambda_2 + \lambda_3
995        """
996        return trace(self.evals)
997
998    @auto_attr
999    def planarity(self):
1000        r"""
1001        Returns
1002        -------
1003        sphericity : array
1004            Calculated sphericity of the diffusion tensor [1]_.
1005
1006        Notes
1007        -----
1008        Sphericity is calculated with the following equation:
1009
1010        .. math::
1011
1012            Sphericity =
1013            \frac{2 (\lambda_2 - \lambda_3)}{\lambda_1+\lambda_2+\lambda_3}
1014
1015        References
1016        ----------
1017        .. [1] Westin C.-F., Peled S., Gubjartsson H., Kikinis R., Jolesz
1018            F., "Geometrical diffusion measures for MRI from tensor basis
1019            analysis" in Proc. 5th Annual ISMRM, 1997.
1020
1021        """
1022        return planarity(self.evals)
1023
1024    @auto_attr
1025    def linearity(self):
1026        r"""
1027        Returns
1028        -------
1029        linearity : array
1030            Calculated linearity of the diffusion tensor [1]_.
1031
1032        Notes
1033        -----
1034        Linearity is calculated with the following equation:
1035
1036        .. math::
1037
1038            Linearity =
1039            \frac{\lambda_1-\lambda_2}{\lambda_1+\lambda_2+\lambda_3}
1040
1041        References
1042        ----------
1043        .. [1] Westin C.-F., Peled S., Gubjartsson H., Kikinis R., Jolesz
1044            F., "Geometrical diffusion measures for MRI from tensor basis
1045            analysis" in Proc. 5th Annual ISMRM, 1997.
1046
1047        """
1048        return linearity(self.evals)
1049
1050    @auto_attr
1051    def sphericity(self):
1052        r"""
1053        Returns
1054        -------
1055        sphericity : array
1056            Calculated sphericity of the diffusion tensor [1]_.
1057
1058        Notes
1059        -----
1060        Sphericity is calculated with the following equation:
1061
1062        .. math::
1063
1064            Sphericity = \frac{3 \lambda_3}{\lambda_1+\lambda_2+\lambda_3}
1065
1066        References
1067        ----------
1068        .. [1] Westin C.-F., Peled S., Gubjartsson H., Kikinis R., Jolesz
1069            F., "Geometrical diffusion measures for MRI from tensor basis
1070            analysis" in Proc. 5th Annual ISMRM, 1997.
1071
1072        """
1073        return sphericity(self.evals)
1074
1075    def odf(self, sphere):
1076        r"""
1077        The diffusion orientation distribution function (dODF). This is an
1078        estimate of the diffusion distance in each direction
1079
1080        Parameters
1081        ----------
1082        sphere : Sphere class instance.
1083            The dODF is calculated in the vertices of this input.
1084
1085        Returns
1086        -------
1087        odf : ndarray
1088            The diffusion distance in every direction of the sphere in every
1089            voxel in the input data.
1090
1091        Notes
1092        -----
1093        This is based on equation 3 in [1]_. To re-derive it from
1094        scratch, follow steps in [2]_, Section 7.9 Equation
1095        7.24 but with an $r^2$ term in the integral.
1096
1097        References
1098        ----------
1099        .. [1] Aganj, I., Lenglet, C., Sapiro, G., Yacoub, E., Ugurbil,
1100            K., & Harel, N. (2010). Reconstruction of the orientation
1101            distribution function in single- and multiple-shell q-ball imaging
1102            within constant solid angle. Magnetic Resonance in Medicine, 64(2),
1103            554-566. doi:DOI: 10.1002/mrm.22365
1104
1105        .. [2] Descoteaux, M. (2008). PhD Thesis: High Angular
1106           Resolution Diffusion MRI: from Local Estimation to Segmentation and
1107           Tractography.
1108           ftp://ftp-sop.inria.fr/athena/Publications/PhDs/descoteaux_thesis.pdf
1109
1110        """
1111        odf = np.zeros((self.evals.shape[:-1] + (sphere.vertices.shape[0],)))
1112        if len(self.evals.shape) > 1:
1113            mask = np.where((self.evals[..., 0] > 0) &
1114                            (self.evals[..., 1] > 0) &
1115                            (self.evals[..., 2] > 0))
1116            evals = self.evals[mask]
1117            evecs = self.evecs[mask]
1118        else:
1119            evals = self.evals
1120            evecs = self.evecs
1121        lower = 4 * np.pi * np.sqrt(np.prod(evals, -1))
1122        projection = np.dot(sphere.vertices, evecs)
1123        projection /= np.sqrt(evals)
1124        result = ((vector_norm(projection) ** -3) / lower).T
1125        if len(self.evals.shape) > 1:
1126            odf[mask] = result
1127        else:
1128            odf = result
1129        return odf
1130
1131    def adc(self, sphere):
1132        """
1133        Calculate the apparent diffusion coefficient (ADC) in each direction on
1134        the sphere for each voxel in the data
1135
1136        Parameters
1137        ----------
1138        sphere : Sphere class instance
1139
1140        Returns
1141        -------
1142        adc : ndarray
1143           The estimates of the apparent diffusion coefficient in every
1144           direction on the input sphere
1145
1146        Notes
1147        -----
1148        The calculation of ADC, relies on the following relationship:
1149
1150        .. math ::
1151
1152            ADC = \vec{b} Q \vec{b}^T
1153
1154        Where Q is the quadratic form of the tensor.
1155        """
1156        return apparent_diffusion_coef(self.quadratic_form, sphere)
1157
1158    def predict(self, gtab, S0=None, step=None):
1159        """
1160        Given a model fit, predict the signal on the vertices of a sphere
1161
1162        Parameters
1163        ----------
1164        gtab : a GradientTable class instance
1165            This encodes the directions for which a prediction is made
1166
1167        S0 : float array
1168           The mean non-diffusion weighted signal in each voxel. Default:
1169           The fitted S0 value in all voxels if it was fitted. Otherwise 1 in
1170           all voxels.
1171
1172        step : int
1173            The chunk size as a number of voxels. Optional parameter with
1174            default value 10,000.
1175
1176            In order to increase speed of processing, tensor fitting is done
1177            simultaneously over many voxels. This parameter sets the number of
1178            voxels that will be fit at once in each iteration. A larger step
1179            value should speed things up, but it will also take up more memory.
1180            It is advisable to keep an eye on memory consumption as this value
1181            is increased.
1182
1183        Notes
1184        -----
1185        The predicted signal is given by:
1186
1187        .. math ::
1188
1189            S(\theta, b) = S_0 * e^{-b ADC}
1190
1191        Where:
1192        .. math ::
1193            ADC = \theta Q \theta^T
1194
1195        $\theta$ is a unit vector pointing at any direction on the sphere for
1196        which a signal is to be predicted and $b$ is the b value provided in
1197        the GradientTable input for that direction
1198        """
1199        if S0 is None:
1200            S0 = self.model_S0
1201            if S0 is None:  # if we didn't input or estimate S0 just use 1
1202                S0 = 1.
1203        shape = self.model_params.shape[:-1]
1204        size = np.prod(shape)
1205        if step is None:
1206            step = self.model.kwargs.get('step', size)
1207        if step >= size:
1208            return tensor_prediction(self.model_params[..., 0:12], gtab, S0=S0)
1209        params = np.reshape(self.model_params,
1210                            (-1, self.model_params.shape[-1]))
1211        predict = np.empty((size, gtab.bvals.shape[0]))
1212        if isinstance(S0, np.ndarray):
1213            S0 = S0.ravel()
1214        for i in range(0, size, step):
1215            if isinstance(S0, np.ndarray):
1216                this_S0 = S0[i:i + step]
1217            else:
1218                this_S0 = S0
1219            predict[i:i + step] = tensor_prediction(params[i:i + step], gtab,
1220                                                    S0=this_S0)
1221        return predict.reshape(shape + (gtab.bvals.shape[0], ))
1222
1223
1224def iter_fit_tensor(step=1e4):
1225    """Wrap a fit_tensor func and iterate over chunks of data with given length
1226
1227    Splits data into a number of chunks of specified size and iterates the
1228    decorated fit_tensor function over them. This is useful to counteract the
1229    temporary but significant memory usage increase in fit_tensor functions
1230    that use vectorized operations and need to store large temporary arrays for
1231    their vectorized operations.
1232
1233    Parameters
1234    ----------
1235    step : int
1236        The chunk size as a number of voxels. Optional parameter with default
1237        value 10,000.
1238
1239        In order to increase speed of processing, tensor fitting is done
1240        simultaneously over many voxels. This parameter sets the number of
1241        voxels that will be fit at once in each iteration. A larger step value
1242        should speed things up, but it will also take up more memory. It is
1243        advisable to keep an eye on memory consumption as this value is
1244        increased.
1245    """
1246
1247    def iter_decorator(fit_tensor):
1248        """Actual iter decorator returned by iter_fit_tensor dec factory
1249
1250        Parameters
1251        ----------
1252        fit_tensor : callable
1253            A tensor fitting callable (most likely a function). The callable
1254            has to have the signature:
1255              fit_method(design_matrix, data, *args, **kwargs)
1256        """
1257
1258        @functools.wraps(fit_tensor)
1259        def wrapped_fit_tensor(design_matrix, data, return_S0_hat=False,
1260                               step=step, *args, **kwargs):
1261            """Iterate fit_tensor function over the data chunks
1262
1263            Parameters
1264            ----------
1265            design_matrix : array (g, 7)
1266                Design matrix holding the covariants used to solve for the
1267                regression coefficients.
1268            data : array ([X, Y, Z, ...], g)
1269                Data or response variables holding the data. Note that the last
1270                dimension should contain the data. It makes no copies of data.
1271            return_S0_hat : bool
1272                Boolean to return (True) or not (False) the S0 values for the
1273                fit.
1274            step : int
1275                The chunk size as a number of voxels. Overrides `step` value
1276                of `iter_fit_tensor`.
1277            args : {list,tuple}
1278                Any extra optional positional arguments passed to `fit_tensor`.
1279            kwargs : dict
1280                Any extra optional keyword arguments passed to `fit_tensor`.
1281            """
1282            shape = data.shape[:-1]
1283            size = np.prod(shape)
1284            step = int(step) or size
1285            if step >= size:
1286                return fit_tensor(design_matrix, data,
1287                                  return_S0_hat=return_S0_hat,
1288                                  *args, **kwargs)
1289            data = data.reshape(-1, data.shape[-1])
1290            dtiparams = np.empty((size, 12), dtype=np.float64)
1291            if return_S0_hat:
1292                S0params = np.empty(size, dtype=np.float64)
1293            for i in range(0, size, step):
1294                if return_S0_hat:
1295                    dtiparams[i:i + step], S0params[i:i + step] \
1296                        = fit_tensor(design_matrix,
1297                                     data[i:i + step],
1298                                     return_S0_hat=return_S0_hat,
1299                                     *args, **kwargs)
1300                else:
1301                    dtiparams[i:i + step] = fit_tensor(design_matrix,
1302                                                       data[i:i + step],
1303                                                       *args, **kwargs)
1304            if return_S0_hat:
1305                return (dtiparams.reshape(shape + (12, )),
1306                        S0params.reshape(shape + (1, )))
1307            else:
1308                return dtiparams.reshape(shape + (12, ))
1309
1310        return wrapped_fit_tensor
1311
1312    return iter_decorator
1313
1314
1315@iter_fit_tensor()
1316def wls_fit_tensor(design_matrix, data, return_S0_hat=False):
1317    r"""
1318    Computes weighted least squares (WLS) fit to calculate self-diffusion
1319    tensor using a linear regression model [1]_.
1320
1321    Parameters
1322    ----------
1323    design_matrix : array (g, 7)
1324        Design matrix holding the covariants used to solve for the regression
1325        coefficients.
1326    data : array ([X, Y, Z, ...], g)
1327        Data or response variables holding the data. Note that the last
1328        dimension should contain the data. It makes no copies of data.
1329    return_S0_hat : bool
1330        Boolean to return (True) or not (False) the S0 values for the fit.
1331
1332    Returns
1333    -------
1334    eigvals : array (..., 3)
1335        Eigenvalues from eigen decomposition of the tensor.
1336    eigvecs : array (..., 3, 3)
1337        Associated eigenvectors from eigen decomposition of the tensor.
1338        Eigenvectors are columnar (e.g. eigvecs[:,j] is associated with
1339        eigvals[j])
1340
1341
1342    See Also
1343    --------
1344    decompose_tensor
1345
1346    Notes
1347    -----
1348    In Chung, et al. 2006, the regression of the WLS fit needed an unbiased
1349    preliminary estimate of the weights and therefore the ordinary least
1350    squares (OLS) estimates were used. A "two pass" method was implemented:
1351
1352        1. calculate OLS estimates of the data
1353        2. apply the OLS estimates as weights to the WLS fit of the data
1354
1355    This ensured heteroscedasticity could be properly modeled for various
1356    types of bootstrap resampling (namely residual bootstrap).
1357
1358    .. math::
1359
1360        y = \mathrm{data} \\
1361        X = \mathrm{design matrix} \\
1362        \hat{\beta}_\mathrm{WLS} =
1363        \mathrm{desired regression coefficients (e.g. tensor)}\\
1364        \\
1365        \hat{\beta}_\mathrm{WLS} = (X^T W X)^{-1} X^T W y \\
1366        \\
1367        W = \mathrm{diag}((X \hat{\beta}_\mathrm{OLS})^2),
1368        \mathrm{where} \hat{\beta}_\mathrm{OLS} = (X^T X)^{-1} X^T y
1369
1370    References
1371    ----------
1372    .. [1] Chung, SW., Lu, Y., Henry, R.G., 2006. Comparison of bootstrap
1373       approaches for estimation of uncertainties of DTI parameters.
1374       NeuroImage 33, 531-541.
1375
1376    """
1377    tol = 1e-6
1378    data = np.asarray(data)
1379    ols_fit = _ols_fit_matrix(design_matrix)
1380    log_s = np.log(data)
1381    w = np.exp(np.einsum('...ij,...j', ols_fit, log_s))
1382    fit_result = np.einsum('...ij,...j',
1383                           pinv(design_matrix * w[..., None]),
1384                           w * log_s)
1385    if return_S0_hat:
1386        return (eig_from_lo_tri(fit_result,
1387                                min_diffusivity=tol / -design_matrix.min()),
1388                np.exp(-fit_result[:, -1]))
1389    else:
1390        return eig_from_lo_tri(fit_result,
1391                               min_diffusivity=tol / -design_matrix.min())
1392
1393
1394@iter_fit_tensor()
1395def ols_fit_tensor(design_matrix, data, return_S0_hat=False):
1396    r"""
1397    Computes ordinary least squares (OLS) fit to calculate self-diffusion
1398    tensor using a linear regression model [1]_.
1399
1400    Parameters
1401    ----------
1402    design_matrix : array (g, 7)
1403        Design matrix holding the covariants used to solve for the regression
1404        coefficients.
1405    data : array ([X, Y, Z, ...], g)
1406        Data or response variables holding the data. Note that the last
1407        dimension should contain the data. It makes no copies of data.
1408    return_S0_hat : bool
1409        Boolean to return (True) or not (False) the S0 values for the fit.
1410
1411    Returns
1412    -------
1413    eigvals : array (..., 3)
1414        Eigenvalues from eigen decomposition of the tensor.
1415    eigvecs : array (..., 3, 3)
1416        Associated eigenvectors from eigen decomposition of the tensor.
1417        Eigenvectors are columnar (e.g. eigvecs[:,j] is associated with
1418        eigvals[j])
1419
1420
1421    See Also
1422    --------
1423    WLS_fit_tensor, decompose_tensor, design_matrix
1424
1425    Notes
1426    -----
1427    .. math::
1428
1429        y = \mathrm{data} \\
1430        X = \mathrm{design matrix} \\
1431
1432        \hat{\beta}_\mathrm{OLS} = (X^T X)^{-1} X^T y
1433
1434    References
1435    ----------
1436    ..  [1] Chung, SW., Lu, Y., Henry, R.G., 2006. Comparison of bootstrap
1437        approaches for estimation of uncertainties of DTI parameters.
1438        NeuroImage 33, 531-541.
1439    """
1440    tol = 1e-6
1441    data = np.asarray(data)
1442    fit_result = np.einsum('...ij,...j', np.linalg.pinv(design_matrix),
1443                           np.log(data))
1444    if return_S0_hat:
1445        return (eig_from_lo_tri(fit_result,
1446                                min_diffusivity=tol / -design_matrix.min()),
1447                np.exp(-fit_result[:, -1]))
1448    else:
1449        return eig_from_lo_tri(fit_result,
1450                               min_diffusivity=tol / -design_matrix.min())
1451
1452
1453def _ols_fit_matrix(design_matrix):
1454    """
1455    Helper function to calculate the ordinary least squares (OLS)
1456    fit as a matrix multiplication. Mainly used to calculate WLS weights. Can
1457    be used to calculate regression coefficients in OLS but not recommended.
1458
1459    See Also:
1460    ---------
1461    wls_fit_tensor, ols_fit_tensor
1462
1463    Examples
1464    ---------
1465    ols_fit = _ols_fit_matrix(design_mat)
1466    ols_data = np.dot(ols_fit, data)
1467    """
1468
1469    U, S, V = np.linalg.svd(design_matrix, False)
1470    return np.dot(U, U.T)
1471
1472
1473def _nlls_err_func(tensor, design_matrix, data, weighting=None,
1474                   sigma=None):
1475    r"""
1476    Error function for the non-linear least-squares fit of the tensor.
1477
1478    Parameters
1479    ----------
1480    tensor : array (3,3)
1481        The 3-by-3 tensor matrix
1482
1483    design_matrix : array
1484        The design matrix
1485
1486    data : array
1487        The voxel signal in all gradient directions
1488
1489    weighting : str (optional).
1490         Whether to use the Geman-McClure weighting criterion (see [1]_
1491         for details)
1492
1493    sigma : float or float array (optional)
1494        If 'sigma' weighting is used, we will weight the error function
1495        according to the background noise estimated either in aggregate over
1496        all directions (when a float is provided), or to an estimate of the
1497        noise in each diffusion-weighting direction (if an array is
1498        provided). If 'gmm', the Geman-Mclure M-estimator is used for
1499        weighting (see below).
1500
1501    Notes
1502    -----
1503    The Geman-McClure M-estimator is described as follows [1]_ (page 1089):
1504    "The scale factor C affects the shape of the GMM
1505    [Geman-McClure M-estimator] weighting function and represents the expected
1506    spread of the residuals (i.e., the SD of the residuals) due to Gaussian
1507    distributed noise. The scale factor C can be estimated by many robust scale
1508    estimators. We used the median absolute deviation (MAD) estimator because
1509    it is very robust to outliers having a 50% breakdown point (6,7).
1510    The explicit formula for C using the MAD estimator is:
1511
1512    .. math ::
1513
1514            C = 1.4826 x MAD = 1.4826 x median{|r1-\hat{r}|,... |r_n-\hat{r}|}
1515
1516    where $\hat{r} = median{r_1, r_2, ..., r_3}$ and n is the number of data
1517    points. The multiplicative constant 1.4826 makes this an approximately
1518    unbiased estimate of scale when the error model is Gaussian."
1519
1520
1521    References
1522    ----------
1523    .. [1] Chang, L-C, Jones, DK and Pierpaoli, C (2005). RESTORE: robust
1524    estimation of tensors by outlier rejection. MRM, 53: 1088-95.
1525    """
1526    # This is the predicted signal given the params:
1527    y = np.exp(np.dot(design_matrix, tensor))
1528
1529    # Compute the residuals
1530    residuals = data - y
1531
1532    # If we don't want to weight the residuals, we are basically done:
1533    if weighting is None:
1534        # And we return the SSE:
1535        return residuals
1536    se = residuals ** 2
1537    # If the user provided a sigma (e.g 1.5267 * std(background_noise), as
1538    # suggested by Chang et al.) we will use it:
1539    if weighting == 'sigma':
1540        if sigma is None:
1541            e_s = "Must provide sigma value as input to use this weighting"
1542            e_s += " method"
1543            raise ValueError(e_s)
1544        w = 1 / (sigma**2)
1545
1546    elif weighting == 'gmm':
1547        # We use the Geman-McClure M-estimator to compute the weights on the
1548        # residuals:
1549        C = 1.4826 * np.median(np.abs(residuals - np.median(residuals)))
1550        with warnings.catch_warnings():
1551            warnings.simplefilter("ignore")
1552            w = 1 / (se + C**2)
1553            # The weights are normalized to the mean weight (see p. 1089):
1554            w = w / np.mean(w)
1555
1556    # Return the weighted residuals:
1557    with warnings.catch_warnings():
1558        warnings.simplefilter("ignore")
1559        return np.sqrt(w * se)
1560
1561
1562def _nlls_jacobian_func(tensor, design_matrix, data, *arg, **kwargs):
1563    """The Jacobian is the first derivative of the error function [1]_.
1564
1565    Notes
1566    -----
1567    This is an implementation of equation 14 in [1]_.
1568
1569    References
1570    ----------
1571    .. [1] Koay, CG, Chang, L-C, Carew, JD, Pierpaoli, C, Basser PJ (2006).
1572        A unifying theoretical and algorithmic framework for least squares
1573        methods of estimation in diffusion tensor imaging. MRM 182, 115-25.
1574
1575    """
1576    pred = np.exp(np.dot(design_matrix, tensor))
1577    return -pred[:, None] * design_matrix
1578
1579
1580def _decompose_tensor_nan(tensor, tensor_alternative, min_diffusivity=0):
1581    """ Helper function that expands the function decompose_tensor to deal
1582    with tensor with nan elements.
1583
1584    Computes tensor eigen decomposition to calculate eigenvalues and
1585    eigenvectors (Basser et al., 1994a). Some fit approaches can produce nan
1586    tensor elements in background voxels (particularly non-linear approaches).
1587    This function avoids the eigen decomposition errors of nan tensor elements
1588    by replacing tensor with nan elements by a given alternative tensor
1589    estimate.
1590
1591    Parameters
1592    ----------
1593    tensor : array (3, 3)
1594        Hermitian matrix representing a diffusion tensor.
1595    tensor_alternative : array (3, 3)
1596        Hermitian matrix representing a diffusion tensor obtain from an
1597        approach that does not produce nan tensor elements
1598    min_diffusivity : float
1599        Because negative eigenvalues are not physical and small eigenvalues,
1600        much smaller than the diffusion weighting, cause quite a lot of noise
1601        in metrics such as fa, diffusivity values smaller than
1602        `min_diffusivity` are replaced with `min_diffusivity`.
1603
1604    Returns
1605    -------
1606    eigvals : array (3)
1607        Eigenvalues from eigen decomposition of the tensor. Negative
1608        eigenvalues are replaced by zero. Sorted from largest to smallest.
1609    eigvecs : array (3, 3)
1610        Associated eigenvectors from eigen decomposition of the tensor.
1611        Eigenvectors are columnar (e.g. eigvecs[..., :, j] is associated with
1612        eigvals[..., j])
1613
1614    """
1615    try:
1616        evals, evecs = decompose_tensor(tensor[:6],
1617                                        min_diffusivity=min_diffusivity)
1618
1619    except np.linalg.LinAlgError:
1620        evals, evecs = decompose_tensor(tensor_alternative[:6],
1621                                        min_diffusivity=min_diffusivity)
1622    return evals, evecs
1623
1624
1625def nlls_fit_tensor(design_matrix, data, weighting=None,
1626                    sigma=None, jac=True, return_S0_hat=False):
1627    """
1628    Fit the cumulant expansion params (e.g. DTI, DKI) using non-linear
1629    least-squares.
1630
1631    Parameters
1632    ----------
1633    design_matrix : array (g, Npar)
1634        Design matrix holding the covariants used to solve for the regression
1635        coefficients. First six parameters of design matrix should correspond
1636        to the six unique diffusion tensor elements in the lower triangular
1637        order (Dxx, Dxy, Dyy, Dxz, Dyz, Dzz), while last parameter to -log(S0)
1638
1639    data : array ([X, Y, Z, ...], g)
1640        Data or response variables holding the data. Note that the last
1641        dimension should contain the data. It makes no copies of data.
1642
1643    weighting: str
1644           the weighting scheme to use in considering the
1645           squared-error. Default behavior is to use uniform weighting. Other
1646           options: 'sigma' 'gmm'
1647
1648    sigma: float
1649        If the 'sigma' weighting scheme is used, a value of sigma needs to be
1650        provided here. According to [Chang2005]_, a good value to use is
1651        1.5267 * std(background_noise), where background_noise is estimated
1652        from some part of the image known to contain no signal (only noise).
1653
1654    jac : bool
1655        Use the Jacobian? Default: True
1656
1657    return_S0_hat : bool
1658        Boolean to return (True) or not (False) the S0 values for the fit.
1659
1660    Returns
1661    -------
1662    nlls_params: the eigen-values and eigen-vectors of the tensor in each
1663        voxel.
1664    """
1665    # Detect number of parameters to estimate from design_matrix length plus
1666    # 5 due to diffusion tensor conversion to eigenvalue and eigenvectors
1667    npa = design_matrix.shape[-1] + 5
1668
1669    # Detect if number of parameters corresponds to dti
1670    dti = (npa == 12)
1671
1672    # Flatten for the iteration over voxels:
1673    flat_data = data.reshape((-1, data.shape[-1]))
1674    # Use the OLS method parameters as the starting point for the optimization:
1675    inv_design = np.linalg.pinv(design_matrix)
1676    log_s = np.log(flat_data)
1677    D = np.dot(inv_design, log_s.T).T
1678
1679    # Flatten for the iteration over voxels:
1680    ols_params = np.reshape(D, (-1, D.shape[-1]))
1681
1682    # Initialize parameter matrix
1683    params = np.empty((flat_data.shape[0], npa))
1684
1685    if return_S0_hat:
1686        model_S0 = np.empty((flat_data.shape[0], 1))
1687    for vox in range(flat_data.shape[0]):
1688        if np.all(flat_data[vox] == 0):
1689            raise ValueError("The data in this voxel contains only zeros")
1690
1691        start_params = ols_params[vox]
1692        # Do the optimization in this voxel:
1693        if jac:
1694            this_param, status = opt.leastsq(_nlls_err_func, start_params,
1695                                             args=(design_matrix,
1696                                                   flat_data[vox],
1697                                                   weighting,
1698                                                   sigma),
1699                                             Dfun=_nlls_jacobian_func)
1700        else:
1701            this_param, status = opt.leastsq(_nlls_err_func, start_params,
1702                                             args=(design_matrix,
1703                                                   flat_data[vox],
1704                                                   weighting,
1705                                                   sigma))
1706
1707        # Convert diffusion tensor parameters to the evals and the evecs:
1708        try:
1709            evals, evecs = decompose_tensor(
1710                from_lower_triangular(this_param[:6]))
1711            params[vox, :3] = evals
1712            params[vox, 3:12] = evecs.ravel()
1713
1714        # If leastsq failed to converge and produced nans, we'll resort to the
1715        # OLS solution in this voxel:
1716        except np.linalg.LinAlgError:
1717            this_params = start_params
1718            evals, evecs = decompose_tensor(
1719                from_lower_triangular(this_params[:6]))
1720            params[vox, :3] = evals
1721            params[vox, 3:12] = evecs.ravel()
1722
1723        if return_S0_hat:
1724            model_S0[vox] = np.exp(-this_param[-1])
1725        if not dti:
1726            md2 = evals.mean(0) ** 2
1727            params[vox, 12:] = this_param[6:-1] / md2
1728
1729    params.shape = data.shape[:-1] + (npa,)
1730    if return_S0_hat:
1731        model_S0.shape = data.shape[:-1] + (1,)
1732        return (params, model_S0)
1733    else:
1734        return params
1735
1736
1737def restore_fit_tensor(design_matrix, data, sigma=None, jac=True,
1738                       return_S0_hat=False):
1739    """
1740    Use the RESTORE algorithm [1]_ to calculate a robust tensor fit
1741
1742    Parameters
1743    ----------
1744
1745    design_matrix : array of shape (g, 7)
1746        Design matrix holding the covariants used to solve for the regression
1747        coefficients.
1748
1749    data : array of shape ([X, Y, Z, n_directions], g)
1750        Data or response variables holding the data. Note that the last
1751        dimension should contain the data. It makes no copies of data.
1752
1753    sigma : float
1754        An estimate of the variance. [1]_ recommend to use
1755        1.5267 * std(background_noise), where background_noise is estimated
1756        from some part of the image known to contain no signal (only noise).
1757
1758    jac : bool, optional
1759        Whether to use the Jacobian of the tensor to speed the non-linear
1760        optimization procedure used to fit the tensor parameters (see also
1761        :func:`nlls_fit_tensor`). Default: True
1762
1763    return_S0_hat : bool
1764        Boolean to return (True) or not (False) the S0 values for the fit.
1765
1766
1767    Returns
1768    -------
1769    restore_params : an estimate of the tensor parameters in each voxel.
1770
1771    References
1772    ----------
1773    .. [1] Chang, L-C, Jones, DK and Pierpaoli, C (2005). RESTORE: robust
1774    estimation of tensors by outlier rejection. MRM, 53: 1088-95.
1775
1776    """
1777    # Detect number of parameters to estimate from design_matrix length plus
1778    # 5 due to diffusion tensor conversion to eigenvalue and eigenvectors
1779    npa = design_matrix.shape[-1] + 5
1780
1781    # Detect if number of parameters corresponds to dti
1782    dti = (npa == 12)
1783
1784    # Flatten for the iteration over voxels:
1785    flat_data = data.reshape((-1, data.shape[-1]))
1786
1787    # Use the OLS method parameters as the starting point for the optimization:
1788    inv_design = np.linalg.pinv(design_matrix)
1789    log_s = np.log(flat_data)
1790    D = np.dot(inv_design, log_s.T).T
1791
1792    # Flatten for the iteration over voxels:
1793    ols_params = np.reshape(D, (-1, D.shape[-1]))
1794
1795    # Initialize parameter matrix
1796    params = np.empty((flat_data.shape[0], npa))
1797
1798    if return_S0_hat:
1799        model_S0 = np.empty((flat_data.shape[0], 1))
1800    for vox in range(flat_data.shape[0]):
1801        if np.all(flat_data[vox] == 0):
1802            raise ValueError("The data in this voxel contains only zeros")
1803
1804        start_params = ols_params[vox]
1805        # Do nlls using sigma weighting in this voxel:
1806        if jac:
1807            this_param, status = opt.leastsq(_nlls_err_func, start_params,
1808                                             args=(design_matrix,
1809                                                   flat_data[vox],
1810                                                   'sigma',
1811                                                   sigma),
1812                                             Dfun=_nlls_jacobian_func)
1813        else:
1814            this_param, status = opt.leastsq(_nlls_err_func, start_params,
1815                                             args=(design_matrix,
1816                                                   flat_data[vox],
1817                                                   'sigma',
1818                                                   sigma))
1819
1820        # Get the residuals:
1821        pred_sig = np.exp(np.dot(design_matrix, this_param))
1822        residuals = flat_data[vox] - pred_sig
1823
1824        # If any of the residuals are outliers (using 3 sigma as a criterion
1825        # following Chang et al., e.g page 1089):
1826        if np.any(np.abs(residuals) > 3 * sigma):
1827            # Do nlls with GMM-weighting:
1828            if jac:
1829                this_param, status = opt.leastsq(_nlls_err_func,
1830                                                 start_params,
1831                                                 args=(design_matrix,
1832                                                       flat_data[vox],
1833                                                       'gmm'),
1834                                                 Dfun=_nlls_jacobian_func)
1835            else:
1836                this_param, status = opt.leastsq(_nlls_err_func,
1837                                                 start_params,
1838                                                 args=(design_matrix,
1839                                                       flat_data[vox],
1840                                                       'gmm'))
1841
1842            # How are you doin' on those residuals?
1843            pred_sig = np.exp(np.dot(design_matrix, this_param))
1844            residuals = flat_data[vox] - pred_sig
1845            if np.any(np.abs(residuals) > 3 * sigma):
1846                # If you still have outliers, refit without those outliers:
1847                non_outlier_idx = np.where(np.abs(residuals) <= 3 * sigma)
1848                clean_design = design_matrix[non_outlier_idx]
1849                clean_sig = flat_data[vox][non_outlier_idx]
1850                if np.iterable(sigma):
1851                    this_sigma = sigma[non_outlier_idx]
1852                else:
1853                    this_sigma = sigma
1854
1855                if jac:
1856                    this_param, status = opt.leastsq(_nlls_err_func,
1857                                                     start_params,
1858                                                     args=(clean_design,
1859                                                           clean_sig,
1860                                                           'sigma',
1861                                                           this_sigma),
1862                                                     Dfun=_nlls_jacobian_func)
1863                else:
1864                    this_param, status = opt.leastsq(_nlls_err_func,
1865                                                     start_params,
1866                                                     args=(clean_design,
1867                                                           clean_sig,
1868                                                           'sigma',
1869                                                           this_sigma))
1870
1871        # Convert diffusion tensor parameters to the evals and the evecs:
1872        try:
1873            evals, evecs = decompose_tensor(
1874                from_lower_triangular(this_param[:6]))
1875            params[vox, :3] = evals
1876            params[vox, 3:12] = evecs.ravel()
1877
1878        # If leastsq failed to converge and produced nans, we'll resort to the
1879        # OLS solution in this voxel:
1880        except np.linalg.LinAlgError:
1881            this_params = start_params
1882            evals, evecs = decompose_tensor(
1883                from_lower_triangular(this_params[:6]))
1884            params[vox, :3] = evals
1885            params[vox, 3:12] = evecs.ravel()
1886
1887        if return_S0_hat:
1888            model_S0[vox] = np.exp(-this_param[-1])
1889        if not dti:
1890            md2 = evals.mean(0) ** 2
1891            params[vox, 12:] = this_param[6:-1] / md2
1892
1893    params.shape = data.shape[:-1] + (npa,)
1894    if return_S0_hat:
1895        model_S0.shape = data.shape[:-1] + (1,)
1896        return (params, model_S0)
1897    else:
1898        return params
1899
1900
1901_lt_indices = np.array([[0, 1, 3],
1902                        [1, 2, 4],
1903                        [3, 4, 5]])
1904
1905
1906def from_lower_triangular(D):
1907    """ Returns a tensor given the six unique tensor elements
1908
1909    Given the six unique tensor elements (in the order: Dxx, Dxy, Dyy, Dxz,
1910    Dyz, Dzz) returns a 3 by 3 tensor. All elements after the sixth are
1911    ignored.
1912
1913    Parameters
1914    -----------
1915    D : array_like, (..., >6)
1916        Unique elements of the tensors
1917
1918    Returns
1919    -------
1920    tensor : ndarray (..., 3, 3)
1921        3 by 3 tensors
1922
1923    """
1924    return D[..., _lt_indices]
1925
1926
1927_lt_rows = np.array([0, 1, 1, 2, 2, 2])
1928_lt_cols = np.array([0, 0, 1, 0, 1, 2])
1929
1930
1931def lower_triangular(tensor, b0=None):
1932    """
1933    Returns the six lower triangular values of the tensor and a dummy variable
1934    if b0 is not None
1935
1936    Parameters
1937    ----------
1938    tensor : array_like (..., 3, 3)
1939        a collection of 3, 3 diffusion tensors
1940    b0 : float
1941        if b0 is not none log(b0) is returned as the dummy variable
1942
1943    Returns
1944    -------
1945    D : ndarray
1946        If b0 is none, then the shape will be (..., 6) otherwise (..., 7)
1947
1948    """
1949    if tensor.shape[-2:] != (3, 3):
1950        raise ValueError("Diffusion tensors should be (..., 3, 3)")
1951    if b0 is None:
1952        return tensor[..., _lt_rows, _lt_cols]
1953    else:
1954        D = np.empty(tensor.shape[:-2] + (7,), dtype=tensor.dtype)
1955        D[..., 6] = -np.log(b0)
1956        D[..., :6] = tensor[..., _lt_rows, _lt_cols]
1957        return D
1958
1959
1960def decompose_tensor(tensor, min_diffusivity=0):
1961    """ Returns eigenvalues and eigenvectors given a diffusion tensor
1962
1963    Computes tensor eigen decomposition to calculate eigenvalues and
1964    eigenvectors (Basser et al., 1994a).
1965
1966    Parameters
1967    ----------
1968    tensor : array (..., 3, 3)
1969        Hermitian matrix representing a diffusion tensor.
1970    min_diffusivity : float
1971        Because negative eigenvalues are not physical and small eigenvalues,
1972        much smaller than the diffusion weighting, cause quite a lot of noise
1973        in metrics such as fa, diffusivity values smaller than
1974        `min_diffusivity` are replaced with `min_diffusivity`.
1975
1976    Returns
1977    -------
1978    eigvals : array (..., 3)
1979        Eigenvalues from eigen decomposition of the tensor. Negative
1980        eigenvalues are replaced by zero. Sorted from largest to smallest.
1981    eigvecs : array (..., 3, 3)
1982        Associated eigenvectors from eigen decomposition of the tensor.
1983        Eigenvectors are columnar (e.g. eigvecs[..., :, j] is associated with
1984        eigvals[..., j])
1985
1986    """
1987    # outputs multiplicity as well so need to unique
1988    eigenvals, eigenvecs = np.linalg.eigh(tensor)
1989
1990    # need to sort the eigenvalues and associated eigenvectors
1991    if eigenvals.ndim == 1:
1992        # this is a lot faster when dealing with a single voxel
1993        order = eigenvals.argsort()[::-1]
1994        eigenvecs = eigenvecs[:, order]
1995        eigenvals = eigenvals[order]
1996    else:
1997        # temporarily flatten eigenvals and eigenvecs to make sorting easier
1998        shape = eigenvals.shape[:-1]
1999        eigenvals = eigenvals.reshape(-1, 3)
2000        eigenvecs = eigenvecs.reshape(-1, 3, 3)
2001        size = eigenvals.shape[0]
2002        order = eigenvals.argsort()[:, ::-1]
2003        xi, yi = np.ogrid[:size, :3, :3][:2]
2004        eigenvecs = eigenvecs[xi, yi, order[:, None, :]]
2005        xi = np.ogrid[:size, :3][0]
2006        eigenvals = eigenvals[xi, order]
2007        eigenvecs = eigenvecs.reshape(shape + (3, 3))
2008        eigenvals = eigenvals.reshape(shape + (3, ))
2009    eigenvals = eigenvals.clip(min=min_diffusivity)
2010    # eigenvecs: each vector is columnar
2011
2012    return eigenvals, eigenvecs
2013
2014
2015def design_matrix(gtab, dtype=None):
2016    """  Constructs design matrix for DTI weighted least squares or
2017    least squares fitting. (Basser et al., 1994a)
2018
2019    Parameters
2020    ----------
2021    gtab : A GradientTable class instance
2022
2023    dtype : string
2024        Parameter to control the dtype of returned designed matrix
2025
2026    Returns
2027    -------
2028    design_matrix : array (g,7)
2029        Design matrix or B matrix assuming Gaussian distributed tensor model
2030        design_matrix[j, :] = (Bxx, Byy, Bzz, Bxy, Bxz, Byz, dummy)
2031    """
2032    if gtab.btens is None:
2033        B = np.zeros((gtab.gradients.shape[0], 7))
2034        B[:, 0] = gtab.bvecs[:, 0] * gtab.bvecs[:, 0] * 1. * gtab.bvals   # Bxx
2035        B[:, 1] = gtab.bvecs[:, 0] * gtab.bvecs[:, 1] * 2. * gtab.bvals   # Bxy
2036        B[:, 2] = gtab.bvecs[:, 1] * gtab.bvecs[:, 1] * 1. * gtab.bvals   # Byy
2037        B[:, 3] = gtab.bvecs[:, 0] * gtab.bvecs[:, 2] * 2. * gtab.bvals   # Bxz
2038        B[:, 4] = gtab.bvecs[:, 1] * gtab.bvecs[:, 2] * 2. * gtab.bvals   # Byz
2039        B[:, 5] = gtab.bvecs[:, 2] * gtab.bvecs[:, 2] * 1. * gtab.bvals   # Bzz
2040        B[:, 6] = np.ones(gtab.gradients.shape[0])
2041    else:
2042        B = np.zeros((gtab.gradients.shape[0], 7))
2043        B[:, 0] = gtab.btens[:, 0, 0]   # Bxx
2044        B[:, 1] = gtab.btens[:, 0, 1] * 2  # Bxy
2045        B[:, 2] = gtab.btens[:, 1, 1]   # Byy
2046        B[:, 3] = gtab.btens[:, 0, 2] * 2  # Bxz
2047        B[:, 4] = gtab.btens[:, 1, 2] * 2  # Byz
2048        B[:, 5] = gtab.btens[:, 2, 2]   # Bzz
2049        B[:, 6] = np.ones(gtab.gradients.shape[0])
2050
2051    return -B
2052
2053
2054def quantize_evecs(evecs, odf_vertices=None):
2055    """ Find the closest orientation of an evenly distributed sphere
2056
2057    Parameters
2058    ----------
2059    evecs : ndarray
2060    odf_vertices : None or ndarray
2061        If None, then set vertices from symmetric362 sphere.  Otherwise use
2062        passed ndarray as vertices
2063
2064    Returns
2065    -------
2066    IN : ndarray
2067    """
2068    max_evecs = evecs[..., :, 0]
2069    if odf_vertices is None:
2070        odf_vertices = get_sphere('symmetric362').vertices
2071    tup = max_evecs.shape[:-1]
2072    mec = max_evecs.reshape(np.prod(np.array(tup)), 3)
2073    IN = np.array([np.argmin(np.dot(odf_vertices, m)) for m in mec])
2074    IN = IN.reshape(tup)
2075    return IN
2076
2077
2078def eig_from_lo_tri(data, min_diffusivity=0):
2079    """
2080    Calculates tensor eigenvalues/eigenvectors from an array containing the
2081    lower diagonal form of the six unique tensor elements.
2082
2083    Parameters
2084    ----------
2085    data : array_like (..., 6)
2086        diffusion tensors elements stored in lower triangular order
2087    min_diffusivity : float
2088        See decompose_tensor()
2089
2090    Returns
2091    -------
2092    dti_params : array (..., 12)
2093        Eigen-values and eigen-vectors of the same array.
2094    """
2095    data = np.asarray(data)
2096    evals, evecs = decompose_tensor(from_lower_triangular(data),
2097                                    min_diffusivity=min_diffusivity)
2098    dti_params = np.concatenate((evals[..., None, :], evecs), axis=-2)
2099    return dti_params.reshape(data.shape[:-1] + (12, ))
2100
2101
2102common_fit_methods = {'WLS': wls_fit_tensor,
2103                      'LS': ols_fit_tensor,
2104                      'OLS': ols_fit_tensor,
2105                      'NLLS': nlls_fit_tensor,
2106                      'RT': restore_fit_tensor,
2107                      'restore': restore_fit_tensor,
2108                      'RESTORE': restore_fit_tensor
2109                      }
2110