BrainGB is a unified, modular, scalable, and reproducible framework established for brain network analysis with GNNs. It is designed to enable fair evaluation with accessible datasets, standard settings, and baselines to foster a collaborative environment within computational neuroscience and other related communities. This library is built upon PyTorch and PyTorch Geometric.

The BrainGB paper is published in IEEE Transactions on Medical Imaging (TMI, IF: 11.037) - Special Issue on Geometric Deep Learning in Medical Imaging. To foster research, we also provide an out-of-box package that can be installed directly using pip, with detailed tutorials available on our hosted BrainGB website.

Acknowledgement

The development of this toolbox is partially supported by NIH (R01MH105561, R01MH118771, R01AG071243, R01MH125928, U01AG068057), NSF (IIS 2045848, IIS 1837956) and ONR (N00014-18-1-2009).

Library Highlights

Our BrainGB implements four main modules of GNN models for brain network analysis:

Node feature construction: studies practical and effective methods to initialize either positional or structural node features for each brain region.
Message passing mechanisms: update the node representation of each brain region iteratively by aggregating neighbor features through local connections.
Attention-enhanced message passing: incorporates attention mechanism to enhance the message passing scheme of GNNs.
Pooling strategies: operate on the set of node vectors to get a graph-level representation.

BrainGB also implements utility functions for model training, performance evaluation, and experiment management.

Using BrainGB

There are two ways to use BrainGB: the first is running direct experiments with BrainGB, and the second is integrating BrainGB into your existing research projects. Follow the sections below to learn more.

Direct Experiments with BrainGB

1. Obtaining Datasets

ABIDE Dataset

We understand the challenges faced by researchers in accessing certain datasets due to restrictions. To facilitate your experimentation with BrainGB, we provide the Autism Brain Imaging Data Exchange (ABIDE) dataset, which is publicly accessible and does not require special access permissions.

Download and Preprocess ABIDE dataset here

Datasets Requiring Access

For a detailed exploration of other datasets like PNC, PPMI, and ABCD utilized in our BrainGB studies, which are not publicly accessible and require specific access permissions, please refer to the following:

You can also construct your own datasets by following the instructions on neuroimaging preprocessing and functional or structural brain network construction on our website.

2. Quick Setup

Clone the repository and Install required dependencies:

git clone https://github.com/HennyJie/BrainGB.git

It is recommended to use commands such as 'virtualenv' to create separate Python environments, in order to prevent conflicts in environment configurations:

virtualenv -p python3 venv
source venv/bin/activate

Install Dependencies: BrainGB depends on the following packages. The packages listed below are dependencies for systems with CUDA version 10.1. If you are using a different version of CUDA, please ensure to install the respective versions of these dependencies that are compatible with your CUDA version. See Pytorch version with different CUDA versions:

torch~=1.8.1+cu101
torch-cluster~=1.5.9
torch-scatter~=2.0.8       
torch-sparse~=0.6.12
torch-spline-conv~=1.2.1
torch-geometric~=2.0.4
numpy~=1.22.2
nni~=2.4
scikit-learn~=1.0.2
networkx~=2.6.2
scipy~=1.7.3
pandas~=1.2.3

First, install some of the dependencies with:

pip install -r requirements.txt

Next, install Pytorch:

pip install torch==1.8.1+cu101 torchvision==0.9.1+cu101 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

Finally, install torch-cluster, torch-scatter, torch-sparse, torch-spline-conv and torch-geometric

pip install https://data.pyg.org/whl/torch-1.8.0%2Bcu101/torch_cluster-1.5.9-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.8.0%2Bcu101/torch_scatter-2.0.8-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.8.0%2Bcu101/torch_sparse-0.6.12-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.8.0%2Bcu101/torch_spline_conv-1.2.1-cp38-cp38-linux_x86_64.whl
pip install torch-geometric~=2.0.4

Note: If you face problems when installing dependencies like torch-cluster, torch-scatter, torch-sparse, torch-spline-conv, and torch-geometric, it is recommended to manually install the respective version of these packages using the '.whl' files available on their official website.

3. Running Example

Use the ABIDE dataset as an example, you should first place the dataset file "abide.npy" (genereated from step 1) in the datasets folder under the examples folder (Create the folder if it does not exist). The abide.npy file contains the following contents:

timeseries: Represents the BOLD time series data for each subject. It's a numpy array with the shape (#sub, #ROI, #timesteps).
Label: Provides the diagnosis label for Autism spectrum disorder for each subject. '0' denotes negative, and '1' indicates positive. It's a numpy array of shape (#sub).
corr: The correlation matrix calculated from BOLD time series data. It's a numpy array with the shape (#sub, #ROIs, #ROIs).
pcorr: Represents the partial correlation matrix derived from the BOLD time series data. It's a numpy array with dimensions (#sub, #ROIs, #ROIs).
site: Specifies where the data was collected for each subject. It's a numpy array with shape (#sub).

Important Note: "Label" and "corr matrix" are the actual inputs for BrainGB. Label represents the target outcome we are interested in predicting, often indicating the diagnosis or condition of a subject in a brain study. corr matrix describes the associated Brain Network. If you are considering running BrainGB using your own dataset, it's important to format your Label and corr matrix similarly to ensure compatibility and accurate results. Ensure that Label is in a numpy array of shape (#sub) and corr matrix is structured as a numpy array with the shape (#sub, #ROIs, #ROIs).<br><br>

Run the BrainGB code, execute the following command:

python -m examples.example_main --dataset_name ABIDE --pooling concat --gcn_mp_type edge_node_concate --hidden_dim 256

The parameter pooling specifies the pooling strategy to get a graph-level representation for each subject and gcn_mp_type sets a message vector design for the gcn model. If you choose gat as the backbone model, you can use gat_mp_type to set an attention-enhancing mechanism.

For other hyper-parameters like --n_GNN_layer, --n_MLP_layers, --hidden_dim, --epochs, etc., you can modify them to adjust the detailed model design or control the training process. If you'd like to automatically search and optimize these hyper-parameters, use the AutoML tool NNI with the --enable_nni command.

Upon successful execution, you should observe an output similar to this:

Processing...
Done!
2023-09-10 15:54:28,486 - Loaded dataset: ABIDE
...
2023-09-10 15:56:29,493 - (Train Epoch 9), test_micro=66.34, test_macro=65.10, test_auc=72.91
...
2023-09-10 17:37:46,561 - (Train Epoch 99), test_micro=64.68, test_macro=64.59, test_auc=70.03
2023-09-10 17:37:47,489 - (Initial Performance Last Epoch) | test_micro=64.68, test_macro=64.59, test_auc=70.03
2023-09-10 17:37:47,489 - (K Fold Final Result)| avg_acc=65.31 +-  1.58, avg_auc=71.29 +- 2.89, avg_macro=64.43 +- 1.87

4. Customizing Your Own GNN Models

Node Feature Construction

In src.dataset.tranforms, BrainGB provides the BaseTransform base class, which offers a universal interface for node feature initialization for each brain region. S

BrainGB

Install / Use

README