# Supplemental Material for
”Leveraging Normalizing Flows for Orbital-Free Density Functional Theory”

Alexandre de Camargo  [  Ricky T. Q. Chen  [  Rodrigo A. Vargas-Hernández  [
<EMAIL_ADDRESS>
###### Abstract

The purpose of this supplemental material is to provide more details about the
proposed work in the main draft. Section 1 presents an introduction to
normalizing flows, and Sections 2 and 3 describe the numerical details of the
simulations and physical systems. In Section 4, we present additional results
from the ones presented in the main text.

###### keywords:

Machine learning, normalizing flows, orbital free density functional theory

McMaster] McMaster University, Hamilton, ON, L8S 4M1, Canada meta]FAIR Meta,
NY McMaster] McMaster University, Hamilton, ON, L8S 4M1, Canada ML, OF-DFT,
DFT, NF, NN

##  1 Normalizing Flows

The central goal of our work is to introduce an alternative ansatz for
parameterizing the electron density $\operatorname{\rho_{{\cal M}}}$ using
Normalizing Flows (NFs)1 $\operatorname{\rho_{\phi}}$. For convenience, we
define $\operatorname{\rho_{{\cal M}}}$ as,

$\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}}):=\operatorname{N_{e}}\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}),$
(1)

where $\rho_{\phi}(\operatorname{\boldsymbol{x}})$ is also known as _shape
factor_ 2, 3. Eq. 1 guarantees the normalization to the number of electrons
$\operatorname{N_{e}}$, if
$\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}})$ normalizes to one;
$\int\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}}=1$.
As explained in the main text, we parametrize
$\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}})$ using a
normalizing flow 1, 4.

Normalizing flows provides a general way of constructing complex probability
distributions from simple ones, using the change of variable formula 4,

$\rho_{\phi}(\mathbf{x})=\operatorname{\rho_{0}}(\mathbf{z})\left|\text{det}\nabla_{\operatorname{\mathbf{z}}}T_{\phi}(\operatorname{\mathbf{z}})\right|^{-1},$
(2)

where $T_{\phi}(\cdot)$ is a bijective differentiable transformation and
$\operatorname{\rho_{0}}(\operatorname{\mathbf{z}})$ is the base distribution.
Eq. 2 guarantees volume preservation in the density transformation. One can
parametrize $T_{\phi}(\cdot)$ through a series of composable functions 4, 5,
1,

$T_{\phi}(\cdot)=T_{K}(\cdot)\circ\cdots\circ T_{1}(\cdot).$ (3)

These composable transformations can be seen as a flow discretized over time,
and via Eq. 3 we can induce further complex distributions.

###  1.1 Continuous Normalizing flow

An alternative formulation of (Eq. 2) is to construct a flow that operates in
the continuous domain 6, 7, assuming that the state transition is governed by
an ordinary differential equation (ODE). This alternative NF framework is
called _continuous normalizing flow_ (CNF), and it computes the
$T_{\phi}(\cdot)$, the $\log$ density, and the score function by solving a
joint ODE given by,

$\partial_{t}\begin{bmatrix}\mathbf{z}(t)\\\
\log\operatorname{\rho_{\phi}}(\mathbf{z}(t))\\\
\nabla\log\rho_{\phi}\end{bmatrix}=\begin{bmatrix}\operatorname{g_{\phi}}(\mathbf{z}(t),t)\\\
-\nabla_{\operatorname{\boldsymbol{x}}}\cdot\operatorname{g_{\phi}}(\mathbf{z}(t),t)\\\
-\nabla^{2}\operatorname{g_{\phi}}-\left(\nabla\log\operatorname{\rho_{\phi}}\right)^{T}\left(\nabla\operatorname{g_{\phi}}(\mathbf{z}(t),t)\right),\end{bmatrix},$
(4)

where “$\nabla\cdot$” denotes the divergence operator, $\nabla^{2}$ is
$\nabla\cdot\nabla$, and
$\nabla_{\operatorname{\boldsymbol{x}}}\log\rho(\operatorname{\boldsymbol{x}})$
represents the score function 7, which is relevant for computing
$\nabla_{\operatorname{\boldsymbol{x}}}\rho(\operatorname{\boldsymbol{x}})$
using the ”log-derivative trick”,

$\nabla_{\operatorname{\boldsymbol{x}}}\log\rho(\operatorname{\boldsymbol{x}})=\frac{\nabla_{\operatorname{\boldsymbol{x}}}\rho(\operatorname{\boldsymbol{x}})}{\rho(\operatorname{\boldsymbol{x}})}.$
(5)

We compute all required gradients by solving a second augmented ODE 6. This
approach applies to all ODE solvers and it’s called the adjoint sensitivity
method 8. For all results presented in this work, we used the mixed 4th/5th
order Runge-Kutta integration method9 implemented in
jax.experimental.ode.odeint.

For the one-dimensional simulations, Section 4.1, of the diatomic molecules,
the architecture of $\operatorname{g_{\phi}}$ is a standard feed-forward
neural network (NN),

$\displaystyle\operatorname{g_{\phi}}=\sum_{\ell}^{M}f_{\ell}(\mathbf{z}_{\ell}(t)),$
(6)

where $f_{\ell}(\cdot)$ is a linear layer followed by an activation function,
and $M$ is the number of layers. For this work, $\operatorname{g_{\phi}}$ has
3 layers, each with 512 neurons, and the $\tanh$ activation function. For the
simulation in three dimensions, Section 4.2, $\operatorname{g_{\phi}}$ is
parametrized by a permutation equivariant graph NN (GNN), 10, 11

$\operatorname{g_{\phi}}(\mathbf{z},t)=\sum_{i}^{N_{a}}f_{\ell}(\|\mathbf{z}(t)-\mathbf{R}_{i}\|_{2},\tilde{Z}_{i})(\mathbf{z}(t)-\mathbf{R}_{i}),$
(7)

where $\tilde{Z}_{i}$ is the atomic number of the $i^{th}$-nucleus, encoded as
a one-hot vector ($[0,\cdots,1_{i},\cdots,0]$), $N_{a}$ is the number of
nucleus in the molecule, and $f_{\ell}(\cdot)$ is a feed-forward NN with $64$
neurons per layer, also with the $\tanh$ activation function. This GNN
architecture was chosen due to its ability to handle input atom permutations
in an equivariant manner, effectively capturing the molecule’s crucial spatial
and chemical characteristics without being affected by the sequence of the
atoms.

##  2 Expectation values of density functionals

In this section, we present the general framework used to compute the value of
the total energy functional $E[\operatorname{\rho_{{\cal M}}}]$ for the one-
dimensional and three-dimensional simulations portrayed in the main text. For
any of the examples considered in this work, the total energy functional
$E[\operatorname{\rho_{{\cal M}}}]$ is given by 12, 2,

$E[\operatorname{\rho_{{\cal M}}}]=T[\operatorname{\rho_{{\cal
M}}}]+V_{\text{H}}[\operatorname{\rho_{{\cal
M}}}]+V_{\text{e-N}}[\operatorname{\rho_{{\cal
M}}}]+E_{\text{XC}}[\operatorname{\rho_{{\cal M}}}],$ (8)

where we approximate the total kinetic energy ($T[\operatorname{\rho_{{\cal
M}}}]$) by the sum of the Thomas-Fermi
($T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$) and the Weizsäcker
($T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$) functionals 2.
$V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$ is the Hartree potential that
describes the classical electron-electron repulsion,
$V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$ is the external potential,
and $E_{\text{XC}}[\operatorname{\rho_{{\cal M}}}]$ is the exchange-
correlation functional 2.

We use a Monte Carlo (MC) method to estimate the value of all individual
functionals, through the following generalization,

$\displaystyle F[\operatorname{\rho_{{\cal M}}}]$ $\displaystyle=$
$\displaystyle\int f(\operatorname{\boldsymbol{x}},\operatorname{\rho_{{\cal
M}}},\nabla\operatorname{\rho_{{\cal M}}})\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}}=(\operatorname{N_{e}})^{p}\int
f(\operatorname{\boldsymbol{x}},\operatorname{\rho_{\phi}},\nabla\operatorname{\rho_{\phi}})\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}},$
(9)

where $(\operatorname{N_{e}})^{p}$, $p\in\mathbb{R}$, is the constant factor
related to the number of electrons $\operatorname{N_{e}}$, due to our
definition of $\operatorname{\rho_{{\cal M}}}$, (Eq. 1), and
$f(\operatorname{\boldsymbol{x}},\rho_{\phi},\nabla\rho_{\phi})$ is the
integrand of the functional $F[\operatorname{\rho_{{\cal M}}}]$. The
expectation value of $F[\operatorname{\rho_{{\cal M}}}]$ is taken with respect
to our base distribution $\operatorname{\rho_{0}}$ 5,

$\displaystyle F[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=(\operatorname{N_{e}})^{p}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}[f(\operatorname{T_{\phi}}(\operatorname{\mathbf{z}}),\rho_{\phi},\nabla\rho_{\phi})]\approx(\operatorname{N_{e}})^{p}\frac{1}{N}\sum^{N}_{i=1}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[f(\operatorname{T_{\phi}}(\operatorname{\mathbf{z}}_{i}),\operatorname{\rho_{\phi}},\nabla\operatorname{\rho_{\phi}})\right],$
(10)

where $N$ is the samples drawn from $\operatorname{\rho_{0}}$ and transformed
by the CNF ($T_{\phi}$), (Eq. 4),

$\operatorname{\boldsymbol{x}}=\operatorname{T_{\phi}}(\operatorname{\mathbf{z}}):=\operatorname{\mathbf{z}}+\int_{t_{0}}^{T}\operatorname{g_{\phi}}(\operatorname{\mathbf{z}}(t),t)\mathrm{d}t.$
(11)

In the following sections, we present the analytic expressions for each
functional used in this work for the one-dimensional and three-dimensional
cases discussed in the main text.

###  2.1 One-dimensional density functionals

For the one-dimensional systems, the kinetic energy functional is the sum of
the Thomas-Fermi ($T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$)13 and the
Weizsäcker ($T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$) 2 functionals,

$\displaystyle T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$ $\displaystyle=$
$\displaystyle\frac{\pi^{2}}{24}\int\left(\operatorname{\rho_{{\cal
M}}}(x)\right)^{3}\mathrm{d}x$ (12) $\displaystyle
T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$ $\displaystyle=$
$\displaystyle\frac{\operatorname{\lambda_{0}}}{8}\int\frac{(\nabla\operatorname{\rho_{{\cal
M}}}(x))^{2}}{\operatorname{\rho_{{\cal
M}}}}\mathrm{d}x=\frac{\operatorname{\lambda_{0}}}{8}\int\left(\nabla\log\operatorname{\rho_{{\cal
M}}}(x)\right)^{2}\operatorname{\rho_{{\cal M}}}(x)\mathrm{d}x,$ (13)

where the phenomenological parameter $\operatorname{\lambda_{0}}$ was set to
0.2 14. For the $V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$ and
$V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$ functionals, we used the soft
approximation 15,

$\displaystyle V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\int\int v_{\text{H}}(x)\operatorname{\rho_{{\cal
M}}}(x)\operatorname{\rho_{{\cal
M}}}(x^{\prime})\mathrm{d}x\mathrm{d}x^{\prime}=\int\int\frac{\operatorname{\rho_{{\cal
M}}}(x)\operatorname{\rho_{{\cal
M}}}(x^{\prime})}{\sqrt{1+|x-x^{\prime}|^{2}}}\mathrm{d}x\mathrm{d}x^{\prime},$
(14) $\displaystyle V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\int v_{\text{e-N}}(x)\operatorname{\rho_{{\cal
M}}}(x)\mathrm{d}x$
$\displaystyle=-\int\left(\frac{Z_{\alpha}}{\sqrt{1+|x-R/2|^{2}}}+\frac{Z_{\beta}}{\sqrt{1+|x+R/2|^{2}}}\right)\operatorname{\rho_{{\cal
M}}}(x)\mathrm{d}x,$ (15)

where $Z_{\alpha}$ and $Z_{\beta}$ are the atomic numbers and $R$ is the
distance between the two nuclei. The $E_{\text{XC}}[\operatorname{\rho_{{\cal
M}}}]$ functional’s form is given by 16,

$E_{\text{XC}}[\operatorname{\rho_{{\cal
M}}}]=\int\epsilon_{\text{XC}}\operatorname{\rho_{{\cal M}}}(x)\mathrm{d}x,$
(16)

where $\epsilon_{\text{XC}}$ is 17,

$\epsilon_{\text{XC}}(\operatorname{r_{\text{s}}},\zeta)=\frac{\operatorname{{a}_{\zeta}}+\operatorname{{b}_{\zeta}}\operatorname{r_{\text{s}}}+\operatorname{{c}_{\zeta}}\operatorname{r_{\text{s}}}^{2}}{1+\operatorname{{d}_{\zeta}}\operatorname{r_{\text{s}}}+\operatorname{{e}_{\zeta}}\operatorname{r_{\text{s}}}^{2}+\operatorname{{f}_{\zeta}}\operatorname{r_{\text{s}}}^{3}}+\frac{\operatorname{{g}_{\zeta}}\operatorname{r_{\text{s}}}\ln[{\operatorname{r_{\text{s}}}+\operatorname{\alpha_{\zeta}}\operatorname{r_{\text{s}}}^{\operatorname{\beta_{\zeta}}}}]}{1+\operatorname{{h}_{\zeta}}\operatorname{r_{\text{s}}}^{2}}.$
(17)

For all one-dimensional simulations, we used,
$\operatorname{r_{\text{s}}}=\tfrac{1}{2\operatorname{\rho_{{\cal M}}}}$
(Wigner-Seitz radius17) and $\zeta=0$ (unpolarized density). All the
parameters of $\epsilon_{\text{XC}}$ are defined in Table 1. The expectation
values of $T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$,
$T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$,
$V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$,
$V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$ and
$\operatorname{E_{\text{XC}}}[\operatorname{\rho_{{\cal M}}}]$, accordingly to
Eq. 10, are,

$\displaystyle T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\frac{\pi^{2}}{24}\operatorname{N_{e}}^{3}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[(\operatorname{\rho_{\phi}}(x))^{2}\right],$
(18) $\displaystyle T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\frac{\operatorname{\lambda_{0}}}{8}\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\left(\frac{\nabla\rho_{\phi}(x)}{\rho_{\phi}(x)}\right)^{2}\right]=\frac{\operatorname{\lambda_{0}}}{8}\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\left(\nabla\log\operatorname{\rho_{\phi}}(x)\right)^{2}\right],$
(19) $\displaystyle V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle={\operatorname{N_{e}}^{2}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\frac{1}{\sqrt{1+|x-x^{\prime}|^{2}}}\right],$
(20) $\displaystyle V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[-\frac{Z_{\alpha}}{\sqrt{1+|x+R/2|^{2}}}-\frac{Z_{\beta}}{\sqrt{1+|x+R/2|^{2}}}\right],$
(21) $\displaystyle\operatorname{E_{\text{XC}}}[\operatorname{\rho_{{\cal
M}}}]$
$\displaystyle=\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}[\epsilon_{\text{XC}}(\operatorname{r_{\text{s}}},\zeta)].$
(22)

$\operatorname{{a}_{0}}$ | $-0.8862269$
---|---
$\operatorname{{b}_{0}}$ | $-2.1414101$
$\operatorname{{c}_{0}}$ | $0.4721355$
$\operatorname{{d}_{0}}$ | $2.81423$
$\operatorname{{e}_{0}}$ | $0.529891$
$\operatorname{{f}_{0}}$ | $0.458513$
$\operatorname{{g}_{0}}$ | $-0.202642$
$\operatorname{{h}_{0}}$ | $0.470876$
$\operatorname{\alpha_{0}}$ | $0.104435$
$\operatorname{\beta_{0}}$ | $4.11613$
Table 1: Parameter values of $\epsilon_{\text{XC}}$ (Eq. 17) in a.u., obtained
from Ref. 17.

###  2.2 Three-dimensional density funcitonals

In this section, we present the functionals used for the full three-
dimensional simulations for the H2, LiH, and H2O molecules. The analytic forms
of the used density functionals are,

$\displaystyle T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\frac{3}{10}(3\pi^{2})^{\frac{2}{3}}\int\left(\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\right)^{5/3}\mathrm{d}\operatorname{\boldsymbol{x}},$
(23) $\displaystyle T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\frac{\lambda}{8}\int\frac{(\nabla\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}}))^{2}}{\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})}\mathrm{d}\operatorname{\boldsymbol{x}}=\frac{\lambda}{8}\int\left(\nabla\log\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\right)^{2}\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}},$
(24) $\displaystyle V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\frac{1}{2}\int\int
v_{\text{H}}(\operatorname{\boldsymbol{x}})\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}}^{\prime})\mathrm{d}\mathbf{x}\mathrm{d}\mathbf{x}^{\prime}=\frac{1}{2}\int\int\frac{\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}}^{\prime})}{\sqrt{|\operatorname{\boldsymbol{x}}-\operatorname{\boldsymbol{x}}^{\prime}|^{2}}}\mathrm{d}\operatorname{\boldsymbol{x}}\mathrm{d}\operatorname{\boldsymbol{x}}^{\prime},$
(25) $\displaystyle V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=\int
v_{\text{e-N}}(\operatorname{\boldsymbol{x}})\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}}=-\int\left(\sum_{i}\frac{Z_{i}}{\|\operatorname{\boldsymbol{x}}-\mathbf{R}_{i}\|}\right)\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}},$
(26)

where $\mathbf{R}_{i}$ is the position of the $i^{\text{th}}-$nucleus. We took
advantage of the ”log-derivative trick” (Eq. 5) and defined
$T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$ (Eq. 24) in terms on
$\nabla\log\operatorname{\rho_{{\cal M}}}$. Furthermore, the expectation
values of the above functions (Eqs. 23–26) are,

$\displaystyle T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$ $\displaystyle=$
$\displaystyle\frac{3}{10}(3\pi^{2})^{\frac{2}{3}}\operatorname{N_{e}}^{2/3}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\rho_{\phi}(\operatorname{\boldsymbol{x}})\right],$
(27) $\displaystyle T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=$
$\displaystyle\frac{\lambda}{8}\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\left(\frac{\nabla\rho_{\phi}(\operatorname{\boldsymbol{x}})}{\rho_{\phi}(\operatorname{\boldsymbol{x}})}\right)^{2}\right]=\frac{\lambda}{8}\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\left(\nabla\log\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}})\right)^{2}\right],$
(28) $\displaystyle V_{\text{H}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=$
$\displaystyle\operatorname{N_{e}}^{2}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\frac{1}{\sqrt{1+|\operatorname{\boldsymbol{x}}-\operatorname{\boldsymbol{x}}^{\prime}|^{2}}}\right],$
(29) $\displaystyle V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$
$\displaystyle=$
$\displaystyle-\operatorname{N_{e}}\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}\left[\sum_{i}\frac{Z_{i}}{\|\operatorname{\boldsymbol{x}}-\mathbf{R}_{i}\|}\right],$
(30)

The exchange-correlation functional is composed of the sum of the exchange (X)
and correlation (C) terms,

$E_{\text{XC}}[\operatorname{\rho_{{\cal
M}}}]=\int\epsilon_{\text{XC}}\operatorname{\rho_{{\cal
M}}}(x)\mathrm{d}\operatorname{\boldsymbol{x}}=\int\epsilon_{\text{X}}\operatorname{\rho_{{\cal
M}}}(x)\mathrm{d}\operatorname{\boldsymbol{x}}+\int\epsilon_{\text{C}}\operatorname{\rho_{{\cal
M}}}(x)\mathrm{d}\operatorname{\boldsymbol{x}}.$ (31)

We report all different $\operatorname{\epsilon_{\text{X}}}$ and
$\operatorname{\epsilon_{\text{C}}}$ used in the simulations,

$\displaystyle\epsilon_{\text{X}}^{\text{LDA}}$
$\displaystyle=-\frac{3}{4}\left(\frac{3}{\pi}\right)^{1/3}\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})^{1/3}$ (32)
$\displaystyle\epsilon_{\text{X}}^{\text{B88}}$
$\displaystyle=-\beta\frac{X^{2}}{\left(1+6\beta
X\sinh^{-1}(X)\right)}\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})^{1/3},$ (33)
$\displaystyle\epsilon_{\text{C}}^{\text{VWN}}$
$\displaystyle=\frac{A}{2}\left\\{\ln\left(\frac{y^{2}}{Y(y)}\right)+\frac{2b}{Q}\tan^{-1}\left(\frac{Q}{2y+b}\right)+\right.$
$\displaystyle-\left.\frac{by_{0}}{Y(y_{0})}\left[\ln\left(\frac{(y-y_{0})^{2}}{Y(y)}\right)+\frac{2(b+2y_{0})}{Q}\tan^{-1}\left(\frac{Q}{2y+b}\right)\right]\right\\}$
(34) $\displaystyle\epsilon_{\text{C}}^{\text{PW92}}$
$\displaystyle=-2A(1+\alpha_{1}\operatorname{r_{\text{s}}})\ln\left[{1+\frac{1}{2A(\beta_{1}\operatorname{r_{\text{s}}}^{1/2}+\beta_{2}\operatorname{r_{\text{s}}}+\beta_{3}\operatorname{r_{\text{s}}}^{3/2}+\beta_{4}\operatorname{r_{\text{s}}}^{2}}}\right]\
$ (35)

where
$\operatorname{r_{\text{s}}}=\left(\frac{3}{4\pi\operatorname{\rho_{{\cal
M}}}}\right)^{\frac{1}{3}}$ from Ref.18. For
$\epsilon_{\text{C}}^{\text{VWN}}$, $y=\operatorname{r_{\text{s}}}^{1/2}$,
$Y(y)=y^{2}+by+c$, $Q=\sqrt{4c-b^{2}}$, and the constants $b$, $c$ and $y_{0}$
are given in the Table 2. The PW92’s parameters are reported in Table 3. For
$\epsilon_{\text{X}}^{\text{B88}}$, $\beta$ is $0.0042$ a.u and
$X=\tfrac{|\nabla\operatorname{\rho_{{\cal M}}}|}{\operatorname{\rho_{{\cal
M}}}^{4/3}}$ 19, where we use the ”log-derivative trick” (Eq. 5) and expand in
terms of the score function $X=\tfrac{|\operatorname{\rho_{{\cal
M}}}\nabla\log{\operatorname{\rho_{{\cal M}}}}|}{\operatorname{\rho_{{\cal
M}}}^{4/3}}$.

$A$ | $0.0621814$
---|---
$y_{0}$ | $-0.10498$
$b$ | $3.72744$
$c$ | $12.9352$
Table 2: Parameter values of $\epsilon_{\text{C}}^{\text{VWN}}$ (Eq. 34) in a.u., obtained from Ref.20. $A$ | $0.031091$
---|---
$\alpha_{1}$ | $0.21370$
$\beta_{1}$ | $7.5957$
$\beta_{2}$ | $3.5876$
$\beta_{3}$ | $1.6382$
$\beta_{4}$ | $0.49294$
Table 3: Parameter values of $\epsilon_{\text{C}}^{\text{PW92}}$ (Eq. 35) in
a.u., obtained from Ref.18.

##  3 Optimization algorithm

In this section, we present the optimization algorithm presented in the main
text.

The optimization of $\operatorname{\rho_{{\cal M}}}$ was performed through a
MC scheme where the parameters of the normalizing flow ($\phi$) are updated
via a stochastic gradient optimization,

$\phi=\phi-\alpha\nabla_{\phi}E[\operatorname{\rho_{{\cal M}}}],$ (36)

where $\alpha$ is the learning rate, and
$\nabla_{\phi}E[\operatorname{\rho_{{\cal M}}}]$ is the gradient of the energy
with respect to the parameters $\phi$. As mentioned in the main text,
$\nabla_{\phi}E[\operatorname{\rho_{{\cal M}}}]$ is also estimated through a
MC scheme 21,

$\nabla_{\phi}E[\operatorname{\rho_{{\cal
M}}}]=\operatorname{\mathbb{E}}_{\operatorname{\rho_{0}}}[\nabla_{\phi}f_{E}(\operatorname{T_{\phi}}(\operatorname{\mathbf{z}}),\operatorname{\rho_{\phi}},\nabla\operatorname{\rho_{\phi}})]\approx\tfrac{1}{N}\sum_{i}^{N}\nabla_{\phi}f_{E}(\operatorname{T_{\phi}}(\operatorname{\mathbf{z}}_{i}),\operatorname{\rho_{\phi}}(\operatorname{\mathbf{z}}_{i}),\nabla_{\operatorname{\mathbf{z}}_{i}}\operatorname{\rho_{\phi}}).$
(37)

For this work, all required gradients were computed using automatic
differentiation, as it is a common practice in computational chemistry 22, 23,
24, 25.

A representation of the proposed algorithm is depicted in Algorithm 1. For all
simulations, we employed the RMSProp optimizer 26, featuring a learning rate
schedule with initial and final values set at $3\times 10^{-4}$ and $10^{-7}$,
respectively. Although other optimizer, such as Adam, were explored, they were
found to be sub-optimal in comparison. Code was developed using JAX Ecosystem
27, 28, 29 and for all simulations, we used an NVIDIA Tesla P100 GPU. The code
developed for this work is available in the following repository.

Algorithm 1 Optimization Algorithm

CNF parameters $\phi$, base distribution $\operatorname{\rho_{0}}$, energy
functional $E[\cdot]$

while not converged do
$\\{\operatorname{\mathbf{z}}_{i},\log\operatorname{\rho_{0}}(\operatorname{\mathbf{z}}_{i}),\nabla_{z}\log\operatorname{\rho_{0}}(\operatorname{\mathbf{z}}_{i})\\}_{i}^{N}$
$\sim\operatorname{\rho_{0}}(\operatorname{\mathbf{z}})$ $\triangleright$
Sample the base distribution
$[\operatorname{\boldsymbol{x}}_{i},\log\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}_{i}),\nabla_{\operatorname{\boldsymbol{x}}}\log\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}_{i})]=\texttt{ODESolve}([\operatorname{\mathbf{z}}_{i},\log\operatorname{\rho_{0}}(\operatorname{\mathbf{z}}_{i}),\nabla_{z}\log\operatorname{\rho_{0}}(\operatorname{\mathbf{z}}_{i})],t_{0},t_{1},\phi$)$\triangleright$
Solve ODE Eq.( 4)
$E=\operatorname{N_{e}}^{p}\operatorname{\mathbb{E}}_{\operatorname{\rho_{\phi}}}[E[\operatorname{\boldsymbol{x}},\operatorname{\rho_{{\cal
M}}},\nabla\operatorname{\rho_{{\cal
M}}}]]\approx\frac{\operatorname{N_{e}}^{p}}{N}\sum^{N}_{i}f_{E}(\operatorname{\boldsymbol{x}}_{i},\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}_{i}),\nabla\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}_{i}))$
$\triangleright$ Compute energy, Eq. ( 9)
$\nabla_{\phi}E\approx\frac{(\operatorname{N_{e}})^{p}}{N}\sum^{N}_{i}\nabla_{\phi}f_{E}(\operatorname{\boldsymbol{x}}_{i},\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}_{i}),\nabla\operatorname{\rho_{\phi}}(\operatorname{\boldsymbol{x}}_{i}))$
$\triangleright$ Evaluate gradients, Eq.( 37) $\phi^{\prime}$ = optmizer
step($\phi$, $\nabla_{\phi}E$) $\triangleright$ Update parameters

end while

return parameters $\phi^{\prime}$

##  4 Results

In this section, we present the results of the 1-D and 3-D simulations. The
analytic equations of all functionals are presented in Sections 2.1 and 2.2,
and Section 3 details the information of the optimization algorithm.

###  4.1 1-D: LiH

In this section, we present additional results for the one-dimensional (1D)
model of LiH at different inter-atomic distances ($R$). For all simulations,
$\operatorname{g_{\phi}}$ was parametrized by a three-layer neural network
with 512 neurons per layer, and a hyperbolic tangent activation function,
$\tanh{(\cdot)}$. A maximum of 10,000 iterations was allowed and at each
gradient iteration, a batch size of 512 samples was used. We only considered
as based distribution $\operatorname{\rho_{0}}$ a 1D Gaussian distribution
centered at $0$; ${\cal N}(x,\sigma=1)$.

The energy functional for different values of $R$ was composed of the sum of
$T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$,
$T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$,
$V_{H}[\operatorname{\rho_{{\cal M}}}]$ and
$V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$ (Eqs. 18-21). For the
exchange-correlation functional, we used the one from Ref. 16 (Eq. 17). In
Fig. 1, we presented the value of each functional through the optimization. To
validate the total energy estimated with MC, we also used quadrature
integration (trapezoidal rule), and as we can observe, it follows the same
trend as the estimated with MC. For low values of $R$, the optimization
converges rapidly as the $\operatorname{\rho_{{\cal M}}}$ mimics a unimodal
distribution. However, for larger values of $R$, a greater number of
iterations is needed for $\operatorname{g_{\phi}}$ to split
$\operatorname{\rho_{{\cal M}}}$ into a bimodal distribution properly.

Figure 1: The value of the total energy and each functional, computed with MC
(Eq. 10), at each iteration. Each panel represents an independent simulation
for different values of $R$. The  $\pentagofill$-markers represent the value
of the total energy computed with quadrature integration. For all simulations,
we used the same architecture for $\operatorname{g_{\phi}}$ and optimizer, see
the text for more details.

###  4.2 3-D: H2, H2O and LiH

In this section, we present additional results for the three-dimensional (3D)
simulations for the H2, LiH, and H2O molecules. For these simulations, we
considered two distinct base distributions, (i) a single Gaussian distribution
($\operatorname{\rho_{0}}$), and (ii) a promolecular density
($\operatorname{\tilde{\rho}_{0}}$),

$\displaystyle\operatorname{\tilde{\rho}_{0}}=\sum_{i}c_{i}{\cal
N}_{i}(\mathbf{R}_{i},\sigma=1),$ (38)

where ${\cal N}_{i}(\cdot)$ is a 1S orbital centered at the nucleus position,
$\mathbf{R}_{i}$. The coefficients $c_{i}$ represent the proportional
influence of each component on the overall density,
$c_{i}=\tfrac{Z_{i}}{\sum_{j}Z_{j}}$, and it depends on $Z_{i}$, the atomic
number of the $i^{th}$-nucleus. To account for possible symmetries in the
$V_{\text{e-N}}$ potential, $\operatorname{g_{\phi}}$ is a permutation
equivariant graph neural network (GNN)10, 11, Eq. 7.

Figure 2: The value of $E[\operatorname{\rho_{{\cal M}}}]$ and each
functional, computed with MC (Eq. 10), at each iteration of the optimization.
In this case, we considered a promolecular density
($\operatorname{\tilde{\rho}_{0}}$) as our base distribution. The
$\pentagofill$-markers represent the value of the total energy computed with
quadrature integration. For all simulations, we used the same architecture for
$\operatorname{g_{\phi}}$ and optimizer, see the text for more details.

Figure 3: The value of $E[\operatorname{\rho_{{\cal M}}}]$ and each
functional, computed with MC (Eq. 10), at each iteration of the optimization.
In this case, we considered a single Gaussian distribution
($\operatorname{\rho_{0}}$) as our base distribution. The
$\pentagofill$-markers represent the value of the total energy computed with
quadrature integration. For all simulations, we used the same architecture for
$\operatorname{g_{\phi}}$ and optimizer, see the text for more details.

The total energy functional was composed of the sum of
$T_{\text{TF}}[\operatorname{\rho_{{\cal M}}}]$,
$T_{\text{W}}[\operatorname{\rho_{{\cal M}}}]$,
$V_{H}[\operatorname{\rho_{{\cal M}}}]$, and
$V_{\text{e-N}}[\operatorname{\rho_{{\cal M}}}]$ (Eqs. 27-30). For the
exchange-correlation functional, we combine
$\operatorname{\epsilon_{\text{X}}^{\text{LDA}}}$ and
$\operatorname{\epsilon_{\text{X}}^{\text{B88}}}$, Eqs. 32 and 33
respectively, as the exchange component, and intercalate
$\epsilon_{\text{C}}^{\text{VWN}}$ (Eq. 34) and
$\epsilon_{\text{C}}^{\text{PW92}}$ (Eq. 35) as the correlation term. The
values of the total energy estimated with MC and standard quadrature
integration are presented in Table 4, combined with the average optimization
step time. We only report the values for $\operatorname{\tilde{\rho}_{0}}$, as
we found the energy is always lower than when using a single Gaussian
distribution as $\operatorname{\rho_{0}}$.

Figs. 2 and 3 shows the value of each functional for H2, LiH, and H2O as a
function of the iteration steps. To validate the MC scheme, we also computed
the full energy using quadrature integration using the proposed CNF ansatz.
For Fig. 2, the base distribution was parametrized using the promolecular
density, and for Fig. 3 $\operatorname{\rho_{0}}$ us a single multi-variate
Gaussian distribution. For both sets of simulations,
$\operatorname{g_{\phi}}$’s architecture (Eq. 7) comprised two layers. For
results with a different number of layers, consult Table 4.

To validate the normalization guarantee, Fig. 4 shows the value of
$\Delta\operatorname{N_{e}}=\int\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}}-\operatorname{N_{e}}$
through the optimization. Note, that we used a logarithmic scale, and the
difference between the normalization starting with a promolecular base
distribution ($\operatorname{\tilde{\rho}_{0}}$) and a single multi-variate
Gaussian ($\operatorname{\rho_{0}}$). We found that in simulations when
$\operatorname{\tilde{\rho}_{0}}$ is used, the value of
$|\Delta\operatorname{N_{e}}|$ is lower than when $\operatorname{\rho_{0}}$ is
used.

Total Energy [a.u.]. |  |
---|---|---
| MC | Integration | Iteration time [s]
Molecule ($\operatorname{N_{e}}$) | # Layers | VWN | PW92 | VWN | PW92 | VWN | PW92
H2 ($2$) | $1$ | $-1.2728$ | $-1.2729$ | $-1.2764$ | $-1.2765$ | $0.11(3)$ | $0.10(3)$
$2$ | $-1.2798$ | $-1.2799$ | $-1.2835$ | $-1.2837$ | $0.26(7)$ | $0.26(7)$
$3$ | $-1.28249$ | $-1.28250$ | $-1.28618$ | $-1.28619$ | $0.39(9)$ | $0.40(9)$
LiH ($4$) | $1$ | $-7.81768$ | $-7.81765$ | $-7.7786$ | $-7.7780$ | $0.22(8)$ | $0.22(8)$
$2$ | $-7.9557$ | $-7.9588$ | $-7.9338$ | $-7.9367$ | $0.5(2)$ | $0.5(2)$
$3$ | $-7.9700$ | $-7.9698$ | $-7.9529$ | $-7.9530$ | $0.8(2)$ | $0.8(2)$
H2O ($10$) | $1$ | $-78.2742$ | $-78.2396$ | $-76.3132$ | $-76.3461$ | $0.31(6)$ | $0.32(6)$
$2$ | $-83.1841$ | $-83.0956$ | $-81.6134$ | $-81.5315$ | $0.8(2)$ | $0.8(2)$
$3$ | $-84.0374$ | $-83.7963$ | $-82.3544$ | $-82.2378$ | $1.3(3)$ | $1.3(3)$
Table 4: The total energy computed with MC and quadrature integration for two
different $\epsilon_{XC}$ functionals. We also present the average iteration
step during optimization.

Figure 4: The absolute difference between the number of electrons
($\operatorname{N_{e}}$) and the integral of $\operatorname{\rho_{{\cal M}}}$
at each iteration during the optimization. We define
$\Delta\operatorname{N_{e}}$ as,
$\Delta\operatorname{N_{e}}=\int\operatorname{\rho_{{\cal
M}}}(\operatorname{\boldsymbol{x}})\mathrm{d}\operatorname{\boldsymbol{x}}-\operatorname{N_{e}}$.

## References

  * Kobyzev et al. 2020 Kobyzev, I.; Prince, S. J.; Brubaker, M. A. Normalizing flows: An introduction and review of current methods. _IEEE transactions on pattern analysis and machine intelligence_ 2020, _43_ , 3964–3979
  * Parr and Yang 1980 Parr, R. G.; Yang, W. Density Functional Theory of Atoms and Molecules. Horizons of Quantum Chemistry. Dordrecht, 1980; pp 5–15
  * Parr and Bartolotti 1983 Parr, R. G.; Bartolotti, L. J. Some remarks on the density functional theory of few-electron systems. _The Journal of Physical Chemistry_ 1983, _87_ , 2810–2815
  * Papamakarios et al. 2021 Papamakarios, G.; Nalisnick, E.; Rezende, D. J.; Mohamed, S.; Lakshminarayanan, B. Normalizing flows for probabilistic modeling and inference. _The Journal of Machine Learning Research_ 2021, _22_ , 2617–2680
  * Rezende and Mohamed 2015 Rezende, D.; Mohamed, S. Variational inference with normalizing flows. International conference on machine learning. 2015; pp 1530–1538
  * Chen et al. 2018 Chen, R. T.; Rubanova, Y.; Bettencourt, J.; Duvenaud, D. K. Neural ordinary differential equations. _Advances in neural information processing systems_ 2018, _31_
  * Shuangshuang et al. 2023 Shuangshuang, C.; Sihao, D.; Yiannis, K.; Mårten, B. Learning Continuous Normalizing Flows For Faster Convergence To Target Distribution via Ascent Regularizations. The Eleventh International Conference on Learning Representations. 2023
  * Pontryagin 2018 Pontryagin, L. S. _Mathematical theory of optimal processes_ ; Routledge, 2018
  * Shampine 1986 Shampine, L. F. Some practical runge-kutta formulas. _Mathematics of computation_ 1986, _46_ , 135–150
  * Köhler et al. 2020 Köhler, J.; Klein, L.; Noe, F. Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities. Proceedings of the 37th International Conference on Machine Learning. 2020; pp 5361–5370
  * Zwartsenberg et al. 2023 Zwartsenberg, B.; Scibior, A.; Niedoba, M.; Lioutas, V.; Sefas, J.; Liu, Y.; Dabiri, S.; Lavington, J. W.; Campbell, T.; Wood, F. Conditional Permutation Invariant Flows. _Transactions on Machine Learning Research_ 2023,
  * Hohenberg and Kohn 1964 Hohenberg, P.; Kohn, W. Inhomogeneous electron gas. _Physical review_ 1964, _136_ , B864
  * Snyder et al. 2013 Snyder, J. C.; Rupp, M.; Hansen, K.; Blooston, L.; Müller, K.; Burke, K. Orbital-free bond breaking via machine learning. _The Journal of chemical physics_ 2013, _139_
  * Chan et al. 2001 Chan, G. K.-L.; Cohen, A. J.; Handy, N. C. Thomas-Fermi-Dirac-von Weizsäcker models in finite systems. _The Journal of Chemical Physics_ 2001, _114_ , 631–638
  * Snyder et al. 2012 Snyder, J. C.; Rupp, M.; Hansen, K.; Müller, K.-R.; Burke, K. Finding density functionals with machine learning. _Physical review letters_ 2012, _108_ , 253002
  * Helbig et al. 2011 Helbig, N.; Fuks, J. I.; Casula, M.; Verstraete, M. J.; Marques, M. A.; Tokatly, I.; Rubio, A. Density functional theory beyond the linear regime: Validating an adiabatic local density approximation. _Physical Review A_ 2011, _83_ , 032503
  * Shulenburger et al. 2009 Shulenburger, L.; Casula, M.; Senatore, G.; Martin, R. M. Spin resolved energy parametrization of a quasi-one-dimensional electron gas. _Journal of Physics A: Mathematical and Theoretical_ 2009, _42_ , 214021
  * Perdew and Wang 1992 Perdew, J. P.; Wang, Y. Accurate and simple analytic representation of the electron-gas correlation energy. _Physical review B_ 1992, _45_ , 13244
  * Becke 1988 Becke, A. D. Density-functional exchange-energy approximation with correct asymptotic behavior. _Physical review A_ 1988, _38_ , 3098
  * Vosko et al. 1980 Vosko, S. H.; Wilk, L.; Nusair, M. Accurate spin-dependent electron liquid correlation energies for local spin density calculations: a critical analysis. _Canadian Journal of physics_ 1980, _58_ , 1200–1211
  * Mohamed et al. 2020 Mohamed, S.; Rosca, M.; Figurnov, M.; Mnih, A. Monte carlo gradient estimation in machine learning. _Journal of Machine Learning Research_ 2020, _21_ , 1–62
  * M. Casares et al. 2024 M. Casares, P. A.; Baker, J. S.; Medvidović, M.; Reis, R. d.; Arrazola, J. M. GradDFT. A software library for machine learning enhanced density functional theory. 2024; https://doi.org/10.1063/5.0181037
  * Kasim and Vinko 2021 Kasim, M. F.; Vinko, S. M. Learning the Exchange-Correlation Functional from Nature with Fully Differentiable Density Functional Theory. _Phys. Rev. Lett._ 2021, _127_ , 126403
  * Tamayo-Mendoza et al. 2018 Tamayo-Mendoza, T.; Kreisbeck, C.; Lindh, R.; Aspuru-Guzik, A. Automatic Differentiation in Quantum Chemistry with Applications to Fully Variational Hartree–Fock. _ACS Central Science_ 2018, _4_ , 559–566, PMID: 29806002
  * Vargas–Hernández et al. 2023 Vargas–Hernández, R. A.; Jorner, K.; Pollice, R.; Aspuru–Guzik, A. Inverse molecular design and parameter optimization with Hückel theory using automatic differentiation. _The Journal of Chemical Physics_ 2023, _158_ , 104801
  * Tieleman et al. 2012 Tieleman, T.; Hinton, G., et al. Lecture 6.5-RMSProp: Divide the gradient by a running average of its recent magnitude. _COURSERA: Neural networks for machine learning_ 2012, _4_ , 26–31
  * Bradbury et al. 2018 Bradbury, J.; Frostig, R.; P.Hawkins,; Johnson, M.; Leary, C.; Maclaurin, D.; Necula, G.; Paszke, A.; VanderPlas, J.; Wanderman-Milne, S.; Zhang, Q. JAX: composable transformations of Python+NumPy programs. 2018; http://github.com/google/jax
  * Heek et al. 2023 Heek, J.; A.Levskaya,; Oliver, A.; Ritter, M.; Rondepierre, B.; Steiner, A.; van Zee, M. Flax: A neural network library and ecosystem for JAX. 2023; http://github.com/google/flax
  * DeepMind et al. 2020 DeepMind, et al. The DeepMind JAX Ecosystem. 2020; http://github.com/deepmind