{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# GPU-accelerated tractography\n", "\n", "## Introduction\n", "\n", "The process of tractography is a major bottleneck in executing tractometry pipelines. GPUs provide an opportunity for massive parallelization and opportunities to accelerate tractography substantially, reducing this bottleneck and enabling rapid large-scale analysis. We have written the `GPUStreamlines` library for running tractography on NVIDIA GPUs.\n", "\n", "The code includes Python bindings, and if the CUDA Toolkit and gcc are already installed, you can install GPUStreamlines by running the standard:\n", "\n", "```bash\n", "pip install git+https://github.com/dipy/GPUStreamlines.git\n", "```\n", "\n", "GPUStreamlines can also be run from pyAFQ itself if it is installed. It is most conveniently installed using Docker containers. We build and distribute containers for each pyAFQ version, which you can find here:\n", "\n", ":::{seealso} Docker containers\n", "https://github.com/orgs/nrdg/packages/container/package/pyafq_gpu_cuda_12\n", ":::\n", "\n", "This container includes both the GPUStreamlines software as well as pyAFQ. GPU-accelerated tractography can then be run by adding the `tractography_ngpus` argument to either the `GroupAFQ` or `ParticipantAFQ` APIs. The scripts included below demonstrate how to run GPUStreamlines directly if you do not want to run it through pyAFQ. The example is based on the script found here:\n", "\n", ":::{seealso} Scripts to run GPUStreamlines directly\n", "https://github.com/dipy/GPUStreamlines/blob/master/run_gpu_streamlines.py\n", ":::" ] }, { "cell_type": "markdown", "id": "2dff2131", "metadata": {}, "source": [ "## Importing Required Libraries\n", "\n", "```python\n", "import argparse\n", "import random\n", "import time\n", "import zipfile\n", "\n", "import numpy as np\n", "import numpy.linalg as npl\n", "\n", "import dipy.reconst.dti as dti\n", "from dipy.io import read_bvals_bvecs\n", "from dipy.io.stateful_tractogram import Origin, Space, StatefulTractogram\n", "from dipy.io.streamline import save_tractogram\n", "from dipy.tracking import utils\n", "from dipy.core.gradients import gradient_table, unique_bvals_magnitude\n", "from dipy.data import default_sphere\n", "from dipy.direction import (BootDirectionGetter, ProbabilisticDirectionGetter, PTTDirectionGetter)\n", "from dipy.reconst.shm import OpdtModel, CsaOdfModel\n", "from dipy.reconst.csdeconv import ConstrainedSphericalDeconvModel, auto_response_ssst\n", "from dipy.tracking.local_tracking import LocalTracking\n", "from dipy.tracking.stopping_criterion import ThresholdStoppingCriterion\n", "from dipy.reconst import shm\n", "from dipy.data import get_fnames\n", "from dipy.data import read_stanford_pve_maps\n", "\n", "import nibabel as nib\n", "from nibabel.orientations import aff2axcodes\n", "\n", "from trx.trx_file_memmap import TrxFile, zip_from_folder\n", "\n", "# Import custom module\n", "import cuslines\n", "```" ] }, { "cell_type": "markdown", "id": "0149dad4", "metadata": {}, "source": [ "## Argument Parsing\n", "\n", "```python\n", "print(\"parsing arguments\")\n", "parser = argparse.ArgumentParser()\n", "parser.add_argument(\"nifti_file\", nargs='?', default='hardi', help=\"path to the DWI nifti file\")\n", "parser.add_argument(\"bvals\", nargs='?', default='hardi', help=\"path to the bvals\")\n", "parser.add_argument(\"bvecs\", nargs='?', default='hardi', help=\"path to the bvecs\")\n", "parser.add_argument(\"mask_nifti\", nargs='?', default='hardi', help=\"path to the mask file\")\n", "parser.add_argument(\"roi_nifti\", nargs='?', default='hardi', help=\"path to the ROI file\")\n", "parser.add_argument(\"--device\", type=str, default='gpu', choices=['cpu', 'gpu'], help=\"Whether to use CPU or GPU\")\n", "parser.add_argument(\"--output-prefix\", type=str, default='', help=\"Path to the output file\")\n", "parser.add_argument(\"--chunk-size\", type=int, default=100000, help=\"Number of seeds to process per sweep, per GPU\")\n", "parser.add_argument(\"--nseeds\", type=int, default=100000, help=\"Total number of seeds to process\")\n", "parser.add_argument(\"--ngpus\", type=int, default=1, help=\"Number of GPUs to use\")\n", "parser.add_argument(\"--write-method\", type=str, default=\"fast\", help=\"Can be trx, fast, or standard\")\n", "parser.add_argument(\"--max-angle\", type=float, default=60, help=\"Maximum angle (in degrees)\")\n", "parser.add_argument(\"--min-signal\", type=float, default=1.0, help=\"Minimum signal threshold\")\n", "parser.add_argument(\"--step-size\", type=float, default=0.5, help=\"Step size\")\n", "parser.add_argument(\"--sh-order\", type=int, default=4, help=\"Spherical harmonics order\")\n", "parser.add_argument(\"--fa-threshold\", type=float, default=0.1, help=\"FA threshold\")\n", "parser.add_argument(\"--relative-peak-threshold\", type=float, default=0.25, help=\"Relative peak threshold\")\n", "parser.add_argument(\"--min-separation-angle\", type=float, default=45, help=\"Minimum separation angle (in degrees)\")\n", "parser.add_argument(\"--sm-lambda\", type=float, default=0.006, help=\"Smoothing lambda\")\n", "parser.add_argument(\"--model\", type=str, default=\"opdt\", choices=['opdt', 'csa', 'csd'], help=\"Model to use\")\n", "parser.add_argument(\"--dg\", type=str, default=\"boot\", choices=['boot', 'prob', 'ptt'], help=\"Direction getting scheme\")\n", "\n", "args = parser.parse_args()\n", "```" ] }, { "cell_type": "markdown", "id": "a2596dd4", "metadata": {}, "source": [ "## Processing Diffusion Data\n", "\n", "```python\n", "t0 = time.time()\n", "\n", "np.random.seed(0)\n", "random.seed(0)\n", "\n", "def get_gtab(fbval, fbvec):\n", " bvals, bvecs = read_bvals_bvecs(fbval, fbvec)\n", " return gradient_table(bvals, bvecs)\n", "\n", "def get_img(ep2_seq):\n", " return nib.load(ep2_seq)\n", "\n", "if 'hardi' in [args.nifti_file, args.bvals, args.bvecs, args.mask_nifti, args.roi_nifti]:\n", " if not all(arg == 'hardi' for arg in [args.nifti_file, args.bvals, args.bvecs, args.mask_nifti, args.roi_nifti]):\n", " raise ValueError(\"If any argument is 'hardi', all must be 'hardi'\")\n", "\n", " hardi_nifti_fname, hardi_bval_fname, hardi_bvec_fname = get_fnames('stanford_hardi')\n", " csf, gm, wm = read_stanford_pve_maps()\n", " wm_data = wm.get_fdata()\n", "\n", " img = get_img(hardi_nifti_fname)\n", " voxel_order = \"\".join(aff2axcodes(img.affine))\n", " gtab = get_gtab(hardi_bval_fname, hardi_bvec_fname)\n", " data = img.get_fdata()\n", " roi_data = (wm_data > 0.5)\n", " mask = roi_data\n", "else:\n", " img = get_img(args.nifti_file)\n", " voxel_order = \"\".join(aff2axcodes(img.affine))\n", " gtab = get_gtab(args.bvals, args.bvecs)\n", " roi = get_img(args.roi_nifti)\n", " mask = get_img(args.mask_nifti)\n", " data = img.get_fdata()\n", " roi_data = roi.get_fdata()\n", " mask = mask.get_fdata()\n", "```\n", "\n", "This section loads the diffusion data and extracts relevant attributes such as voxel ordering and data arrays. We provide a Stanford HARDI dataset for example purposes if no diffusion data is given." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Tensor Model Fitting\n", "\n", "```python\n", "tenmodel = dti.TensorModel(gtab, fit_method='WLS')\n", "print('Fitting Tensor')\n", "tenfit = tenmodel.fit(data, mask)\n", "print('Computing anisotropy measures (FA,MD,RGB)')\n", "FA = tenfit.fa\n", "FA[np.isnan(FA)] = 0\n", "\n", "tissue_classifier = ThresholdStoppingCriterion(FA, args.fa_threshold)\n", "metric_map = np.asarray(FA, 'float64')\n", "\n", "seeds = np.asarray(utils.random_seeds_from_mask(\n", " roi_data, seeds_count=args.nseeds,\n", " seed_count_per_voxel=False,\n", " affine=np.eye(4)))\n", "```\n", "\n", "We fit a tensor model to the diffusion data using the Weighted Least Squares (WLS) method. Then, Fractional Anisotropy (FA) is computed and NaN values are replaced with zeros. This will be used as a proxy for a white matter mask, as a seed/stop mask for the tractography. Of course, you can pass your own numpy array in for the `metric_map` parameter and provide your own seeds if you have a better white matter mask and/or seeding strategy." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Model Selection\n", "\n", "```python\n", "# Setup model\n", "sphere = default_sphere\n", "if args.model == \"opdt\":\n", " model_type = cuslines.ModelType.OPDT\n", " print(\"Running OPDT model...\")\n", " model = OpdtModel(gtab, sh_order=args.sh_order, smooth=args.sm_lambda, min_signal=args.min_signal)\n", " fit_matrix = model._fit_matrix\n", " delta_b, delta_q = fit_matrix\n", "elif args.model == \"csa\":\n", " model_type = cuslines.ModelType.CSA\n", " print(\"Running CSA model...\")\n", " model = CsaOdfModel(gtab, sh_order=args.sh_order, smooth=args.sm_lambda, min_signal=args.min_signal)\n", " fit_matrix = model._fit_matrix\n", " # Unlike OPDT, CSA has a single matrix used for fit_matrix. Populating delta_b and delta_q with necessary values for\n", " # now.\n", " delta_b = fit_matrix\n", " delta_q = fit_matrix\n", "else:\n", " print(\"Running CSD model...\")\n", " unique_bvals = unique_bvals_magnitude(gtab.bvals)\n", " if len(unique_bvals[unique_bvals > 0]) > 1:\n", " low_shell_idx = gtab.bvals <= unique_bvals[unique_bvals > 0][0]\n", " response_gtab = gradient_table( # reinit as single shell for this CSD\n", " gtab.bvals[low_shell_idx],\n", " gtab.bvecs[low_shell_idx])\n", " data = data[..., low_shell_idx]\n", " else:\n", " response_gtab = gtab\n", " response, _ = auto_response_ssst(\n", " response_gtab,\n", " data,\n", " roi_radii=10,\n", " fa_thr=0.7)\n", " model = ConstrainedSphericalDeconvModel(gtab, response, sh_order=args.sh_order)\n", " delta_b = model._X\n", " delta_q = model.B_reg\n", "\n", "if args.dg != \"boot\":\n", " if args.dg == \"prob\":\n", " model_type = cuslines.ModelType.PROB\n", " dg = ProbabilisticDirectionGetter\n", " else:\n", " model_type = cuslines.ModelType.PTT\n", " dg = PTTDirectionGetter\n", " fit = model.fit(data, mask=(metric_map >= args.fa_threshold))\n", " data = fit.odf(sphere).clip(min=0)\n", "else:\n", " dg = BootDirectionGetter\n", "\n", "global_chunk_size = args.chunk_size\n", "\n", "# Setup direction getter args\n", "if args.device == \"cpu\":\n", " if args.dg != \"boot\":\n", " dg = dg.from_pmf(data, max_angle=args.max_angle, sphere=sphere, relative_peak_threshold=args.relative_peak_threshold, min_separation_angle=args.min_separation_angle)\n", " else:\n", " dg = BootDirectionGetter.from_data(data, model, max_angle=args.max_angle, sphere=sphere, sh_order=args.sh_order, relative_peak_threshold=args.relative_peak_threshold, min_separation_angle=args.min_separation_angle)\n", "else:\n", " # Setup direction getter args\n", " b0s_mask = gtab.b0s_mask\n", " dwi_mask = ~b0s_mask\n", "\n", " # setup sampling matrix\n", " theta = sphere.theta\n", " phi = sphere.phi\n", " sampling_matrix, _, _ = shm.real_sym_sh_basis(args.sh_order, theta, phi)\n", "\n", " ## from BootPmfGen __init__\n", " # setup H and R matrices\n", " x, y, z = model.gtab.gradients[dwi_mask].T\n", " r, theta, phi = shm.cart2sphere(x, y, z)\n", " B, _, _ = shm.real_sym_sh_basis(args.sh_order, theta, phi)\n", " H = shm.hat(B)\n", " R = shm.lcr_matrix(H)\n", "\n", " # create floating point copy of data\n", " dataf = np.asarray(data, dtype=np.float64)\n", "```\n", "\n", "In this long section we setup the model we are going to use for tractography. We have bootstrapping and probabilistic direction getting implemented as of time of writing, and are working on parallel transport tractography (ptt) direction getting. When using probabilistic or ptt, the data you pass into GPUStreamlines will be the fit model. When using bootstrapped direciton getting, you pass in the original DWI data as well as some model parameters. Note that the delta_b/delta_q parameters are unused in the probabilistic and ptt case. If you fit your own fODFs, they can be passed in to GPUStreamlines as long as they are in the DIPY format." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Initializing GPUStreamlines\n", "\n", "```python\n", "gpu_tracker = cuslines.GPUTracker(\n", " model_type,\n", " args.max_angle * np.pi/180,\n", " args.min_signal,\n", " args.fa_threshold,\n", " args.step_size,\n", " args.relative_peak_threshold,\n", " args.min_separation_angle * np.pi/180,\n", " dataf.astype(np.float64), H.astype(np.float64), R.astype(np.float64), delta_b.astype(np.float64), delta_q.astype(np.float64),\n", " b0s_mask.astype(np.int32), metric_map.astype(np.float64), sampling_matrix.astype(np.float64),\n", " sphere.vertices.astype(np.float64), sphere.edges.astype(np.int32),\n", " ngpus=args.ngpus, rng_seed=0)\n", "```\n", "\n", "Finally, we pass all of these parameters we have set up into the cuslines.GPUTracker constructor. Note that typing must be correct, with each numpy array using either float64s or int32s as appropriate. Also, data must be in contiguous C memory order. So, data loaded directly from a Nibabel file (which stores data in fortran's memory order) must be ordered using numpy's `ascontiguousarray`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Running GPUStreamlines and Saving Results\n", "\n", "```python\n", "print('streamline gen')\n", "nchunks = (seed_mask.shape[0] + global_chunk_size - 1) // global_chunk_size\n", "\n", "t1 = time.time()\n", "streamline_time = 0\n", "io_time = 0\n", "\n", "if args.output_prefix and write_method == \"trx\":\n", " # Will resize by a factor of 2 if these are exceeded\n", " sl_len_guess = 100\n", " sl_per_seed_guess = 3\n", " n_sls_guess = sl_per_seed_guess*len(seed_mask)\n", "\n", " # trx files use memory mapping\n", " trx_file = TrxFile(\n", " reference=hardi_nifti_fname,\n", " nb_streamlines=n_sls_guess,\n", " nb_vertices=n_sls_guess*sl_len_guess)\n", " offsets_idx = 0\n", " sls_data_idx = 0\n", "\n", "for idx in range(int(nchunks)):\n", " # Main streamline computation\n", " ts = time.time()\n", " if args.device == \"cpu\":\n", " streamline_generator = LocalTracking(dg, tissue_classifier, seed_mask[idx*global_chunk_size:(idx+1)*global_chunk_size], affine=np.eye(4), step_size=args.step_size)\n", " streamlines = [s for s in streamline_generator]\n", " else:\n", " streamlines = gpu_tracker.generate_streamlines(seed_mask[idx*global_chunk_size:(idx+1)*global_chunk_size])\n", " te = time.time()\n", " streamline_time += (te-ts)\n", " print(\"Generated {} streamlines from {} seeds, time: {} s\".format(\n", " len(streamlines),\n", " seed_mask[idx*global_chunk_size:(idx+1)*global_chunk_size].shape[0],\n", " te-ts))\n", "\n", " # Save tracklines file\n", " if args.output_prefix:\n", " ts = time.time()\n", " if write_method == \"standard\":\n", " fname = \"{}.{}_{}.trk\".format(args.output_prefix, idx+1, nchunks)\n", " sft = StatefulTractogram(streamlines, args.nifti_file, Space.VOX)\n", " save_tractogram(sft, fname)\n", " te = time.time()\n", " print(\"Saved streamlines to {}, time {} s\".format(fname, te-ts))\n", " elif write_method == \"trx\":\n", " tractogram = nib.streamlines.Tractogram(streamlines, affine_to_rasmm=img.affine)\n", " tractogram.to_world()\n", " sls = tractogram.streamlines\n", "\n", " new_offsets_idx = offsets_idx + len(sls._offsets)\n", " new_sls_data_idx = sls_data_idx + len(sls._data)\n", "\n", " if new_offsets_idx > trx_file.header[\"NB_STREAMLINES\"]\\\n", " or new_sls_data_idx > trx_file.header[\"NB_VERTICES\"]:\n", " print(\"TRX resizing...\")\n", " trx_file.resize(nb_streamlines=new_offsets_idx*2, nb_vertices=new_sls_data_idx*2)\n", "\n", " # TRX uses memmaps here\n", " trx_file.streamlines._data[sls_data_idx:new_sls_data_idx] = sls._data\n", " trx_file.streamlines._offsets[offsets_idx:new_offsets_idx] = offsets_idx + sls._offsets\n", " trx_file.streamlines._lengths[offsets_idx:new_offsets_idx] = sls._lengths\n", "\n", " offsets_idx = new_offsets_idx\n", " sls_data_idx = new_sls_data_idx\n", "\n", " te = time.time()\n", " print(\"Streamlines to TRX format, time {} s\".format(te-ts))\n", " else:\n", " fname = \"{}.{}_{}\".format(args.output_prefix, idx+1, nchunks)\n", " gpu_tracker.dump_streamlines(fname, voxel_order, wm.shape, wm.header.get_zooms(), img.affine)\n", " te = time.time()\n", " print(\"Saved streamlines to {}, time {} s\".format(fname, te-ts))\n", "\n", " io_time += (te-ts)\n", "\n", "if args.output_prefix and write_method == \"trx\":\n", " ts = time.time()\n", " fname = \"{}.trx\".format(args.output_prefix)\n", " trx_file.resize()\n", " zip_from_folder(\n", " trx_file._uncompressed_folder_handle.name,\n", " fname,\n", " zipfile.ZIP_STORED)\n", " trx_file.close()\n", " te = time.time()\n", " print(\"Saved streamlines to {}, time {} s\".format(fname, te-ts))\n", " io_time += (te-ts)\n", "\n", "t2 = time.time()\n", "\n", "print(\"Completed processing {} seeds.\".format(seed_mask.shape[0]))\n", "print(\"Initialization time: {} sec\".format(t1-t0))\n", "print(\"Streamline generation total time: {} sec\".format(t2-t1))\n", "print(\"\\tStreamline processing: {} sec\".format(streamline_time))\n", "if args.output_prefix:\n", " print(\"\\tFile writing: {} sec\".format(io_time))\n", "```\n", "\n", "Finally, we batch the seeds that we want to track from into the GPUStreamlines class we constructed. The batch size is set manually, and can be set as large as possible without exceeding the GPU's memory. The default batch size should work well in most cases. GPUStreamlines returns the streamlines for each batch, and you can work with them as you please. In this example, we save them as either TRK or TRX files. In the TRK case, we save multiple TRK files that we can merge later. In the TRX case, we save them to one large TRX file that we save out in the end." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusion\n", "\n", "Using this script as reference, GPUstreamlines should slot nicely into any diffusion MRI processing pipeline, and dramatically speedup tractography generation. We allow flexibility in choosing different models and saving methods for tractography results. \n", "\n", ":::{tip}\n", "Feel free to post issues and ask questions at the `https://github.com/dipy/GPUStreamlines` repository.\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 }