Neural network interatomic potentials (NNPs) have recently proven to be powerful tools to accurately model complex molecular systems while bypassing the high numerical cost of ab initio molecular dynamics simulations. In recent years, numerous advances in model architectures as well as the development of hybrid models combining machine-learning (ML) with more traditional, physically motivated, force-field interactions have considerably increased the design space of ML potentials. In this paper, we present FeNNol, a new library for building, training, and running force-field-enhanced neural network potentials. It provides a flexible and modular system for building hybrid models, allowing us to easily combine state-of-the-art embeddings with ML-parameterized physical interaction terms without the need for explicit programming. Furthermore, FeNNol leverages the automatic differentiation and just-in-time compilation features of the Jax Python library to enable fast evaluation of NNPs, shrinking the performance gap between ML potentials and standard force-fields. This is demonstrated with the popular ANI-2x model reaching simulation speeds nearly on par with the AMOEBA polarizable force-field on commodity GPUs (graphics processing units). We hope that FeNNol will facilitate the development and application of new hybrid NNP architectures for a wide range of molecular simulation problems.
REFERENCES
1.
J.
Wang
, R. M.
Wolf
, J. W.
Caldwell
, P. A.
Kollman
, and D. A.
Case
, J. Comput. Chem.
25
, 1157
(2004
).2.
K.
Vanommeslaeghe
, E.
Hatcher
, C.
Acharya
, S.
Kundu
, S.
Zhong
, J.
Shim
, E.
Darian
, O.
Guvench
, P.
Lopes
, I.
Vorobyov
, and A. D.
Mackerell
, Jr., J. Comput. Chem.
31
, 671
(2010
).3.
W. L.
Jorgensen
, D. S.
Maxwell
, and J.
Tirado-Rives
, J. Am. Chem. Soc.
118
, 11225
(1996
).4.
N.
Gresh
, G. A.
Cisneros
, T. A.
Darden
, and J.-P.
Piquemal
, J. Chem. Theory Comput.
3
, 1960
(2007
).5.
S.
Naseem-Khan
, L.
Lagardère
, C.
Narth
, G. A.
Cisneros
, P.
Ren
, N.
Gresh
, and J.-P.
Piquemal
, J. Chem. Theory Comput.
18
(6
), 3607
–3621
(2022
).6.
J.-P.
Piquemal
, G. A.
Cisneros
, P.
Reinhardt
, N.
Gresh
, and T. A.
Darden
, J. Chem. Phys.
124
, 104101
(2006
).7.
R. E.
Duke
, O. N.
Starovoytov
, J.-P.
Piquemal
, and G. A.
Cisneros
, J. Chem. Theory Comput.
10
, 1361
(2014
).8.
S. K.
Reddy
, S. C.
Straight
, P.
Bajaj
, C.
Huy Pham
, M.
Riera
, D. R.
Moberg
, M. A.
Morales
, C.
Knight
, A. W.
Götz
, and F.
Paesani
, J. Chem. Phys.
145
, 194504
(2016
).9.
X.
Zhu
, M.
Riera
, E. F.
Bull-Vulpe
, and F.
Paesani
, J. Chem. Theory Comput.
19
, 3551
(2023
).10.
P.
Ren
and J. W.
Ponder
, J. Phys. Chem. B
107
, 5933
(2003
).11.
J. W.
Ponder
, C.
Wu
, P.
Ren
, V. S.
Pande
, J. D.
Chodera
, M. J.
Schnieders
, I.
Haque
, D. L.
Mobley
, D. S.
Lambrecht
, R. A.
DiStasio
, Jr. et al, J. Phys. Chem. B
114
, 2549
(2010
).12.
C.
Liu
, J.-P.
Piquemal
, and P.
Ren
, J. Chem. Theory Comput.
15
, 4122
(2019
).13.
C.
Liu
, J.-P.
Piquemal
, and P.
Ren
, J. Phys. Chem. Lett.
11
(2
), 419
(2020
).14.
G.
Nawrocki
, I.
Leontyev
, S.
Sakipov
, M.
Darkhovskiy
, I.
Kurnikov
, L.
Pereyaslavets
, G.
Kamath
, E.
Voronina
, O.
Butin
, A.
Illarionov
et al, J. Chem. Theory Comput.
18
, 7751
(2022
).15.
J.
Melcr
and J.-P.
Piquemal
, Front. Mol. Biosci.
6
, 143
(2019
).16.
Y.
Shi
, P.
Ren
, M.
Schnieders
, and J.-P.
Piquemal
, “Polarizable force fields for biomolecular modeling
,” in Reviews in Computational Chemistry
(John Wiley and Sons, Ltd.
, 2015
), Vol. 28
, Chap. 2, pp. 51
–86
.17.
Z.
Jing
, C.
Liu
, S. Y.
Cheng
, R.
Qi
, B. D.
Walker
, J.-P.
Piquemal
, and P.
Ren
, Annu. Rev. Biophys.
48
, 371
(2019
).18.
19.
B.
Walker
, C.
Liu
, E.
Wait
, and P.
Ren
, J. Comput. Chem.
43
, 1530
(2022
).20.
Y.
Wang
, J.
Fass
, B.
Kaminow
, J. E.
Herr
, D.
Rufa
, I.
Zhang
, I.
Pulido
, M.
Henry
, H. E. B.
Macdonald
, K.
Takaba
, and J. D.
Chodera
, Chem. Sci.
13
, 12016
(2022
).21.
G.
Chen
, T. J.
Inizan
, T.
Plé
, L.
Lagardere
, J.-P.
Piquemal
, and Y.
Maday
, J. Chem. Theory Comput.
20
(13
), 5558
–5569
(2024
).22.
A. P.
Thompson
, H. M.
Aktulga
, R.
Berger
, D. S.
Bolintineanu
, W. M.
Brown
, P. S.
Crozier
, P. J.
in’t Veld
, A.
Kohlmeyer
, S. G.
Moore
, T. D.
Nguyen
et al, Comput. Phys. Commun.
271
, 108171
(2022
).23.
R.
Salomon-Ferrer
, D. A.
Case
, and R. C.
Walker
, Wiley Interdiscip. Rev.: Comput. Mol. Sci.
3
(2
), 198
–210
(2013
).24.
D.
Van Der Spoel
, E.
Lindahl
, B.
Hess
, G.
Groenhof
, A. E.
Mark
, and H. J.
Berendsen
, J. Comput. Chem.
26
, 1701
(2005
).25.
J. C.
Phillips
, R.
Braun
, W.
Wang
, J.
Gumbart
, E.
Tajkhorshid
, E.
Villa
, C.
Chipot
, R. D.
Skeel
, L.
Kale
, and K.
Schulten
, J. Comput. Chem.
26
, 1781
(2005
).26.
L.
Lagardère
, L.-H.
Jolly
, F.
Lipparini
, F.
Aviat
, B.
Stamm
, Z. F.
Jing
, M.
Harger
, H.
Torabifard
, G. A.
Cisneros
, M. J.
Schnieders
, N.
Gresh
, Y.
Maday
, P. Y.
Ren
, J. W.
Ponder
, and J.-P.
Piquemal
, Chem. Sci.
9
, 956
(2018
).27.
O.
Adjoua
, L.
Lagardère
, L.-H.
Jolly
, A.
Durocher
, T.
Very
, I.
Dupays
, Z.
Wang
, T.
Jaffrellot Inizan
, F.
Célerse
, P.
Ren
, J. W.
Ponder
, and J.-P.
Piquemal
, J. Chem. Theory Comput.
17
, 2034
(2021
).28.
J.
Behler
and M.
Parrinello
, Phys. Rev. Lett.
98
, 146401
(2007
).29.
A. P.
Bartók
, M. C.
Payne
, R.
Kondor
, and G.
Csányi
, Phys. Rev. Lett.
104
, 136403
(2010
).30.
S.
Chmiela
, H. E.
Sauceda
, K.-R.
Müller
, and A.
Tkatchenko
, Nat. Commun.
9
, 3887
(2018
).31.
32.
K.
Shakouri
, J.
Behler
, J.
Meyer
, and G.-J.
Kroes
, J. Phys. Chem. Lett.
8
, 2131
(2017
).33.
J. S.
Smith
, O.
Isayev
, and A. E.
Roitberg
, Chem. Sci.
8
, 3192
(2017
).34.
K. T.
Schütt
, F.
Arbabzadah
, S.
Chmiela
, K. R.
Müller
, and A.
Tkatchenko
, Nat. Commun.
8
, 13890
(2017
).35.
K.
Schütt
, P.-J.
Kindermans
, H. E.
Sauceda Felix
, S.
Chmiela
, A.
Tkatchenko
, and K.-R.
Müller
, in Advances in Neural Information Processing Systems
[NIPS (NeurIPS Proceedings)
, 2017
], Vol. 30
.36.
J.
Gilmer
, S. S.
Schoenholz
, P. F.
Riley
, O.
Vinyals
, and G. E.
Dahl
, in International Conference on Machine Learning
(PMLR
, 2017
), pp. 1263
–1272
.37.
N.
Lubbers
, J. S.
Smith
, and K.
Barros
, J. Chem. Phys.
148
, 241715
(2018
).38.
R.
Zubatyuk
, J. S.
Smith
, J.
Leszczynski
, and O.
Isayev
, Sci. Adv.
5
, eaav6490
(2019
).39.
V.
Zaverkin
and J.
Kästner
, J. Chem. Theory Comput.
16
, 5410
(2020
).40.
Z.
Qiao
, A. S.
Christensen
, M.
Welborn
, F. R.
Manby
, A.
Anandkumar
, and T. F.
Miller
III, Proc. Natl. Acad. Sci. U. S. A.
119
(31
), e2205221119
(2022
).41.
O. T.
Unke
, S.
Chmiela
, M.
Gastegger
, K. T.
Schütt
, H. E.
Sauceda
, and K.-R.
Müller
, Nat. Commun.
12
, 7273
(2021
).42.
J.
Gasteiger
, F.
Becker
, and S.
Günnemann
, in Advances in Neural Information Processing Systems
[NIPS (NeurIPS Proceedings)
, 2021
], Vol. 34
, p. 6790
.43.
S.
Batzner
, A.
Musaelian
, L.
Sun
, M.
Geiger
, J. P.
Mailoa
, M.
Kornbluth
, N.
Molinari
, T. E.
Smidt
, and B.
Kozinsky
, Nat. Commun.
13
, 2453
(2022
).44.
A.
Musaelian
, S.
Batzner
, A.
Johansson
, L.
Sun
, C. J.
Owen
, M.
Kornbluth
, and B.
Kozinsky
, Nat. Commun.
14
, 579
(2023
).45.
A.
Grisafi
and M.
Ceriotti
, J. Chem. Phys.
151
, 204105
(2019
).46.
A.
Grisafi
, J.
Nigam
, and M.
Ceriotti
, Chem. Sci.
12
, 2078
(2021
).47.
C. M.
Handley
and P. L.
Popelier
, J. Chem. Theory Comput.
5
, 1474
(2009
).48.
N.
Artrith
, T.
Morawietz
, and J.
Behler
, Phys. Rev. B
83
, 153101
(2011
).49.
T. W.
Ko
, J. A.
Finkler
, S.
Goedecker
, and J.
Behler
, Nat. Commun.
12
, 398
(2021
).50.
O. T.
Unke
and M.
Meuwly
, J. Chem. Theory Comput.
15
, 3678
(2019
).51.
C.
Qu
, Q.
Yu
, P. L.
Houston
, R.
Conte
, A.
Nandi
, and J. M.
Bowman
, “Interfacing q-AQUA with a polarizable force field: The best of both worlds
,” J. Chem. Theory Comput.
19
(12
), 3446
–3459
(2023
).52.
L.
Yang
, J.
Li
, F.
Chen
, and K.
Yu
, J. Chem. Phys.
157
, 214108
(2022
).53.
L.
Zhang
, H.
Wang
, M. C.
Muniz
, A. Z.
Panagiotopoulos
, R.
Car
, and W.
E
, J. Chem. Phys.
156
, 124107
(2022
).54.
J. M.
Bowman
, C.
Qu
, R.
Conte
, A.
Nandi
, P. L.
Houston
, and Q.
Yu
, J. Chem. Theory Comput.
19
, 1
(2022
).55.
J.
Chen
and K.
Yu
, J. Chem. Theory Comput.
20
, 253
(2023
).56.
Y.
Wang
, T. J.
Inizan
, C.
Liu
, J.-P.
Piquemal
, and P.
Ren
, J. Phys. Chem. B
128
, 2381
(2024
).57.
T.
Plé
, L.
Lagardère
, and J.-P.
Piquemal
, Chem. Sci.
14
, 12554
(2023
).58.
M.
Pinheiro
, F.
Ge
, N.
Ferré
, P. O.
Dral
, and M.
Barbatti
, Chem. Sci.
12
, 14396
(2021
).59.
E.
Kocer
, T. W.
Ko
, and J.
Behler
, Annu. Rev. Phys. Chem.
73
, 163
(2022
).60.
K. T.
Schütt
, S. S.
Hessmann
, N. W.
Gebauer
, J.
Lederer
, and M.
Gastegger
, J. Chem. Phys.
158
, 144801
(2023
).61.
P. O.
Dral
, F.
Ge
, Y.-F.
Hou
, P.
Zheng
, Y.
Chen
, M.
Barbatti
, O.
Isayev
, C.
Wang
, B.-X.
Xue
, M.
Pinheiro
, Jr. et al, J. Chem. Theory Comput.
20
, 1193
(2024
).62.
J.
Zeng
, D.
Zhang
, D.
Lu
, P.
Mo
, Z.
Li
, Y.
Chen
, M.
Rynik
, L.
Huang
, Z.
Li
, S.
Shi
et al, J. Chem. Phys.
159
, 054801
(2023
).63.
T.
Bereau
, D.
Andrienko
, and O. A.
Von Lilienfeld
, J. Chem. Theory Comput.
11
, 3225
(2015
).64.
Z. L.
Glick
, A.
Koutsoukas
, D. L.
Cheney
, and C. D.
Sherrill
, J. Chem. Phys.
154
, 224103
(2021
).65.
M.
Thurlemann
, L.
Boselt
, and S.
Riniker
, J. Chem. Theory Comput.
18
, 1701
(2022
).66.
J.
Bradbury
, R.
Frostig
, P.
Hawkins
, M. J.
Johnson
, C.
Leary
, D.
Maclaurin
, G.
Necula
, A.
Paszke
, J.
VanderPlas
, S.
Wanderman-Milne
, and Q.
Zhang
, JAX: composable transformations of Python + NumPy programs
, 2018
.67.
I. B.
DeepMind
, K.
Baumli
, A.
Bell
, S.
Bhupatiraju
, J.
Bruce
, P.
Buchlovsky
, D.
Budden
, T.
Cai
, A.
Clark
, I.
Danihelka
, A.
Dedieu
, C.
Fantacci
, J.
Godwin
, C.
Jones
, R.
Hemsley
, T.
Hennigan
, M.
Hessel
, S.
Hou
, S.
Kapturowski
, T.
Keck
, I.
Kemaev
, M.
King
, M.
Kunesch
, L.
Martens
, H.
Merzic
, V.
Mikulik
, T.
Norman
, G.
Papamakarios
, J.
Quan
, R.
Ring
, F.
Ruiz
, A.
Sanchez
, L.
Sartran
, R.
Schneider
, E.
Sezener
, S.
Spencer
, S.
Srinivasan
, M.
Stanojević
, W.
Stokowiec
, L.
Wang
, G.
Zhou
, and F.
Viola
, The DeepMind JAX Ecosystem
, 2020
.68.
C.
Devereux
, J. S.
Smith
, K. K.
Huddleston
, K.
Barros
, R.
Zubatyuk
, O.
Isayev
, and A. E.
Roitberg
, J. Chem. Theory Comput.
16
, 4192
(2020
).69.
T.
Jaffrelot Inizan
, T.
Plé
, O.
Adjoua
, P.
Ren
, H.
Gökcan
, O.
Isayev
, L.
Lagardère
, and J.-P.
Piquemal
, Chem. Sci.
14
, 5438
(2023
).70.
J.
Heek
, A.
Levskaya
, A.
Oliver
, M.
Ritter
, B.
Rondepierre
, A.
Steiner
, and M.
van Zee
, Flax: A neural network library and ecosystem for JAX
, 2023
.71.
P. P.
Ewald
, Ann. Phys.
369
, 253
(1921
).72.
73.
M.
Rupp
, A.
Tkatchenko
, K.-R.
Müller
, and O. A.
Von Lilienfeld
, Phys. Rev. Lett.
108
, 058301
(2012
).74.
M.
Eckhoff
and M.
Reiher
, J. Chem. Theory Comput.
19
, 3509
(2023
).75.
L.
Zhang
, J.
Han
, H.
Wang
, W.
Saidi
, R.
Car
et al, in Advances in Neural Information Processing Systems
[NIPS (NeurIPS Proceedings)
, 2018
], Vol. 31
.76.
K.
Schütt
, O.
Unke
, and M.
Gastegger
, in International Conference on Machine Learning
(PMLR
, 2021
), pp. 9377
–9388
.77.
M.
Haghighatlari
, J.
Li
, X.
Guan
, O.
Zhang
, A.
Das
, C. J.
Stein
, F.
Heidar-Zadeh
, M.
Liu
, M.
Head-Gordon
, L.
Bertels
et al, Digital Discovery
1
, 333
(2022
).78.
79.
B.
Kozinsky
, A.
Musaelian
, A.
Johansson
, and S.
Batzner
, in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis
[ACM (Association for Computing Machinery)
, 2023
], pp. 1
–12
.80.
S.
Takamoto
, S.
Izumi
, and J.
Li
, Comput. Mater. Sci.
207
, 111280
(2022
).81.
U.
Essmann
, L.
Perera
, M. L.
Berkowitz
, T.
Darden
, H.
Lee
, and L. G.
Pedersen
, J. Chem. Phys.
103
, 8577
(1995
).82.
A.
Khabibrakhmanov
, D. V.
Fedorov
, and A.
Tkatchenko
, J. Chem. Theory Comput.
19
, 7895
(2023
).83.
E.
Caldeweyher
, S.
Ehlert
, A.
Hansen
, H.
Neugebauer
, S.
Spicher
, C.
Bannwarth
, and S.
Grimme
, J. Chem. Phys.
150
, 154122
(2019
).84.
J. F.
Ziegler
and J. P.
Biersack
, in Treatise on Heavy-Ion Science: Volume 6: Astrophysics, Chemistry, and Condensed Matter
(Springer
, 1985
), pp. 93
–129
.85.
E. J.
Hu
, Y.
Shen
, P.
Wallis
, Z.
Allen-Zhu
, Y.
Li
, S.
Wang
, L.
Wang
, and W.
Chen
, arXiv:2106.09685 (2021
).86.
I.
Batatia
, D. P.
Kovacs
, G.
Simm
, C.
Ortner
, and G.
Csányi
, in Advances in Neural Information Processing Systems
[NIPS (NeurIPS Proceedings)
, 2022
], Vol. 35
, p. 11423
.87.
J.
Zhuang
, T.
Tang
, Y.
Ding
, S. C.
Tatikonda
, N.
Dvornek
, X.
Papademetris
, and J.
Duncan
, in Advances in Neural Information Processing Systems
[NIPS (NeurIPS Proceedings)
, 2020
], Vol. 33
, p. 18795
.88.
89.
A.
Brock
, S.
De
, S. L.
Smith
, and K.
Simonyan
, in International Conference on Machine Learning
(PMLR
, 2021
), pp. 1059
–1071
.90.
L. N.
Smith
and N.
Topin
, “Artificial intelligence and machine learning for multi-domain operations applications
,” Proc. SPIE
11006
, 369
–386
(2019
).91.
M.
Kellner
and M.
Ceriotti
, Mach. Learn.: Sci. Technol.
5
, 035006
(2024
).92.
A.
Amini
, W.
Schwarting
, A.
Soleimany
, and D.
Rus
, in Advances in Neural Information Processing Systems
[NIPS (NeurIPS Proceedings)
, 2020
], Vol. 33
, p. 14927
.93.
94.
A. H.
Larsen
, J. J.
Mortensen
, J.
Blomqvist
, I. E.
Castelli
, R.
Christensen
, M.
Dułak
, J.
Friis
, M. N.
Groves
, B.
Hammer
, C.
Hargus
et al, J. Phys.: Condens. Matter
29
, 273002
(2017
).95.
V.
Kapil
, M.
Rossi
, O.
Marsalek
, R.
Petraglia
, Y.
Litman
, T.
Spura
, B.
Cheng
, A.
Cuzzocrea
, R. H.
Meißner
, D. M.
Wilkins
et al, Comput. Phys. Commun.
236
, 214
(2019
).96.
T.
Plé
, N.
Mauger
, O.
Adjoua
, T. J.
Inizan
, L.
Lagardère
, S.
Huppert
, and J.-P.
Piquemal
, J. Chem. Theory Comput.
19
, 1432
(2023
).97.
S.-L. J.
Lahey
and C. N.
Rowley
, Chem. Sci.
11
, 2362
(2020
).98.
R.
Galvelis
, A.
Varela-Rial
, S.
Doerr
, R.
Fino
, P.
Eastman
, T. E.
Markland
, J. D.
Chodera
, and G.
De Fabritiis
, J. Chem. Inf. Model.
63
, 5701
(2023
).99.
A.
Illarionov
, S.
Sakipov
, L.
Pereyaslavets
, I. V.
Kurnikov
, G.
Kamath
, O.
Butin
, E.
Voronina
, I.
Ivahnenko
, I.
Leontyev
, G.
Nawrocki
et al, J. Am. Chem. Soc.
145
, 23620
(2023
).100.
B.
Leimkuhler
and C.
Matthews
, Appl. Math. Res. eXpress
2013
(1
), 34
–56
.101.
E.
Mangaud
, S.
Huppert
, T.
Plé
, P.
Depondt
, S.
Bonella
, and F.
Finocchi
, J. Chem. Theory Comput.
15
, 2863
(2019
).102.
S.
Habershon
, D. E.
Manolopoulos
, T. E.
Markland
, and T. F.
Miller
III, Annu. Rev. Phys. Chem.
64
, 387
(2013
).103.
N.
Mauger
, T.
Plé
, L.
Lagardère
, S.
Bonella
, E.
Mangaud
, J.-P.
Piquemal
, and S.
Huppert
, J. Phys. Chem. Lett.
12
, 8285
(2021
).104.
X.
Gao
, F.
Ramezanghorbani
, O.
Isayev
, J. S.
Smith
, and A. E.
Roitberg
, J. Chem. Inf. Model.
60
, 3408
(2020
).105.
L.
Lagardere
, F.
Aviat
, and J.-P.
Piquemal
, J. Phys. Chem. Lett.
10
, 2593
(2019
).106.
J.
Hoja
, L.
Medrano Sandonas
, B. G.
Ernst
, A.
Vazquez-Mayagoitia
, R. A.
DiStasio
, Jr., and A.
Tkatchenko
, Sci. Data
8
, 43
(2021
).107.
J. S.
Smith
, R.
Zubatyuk
, B.
Nebgen
, N.
Lubbers
, K.
Barros
, A. E.
Roitberg
, O.
Isayev
, and S.
Tretiak
, Sci. Data
7
, 134
(2020
).108.
P.
Eastman
, P. K.
Behara
, D. L.
Dotson
, R.
Galvelis
, J. E.
Herr
, J. T.
Horton
, Y.
Mao
, J. D.
Chodera
, B. P.
Pritchard
, Y.
Wang
et al, Sci. Data
10
, 11
(2023
).109.
T.
Verstraelen
, S.
Vandenbrande
, F.
Heidar-Zadeh
, L.
Vanduyfhuys
, V.
Van Speybroeck
, M.
Waroquier
, and P. W.
Ayers
, J. Chem. Theory Comput.
12
, 3894
(2016
).110.
R.
Zubatyuk
, J. S.
Smith
, B. T.
Nebgen
, S.
Tretiak
, and O.
Isayev
, Nat. Commun.
12
, 4870
(2021
).© 2024 Author(s). Published under an exclusive license by AIP Publishing.
2024
Author(s)
You do not currently have access to this content.
