# Gaussian Plane-Wave Neural Operator for Electron Density Estimation

Seongsu Kim  Sungsoo Ahn

###### Abstract

This work studies machine learning for electron density prediction, which is
fundamental for understanding chemical systems and density functional theory
(DFT) simulations. To this end, we introduce the Gaussian plane-wave neural
operator (GPWNO), which operates in the infinite-dimensional functional space
using the plane-wave and Gaussian-type orbital bases, widely recognized in the
context of DFT. In particular, both high- and low-frequency components of the
density can be effectively represented due to the complementary nature of the
two bases. Extensive experiments on QM9, MD, and material project datasets
demonstrate GPWNO’s superior performance over ten baselines.

Machine Learning, ICML

##  1 Introduction

The study of electron density is fundamental for understanding the properties
and behaviors of chemical systems. Its estimation is a crucial component of
Kohn-Sham density functional theory (Kohn & Sham, 1965, DFT) for simulations
like ground-state energy estimation with applications to the development of
new battery cathodes (Kondrakov et al., 2017), solar cell materials (Chang et
al., 2021) or drug design (Adekoya et al., 2022). However, the computational
complexity of the DFT is cubic in the system size and is often prohibitively
expensive for real-world problems.

Machine learning has shown promise in predicting electron density with reduced
computational complexity. Pioneering works investigated methods based on
kernel ridge regression (Grisafi et al., 2018; Bogojeski et al., 2018) and
Gaussian process-based regression (Grisafi et al., 2018; Fabrizio et al.,
2019). Recently, deep learning methods have been increasingly investigated due
to their remarkable capability to capture complex relationships within data.
Sinitskiy & Pande (2018) utilized a voxel-based 3D convolutional net with a
U-Net architecture to predict density at a voxel level. Jørgensen & Bhowmik
(2020) and Jørgensen & Bhowmik (2022) aimed to predict the electron density
using graph neural networks, where the graph consists of interacting atom
vertices and a special query point vertex for which the density is predicted.

More recently, Cheng & Peng (2023) tackled the problem with neural operator,
which is an attractive paradigm that can map the given molecule into an
function space of electron density whose output can be evaluated at an
arbitrary point. In particular, they proposed to predict the density using
atom-centered Gaussian-type orbital (Hehre et al., 1972, GTO). They formulated
their message passing scheme as a graphon convolution.

Notably, the exploration of the plane-wave (PW) basis, a cornerstone in DFT
solvers like VASP (Hafner, 2008), Quantum Espresso (Giannozzi et al., 2017),
and ABINIT (Gonze et al., 2009), remains absent within the domain of neural
operators. The PW basis comprises periodic wavefunctions that meet the
periodic boundary condition of the molecule. It is an attractive approach due
to its ease of manipulation through fast Fourier transformation, and its
ability to effectively capture low-frequency signals such as the ones
associated with long-range interactions, e.g., Coloumb potentials. The most
similar work is the Fourier neural operator (FNO) baseline considered by Cheng
& Peng (2023), which uses the Fourier basis consisting of periodic functions.
The difference is how the Fourier basis does not respect the periodic boundary
condition while the PW basis does.

Motivated by this lack of exploration, we aim to predict the electron density
using neural operators that operate on the PW basis. However, using a
moderately-sized PW basis, e.g., with $10^{3}$ elements, alone cannot
accurately express the electron density, which is in agreement with the poor
performance of FNO reported by Cheng & Peng (2023). While DFT solvers overcome
this by increasing the size, e.g., up to $8.9\times 10^{7}$, the neural
operators are not scalable to such high frequencies. Our focus is to enjoy the
benefits of PW basis while keeping the basis size realistic for neural
operators. Prior works on DFT, e.g., pseudopotential (Harrison, 1966; Bachelet
et al., 1982), linearized augmented plane-wave (Koelling & Arbman, 1975),
projector augmented wave (Blöchl, 1994), and mixed density fitting (Sun et
al., 2017), hint that it is possible to reduce the required basis size by
decoupling vicinity of atoms from the rest of the regions and mixing local
orbitals with PWs.

(a)

(b)

Figure 1: Illustration of Gasussian plane wave neural operator (GPWNO). (left)
Our GPWNO decomposes the electron density estimation into two regions: high-
frequency and fast-decaying regions for GTO basis and low-frequency and slow-
decaying regions for the PW basis. (right) Our GPWNO makes two predictions
based on the GTO and the PW bases. The PW-based prediction is from the
lattice-based discretization via the probe nodes, while the GTO-based
prediction is from the atom-wise message-passing scheme. Final output is
evaluated at an arbitrary query point.

Contribution. In this work, we propose Gaussian plane-wave neural operator
(GPWNO) for estimating the electron density. Our main idea is to decompose the
electron density into local atom-wise Gaussian-type orbital (GTO) and global
plane-wave (PW) bases, where the bases complement each other for the
expression of high- and low-frequency signals. In addition, our design choice
allows prediction that is both SE(3) equivariant and satisfies the periodic
boundary condition. We provide an overview of GPWNO in Figure 1.

Our main component is the PW-based electron density prediction, which is based
on PW neural operator layers with the convolution parameterized in the
reciprocal space of the molecule, i.e., PW coefficients. To this end, we
discretize the signal via probe nodes regularly placed along the lattice
structure and apply discrete Fourier transformations to compute the
convolution. Finally, the density at query nodes is evaluated via message
passing. Our main idea on the lattice-based discretization, which allows
convolution of the signals in the reciprocal space of molecules while
satisfying the periodic boundary condition.

Our next component is the GTO-based electron density prediction, which can be
defined as a neural operator in the graphon space (Cheng & Peng, 2023). We
also introduce high-frequency masking scheme that allows the PW basis to
concentrate on regions outside the vicinity of atoms. Our main idea here is to
exploit the nature of GTO complementing the PW basis, which was inspired by
DFT solvers like Gaussian plane-wave methods (Lippert et al., 1997), LAPW with
local orbitals (Koelling & Arbman, 1975), and mixed density fitting (Sun et
al., 2017).

We empirically verify the effectiveness of our methods on three datasets: QM9
(Ruddigkeit et al., 2012; Ramakrishnan et al., 2014; Jørgensen & Bhowmik,
2022), MD17 (Brockherde et al., 2017; Bogojeski et al., 2020) datasets. We
additionally benchmark our algorithm on the newly curated datasets from the
material project (Jain et al., 2013; Shen et al., 2021, MP), categorized by
seven different crystal family types, e.g., the crystal being triclinic and
hexagonal. We show that our algorithm consistently outperforms the prior works
by a large margin. We validate the effectiveness of our algorithmic components
through an ablation study.

To summarize, our contribution is twofold:

  * •

We propose GPWNO for electron density prediction based on the mixture of GTO
and PW bases which complement each other for the expression of high- and low-
frequency components in the density. We design the framework to satisfy
SE(3)-equivariance and periodic boundary conditions.

  * •

We empirically validate our GPWNO for the QM9, MD17, and newly curated MP
datasets with comparison against ten baselines. We also empirically verify the
effectiveness of each component of our algorithm through ablation studies.

##  2 Related Works

Machine learning for electron density estimation. Initially, researchers have
resorted to statistical approaches such as kernel ridge regression (Grisafi et
al., 2018; Bogojeski et al., 2018) and symmetry-adapted Gaussian processes
(Grisafi et al., 2018; Fabrizio et al., 2019) for electron density estimation.
Recent advance in deep learning have opened new avenues for modeling these
functions. For example, Sinitskiy & Pande (2018) employed a voxel-based 3D
U-Net convolutional network architecture, to predict electron densities at the
voxel level.

Furthermore, Jørgensen & Bhowmik (2020) and Jørgensen & Bhowmik (2022)
approached electron density prediction using neural message passing on graphs.
Finally, Cheng & Peng (2023) embraced the idea of coefficient learning within
the multicentric approximation framework and developed an equivariant graph
neural network (GNN) based on tensor products for predicting density spectra.

Neural operator. Neural operators (Li et al., 2020b, a; Kovachki et al., 2023;
Tran et al., 2023, NOs) form a paradigm to map data into infinite-dimensional
function spaces where these spaces and the data share a common domain, e.g.,
from 3D vehicle geometries to pressure on car surfaces (Li et al., 2023), from
molecular systems to electron densities (Cheng & Peng, 2023) and from seismic
wavefield to probability to P- and S- wave arrivals (Sun et al., 2023). The
input function for neural operators can be represented with any level of
discretization, grid, resolution, or mesh. The output function, in turn, can
be evaluated at any arbitrary point.

Initially, the graph neural operator (Li et al., 2020b, GNO) proposed to
integrate kernel operation within graph structures, drawing similarities to
the message-passing networks. Next, the Fourier neural operator (Li et al.,
2020a, FNO) aims to parameterize and learn operators within the Fourier
domain. This is done by discretizing the input domain and applying the fast
Fourier transform (FFT). To improve the scalability of FNO, factorized-FNO
(Tran et al., 2023, F-FNO), employs Fourier factorization, learning features
in the Fourier space independently across each dimension. Finally, to overcome
the limits of FNO being restricted to regular grids for using FFT operations,
the geometry-informed neural operator (Li et al., 2023, GINO) maps irregular
grids into a latent regular grid and then applies FFT.

##  3 Preliminary

###  3.1 Problem Definition

Here, we formulate the electron density prediction problem as learning a
mapping from a three-dimensional Cartesian coordinate to a real number, given
a molecule. We consider periodic molecules, where the periodicity is defined
by the lattice structure. We also discuss extending our framework to aperiodic
molecules.

To be specific, we represent a molecule of $N$ atoms as
$\mathcal{M}=(\mathcal{X},\mathcal{A},{\bm{L}})$ with atom positions
$\mathcal{X}$, atomic features $\mathcal{A}$, and unit cell lattice
${\bm{L}}$. The atomic positions
$\mathcal{X}=\\{{\bm{x}}_{u}\mathrel{\mathop{\mathchar
58\relax}}u\in\mathcal{V}\\}$ and the atom features
$\mathcal{A}=\\{{\bm{z}}_{u}\mathrel{\mathop{\mathchar
58\relax}}u\in\mathcal{V}\\}$ where we use vertices $\mathcal{V}$ to index the
atoms. Finally, we let the unit cell lattice
${\bm{L}}=[{\bm{a}}_{1},{\bm{a}}_{2},{\bm{a}}_{3}]^{\top}$ describe the
periodicity where each ${\bm{a}}_{i}$ representing a lattice vector.

Given the molecule $\mathcal{M}=(\mathcal{X},\mathcal{A},{\bm{L}})$, the
infinite molecular representation $(\bar{\mathcal{X}},\bar{\mathcal{A}})$ with
atom positions $\bar{\mathcal{X}}$ and atom features $\bar{\mathcal{A}}$ is
defined as follows:

$\displaystyle\bar{\mathcal{X}}$
$\displaystyle=\bigg{\\{}{\bm{x}}_{{\bm{t}},u}\bigg{|}{\bm{x}}_{{\bm{t}},u}={\bm{x}}_{u}+\sum_{d=1}^{3}t_{d}{\bm{a}}_{d},{\bm{t}}\in\mathbb{Z}^{3},u\in\mathcal{V}\bigg{\\}},$
(1) $\displaystyle\bar{\mathcal{A}}$
$\displaystyle=\\{{\bm{z}}_{{\bm{t}},u}|{\bm{z}}_{{\bm{t}},u}={\bm{z}}_{u},{\bm{t}}\in\mathbb{Z}^{3},u\in\mathcal{V}\\},$
(2)

where ${\bm{t}}=[t_{1},t_{2},t_{3}]\in\mathbb{Z}^{3}$ indicates the change of
unit cell over the infinitely repeating periodic patterns.

Consequently, the true electron density function
$\rho\mathrel{\mathop{\mathchar 58\relax}}\mathbb{R}^{3}\rightarrow\mathbb{R}$
also satisfies the following periodic boundary condition:

$\rho({\bm{x}})=\rho\bigg{(}{\bm{x}}+\sum_{d=1}^{3}t_{d}{\bm{a}}_{d}\bigg{)},$
(3)

where ${\bm{t}}=[t_{1},t_{2},t_{3}]\in\mathbb{Z}^{3}$ indicates the change of
unit cell over the infinitely repeating periodic patterns.

Extension to aperiodic molecules. For aperiodic molecules, we follow the
common practice of formulating them as periodic molecules with a sufficiently
large unit cell lattice ${\bm{L}}$ with zero-padded boundaries. Examples
include DFT solvers (Blöchl, 1995; Hafner, 2008; Ulian et al., 2013) and graph
neural networks (Kosmala et al., 2023).

###  3.2 Plane-Wave Basis

The plane-wave (PW) basis is widely used in electronic structure computation
of solid-state systems (Martin, 2004; Kittel, 2005) because it satisfies the
periodic boundary condition by construction and is easily manipulated through
fast Fourier transformations. We consider the truncated version of the
(infinite) PW basis determined by a cutoff in the maximum frequency (or
energy), as commonly employed by DFT solvers and Fourier neural operators.

Formally, given a lattice
${\bm{L}}=[{\bm{a}}_{1},{\bm{a}}_{2},{\bm{a}}_{3}]^{\top}$, the wavevectors
are constructed from reciprocal primitive vectors
$({\bm{b}}_{1},{\bm{b}}_{2},{\bm{b}}_{3})$ where ${\bm{b}}_{1}$ is defined by
$({\bm{a}}_{1},{\bm{a}}_{2},{\bm{a}}_{3})$ as:

${\bm{b}}_{1}=2\pi\frac{{\bm{a}}_{2}\times{\bm{a}}_{3}}{{\bm{a}}_{1}\cdot({\bm{a}}_{2}\times{\bm{a}}_{3})}$
(4)

and ${\bm{b}}_{2}$ and ${\bm{b}}_{3}$ are defined similarly by
$({\bm{a}}_{2},{\bm{a}}_{3},{\bm{a}}_{1})$ and
$({\bm{a}}_{3},{\bm{a}}_{1},{\bm{a}}_{2})$, respectively. Note the dual
relationship between the lattice vectors and the primitive vectors, e.g.,
${\bm{a}}_{1}\cdot{\bm{b}}_{1}=2\pi$ and ${\bm{a}}_{1}\cdot{\bm{b}}_{2}=0$.

Finally, the PW basis element $\phi_{\bm{\lambda}}({\bm{x}})$ indexed by a
tuple of integers
$\bm{\lambda}=[\lambda_{1},\lambda_{2},\lambda_{3}]^{\top}\in\mathbb{Z}^{3}$
is defined as follows:

$\phi_{\bm{\lambda}}({\bm{x}})=\exp\bigg{(}j\sum_{d=1}^{3}\lambda_{d}{\bm{b}}_{d}\cdot{\bm{x}}\bigg{)},$
(5)

where $j$ is the imaginary number. It should be noted that the plane wave
basis element for any
$\bm{\lambda}=[\lambda_{1},\lambda_{2},\lambda_{3}]^{\top}\in\mathbb{Z}^{3}$
satisfies the periodic boundary condition, i.e. Equation 3. We also point out
how the PW basis was originally designed for the wavefunctions to solve
Schrödinger’s equation (Kittel, 2005). However, in this work, we use them for
parameterization of the electron density functions which is the square of the
magnitude of the wavefunction.

###  3.3 Gaussian-Type Orbital Basis

Gaussian-type orbitals (Hehre et al., 1972, GTO) offer an alternative to the
PW basis for the expression of electronic structures. We consider GTOs in the
following form:

$\phi_{nlm}({\bm{x}},{\bm{x}}_{0})=C_{nlm}R_{nl}(\lVert{\bm{r}}\rVert)Y_{lm}\bigg{(}\frac{{\bm{r}}}{\lVert{\bm{r}}\rVert}\bigg{)},$
(6)

where ${\bm{r}}={\bm{x}}-{\bm{x}}_{0}$, ${\bm{x}}_{0}$ is center of the GTO,
$C_{nlm}$ is a normalization constant, $R_{nl}(\cdot)$ is the Gaussian radial
basis function (RBF), and $Y_{lm}(\cdot)$ is the spherical harmonics (SH). The
index of radial basis $n$ and degree $l$ of the spherical harmonics is a non-
negative integer, and the magnetic order $m$ is an integer satisfying $-l\leq
m\leq l$. The GTO basis does not satisfy the periodic boundary condition
Equation 3.

The use of Gaussian-type orbitals in quantum chemistry has become popular due
to their mathematical simplicity and efficiency in electronic structure
computation (Hehre et al., 1972). They allow straightforward integrations and
often result in faster convergence in numerical simulations compared to the
Slater-type orbitals (Gill, 1994).

Figure 2: Neural architecture for the PWNO layer. Starting from a molecule
${\mathcal{M}}$, our framework makes a PW-based prediction $\rho_{\text{PW}}$
for given query points ${\mathcal{Q}}$. The PWNO layer processes the
discretized signal using convolution, which is expressed as applying FT,
frequency-wise linear transformation $\mathcal{R}$ and inverse FT in sequence.
Point-wise activation function $\sigma$ is introduced for the non-linear
transformation of the discretized signals.

##  4 Method

###  4.1 Overview

Here, we provide a high-level description of our Gaussian plane wave neural
operator (GPWNO). Our main idea is to introduce the plane-wave basis set for
effectively representing the electron density functions. However, since the
electron density function consists of high-frequency elements near the atoms,
we additionally introduce the GTO-based basis set for representing such
elements. We depict the visualization of the overall architecture in Appendix
A.

Decomposition of prediction. We parameterize our electron density prediction
$\rho_{\text{pred}}\mathrel{\mathop{\mathchar
58\relax}}\mathbb{R}^{3}\rightarrow\mathbb{R}$ as follows:

$\rho_{\text{pred}}({\bm{x}})=\rho_{\text{PW}}({\bm{x}})+\rho_{\text{GTO}}({\bm{x}}),$
(7)

where $\rho_{\text{PW}}$ and $\rho_{\text{GTO}}$ corresponds to our prediction
using the plane-wave basis set and the GTO-based basis set, respectively. For
training, we use the mean-squared error:

$\mathcal{L}_{\mathcal{M}}=\frac{1}{|\mathcal{D}_{\mathcal{M}}|}\sum_{{\bm{x}}\in\mathcal{D}_{\mathcal{M}}}\lVert\rho({\bm{x}})-\rho_{\text{pred}}({\bm{x}})\rVert^{2},$
(8)

where we train over different molecules $\mathcal{M}\in\mathcal{D}$ and the
locations ${\bm{x}}\in\mathcal{D}_{\mathcal{M}}$ is the dataset of locations
where the electron density is known, e.g., computed by DFT simulations.

PBC implementation. Importantly, our parameterization gaurantees the
predictions to satisfy the periodic boundary conditions on PW and GTO-based
prediction. Specifically, we utilized the crystal graph framework introduced
by Xie & Grossman (2018), which emulates periodic boundary conditions common
in computational physics, as outlined in Landau & Binder (2021). In simple
terms, we consider each atom to have multiple coordinates, each translated by
integer multiples of lattice vector, while maintaining identical node features
during message passing.

Equivariance. Both the electron density function and our predictions are
scalar fields that are equivariant under SE(3) transformations. In particular,
our PW-based prediction is a scalar field since the PW basis aligns with the
lattice, which transforms along with the molecule and hence the PW
coefficients are invariant.111We provide proof for the aperioic molecule in
Appendix C The proof for the GTO-based prediction being a scalar field follows
that of Cheng & Peng (2023).

Masking for decomposition. Moreover, the density prediction
$\rho_{\text{PW}}({\bm{x}})$, as characterized by the PW basis set, is
designed to be zero within a specific cutoff radius for certain atoms. This
approach aids in segmenting the target region for each basis and integrates
the physical understanding that high-frequency electron density components are
indicative of core (non-valence) electrons, which are typically located near
an atom’s center.

###  4.2 PW-based Prediction

Here, we further describe our PW-based prediction $\rho$ of the electron
density. To this end, following the Fourier neural operator (Li et al., 2020a,
FNO), we let each layer of our network implicitly represent a signal updated
in both the Fourier space and the Euclidean space. However, unlike the FNO
approach, our signal is represented using the PW basis and the corresponding
coefficients, which is a subset of Fourier series that satisfies the periodic
boundary condition.

Signal initialization. To be specific, we initialize the first layer using
composition of tensor field network (Thomas et al., 2018, TFN), grid point-
wise embedding, and discrete Fourier transformation. We first obtain the atom-
wise embedding of the molecule as
$\mathbf{f}_{u}^{(0)}=\operatorname{TFN}_{u}(\mathcal{M})$, where
$\operatorname{TFN}_{u}$ denotes the TFN output associated with the atom
$u\in\mathcal{V}$. We provide more details of our TFN architecture in Appendix
B.

Lattice-based discretization. For point-wise embedding, we locate probe nodes
$\mathcal{P}$ on a three-dimensional space inside the lattice
$\mathbf{L}=[{\bm{a}}_{1},{\bm{a}}_{2},{\bm{a}}_{3}]^{\top}$. Specifically, we
uniformly discretized the space inside the lattice and located the probe node
$p$ at each point. i.e., a probe node $p\in\mathcal{P}$ is placed at

${\bm{x}}_{p}=\frac{d_{1}}{D}{\bm{a}}_{1}+\frac{d_{2}}{D}{\bm{a}}_{2}+\frac{d_{3}}{D}{\bm{a}}_{3},$
(9)

where $d_{1},d_{2},d_{3}<D$ are non-negative integers and the total number of
probe nodes is $M=D^{3}$. Note that this lattice generalizes the regular voxel
grids with cubic cells.

Then, we aggregate messages from the atoms $u\in\mathcal{V}$ in vicinity of a
probe point $p\in\mathcal{P}$:

$\mathbf{f}_{p}^{(1)}=\operatorname{UPD}\circ\operatornamewithlimits{AGG}(\\{(r_{up},\mathbf{f}_{u}^{(0)})\mathrel{\mathop{\mathchar
58\relax}}u\in{\mathcal{N}_{\mathcal{V}}}(p)\\}),$ (10)

where $r_{up}=\|{\bm{x}}_{p}-{\bm{x}}_{u}\|$ and
$\mathcal{N}_{\mathcal{V}}(p)\subseteq\mathcal{V}$ is the set of atoms within
a certain radius for the probe node $p\in\mathcal{P}$. Furthermore,
$\operatorname{AGG}$ and $\operatorname{UPD}$ are the aggregation and the
update function typically used for graph neural networks. The grid embeddings
$\\{\mathbf{f}_{p}^{(1)}\mathrel{\mathop{\mathchar
58\relax}}p\in\mathcal{P}\\}$ implicitly represent a signal $f^{(1)}$
evaluated at location ${\bm{x}}_{p}$, i.e.,
$\mathbf{f}_{p}^{(1)}=f^{(1)}({\bm{x}}_{p})$.

PW neural operator layer. To represent the interaction between the PW bases,
we employ the PW neural operator layer (PWNO). PWNO iteratively updates the
signal $f^{(h)}\mathrel{\mathop{\mathchar
58\relax}}\mathbb{R}^{3}\to\mathbb{R}^{C}$ by a convolution operator using the
PW basis where $C$ is the channel dimension of signal, i.e., sequentially
applying Fourier transformation, frequency-wise linear transformation, and
inverse Fourier transformation:

$f^{(h+1)}=\sigma(\mathcal{W}\cdot f^{(h)}+\mathcal{K}\ast f^{(h)}),$ (11)

where $\ast$ is the convolution operator, $\mathcal{K}$ is the kernel,
$\mathcal{W}$ is the point-wise linear operator, and $\sigma$ is the point-
wise non-linearities. i.e., $\operatorname{GELU}$ (Hendrycks & Gimpel, 2016).
We implement the convolution operator via linear transformation in the Fourier
space:

$\mathcal{K}\ast
f^{(h)}=\operatorname{FT}^{-1}(\mathcal{R}\cdot\operatorname{FT}(f^{(h)})),$
(12)

where $\mathcal{R}$ is the frequency-wise linear transformation,
$\operatorname{FT}$ is the Fourier transformation that maps point-wise
representation into the PW basis coefficients, and $\operatorname{FT}^{-1}$
its inverse.

To implement the Fourier transformation in discrete space, we operate with
point-wise representations
$\\{\mathbf{f}_{p}^{(h)}\in\mathbb{R}^{C}\mathrel{\mathop{\mathchar
58\relax}}p\in\mathcal{P}\\}$ and frequency-wise representations
$\\{\hat{\mathbf{f}}_{\bm{\lambda}}^{(h)}\in\mathbb{R}^{C}\mathrel{\mathop{\mathchar
58\relax}}\bm{\lambda}\in\Lambda\\}$ of the underlying signal $f^{(h)}$, e.g.,
$f^{(h)}(\bm{x})=\sum_{\bm{\lambda}\in\Lambda}\hat{\mathbf{f}}_{\bm{\lambda}}^{(h)}\phi_{\bm{\lambda}}(\bm{x})$
where $\phi_{\bm{\lambda}}(\bm{x})$ is the PW with wavevector $\bm{\lambda}$.
Our framework operates only on the wavevectors $\Lambda$ of the PW basis,
incorporating the periodicity constraints regarding the lattice structure
$\Lambda$. We depicted the architecture of the PWNO layer in Figure 2.

Prediction network. At the final $H$-th layer, given a query point
$\bm{x}_{q}$, we compute the output as follows:

$\mathbf{f}_{q}^{(H)}=\operatorname{UPD}\circ\operatornamewithlimits{AGG}(\\{(r_{pq},\mathbf{f}_{p}^{(H-1)})\mathrel{\mathop{\mathchar
58\relax}}p\in{\mathcal{N}}_{\mathcal{P}}(q)\\}),$ (13)

where $r_{pq}=\|{\bm{x}}_{q}-{\bm{x}}_{p}\|$ and
${\mathcal{N}}_{\mathcal{P}}(q)\subseteq{\mathcal{P}}$ is the set of probe
nodes within a certain radius for the query node $q$. Note that the query node
$p$ can be placed at any points in the continuous space unlike the probe node
$p$ which is regularly placed along the lattice. One can also process multiple
queries in parallel without re-computing the final probe-wise features
$\mathbf{f}_{p}^{(H-1)}$ for different queries.

Table 1: Evaluation of GPWNO for aperiodic materials. We report the
performance for QM9 and MD dataset in NMAE (%). The best number is highlighted
in bold. The baseline results are from Cheng & Peng (2023). Each number is
averaged over three runs. For brevity, we denote DimeNet and DeepDFT by DmNet
and DDFT, respectively.

Dataset | LNO | FNO | GNO | DmNet++ | DmNet | EGNN | DDFT2 | DDFT | CNN | InfGCN | GPWNO
---|---|---|---|---|---|---|---|---|---|---|---
QM9 | 26.14 | 28.83 | 40.86 | 11.69 | 11.97 | 11.92 |  01.03 | 2.95 |  02.01 |  00.93 | 0.73
Ethanol | 43.17 | 31.98 | 82.35 | 14.24 | 13.99 | 13.90 |  08.83 | 7.34 | 13.97 |  08.43 | 4.00
Benzene | 38.82 | 20.05 | 82.46 | 14.34 | 14.48 | 13.49 |  05.49 | 6.61 | 11.98 |  05.11 | 2.45
Phenol | 60.70 | 42.98 | 66.69 | 12.99 | 12.93 | 13.59 |  07.00 | 9.09 | 11.52 |  05.51 | 2.68
Resorcinol | 35.07 | 26.06 | 58.75 | 12.01 | 12.04 | 12.61 |  06.95 | 8.18 | 11.07 |  05.95 | 2.73
Ethane | 77.14 | 26.31 | 71.12 | 12.95 | 13.11 | 15.17 |  06.36 | 8.31 | 14.72 |  07.01 | 3.67
MDA | 47.22 | 34.58 | 84.52 | 16.79 | 18.71 | 12.37 | 10.68 | 9.31 | 18.52 | 10.34 | 5.32

| QM9 (O2C3H4) | Monoclinic (Sr(P3Pt2)2) | Tetragonal(Eu(ZnGe)2)
---|---|---|---
|  GT |  GPWNO |  InfGCN |  GT |  GPWNO |  InfGCN |  GT |  GPWNO |  InfGCN
Pred. |  |  |  |  |  |  |  |  |
Error |  |  |  |  |  |  |  |  |
NMAE |  |  0.47 |  0.59 |  |  5.48 |  5.78 |  |  4.97 |  5.28

| QM9 (C3ONH5) | Triclinic (Li9Mn2Co5O16) | Orthorhombic (Dy5Sn11)
---|---|---|---
Pred. |  |  |  |  |  |  |  |  |
Error |  |  |  |  |  |  |  |  |
NMAE |  |  0.56 |  0.60 |  |  4.00 |  5.89 |  |  4.59 |  5.15

Figure 3: Visualization of ground truth (GT), density prediction, and error in
NMAE (%). For electron density prediction (Pred.) on QM9 and MP datasets, red
and blue colors indicate higher electron density, respectively. For the error
(Error), lighter colors indicates lower error.

High frequency masking. To promote the synergy between the PW and GTO-based
predictions, we employ high frequency mask (HFM), which constrains the PW
prediction to low frequency region that is distant from atoms. To this end, we
propose to mask the PW-based prediction near the atom as follows:

$a({\bm{x}})=\begin{cases}1\quad\text{if
}\min_{u\in\mathcal{V}}d({\bm{x}},{\bm{x}}_{u})>r_{\text{mask}},\\\
0\quad\text{else}.\end{cases}$ (14)

Note that our masking introduces discontinuities near the cutoff boundary,
which we empirically observe to be ignorable.222One could replace the masking
function with an envelope function (Gasteiger et al., 2020b) for the
continuous cutoff. This results in the following prediction:

$\rho_{\text{PW}}({\bm{x}}_{q})=a({\bm{x}})\mathbf{f}_{q}^{(H)},$

where ${\bm{x}}_{q}$ is the input query point. We provide ablation studies on
the effectiveness of HFM in Section 5.4.

###  4.3 GTO-based Prediction

Finally, we describe our prediction using the GTO basis set. Our philosophy
mostly follows the multicentric expansion scheme of InfGCN (Cheng & Peng,
2023). Note that we can design the GTO-based prediction to satisfy the
periodic boundary condition by introducing a periodical graph which is used in
Xie & Grossman (2018).

The GTO-based prediction $\rho_{\text{GTO}}$ is constructed from the linear
combination of GTO basis set as follows:

$\rho_{\text{GTO}}({\bm{x}})=\sum_{u\in\mathcal{V}}\sum_{n=1}^{N}\sum_{l=1}^{L}\sum_{m=-l}^{l}\hat{f}_{u,nlm}\phi_{nlm}({\bm{x}},{\bm{x}}_{u}),$
(15)

where $\phi_{nlm}({\bm{x}},{\bm{x}}_{u})$ is the GTO as described in Equation
6 centered at each atom position ${\bm{x}}_{u}$ for $u\in\mathcal{V}$.
Furthermore, $\hat{f}_{u,nlm}\in\mathbb{R}$ is the atom-wise coefficient term
constructed from our neural operator given the molecule $\mathcal{M}$.
Finally, $N$ and $L$ are the truncated maximum degrees for the radial basis
and the spherical harmonics.

To be specific, the atom-wise coefficient terms are computed by TFN-based
layer and its implementation (Thomas et al., 2018; Geiger & Smidt, 2022) to
achieve the equivariant message passing. The spherical tensor feature
$\hat{\mathbf{f}}_{u,nl}^{(h)}=\\{\hat{f}_{u,nlm}^{(h)}\mathrel{\mathop{\mathchar
58\relax}}m\in\mathbb{Z},-l\leq m\leq l\\}$ is given by the following updates:

$\displaystyle\hat{\mathbf{f}}^{(h+1)}_{u,nl}$
$\displaystyle=\sum_{v\in\mathcal{N}_{{\mathcal{V}}}(u)}\sum_{k\geq
0}W^{(h)}_{nlk}({\bm{x}}_{v}-{\bm{x}}_{u})\hat{\mathbf{f}}^{(h)}_{v,nk},$ (16)
$\displaystyle W^{(h)}_{nlk}({\bm{r}})$
$\displaystyle=\sum^{k+1}_{J=\|k-l\|}\varphi_{nlk}^{J}(r)\sum^{J}_{m=-J}Y_{Jm}\left(\hat{r}\right)Q^{lk}_{Jm},$
(17)

where $r=\lVert{\bm{r}}\rVert$, $\hat{r}=\frac{{\bm{r}}}{r}$,
$\mathcal{N}_{{\mathcal{V}}}(u)$ is the set of neighbors of atom $u$,
$Q^{lk}_{Jm}\in\mathbb{R}^{(2l+2)\times(2k+1)}$ is the Clebsh-Gordan matrix,
and $\varphi_{nlk}^{J}\mathrel{\mathop{\mathchar
58\relax}}\mathbb{R}^{+}\rightarrow\mathbb{R}$ is a learnable scalar function.

##  5 Experiment

###  5.1 Experimental Setting

Baselines. The baseline methods of dataset include the GTO-based InfGCN (Cheng
& Peng, 2023), voxel based model CNN (Çiçek et al., 2016), interpolation-based
models, i.e., DeepDFT (Jørgensen & Bhowmik, 2020), DeepDFT2 (Jørgensen &
Bhowmik, 2022), EGNN (Satorras et al., 2021), DimeNet (Gasteiger et al.,
2020b), DimeNet++ (Gasteiger et al., 2020a), and neural operator based models,
i.e., GNO (Li et al., 2020b), FNO (Li et al., 2020a), and LNO (Lu et al.,
2019). We faithfully compare with the official numbers reported by the prior
work (Cheng & Peng, 2023) when possible. We also outperform on another
dataset, which we have not detailed here. We provide a comprehensive
comparison in Table 8 of Appendix F.

Evaluation metric. For evaluation, we follow prior works (Jørgensen & Bhowmik,
2020; Cheng & Peng, 2023) in using the normalized mean absolute error (NMAE):

$\operatorname{NMAE}(\rho,\rho_{\text{pred}})=\frac{\int_{\mathbb{R}^{3}}\|\rho({\bm{x}})-\rho_{\text{pred}}({\bm{x}})\|\text{d}{\bm{x}}}{\int_{\mathbb{R}^{3}}\|\rho({\bm{x}})\|\text{d}{\bm{x}}}.$
(18)

We provide more details regarding experiments setup, baselines, and
evaluations in Appendix D.

###  5.2 Aperiodic Materials

Dataset. For aperiodic molecules, following the prior work (Cheng & Peng,
2023), we consider electron density prediction for the QM9 dataset (Ruddigkeit
et al., 2012; Ramakrishnan et al., 2014) and the molecular dynamics (MD)
dataset (Bogojeski et al., 2018). Specifically, the electron density dataset
of the QM9 was provided by Jørgensen & Bhowmik (2022), and the MD was provided
by Brockherde et al. (2017); Bogojeski et al. (2020). The MD dataset is
further categorized into six smaller datasets corresponding to six types of
molecules, i.e., ethanol, benzene, phenol, resorcinol, ethane, and
malonaldehyde (MDA). Each smaller dataset contains snapshots of different
geometries of a molecule. More details regarding the train and dataset is in
the Appendix D.

Quantitative results. We present our experimental results in Table 1. The
results demonstrate that our algorithm significantly outperforms existing
baselines. In particular, in the QM9 and MD datasets, our approach
substantially improves over the second-best model, i.e., InfGCN, by
approximately reducing the NMAE by 20% in QM9 and 50% in MD. One can also
observe that FNO performs poorly despite its similarity to our GPWNO in using
Fourier transformation as an intermediate layer. We hypothesize that the FNO
is unable to express the high-frequency details in the true electron density
due to inefficient memory usage.

We also depict the tradeoff between the number of parameters and performance
for the QM9 dataset in Figure 4, where one can observe how GPWNO outperforms
the considered baselines despite using the smallest number of parameters. The
corresponding numbers are reported in Appendix F.

Qualitative result. We additionally visualize the predicted densities for the
QM9 dataset in Figure 3. Here, the performance of InfGCN is worse than GPWNO.

Table 2: Evaluation of GPWNO for periodic materials. NMAE (%) in the MP dataset, categorized by the seven crystal family and their combinations. The best number is highlighted in bold. Each number is averaged over three runs. Model | Mixed | Triclinic | Monoclinic | Orthorhombic | Tetragonal | Trigonal | Hexagonal | Cubic
---|---|---|---|---|---|---|---|---
DeepDFT | 11.50 | 32.33 | 50.54 | 22.30 | 37.68 | 23.89 | 18.61 | 27.30
DeepDFT2 | 15.11 | 30.20 | 49.74 | 20.50 | 41.21 | 22.13 | 12.87 | 27.63
InfGCN |  05.11 |  04.87 |  05.32 |  05.29 |  05.12 |  04.99 |  04.90 |  05.18
GPWNO |  04.83 |  04.69 |  05.20 |  05.14 |  04.81 |  04.86 |  04.77 |  04.89
Figure 4: Comparison between the baselines on QM9 (lower left is better). Table 3: Ablation study on MD and MP. Performance in NMAE (%) for the models using one or more of our algorithmic components, i.e., plane wave (PW), Gaussian-type orbital (GTO), and high-frequency masking (HFM), are reported. The best number is highlighted in bold. Each number is averaged over three runs. Model | MD | MP
---|---|---
PW | GTO | HFM | Benzene | Phenol | Resorcinol | Mixed | Triclinic | Cubic
✓ | ✗ | ✗ | 6.92 | 8.27 | 8.12 | 48.69 | 49.20 | 48.12
✗ | ✓ | ✗ | 2.90 | 3.24 | 2.96 |  05.11 |  04.87 |  05.18
✓ | ✓ | ✗ | 2.62 | 2.83 | 2.77 |  04.81 |  04.88 |  05.11
✓ | ✓ | ✓ | 2.45 | 2.68 | 2.73 |  04.83 |  04.69 |  04.89

###  5.3 Periodic Systems

Dataset. To evaluate our model on periodic systems with diverse lattice
structures, we newly conduct a benchmark from the materials project (Jain et
al., 2013; Shen et al., 2021, MP). Precisely, we categorize the electron
density for 117,535 molecules provided by Shen et al. (2021) into seven
crystal families: triclinic, monoclinic, orthorhombic, tetragonal, trigonal,
hexagonal, and cubic. Additionally, we conduct an extra dataset as a mixture
of these crystal families. The size of the datasets range from 7,681 to
26,081. Depending on the dataset size, we use 1,000 or 500 samples for the
test and validation, respectively. For this dataset, we compare with the
baseline results obtained from running the official repository
codes333https://github.com/peterbjorgensen/DeepDFT on the newly curated
datasets.

Quantitative results. We report the experimental results in Table 2. Here, one
can observe how our algorithm significantly outperforms the baselines,
regardless of the crystal family. We also note that DeepDFT and DeepDFT2
perform poorly for the dataset, despite using the official repository.

Qualitative results. We provide visualization of the predicted densities on
the MP dataset in Figure 3. Here, one can observe how the overall error and
near atom region error both are reduced on GPWNO compared to InfGCN.

###  5.4 Ablation Study

We conduct extensive ablation studies on the MD and MP to verify the
effectiveness of our method. We provide the results in Table 3 and Figure 5.
In Table 3, we report the performance of our framework with and without each
algorithmic component. In Figure 5, we check whether if introducing high-
frequency components to the PW-only prediction eliminates the necessity of GTO
basis. The detailed number of Figure 5 is reported in Table 9 of Appendix F.

Independently evaluating PW and GTO basis. In Table 3, one can observe how PW-
and GTO-based predictions perform without additional algorithmic components.
Intriguingly, one can observe how the PW performs poorly without the GTO
basis. We hypothesize this to be since the PW basis cannot capture high
frequency components in the true signal. This is additionally supported by
Figure 5, where PW demonstrates reasonable performance for the MD dataset when
increasing the size of the PW basis.

Figure 5: Increasing number of probe nodes $M$ for PW-only prediction and
comparing with GPWNO.

Synergy between PW and GTO basis. One can verify the complementary behavior of
PW and GTO basis from Table 3 since combining the two predictions consistently
improve the performance. This verifies our idea that employing a dual
representation is more effective than relying on a single basis alone,
demonstrating the synergistic advantages of utilizing both representations.

Effectiveness of HFM. When comparing our algorithm with and without HFM in
Table 3, one can verify that HFM indeed enables the GTO-based prediction to
concentrate on regions near atoms and PW-based prediction focus on regions far
from atoms. Notably, HFM achieves this improvement in efficiency without
introducing any additional model parameters.

##  6 Conclusion

In this work, we propose a new machine learning algorithm for electron density
estimation. Through extensive experiments, we demonstrated the effectiveness
and generalizability of our model. While our experiments clearly demonstrate
the superiority our GPWNO over the baselines, the evaluation metric is not
fully aligned with our potential applications, e.g., errors in predicting the
exchange-correlation energy. Since DFT computations required for this
evaluation are expensive and rely on commercial licenses, we leave this for
future work.

## Broader Impact

Our framework can advance drug design and material discovery by accelerating
molecular simulation. Such improvements often improve the quality of human
life. However, they should also be approached with caution due to the
potential misuse in the development of hazardous products.

## Acknowledgement

We thank Seungjin Kang, Suhwan Song, Byunggook Na, Yongdeok Kim, Minjong Lee,
Seonghyun Park, Hyosoon Jang, Juwon Hwang, Kiyoung Seong, and Dongyeop Woo for
providing helpful feedbacks and suggestions in preparing the early version of
the manuscript.

## References

  * Adekoya et al. (2022) Adekoya, O. C., Adekoya, G. J., Sadiku, E. R., Hamam, Y., and Ray, S. S.  Application of dft calculations in designing polymer-based drug delivery systems: An overview.  _Pharmaceutics_ , 14(9):1972, 2022.
  * Bachelet et al. (1982) Bachelet, G. B., Hamann, D. R., and Schlüter, M.  Pseudopotentials that work: From h to pu.  _Physical Review B_ , 26(8):4199, 1982.
  * Blöchl (1995) Blöchl, P.  Electrostatic decoupling of periodic images of plane-wave-expanded densities and derived atomic point charges.  _The Journal of chemical physics_ , 103(17):7422–7428, 1995.
  * Blöchl (1994) Blöchl, P. E.  Projector augmented-wave method.  _Physical review B_ , 50(24):17953, 1994.
  * Bogojeski et al. (2018) Bogojeski, M., Brockherde, F., Vogt-Maranto, L., Li, L., Tuckerman, M. E., Burke, K., and Müller, K.-R.  Efficient prediction of 3d electron densities using machine learning.  _arXiv preprint arXiv:1811.06255_ , 2018.
  * Bogojeski et al. (2020) Bogojeski, M., Vogt-Maranto, L., Tuckerman, M. E., Müller, K.-R., and Burke, K.  Quantum chemical accuracy from density functional approximations via machine learning.  _Nature communications_ , 11(1):5223, 2020.
  * Brockherde et al. (2017) Brockherde, F., Vogt, L., Li, L., Tuckerman, M. E., Burke, K., and Müller, K.-R.  Bypassing the kohn-sham equations with machine learning.  _Nature communications_ , 8(1):872, 2017.
  * Chang et al. (2021) Chang, J., Jiang, L., Wang, G., Zhao, W., Huang, Y., and Chen, H.  Lead-free perovskite compounds cssn 1- x ge x i 3- y br y explored for superior visible-light absorption.  _Physical Chemistry Chemical Physics_ , 23(26):14449–14456, 2021.
  * Cheng & Peng (2023) Cheng, C. and Peng, J.  Equivariant neural operator learning with graphon convolution.  _arXiv preprint arXiv:2311.10908_ , 2023.
  * Çiçek et al. (2016) Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., and Ronneberger, O.  3d u-net: learning dense volumetric segmentation from sparse annotation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19_ , pp. 424–432. Springer, 2016.
  * Fabrizio et al. (2019) Fabrizio, A., Grisafi, A., Meyer, B., Ceriotti, M., and Corminboeuf, C.  Electron density learning of non-covalent systems.  _Chemical science_ , 10(41):9424–9432, 2019.
  * Gasteiger et al. (2020a) Gasteiger, J., Giri, S., Margraf, J. T., and Günnemann, S.  Fast and uncertainty-aware directional message passing for non-equilibrium molecules.  _arXiv preprint arXiv:2011.14115_ , 2020a.
  * Gasteiger et al. (2020b) Gasteiger, J., Groß, J., and Günnemann, S.  Directional message passing for molecular graphs.  _arXiv preprint arXiv:2003.03123_ , 2020b.
  * Geiger & Smidt (2022) Geiger, M. and Smidt, T.  e3nn: Euclidean neural networks, 2022.
  * Giannozzi et al. (2017) Giannozzi, P., Andreussi, O., Brumme, T., Bunau, O., Nardelli, M. B., Calandra, M., Car, R., Cavazzoni, C., Ceresoli, D., Cococcioni, M., et al.  Advanced capabilities for materials modelling with quantum espresso.  _Journal of physics: Condensed matter_ , 29(46):465901, 2017.
  * Gill (1994) Gill, P. M.  Molecular integrals over gaussian basis functions.  In _Advances in quantum chemistry_ , volume 25, pp. 141–205. Elsevier, 1994.
  * Gonze et al. (2009) Gonze, X., Amadon, B., Anglade, P.-M., Beuken, J.-M., Bottin, F., Boulanger, P., Bruneval, F., Caliste, D., Caracas, R., Côté, M., et al.  Abinit: First-principles approach to material and nanosystem properties.  _Computer Physics Communications_ , 180(12):2582–2615, 2009.
  * Grisafi et al. (2018) Grisafi, A., Fabrizio, A., Meyer, B., Wilkins, D. M., Corminboeuf, C., and Ceriotti, M.  Transferable machine-learning model of the electron density.  _ACS central science_ , 5(1):57–64, 2018.
  * Hafner (2008) Hafner, J.  Ab-initio simulations of materials using vasp: Density-functional theory and beyond.  _Journal of computational chemistry_ , 29(13):2044–2078, 2008.
  * Harrison (1966) Harrison, W. A.  Pseudopotentials in the theory of metals.  _(No Title)_ , 1966.
  * Hehre et al. (1972) Hehre, W. J., Ditchfield, R., and Pople, J. A.  Self—consistent molecular orbital methods. xii. further extensions of gaussian—type basis sets for use in molecular orbital studies of organic molecules.  _The Journal of Chemical Physics_ , 56(5):2257–2261, 1972.
  * Hendrycks & Gimpel (2016) Hendrycks, D. and Gimpel, K.  Gaussian error linear units (gelus).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Jain et al. (2013) Jain, A., Ong, S. P., Hautier, G., Chen, W., Richards, W. D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G., et al.  Commentary: The materials project: A materials genome approach to accelerating materials innovation.  _APL materials_ , 1(1), 2013.
  * Jørgensen & Bhowmik (2020) Jørgensen, P. B. and Bhowmik, A.  Deepdft: Neural message passing network for accurate charge density prediction.  _arXiv preprint arXiv:2011.03346_ , 2020.
  * Jørgensen & Bhowmik (2022) Jørgensen, P. B. and Bhowmik, A.  Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids.  _npj Computational Materials_ , 8(1):183, 2022.
  * Jørgensen & Bhowmik (2022) Jørgensen, P. B. and Bhowmik, A.  Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids.  _npj Computational Materials_ , 8(1), 2022.  doi: 10.1038/s41524-022-00863-y.
  * Kittel (2005) Kittel, C.  _Introduction to solid state physics_.  John Wiley & sons, inc, 2005.
  * Koelling & Arbman (1975) Koelling, D. and Arbman, G.  Use of energy derivative of the radial solution in an augmented plane wave method: application to copper.  _Journal of Physics F: Metal Physics_ , 5(11):2041, 1975.
  * Kohn & Sham (1965) Kohn, W. and Sham, L. J.  Self-consistent equations including exchange and correlation effects.  _Physical review_ , 140(4A):A1133, 1965.
  * Kondrakov et al. (2017) Kondrakov, A. O., Geßwein, H., Galdina, K., De Biasi, L., Meded, V., Filatova, E. O., Schumacher, G., Wenzel, W., Hartmann, P., Brezesinski, T., et al.  Charge-transfer-induced lattice collapse in ni-rich ncm cathode materials during delithiation.  _The journal of physical chemistry C_ , 121(44):24381–24388, 2017.
  * Kosmala et al. (2023) Kosmala, A., Gasteiger, J., Gao, N., and Günnemann, S.  Ewald-based long-range message passing for molecular graphs, 2023.
  * Kovachki et al. (2023) Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhattacharya, K., Stuart, A., and Anandkumar, A.  Neural operator: Learning maps between function spaces, 2023.
  * Landau & Binder (2021) Landau, D. and Binder, K.  _A guide to Monte Carlo simulations in statistical physics_.  Cambridge university press, 2021.
  * Li et al. (2020a) Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, A.  Fourier neural operator for parametric partial differential equations.  _arXiv preprint arXiv:2010.08895_ , 2020a.
  * Li et al. (2020b) Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, A.  Neural operator: Graph kernel network for partial differential equations.  _arXiv preprint arXiv:2003.03485_ , 2020b.
  * Li et al. (2023) Li, Z., Kovachki, N. B., Choy, C., Li, B., Kossaifi, J., Otta, S. P., Nabian, M. A., Stadler, M., Hundt, C., Azizzadenesheli, K., et al.  Geometry-informed neural operator for large-scale 3d pdes.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Lippert et al. (1997) Lippert, B. G., PARRINELLO, J. H., and MICHELE.  A hybrid gaussian and plane wave density functional scheme.  _Molecular Physics_ , 92(3):477–488, 1997.
  * Lu et al. (2019) Lu, L., Jin, P., and Karniadakis, G. E.  Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators.  _arXiv preprint arXiv:1910.03193_ , 2019.
  * Martin (2004) Martin, R. M.  _Electronic Structure: Basic Theory and Practical Methods_.  Cambridge University Press, 2004.  doi: 10.1017/CBO9780511805769.
  * Ramakrishnan et al. (2014) Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific data_ , 1(1):1–7, 2014.
  * Ruddigkeit et al. (2012) Ruddigkeit, L., Van Deursen, R., Blum, L. C., and Reymond, J.-L.  Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17.  _Journal of chemical information and modeling_ , 52(11):2864–2875, 2012.
  * Satorras et al. (2021) Satorras, V. G., Hoogeboom, E., and Welling, M.  E (n) equivariant graph neural networks.  In _International conference on machine learning_ , pp. 9323–9332. PMLR, 2021.
  * Schütt et al. (2021) Schütt, K., Unke, O., and Gastegger, M.  Equivariant message passing for the prediction of tensorial properties and molecular spectra.  In _International Conference on Machine Learning_ , pp. 9377–9388. PMLR, 2021.
  * Shen et al. (2021) Shen, J.-X., Munro, J. M., Horton, M. K., Huck, P., Dwaraknath, S., and Persson, K. A.  A representation-independent electronic charge density database for crystalline materials, 2021.
  * Sinitskiy & Pande (2018) Sinitskiy, A. V. and Pande, V. S.  Deep neural network computes electron densities and energies of a large set of organic molecules faster than density functional theory (dft).  _arXiv preprint arXiv:1809.02723_ , 2018.
  * Sun et al. (2023) Sun, H., Ross, Z. E., Zhu, W., and Azizzadenesheli, K.  Phase neural operator for multi-station picking of seismic arrivals.  _Geophysical Research Letters_ , 50(24):e2023GL106434, 2023.
  * Sun et al. (2017) Sun, Q., Berkelbach, T. C., McClain, J. D., and Chan, G. K.  Gaussian and plane-wave mixed density fitting for periodic systems.  _The Journal of chemical physics_ , 147(16), 2017.
  * Thomas et al. (2018) Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P.  Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds.  _arXiv preprint arXiv:1802.08219_ , 2018.
  * Tran et al. (2023) Tran, A., Mathews, A., Xie, L., and Ong, C. S.  Factorized fourier neural operators.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=tmIiMPl4IPa.
  * Ulian et al. (2013) Ulian, G., Tosoni, S., and Valdrè, G.  Comparison between gaussian-type orbitals and plane wave ab initio density functional theory modeling of layer silicates: Talc [mg3si4o10 (oh) 2] as model system.  _The Journal of chemical physics_ , 139(20), 2013.
  * Wang et al. (2022) Wang, F. Q., Choudhary, K., Liu, Y., Hu, J., and Hu, M.  Large scale dataset of real space electronic charge density of cubic inorganic materials from density functional theory (dft) calculations.  _Scientific Data_ , 9(1):59, 2022.
  * Xie & Grossman (2018) Xie, T. and Grossman, J. C.  Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties.  _Phys. Rev. Lett._ , 120:145301, Apr 2018.  doi: 10.1103/PhysRevLett.120.145301.  URL https://link.aps.org/doi/10.1103/PhysRevLett.120.145301.

##  Appendix A Overall Architecture

In Figure 6, we provide a detailed description of our model.

Figure 6: Architecture of Gaussian plane-wave neural operator (GPWNO).
Starting from molecule ${\mathcal{M}}=({\mathcal{X}},{\mathcal{A}},{\bm{L}})$
with the atomic positions ${\mathcal{X}}$, atomic features ${\mathcal{A}}$ and
the lattice vectors ${\bm{L}}$, GPWNO aims to predict a scalar field
$\rho_{\text{pred}}$ for given query points
${\mathcal{Q}}\subset\mathbb{R}^{3}$. Notably, we decompose the target
$\rho_{\text{pred}}$ into two types of the bases: Gaussian-type orbital (GTO)
basis and Plane wave (PW) basis. (top) GTO basis layer predict the
$\rho_{\text{GTO}}$ by the direct coefficients learning of spherical
harmonics. (bottom) PW basis layer predict the signals of reciprocal space by
using discretized voxel points ${\mathcal{P}}$ and the FNO layers, and predict
the $\rho_{\text{PW}}$ the message aggregation from voxel points
${\mathcal{P}}$.

##  Appendix B Signal Initialization by Tensor Field Network

Here, we describe our TFN-based architecture for initializing a signal as an
input to our PW neural operator. For a given molecule
${\mathcal{M}}=({\mathcal{X}},{\mathcal{A}},{\bm{L}})$, we first initialize
the atomic feature ${\bm{f}}_{u}^{(0)}$ of atom $u\in{\mathcal{V}}$ by the
linear-embedding layer that takes input as the one-hot encoding of the atomic
number of each atom. Then, we update the atomic feature with the TFN layer by
treating ${\bm{f}}_{u}^{(0)}\in\mathbb{R}^{C}$ as a multi-channel rank-0
tensor feature of dimension $C$. i.e.,
${\bm{f}}_{u}^{(0)}=\\{f^{1}_{u,00},\ldots,f^{C}_{u,00}\\}$. Specifically,
each embedding layer updates the spherical node feature
${{\bm{f}}}_{u,l}^{(h)}=\\{\hat{f}_{u,lm}^{(h)}\mathrel{\mathop{\mathchar
58\relax}}m\in\mathbb{Z},-l\leq m\leq l\\}$ from the near atoms spherical node
feature ${{\bm{f}}}_{v,l}^{(h)}$ by tensor product:

$\displaystyle\hat{\mathbf{f}}^{(h+1)}_{u,l}$
$\displaystyle=\sum_{v\in\mathcal{N}(u)}\sum_{k\geq
0}W^{(h)}_{lk}({\bm{x}}_{v}-{\bm{x}}_{u})\hat{\mathbf{f}}^{(h)}_{v,k},$ (19)
$\displaystyle W^{(h)}_{lk}({\bm{r}})$
$\displaystyle=\sum^{k+1}_{J=\|k-l\|}\varphi_{lk}^{J}(\lVert{\bm{r}}\rVert)\sum^{J}_{m=-J}Y_{Jm}\left(\frac{{\bm{r}}}{\lVert{\bm{r}}\rVert}\right)Q^{lk}_{Jm}.$
(20)

Here, $\mathcal{N}(u)$ is the set of neighbors of atom $u$,
$Q^{lk}_{Jm}\in\mathbb{R}^{(2l+2)\times(2k+1)}$ is the Clebsh-Gordan matrix,
and $\varphi_{nlk}^{J}\mathrel{\mathop{\mathchar
58\relax}}\mathbb{R}^{+}\rightarrow\mathbb{R}$ is a learnable scalar function.
At the final embedding layer $H$, we set the output feature of the TFN as the
multi-channel rank-0 tensor with dimension $K$. i.e.,
${\bm{f}}_{u}^{(H)}\in\mathbb{R}^{K}$. Finally, we use this
${\bm{f}}_{u}^{(H)}$ as the input of the Equation 10.

##  Appendix C Proof of Equivariance for the Lattice Construction on Aperiodic
Molecule

We define the lattice $\bm{L}$ of the aperiodic molecule based on the
principal component analysis (PCA). For a given molecule ${\mathcal{M}}$, let
$n$ is the number of the atoms and ${\mathbf{X}}$ is the atom positon matrix.
i.e.,
${\mathbf{X}}=[{\bm{x}}_{1},\ldots,{\bm{x}}_{n}]^{\top}\in\mathbb{R}^{n\times
3}$. Then, let
${\bm{t}}=\frac{1}{n}\sum_{u\in{\mathcal{V}}}{\bm{x}}_{u}\in\mathbb{R}^{3}$ be
the centroid of the molecule ${\mathbf{X}}$ with $n$ atoms and
${\mathbf{C}}=({\mathbf{X}}-\mathbb{1}{\bm{t}}^{\top})^{\top}({\mathbf{X}}-\mathbb{1}^{\top})\in\mathbb{R}^{3\times
3}$ be the covariance matrix computed after removing the centroid from
${\mathbf{X}}$. In the generic case, the eigenvalues of ${\mathbf{C}}$ satisfy
$\lambda_{1}<\lambda_{2}<\lambda_{3}$. Let
${\bm{v}}_{1},{\bm{v}}_{2},{\bm{v}}_{3}$ be the unit length corresponding
eigenvectors. Then we define $\bm{L}=[{\bm{v}}_{1},{\bm{v}}_{2},{\bm{v}}_{3}]$
Note that eigenvectors are orthogonal to each other, so $\bm{L}$ can be used
to the cubical lattice vector.

Next, we will show this $\bm{L}$ is SE(3) equivariant, meaning that the
rotation of the data naturally transforms the lattice vector, i.e,
$\bm{L}^{\prime}=[{\bm{v}}_{1}{\mathbf{R}}^{\top},{\bm{v}}_{2}{\mathbf{R}}^{\top},{\bm{v}}_{3}{\mathbf{R}}^{\top}]$.
Note that when we think about the direction of PCA eigenvector points out the
maximal variance axis, it is natural that eigenvectors rotatated equivalently
to the frame.

Let $({\mathbf{R}},{\bm{k}})\in$ SO(3); then the position of molecule are
transformed to
${\mathbf{X}}^{\prime}={\mathbf{X}}{\mathbf{R}}+\mathbb{1}{\bm{k}}^{\top}$.
The centroid of the molecule is
${\bm{t}}^{\prime}={\mathbf{R}}^{\top}{\bm{t}}+{\bm{k}}$, and the covariance
matrix becomes

$\displaystyle{\mathbf{C}}^{\prime}$
$\displaystyle=({\mathbf{X}}^{\prime}-\mathbb{1}{\bm{t}}^{\prime\top})^{\top}({\mathbf{X}}^{\prime}-\mathbb{1}{\bm{t}}^{\prime\top})$
$\displaystyle=({\mathbf{X}}{\mathbf{R}}+\mathbb{1}{\bm{k}}^{\top}-\mathbb{1}{\bm{t}}^{\top}{\mathbf{R}}-\mathbb{1}{\bm{k}}^{\top})^{\top}({\mathbf{X}}{\mathbf{R}}+\mathbb{1}{\bm{k}}^{\top}-\mathbb{1}{\bm{t}}^{\top}{\mathbf{R}}+\mathbb{1}{\bm{k}}^{\top})$
$\displaystyle={\mathbf{R}}^{\top}({\mathbf{X}}-\mathbb{1}{\bm{t}}^{\top})^{\top}({\mathbf{X}}-\mathbb{1}{\bm{t}}^{\top}){\mathbf{R}}$
$\displaystyle={\mathbf{R}}^{\top}{\mathbf{C}}{\mathbf{R}}$

By the definition of the eigen decomposition, the eigenvector ${\bm{v}}$ of
covaraince matrix ${\mathbf{C}}$ with eigenvalue $\lambda$ satisfies
${\mathbf{C}}{\bm{v}}=\lambda{\bm{v}}$. Similary, the transformed covariance
matrix ${\mathbf{C}}^{\prime}$ has the new eigenvector
${\bm{v}}^{\prime}={\mathbf{R}}^{\top}{\bm{v}}$. We show that
${\bm{v}}^{\prime}$ is the eigenvector of ${\mathbf{C}}^{\prime}$ and its
eigenvalue is $\lambda$:

$\displaystyle{\mathbf{C}}^{\prime}{\bm{v}}^{\prime}$
$\displaystyle=({\mathbf{R}}^{\top}{\mathbf{C}}{\mathbf{R}})({\mathbf{R}}^{\top}{\bm{v}})$
$\displaystyle={\mathbf{R}}^{\top}{\mathbf{C}}{\bm{v}}={\mathbf{R}}^{\top}\lambda{\bm{v}}$
$\displaystyle=\lambda{\bm{v}}^{\prime}$

To this end, the lattice system of the rotated system is naturally rotated,
that is
$\bm{L}^{\prime}=[{\bm{v}}^{\prime}_{1},{\bm{v}}^{\prime}_{2},{\bm{v}}^{\prime}_{3}]=[{\bm{v}}_{1}{\mathbf{R}}^{\top},{\bm{v}}_{2}{\mathbf{R}}^{\top},{\bm{v}}_{3}{\mathbf{R}}^{\top}]$
as mentioned above.

##  Appendix D Dataset and Hyperparameters

In this section, we provide more details about our dataset and train
configuration. All the dataset contains the molecule graph of each species and
its voxelized 3D electron charge density data on its pre-defined grid.

QM9. The dataset is provided by Jørgensen & Bhowmik (2022, 2020) The electron
density of the dataset are calculated by VASP. To be specific, VASP with
Perdew-Burke-Ernzerhof exchange correlation (PBE XC) functional and projector-
augemented wave (PAW) is used.

MD. The dataset is provided by Brockherde et al. (2017); Bogojeski et al.
(2018, 2020) and it is calculated by PBE XC functional with Quantum ESPRESSO
using the PBE functional.

Cubic2. Cubic2 is not directly reported, since it shares the same materials
structure with the cubic of the Materials Project. We report this to point out
our model outperformed than previous works on this dataset, too. The data is
provided by (Wang et al., 2022) and calculated by the VASP with the PAW. The
performance on Cubic2 is in Table 8.

Materials project. This data is set of the periodic materials including cubic,
hexagonal and so on. The detailed categorized condition is in Table 6.

Common statistics of the QM9, Cubic2 and MD are summarized in Table 4. MD
dataset does not have a validation split, and the number of grids and grid
lengths is for a single dimension so the number of voxels cubically with
respect to it.

Table 4: Statistics of the dataset. We round up the number of grids and nodes to the nearest integer for readability. Dataset | QM9 | Cubic2 | MD
---|---|---|---
train/val/test split | 123835/50/1600 | 14421/1000/1000 | 1000(2000)/500(400)
min/mean/max #grid | 40/88/160 | 32/94/448 | 20/20/20
min/mean/max #node | 3/18/29 | 1/10/64 | 8/11/14
min/mean/max length (Bohr) | 4.00/8.65/15.83 | 1.78/5.82/26.20 | 20/20/20
# node type | 5 | 84 | 3

Additional statistics of the MP dataset are detailed in Table 5. The mixed
dataset is composed of samples from each crystal family, with a distribution
of 2000 for training, 150 for validation, and 150 for testing, resulting in a
total of 14,000 training, 1,050 validation, and 1,050 testing samples. Note
that the MP dataset may include non-cubical voxels, as its lattice structure
is not confined to cubic forms. Differing from cubic structures, the
generalized lattice types vary in length and grid dimensions along the
${\bm{a}}_{1},{\bm{a}}_{2},$ and ${\bm{a}}_{3}$ directions of the lattice
vector. Each number of the grid binning and lattice length are denoted by ”#
grid ${\bm{a}}_{i}$” and ”length ${\bm{a}}_{i}$”, respectivly. Accordingly, we
provide statistics for each side, as well as the average binning of grid and
length, which are represented as ”# grid avg”, and ”length avg”. These
averages are calculated using the cube root of the product of the values along
${\bm{a}}_{1}$, ${\bm{a}}_{2}$, and ${\bm{a}}_{3}$.

Table 5: Dataset details of materials project. We round up the number of grids
to the nearest integer for readbility.

| Triclinic | Monoclinic | Orthorhombic | Tetragonal | Trigonal | Hexagonal | Cubic
---|---|---|---|---|---|---|---
train/val/test split | 12843/1000/1000 | 26081/1000/1000 | 22115/1000/1000 | 12031/1000/1000 | 9651/500/500 | 7681/500/500 | 15133/1000/1000
# atom type | 85 | 85 | 87 | 87 | 86 | 85 | 87
# grid ${\bm{a}}_{1}$ | 32/102/576 | 32/102/972 | 24/102/1120 | 28/102/576 | 24/102/480 | 36/102/540 | 24/103/480
# grid ${\bm{a}}_{2}$ | 36/108/384 | 32/108/640 | 28/108/1120 | 28/108/500 | 24/108/480 | 36/108/756 | 20/109/480
# grid ${\bm{a}}_{3}$ | 16/143/1344 | 32/143/1512 | 28/143/1372 | 36/143/1500 | 36/142/1344 | 36/143/1500 | 24/143/1372
# grid avg. | 29/112/360 | 32/112/500 | 28/112/540 | 35/112/480 | 36/112/480 | 40/112/448 | 31/112/480
Length ${\bm{a}}_{1}$ | 2.05/6.63/38.47 | 2.17/6.67/65.40 | 1.52/6.65/74.23 | 1.71/6.65/38.78 | 1.63/6.68/31.52 | 2.28/6.67/34.83 | 1.61/6.70/31.54
Length ${\bm{a}}_{2}$ | 2.38/7.03/25.27 | 2.17/7.05/42.85 | 1.77/7.02/74.88 | 1.71/7.04/32.61 | 1.63/7.06/31.52 | 2.40/7.05/50.18 | 1.28/7.09/31.54
Length ${\bm{a}}_{3}$ | 0.98/9.34/88.15 | 2.13/9.35/101.38 | 1.77/9.39/90.44 | 2.34/9.38/100.46 | 2.45/9.28/88.13 | 2.48/9.35/98.70 | 1.54/9.38/91.10
Length avg. | 2.07/7.67/31.95 | 2.17/7.69/35.94 | 1.77/7.68/41.50 | 2.37/7.69/35.59 | 2.45/7.67/31.94 | 2.52/7.69/34.96 | 2.08/7.72/32.42
# node | 1/23.90/144 | 1/24/144 | 1/23.74/136 | 1/23.99/154 | 1/23.71/144 | 1/23.79/136 | 1/24.36/144

Table 6: The condition of crystal familiy. Here, each $a,b,c$ represents the length of the corresponding lattice vector, i.e., $a=|{\bm{a}}_{1}|$, $b=|{\bm{a}}_{2}|$, and $c=|{\bm{a}}_{3}|$. Additionally, $\alpha$, $\beta$, and $\gamma$ denote the angles between ${\bm{a}}_{2}$ and ${\bm{a}}_{3}$, ${\bm{a}}_{1}$ and ${\bm{a}}_{3}$, and ${\bm{a}}_{2}$ and ${\bm{a}}_{1}$, respectively.  Crystal Famlily | Edge | Angle
---|---|---
Triclinic | $a\neq b\neq c$ | $\alpha\neq\beta\neq\gamma\neq 90^{\circ}$
Monoclinic | $a\neq b\neq c$ | $\alpha=\gamma=90^{\circ},\beta\neq 120^{\circ}$
Orthorhombic | $a\neq b\neq c$ | $\alpha=\beta=\gamma=90^{\circ}$
Tetragonal | $a=b\neq c$ | $\alpha=\beta=\gamma=90^{\circ}$
Trigonal | $a=b=c$ | $\alpha=\beta=\gamma\neq 90^{\circ}$
Hexagonal | $a=b\neq c$ | $\alpha=\beta=90^{\circ},\gamma=120^{\circ}$
Cubic | $a=b=c$ | $\alpha=\beta=\gamma=90^{\circ}$

##  Appendix E Model and Training Details

###  E.1 Model Details

We provide more details about the proposed GWPNO model and the baseline.

GPWNO.444We plan to release our code upon acceptance. Our implementation of
GTO layer is inspired by the InfGCN555https://github.com/ccr-cheng/infgcn-
pytorch. We used the spherical degree $l=3$ or $l=4$. The number of radio
bases is settled $n=16$ with the Gaussian parameters starting at 0.5 (Bohr)
and ending at 5.0. We used 3 layers of the GTO layers and the number of the
trainable parameter of GTO layer is 223K when $l=3$ and 353K when $l=4$.

In the implementation of the PW layer, we set the spherical degree $l=3$ for
embedding purposes. The number of probe nodes $M$ was configured to either 20
or 40. To enhance the scalability of the PW layer, we incorporated the
approach outlined by Tran et al. (2023). However, our empirical observations
indicate that the performance of the model remains nearly unchanged with or
without the use of Tran et al. (2023), provided that the number of probe nodes
remains consistent. Regarding the trainable parameters in the PW neural
operator layer, the count stands at 19.3K for $M=20$ and increases to 47.0K
when $M=40$.

InfGCN. InfGCN is a GTO-based model designed to learn the direct coefficients
of the GTO basis. We used the official code implementation of InfGCN by Cheng
& Peng (2023).

DeepDFT. and DeepDFT2. DeepDFT is a GNN based model that uses message passing
to the query node where the charge density is predicted. DeepDFT2 uses PaiNN
(Schütt et al., 2021) as the GNN architecture. We used the official code
implementation of DeepDFT and DeepDFT2 by Jørgensen & Bhowmik (2020, 2022).

###  E.2 Training Details

In this section, we outline our model specifications and those of the baseline
models. Additionally, we detail the hyperparameters related to training and
testing utilized in our experiments. As previously mentioned, we make direct
comparisons with the official results reported in prior work by Cheng & Peng
(2023) whenever possible.

Table 7: Hyperaparameter details

Model | Dataset | sp | cutoff | q.cutoff | p.cutoff | $M$ | num.grad | learning rate | batch size | learning rate decay | query points
---|---|---|---|---|---|---|---|---|---|---|---
GPWNO | QM9 | 4 | 3.0 | 1.5 | 0.75 | 40 | 80k | 1e-3 | 32 | 0.5 | 1024
MD | 3 | 3.0 | 1.5 | 0.75 | 40 | 4k | 5e-3
Cubic | 4 | 5.0 | 2.5 | 1.25 | 20 | 10k | 5e-3
MP | 3 | 3.0 | 2.5 | 1.25 | 20 | 10k | 5e-3
InfGCN | MP | 4 | 3.0 | - | - | - | 10k | 5e-3 | 32 | 0.5 | 1024
DeepDFT | MP | 4 | 4.0 | - | - | - | 10k | 3e-4 | 32 | 0.5 | 1024
DeepDFT2

We list the training specifications and major hyperparameters of our
experiments in Table 7. Following the approach of Cheng & Peng (2023) and
Jørgensen & Bhowmik (2020), we employ a sampling training scheme. In each
training iteration, we sample multiple query points from the given data batch.
For instance, we sample 1024 query points from a batch of 32 data points,
resulting in a target size of $(32\times 1024)$ query points during training.
For testing, we evaluate all samples in the test dataset to avoid randomness.
Minibatch inference does not impact the results, so we segment the full grid
into small sections, evaluate the points sequentially as mini-batch, and
compute the normalized mean absolute error (NMAE) across the entire inference
dataset.

We utilized the ADAM optimizer, and the learning rate and its decay also
reported in Table 7. In addition the number of the gradient step is also
reported as num.grad.

Key hyperparameters in our model include the spherical harmonics (sp) order in
the GTO layer and the radius cutoff distance for each message-passing
operation. The cutoff parameter defines the neighborhood radius for the GTO
layer. Additionally, the probe cutoff (p.cutoff) and query cutoff (q.cutoff)
parameters specify the radius distances used in Equation 10 and Equation 13,
respectively.

##  Appendix F Additional Results

We provide the additional results that is not included in the main context for
the space.

###  F.1 Additional Comparison

We provide the performance on another data set that is not reported in the
main context. For QM9, we also evaluate the performance on the randomly
rotated test set. We observe that our GPWNO also outperform than previous
works on the rotated QM9. For Cubic2, GPWNO also achieve the state-of-the-art
performance among the baseline.

Table 8: The comparison of NMAE (%) reported in (Cheng & Peng, 2023).

Dateset / Model | Ours | InfGCN | CNN | DeepDFT | DeepDFT2 | EGNN | DimeNet | DimeNet++ | GNO | FNO | LNO
---|---|---|---|---|---|---|---|---|---|---|---
QM9 | rot | 4.68 | 4.73 | 5.89 | 5.87 | 4.98 | 12.13 | 12.98 | 12.75 | 46.90 | 33.25 | 24.13
unrot | 0.73 | 0.93 | 2.01 | 2.95 | 1.03 | 11.92 | 11.97 | 11.69 | 40.86 | 28.83 | 26.14
MD | Ethanol | 4.00 | 8.43 | 13.97 | 7.34 | 8.83 | 13.90 | 13.99 | 14.24 | 82.35 | 31.98 | 43.17
Benzene | 2.45 | 5.11 | 11.98 | 6.61 | 5.49 | 13.49 | 14.48 | 14.34 | 82.46 | 20.05 | 38.82
Phenol | 2.68 | 5.51 | 11.52 | 9.09 | 7.00 | 13.59 | 12.93 | 12.99 | 66.69 | 42.98 | 60.70
Resorcinol | 2.73 | 5.95 | 11.07 | 8.18 | 6.95 | 12.61 | 12.04 | 12.01 | 58.75 | 26.06 | 35.07
Ethane | 3.67 | 7.01 | 14.72 | 8.31 | 6.36 | 15.17 | 13.11 | 12.95 | 71.12 | 26.31 | 77.14
Malonaldehyde | 5.32 | 10.34 | 18.52 | 9.31 | 10.68 | 12.37 | 18.71 | 16.79 | 84.52 | 34.58 | 47.22
Cubic2 | 7.69 | 8.98 | - | 14.08 | 10.37 | 11.74 | 12.51 | 12.18 | 53.55 | 48.08 | 46.33

###  F.2 Numbers for Figure 4

The numbers for the performance of Figure 4 is reported in Table 9.

Table 9: Numbers for Figure 4. # probes $M$ | Benzene | Phenol | Resorcinol
---|---|---|---
$10$ | 93.81 | 94.06 | 96.11
$20$ | 36.71 | 38.41 | 39.70
$30$ | 14.04 | 15.18 | 15.36
$40$ | 6.92 | 8.27 | 8.12
$50$ | 4.77 | 5.66 | 5.18