# ParticleGrid: Enabling Deep Learning using 3D Representation of Materials
††thanks: *Equal contribution

Shehtab Zaman  Binghamton University
Binghamton, New York, USA
<EMAIL_ADDRESS>Ethan Ferguson  Binghamton University
Binghamton, New York, USA
<EMAIL_ADDRESS>Cécile Pereira  TotalEnergies Marketing and Services
Paris, France
<EMAIL_ADDRESS>Denis Akhiyarov  TotalEnergies EP Research &
Technologies US
Houston, Texas, USA
<EMAIL_ADDRESS>Mauricio Araya-Polo  TotalEnergies EP
Research & Technologies US
Houston, Texas, USA
<EMAIL_ADDRESS>Kenneth Chiu  Binghamton University
Binghamton, New York, USA
<EMAIL_ADDRESS>
###### Abstract

From AlexNet to Inception, autoencoders to diffusion models, the development
of novel and powerful deep learning models and learning algorithms has
proceeded at breakneck speeds. In part, we believe that rapid iteration of
model architecture and learning techniques by a large community of researchers
over a common representation of the underlying entities has resulted in
transferable deep learning knowledge. As a result, model scale, accuracy,
fidelity, and compute performance have dramatically increased in computer
vision and natural language processing. On the other hand, the lack of a
common representation for chemical structure has hampered similar progress. To
enable transferable deep learning, we identify the need for a robust
3-dimensional representation of materials such as molecules and crystals. The
goal is to enable both materials property prediction and materials generation
with 3D structures. While computationally costly, such representations can
model a large set of chemical structures. We propose ParticleGrid, a SIMD-
optimized library for 3D structures, that is designed for deep learning
applications and to seamlessly integrate with deep learning frameworks. Our
highly optimized grid generation allows for generating grids on the fly on the
CPU, reducing storage and GPU compute and memory requirements. We show the
efficacy of 3D grids generated via ParticleGrid and accurately predict
molecular energy properties using a 3D convolutional neural network. Our model
is able to get 0.006 mean square error and nearly match the values calculated
using computationally costly density functional theory at a fraction of the
time.

###### Index Terms:

Deep Learning, Scientific Machine Learning, High-Performance Computing, 3D
Structures

##  I Introduction [1 pg]

Deep learning (DL) has had a transformative effect on disciplines from natural
language processing, computer vision to reinforcement learning. It is
expanding the way we perform science. Whether it is news articles generated by
GPT-3[1], paintings in the style of Van Gogh by StyleGAN [2, 3], or
conditional images generated by Dall-E [4], generative modeling is one of the
most striking uses of DL.

The aforementioned models achieve such amazing results by building on the work
of a wide community of researchers consistently innovating on DL algorithms.
For image classification, iteration over convolutional neural networks (CNNs)
resulted in rapid improvements in both predictive and computing performance
from AlexNet to ResNet and InceptionNet to EfficientNets. Similarly, for image
segmentation we’ve seen improvements from R-CNN, Faster R-CNN, to YOLO,
allowing for real-time image segmentation. We have seen similar trends in
image generation with Generative Adversarial Networks (GANs), Variational
Autoencoders (VAEs), Normalizing Flows, and Diffusion models. Due to the
standardization of input, we see a trend of rapid discovery of models, layers,
and learning techniques to facilitate performance improvements. Images have a
natural and well-established representation. As a result, we’ve seen
significant improvements in DL-based tasks in a relatively short amount of
time.

Yet, DL-based generative models are still at a nascent stage for scientific
principles. Applications in molecular and materials discovery have yet to be
utilized in wide applications. Generative models of chemical structures
especially have yet to experience similar growth as in image and sequence
generation. This can be largely traced to the difficulty in identifying a
common representation for chemical structures. Molecules and materials cover a
large set of possible combinations of atomic elements with varying structures,
and varying representations, each with its own strengths and drawbacks. Some
classes of materials such as molecules have textual representations such as
SMILES [5] and SELFIES [6]. The same molecules are also succinctly represented
using graphs and graph-based structures such as coulomb matrices [7, 8].
Materials such as crystals on the other hand do not share such
representations. The difficulty in the representation of these materials is
further compounded due to the spatial symmetry requirements and periodicity.
We can see the various representations of molecular structures in Fig. 1.

Figure 1: The many machine-readable views of the molecule $C_{9}H_{8}O_{4}$,
commonly known as Aspirin. (a) is the canonical SMILES retrieved from the
PubChem database. (b) is the INchI (International Chemical Identifier)
representation from the same database. (c) is the SELFIES string obtained from
Krenn et al. [6]. The graph-like representation in (d) is visualized using the
ball-and-stick model. Each atom corresponds to a vertex, and each bond
corresponds to an edge in the graph. (e) (f) is the visualization of a
ParticleGrid generated grid using Mayavi [9].

In recent years, due to larger GPU memory and faster compute capabilities,
three-dimensional grid-like representations have been explored. The importance
of three-dimensional structural information can be highlighted as large-scale
datasets such as the open graph benchmark - large-scale challenge (OGB-LSC)
[10] updated their datasets to include 3D coordinate information. This follows
the trend as [11] employed 3D conformer models to extend PubChem [12], an open
repository for small molecules designed to study biological activity. 3D grids
can encode both geometric and chemical information, while also being
generalizable to represent many different classes of materials. Including
geometric and chemical information to the input of deep learning models
provides inductive bias to improve learning results [13, 14]. Representations
that can be shared across many different classes enable the use of transfer
learning [15]. Significantly, this enables the use of deep learning,
traditionally requiring large datasets, to be used on classes of materials
with few samples by pre-training on larger datasets.

3D generative models as in [16, 17] show that there are promising DL models
that can be utilized to augment molecule and crystal generation workflows. The
ability to represent different classes of materials such as organic molecules
and crystals shows that grid-like representations may be promising as a
representation for DL tasks. As grids are significantly larger than more
concise objects such as SMILES, and graphs, the computation and storage costs
can be prohibitive. Therefore we propose, ParticleGrid, a SIMD-optimized grid
generation for molecules and crystals. Our goal is to provide on the fly grid
generation abilities from 3D coordinates and seamlessly integrate with data
ingestion pipelines of deep learning frameworks such as PyTorch and
Tensorflow. We also propose a general-purpose transformation from grids to 3D
coordinates.

We discuss the various strengths and shortcomings of text-based, graph-based,
and grid-based representations in Section II. In Section III, we present the
mathematical foundation of our grid representations. In Section IV, we show
the reversibility of the grid to discrete coordinates. Section V and Section
VI show the programming interface and detail the optimizations in
ParticleGrid, respectively. Section VII, we show an application of 3D
structures for molecular property prediction, and finally, Section VIII,
discusses the results of our findings.

We publicly release our implementation as an open-source library, and the
performance and application codes can be found for reproducibility at
www.github.com/ParticleGrid/ParticleGrid

##  II Related Work

Deep learning algorithms have been used in materials science applications to
both perform supervised prediction of material properties such as energy, bond
forces, and others, and also the unsupervised generation and reinforcement
learning of new molecules and crystals for drug discovery and delivery, carbon
capture and storage, and batteries among others. This work has used a variety
of different representations for materials. Broadly, we categorize them as
string-based, graph-based, and grid-based.

Gomez-Bombarelli et al. [18] use deep generative models to generate new SMILES
(Simplified Molecular Input Line Entry Specification) strings. Kusner et al.
[19] extend such generative algorithms by imposing syntactic and semantic
constraints on the decoder. While generating valid structures, the SMILES
representation is inherently unable to capture many of the physical
invariances present in molecules [20]. Furthermore, similar structures may
have markedly different SMILES representations. In order to mitigate the
structural representation limitations of SMILES, recently the SELFIES (SELF-
referencIng Embedded Strings) representation has been proposed [6]. Using
formal grammar, Krenn et al. derive a representation that is highly robust.
Even a random arrangement of the symbols defined in the grammar result in a
valid molecule. This is a significant characteristic, as generative models
benefit greatly from such a robust representation. While using string
representations is very attractive for generative modeling, they suffer from
some drawbacks. Other classes of materials such as periodic crystals also pose
difficulties for SMILES and SELFIES. Whilst future work can enable these
character-based representations for a wider variety of materials, currently,
these representations are limited to molecules [21]. Yao et al. developed
RFCode, a string representation of crystals based on SMILES and InChI for
reticular frameworks such as metal-organic frameworks (MOFs). Reticular
frameworks are nanoporous crystals with many constituent atoms and use in
diverse fields such as gas storage, separation, and catalysis [22]. The
authors designed a data-driven inverse design framework to propose new MOFs
for optimized methane storage using a VAE.

Graph-based representations of molecules have also been used for molecular
generation tasks. Jin et al. [23] use graph-structured encoder and junction
tree-based decoder to generate molecular structures. While graph-structured
generative models allow for greater representation of long-range atomic
interactions and provide geometric information of the molecule, graph-
structured generative models have been restricted to a small and limited set
of molecules [24]. On the other hand. MPNN, CGCNN, SchNet, and GemNet, achieve
state-of-the-art accuracy in molecular property prediction, and molecular
dynamics are achieved by graph-based neural networks [25, 26, 27, 28]. Yet,
these models are not suitable for use in unsupervised training and generative
modeling. Graph-based generators are required to sample discrete structures,
i.e graphs, from a continuous latent space. This results in poor performance
when the discrete structure is too large [24].

Grid-based modeling of materials can provide both geometric information and
generalizability. Hoffman et al. explore the possibility of performing
dimensionality reduction using deep learning models to encode and decode 3D
crystal structures [29]. Ragoza et al. [17] use grids with Gaussian-like
density functions for voxels to learn continuous representations of 3D
molecular conformations. Our formalism builds on this by allowing for grids
with variable spatial resolutions. Each grid can represent a different volume
of space depending on the size of the material. This enables the model to
learn multi-resolution representations of molecules. This allows us to
represent different classes of materials that may have different numbers of
atoms required to scale. As we integrate the atomic probability densities over
voxel regions, we are also able to encode global molecule level features such
as the total number of atoms, and the number of atoms per element.
Furthermore, the discretization step described by Ragoza et al. utilizes a
solver that uses gradient descent for atomic structure detection. They
generate multiple possible structures per grid. They subsequently use beam
search to find the best possible 3D structure. Similarly, Kim et al. use a
3D-GAN to generate crystal structures called, Zeolites [16].

A significant issue with using molecular grids is the compute and memory
requirements. Deep learning models require high data throughput, especially in
high-performance setting. Generating grids and storing them on disks is
infeasible, as datasets can be scaled up to millions and billions of samples.
As a result, the grids must be generated on the fly. Yet, high-performance
vectorized libraries such as Numpy [30] cannot provide sufficient throughput.
Therefore a more low-level implementation is required, such as ParticleGrid.
We provide direct comparisons with the grid generation times in Section VI.

##  III Grid Representation of Molecules

We define three objectives for a suitable grid representation for deep
learning (DL) for material science:

  1. 1.

Physically inspired

  2. 2.

Generalizable

  3. 3.

Invertible

A suitable grid needs to provide physically relevant information not only to
provide suitable inductive bias [14, 13], but also to provide important
information to practitioners. Inductive biases can be physical laws,
assumptions, or human knowledge that can be inscribed into the data, model,
and loss function to prioritize a specific solution and improve learning
performance. Visually inspecting the grid should provide significant
information to a scientist. The generalizability criterion is significant for
using the grid for DL. As we have seen before, DL models using the same
representation allow multiple different types of models to be used in
conjunction. Also significantly, re-using learned models in transfer learning
is imperative in low-data settings. Data availability in scientific research
can be bimodal. There exist large datasets with millions and billions of
samples as is the case for small organic and drug-like molecules. On the other
hand novel, exotic materials may only have a hundred or thousand known
samples. Therefore, a general representation that can cover many subdomains of
material science is important. Finally, the invertibility of the
representation is the ability to transform the structure to other canonical
representations. First principle methods have long-established algorithms to
quantify and classify materials. DL methods, especially generative models, are
designed to augment the established scientific discovery pipeline. Therefore
the data representation requires interoperability between representations
suited for DL and representations required by other computational and
experimental methods.

###  III-A Voxel-filling Function

With the aforementioned criterion in mind, we propose a grid representation
discretizing the real space realization of an atom in to grids.

A simple realization of our approach is a grid representing a single molecule
is a 4-dimensional tensor of the form (Channel, H, W, D). Each channel of the
data cube represents the density of the element at the point. Generally, we
use a cube grid, there for H=W=D, but not required. Each voxel represents
$\left(\frac{1}{N}\right)^{3}(\text{ABC})$ volume of 3D space. $N$ is the
number of voxels, usually $N=H\times W\times D$. Here, A, B, and C represent
the height, width, and depth of a bounding box of the molecule. The user has
control over choosing the orientation and size of the bounding box. It is
important to note that the voxel need not be cubic, and the grids can have
different resolutions between molecules.

For a molecule at point $\vec{\mu}=(\mu_{x},\mu_{y},\mu_{z})$, the function
defining the density of the atom existing at point $\vec{p}=(x,y,z)$ is given
by $f_{\vec{\mu}}(\vec{p})$, the 3D Gaussian probability distribution
function.

$f_{\vec{\mu}}(\vec{p})=\frac{1}{\sigma^{3}\left(2\pi\right)^{\frac{3}{2}}}e^{-\frac{d(\vec{(}\mu),\vec{p})}{2\sigma^{2}}}$
(1)

Where, $d(\vec{(}\mu),\vec{p})$, is the element-wise distance of the functions
of the inputs. For a non-periodic molecule, we use the squared euclidean
distance. For the case of periodic crystals, the nearest image distance may be
used. The nearest image distance of an atom on a periodic lattice is
$\text{min}(d(\mu^{i}_{x},x_{0}))$ where, $\mu^{i}_{x}\in T(\mu_{x})$, where
$T(\mu_{x})$ is the set of coordinates of the lattice vector translations.
$\sigma$ is the variance or the ”spread” of the Gaussian.

Gaussians are a natural choice for a function to use for spatial
representation molecules. Furthermore, since we used a correctly normalized
the sum of the grid values for all the channels should equal the number of
atoms, and the sum over each channel should equal the number of atoms of that
type.

For a given atom at position $(\mu_{x},\mu_{y},\mu_{z})$, on channel q, the
value of the the region
$\mathcal{D}^{q}_{i,j,k}=[(x_{i},y_{j},z_{k}),(x_{i}+\delta_{x},y_{j}+\delta_{y},z_{k}+\delta_{z})]$
is defined by Eq. 2. Here, the $\mathcal{D}^{q}_{i,j,k}$ is the value of the
$i,j,k$-th voxel. $x_{i},y_{j},$ and $z_{k}$ are the left endpoints of the
real space locations of the $i,j,k$-th voxel. $\delta_{x},\delta_{y},$ and
$\delta_{z}$ are the size of side of the voxels at their respective
dimensions.

$\displaystyle F_{\vec{\mu}}(\vec{p})$
$\displaystyle=\left(\frac{1}{2}\right)^{3}G_{\mu_{x}}(x_{a})H_{\mu_{y}}(y_{b})L_{\mu_{z}}(z_{c})$
(2) $\displaystyle G_{\mu_{x}}(x)$
$\displaystyle=\mathrm{erf}\left(\frac{\sqrt{2}}{2\sigma}\left(\mu_{i}-x\right)\right)\biggr{\rvert}^{x_{i}+\delta_{x}}_{x_{i}}$
(3) $\displaystyle H_{\mu_{y}}(y)$
$\displaystyle=\mathrm{erf}\left(\frac{\sqrt{2}}{2\sigma}\left(\mu_{j}-y\right)\right)\biggr{\rvert}^{y_{j}+\delta_{y}}_{y_{j}}$
(4) $\displaystyle L_{\mu_{z}}(z)$
$\displaystyle=\mathrm{erf}\left(\frac{\sqrt{2}}{2\sigma}\left(\mu_{k}-z\right)\right)\biggr{\rvert}^{z_{k}+\delta_{z}}_{z_{k}}$
(5)

Since the Gaussian in Eq. 1 is separable, we can separate the integral
evaluated in Eq. 2 into three functions Eq 3, 4, and 5.

Given a molecule as a coordinate set, $\mathcal{M}$, such that each molecule
at coordinate $\vec{\mu}=(\mu_{x},\mu_{j},\mu_{k})$ is in set $\mathcal{M}$.
The grid is generated by evaluating $F_{\vec{\mu}}(\vec{p})$ over all
$(H\times W\times D)$ regions for all molecules
$(\mu_{i},\mu_{j},\mu_{k})\in\mathcal{M}$.

We can also see that due to this formulation, we have the condition,

$\sum^{\text{channel}}_{q}\sum^{N}_{i,j,k}\mathcal{D}^{q}_{i,j,k}=\text{Number
of Atoms}$ (6)

It is important to note the voxel-filling process does not require a fixed
bounding and therefore resolution. Similar to images, the same molecule can
have multiple grid representations with differing resolutions. This allows us
to use the same sized grid for molecules with varying volumes and numbers of
elements.

Equation 2. can be extended to augmented to produce weighted grids. Also,
other grid-like structures may be used in conjunction with the probability
grids such as the potential energy grids derived from the Lennard-Jones
potentials used in [16].

##  IV Maximum Likelihood Estimation of Atomic Coordinates

Figure 2: Reversing the gridding process to obtain molecular coordinates. We
use a 2D example for ease of visualization.

In order for the grids to be useful for generative modeling, the grids must be
invertible to retrieve atomic coordinates. Usually, generative models work on
grid space and have grids as inputs and outputs. It is desirable to obtain the
coordinates from the grids for further processing with traditional
cheminformatics and other ML methods. This allows ParticleGrid to be a middle
layer that is interoperable with string and graph-based representations and
can work in conjunction with the current state of the practice.

###  IV-A Algorithm for grid reversal

The deduction of coordinates follow a two- stage process, as seen in Fig. 2.

For a given point $\mathbf{P}_{0}$ with coordinates $(x_{0},y_{0},z_{0})$
generating grid $\mathbf{G}$, we wish find $\Theta$, with coordinates
$(x^{\prime}_{0},y^{\prime}_{0},z^{\prime}_{0})$, such that
$\Theta\approx\mathbf{P}_{0}$. At any given time, only the original grid
$\mathbf{G}$ is available.

  1. 1.

Approximate $\Theta$ from $\mathbf{G}$

  2. 2.

Generated approximated grid $\mathbf{G}^{\prime}$ using Eq 2 with $\Theta$ as
the input

  3. 3.

Calculate the difference, $\mathcal{L}$, using mean squared error between
$\mathbf{G}$ and $\mathbf{G}^{\prime}$

  4. 4.

Update
$\Theta\leftarrow\Theta-\lambda\frac{\partial\mathcal{L}}{\partial\Theta}$,
where $\lambda$ is the stochastic gradient descent step size (learning rate)

  5. 5.

Repeat steps 2-4 until $\mathcal{L}<\epsilon$, for some small value of
$\epsilon\in\mathbf{R}$

###  IV-B Optimal initialization of approximate coordinates

As a result of our grid structure, atomic centers display a clearly defined
local maxima nearest to the atomic centers. Reversing the grid to obtain
possible coordinates of generated grids constitutes a maximum likelihood
estimation of the local maxima across the grid. This requires mode detection
on the 3D grids. While there are many possible ways to estimate peaks on a 3D
volume, we explore the use of persistent homology [31]. Persistent homology is
a topological data analysis (TDA) tool to extract information from simplicial
complexes. Defining the 3D grids as 3D cubical complex, we know the
topologically significant structures will be atomic centers. We use the
efficient CubicalRipser library [32] to compute 2-cycles on our 3D grids by
computing the cohomology by matrix reduction. For further details on the
algorithm refer to [33].

###  IV-C Optimizing the coordinates

After obtaining the approximate coordinate, $\Theta$, we can use stochastic
gradient descent to optimize the approximated coordinates to match the
original grid.

Therefore, we need to be able to calculate the gradient update value
$\frac{\partial\mathcal{L}}{\partial\Theta}$ efficiently to update the our
initial set of coordinates $\Theta$.

We know the loss, given positions $\Theta$ or $(x,y,z)$ is

$\displaystyle\mathcal{L}$
$\displaystyle=||\mathbf{G}-\mathbf{G}^{\prime}||^{2}_{2}$ (7)
$\displaystyle\mathcal{L}$
$\displaystyle=\sum_{i,j,k}||\mathbf{G}-F_{\Theta}(x_{i},y_{j},z_{k})||^{2}_{2}$
(8)

Here, $\mathbf{G}^{\prime}$ is the grid generated using the approximated For
simplicity of notation, we use $F_{ijk}$ to refer to the function
$F_{\Theta}(x_{i},y_{j},z_{k})$ from here on.

Updates to $\Theta$ can be separately achieved for each coordinate so we have
the set of equations,

$\displaystyle x^{t+1}_{0}$ $\displaystyle\leftarrow
x^{t}_{0}-\lambda\frac{\partial\mathcal{L}}{\partial x^{t}_{0}}$ (9)

So the partial derivative with respect to the coordinate $x_{0}$ is given by

$\displaystyle\frac{\partial\mathcal{L}}{\partial x_{0}}$
$\displaystyle=\sum_{i,j,k}\frac{\partial\mathcal{L}}{\partial
F_{ijk}}\frac{\partial F_{ijk}}{\partial x_{0}}$ (10)
$\displaystyle=\sum_{i,j,k}2\left(F_{ijk}-G_{ijk}\right)\frac{\partial
F_{ijk}}{\partial x_{0}}$ (11)

At this point, we can use the Fundamental theorem of calculus to evaluate
$\frac{\partial F_{ijk}}{\partial x}$.

It is important to note that Eq. 1 and Eq. 2 are separable in x, y, and z
dimensions. We know that for some separable multivariate function

$\displaystyle F$ $\displaystyle=G(x)H(y)P(z)$ (12)
$\displaystyle\frac{\partial F}{\partial x}$
$\displaystyle=H(y)P(z)\frac{\partial G}{\partial x}$ (13)

and $F=\int fdx$ then $\frac{dF}{dx}=f$. So, the partial derivative term for
the x coordinate is,

$\frac{\partial F_{ijk}}{\partial
x_{0}}=\left(\frac{1}{8}H_{j}(y_{0})L_{k}(z_{0})\right)\frac{\partial
G_{i}(x_{0})}{\partial x_{0}}$ (14)

Since we know that,

$\frac{d}{dt}\mathrm{erf}(t)=\frac{2}{\sqrt{\pi}}e^{t^{2}}$ (15)

We have,

$\frac{\partial G_{i}(x_{0})}{\partial
x_{0}}=\frac{\sqrt{2}}{\sigma\sqrt{\pi}}\left(e^{-\frac{(i-x_{0})^{2}}{2\sigma^{2}}}-e^{-\frac{(i+\delta_{x}-x_{0})^{2}}{2\sigma^{2}}}\right)$
(16)

The other derivatives follow the same pattern.

##  V ParticleGrid’s Programming Interface

ParticleGrid, is designed to enable SciML111SciML: Scientific Machine Learning
practitioners to quickly and easily generate grids for molecular structures.
In order to minimize the friction of usage, a high-level interface developed
in Python binds to a C++ backend. The entire library can be installed using
the Python package manager Pip, with command:

⬇

1pip install git+https://github.com/ParticleGrid/ParticleGrid.git

Listing 1: Installation instructions

We ensure low installation overhead by only requiring a C++17 enabled compiler
and the Numpy [30]package.

One of the major considerations for ParticleGrid is to smoothly integrate with
deep learning frameworks such as PyTorch [34] and TensorFlow [35]. We designed
our implementation to work seamlessly with the data loading mechanism for
these frameworks, and therefore optimized for CPU performance, in order to
generate grids on background processes while training on accelerators. Also,
Numpy arrays holding grids are produced, in order to facilitate zero-copy
transfers to PyTorch and TensorFlow data containers. A simple grid generation
sample can be found in Listing 2.

⬇

1import numpy as np

2import torch

3import tensorflow as tf

4from ParticleGrid import coord_to_grid

5

6# Points are in the format (channel, x, y, z)

7test_points = np.array([0, 0.5, 0.5, 0.5],

8 [1, 0.0, 0.1, 0.2])

9# Generates a (2,32,32,32) grid

10grid = coord_to_grid(test_points,

11 width=1,

12 height=1,

13 depth=1,

14 num_channels=2,

15 grid_size=32,

16 variance=0.05)

17

18# Convert to PyTorch tensor

19grid_torch_tensor = torch.from_numpy(grid)

20# Convert to TensorFlow tensor

21grid_tf_tensor = tf.convert_to_tensor(grid)

Listing 2: Grid generation example

Similarly, we also provide an easy and extendable workflow to reverse the
grid. In Section V.B and in Fig. 2 we use the Cubical Ripser library in order
to use persistence homology to determine voxels with atoms.

##  VI Grid Generation Implementation

The following development and testing are performed on an Intel Xeon Silver
4114 CPU @ 2.20 Ghz, with 32k L1i and L1d cache, 1024K L2 cache, and 14080k L3
cache, 128 GB of system RAM. We used the GCC 7.5.0 compiler.

###  VI-A Baseline

Because the value at each grid point determined by Eq. 2. is separable into
$G$, $H$, and $L$, these values can be computed and stored for their
respective axes ($x$, $y$, and $z$) to later compute the value at a given grid
point. This already is a significant speedup over the naive approach of
calculating $F_{i,j,k}(x,y,z)$ at each grid point. Furthermore, because each
part of the separable function is the difference of a function calculated at
two points, and the minuend of one calculation is the subtrahend of the next,
the value of the erf can be calculated half as many times. These optimizations
provide much faster algorithmic performance over a naive baseline
implementation, but the performance still leaves much to be desired, and
exploiting the fact that much of the grid can be effectively ignored provides
fewer memory accesses and computations without degrading accuracy.

###  VI-B Exploiting Truncation

Since each of Eq. 3, 4, and 5 rapidly approach zero when calculated at
increasing distance from an atom. Because $F=G(x)H(y)L(z)$, if the values of
$G$, $H$, or $L$ are sufficiently small, they can be ignored. And, where these
functions are insignificant can be known before calculating $F(x,y,z)$ for
each point by simply searching through all $x$ for $G(x)$, $y$ for $H(y)$, and
$z$ for $L(z)$. The search space for these functions grows with $N$, whereas
the number of points to be calculated grows with $N^{3}$, where $N$ is the
size of the grid. So, it is very economical to add a limited computation ahead
of time to be able to skip many calculations of $F$. In fact, we found through
performance profiling that over 99% of the time taken to compute a grid was in
the multiplication of the three functions, so any optimizations there will be
greatly beneficial. Not only can many calculations of $F$ be skipped, but so
can writing all of the zero-valued points of $F$, because the tensor
representing the grid is initialized to zero.

How quickly the functions shrink depends on the variance, but the variance can
be tuned so that the falloff is over any desired distance. In our testing (See
VI-D), we found that the specific variance chosen is relatively
inconsequential so long as it is “reasonable”, as can be seen in 4. That is,
as long as individual atoms are readily distinguished, the speedup while
exploiting sparsity stays approximately the same. This is particularly useful
because with floating point numbers, the value of these functions rapidly
saturates to zero in practice. For example, using 32-bit floating point
numbers, a variance of 0.5, and a grid size of 64 on a group of 100 molecules
from the MOSES dataset [36] showed that out of the total $64^{3}=262144$ cells
in each grid, an average of only $2623$ cells were non-zero. That is $1\%$ of
all cells, allowing for significant skipping of calculation and memory writes.

As suggested before, saturation to zero is not necessary but it is convenient.
A threshold could be used instead of relying on floating point imprecision. In
this case, instead of seeing where $G$, $H$, and $P$ are zero before
calculating $F$, we instead see where they are below some threshold, and do
not bother doing the multiplication beyond that threshold. However, we did not
implement this because allowing floating point values to naturally saturate to
zero is very effective, and largely maintains precision. Of the various
optimization strategies used, this was the most significant. On average, we
see about a 7x speedup due to this strategy in our testing, using many
different grid sizes and variances. Specific speedup can be seen in Fig. 4.
The speedup varies quite significantly, particularly due to grid size. The
greatest speedup is seen with $32^{3}$ grids. This is likely because they have
a good mix of being able to fit in cache and benefit significantly from
sparsity. $16^{3}$ grids do not benefit as much from sparsity as calculating a
$16^{3}$ grid row requires accessing many cache lines representing a grid, a
cache line of 64 bytes holding 8 floats is already half of a full grid row.
And, non-sparse accesses all cache lines in order, which is easy for
processors to predict. On the other hand, the larger grids require significant
data movement. For example, a single $64^{3}$ grid is 1MB, thus performing
calculations on larger grids with many channels easily saturates the cache and
compute time is dominated by memory access.

However, there is still significant hardware speedup left on the table which
can be used to further enhance efficiency.

Figure 3: Speedup over baseline C++ implementation using Sparsity, AVX, and
both Sparsity + AVX.  Figure 4: Speedup over baseline C++ implementation using
Sparsity with different variances. The higher variance values result in
increased number of non-zero voxels.

###  VI-C Instructions-level parallelization of main functions

Grid generation is a prime candidate for advanced math instructions that
compute multiple values at the same time (SIMD). The functions $G$, $H$, $L$,
and $F$ are readily SIMDizable. We added SIMD intrinsics to ParticleGrid in
order to force the compiler to use AVX instructions. As can be seen when
looking at the speedup of AVX over baseline, and the speedup of AVX+Sparsity
over only Sparsity (See 4), the manual addition has a significant effect
beyond optimizations the compiler did automatically.

Unfortunately, there is no Gauss error function (erf) intrinsic on an 8-float
vector in any of Intel’s AVX iterations (as far as we know), meaning there is
no built-in instruction to compute erf efficiently. Fortunately, the Bürmann
series [37] can be used to approximate it with little error. We use the first
three terms in our approximation with $c_{1}=\frac{31}{200}$ and
$c_{2}=\frac{341}{8000}$ This is also used in our baseline implementation.

Other than this, it is a simple case of calculating the bounds for $G$, $H$,
and $L$ in a SIMD form, and doing the same for the $i$ dimension of $F$.

###  VI-D Performance Evaluation

Figure 5: Throughput using NumPy, Numba, and ParticleGrid for calculating
molecules, in Molecules/s

In order to test performance, we use the MOSES dataset and conformed the
samples to get a set of atom positions using RDKit. Then, we took a random
sample of 100 of the resulting molecules and conduct experiments with our
NumPy, Numba [38], Base ParticleGrid, and Optimized ParticleGrid
implementations on each molecule in the set, repeating the calculations 10
times for each variance we test, keeping track of the total time to do the
calculations. We tested the variances of 0.05, 0.1, 0.25, and 0.5, since they
seem to represent visually acceptable results when inspecting ParticleGrid
outputs. This results in each grid generating function being run 4,000 times
for each grid size. Then we calculate the average time to compute a molecule
and its inverse, the throughput. The following is an analysis of the times we
observed from these runs.

The performance impact of hardware optimizations is less than expected.
Nevertheless, the overall performance gain from using C++ and avoiding
unnecessary computation is significant (See 5). In the most recent iteration
of ParticleGrid, enabling vector instructions with sparsity already enabled
only resulted in an average speedup of less than 2% in our tests for grids
with $N\geq 64$. Smaller grids do see a more significant speedup, and enabling
vector instructions without sparsity shows an over 2x speedup. However, a
speedup of 8x would be expected from a perfect vectorization implementation.
This suggests either a problem with the implementation or a hardware
limitation, further analysis is underway. The implementation is both simple
and works correctly, so it would seem the implementation runs into a hardware
bottleneck of some kind. We believe ParticleGrid is memory bound on our
system, and so will minimally benefit from further hardware optimization.
However, more research and profiling is needed to pinpoint the cause. The
overall speedup from a baseline implementation can still be over 15x despite
these limitations, thanks mostly to the benefit of sparsity. The exact testing
parameters change this significantly, and it appears cache pollution causes
very large fluctuations in speedup, as much as 10x in our testing when
repeatedly calculating a grid for only one molecule, instead of calculating a
grid for many. This would also point to a memory bottleneck, since a compute-
bound process would be unlikely to see such a significant impact from memory
bandwidth. There is certainly more research to be done in this area to further
identify the exact causes of various performance changes, but even without
this ParticleGrid outperforms other implementations.

We originally implemented the grid generation in Python using NumPy, and also
using Numba. Compared to our optimized C++ implementations, these
implementations are as much as 4 orders of magnitude slower (See 5). Numba
gives a faster result than NumPy, but is still routinely 3 orders of magnitude
slower than the C++ ParticleGrid implementation. All of these implementations
are available in the ParticleGrid repository.

One final consideration is in the context of machine learning, many molecules
will be needed in quick succession. Because separate molecules in the same
batch can all be learned on in parallel, an ML framework can do grid
computation for all of them in parallel. Commonly used frameworks such as
PyTorch can conveniently use multiprocessing to request multiple molecules
simultaneously as they are needed in the training process, freeing the
ParticleGrid implementation from the need to implement this. If ParticleGrid
is indeed memory bound, this will have little impact on performance. We have
yet to fully characterize the impact of multi-threading on the implementation,
but preliminary exploration suggests multi-threading is indeed less effective
than would be naively expected. There is more research to be done in this
area. Despite some unknowns in hardware optimization, ParticleGrid is easily
mature enough to be part of useful, performant applications.

##  VII Application: Property Prediction: Structure is all you need

It is natural to question whether the information contained in 3D structures
is sufficient and required to perform scientific deep learning. In order to
empirically show the importance of the 3D structure for learning on molecular
properties, we train a 3D-CNN to predict molecular properties.

We perform the training and testing on the hardware mentioned in Table I.

CPU | AMD EPYC 7F52 (16 cores) / Node
---|---
Main Memory | 32 GB / Node
GPU | NVIDIA A100 40 GB (x4) / Node
Compiler | GCC 8.3.1
ML framework & Python | Python 3.7
| PyTorch 1.10
| PyTorch-Lightning 1.7.0
TABLE I:  Hardware and software configuration for the Cypress cluster, where
all of the experiments presented in Section VII were performed.

###  VII-A Dataset

We use the Open Graph Benchmarks Large-Scale Challenge (OGB-LSC) [10] for our
exploration on the importance of 3D structure for molecular properties. The
PubChem Quantum Mechanics for Molecules (PCQM4M), a part of the OGB-LSC, is a
quantum chemistry dataset with about 4 million molecular samples created by
the PubChemQC project [39, 40]. This is a subset of the larger PubChem
database, which hosts more than 60 million chemical structures [12]. The
training subset of the OGB-LSC-PCQM4Mv2 dataset, with 3.3 million molecules,
contains 3D coordinate information derived from DFT. Each molecule sample are
of varying sizes and a wide array of elements. The dataset sample includes the
3D structures of the molecule and it’s accompanying HOMO-LUMO energy gap
calculated using DFT. We use 3D coordinates to generate the grids for our
model.

The HOMO-LUMO gap is of interest to domain scientists since it describes the
chemical and electrical properties of the molecule and is a useful indicator
for reactivity and stability.

Physics-based simulations using DFT are computationally demanding and time-
consuming, taking up to several hours for even small molecules. However, high
throughput characterization is important for practical workflows, so efficient
machine learning surrogates for DFT simulations have the potential for high
impact.

Only the training split of the original dataset contains the 3D structural
data, therefore we use that as the basis for our training-validation-test
datasets. We perform a 90-5-5 random split of the original dataset. Since our
custom dataset splits are not publicly available, we release the code and
dataset publicly for reproducibility 222The model and dataset code can be
found here:
www.github.com/ParticleGrid/ParticleGrid/tree/main/examples/PCQM4M-3D.

Figure 6: CDF plot of the molecular weight distribution to verify that the
dataset has sufficient variety of molecules.

###  VII-B Data Representation

We are interested in understanding the information stored in the 3D structure
of the molecules. Therefore we generate grids where all atoms are treated as
the same, and each molecule is represented via a single channel, $(1\times
32\times 32\times 32)$ grid. The grids have adaptive spatial resolution such
that the volume covered by each voxel depends on the volume of the bounding
box the molecule.

###  VII-C Experiment Details

A 10-layer 3D-ResNet with skip connections and batch normalization
architecture is trained for our experiments. The model is relatively small
with only 14.4M trainable parameters. We use the RMSProp optimizer with
learning rate $10^{-3}$ and an exponentially decaying learning rate scheduler
with $\gamma=0.999$. During training, we set the batch size to 256 and train
concurrently on 4 A100s. We use ParticleGrid with grid size of 32, and
variance of 0.4 and generate the grids on the fly during training, validation,
and testing. We train for a total of 400 epochs, and save the model with the
highest performing validation error. The computation cost is only 166 seconds
per epoch, and thus we are able to train the model to completion in less than
20 hours on four A100s. The exact details of the hardware and software stack
can be found in Table. I.

###  VII-D Results

We use the mean squared error (MSE) as our metric of evaluation. We achieve
0.006 MSE on the test dataset. So even though only using the structural
information and single channel grid, we are able to predict the HOMO-LUMO
energy gap to DFT-level accuracy. We can see the prediction accuracy in Fig.
7.

Figure 7: We achieve highly accurate predictions of the HOMO-LUMO energy gap
using our 3D-CNN.

Furthermore, the computation cost still remains negligible compared DFT.

As we can see, three dimensional structures contain significant information
that can help scientists understand the behavior of the molecules.
Significantly, we used no atomic information other than the atomic locations.
We can therefore take other molecules and crystals, represent them in the
exact representation and use transfer learning for multi-task learning.

##  VIII Discussion

DFT is not the only method to obtain 3D structures materials. Semi-empirical
methods such PM and PM7 can provide fast and accurate 3D representation of
large scale molecules. Small molecular conformations can be efficiently
generated with ETKDG [41] with popular cheminformatics tools such as RDKit
[42].

We do note that our formulation in Section III uses a single element for each
element type. While this may be feasible for organic molecules which generally
consists of only a few different types elements, it will be inefficient for
some materials of interest with a wide variety of possible constituent
elements. Since the grid generation process is differentiable, a fact we have
used for the discretizer, we can also attach invertible embedding layers to
“learn” molecular embeddings and generate grids using the molecular
embeddings. We leave this for future work to enable efficient learnable grid
generation on ParticleGrid.

Deep learning kernels for computer vision and natural language processing,
such 2D convolutions, dense matrix multiplication have been highly optimized
for optimal performance, especially on hardware accelerators, such as GPUs,
TPUs, and other ASICs. Optimized GPU kernels in the cuDNN [43] and ROCm
libraries enable fast and efficient processing of images and text. 3D layers
such convolutions and pooling have not been optimized to the same extent due
to utility. Wide scale use of 3D models can spur on hardware and software
advancements to further improve performance.

##  IX Conclusion

We introduce ParticleGrid as a middle layer to facilitate deep learning of a
wide range of material structures. The library allows us to generate physics
inspired grid representations of molecules and crystals that can be used with
deep learning to perform data-driven structure property quantification, or
even inverse design.

Practitioners can take advantage of the gridding functionality to train high-
capacity models large scale 3D datasets such as the PubChemQC PM6 [44] and
OGB-LSC-PCQM4M [10]. We also propose a transformation to 3D coordinate space
to enable applications in generative models and materials discovery.

ParticleGrid, and the 3D grid-based representations, can be used in
conjunction with graph-based and string-based representations to perform
predictive and generative modeling of materials. ParticleGrid provides a
simple Python API for ease of use, and to integrate into the data pipelines
for deep learning frameworks. The CPU optimized algorithm to enable real-time
grid generation for deep learning applications. We achieve up to 9000x, 79.5x,
and 14x over Numpy, Numba, and baseline C++ implementations, respectively.
Using ParticleGrid, we train a 3D CNN to predict molecular HOMO-LUMO energy
gap on a dataset of more than 3.3 million in under 20 hours on 4 GPUs, and
achieve DFT level accuracy. We hope fast and efficient 3D grid generating
libraries can help further 3D modelling of material structures.

## Acknowledgements

We would like to thank TotalEnergies EP Research & Technology for supporting
and allowing this material to be shared.

## References

  * [1] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell _et al._ , “Language models are few-shot learners,” _Advances in neural information processing systems_ , vol. 33, pp. 1877–1901, 2020.
  * [2] T. Karras, S. Laine, and T. Aila, “A style-based generator architecture for generative adversarial networks,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2019, pp. 4401–4410.
  * [3] T. Karras, S. Laine, M. Aittala, J. Hellsten, J. Lehtinen, and T. Aila, “Analyzing and improving the image quality of stylegan,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2020, pp. 8110–8119.
  * [4] A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever, “Zero-shot text-to-image generation,” in _International Conference on Machine Learning_. PMLR, 2021, pp. 8821–8831.
  * [5] D. Weininger, “Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules,” _Journal of chemical information and computer sciences_ , vol. 28, no. 1, pp. 31–36, 1988.
  * [6] M. Krenn, F. Häse, A. Nigam, P. Friederich, and A. Aspuru-Guzik, “Self-referencing embedded strings (selfies): A 100% robust molecular string representation,” _Machine Learning: Science and Technology_ , vol. 1, no. 4, p. 045024, 2020.
  * [7] M. Rupp, A. Tkatchenko, K.-R. Müller, and O. A. Von Lilienfeld, “Fast and accurate modeling of molecular atomization energies with machine learning,” _Physical review letters_ , vol. 108, no. 5, p. 058301, 2012.
  * [8] G. Montavon, K. Hansen, S. Fazli, M. Rupp, F. Biegler, A. Ziehe, A. Tkatchenko, A. Lilienfeld, and K.-R. Müller, “Learning invariant representations of molecules for atomization energy prediction,” _Advances in neural information processing systems_ , vol. 25, 2012.
  * [9] P. Ramachandran and G. Varoquaux, “Mayavi: 3d visualization of scientific data,” _Computing in Science & Engineering_, vol. 13, no. 2, pp. 40–51, 2011.
  * [10] W. Hu, M. Fey, H. Ren, M. Nakata, Y. Dong, and J. Leskovec, “Ogb-lsc: A large-scale challenge for machine learning on graphs,” _arXiv preprint arXiv:2103.09430_ , 2021.
  * [11] E. E. Bolton, J. Chen, S. Kim, L. Han, S. He, W. Shi, V. Simonyan, Y. Sun, P. A. Thiessen, J. Wang _et al._ , “Pubchem3d: a new resource for scientists,” _Journal of cheminformatics_ , vol. 3, no. 1, pp. 1–15, 2011\.
  * [12] S. Kim, P. A. Thiessen, E. E. Bolton, J. Chen, G. Fu, A. Gindulyte, L. Han, J. He, S. He, B. A. Shoemaker _et al._ , “PubChem substance and compound databases,” _Nucleic acids research_ , vol. 44, no. D1, pp. D1202–D1213, 2016.
  * [13] T. L. Griffiths, N. Chater, C. Kemp, A. Perfors, and J. B. Tenenbaum, “Probabilistic models of cognition: Exploring representations and inductive biases,” _Trends in cognitive sciences_ , vol. 14, no. 8, pp. 357–364, 2010\.
  * [14] P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. Zambaldi, M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner _et al._ , “Relational inductive biases, deep learning, and graph networks,” _arXiv preprint arXiv:1806.01261_ , 2018.
  * [15] F. Zhuang, Z. Qi, K. Duan, D. Xi, Y. Zhu, H. Zhu, H. Xiong, and Q. He, “A comprehensive survey on transfer learning,” _Proceedings of the IEEE_ , vol. 109, no. 1, pp. 43–76, 2020.
  * [16] B. Kim, S. Lee, and J. Kim, “Inverse design of porous materials using artificial neural networks,” _Science advances_ , vol. 6, no. 1, p. eaax9324, 2020.
  * [17] M. Ragoza, T. Masuda, and D. R. Koes, “Learning a continuous representation of 3d molecular structures with deep generative models,” _arXiv preprint arXiv:2010.08687_ , 2020.
  * [18] R. Gómez-Bombarelli, J. N. Wei, D. Duvenaud, J. M. Hernández-Lobato, B. Sánchez-Lengeling, D. Sheberla, J. Aguilera-Iparraguirre, T. D. Hirzel, R. P. Adams, and A. Aspuru-Guzik, “Automatic chemical design using a data-driven continuous representation of molecules,” _ACS central science_ , vol. 4, no. 2, pp. 268–276, 2018.
  * [19] M. J. Kusner, B. Paige, and J. M. Hernández-Lobato, “Grammar variational autoencoder,” in _International Conference on Machine Learning_. PMLR, 2017, pp. 1945–1954.
  * [20] G. L. Guimaraes, B. Sanchez-Lengeling, C. Outeiral, P. L. C. Farias, and A. Aspuru-Guzik, “Objective-reinforced generative adversarial networks (organ) for sequence generation models,” _arXiv preprint arXiv:1705.10843_ , 2017.
  * [21] M. Krenn, Q. Ai, S. Barthel, N. Carson, A. Frei, N. C. Frey, P. Friederich, T. Gaudin, A. A. Gayle, K. M. Jablonka _et al._ , “Selfies and the future of molecular string representations,” _arXiv preprint arXiv:2204.00056_ , 2022.
  * [22] Z. Yao, B. Sánchez-Lengeling, N. S. Bobbitt, B. J. Bucior, S. G. H. Kumar, S. P. Collins, T. Burns, T. K. Woo, O. K. Farha, R. Q. Snurr _et al._ , “Inverse design of nanoporous crystalline reticular materials with deep generative models,” _Nature Machine Intelligence_ , vol. 3, no. 1, pp. 76–86, 2021.
  * [23] W. Jin, R. Barzilay, and T. Jaakkola, “Junction tree variational autoencoder for molecular graph generation,” in _International conference on machine learning_. PMLR, 2018, pp. 2323–2332.
  * [24] N. De Cao and T. Kipf, “Molgan: An implicit generative model for small molecular graphs,” _arXiv preprint arXiv:1805.11973_ , 2018.
  * [25] J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl, “Neural message passing for quantum chemistry,” in _International conference on machine learning_. PMLR, 2017, pp. 1263–1272.
  * [26] T. Xie and J. C. Grossman, “Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties,” _Physical review letters_ , vol. 120, no. 14, p. 145301, 2018.
  * [27] K. Schütt, P.-J. Kindermans, H. E. Sauceda Felix, S. Chmiela, A. Tkatchenko, and K.-R. Müller, “Schnet: A continuous-filter convolutional neural network for modeling quantum interactions,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [28] J. Klicpera, F. Becker, and S. Günnemann, “Gemnet: Universal directional graph neural networks for molecules,” _arXiv preprint arXiv:2106.08903_ , 2021.
  * [29] J. Hoffmann, L. Maestrati, Y. Sawada, J. Tang, J. M. Sellier, and Y. Bengio, “Data-driven approach to encoding and decoding 3-d crystal structures,” _arXiv preprint arXiv:1909.00949_ , 2019.
  * [30] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant, “Array programming with NumPy,” _Nature_ , vol. 585, no. 7825, pp. 357–362, Sep. 2020. [Online]. Available: https://doi.org/10.1038/s41586-020-2649-2
  * [31] H. Edelsbrunner, J. Harer _et al._ , “Persistent homology-a survey,” _Contemporary mathematics_ , vol. 453, pp. 257–282, 2008.
  * [32] S. Kaji, T. Sudo, and K. Ahara, “Cubical ripser: Software for computing persistent homology of image and volume data,” _arXiv preprint arXiv:2005.12692_ , 2020.
  * [33] V. De Silva, D. Morozov, and M. Vejdemo-Johansson, “Dualities in persistent (co) homology,” _Inverse Problems_ , vol. 27, no. 12, p. 124003, 2011.
  * [34] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, “Pytorch: An imperative style, high-performance deep learning library,” in _Advances in Neural Information Processing Systems 32_ , H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, Eds. Curran Associates, Inc., 2019, pp. 8024–8035. [Online]. Available: http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * [35] M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Jozefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, I. Sutskever, K. Talwar, P. Tucker, V. Vanhoucke, V. Vasudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng, “TensorFlow: Large-scale machine learning on heterogeneous systems,” 2015, software available from tensorflow.org. [Online]. Available: https://www.tensorflow.org/
  * [36] D. Polykovskiy, A. Zhebrak, B. Sanchez-Lengeling, S. Golovanov, O. Tatanov, S. Belyaev, R. Kurbanov, A. Artamonov, V. Aladinskiy, M. Veselov _et al._ , “Molecular sets (moses): a benchmarking platform for molecular generation models,” _Frontiers in pharmacology_ , vol. 11, p. 1931, 2020.
  * [37] H. Schöpf and P. Supancic, “On bürmann’s theorem and its application to problems of linear and nonlinear heat transfer and diffusion,” _The Mathematica Journal_ , vol. 16, no. 11, 2014.
  * [38] S. K. Lam, A. Pitrou, and S. Seibert, “Numba: A llvm-based python jit compiler,” in _Proceedings of the Second Workshop on the LLVM Compiler Infrastructure in HPC_ , 2015, pp. 1–6.
  * [39] M. Nakata and T. Shimazaki, “PubChemQC project: a large-scale first-principles electronic structure database for data-driven chemistry,” _Journal of chemical information and modeling_ , vol. 57, no. 6, pp. 1300–1308, 2017.
  * [40] N. Maho, “The PubChemQC project: A large chemical database from the first principle calculations,” in _AIP conference proceedings_ , vol. 1702, no. 1. AIP Publishing LLC, 2015, p. 090058\.
  * [41] S. Wang, J. Witek, G. A. Landrum, and S. Riniker, “Improving conformer generation for small rings and macrocycles based on distance geometry and experimental torsional-angle preferences,” _Journal of chemical information and modeling_ , vol. 60, no. 4, pp. 2044–2058, 2020.
  * [42] G. Landrum _et al._ , “Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling,” 2013.
  * [43] S. Chetlur, C. Woolley, P. Vandermersch, J. Cohen, J. Tran, B. Catanzaro, and E. Shelhamer, “cudnn: Efficient primitives for deep learning,” _arXiv preprint arXiv:1410.0759_ , 2014.
  * [44] M. Nakata, T. Shimazaki, M. Hashimoto, and T. Maeda, “Pubchemqc pm6: data sets of 221 million molecules with optimized molecular geometries and electronic properties,” _Journal of Chemical Information and Modeling_ , vol. 60, no. 12, pp. 5891–5899, 2020.