#!/usr/bin/env python
# coding: utf-8

# # Generalized Constrained Spherical Deconvolution

# Dmipy's Constrained Spherical Deconvolution (CSD) implementation completely generalizes the Multi-Shell-Multi-Tissue (MSMT) CSD implementation of *(Jeurissen et al. 2014)* to
# - Any number of compartments,
# - Any acquisition scheme parameters,
# - Being able to use voxel-varying convolution kernels (e.g. as estimated by spherical mean models).
# 
# ## Short Summary of CSD:
# Constrained Spherical Deconvolution (CSD) *(Tournier et al. 2007)* is one of the most well-known ways to estimate Fiber Orientation Distributions (FODs) that describe the orientation of the white matter tissue. In short, CSD's formulation states that any composition of oriented white matter (dispersion/crossings) can be described as the spherical convolution of a *positive* probability density on the sphere and a convolution kernel $K$ that describes one parallel axon bundle. With some abuse of notation:
# 
# $$
# \begin{align}
# \begin{aligned}
#  E_{\textrm{CSD}}= \overbrace{\operatorname{FOD}(\textbf{c})}^{\textrm{Fiber Distribution}}\,*_{\mathbb{S}^2}\,\overbrace{K(\cdot)}^{\textrm{Convolution Kernel}}\quad\textrm{subject to}\quad \operatorname{FOD}(\textbf{c}) > 0.
# \end{aligned}
# \end{align}$$
# 
# Here, the FOD is described in terms of a truncated even spherical harmonics series $FOD=\sum_{l=0}^{lmax}\sum_{m=-l}^l\textbf{c}_{lm}Y_{lm}$, with $l$ and $m$ describing the order and moment of the spherical harmonic *(Descoteaux et al. 2006)*. Furthermore, $*_{\mathbb{S}^2}$ describes a spherical convolution on the $\mathbb{S}^2$ sphere, and $K$ describes the *rotational* harmonics describing a single axon bundle. The kernel $K$ is typically estimated from the data *(Tournier et al. 2007, Tax et al. 2014)*, and CSD is therefore considered a *model-free* approach.

# ## Multi-Shell Single-Compartment CSD
# When using a single convolution kernel (voxel-varying or not), Dmipy's uses the classical (faster) optimizer proposed by *(Tournier et al. 2007)* to estimate the spherical harmonics of the FOD. The single- and multi-shell implementations of CSD estimate a positive and smooth FOD as follows:

# \begin{equation}
# \textbf{f}^* = argmin_f \overbrace{\|\textbf{Af} - \textbf{b}\|^2}^{\textrm{Data Term}} + \lambda_{\textrm{pos}}\overbrace{\|\textbf{Lf}\|^2}^{\textrm{Positivity}} + \lambda_{\textrm{lb}}\overbrace{\|\textbf{Rf}\|^2}^{\textrm{smoothness}},
# \end{equation}
# 
# where $\textbf{f}$ is SH coefficients of length $N_c$. $\textbf{A}$ of size $N_{data}\times N_c$ is observation matrix that maps spherical harmonics coefficients to signal values, $\textbf{b}$ is array of measured DWIs, and $\textrm{L}$ maps $\textbf{f}$ on the sphere and penalizes negative values, and $\textbf{R}$ is the Laplace-Beltrami operator following *(Descoteaux et al. 2006)*.
# 
# The observation matrix $\textbf{A}$ is - *for each shell separately* - formed by the product of a spherical-harmonics transform (SHT) matrix with a diagonal matrix with prepared rotational harmonics of the convolution kernel, meaning they are multiplied by $\sqrt{(4 * \pi) / (2 * l + 1)}$.
# 
# The positivity matrix $\textbf{L}$ is a spherical harmonics transform (SHT) matrix that maps spherical harmonics coefficients to real values on the sphere, for which rows mapping to positive FOD values are zeroed out.
# 
# The smoothness Laplace-Beltrami matrix $\textbf{R}$ is a diagonal matrix with values $l(l+1)$.
# 
# The minimum of this equation is found by taking the derivative w.r.t. $\textbf{f}$ and setting is to zero
# 
# \begin{equation}
# 0 = 2\textbf{A}^T(\textbf{Af}^* - \textbf{b}) + 2\lambda_{\textrm{pos}} \textbf{L}^T\textbf{L}\textbf{f}^* + 2\lambda_{\textrm{lb}} \textbf{R}^T\textbf{R}\textbf{f}^*
# \end{equation}
# 
# \begin{equation}
# \textbf{A}^T\textbf{b} = (\textbf{A}^T\textbf{A} + \lambda_{\textrm{pos}} \textbf{L}^T\textbf{L} + \lambda_{\textrm{lb}} \textbf{R}^T\textbf{R})\textbf{f}^*
# \end{equation}
# 
# \begin{equation}
# \textbf{f}^* = (\textbf{A}^T\textbf{A} + \lambda_{\textrm{pos}} \textbf{L}^T\textbf{L} + \lambda_{\textrm{lb}} \textbf{R}^T\textbf{R})^{-1}\textbf{A}^T\textbf{b}
# \end{equation}
# 
# In the CSD optimization following *(Tournier et al. 2007)*, this equation is solved iteratively by starting with a lower spherical harmonics order FOD (typically 4), and and updating the positivity constraint matrix until convergence.

# ## Multi-Shell Multi-Compartment CSD

# The multi-compartment implementation of CSD (MC-CSD) follows the work of *(Jeurissen et al. 2014)* and is exactly the same as the single-compartment one, but has some constraints:
# - When estimating volume fractions, MC-CSD can only estimate the FODs for one anisotropic kernel and one or more isotropic kernels.
# - An additional constraint must be added that the volume fractions add up to one and are positive.
# 
# For every isotropic compartment, one extra line on $\textbf{A}$ is added that maps the isotropic signal attenuation to the data positions. Similarly, one extra coefficient is added at the end of $\textbf{f}$, which represents the isotropic compartment's only $l=0, m=0$ coefficient. The volume fractions of each compartment are defined by their respects $l=0, m=0$ coefficients, multiplied by the sphere's jacobian $2 * \sqrt{\pi}$. Using cvxpy, the optimization is cast as
# 
# \begin{equation}
# \textbf{f}^* = argmin_f \overbrace{\|\textbf{Af} - \textbf{b}\|^2}^{\textrm{Data Term}} + \lambda_{\textrm{lb}}\overbrace{\|\textbf{Rf}\|^2}^{\textrm{smoothness}},\quad\textrm{subject to}\quad \sum vf=1, vf>0, FOD>0
# \end{equation}

# ## Voxel-Varying Convolution Kernels
# Just like in previous models, it is possible to set voxel-varying initial conditions when fitting models. This can be done by hand, or particularly, a fitted `MultiCompartmentSphericalMeanModel` can call the `MultiCompartmentSphericalHarmonicsModel` with all its parameters fixed voxel-wise to the spherical mean fitted parameters. See e.g. the [Spherical Mean Technique](http://nbviewer.jupyter.org/github/AthenaEPI/dmipy/blob/master/examples/example_spherical_mean_technique.ipynb) or [Multi-Compartment Microscopic Diffusion Imaging examples](http://nbviewer.jupyter.org/github/AthenaEPI/mipy/blob/master/examples/example_multi_compartment_spherical_mean_technique.ipynb) examples.

# ## References
# - Tournier, J-Donald, Fernando Calamante, and Alan Connelly. "Robust determination of the fibre orientation distribution in diffusion MRI: non-negativity constrained super-resolved spherical deconvolution." Neuroimage 35.4 (2007): 1459-1472.
# - Descoteaux, Maxime, et al. "Regularized, fast, and robust analytical Q‐ball imaging." Magnetic resonance in medicine 58.3 (2007): 497-510.
# - Tax, Chantal MW, et al. "Recursive calibration of the fiber response function for spherical deconvolution of diffusion MRI data." Neuroimage 86 (2014): 67-80.
# - Jeurissen, Ben, et al. "Multi-tissue constrained spherical deconvolution for improved analysis of multi-shell diffusion MRI data." NeuroImage 103 (2014): 411-426.