Rigid Body Flows for Sampling Molecular Crystal Structures
This work addresses the challenge of modeling distributions of many interacting molecules for applications in computational chemistry and materials science, representing an incremental step by combining flows on rigid body motions with existing methods for internal degrees of freedom.
The authors tackled the problem of sampling molecular crystal structures by introducing a new normalizing flow tailored for modeling positions and orientations of multiple objects in 3D space, achieving results such as training Boltzmann generators for a tetrahedral system and the ice XI phase in the TIP4P water model.
Normalizing flows (NF) are a class of powerful generative models that have gained popularity in recent years due to their ability to model complex distributions with high flexibility and expressiveness. In this work, we introduce a new type of normalizing flow that is tailored for modeling positions and orientations of multiple objects in three-dimensional space, such as molecules in a crystal. Our approach is based on two key ideas: first, we define smooth and expressive flows on the group of unit quaternions, which allows us to capture the continuous rotational motion of rigid bodies; second, we use the double cover property of unit quaternions to define a proper density on the rotation group. This ensures that our model can be trained using standard likelihood-based methods or variational inference with respect to a thermodynamic target density. We evaluate the method by training Boltzmann generators for two molecular examples, namely the multi-modal density of a tetrahedral system in an external field and the ice XI phase in the TIP4P water model. Our flows can be combined with flows operating on the internal degrees of freedom of molecules and constitute an important step towards the modeling of distributions of many interacting molecules.