Reference implementation of PESNet from

Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions
by Nicholas Gao, Stephan Günnemann
published as Spotlight at ICLR 2022.

and Planet and PESNet++ from

Sampling-free Inference for Ab-Initio Potential Energy Surface Networks
by Nicholas Gao, Stephan Günnemann
published at ICLR 2023

If you're looking for the code of our ICML paper, please check out our globe repository.

First install JAX and the correct CUDA Toolkit and CUDNN, then this package via

pip install -e .

You can now train a model, e.g., H2, via a config file

python train.py with configs/systems/h2.yaml print_progress=True

You can overwrite parameters either via CLI or via the config file. All progress is tracked on tensorboard.

We encourage the use of seml to manage all experiments but we also supply commands to run the experiments directly.

With seml:

seml n2_ablation add train_n2_ablation.yaml start

Without seml:

# PESNet
python train.py with configs/systems/n2.yaml \\
    init_method=pesnet \\
    pesnet.ferminet_params.activation=tanh \\
    pesnet.ferminet_params.input_config.mlp_activation=tanh \\
    pesnet.ferminet_params.jastrow_config=None \\
    pesnet.ferminet_params.determinants=32
# PESNet++ (default config)
python train.py with configs/systems/n2.yaml \\
    init_method=pesnet \\
    pesnet.ferminet_params.activation=silu \\
    pesnet.ferminet_params.input_config.mlp_activation=silu \\
    pesnet.ferminet_params.jastrow_config.n_layers=3 \\
    pesnet.ferminet_params.jastrow_config.activation=silu \\
    pesnet.ferminet_params.determinants=32

To run all experiments from the PlaNet paper with seml simply run:

seml pes add train_pes.yaml start

Please contact [email protected] if you have any questions.

Please cite our paper if you use our method or code in your own works:

@inproceedings{gao_pesnet_2022,
    title = {Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions},
    author = {Gao, Nicholas and G{\"u}nnemann, Stephan}
    booktitle = {International Conference on Learning Representations (ICLR)},
    year = {2022}
}

@inproceedings{gao_planet_2023,
    title = {Sampling-free Inference of Ab-initio Potential Energy Surface Networks},
    author = {Gao, Nicholas and G{\"u}nnemann, Stephan}
    booktitle = {International Conference on Learning Representations (ICLR)},
    year = {2023}
}

Hippocratic License v2.1

The logo is generated by Bing Image Creator.