Accelerating Multi-Shell Multi-Tissue CSD with Ray
Multi-shell multi-tissue constrained spherical deconvolution is a powerful model for reconstructing the configuration of fibers and the volume fraction of different tissue compartments simultaneuosly (Jeurissen et al., 2014. 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.
We demonstrate this functionality here directly with the DIPY library functionality (based on an example in the DIPY documentation). This may not be runnable on this platform due to computational limitations.
from paths import afq_home
import os.path as op
import numpy as np
import matplotlib.pyplot as plt
import AFQ.data.fetch as afd
from AFQ.models.QBallTP import anisotropic_power
from dipy.core.gradients import gradient_table, unique_bvals_tolerance
from dipy.data import get_sphere
from dipy.io.gradients import read_bvals_bvecs
from dipy.io.image import load_nifti
from dipy.reconst.mcsd import (
MultiShellDeconvModel,
mask_for_response_msmt,
multi_shell_fiber_response,
response_from_mask_msmt,
)1Download dataset¶
We will use a multi-shell dataset from the HBN POD2 data-set (Richie-Halford et al., 2022). 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.
sphere = get_sphere(name="symmetric724")
study_dir = afd.fetch_hbn_preproc(["NDARAA948VFH"])[1]
sub_dir = op.join(study_dir, "derivatives/qsiprep/sub-NDARAA948VFH")
fraw = op.join(sub_dir, "ses-HBNsiteRU/dwi/sub-NDARAA948VFH_ses-HBNsiteRU_acq-64dir_space-T1w_desc-preproc_dwi.nii.gz")
fbval = op.join(sub_dir, "ses-HBNsiteRU/dwi/sub-NDARAA948VFH_ses-HBNsiteRU_acq-64dir_space-T1w_desc-preproc_dwi.bval")
fbvec = op.join(sub_dir, "ses-HBNsiteRU/dwi/sub-NDARAA948VFH_ses-HBNsiteRU_acq-64dir_space-T1w_desc-preproc_dwi.bvec")
t1_fname = op.join(sub_dir, "anat/sub-NDARAA948VFH_desc-preproc_T1w.nii.gz")
brain_mask = op.join(sub_dir, "anat/sub-NDARAA948VFH_desc-brain_mask.nii.gz")
gm_seg = op.join(sub_dir, "anat/sub-NDARAA948VFH_space-MNI152NLin2009cAsym_label-GM_probseg.nii.gz")
wm_seg = op.join(sub_dir, "anat/sub-NDARAA948VFH_space-MNI152NLin2009cAsym_label-WM_probseg.nii.gz")
csf_seg = op.join(sub_dir, "anat/sub-NDARAA948VFH_space-MNI152NLin2009cAsym_label-CSF_probseg.nii.gz")data, affine = load_nifti(fraw)
bvals, bvecs = read_bvals_bvecs(fbval, fbvec)
gtab = gradient_table(bvals, bvecs=bvecs)csf = np.where(load_nifti(csf_seg)[0] > 0.5, 1, 0)
gm = np.where(load_nifti(gm_seg)[0] > 0.5, 1, 0)
wm = np.where(load_nifti(wm_seg)[0] > 0.5, 1, 0)2Estimate response functions¶
mask_wm, mask_gm, mask_csf = mask_for_response_msmt(
gtab,
data,
roi_radii=10,
wm_fa_thr=0.7,
gm_fa_thr=0.3,
csf_fa_thr=0.15,
gm_md_thr=0.001,
csf_md_thr=0.0032,
)
response_wm, response_gm, response_csf = response_from_mask_msmt(
gtab, data, mask_wm, mask_gm, mask_csf
)
print(response_wm)
print(response_gm)
print(response_csf)3Reconstruction with MSMT-CSD¶
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).
ubvals = unique_bvals_tolerance(gtab.bvals)
response_mcsd = multi_shell_fiber_response(
sh_order_max=8,
bvals=ubvals,
wm_rf=response_wm,
gm_rf=response_gm,
csf_rf=response_csf,
)
mcsd_model = MultiShellDeconvModel(gtab, response_mcsd)
mcsd_fit = mcsd_model.fit(data[:, :, 50], engine="ray") # Using a subset of the data for speed in this example
# We can use the anisotropic power map to visualize the fit
plt.imshow(anisotropic_power(mcsd_fit.shm_coeff))- Jeurissen, B., Tournier, J.-D., Dhollander, T., Connelly, A., & Sijbers, J. (2014). Multi-tissue constrained spherical deconvolution for improved analysis of multi-shell diffusion MRI data. NeuroImage, 103, 411–426. 10.1016/j.neuroimage.2014.07.061
- Richie-Halford, A., Cieslak, M., Ai, L., Caffarra, S., Covitz, S., Franco, A. R., Karipidis, I. I., Kruper, J., Milham, M., Avelar-Pereira, B., Roy, E., Sydnor, V. J., Yeatman, J. D., Abbott, N. J., Anderson, J. A. E., Gagana, B., Bleile, M., Bloomfield, P. S., Bottom, V., … Rokem, A. (2022). An analysis-ready and quality controlled resource for pediatric brain white-matter research. Scientific Data, 9(1). 10.1038/s41597-022-01695-7