GMN

KO | EN

gitlite — search

#typescript #golang #android #rust #deep-learning #linux #awesome #cli #kubernetes #vue #awesome-list #mcp #openai #open-source

(★ 66)

[ICLR 2022] The implementation for the paper "Equivariant Graph Mechanics Networks with Constraints".

File Explorer

Download Latest Version (.zip)

- - config_gmnoc.yml
  - gmn_temp.py
  - gmnoc.py
  - layer.py
  - README.md
- demo.gif
- mocap.png
- overview.png
- poster.pdf
- simulation.png
- simple_config.json
- simple_config_md17.json
- simple_config_motion.json
- - aspirin_split.pkl
  - benzene_old_split.pkl
  - dataset.py
  - ethanol_split.pkl
  - malonaldehyde_split.pkl
  - naphthalene_split.pkl
  - salicylic_split.pkl
  - toluene_split.pkl
  - uracil_split.pkl
- - __init__.py
  - gcl.py
  - layer.py
- - dataset.py
  - motion.pkl
  - split.pkl
- - __init__.py
  - dataset_nbody.py
  - model.py
- main.py
- main_md17.py
- main_motion.py
- vis.ipynb
- vis_cont.ipynb
.gitignore
LICENSE
README.md
requirements.yml

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Equivariant Graph Mechanics Networks with Constraints (ICLR 2022)

Wenbing Huang\*, Jiaqi Han\*, Yu Rong, Tingyang Xu, Fuchun Sun, Junzhou Huang

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/hanjq17/GMN/blob/main/LICENSE)

[**[OpenReview]**](https://openreview.net/forum?id=SHbhHHfePhP) [**[Paper]**](https://openreview.net/pdf?id=SHbhHHfePhP) [**[Poster]**](assets/poster.pdf)

Graph Mechanics Networks (GMNs) are novel graph neural networks particularly powerful for modeling the dynamics of constrained systems. GMNs are equivariant to translations, rotations, and reflections. The flowchart of our model is provided in the figure below. Please refer to our paper for more details.

![Overview](https://raw.githubusercontent.com/hanjq17/GMN/main/assets/overview.png "Overview")

## Dependencies

```
python==3.8
torch==1.8.0
scikit-learn==0.24.2
networkx==2.5.1
```

A more detailed Python environment is depicted in `requirements.yml`.

## Data Preparation

**1. Simulation Dataset**

Our simulation scripts are placed under `spatial_graph/n_body_system/dataset`.

To generate datasets containing multiple isolated particles, sticks, and hinges, use the following command under the path `spatial_graph/n_body_system/dataset`:

```bash
python -u generate_dataset.py --num-train 5000 --seed 43 --n_isolated 3 --n_stick 2 --n_hinge 1 --n_workers 50
```

where the arguments `n_isolated`, `n_stick`, and `n_hinge` indicate the number of isolated particles, sticks, and hinges, respectively. The argument `n_workers` refers to the number of parallel threads for parallel data generation. For other potential arguments, please refer to `generate_dataset.py`.

By default, the generated data will be placed in a new folder named `data` under the current path.

*Note:* On our CPU machine with 50 parallel workers, the entire data generation process takes from 10 minutes to 1~2 hours, depending on the complexity of the particle system (number of particles, sticks, and hinges).

**2. MD17**

The MD17 dataset can be downloaded from [MD17](http://quantum-machine.org/gdml/#datasets). Download the dataset and place the files under `spatial_graph/MD17`. The splits are also provided in the `MD17` folder.

**3. Motion Capture**

The raw data were obtained via [CMU Motion Capture Database](http://mocap.cs.cmu.edu/search.php?subjectnumber=35). The preprocessed dataset as well as the splits are provided in  `spatial_graph/motion` folder.

## Training and Evaluation

**1. Simulation Dataset**

Under the root path, simply use

```bash
python -u spatial_graph/main.py --config_by_file
```

where the `--config_by_file` option enables loading the hyper-parameters from the config files in `configs` folder.

To run experiments in different scenarios, simply change the hyper-parameters in the config files. For instance, one may change the `n_isolated`, `n_stick`, and `n_hinge` options to evaluate the model under various combinations of isolated particles, sticks, and hinges.

**2. MD17**

```bash
python -u spatial_graph/main_md17.py --config_by_file
```

**3. Motion Capture**

```bash
python -u spatial_graph/main_motion.py --config_by_file
```

## Visualizations

We provide a toy script for visualization in `visualization/vis.ipynb`.

We exhibit some visualizations of the predictions as below (*Left*: initial position, *Middle*: GMN, *Right*: EGNN):

**Simulation dataset**

![Simulation](https://raw.githubusercontent.com/hanjq17/GMN/main/assets/simulation.png "Simulation")

**Motion Capture**

![Mocap](https://raw.githubusercontent.com/hanjq17/GMN/main/assets/mocap.png "Mocap")

Here we also provide some dynamic cases tested with a batch of **continuous** input frames (*Blue*: prediction, *Red*: ground truth). The toy script is deferred to `visualization/vis_cont.ipynb`.

![Demo](https://raw.githubusercontent.com/hanjq17/GMN/main/assets/demo.gif "Demo")

## Citation

If you find our work helpful, please cite as:

```
@inproceedings{
huang2022equivariant,
title={Equivariant Graph Mechanics Networks with Constraints},
author={Wenbing Huang and Jiaqi Han and Yu Rong and Tingyang Xu and Fuchun Sun and Junzhou Huang},
booktitle={International Conference on Learning Representations},
year={2022},
url={https://openreview.net/forum?id=SHbhHHfePhP}
}
```

## New Update

We update the code of our GMN-OC, an enhanced version of GMN for molecular data, to `assets/gmnoc` folder. GMN-OC contributes to our winning the first place of [The Open Catalyst Challenge 2022](https://opencatalystproject.org/challenge.html).

Typo: Delta frame should be 3000 instead of 5000 in the paper (on Page 16). The configuration described in the code (`https://github.com/hanjq17/GMN/blob/main/configs/simple_config_md17.json`) is correct.

## Contact

If you have any questions, feel free to reach us at:

Jiaqi Han: jiaqihan@stanford.edu

GMN

File Explorer

Today's picks

Was this content helpful?

Recent Feedback

GMN

File Explorer

Today's picks

Was this content helpful?

Recent Feedback

Download README