E(n)-equivariant Graph Cellular Automata

This repository is the official implementation of E(n)-equivariant Graph Cellular Automata.

@article{gala2024enequivariant,
    title={E(n)-equivariant Graph Neural Cellular Automata},
    author={Gennaro Gala and Daniele Grattarola and Erik Quaeghebeur},
    journal={Transactions on Machine Learning Research},
    year={2024},
    url={https://openreview.net/forum?id=7PNJzAxkij}
}

N.B. Sometimes GIFs are not properly loaded in README.md. Please refresh the page.

Pattern Formation - Convergence to Geometric Graphs

E(n) Convergence to 3D-Torus	E(n) Regeneration of 3D-Cube

python -m trainers.geometric_graph -ds Grid2d -sdg 0.05 -rdg 1.0 -bsc 0 8 2000 16 4000 32 -pats 800
python -m trainers.geometric_graph -ds Torus  -sdg 0.05 -rdg 1.0 -bsc 0 6 1000 8 2000 16 4000 32 -pats 800
python -m trainers.geometric_graph -ds Cube   -sdg 0.05 -rdg 1.0 -bsc 0 16
python -m trainers.geometric_graph -ds Bunny  -sdg 0.05 -rdg 1.0 -bsc 0 4 1000 8 2000 16 4000 32

For testing, play with notebooks/test_geometric_graph.ipynb.

Graph Autoencoding

First, unzip the datasets in ./data/.

E(n)-GNCA 3D demo:

python -m trainers.gae -ds community -ne 3000 -cd 3 -nd 16 -md 32 -s1 15 -s2 25 -ps 5 -rs 0 5 1000 1000
python -m trainers.gae -ds planar    -ne 3000 -cd 3 -nd 16 -md 32 -s1 15 -s2 25 -ps 5 -rs 0 5 1000 1000 -ng 200 -n1 12 -n2 20

E(n)-GNCA autoencoder:

python -m trainers.gae -ds community -ne 3000 -cd 8  -nd 16 -md 32 -s1 25 -s2 35 -ps 5 -rs 0 5 1000 1000
python -m trainers.gae -ds planar    -ne 3000 -cd 8  -nd 16 -md 32 -s1 25 -s2 35 -ps 5 -rs 0 5 1000 1000 -ng 200 -n1 12 -n2 20
python -m trainers.gae -ds planar    -ne 3000 -cd 8  -nd 16 -md 32 -s1 25 -s2 35 -ps 5 -rs 0 5 1000 1000 -ng 200 -n1 32 -n2 64 -nc 3
python -m trainers.gae -ds proteins  -ne 1000 -cd 16 -nd 16 -md 32 -s1 25 -s2 35 -ps 5 -rs 0 5 100  1000 -pats 20
python -m trainers.gae -ds sbm       -ne 3000 -cd 24 -nd 16 -md 32 -s1 25 -s2 35 -ps 5 -rs 0 5 1000 1000

For testing, play with notebooks/test_gae.ipynb.

Simulation of E(n)-equivariant dynamical systems

Boids

First create the dataset:

python data/boids.py --n_simulations 500 --n_steps 500 --n_boids 100

then run:

python -m trainers.dsystem -ds boids -ne 500 -bs 16 -nd 16 -md 32 -sls 0 20 -lr 1e-3

For testing, play with notebooks/test_dsystems.ipynb.

NBody

First create the dataset:

python data/nbody.py --n_simulations 5000 --n_steps 1000 --n_bodies 5

then run:

python -m trainers.dsystem -ds nbody -ne 500 -nd 16 -md 32 -sls 0 25

For testing, play with notebooks/test_dsystems.ipynb.