MolCrystalFlow: Molecular Crystal Structure Prediction via Flow Matching

TL;DR

MolCrystalFlow tackles molecular crystal CSP by learning a joint, periodic, E(3)-invariant flow over lattice $L$, fractional centroids $F$, and orientations $R$ while treating molecules as rigid bodies. It combines an invariant EGNN-based molecular embedding with a Riemannian flow that evolves $(L,F,R)$ on their respective manifolds, using $u_t$ velocity fields and $ ext{SO}(3)$ geodesics to generate full crystal structures. Benchmarking on Thurlemann and OMC25-MCF demonstrates state-of-the-art performance among flow-based methods and competitive results relative to rule-based approaches, with end-to-end integration with u-MLIP and DFT enabling rapid polymorph screening. The work highlights a scalable, geometry-aware path toward data-driven discovery of molecular crystals, while outlining future improvements like energy-aware training, torsional flexibility, and space-group-constrained manifolds for further gains. $$L\in\mathbb{R}^{3\times3},\ F\in[0,1)^{3},\ R\in SO(3)$$ are jointly modeled through learned velocity fields to produce physically plausible packings that respect periodicity and symmetry. $$

Abstract

Molecular crystal structure prediction represents a grand challenge in computational chemistry due to large sizes of constituent molecules and complex intra- and intermolecular interactions. While generative modeling has revolutionized structure discovery for molecules, inorganic solids, and metal-organic frameworks, extending such approaches to fully periodic molecular crystals is still elusive. Here, we present MolCrystalFlow, a flow-based generative model for molecular crystal structure prediction. The framework disentangles intramolecular complexity from intermolecular packing by embedding molecules as rigid bodies and jointly learning the lattice matrix, molecular orientations, and centroid positions. Centroids and orientations are represented on their native Riemannian manifolds, allowing geodesic flow construction and graph neural network operations that respects geometric symmetries. We benchmark our model against state-of-the-art generative models for large-size periodic crystals and rule-based structure generation methods on two open-source molecular crystal datasets. We demonstrate an integration of MolCrystalFlow model with universal machine learning potential to accelerate molecular crystal structure prediction, paving the way for data-driven generative discovery of molecular crystals.

Figures (4)

• Figure 1: Overview of the MolCrystalFlow framework. a, Schematic view of polymorph energy landscape in molecular crystals.b, Representation of a molecular crystal structure. A molecular crystal is decomposed to a lattice matrix spanning by three vectors $\overrightarrow{l_1}$, $\overrightarrow{l_2}$ and $\overrightarrow{l_3}$, and a constituent molecular conformer building block which can be characterized by a rigid-body centroid and rotational orientation given by three eigenvectors derived from a principal component analysis over the atomic positions of the building block. Hence a coarse-grained crystal includes all the modalities to be generated. c, MolCrystalFlow works by encoding each building block in the crystal (for example two molecules in the schematic) with an equivariant graph neural network (EGNN) that produces an invariant embedding. The invariant embedding is augmented with auxiliary features to form the input for the joint flow-matching generative process that simultaneously produce the lattice matrix using a linear interpolant for each lattice matrix entry, centroid fractional coordinates on a torus geodesic, and rotational orientation on an $SO(3)$ geodesic. Together, the three learned modalities will reconstruct the full atomic structure of the molecular crystal. d. An example generation trajectory with four time steps, showing the co-evolution of centroid positions, molecular orientations and lattice.
• Figure 2: Neural network architecture of MolCrystalFlow. a, Molecular building block (BB) embedding network with EGNN. An atomic node embedding is constructed by the atom interacting with neighbors in the local chemical environment through multiple graph convolution layers. The final building block embedding is a weighted atomic embedding where each weight is learned and comes from the node embedding interacting with invariant atomic distance to the centroid. b, MolCrystalNet parameterizes the joint flow matching for lattice matrix $L$, centroid fractional coordinates ($F$), and rotational orientation ($R$). Given interpolants at time $t$, a concatenation of BB embedding, auxiliary features and time embedding, followed by a multilayer perceptron (MLP), gives rise to the initial node embedding for each molecular building block. In the update block, the initial embedding interacts with a $\chi$ embedding for each building block. Message passing relies on the relative positions and orientations between building blocks, and current lattice matrix at time $t$. The final output is the velocity field for centroid fractional coordinates ($u_t(F_t)$), the denoised rotational orientation ($R_1$) and lattice matrix ($L_1$). Interactions of lattice with fractional coordinates $(F,L)$ and orientations $(R,L)$ are also included in the message passing. c, For rotation-dedicated message passing ($m_R$), we use the axis-angle representation of relative rotational orientation between two building blocks. We embed the rotational angle $\omega$ and azimuthal angle $\rho$ using Fourier-series like trigonometric functions while using an MLP to create the embedding for the inclination angle $\kappa$. d, Concatenation-summation operation to strengthen the signal of each building block's $\chi$. e, $\chi$-grouped optimal transport to reduce cross-links between paths of different building blocks and facilitate inference.
• Figure 3: Performance of MolCrystalFlow against MOFFlow and Genarris-3 baselines. a, Comparison of MolCrystalFlow with MOFFlow and Genarris-3 over 10-sample matching rates evaluated for directly generated structures without any optimization under different levels of site tolerances. b, Comparison of MolCrystalFlow with Genarris-3 over 10-sample matching rates evaluated for generated structures after a rigid-body optimization under different levels of site tolerances. Regarding the box plots, the median is shown as a solid line. The edges of the box correspond to the first and third quartiles, and the whiskers extend to values within 1.5 × interquartile ranges. c--e, Lattice volume deviation of generated structures for MolCrystalFlow (c), MOFFlow (d), and Genarris-3 with rigid-press optimization (e). A kernel density estimation is used to show data distributions for 10 samples. Relative mean absolute deviation (RMAD) is reported as mean and standard deviation across 10 samples. d, Grid search for the optimal velocity annealing scalings for centroid ($s_{u_F}$) and rotation ($s_{u_R}$) using 10 samples and stol=0.8. For figure a and b, five independent runs for 10 samples are generated the statistics and the results for each run are overlaid on the box plots.
• Figure 4: Molecular crystal structure prediction by integrating MolCrystalFlow generator, u-MLIP and DFT stability ranking. a, Four-step pipeline to integrate MolCrystalFlow for structure generation, u-MLIP for structural optimization and obtaining coarse energy landscape, and DFT for final stability ranking. b, Molecular graphs and conformer structures of three open targets from the 3$^{\mathrm{rd}}$ CCDC CSP competition. c--e, Energy ranking results evaluated with u-MLIP, PBE-D3 and PBE-MBD for three generated and relaxed polymorphs whose PBE-MBD energies are the lowest, combined with the experimental structure evaluated at the same level of theory. Figure c, d and e correspond to target VIII, X and XI, respectively. f, Packing similarity analysis of the PBE-MBD lowest-energy polymorph relative to the experimental structure was performed using the COMPACK algorithm as implemented in the CSD Python API.