Lin, Dahiya, Cersonsky propose generalized approach to incorporating geometry and directionality into coarse-grained machine-learned potentials University of Wisconsin–Madison; enables anisotropic coarse-grained simulations

Lin, Dahiya, Cersonsky propose generalized approach to incorporating geometry and directionality into coarse-grained machine-learned potentials University of Wisconsin–Madison; enables anisotropic coarse-grained simulations

A Generalized Approach for Incorporating Geometry and Directionality into Coarse-Grained Machine-Learned Potentials

A Generalized Approach for Incorporating Geometry and Directionality into Coarse-Grained Machine-Learned Potentials Arthur Y. Lin,1 Tejas Dahiya,2 and Rose K. Cersonsky1, 3, 4, a) 1)Department of Chemical and Biological Engineering, University of Wisconsin - Madison, Madison, WI, USA 2)Department of Computer Science, University of Wisconsin - Madison, Madison, WI, USA 3)Department of Materials Science and Engineering, University of Wisconsin - Madison, Madison, WI, USA 4)Data Science Institute, University of Wisconsin - Madison, Madison, WI, USA I. INTRODUCTION Machine-learned interatomic potentials (MLIPs) have dramatically expanded the scope of molecular simulation, enabling the prediction of atomic energies and forces with near ab initio accuracy at costs compatible with large- scale molecular dynamics.1–3 As these models continue to mature, the principal obstacle to computational discovery is increasingly shifting from the accuracy of atomistic interactions to the challenge of accessing the larger length and time scales associated with collective behavior, self- assembly, and phase transitions. Coarse-graining (CG), in which groups of atoms are mapped onto effective particles with reduced degrees of freedom, remains one of the most powerful approaches for bridging this gap.4–8 For any MLIP, its capabilities are determined jointly by the quality of the underlying data, the choice of repre- sentation X, and the architecture f(X), where X varies from simple forms (e.g., distance matrices or graphs) to more complex descriptors (e.g., Smooth Overlap of Atomic Positions [SOAP]9 or Behler-Parinello symme- try functions10). A representation is complete if distinct data objects map to distinct data points; in the atomistic MLIP community there has been considerable discussion on what constitutes a complete descriptor, as this has direct implications to the ceiling of predictive accuracy in subsequent models.11–14. Simply put, while expressive architectures can approximate complex functions f(X), they cannot recover information that is not encoded in X, and degeneracy in data representation can explicitly limit model performance. However, unlike atomistic machine-learned potentials, the variables describing a coarse-grained system are not fixed, but instead chosen. Historically, much of coarse-graining has therefore been built around isotropic particles.4 This choice is attractive for both concep- tual and computational reasons. Interactions become functions of intermolecular separation alone, simulation methodologies are well-established, and many successful coarse-grained models have been developed within this framework.15 Such approaches have proven remarkably a)Electronic mail: rose.cersonsky@wisc.edu successful, particularly in biomolecular systems where coarse-graining has enabled simulations inaccessible at atomistic resolution.15 Yet there are also many examples where isotropic descriptions struggle to reproduce experi- mentally observed behavior, particularly in systems whose organization is governed by packing, local structure, or directional interactions.16,17 One possible explanation of these difficulties is that geometry itself carries critical information. Consider two molecular configurations possessing similar intermolecular separations but different relative orientations. Depending on the system, these configurations may exhibit substan- tially different interaction energies. When orientation is removed from the coarse-grained representation, how- ever, both are represented by the same set of variables R. The resultant coarse-grained configuration is there- fore associated not with a single underlying interaction energy, but with a distribution of possible interaction en- ergies. As this distribution broadens, the construction of a coarse-grained potential becomes increasingly difficult, as distinct molecular environments become indistinguishable within the representation. This observation naturally suggests retaining orienta- tional and geometric information within coarse-grained models. However, doing so raises a separate question: how should such information be represented? While atomistic machine-learned potentials have undergone rapid develop- ment over the past decade, the overwhelming majority of descriptor frameworks and architectures were developed for collections of isotropic atoms. Extending these ap- proaches to anisotropic particles requires representations that remain sensitive to particle geometry while preserv- ing the rotational and translational symmetries of the underlying problem. Several recent approaches have begun to address this challenge through anisotropic extensions to machine- learning architectures.18? ,19 Of particular interest are many-body density expansion methods, which have proven highly successful in atomistic machine learning due to their systematic improvability and clear connection to local structure.9,20,21 AniSOAP extended the SOAP formal- ism to non-spherical particles through anisotropic density expansions, enabling the explicit incorporation of particle shape and orientation into a many-body descriptor,22,23, and similar efforts expanded the Chebyshev polynomial arXiv:2609.01911v1 [physics.chem-ph] 1 Sep 2026

2 suite for nanoscale particles.24–26. In this work, we demonstrate the strategies and subse- quent trade-offs for incorporating geometric information into MLIPs. First, we show how incorporating geometric tensors, both through AniSOAP and as node features in MACE, compare to their isotropic analogs as a function of molecular symmetry and objective weighting. Then, we demonstrate how to successfully incorporate geometric in- formation into MACE-style potentials for arbitrary molec- ular symmetries, in ways that extend to coarse-graining at any resolution, even that of rigid composite bodies typically found in the nanoscale community. We do so for single-bead coarse-grained systems with minimal mapping entropy in order to demonstrate the limitations of poten- tial energy fitting for coarse-graining, but will note that analogous schemes can be applied to learn the potential of mean force. Our coinciding packages, AniSOAP23 and MACE-CG are also included open-source, to allow for adoption of this pipeline for new coarse-graining tasks. II. METHODS A. AniSOAP Descriptors Many-body density expansions provide a natural frame- work for constructing machine-learning representations of local environments. In conventional SOAP-based ap- proaches, each particle is represented by an isotropic Gaussian density centered at its position. AniSOAP22,23 extends this formalism to anisotropic particles by replac- ing isotropic densities with ellipsoidal density fields that explicitly depend on particle geometry and orientation. Each coarse-grained particle is described by a position ri, an orientation matrix Ri, and geometric parameters corresponding to its principal semiaxes (ai, bi, ci). The particle geometry is encoded through the matrix Σi = Ri   a2 i 0 0 0 b2 i 0 0 0 c2 i  RT i . (1) The local density surrounding particle i is then written as ρi(r) = X j exp

−1 2(r −rij)T Σ−1 j (r −rij)

, (2) where rij = rj −ri denotes the position of neighboring particle j relative to particle i. Following the SOAP formalism, the density is expanded in a basis of radial functions and spherical harmonics, ρi(r) = X nlm c(i) nlmgn(r)Ylm(ˆr), (3) producing a set of rotationally equivariant coefficients c(i) nlm. Rotationally invariant features are subsequently ob- tained through contractions of these coefficients. Through- out this work, we employ the AniSOAP power spectrum, p(i) nn′l = X m c(i) nlmc(i)∗ n′lm, (4) which forms the descriptor vector xi used for regression. By construction, the descriptor is invariant to global rotations and translations while remaining sensitive to particle geometry and local orientational environments. We can consider linear energy models acting on the AniSOAP power spectrum written as E = X i xT i w (5) where xi denotes the AniSOAP descriptor associated with particle i and w is a learned coefficient vector. B. Incorporating Rigid-Body Orientation into MACE As a complementary approach to AniSOAP, we extended the MACE equivariant message-passing framework to de- scribe coarse-grained particles with explicit rigid-body orientation. The central objective was to preserve the existing MACE treatment of translational geometry while augmenting each intermolecular edge with orientational information. In this construction, the molecular symme- try is incorporated into the representation of the body orientation before that representation enters the MACE interaction network. 1. Standard MACE edge geometry For a pair of coarse-grained particles i and j, conven- tional MACE begins from the relative displacement rij = rj −ri, (6) which is decomposed into the scalar separation rij = ∥rij∥ (7) and the unit edge direction ˆrij = rij rij . (8) The distance is expanded in a radial basis, while the angular dependence is represented by spherical harmonics, Yℓm(ˆrij). (9) Under a global rotation S ∈SO(3), these functions trans- form according to the ℓth irreducible representation, Yℓ(Sˆrij) = D(ℓ)(S)Yℓ(ˆrij), (10)

3 FIG. 1: Workflow for incorporating molecular geometry and directionality into MACE-style potentials. which allows the edge geometry to be coupled to node features through the Clebsch–Gordan tensor products used throughout MACE. For an atomistic system, the combination of particle identity and relative positions is sufficient to define the local geometry. This is no longer true when an entire molecule is mapped to a single coarse-grained site. Two molecular configurations can possess identical center-of- mass positions, and therefore identical rij and Yℓm(ˆrij), while differing substantially in molecular orientation and interaction energy. A single-site rigid-body model there- fore requires an additional orientational degree of freedom. 2. Rigid-body state and orientation Each coarse-grained molecule is represented by (ri, Ri, si) , (11) where ri is the molecular center-of-mass position, Ri ∈ SO(3) is a rotation matrix mapping a fixed molecular body frame to the simulation frame (we will call this the “simulation-frame” rotation), and si denotes molec- ular species. In practice, orientations are stored as unit quaternions and converted to rotation matrices before constructing equivariant features. Writing the columns of Ri as the simulation-frame body axes, Ri =   | | | ai bi ci | | |  , (12) a global rotation S acts as ri →Sri, Ri →SRi. (13) The body axes therefore transform as ordinary equivariant vectors and can, in principle, be coupled directly to the spherical-harmonic edge features using the same tensor- product machinery already present in MACE. A full-frame rigid-body representation can consequently be constructed from the equivariant product Y (ˆrij) ⊗Ri ⊗Rj, (14) where the three columns of each rotation matrix are treated as ℓ= 1 features. Clebsch–Gordan tensor prod- ucts decompose this product into irreducible SO(3) chan- nels, which are subsequently projected into a chosen set of edge irreducible representations. Although Eq. 14 preserves complete body-frame infor- mation, it contains a redundancy whenever the molecule possesses a nontrivial rotational point-group symmetry. 3. Molecular symmetry and the orientation quotient Let G ⊂SO(3) denote the proper rotational symmetry group of a molecule. For any g ∈G, the two body-frame orientations Ri and Rig (15) describe the same physical molecular configuration. The physical orientational degree of freedom therefore belongs to the quotient space [Ri] ∈SO(3)/G. (16) An orientational feature map ΦG(R) should accordingly satisfy two distinct transformation properties. First, a rotation of the entire physical system should act equivari- ant