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

# # AxCaliber

# To estimate the axon diameter distribution itself, *(Assaf et al. 2008)* proposed a method called AxCaliber.
# In AxCaliber, a distribution of cylinders with different diameters are fitted to the diffusion signal. In particular, a Gamma distribution is used. A requirement of this technique is that the signal is measured perpendicular to the axon bundle direction for different gradient strengths and diffusion times.
# The signal representation in AxCaliber is given as
# 
# \begin{equation}
# E_{\textrm{AxCaliber}}= \underbrace{(1-f_r)\overbrace{E_h(\lambda_{\textrm{h}})}^{\textrm{Ball}}}_{\textrm{Extra-Axonal}} + \underbrace{f_r \overbrace{\Gamma(\alpha,\beta)}^{\textrm{Gamma Distribution}}*_{\mathbb{R}}\overbrace{E_r(\cdot|D_\perp)}^{\textrm{Cylinder}}}_{\textrm{Intra-Axonal}}
# \end{equation}
# 
# where the cylinder was initially represented using the Callaghan model, but was later replaced with the more general Van Gelderen model *(Assaf et al. 2008, Huang et al. 2015, De Santis et al. 2016)*.
# 
# Advantages:
# - Estimates the axon diameter distribution.
# 
# Limitations:
# - Tissue orientation must be known for the signal to be measured perpendicular,
# - Does not take axon dispersion into account,
# - Does not include the effects of extra axonal restriction *(Burcaw et al. 2015)*.

# ## AxCaliber model with Gaussian Phase Approximation

# The process of setting up the AxCaliber model with Gamma-distributed cylinder diameters is the same as setting up an axon dispersed model. We import the needed Ball and Cylinder models, and then distribute the cylinder using the DD1GammaDistributed function from distribute_models.

# In[1]:


from dmipy.distributions import distribute_models
from dmipy.signal_models import gaussian_models, cylinder_models

ball = gaussian_models.G1Ball()
cylinder = cylinder_models.C4CylinderGaussianPhaseApproximation()
gamma_cylinder = distribute_models.DD1GammaDistributed(models=[cylinder])


# Then put the models together in a MultiCompartmentModel and display the parameters

# In[2]:


from dmipy.core import modeling_framework
axcaliber_gamma = modeling_framework.MultiCompartmentModel(models=[ball, gamma_cylinder])
axcaliber_gamma.parameter_cardinality


# Visualize the model:

# In[3]:


from IPython.display import Image
axcaliber_gamma.visualize_model_setup(view=False, cleanup=False, with_parameters=True)
Image('Model Setup.png')


# The data we will fit is measured only along 2 gradient orientations that are perpendicular to the axon axis (along x and y axis). This means there is no data that provides information to fit parallel diffusivity "lambda_par" or the axon orientation "mu". To not fit these superfluous parameters we fix mu along the z-axis and give lambda_par an arbitrary value (it will not affect the results).

# In[3]:


axcaliber_gamma.set_fixed_parameter(
    'DD1GammaDistributed_1_C4CylinderGaussianPhaseApproximation_1_lambda_par', 1.7e-9)
axcaliber_gamma.set_fixed_parameter(
    'DD1GammaDistributed_1_C4CylinderGaussianPhaseApproximation_1_mu', [0, 0])


# ## Load AxCaliber-compatible data

# For the example data, we kindly make use of a cat spinal cord dataset hosted at the White Matter Microscopy Database (https://osf.io/yp4qg/). The data includes 2D Axcaliber data, 3D multi-shell data and the corresponding histology values for axon diameter, intra-axonal volume fraction (and others).

# In[4]:


from dmipy.data import saved_data
scheme_spinal_cord, data_spinal_cord = saved_data.duval_cat_spinal_cord_2d()


# In[5]:


import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
plt.imshow(data_spinal_cord.signal[:,:,0,0])
plt.title('b0 of spinal cord data')
plt.axis('off');


# Printing the acquisition scheme information we can see the Delta-shells of the acquisition.

# In[6]:


scheme_spinal_cord.print_acquisition_info


# The 2D data only has data measured in the xy-plane, where we only use the [1, 0, 0] dummy direction for b0 measurements. We can see the unique gradient directions as follows:

# In[7]:


import numpy as np
unique_directions = np.unique(scheme_spinal_cord.gradient_directions, axis=0)
unique_directions


# It is valuable to check if actually the signal between the two orthogonal directions is the same:

# In[8]:


dir1mask = np.all(scheme_spinal_cord.gradient_directions == unique_directions[0], axis=1)
dir2mask = np.all(scheme_spinal_cord.gradient_directions == unique_directions[1], axis=1)

test_voxel = data_spinal_cord.signal[20, 20, 0]
data_dir1 = test_voxel[dir1mask]
data_dir2 = test_voxel[dir2mask]

Delta_dir1 = scheme_spinal_cord.Delta[dir1mask]
Delta_dir2 = scheme_spinal_cord.Delta[dir2mask]

G = scheme_spinal_cord.gradient_strengths
colors = ['r', 'g', 'b', 'orange']

unique_Deltas = np.unique(scheme_spinal_cord.shell_Delta)[2:-1:2]
for i, Delta_ in enumerate(unique_Deltas):
    mask_Delta1 = Delta_dir1 == Delta_
    mask_Delta2 = Delta_dir2 == Delta_

    plt.plot(G[dir1mask][mask_Delta1], data_dir1[mask_Delta1], c=colors[i], label='dir1 Delta='+str(Delta_))
    plt.plot(G[dir2mask][mask_Delta2], data_dir2[mask_Delta2], c=colors[i], ls='--', label='dir2 Delta='+str(Delta_))
plt.legend()
plt.xlabel('Gradient Strength [T/m]')
plt.ylabel('Signal [-]');


# It can be seen that the two measured perpendicular directions experience slightly different levels of diffusion restriction; the dir2 signal profiles attenuate faster than the dir1 ones. This could implies that that some kind of anisotropic axon dispersion could be happening, or that the axon bundle is slightly tilted towards dir2.
# 
# Dmipy allows us to add axon dispersion into the AxCaliber model to compensate for this, but we'll ignore that for this example and fit the signal as-is.

# ## Fit AxCaliber to spinal cord data

# It is known that a Gamma distribution is tricky to fit since there are several parameter combinations that produce similar distributions. We will use the stochastic MIX optimization *(Farooq et al. 2016)*, which is slow but is likely to robustly give the global minimum.
# 
# #### NOTE: Fitting AxCaliber to this dataset with MIX takes over 7 hours!

# In[9]:


axcaliber_gamma_fit = axcaliber_gamma.fit(
    scheme_spinal_cord, data_spinal_cord.signal,
    mask=data_spinal_cord.mask, solver='mix', maxiter=100)


# ## Visualize fitted parameter maps

# We can visualize the fitted parameters as before:

# In[10]:


fitted_parameters = axcaliber_gamma_fit.fitted_parameters

fig, axs = plt.subplots(2, 3, figsize=[15, 10])
axs = axs.ravel()

counter = 0
for name, values in fitted_parameters.items():
    cf = axs[counter].imshow(values.squeeze().T, origin=True, interpolation='nearest')
    axs[counter].set_title(name)
    axs[counter].set_axis_off()
    fig.colorbar(cf, ax=axs[counter], shrink=0.5)
    counter += 1


# We can see that the maps for alpha and beta are not super smooth. This could be because we're actually fitting two signals experiencing different restriction levels at the same time, but hard to say without fitting each signal separately. We can visualize the distributions themselves as follows:

# In[15]:


axcaliber_alpha = axcaliber_gamma_fit.fitted_parameters['DD1GammaDistributed_1_DD1Gamma_1_alpha']
axcaliber_beta = axcaliber_gamma_fit.fitted_parameters['DD1GammaDistributed_1_DD1Gamma_1_beta']

from dmipy.distributions.distributions import DD1Gamma
gamma = DD1Gamma()

x = 15 # some x-pos in the data
y = 25 # some y-pos in the data

alpha = axcaliber_alpha[x, y, 0]
beta = axcaliber_beta[x, y, 0]

radii, Pgamma = gamma(alpha=alpha, beta=beta)
plt.plot(2 * radii * 1e6, Pgamma)
plt.xlabel('cylinder diameter [$\mu$m]')
plt.ylabel('Gamma probability density')
plt.title('Estimated Gamma distribution');


# ## Comparison between estimated and histology axon diameter

# The mean axon diameter can be estimated from the axon diameter distribution as diameter = 2$\alpha\beta$. We compare the histological axon diameter with the estimated ones:

# In[16]:


mean_diameter = 2 * axcaliber_alpha * axcaliber_beta * 1e6  # in microns
mean_diameter_histology = data_spinal_cord.histology.h1_axonEquivDiameter  # was already in microns


# In[17]:


fig, [ax1, ax2] = plt.subplots(1, 2, figsize=[10, 5])

ax1plot = ax1.imshow(mean_diameter.squeeze().T,
           origin=True, vmax=6, interpolation='nearest')
ax1.set_title('Estimated Mean Axon Diameter')
ax1.set_axis_off()
plt.colorbar(ax1plot, ax=ax1, shrink=0.75)

ax2plot = ax2.imshow(mean_diameter_histology.T, 
           origin=True, vmax=6, interpolation='nearest')
ax2.set_title('Histology Mean Axon Diameter')
ax2.set_axis_off()
plt.colorbar(ax2plot, ax=ax2, shrink=0.75)


# We can see that our estimated mean axon diameter map is completely smooth, aside from a few outliers in areas which actually lie outside of the white matter area (see right histology image). The estimated diameters are overestimated in general. We find between 4-6 $\mu$m, while the histological values are between 0-5 $\mu$m. 
# 
# Using the Seaborn package, we can nicely visualize the correlation between our estimated and the histology diameters. We separate visualize this correlation of axon diameters between 1-2$\mu$m, 2-3$\mu$m and 3+$\mu$m, and remove the exceptional outliers.

# In[18]:


mean_diameter_mask12 = np.all([mean_diameter_histology > 1,
                               mean_diameter_histology < 2,
                               mean_diameter.squeeze() < 10], axis=0)
mean_diameter_mask23 = np.all([mean_diameter_histology > 2,
                               mean_diameter_histology < 3,
                               mean_diameter.squeeze() < 10], axis=0)
mean_diameter_mask3p = np.all([mean_diameter_histology > 3,
                               mean_diameter.squeeze() < 10], axis=0)


# ### Axon diameter correlation 1-2 $\mu$m - not significant

# In[19]:


import seaborn as sns
mask = mean_diameter_mask12
im = sns.jointplot(mean_diameter[mask].squeeze(),
                   mean_diameter_histology[mask].squeeze(), kind="reg")
im.set_axis_labels("Estimated Axon Diameter", "Histology Axon Diameter", fontsize=15)


# ### Axon diameter correlation 2-3 $\mu$m - not significant

# In[20]:


mask = mean_diameter_mask23
im = sns.jointplot(mean_diameter[mask].squeeze(),
                   mean_diameter_histology[mask].squeeze(), kind="reg", color='r')
im.set_axis_labels("Estimated Axon Diameter", "Histology Axon Diameter", fontsize=15)


# ### Axon diameter correlation 3+ $\mu$m - Significant

# In[21]:


mask = mean_diameter_mask3p
im = sns.jointplot(mean_diameter[mask].squeeze(),
                   mean_diameter_histology[mask].squeeze(), kind="reg", color='g')
im.set_axis_labels("Estimated Axon Diameter", "Histology Axon Diameter", fontsize=15)


# The results show that for our current approach to the data set we can only estimate axon diameters that correlate with the histology values for axon diameters of 3+$\mu$m.

# ## Comparison between estimated and histology volume fraction

# Similarly, we can also visualize the estimated intra-axonal volume fraction between our model and the histology values.

# In[22]:


axcaliber_fraction = axcaliber_gamma_fit.fitted_parameters['partial_volume_1'].squeeze()
histology_fraction = data_spinal_cord.histology.h4_fr


# In[23]:


fig, [ax1, ax2] = plt.subplots(1, 2, figsize=[10, 5])

ax1plot = ax1.imshow(axcaliber_fraction.T,
           origin=True,interpolation='nearest', vmax=0.55)
ax1.set_title('Estimated Volume Fraction')
ax1.set_axis_off()
plt.colorbar(ax1plot, ax=ax1, shrink=0.75)

ax2plot = ax2.imshow(histology_fraction.T, 
           origin=True, interpolation='nearest', vmax=0.55)
ax2.set_title('Histology Volume Fraction')
ax2.set_axis_off()
plt.colorbar(ax2plot, ax=ax2, shrink=0.75)


# In[24]:


im = sns.jointplot(axcaliber_fraction[mean_diameter_histology > 1],
              histology_fraction[mean_diameter_histology > 1], kind="reg", color='y')
im.set_axis_labels("Estimated Volume Fraction", "Histology Volume Fraction", fontsize=15);


# The results show that estimating intra-axonal volume fractions that correlate with histology is easier than doing the same for axon diameter. However, the model does systematically underestimate the value compared to histology.

# ## References
# - Assaf, Yaniv, et al. "AxCaliber: a method for measuring axon diameter distribution from diffusion MRI." Magnetic resonance in medicine 59.6 (2008): 1347-1354.
# - Burcaw, Lauren M., Els Fieremans, and Dmitry S. Novikov. "Mesoscopic structure of neuronal tracts from time-dependent diffusion." NeuroImage 114 (2015): 18-37.
# - De Santis, Silvia, Derek K. Jones, and Alard Roebroeck. "Including diffusion time dependence in the extra-axonal space improves in vivo estimates of axonal diameter and density in human white matter." NeuroImage 130 (2016): 91-103.
# - Farooq, Hamza, et al. "Microstructure imaging of crossing (MIX) white matter fibers from diffusion MRI." Scientific reports 6 (2016): 38927.
# - Huang, Susie Y., et al. "The impact of gradient strength on in vivo diffusion MRI estimates of axon diameter." NeuroImage 106 (2015): 464-472.