Running pyAFQ using the GPU for tractography

Running pyAFQ using the GPU for tractography is as simple as (1) Installing GPUStreamlines using pip install and (2) passing in the jit_backend parameter when you create your GroupAFQ object. To install GPUStreamlines, do: pip install git+https://github.com/dipy/GPUStreamlines.git That’s step 1 complete! The rest of this example is the same as the GroupAFQ example except with the jit_backend parameter set.

from AFQ.api.group import GroupAFQ
import AFQ.data.fetch as afd
import os.path as op
import plotly

2026-05-19 01:00:27,790	INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

We start with some example data. The data we will use here is generated from the Stanford HARDI dataset. We then setup our myafq object which we will use to demonstrate the clobber method.

afd.organize_stanford_data()

Set tractography parameters

We make create a tracking_params variable to define the parameters for tractography. The only parameter we need to set to use the GPU is jit_backend, which we set to “cuda”. Other backends include: “metal”, “webgpu”, or “numba”. Numba is the default. Note that the GPU backend will only run for probabilistic tracking, which is the default.

tracking_params = dict(n_seeds=1e7,
                       random_seeds=True,
                       rng_seed=2025,
                       jit_backend="cuda",
                       trx=True)

Running with the GPU

Then, run pyAFQ normally. That’s it!

myafq = GroupAFQ(
    bids_path=op.join(afd.afq_home, 'stanford_hardi'),
    dwi_preproc_pipeline='vistasoft',
    t1_preproc_pipeline='freesurfer',
    tracking_params=tracking_params)

bundle_html = myafq.export("all_bundles_figure")
plotly.io.show(bundle_html["01"][0])

INFO:bidsschematools:No schema path specified, defaulting to the bundled schema, `/opt/hostedtoolcache/Python/3.13.13/x64/lib/python3.13/site-packages/bidsschematools/data/schema.json`.

INFO:AFQ:Using the following files for subject 01 and session 01:

INFO:AFQ:  DWI: /home/runner/AFQ_data/stanford_hardi/derivatives/vistasoft/sub-01/ses-01/dwi/sub-01_ses-01_dwi.nii.gz

INFO:AFQ:  BVAL: /home/runner/AFQ_data/stanford_hardi/derivatives/vistasoft/sub-01/ses-01/dwi/sub-01_ses-01_dwi.bval

INFO:AFQ:  BVEC: /home/runner/AFQ_data/stanford_hardi/derivatives/vistasoft/sub-01/ses-01/dwi/sub-01_ses-01_dwi.bvec

INFO:AFQ:  T1: /home/runner/AFQ_data/stanford_hardi/derivatives/freesurfer/sub-01/ses-01/anat/sub-01_ses-01_T1w.nii.gz

WARNING:AFQ:It is recommended to provide CSF/GM/WM segmentations using PVEImage or PVEImages in AFQ.definitions.image. Otherwise, SynthSeg2 will be used

INFO:AFQ:Calculating _desc-brain_mask.nii.gz...

2026-05-19 01:00:32.147527936 [W:onnxruntime:Default, device_discovery.cc:133 GetPciBusId] Skipping pci_bus_id for PCI path at "/sys/devices/LNXSYSTM:00/LNXSYBUS:00/ACPI0004:00/MSFT1000:00/5620e0c7-8062-4dce-aeb7-520c7ef76171" because filename "5620e0c7-8062-4dce-aeb7-520c7ef76171" did not match expected pattern of [0-9a-f]+:[0-9a-f]+:[0-9a-f]+[.][0-9a-f]+
INFO:AFQ:Calculating _desc-T1w_mask.nii.gz...

INFO:AFQ:Running mindgrab...