/matten

Materials Tensor prediction using equivariant graph neural networks

Primary LanguagePython

eigenn

Equivariant Invariant Graph Enabled Neural Network

Install

This below installing guide should get you started on Mac, without using GPUs.

  • PyTorch

    conda install pytorch -c pytorch

  • PyTorch Geometrics. PyG now supports installing via conda, much easier than before. And the installing guide on e3nn for PyG is not recommended.

    conda install pyg -c pyg -c conda-forge

  • Lightning

    conda install pytorch-lightning==1.5.2 torchmetrics==0.6.0 lightning-bolts -c conda-forge

  • e3nn

    pip install e3nn

  • This repo

    git clone https://github.com/mjwen/eigenn.git
cd eigenn
pip install -e .

  • Weights & Biases

    We use wandb for experiments tracking and management; it is free. Get an account at their website and then the below commands are all you need.

    pip install wandb
wandb login

Examples

Train on atomic property (e.g. NMR tensor)

Run this script train_atomic.py

python train_atomic.py --config <config file>

If config is not provided, the default minimal_atomic.yaml is used. The model used in train_atomic.py is built at atomic_tensor_model.py

Train on structure property

Run this scirpt train.py

python train.py --config <config file>

If config is not provided, the default minimal.yaml is used. The model used in train.py is built at nequip_energy_model.py

To get help

python train[_atomic].py --help

Under the hood, we use Lightning CLI to build the interface, more usage info at here.