# Source code for dipy.reconst.dki_micro

    #!/usr/bin/python
    """Classes and functions for fitting the DKI-based microstructural model"""

    import numpy as np

    from dipy.core.ndindex import ndindex
    import dipy.core.sphere as dps
    from dipy.data import get_sphere
    from dipy.reconst.dki import (
        DiffusionKurtosisFit,
        DiffusionKurtosisModel,
        _positive_evals,
        directional_diffusion,
        directional_kurtosis,
        kurtosis_maximum,
        split_dki_param,
    )
    from dipy.reconst.dti import (
        MIN_POSITIVE_SIGNAL,
        axial_diffusivity,
        decompose_tensor,
        design_matrix as dti_design_matrix,
        from_lower_triangular,
        lower_triangular,
        mean_diffusivity,
        radial_diffusivity,
        trace,
    )
    from dipy.reconst.vec_val_sum import vec_val_vect
    from dipy.testing.decorators import warning_for_keywords

    @warning_for_keywords()
    def axonal_water_fraction(dki_params, *, sphere="repulsion100", gtol=1e-2, mask=None):
        """Computes the axonal water fraction from DKI.

        See :footcite:p:`Fieremans2011` for further details about the method.

        Parameters
        ----------
        dki_params : ndarray (x, y, z, 27) or (n, 27)
            All parameters estimated from the diffusion kurtosis model.
            Parameters are ordered as follows:

                1) Three diffusion tensor's eigenvalues
                2) Three lines of the eigenvector matrix each containing the first,
                   second and third coordinates of the eigenvector
                3) Fifteen elements of the kurtosis tensor
        sphere : Sphere class instance, optional
            The sphere providing sample directions for the initial search of the
            maximal value of kurtosis.
        gtol : float, optional
            This input is to refine kurtosis maxima under the precision of the
            directions sampled on the sphere class instance. The gradient of the
            convergence procedure must be less than gtol before successful
            termination. If gtol is None, fiber direction is directly taken from
            the initial sampled directions of the given sphere object
        mask : ndarray
            A boolean array used to mark the coordinates in the data that should be
            analyzed that has the shape dki_params.shape[:-1]

        Returns
        -------
        awf : ndarray (x, y, z) or (n)
            Axonal Water Fraction

        References
        ----------
        .. footbibliography::
        """
        kt_max = kurtosis_maximum(dki_params, sphere=sphere, gtol=gtol, mask=mask)

        awf = kt_max / (kt_max + 3)

        return awf

    @warning_for_keywords()
    def diffusion_components(dki_params, *, sphere="repulsion100", awf=None, mask=None):
        """Extracts the restricted and hindered diffusion tensors of well aligned
        fibers from diffusion kurtosis imaging parameters.

        See :footcite:p:`Fieremans2011` for further details about the method.

        Parameters
        ----------
        dki_params : ndarray (x, y, z, 27) or (n, 27)
            All parameters estimated from the diffusion kurtosis model.
            Parameters are ordered as follows:

                1) Three diffusion tensor's eigenvalues
                2) Three lines of the eigenvector matrix each containing the first,
                   second and third coordinates of the eigenvector
                3) Fifteen elements of the kurtosis tensor
        sphere : Sphere class instance, optional
            The sphere providing sample directions to sample the restricted and
            hindered cellular diffusion tensors. For more details see
            :footcite:p:`Fieremans2011`.
        awf : ndarray, optional
            Array containing values of the axonal water fraction that has the shape
            dki_params.shape[:-1]. If not given this will be automatically computed
            using :func:`axonal_water_fraction` with function's default precision.
        mask : ndarray, optional
            A boolean array used to mark the coordinates in the data that should be
            analyzed that has the shape dki_params.shape[:-1]

        Returns
        -------
        edt : ndarray (x, y, z, 6) or (n, 6)
            Parameters of the hindered diffusion tensor.
        idt : ndarray (x, y, z, 6) or (n, 6)
            Parameters of the restricted diffusion tensor.

        Notes
        -----
        In the original article of DKI microstructural model
        :footcite:p:`Fieremans2011`, the hindered and restricted tensors were
        defined as the intra-cellular and extra-cellular diffusion compartments
        respectively.

        References
        ----------
        .. footbibliography::
        """
        shape = dki_params.shape[:-1]

        # load gradient directions
        if not isinstance(sphere, dps.Sphere):
            sphere = get_sphere(name=sphere)

        # select voxels where to apply the single fiber model
        if mask is None:
            mask = np.ones(shape, dtype="bool")
        else:
            if mask.shape != shape:
                raise ValueError("Mask is not the same shape as dki_params.")
            else:
                mask = np.asarray(mask, dtype=bool)

        # check or compute awf values
        if awf is None:
            awf = axonal_water_fraction(dki_params, sphere=sphere, mask=mask)
        else:
            if awf.shape != shape:
                raise ValueError("awf array is not the same shape as dki_params.")

        # Initialize hindered and restricted diffusion tensors
        edt_all = np.zeros(shape + (6,))
        idt_all = np.zeros(shape + (6,))

        # Generate matrix that converts apparent diffusion coefficients to tensors
        B = np.zeros((sphere.x.size, 6))
        B[:, 0] = sphere.x * sphere.x  # Bxx
        B[:, 1] = sphere.x * sphere.y * 2.0  # Bxy
        B[:, 2] = sphere.y * sphere.y  # Byy
        B[:, 3] = sphere.x * sphere.z * 2.0  # Bxz
        B[:, 4] = sphere.y * sphere.z * 2.0  # Byz
        B[:, 5] = sphere.z * sphere.z  # Bzz
        pinvB = np.linalg.pinv(B)

        # Compute hindered and restricted diffusion tensors for all voxels
        evals, evecs, kt = split_dki_param(dki_params)
        dt = lower_triangular(vec_val_vect(evecs, evals))
        md = mean_diffusivity(evals)

        index = ndindex(mask.shape)
        for idx in index:
            if not mask[idx]:
                continue
            # sample apparent diffusion and kurtosis values
            di = directional_diffusion(dt[idx], sphere.vertices)
            ki = directional_kurtosis(
                dt[idx], md[idx], kt[idx], sphere.vertices, adc=di, min_kurtosis=0
            )
            edi = di * (1 + np.sqrt(ki * awf[idx] / (3.0 - 3.0 * awf[idx])))
            edt = np.dot(pinvB, edi)
            edt_all[idx] = edt

            # We only move on if there is an axonal water fraction.
            # Otherwise, remaining params are already zero, so move on
            if awf[idx] == 0:
                continue
            # Convert apparent diffusion and kurtosis values to apparent diffusion
            # values of the hindered and restricted diffusion
            idi = di * (1 - np.sqrt(ki * (1.0 - awf[idx]) / (3.0 * awf[idx])))
            # generate hindered and restricted diffusion tensors
            idt = np.dot(pinvB, idi)
            idt_all[idx] = idt

        return edt_all, idt_all

    @warning_for_keywords()
    def dkimicro_prediction(params, gtab, *, S0=1):
        r"""Signal prediction given the DKI microstructure model parameters.

        Parameters
        ----------
        params : ndarray (x, y, z, 40) or (n, 40)
            All parameters estimated from the diffusion kurtosis microstructure
            model.
            Parameters are ordered as follows:

                1) Three diffusion tensor's eigenvalues
                2) Three lines of the eigenvector matrix each containing the
                   first, second and third coordinates of the eigenvector
                3) Fifteen elements of the kurtosis tensor
                4) Six elements of the hindered diffusion tensor
                5) Six elements of the restricted diffusion tensor
                6) Axonal water fraction
        gtab : a GradientTable class instance
            The gradient table for this prediction
        S0 : float or ndarray, optional
            The non diffusion-weighted signal in every voxel, or across all
            voxels

        Returns
        -------
        S : (..., N) ndarray
            Simulated signal based on the DKI microstructure model

        Notes
        -----
        1) The predicted signal is given by:
           .. math::

               S(\theta, b) = S_0 * [f * e^{-b ADC_{r}} + (1-f) * e^{-b ADC_{h}]

           where $ADC_{r}$ and $ADC_{h}$ are the apparent diffusion coefficients of
           the diffusion hindered and restricted compartment for a given direction
           $\theta$, $b$ is the b value provided in the GradientTable input for that
           direction, $f$ is the volume fraction of the restricted diffusion
           compartment (also known as the axonal water fraction).

        2) In the original article of DKI microstructural model
           :footcite:p:`Fieremans2011`, the hindered and restricted tensors were
           defined as the intra-cellular and extra-cellular diffusion compartments
           respectively.

        References
        ----------
        .. footbibliography::
        """

        # Initialize pred_sig
        pred_sig = np.zeros(params.shape[:-1] + (gtab.bvals.shape[0],))

        # Define dti design matrix and region to process
        D = dti_design_matrix(gtab)
        evals = params[..., :3]
        mask = _positive_evals(evals[..., 0], evals[..., 1], evals[..., 2])

        # Prepare parameters
        f = params[..., 27]
        adce = params[..., 28:34]
        adci = params[..., 34:40]

        if isinstance(S0, np.ndarray):
            S0_vol = S0 * np.ones(params.shape[:-1])
        else:
            S0_vol = S0

        # Process pred_sig for all data voxels
        index = ndindex(evals.shape[:-1])
        for v in index:
            if mask[v]:
                pred_sig[v] = (1.0 - f[v]) * np.exp(np.dot(D[:, :6], adce[v])) + f[
                    v
                ] * np.exp(np.dot(D[:, :6], adci[v]))

        return pred_sig * S0_vol

    def tortuosity(hindered_ad, hindered_rd):
        """Computes the tortuosity of the hindered diffusion compartment given
        its axial and radial diffusivities

        Parameters
        ----------
        hindered_ad: ndarray
            Array containing the values of the hindered axial diffusivity.
        hindered_rd: ndarray
            Array containing the values of the hindered radial diffusivity.

        Returns
        -------
        Tortuosity of the hindered diffusion compartment
        """
        if not isinstance(hindered_rd, np.ndarray):
            hindered_rd = np.array(hindered_rd)
        if not isinstance(hindered_ad, np.ndarray):
            hindered_ad = np.array(hindered_ad)

        tortuosity = np.zeros(hindered_rd.shape)

        # mask to avoid divisions by zero
        mask = hindered_rd > 0

        # Check single voxel cases. For numpy versions more recent than 1.7,
        # this if else condition is not required since single voxel can be
        # processed using the same line of code of multi-voxel
        if hindered_rd.size == 1:
            if mask:
                tortuosity = hindered_ad / hindered_rd
        else:
            tortuosity[mask] = hindered_ad[mask] / hindered_rd[mask]

        return tortuosity

    def _compartments_eigenvalues(cdt):
        """Helper function that computes the eigenvalues of a tissue sub
        compartment given its individual diffusion tensor

        Parameters
        ----------
        cdt : ndarray (..., 6)
            Diffusion tensors elements of the tissue compartment stored in lower
            triangular order.

        Returns
        -------
        eval : ndarry (..., 3)
            Eigenvalues of the tissue compartment
        """
        evals, evecs = decompose_tensor(from_lower_triangular(cdt))
        return evals

    class KurtosisMicrostructureModel(DiffusionKurtosisModel):
        """Class for the Diffusion Kurtosis Microstructural Model"""

        def __init__(self, gtab, *args, fit_method="WLS", **kwargs):
            """Initialize a KurtosisMicrostrutureModel class instance.

            See :footcite:p:`Fieremans2011` for further details about the model.

            Parameters
            ----------
            gtab : GradientTable class instance
                Gradient table.
            fit_method : str or callable
                str can be one of the following:

                - 'OLS' or 'ULLS' to fit the diffusion tensor and kurtosis tensor
                  using the ordinary linear least squares solution
                  `:func:dki.ols_fit_dki`
                - 'WLS' or 'UWLLS' to fit the diffusion tensor and kurtosis tensor
                  using the ordinary linear least squares solution
                  :func:`dki.wls_fit_dki`

                callable has to have the signature:
                    ``fit_method(design_matrix, data, *args, **kwargs)``

            args, kwargs : arguments and key-word arguments passed to the
               fit_method. See :func:`dki.ols_fit_dki`, :func:`dki.wls_fit_dki` for
               details

            References
            ----------
            .. footbibliography::
            """
            DiffusionKurtosisModel.__init__(
                self, gtab, fit_method=fit_method, *args, **kwargs
            )

        @warning_for_keywords()
        def fit(self, data, *, mask=None, sphere="repulsion100", gtol=1e-2, awf_only=False):
            """Fit method of the Diffusion Kurtosis Microstructural Model

            Parameters
            ----------
            data : array
                An 4D matrix containing the diffusion-weighted data.

            mask : array
                A boolean array used to mark the coordinates in the data that
                should be analyzed that has the shape data.shape[-1]

            sphere : Sphere class instance, optional
                The sphere providing sample directions for the initial search of
                the maximal value of kurtosis.

            gtol : float, optional
                This input is to refine kurtosis maxima under the precision of the
                directions sampled on the sphere class instance. The gradient of
                the convergence procedure must be less than gtol before successful
                termination. If gtol is None, fiber direction is directly taken
                from the initial sampled directions of the given sphere object

            awf_only : bool, optiomal
                If set to true only the axonal volume fraction is computed from
                the kurtosis tensor. Default = False
            """
            if mask is not None:
                # Check for valid shape of the mask
                if mask.shape != data.shape[:-1]:
                    raise ValueError("Mask is not the same shape as data.")
                mask = np.asarray(mask, dtype=bool)
            data_in_mask = np.reshape(data[mask], (-1, data.shape[-1]))

            if self.min_signal is None:
                self.min_signal = MIN_POSITIVE_SIGNAL

            data_in_mask = np.maximum(data_in_mask, self.min_signal)

            # DKI fit
            dki_params = super().fit(data_in_mask).model_params

            # Computing awf
            awf = axonal_water_fraction(dki_params, sphere=sphere, gtol=gtol)

            if awf_only:
                params_all_mask = np.concatenate((dki_params, np.array([awf]).T), axis=-1)
            else:
                # Computing the hindered and restricted diffusion tensors
                hdt, rdt = diffusion_components(dki_params, sphere=sphere, awf=awf)
                params_all_mask = np.concatenate(
                    (dki_params, np.array([awf]).T, hdt, rdt), axis=-1
                )

            if mask is None:
                out_shape = data.shape[:-1] + (-1,)
                params = params_all_mask.reshape(out_shape)
                # if extra is not None:
                #    self.extra = extra.reshape(data.shape)
            else:
                params = np.zeros(data.shape[:-1] + (params_all_mask.shape[-1],))
                params[mask, :] = params_all_mask
                # if extra is not None:
                #    self.extra = np.zeros(data.shape)
                #    self.extra[mask, :] = extra

            return KurtosisMicrostructuralFit(self, params)

        @warning_for_keywords()
        def predict(self, params, *, S0=1.0):
            """Predict a signal for the DKI microstructural model class instance
            given parameters.

            Parameters
            ----------
            params : ndarray (x, y, z, 40) or (n, 40)
                All parameters estimated from the diffusion kurtosis
                microstructural model.
                Parameters are ordered as follows:

                    1) Three diffusion tensor's eigenvalues
                    2) Three lines of the eigenvector matrix each containing the
                       first, second and third coordinates of the eigenvector
                    3) Fifteen elements of the kurtosis tensor
                    4) Six elements of the hindered diffusion tensor
                    5) Six elements of the restricted diffusion tensor
                    6) Axonal water fraction

            S0 : float or ndarray, optional
                The non diffusion-weighted signal in every voxel, or across all
                voxels.

            Notes
            -----
            In the original article of DKI microstructural model
            :footcite:p:`Fieremans2011`, the hindered and restricted tensors were
            defined as the intra-cellular and extra-cellular diffusion compartments
            respectively.

            References
            ----------
            .. footbibliography::
            """
            return dkimicro_prediction(params, self.gtab, S0=S0)

    class KurtosisMicrostructuralFit(DiffusionKurtosisFit):
        """Class for fitting the Diffusion Kurtosis Microstructural Model"""

        def __init__(self, model, model_params):
            """Initialize a KurtosisMicrostructural Fit class instance.

            Parameters
            ----------
            model : DiffusionKurtosisModel Class instance
                Class instance containing the Diffusion Kurtosis Model for the fit
            model_params : ndarray (x, y, z, 40) or (n, 40)
                All parameters estimated from the diffusion kurtosis
                microstructural model.
                Parameters are ordered as follows:

                    1) Three diffusion tensor's eigenvalues
                    2) Three lines of the eigenvector matrix each containing the
                       first, second and third coordinates of the eigenvector
                    3) Fifteen elements of the kurtosis tensor
                    4) Six elements of the hindered diffusion tensor
                    5) Six elements of the restricted diffusion tensor
                    6) Axonal water fraction

            Notes
            -----
            In the original article of DKI microstructural model
            :footcite:p:`Fieremans2011`, the hindered and restricted tensors were
            defined as the intra-cellular and extra-cellular diffusion compartments
            respectively.

            References
            ----------
            .. footbibliography::
            """
            DiffusionKurtosisFit.__init__(self, model, model_params)

        @property
        def awf(self):
            """Returns the volume fraction of the restricted diffusion compartment
            also known as axonal water fraction.

            Notes
            -----
            The volume fraction of the restricted diffusion compartment can be seen
            as the volume fraction of the intra-cellular compartment
            :footcite:p:`Fieremans2011`.

            References
            ----------
            .. footbibliography::
            """
            return self.model_params[..., 27]

        @property
        def restricted_evals(self):
            """Returns the eigenvalues of the restricted diffusion compartment.

            Notes
            -----
            The restricted diffusion tensor can be seen as the tissue's
            intra-cellular diffusion compartment :footcite:p:`Fieremans2011`.

            References
            ----------
            .. footbibliography::
            """
            self._is_awfonly()
            return _compartments_eigenvalues(self.model_params[..., 34:40])

        @property
        def hindered_evals(self):
            """Returns the eigenvalues of the hindered diffusion compartment.

            Notes
            -----
            The hindered diffusion tensor can be seen as the tissue's
            extra-cellular diffusion compartment :footcite:p:`Fieremans2011`.

            References
            ----------
            .. footbibliography::
            """
            self._is_awfonly()
            return _compartments_eigenvalues(self.model_params[..., 28:34])

        @property
        def axonal_diffusivity(self):
            """Returns the axonal diffusivity defined as the restricted diffusion
            tensor trace.

            See :footcite:p:`Fieremans2011` for further details about the method.

            References
            ----------
            .. footbibliography::
            """
            return trace(self.restricted_evals)

        @property
        def hindered_ad(self):
            """Returns the axial diffusivity of the hindered compartment.

            Notes
            -----
            The hindered diffusion tensor can be seen as the tissue's
            extra-cellular diffusion compartment :footcite:p:`Fieremans2011`.

            References
            ----------
            .. footbibliography::
            """
            return axial_diffusivity(self.hindered_evals)

        @property
        def hindered_rd(self):
            """Returns the radial diffusivity of the hindered compartment.

            Notes
            -----
            The hindered diffusion tensor can be seen as the tissue's
            extra-cellular diffusion compartment :footcite:p:`Fieremans2011`.

            References
            ----------
            .. footbibliography::
            """
            return radial_diffusivity(self.hindered_evals)

        @property
        def tortuosity(self):
            """Returns the tortuosity of the hindered diffusion which is defined
            by ADe / RDe, where ADe and RDe are the axial and radial diffusivities
            of the hindered compartment.

            See :footcite:p:`Fieremans2011` for further details about the method.

            Notes
            -----
            The hindered diffusion tensor can be seen as the tissue's
            extra-cellular diffusion compartment :footcite:p:`Fieremans2011`.

            References
            ----------
            .. footbibliography::
            """
            return tortuosity(self.hindered_ad, self.hindered_rd)

        def _is_awfonly(self):
            """To raise error if only the axonal water fraction was computed"""
            if self.model_params.shape[-1] < 39:
                raise ValueError(
                    "Only the awf was processed! Rerun model fit "
                    "with input parameter awf_only set to False"
                )

        @warning_for_keywords()
        def predict(self, gtab, *, S0=1.0):
            r"""Given a DKI microstructural model fit, predict the signal on the
            vertices of a gradient table

            gtab : a GradientTable class instance
                The gradient table for this prediction

            S0 : float or ndarray, optional
                The non diffusion-weighted signal in every voxel, or across all
                voxels.

            Notes
            -----
            The predicted signal is given by:

            .. math::

                S(\theta, b) = S_0 * [f * e^{-b ADC_{r}} + (1-f) * e^{-b ADC_{h}]

            where $ADC_{r}$ and $ADC_{h}$ are the apparent diffusion coefficients
            of the diffusion hindered and restricted compartment for a given
            direction $\theta$, $b$ is the b value provided in the GradientTable
            input for that direction, $f$ is the volume fraction of the restricted
            diffusion compartment (also known as the axonal water fraction).
            """
            self._is_awfonly()
            return dkimicro_prediction(self.model_params, gtab, S0=S0)