Leveraging Taichi Lang to Customize Brain Dynamics Operators

<p align="center"> <img alt="Header image of braintaichi." src="https://github.com/chaoming0625/braintaichi/blob/main/docs/_static/braintaichi.png" width=50%> </p> <p align="center"> <a href="https://pypi.org/project/braintaichi/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/braintaichi"></a> <a href="https://github.com/chaoming0625/braintaichi/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a> <a href='https://braintaichi.readthedocs.io/en/latest/?badge=latest'> <img src='https://readthedocs.org/projects/braintaichi/badge/?version=latest' alt='Documentation Status' /> </a> <a href="https://badge.fury.io/py/braintaichi"><img alt="PyPI version" src="https://badge.fury.io/py/braintaichi.svg"></a> <a href="https://github.com/chaoming0625/braintaichi/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaoming0625/braintaichi/actions/workflows/CI.yml/badge.svg"></a> </p>

braintaichi leverages Taichi Lang to customize brain dynamics operators.

Quick Start

import taichi as ti

import braintaichi as bti


# define the custom kernel

@ti.kernel
def transpose_bool_homo_kernel(
    values: ti.types.ndarray(ndim=1),
    indices: ti.types.ndarray(ndim=1),
    indptr: ti.types.ndarray(ndim=1),
    events: ti.types.ndarray(ndim=1),
    out: ti.types.ndarray(ndim=1)
):
    value = values[0]
    ti.loop_config(serialize=True)
    for row_i in range(indptr.shape[0] - 1):
        if events[row_i]:
            for j in range(indptr[row_i], indptr[row_i + 1]):
                out[indices[j]] += value


kernel = bti.XLACustomOp(
    cpu_kernel=transpose_bool_homo_kernel,
    gpu_kernel=transpose_bool_homo_kernel,
)

# run with the sample data

import numpy as np
import jax
import jax.numpy as jnp
from scipy.sparse import csr_matrix

csr = csr_matrix((np.random.rand(10, 10) < 0.5).astype(float))
events = np.random.rand(10) < 0.5

out = kernel(
    jnp.array(csr.data),
    jnp.array(csr.indices),
    jnp.array(csr.indptr),
    events,
    outs=[jax.ShapeDtypeStruct([10], dtype=jnp.float32)]
)
print(out)

Installation

From PyPI

You can install braintaichi via pip:

pip install braintaichi --upgrade

From Source

To compile and install from source code:

Prerequisites

1. Install dependencies:

   pip install taichi==1.7.3 pybind11
   

Platform-specific Requirements

Linux with GPU support:

• NVIDIA GPU drivers
• CUDA Toolkit (version 12.8 recommended, see ci/linux/gpu/prepare_build_wheel_on_cuda.sh)
• Development packages:

  # Example for RHEL/CentOS
  yum install -y python3-devel cuda-cudart-devel
  

Windows:

• Visual Studio with C++ development tools (Installation Guide)

Build and Install

git clone https://github.com/chaoming0625/braintaichi.git
cd braintaichi
pip install -e .

Or alternatively:

python setup.py install

Documentation

The official documentation is hosted on Read the Docs: https://braintaichi.readthedocs.io

See also the ecosystem

We are building the brain modeling ecosystem: https://brainmodeling.readthedocs.io/

Citation

If you think braintaichi is significant in your work, please consider to cite the following pubilication:

@inproceedings{wang2024brainpy,
    title={A differentiable brain simulator bridging brain simulation and brain-inspired computing},
    author={Wang, Chaoming and Zhang, Tianqiu and He, Sichao and Gu, Hongyaoxing and Li, Shangyang and Wu, Si},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024}
}