MOLPIPx: an end-to-end differentiable package for permutationally invariant polynomials in Python and Rust

TL;DR

MOLPIPx is a new Python and Rust library that enables efficient, differentiable permutationally invariant polynomials for molecular modeling, integrating with machine learning frameworks for energy surface parameterization.

Contribution

It introduces MOLPIPx, a versatile, automatic differentiation-enabled library for PIPs, facilitating advanced molecular modeling and simulations.

Findings

Supports energy gradient computations with JAX and EnzymeAD-Rust.

Enables integration of PIPs with neural networks and Gaussian processes.

Improves efficiency and flexibility in potential energy surface modeling.

Abstract

In this work, we present MOLPIPx, a versatile library designed to seamlessly integrate Permutationally Invariant Polynomials (PIPs) with modern machine learning frameworks, enabling the efficient development of linear models, neural networks, and Gaussian process models. These methodologies are widely employed for parameterizing potential energy surfaces across diverse molecular systems. MOLPIPx leverages two powerful automatic differentiation engines -JAX and EnzymeAD-Rust- to facilitate the efficient computation of energy gradients and higher-order derivatives, which are essential for tasks such as force field development and dynamic simulations. MOLPIPx is available at https://github.com/ChemAI-Lab/molpipx.