This repository contains the training code and model weights presented in a companion paper,polyGNN: Multitask graph neural networks for polymer informatics.
This repository was also used to train a model in a subsequent paper,Gas permeability, diffusivity, and solubility in polymers.Those model weights are provided in this repository, while the training data is in a differentrepository.
This repository is currently set up to run on 1) Mac OSX and 2) Linux/Windows machines with CUDA 10.2. Please raise a GitHub issue if you want to use this repo with a different configuration. Otherwise, follow these steps for installation:
- Installpoetryon your machine.
- If Python3.7 is installed on your machine skip to step 3, if not you will need to install it. There are many ways to do this, one option is detailed below:
- InstallHomebrewon your machine.
- Run
brew install Python @3.7.Take note of the path to the Python executable.
- Clone this repo on your machine.
- Open a terminal at the root directory of this repository.
- Run
poetry env use /path/to/ Python 3.7/executable.If you installed Python3.7 with Homebrew, the path may be something like/usr/local/Cellar/ Python \@3.7/3.7.13_1/bin/ Python 3.7. - Run
poetry install. - If your machine is a Mac, run
poetry run poe torch-osx.If not, runpoetry run poe torch-linux_win-cuda102. - If your machine is a Mac, run
poetry run poe pyg-osx.If not, runpoetry run poe pyg-linux_win-cuda102.
As can be seen inpyproject.toml,polygnndepends on several other packages, includingpolygnn_trainer,
polygnn_kit,andnndebugger.The functional relationships between these libraries are described briefly below and inexample.py.
polygnncontains the polyGNN architecture developed in the companion paper. The architecture relies onpolygnn_kit,which is a library for performing operations on polymer SMILES strings. Meanwhile,polygnn_traineris a library for training neural network architectures, and was used in the companion paper to train the polyGNN architectures. Part of the training process utilizednndebugger,a library for debugging neural networks.
The fileexample.pycontains example code that illustrates how this package was used to the train models in the companion paper. The code uses training data located in the directorysample_datato train an ensemble model (composed of several submodels). The submodels, by default, are saved in a directory namedexample_models.The data insample_datais a small subset of the DFT data used to train the models in the companion paper. A complete set of the DFT data can be found atKhazana.
To train polygnn models run:poetry run Python example.py --polygnn.To train polygnn2 models run:poetry run Python example.py --polygnn2.Running either line on a machine with at least 8GB of free GPU memory should not take longer than 3 minutes. To manually specify the device you want to use for training, set the device flag. For examplepoetry run Python example.py --polygnn --device cpu.Otherwise, the device will automatically be chosen.
Looking atsample_data/sample.csv,you will notice that this dataset contains multiple different properties (e.g., band gap, electron affinity, etc.). Inexample.py,we use this data to train a multitask model, capable of predicting each property. To train your own multitask model, you can replacesample_data/sample.csvwith your own dataset containing multiple properties. Single task models are also supported.
example.pyis an example of how to train a multitask model with only SMILES strings as features.example2.pyis an example of how to train a multitask model containing both SMILES andnon-SMILESfeatures.example.pyandexample2.pyshare the same flags. Read the comments inexample2.pyfor more details.
A directory containing more example files.
more_examples/example_predict_trained_models.pyan example of how to just do prediction using one of the models trained in thecompanion paper.To run the file docd more_examples && poetry run Python example_predict_trained_models.py.The file shows how to get predictions for 36 different properties of polyethylene. Of course, you can change the polymer that you want to get predictions for, but the properties that can be predicted are fixed. If you want to make a prediction for a different property, then you'll need to train your own model (seeexample.pyandexample2.pyfor more details on training your own model).more_examples/example_predict.pyan example of how to just do prediction using a previously-trained model not included in thecompanion paper.This file requires that an unmodifiedexample.pybe run first. This file shares the same flags asexample.py.more_examples/example_predict_phanEtAl.pyan example of how to do predictions using the model trained in Phan et. al's paper,Gas permeability, diffusivity, and solubility in polymers.
I (@rishigurnani) am more than happy to answer any questions about this codebase. If you encounter any troubles, please open a new Issue in the "Issues" tab and I will promptly respond. In addition, if you discover any bugs or have any suggestions to improve the codebase (documentation, features, etc.) please also open a new Issue. This is the power of open source!
If you use this repository in your work please consider citing the origial polyGNN paper.
@article{Gurnani2023,
annote = {doi: 10.1021/acs.chemmater.2c02991},
author = {Gurnani, Rishi and Kuenneth, Christopher and Toland, Aubrey and Ramprasad, Rampi},
doi = {10.1021/acs.chemmater.2c02991},
issn = {0897-4756},
journal = {Chemistry of Materials},
month = {feb},
number = {4},
pages = {1560--1567},
publisher = {American Chemical Society},
title = {{Polymer Informatics at Scale with Multitask Graph Neural Networks}},
url = {https://doi.org/10.1021/acs.chemmater.2c02991},
volume = {35},
year = {2023}
}
And, if relevant, the subsequent papers which have used polyGNN
- Gas permeability, diffusivity, and solubility in polymers
- AI-assisted discovery of high-temperature dielectrics for energy storage
The results shown in thecompanion paperwere generated using v0.2.0 of this package.
This repository is protected under a General Public Use License Agreement, the details of which can be found inGT Open Source General Use License.pdf.