ParticleGrid: A Library for 3D Molecular Representation for Deep Learning

Machine learning and especially deep learning have recently seen exponential growth in fields such as computer vision, natural language processing, and the physical sciences. Deep learning has been used to perform predictive and generative modeling for a wide range of scientific problems ranging from quantum physics, computational biology, astrophysics, and material sciences. The application of deep learning to generative tasks for novel materials provides a paradigm shift in the traditional discovery process. An ideal representation for deep learning-based generative workflows for molecules requires a structured representation that preserves the geometric structure and encodes physical constraints. A 3D representation provides an ideal input to neural networks and also preserves structural information compared to character-based representations such as SMILES. We present ParticleGrid, a fast, and reversible 3D molecular grid generation library designed to seamlessly attach to 3D generative workflows. Reversible grids also allow for retrieving discrete atomic information, and geometric optimizations using traditional methods. Our highly optimized implementation allows integration with deep learning frameworks such as PyTorch without adding significant computational overhead.