{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cb2c2324",
   "metadata": {},
   "source": [
    "# Accelerating Multi-Shell Multi-Tissue CSD with Ray\n",
    "\n",
    "Multi-shell multi-tissue constrained spherical deconvolution is a powerful model for reconstructing the configuration \n",
    "of fibers and the volume fraction of different tissue compartments simultaneuosly ([Jeurissen et al., 2014](https://doi.org/10.1016/j.neuroimage.2014.07.061). However, because it requires convex optimization to be executed at every voxel, it can also be a performance bottleneck. This example demonstrates how to fit Multi-Shell Multi-Tissue Constrained Spherical Deconvolution (MSMT-CSD), while using Ray for parallelization to accelerate processing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77d26e18-2402-458c-8b07-3224bdac9df6",
   "metadata": {},
   "source": [
    "We demonstrate this functionality here directly with the DIPY library functionality (based on an [example](https://docs.dipy.org/stable/examples_built/reconstruction/reconst_mcsd.html) in the DIPY documentation). This may not be runnable on this platform due to computational limitations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6a43323",
   "metadata": {},
   "source": [
    "```python\n",
    "from paths import afq_home\n",
    "\n",
    "import os.path as op\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import AFQ.data.fetch as afd\n",
    "from AFQ.models.QBallTP import anisotropic_power\n",
    "\n",
    "from dipy.core.gradients import gradient_table, unique_bvals_tolerance\n",
    "from dipy.data import get_sphere\n",
    "from dipy.io.gradients import read_bvals_bvecs\n",
    "from dipy.io.image import load_nifti\n",
    "from dipy.reconst.mcsd import (\n",
    "    MultiShellDeconvModel,\n",
    "    mask_for_response_msmt,\n",
    "    multi_shell_fiber_response,\n",
    "    response_from_mask_msmt,\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c9e9e44",
   "metadata": {},
   "source": [
    "## Download dataset\n",
    "\n",
    "We will use a multi-shell dataset from the HBN POD2 data-set ([Richie-Halford et al., 2022](https://doi.org/10.1038/s41597-022-01695-7)). This dataset also includes T1-weighted data and tissue-type segmentations that can be used to constrain the response function that is calculated for MSMT. For simplicity, we will use the functionality of DIPY without using this information, but for completeness, we point out here that it could be used to restrict the regions accessed by the code that computes the response function for a more refined response function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f02ded2-7333-41d8-87ba-0bfa055dc8b4",
   "metadata": {},
   "source": [
    "```python\n",
    "sphere = get_sphere(name=\"symmetric724\")\n",
    "study_dir = afd.fetch_hbn_preproc([\"NDARAA948VFH\"])[1]\n",
    "sub_dir = op.join(study_dir, \"derivatives/qsiprep/sub-NDARAA948VFH\")\n",
    "\n",
    "fraw = op.join(sub_dir, \"ses-HBNsiteRU/dwi/sub-NDARAA948VFH_ses-HBNsiteRU_acq-64dir_space-T1w_desc-preproc_dwi.nii.gz\")\n",
    "fbval = op.join(sub_dir, \"ses-HBNsiteRU/dwi/sub-NDARAA948VFH_ses-HBNsiteRU_acq-64dir_space-T1w_desc-preproc_dwi.bval\")\n",
    "fbvec = op.join(sub_dir, \"ses-HBNsiteRU/dwi/sub-NDARAA948VFH_ses-HBNsiteRU_acq-64dir_space-T1w_desc-preproc_dwi.bvec\")\n",
    "t1_fname = op.join(sub_dir, \"anat/sub-NDARAA948VFH_desc-preproc_T1w.nii.gz\")\n",
    "brain_mask = op.join(sub_dir, \"anat/sub-NDARAA948VFH_desc-brain_mask.nii.gz\")\n",
    "gm_seg = op.join(sub_dir, \"anat/sub-NDARAA948VFH_space-MNI152NLin2009cAsym_label-GM_probseg.nii.gz\")\n",
    "wm_seg = op.join(sub_dir, \"anat/sub-NDARAA948VFH_space-MNI152NLin2009cAsym_label-WM_probseg.nii.gz\")\n",
    "csf_seg = op.join(sub_dir, \"anat/sub-NDARAA948VFH_space-MNI152NLin2009cAsym_label-CSF_probseg.nii.gz\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f7eca50-6186-41d1-8b68-105fbdf31090",
   "metadata": {},
   "source": [
    "```python\n",
    "data, affine = load_nifti(fraw)\n",
    "bvals, bvecs = read_bvals_bvecs(fbval, fbvec)\n",
    "gtab = gradient_table(bvals, bvecs=bvecs)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c2671c",
   "metadata": {},
   "source": [
    "```python\n",
    "csf = np.where(load_nifti(csf_seg)[0] > 0.5, 1, 0)\n",
    "gm = np.where(load_nifti(gm_seg)[0] > 0.5, 1, 0)\n",
    "wm = np.where(load_nifti(wm_seg)[0] > 0.5, 1, 0)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf9eed8f",
   "metadata": {},
   "source": [
    "## Estimate response functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8722ab2",
   "metadata": {},
   "source": [
    "```python\n",
    "mask_wm, mask_gm, mask_csf = mask_for_response_msmt(\n",
    "    gtab,\n",
    "    data,\n",
    "    roi_radii=10,\n",
    "    wm_fa_thr=0.7,\n",
    "    gm_fa_thr=0.3,\n",
    "    csf_fa_thr=0.15,\n",
    "    gm_md_thr=0.001,\n",
    "    csf_md_thr=0.0032,\n",
    ")\n",
    "\n",
    "response_wm, response_gm, response_csf = response_from_mask_msmt(\n",
    "    gtab, data, mask_wm, mask_gm, mask_csf\n",
    ")\n",
    "\n",
    "print(response_wm)\n",
    "print(response_gm)\n",
    "print(response_csf)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "435f52c4",
   "metadata": {},
   "source": [
    "## Reconstruction with MSMT-CSD\n",
    "\n",
    "Finally, this code fits the MSMT-CSD model to the data. Using `engine=\"ray\"` tells DIPY that the fit should be parallelized across chunks of voxels. This can result in substantial speedup (see article figures for details)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd78a0d",
   "metadata": {},
   "source": [
    "```python\n",
    "ubvals = unique_bvals_tolerance(gtab.bvals)\n",
    "response_mcsd = multi_shell_fiber_response(\n",
    "    sh_order_max=8,\n",
    "    bvals=ubvals,\n",
    "    wm_rf=response_wm,\n",
    "    gm_rf=response_gm,\n",
    "    csf_rf=response_csf,\n",
    ")\n",
    "\n",
    "mcsd_model = MultiShellDeconvModel(gtab, response_mcsd)\n",
    "mcsd_fit = mcsd_model.fit(data[:, :, 50], engine=\"ray\") # Using a subset of the data for speed in this example\n",
    "\n",
    "# We can use the anisotropic power map to visualize the fit\n",
    "plt.imshow(anisotropic_power(mcsd_fit.shm_coeff))\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}