Machine learning configuration-dependent friction tensors in Langevin heatbaths

Dynamics of coarse-grained particle systems derived via the Mori-Zwanzig projection formalism commonly take the form of a (generalized) Langevin equation with configuration-dependent friction tensor and diffusion coefficient matrix. In this article, we introduce a class of equivariant representations of tensor-valued functions based on the Atomic Cluster Expansion framework that allows for efficient learning of such configuration-dependent friction tensors from data. Besides satisfying the correct equivariance properties with respect to the Euclidean group E(3), the resulting heat bath models satisfy a fluctuation-dissipation relation. We demonstrate the capabilities of the model approach by fitting a model of configuration-dependent tensorial electronic friction calculated from first principles that arises during reactive molecular dynamics at metal surfaces.