references

paper/paper.bib

@@ -67,3 +67,102 @@ @article{Thompson-22-02
   doi     = {10.1016/j.cpc.2021.108171},
   urldate = {2022-11-08}
 }
+
+
+@misc{Simeon-23-06,
+  title         = {{{TensorNet}}: {{Cartesian Tensor Representations}} for {{Efficient Learning}} of {{Molecular Potentials}}},
+  author        = {Simeon, Guillem and {de Fabritiis}, Gianni},
+  year          = {2023},
+  number        = {arXiv:2306.06482},
+  eprint        = {2306.06482},
+  primaryclass  = {physics},
+  archiveprefix = {arXiv}
+}
+
+@misc{Schutt-21-06,
+  title         = {Equivariant Message Passing for the Prediction of Tensorial Properties and Molecular Spectra},
+  author        = {Sch{\"u}tt, Kristof T. and Unke, Oliver T. and Gastegger, Michael},
+  year          = {2021},
+  number        = {arXiv:2102.03150},
+  eprint        = {2102.03150},
+  primaryclass  = {physics},
+  doi           = {10.48550/arXiv.2102.03150},
+  archiveprefix = {arXiv}
+}
+
+@misc{data2objects,
+  author    = {John L. A. Gardner},
+  title     = {data2objects: self-documenting config files for use with Python},
+  year      = {2024},
+  publisher = {GitHub},
+  journal   = {GitHub repository},
+  url       = {https://github.com/jla-gardner/data2objects}
+}
+
+@article{Wang-18-07,
+  title   = {{{DeePMD-kit}}: {{A}} Deep Learning Package for Many-Body Potential Energy Representation and Molecular Dynamics},
+  author  = {Wang, Han and Zhang, Linfeng and Han, Jiequn and E, Weinan},
+  year    = {2018},
+  journal = {Computer Physics Communications},
+  volume  = {228},
+  pages   = {178--184},
+  doi     = {10.1016/j.cpc.2018.03.016}
+}
+
+@article{Zeng-23-08,
+  title   = {{{DeePMD-kit}} v2: {{A}} Software Package for Deep Potential Models},
+  author  = {Zeng, Jinzhe and Zhang, Duo and Lu, Denghui and Mo, Pinghui and Li, Zeyu and Chen, Yixiao and Rynik, Mari{\'a}n and Huang, Li'ang and Li, Ziyao and Shi, Shaochen and Wang, Yingze and Ye, Haotian and Tuo, Ping and Yang, Jiabin and Ding, Ye and Li, Yifan and Tisi, Davide and Zeng, Qiyu and Bao, Han and Xia, Yu and Huang, Jiameng and Muraoka, Koki and Wang, Yibo and Chang, Junhan and Yuan, Fengbo and Bore, Sigbj{\o}rn L{\o}land and Cai, Chun and Lin, Yinnian and Wang, Bo and Xu, Jiayan and Zhu, Jia-Xin and Luo, Chenxing and Zhang, Yuzhi and Goodall, Rhys E. A. and Liang, Wenshuo and Singh, Anurag Kumar and Yao, Sikai and Zhang, Jingchao and Wentzcovitch, Renata and Han, Jiequn and Liu, Jie and Jia, Weile and York, Darrin M. and E, Weinan and Car, Roberto and Zhang, Linfeng and Wang, Han},
+  year    = {2023},
+  journal = {The Journal of Chemical Physics},
+  volume  = {159},
+  number  = {5},
+  pages   = {054801},
+  doi     = {10.1063/5.0155600}
+}
+
+@misc{TorchMDNet,
+  title         = {TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations},
+  author        = {Raul P. Pelaez and Guillem Simeon and Raimondas Galvelis and Antonio Mirarchi and Peter Eastman and Stefan Doerr and Philipp Thölke and Thomas E. Markland and Gianni De Fabritiis},
+  year          = {2024},
+  eprint        = {2402.17660},
+  archiveprefix = {arXiv},
+  primaryclass  = {cs.LG}
+}
+
+@article{schutt2023schnetpack,
+  author  = {Sch{\"u}tt, Kristof T. and Hessmann, Stefaan S. P. and Gebauer, Niklas W. A. and Lederer, Jonas and Gastegger, Michael},
+  title   = {{SchNetPack 2.0: A neural network toolbox for atomistic machine learning}},
+  journal = {The Journal of Chemical Physics},
+  volume  = {158},
+  number  = {14},
+  pages   = {144801},
+  year    = {2023},
+  month   = {04},
+  issn    = {0021-9606},
+  doi     = {10.1063/5.0138367},
+  url     = {https://doi.org/10.1063/5.0138367},
+  eprint  = {https://pubs.aip.org/aip/jcp/article-pdf/doi/10.1063/5.0138367/16825487/144801\_1\_5.0138367.pdf}
+}
+
+@article{schutt2019schnetpack,
+  author  = {Sch{\"u}tt, Kristof T. and Kessel, Pan and Gastegger, Michael and Nicoli, Kim A. and Tkatchenko, Alexandre and Müller, Klaus-Robert},
+  title   = {{SchNetPack: A Deep Learning Toolbox For Atomistic Systems}},
+  journal = {Journal of Chemical Theory and Computation},
+  volume  = {15},
+  number  = {1},
+  pages   = {448-455},
+  year    = {2019},
+  doi     = {10.1021/acs.jctc.8b00908},
+  url     = {https://doi.org/10.1021/acs.jctc.8b00908},
+  eprint  = {https://doi.org/10.1021/acs.jctc.8b00908}
+}
+
+@software{Lightning,
+  author  = {Falcon, William and {The PyTorch Lightning team}},
+  doi     = {10.5281/zenodo.3828935},
+  license = {Apache-2.0},
+  title   = {{PyTorch Lightning}},
+  url     = {https://github.com/Lightning-AI/lightning},
+  version = {1.4},
+  year    = {2019}
+}

paper/paper.md

@@ -27,7 +27,7 @@ bibliography: paper.bib
 
 We present `graph-pes`, an open-source toolkit for accelerating the development, training and deployment of machine-learned potential energy surfaces (ML-PESs) that act on graph-based representations of atomic structures. The `graph-pes` toolkit comprises three components:
 
-1. **the `graph_pes` Python package**: a modular Python package containing all the functionality required to build and train graph-based ML-PES models. This includes a mature data pipeline for converting atomic structures into graph representations (`AtomicGraph`s), a fully featured ML-PES base class with automatic force and stress calcuations (`GraphPESModel`), and a suite of common data manipulations and model building blocks.
+1. **the `graph_pes` Python package**: a modular Python framework containing all the functionality required to build and train graph-based ML-PES models. This includes a mature data pipeline for converting atomic structures into graph representations (`AtomicGraph`s), a fully featured ML-PES base class with automatic force and stress calcuations (`GraphPESModel`), and a suite of common data manipulations and model building blocks.
 
 2. **the `graph-pes-train` command-line interface**: a unified tool for training graph-based ML-PES models on datasets of labelled atomic structures. Several popular model architectures are provided out of the box, but the interface is designed so that custom, end-user defined architectures can also easily be used (alongside custom loss functions, optimizers, datasets, etc.)
 
@@ -36,30 +36,27 @@ We present `graph-pes`, an open-source toolkit for accelerating the development,
 
 # Statement of need
 
-In recent years, machine-learned potential energy surfaces (ML-PESs) have become an indispensable tool for computational chemists. Their high accuracy, favourable scaling with system size, and ease of parallelisation allows ML-PES models to faithfully simulate the dynamics of large systems over long timescales. ML-PESs are thus facilitating the study of complex chemical phenomena at the atomic scale, and hence the generation of novel insight and understanding. [cite]
+In recent years, machine-learned potential energy surfaces (ML-PESs) have become an indispensable tool for computational chemists. Their high accuracy, favourable scaling with system size, and ease of parallelisation allows ML-PES models to faithfully simulate the dynamics of large systems over long timescales. ML-PESs are thus facilitating the study of complex chemical phenomena at the atomic scale, in turn driving the generation of novel insight and understanding. [cite]
 
-Many "flavours" of ML-PES exist, and with them have arisen a variety of softwares that are tailored to training specific architectures (see below). Given their unique specialisations, these varied softwares fail to conform to a common interface, making it difficult for practitioners to migrate their training and validation setups between different architectures.
+Many "flavours" of ML-PES exist, and with them have arisen a variety of softwares that are tailored to training specific architectures (see below). Given their unique specialisations, these varied softwares fail to conform to a common interface, making it difficult for practitioners to migrate their training and validation setups between different model architectures.
 
 `graph_pes` provides a **unified interface and framework** for defining, training, and  working with all ML-PES models that act on graph-based representations of atomic structures. This unified interface has several advantages:
 
-- **graph representations are genral and useful.** A chemical structure is completely defined by the positions of its atoms ($\mathbf{R}$), and their chemical identities ($Z$).[^1] A graph representation of the atomic structure incorporates this complete description, together with a neighbour list ($\mathbf{N}$) indicating which atoms are within the locality of others (for instance using a fixed cutoff radius). The resulting graph ($G = \{\mathbf{R}, Z, \mathbf{N}\}$), is thus an extremely general representation of chemical structure that does not impose any constraints on the architecture of the ML-PES model while easily facilitating the implementation of both local and non-local models.
+- **graph representations are general and useful.** A chemical structure is completely defined by the positions of its atoms ($\mathbf{R}$), and their chemical identities ($Z$).[^1] A graph representation of the atomic structure incorporates this complete description, together with a neighbour list ($\mathbf{N}$) indicating which atoms are within the locality of others (for instance using a fixed cutoff radius). The resulting graph ($G = \{\mathbf{R}, Z, \mathbf{N}\}$), is thus an extremely general representation of chemical structure that does not impose any constraints on the architecture of the ML-PES model while easily facilitating the implementation of both local and non-local models.
 
 [^1]: isolated structure only - ignoring cell and periodic boundary conditions. Plus gorund state, uncharged.
 
-- **transferable training and validation pipelines.** The unified interface of the `graph_pes.GraphPESModel` class and the `graph-pes-train` CLI tool allow researchers to easily transfer their training and validation pipelines between different model architectures. For convenience, we have implemented several popular model architectures out of the box (including PaiNN, NequIP, MACE and TensorNet). Training scripts require as little as 2 lines of code to change to the model architecture. LAMMPS input scripts require no changes other than pointing to the new model's file.
+- **transferable training and validation pipelines.** The unified interface of the `graph_pes.GraphPESModel` class and the `graph-pes-train` CLI tool allow researchers to easily transfer their training and validation pipelines between different model architectures. For convenience, we have implemented several popular model architectures out of the box (including PaiNN [@Schutt-21-06], NequIP [@Batzner-22-05], MACE [@Batatia-23-01] and TensorNet [@Simeon-23-06]). Training scripts require as little as 2 lines of code to change to the model architecture. LAMMPS input scripts require no changes other than pointing to the new model's file.
 
 - **accelerated ML-PES development.** `graph-pes` provides all the functionality required to quickly design new ML-PES architectures, including common and well-documented data manipulations, message passing operations and model building blocks such as distance expansions and neighbour summations. By inheriting from the `graph_pes.GraphPESModel` class, these new architectures can instantly be trained using the `graph-pes-train` CLI tool, and validated using the `pair style graph_pes` LAMMPS pair style.
 
-ML-PES research is not just limited to the development of new model architectures. 
-Among other important research topics, `graph-pes` aims to facilitate:
+ML-PES research is not just limited to the development of new model architectures. Among other important research topics, `graph-pes` aims to address:
 
 - **model fine-tuning.** Various pre-train/fine-tune strategies have been proposed for improving the accuraacy and robustness of ML-PES models. Model fine-tuning in `graph-pes` can be performed using the `graph-pes-train` CLI tool, where users can optionally specify which of the model's parameters should be frozen during training.
 
 - **universal force-fields.** A topical and recent area of research is the development of universal force-fields that can be used to describe the potential energy surface of a wide range of systems. `graph-pes` provides access to the `MACE-OFF` and `MACE-MP0` fuondation models for ease of use and fine-tuning. We use these models as an example of how to incorporate externally trained models into the `graph-pes` framework, either by porting the model's implementation into the `graph_pes` package, or by constructing a custom, lightweight interface to a 3rd party package (in this case, `mace-torch`).
 
-<!-- - **customised training procedures.** Custom optimziers, loss functions etc in the `graph-pes-train` CLI tool, but also well documented examples of how to implement custom training procedures. -->
-
-- **customised training procedures.** The `graph-pes-train` CLI tool supports user-defined optimizers, loss functions, and datasets through the flexible plugin system provided by the `data2objects` package. 
+- **customised training procedures.** The `graph-pes-train` CLI tool supports user-defined optimizers, loss functions, and datasets through the flexible plugin system provided by the `data2objects` package. [@data2objects]
 Beyond this, we provide well documented examples of how to implement custom, architecture-agnostic training procedures in simple Python scripts.
 
 - **beyond numerical validation.** While it is common to benchmark ML-PES models using numerical validation metrics (such as energy and force RMSEs), more extensive and physically motivated validation routines are important. The *unified*, *varied* and *architecture-agnostic* functionalities that `graph-pes` provide (including an `ase` Calculator interface and `LAMMPS` pair style) allow researchers to define a validation procedure once, and then to use it for all their ML-PES models, and to share this validation procedure with the wider community. To aid in these validation procedures, we provide some utility functions for plotting dimer curves, energy volume scans and parity plots. 
@@ -69,21 +66,11 @@ Beyond this, we provide well documented examples of how to implement custom, arc
 
 `graph-pes` is driving a significant number of projects within the Deringer group, and has already been cited in @Liu-24-12.
 
-The core functionality of `graph-pes` builds upon several existing, open-source packages. We use the `Tensor` data structure, `nn` module and `autograd` functionality from `PyTorch` [@Paszke-19] for data representation, model definition and automatic differentiation respectively. We use the `ase` [@HjorthLarsen-17-06] package for reading serialised atomic structures from disc, and for converting them into graph representations. We use the `LAMMPS` [@Thompson-22-02] framework for creating the `pair style graph_pes` command for molecular dynamics simulations.
-
+The core functionality of `graph-pes` builds upon several existing, open-source packages. We use the `Tensor` data structure, `nn` module and `autograd` functionality from `PyTorch` [@Paszke-19] for data representation, model definition and automatic differentiation respectively. We use the `PyTorch Lightning` [@Lightning] to implement our core training loop, as well as for advanced features such as learning rate scheduling, stochastic weight averaging and more.
+We use the `ase` [@HjorthLarsen-17-06] package for reading serialised atomic structures from disc, and for converting them into graph representations. We use the `LAMMPS` [@Thompson-22-02] framework for creating the `pair style graph_pes` command for molecular dynamics simulations.
 We also use the `e3nn` [@Geiger-22-07] package for implementing the `NequIP` [@Batzner-22-05] and `MACE` [@Batatia-23-01] architectures.
 
-Note-worthy softwares that offer training and validation functionaltiy for specific ML-PES architectures include:
-
-- `nequip`
-
-- `mace-torch`
-
-- `deepmd-kit`
-
-- `schnetpack`
-
-- `torchmd`
+Relevant packages that offer training and validation functionaltiy for specific ML-PES architectures include: `nequip` [@Batzner-22-05], `mace-torch` [@Batatia-23-01], `deepmd-kit` [@Wang-18-07; @Zeng-23-08], `schnetpack` [@schutt2019schnetpack; @schutt2023schnetpack] and `torchmd-net` [@TorchMDNet].
 
 
 # Acknowledgements