PyTorch implementation of the model used in Chem. Sci., 2019,10, 7503-7515 ( https://doi.org/10.1039/C9SC02097G )
The original TensorFlow-based model is located at https://github.com/rsdefever/GenStrIde
Below are the required packages and versions used:
- PyTorch (2.5.0)
- Lightning (2.5.2)
- e3nn (0.5.1)
- Tensorboard (2.20.0)
Each input sample is expected to be a set of coordinates of neighboring particles in 3D space centered around the particle of interest at (0,0,0). The overall shape of the input sample file is (n_samples, n_neighbors, 3). In the original work, coordinates were scaled such that the closest neighbor was at unit distance. Class labels are one-hot encoded and the expected shape is (n_samples, n_classes). A small sample of input point clouds is provided in the samples/ directory.
See the example/ directory for scripts to train, test, and run inference with the sample data in samples/. During training, log files are by default written to current_dir/lightning_logs and can be monitored with Tensorboard. Model checkpoint files are written to current_dir/ckpts. See example/runtrain.sh for an example set of options. For the data provided in samples/, the validation accuracy should reach approximately 85% to 90% in 100 epochs.
Trained models can be loaded using the a saved checkpoint and used for inference. See scripts/predict.py and example/runpredict.sh for an example.