Data-Driven Approach to Encoding and Decoding 3-D Crystal Structures

Generative models have achieved impressive results in many domains including image and text generation . In the natural sciences , generative models have lead to rapid progress in automated drug discovery . Many of the current methods focus on either 1-D or 2-D representations of typically small , drug-like molecules . However , many molecules require 3-D descriptors and exceed the chemical complexity of commonly used datasets . We present a method to encode and decode the position of atoms in 3-D molecules along with a dataset of nearly 50,000 stable crystal unit cells that vary from containing 1 to over 100 atoms . We construct a smooth and continuous 3-D density representation of each crystal based on the positions of different atoms . Two different neural networks were trained on a dataset of over 120,000 three-dimensional samples of single and repeating crystal structures . The first , an Encoder-Decoder pair , constructs a compressed latent space representation of each molecule and then decodes this description into an accurate reconstruction of the input . The second network segments the resulting output into atoms and assigns each atom an atomic number . By generating compressed , continuous latent spaces representations of molecules we are able to decode random samples , interpolate between two molecules , and alter known molecules . 1 INTRODUCTION . Generative models have recently seen tremendous success in generating 2-D images of every day objects ( Kingma & Welling , 2013 ; Goodfellow et al. , 2014 ; Brock et al. , 2018 ; Razavi et al. , 2019 ) . The size and accuracy of generated results has greatly improved to the point where samples from the latent space decode to photo-realistic samples ( Brock et al. , 2018 ; Razavi et al. , 2019 ) . A very exciting and important future avenue for generative models is the generation of 3-D structures , like in the world around us . Adversarial networks and autoencoders have been extended into 3-D and have shown they are able to encode useful representations of everyday objects ( Wu et al. , 2016 ; Zhu et al. , 2018 ; Brock et al. , 2016 ; Achlioptas et al. , 2017 ; Valsesia et al. , 2018 ; Meng et al. , 2019 ; Maturana & Scherer , 2015 ) . Representations using point clouds have gained popularity as a way to capture the distribution of different types of objects ( Achlioptas et al. , 2017 ; Yang et al. , 2019 ; Wu et al. , 2018 ) . As machine learning approaches are able to understand and recreate the underlying distributions for many different types of objects in our world , the application of these tools in the physical sciences is a very exciting direction . A clear field where generative models can have a tremendous impact is in material discovery , which is very important for designing new batteries or carbon capture devices for fighting climate change . One very successful area of applying generative models to the sciences has been the field of drug discovery ( Jin et al. , 2018 ; Gómez-Bombarelli et al. , 2018 ; Assouel et al. , 2018 ) ( see Schwalbe-Koda & Gómez-Bombarelli ( 2019 ) for an excellent summary of many methods ) . The success of machine learning in this domain has been enormously helpful as a process that was often painstakingly slow has been rapidly accelerated in searching an unimaginably large space of possible drug compounds , estimate to be up to 1060 ( Polishchuk et al. , 2013 ) . Additionally , excellent datasets such as QM9 ( Ramakrishnan et al. , 2014 ) and ZINC ( Irwin et al. , 2012 ) which contain molecules of interest along with their precomputed properties have allowed for comparing different methods and developing state-of-the-art tools . By learning a compressed , useful , latent space representation , this enormous space of molecules can be embedded in a simpler latent space that is easier to search Kusner et al . ( 2017 ) . Using generative models , compounds hypothesized to have specific properties can be rapidly generated and then a targeted subset of these can be experimentally tested . For example , in Gómez-Bombarelli et al . ( 2018 ) , an auxiliary property prediction task is introduced for a network separate from the encoder/decoder . This allows for optimization of properties in the latent space and then the resulting latent space vector is decoded into a candidate molecule that can undergo more rigorous computational testing before it is experimentally synthesized . Using a graph representation and reinforcement learning , You et al propose a novel method for molecule generation with specific properties You et al . ( 2018 ) . Most of the work combining generative models with chemical discovery has focused on molecules that can be represented in either 1 dimension ( such as a SMILES string ( Weininger , 1988 ) ) or by leveraging a 2 dimensional representation of a molecule , such as a graph . While 1-D and 2-D representations have been very successful for many drug compounds , these representations are not sufficient to describe all molecules . For example , it is possible for different molecules , with very different properties , to have identical graphs ( Gebauer et al. , 2019 ) . Additionally , there are more complex classes of compounds where the 3-D structure is integral to the molecule ’ s properties . Therefore , in this work we focus on generating representations of molecules where the 3-D information is essential for reconstruction . Generating 3-D structures is still a relatively nascent field , when compared to image and text generation . The data requirements for modeling 3-D structures are larger than their 2-D counterparts and there are fewer standard datasets ( Nguyen-Phuoc et al. , 2019 ) . Additionally , to measure the properties of inorganic materials ( like those considered here ) , costly DFT calculations are needed which require precise 3-D locations . There is an extra dimension in 3-D problems , providing additional symmetries that often need to be learned ( Weiler et al. , 2018 ) . New work from generates adjacency matrices which are , by construction , invariant to rotation Hoffmann & Noé ( 2019 ) . In either case , for the automatic generation of viable complex 3-D structures of inorganic compounds one needs to accurately generate the locations of many atoms in 3-D space . In this work , we focus specifically on proposing an effective 3-D representation for this class of molecules that can not be represented in the same ways as many smaller molecules . We hope to promote work in this still nascent field which has yet to see the same level of success as the generation of smaller molecules . In the Appendix , we have provided a clear comparison with 1-D and 2-D representations . Currently , no direct comparisons are possible . In addition to drug discovery , a promising avenue for molecular design is for minimizing environmental impact through the design of more efficient or more environmentally friendly compounds for a variety of applications . Designing more efficient materials for photo-voltaic panels and more environmentally friendly materials for batteries are both very exciting research directions for using machine learning as a tool towards mitigating climate change ( Niu et al. , 2015 ; Tabor et al. , 2018 ; Gebauer et al. , 2019 ; Rolnick et al. , 2019 ) . Many of the compounds of interest are crystal structures made of single small sub-units that repeat in all directions . This sub-unit is referred to as a “ unit cell ” and can vary in size and shape as well as internal arrangement and chemical composition . In this work , we present a method for encoding and decoding a very important set of 3-D molecules that can not be represented using standard methods used for drug-like molecules . We directly encode and decode 3-D volumes of density from two different data representations . In the first , we use single unit cells , as shown in Fig . 1 . Each unit cell is centered in the cube , but randomly rotated . In the second representation , we repeat the unit cell along all three axes such that the resulting sampled cube contains repeated unit cells of crystal structures , like those shown in Fig . 2 ( and Fig . S6C ) . Many very interesting and important molecules , such as materials for solar panels or batteries , are composed of crystalline units . We train a variational autoencoder ( VAE ) and a network to segment the decoded output based on the true locations of atoms in tandem . By coupling these two tasks , we are able to accurately encode and decode 3-D atomic positions and species . As far as we are aware , this is not an explored direction as a means to represent molecules for encoding and decoding their 3-D arrangement . We summarize our contribution as follows : • We propose a method for encoding and decoding 3-D molecules that can not be represented by traditional string or graph based approaches . By jointly training a VAE and a segmentation network on the output of the decoder , we train the entire network in an end-to-end fashion . • We consider two different problems : ( 1 ) encoding and decoding single unit cells and ( 2 ) encoding and decoding repeated unit cells . In both cases we accurately reconstruct the locations of atoms . We also achieve good results in atomic species identification and future work will improve this front . By sampling from z ∼ N ( 0 , 1 ) , we generate complex structures that have physically realistic spacing between atoms . 2 PRELIMINARIES AND RELATED WORK . There are many classes of compounds that are of interest for data-driven discovery . There has been significant work exploring the generation of molecules with potential medicinal properties . These molecules tend to be organic molecules and can be represented efficiently using a string or a graph . A more complex class of molecules are inorganic crystal structures which vary in complexity across many different axes . Crystals are materials that are made up of a repeating pattern of a simpler “ unit cell. ” Crystal structures are of key interest for many environmental problems , such as materials for solar panels and batteries ( Niu et al. , 2015 ) . In contrast to problems in drug discovery , many crystals are not composed solely of organic molecules . Crystals commonly contain many heavy ( non-Hydrogen ) atoms which make various quantum mechanical calculations of their energetic properties very difficult and time consuming . For example , the common benchmark dataset QM9 for predicting quantum mechanical properties includes over 130,000 molecules contains fewer than 9 “ heavy ” atoms ( Ramakrishnan et al. , 2014 ) . There have been two parallel lines of work applying machine learning to material sciences . The first deals with the prediction of physical properties from a compound without performing a computationally expensive density functional theory ( DFT ) calculation . Recently , Xie and Grossman introduced crystal graph convolutional networks ( CGCNN ) for accurate prediction of 8 different DFT-calculated properties of crystal structures ( Xie & Grossman , 2018 ) . MatErials Graph Network ( MEGNet ) uses graph neural networks ( Chen et al. , 2019 ) to predict a variety of energetic properties on both the QM9 dataset ( Ramakrishnan et al. , 2014 ) as well 69,000 crystal structures from the Materials Project ( Jain et al. , 2013 ) . Another network , SchNet , uses 3-D spatial information to directly leverage the interactions between separate atoms . They then use a continuous-filter convolution for accurate property prediction ( Schütt et al. , 2018 ) . These methods , and many more , have all achieved impressive results on many standard benchmark datasets . Cubuk , et al . use transfer learning to search an enormous space of molecules for promising Lithium ion conductors ( Cubuk et al. , 2019 ) . Another line of work is concerned with generating molecules , often molecules that have specific properties . In the past decade , generative models for drug discovery have achieved impressive results . Typically a representation of a compound , such as a SMILES string ( Gómez-Bombarelli et al. , 2018 ; Segler et al. , 2017 ; Alperstein et al. , 2019 ) or a graph ( Jin et al. , 2018 ; Assouel et al. , 2018 ; Mansimov et al. , 2019 ) is encoded and decoded . Then , by sampling the resulting latent space , novel molecules can be generated . Either by performing an auxiliary task with the latent space ( as done in Gómez-Bombarelli et al . ( 2018 ) ) and then performing optimization or by conditioning the latent space on desirable properties , novel chemical structures are obtained . One possible representation of drug-like molecules is a SMILES string . However , a difficulty faced using the SMILES representation is ensuring that the decoder decoded a valid SMILES string ( Gómez-Bombarelli et al. , 2018 ; Kusner et al. , 2017 ) . Additionally , there is not a strong notion of distance between molecules and their SMILES representation ( Jin et al. , 2018 ) . One approach , ChemTS , uses SMILES strings along with recurrent neural networks and Monte Carlo Tree Search to better ensure generated strings decode to real molecules ( Yang et al. , 2017 ) . By explicitly penalizing invalid SMILES strings in their reward function they are able to bypass the difficulty in their decoded molecules being non-physical . More recently , graph based methods have proven very successful in generating synthetically attainable molecules . By building a database of molecule fragments ( like LEGO bricks ) , a graph is formed based on the connectivity of different structures ( Jin et al. , 2018 ) . These methods have far fewer difficulties ensuring that the decoder produces physically valid molecules . For generation of physical molecules , this is a very important consideration , as there are many hard constrains that most be obeyed . Other graph based approaches include flow based models ( Madhawa et al. , 2019 ) . Also using graph neural networks , E. Manismov and collaborators developed a method to , given a graph , recreate conformations of a 3-D molecule and predict its energetic properties ( Mansimov et al. , 2019 ) . In Mansimov et al . ( 2019 ) , the authors propose a graph-based generative method that is able to generate molecules with desired target properties . Impressively , they recover 3-D positional information for the atoms and show they achieve the required accuracy for relaxation . An additional challenge for generative models in the physical sciences is ensuring that samples from the latent space , z , decode into physically plausible objects . The decoded objects need to obey physical constraints and an object that may appear physical is not certain to actually be experimentally realizable . In many domains this is not of specific concern , for example a generated animal is scored on some likelihood based on its visual appearance , not whether or not the generated creature is genetically possible . Recently , there have been generative models leveraging the 3-D nature of the problem . G-SchNet , a generative model for 3-D molecules ( Gebauer et al. , 2019 ) leverages the SchNet architecture- a start-of-the-art property prediction network ( Schütt et al. , 2018 ) . They show they are able to generate molecules , placing atoms in 3-D space in a rotationally invariant manner . By appropriately conditioning the placement of new molecules based on the location and identity of previous molecules , they ensure the symmetries that are required for the molecules to relax in a DFT calculation . They propose expanding G-SchNet to generating crystal structures as a future direction . While there is a rich literature in the generation of organic compounds , recently there has been work in the generation of more complex crystal structures , such as CrystalGAN ( Nouira et al. , 2018 ) . The authors of CrystalGAN introduce geometric constraints and show that their method is able to produce stable structures compared to methods that do not include such domain knowledge such as DiscoGAN ( Kim et al. , 2017 ) . CrystalGAN uses a dataset with a much larger selection of elements than much previous work , but they limit themselves to molecules of particular chemical structure , nor do they generate a 3-D arrangement . In this work , we do not limit the generated crystals in this way .

Summary: The authors describe a method to encode and decode the position of atoms in 3-D molecules. An encoder-decoder architecture is used to create a representation of a molecule and to reconstruct the molecule from its representation. Then a second Neural Network segments the output and assigns an atomic number. Prior work on this task has used 1D (SMILES) and 2D (Graph) representations. The authors argues that exploiting 3D structure can create better representations.