https://github.com/chaobrain/braintaichi
Leveraging Taichi Lang to customize brain dynamics operators.
https://github.com/chaobrain/braintaichi
brain-dynamics event-driven-computation sparse-computation
Last synced: 2 months ago
JSON representation
Leveraging Taichi Lang to customize brain dynamics operators.
- Host: GitHub
- URL: https://github.com/chaobrain/braintaichi
- Owner: chaobrain
- License: apache-2.0
- Created: 2024-08-24T19:58:48.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-02-06T08:49:08.000Z (about 1 year ago)
- Last Synced: 2025-02-06T09:36:15.549Z (about 1 year ago)
- Topics: brain-dynamics, event-driven-computation, sparse-computation
- Language: Python
- Homepage: https://braintaichi.readthedocs.io
- Size: 1.75 MB
- Stars: 0
- Watchers: 2
- Forks: 0
- Open Issues: 8
-
Metadata Files:
- Readme: README.md
- Changelog: changelog.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Security: SECURITY.md
Awesome Lists containing this project
- awesome-jax - braintaichi - Leveraging Taichi Lang to customize brain dynamics operators. <img src="https://img.shields.io/github/stars/chaobrain/braintaichi?style=social" align="center"> (Libraries / New Libraries)
README
# Leveraging Taichi Lang to Customize Brain Dynamics Operators
[``braintaichi``](https://github.com/chaoming0625/braintaichi) leverages Taichi Lang to customize brain dynamics operators.
## Quick Start
```python
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:
```bash
pip install braintaichi --upgrade
```
### From Source
To compile and install from source code:
#### Prerequisites
1. **Install dependencies**:
```bash
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](ci/linux/gpu/prepare_build_wheel_on_cuda.sh))
- Development packages:
```bash
# Example for RHEL/CentOS
yum install -y python3-devel cuda-cudart-devel
```
**Windows**:
- Visual Studio with C++ development tools ([Installation Guide](https://docs.microsoft.com/en-us/cpp/build/vscpp-step-0-installation))
#### Build and Install
```bash
git clone https://github.com/chaoming0625/braintaichi.git
cd braintaichi
pip install -e .
```
Or alternatively:
```bash
python setup.py install
```
## Documentation
The official documentation is hosted on Read the Docs: [https://braintaichi.readthedocs.io](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:
```bibtex
@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}
}
```