diff --git a/README.md b/README.md
index 1bcb80b..363b434 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # OF-DFT with Continuous-time Normalizing Flows
 
-This repository contains the original implementation of the experiments for ["Orbital-Free Density Functional Theory with Continuous Normalizing Flows"](archive).
+This repository contains the original implementation of the experiments for ["Leveraging Normalizing Flows for Orbital-Free Density Functional Theory"](https://arxiv.org/abs/2404.08764).
 
 --------------------
 
@@ -19,6 +19,16 @@ $$\min_{\rho_\mathcal{M}}  E[\rho_\mathcal{M}],$$
 
 where we parameterize the electron density $\rho_\mathcal{M} := N_{e}  \rho_{\phi}(\mathbf{x})$, where $\rho_{\phi}$ is a NF, this form is also referred to as the *shape factor*. The samples are drawn from the base distribution $\rho_0$ and transformed by, $$\mathcal{x} = T_\phi(\mathcal{z}) := \mathcal{z} + \int_{t_{0}}^{T} g_\phi(\mathcal{z}(t),t) \mathrm{d}t.$$
 
+For the one-dimensional simulations, the architecture of $g_\phi$ is a standard feed-forward neural network (NN),
+$$g_\phi = \sum_\ell^M f_\ell(\mathbf{z}_\ell(t)),$$
+
+where $f_\ell(\cdot)$ is a linear layer followed by an activation function, and $M$ is the number of layers. For this work, $g_\phi$ has 3 layers, each with 512 neurons, and the $\tanh$ activation function. For the simulation in three dimensions, $g_\phi$ is parametrized by a permutation equivariant graph NN (GNN), 
+
+$$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),$$
+
+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.
+
+
 ## Results
 
 We successfully replicate the electronic density for the one-dimensional Lithium hydride molecule with varying interatomic distances, as well as comprehensive simulations of hydrogen and water molecules, all conducted in
@@ -26,7 +36,7 @@ Cartesian space.
 
 ## Running the code 
 
-# 1-D 
+## 1-D 
 For Lithium hydride ($\texttt{LiH}$) molecule, simulations can be run in the following way, 
 ```
 python LiH.py
@@ -37,13 +47,13 @@ python LiH.py
                 --R <interatomic distances>
                 --Z <atomic number> 
 ```
-The default functionals can be found in the directory [ofdft_normflows](https://github.com/RodrigoAVargasHdz/ofdft_normflows/tree/ml4phys2023/ofdft_normflows#readme)
+The default functionals can be found in the directory [ofdft_normflows](https://github.com/RodrigoAVargasHdz/ofdft_normflows/tree/ml4phys2023/ofdft_normflows#readme).
 
 |$\rho_{{\cal M}}$ of $\texttt{LiH}$ for various inter-atomic distances.|The change of $\rho_{{\cal M}}$ and $T_\phi(\mathcal{z})$ during training.|
 |:--:|:--:|
 |![](https://github.com/RodrigoAVargasHdz/ofdft_normflows/blob/ml4phys2023/Assets/Figure_1.png)|![](https://github.com/RodrigoAVargasHdz/ofdft_normflows/blob/ml4phys2023/Assets/neural_ode_2_gif.gif)|
 
-# 3-D 
+## 3-D 
 For water ($\texttt{H2O}$) and hydrogen ($\texttt{H2}$) molecules, simulations can be run in the following way,
 ```
 python H2_mol_ofdft_min.py
@@ -58,4 +68,6 @@ The default kinetic energy functional is the sum of the [Thomas-Fermi and Weizs
 
 1. [DeepMind JAX Ecosystem]([jax.readthedocs.io/](https://deepmind.google/discover/blog/using-jax-to-accelerate-our-research/))
 2. [Flax](flax.readthedocs.io/)
-6. [PySCF](pyscf.org/)
+3. [PySCF](pyscf.org/)
+
+This is a library that is currently being built.
diff --git a/ofdft_normflows/README.md b/ofdft_normflows/README.md
index 533c364..43ea002 100644
--- a/ofdft_normflows/README.md
+++ b/ofdft_normflows/README.md
@@ -5,9 +5,9 @@ This `ofdft_normflows` directory contains clean up code regarding the energy fun
 # Energy Functionals 
 
 The `functionals.py` file contains the codes regarding the total energy functional,
-    $$E[\rho_{\mathcal{M}}] = T[\rho_{\mathcal{M}}] + V_{\text{H}}[\rho_{\mathcal{M}}] +  V_{\text{e-N}}[\rho_{\mathcal{M}}]  + E_{X}[\rho_{\mathcal{M}}],$$
+    $$E[\rho_{\mathcal{M}}] = T[\rho_{\mathcal{M}}] + V_{\text{H}}[\rho_{\mathcal{M}}] +  V_{\text{e-N}}[\rho_{\mathcal{M}}]  + E_{\text{XC}}[\rho_{\mathcal{M}}],$$
 
-where $\rho_{\mathcal{M}}(\mathbf{x})$ is already defined [here](https://github.com/RodrigoAVargasHdz/ofdft_normflows/blob/ml4phys2023/README.md). The total energy functional ($E[\rho_{\mathcal{M}}]$) is composed with the total kinetic energy ($T[\rho_{\mathcal{M}}]$), the Hartree potential ($V_{\text{H}}[\rho_{\mathcal{M}}]$), the external potential ($V_{\text{e-N}}[\rho_{\mathcal{M}}]$) and the Dirac exchange ($E_{X}[\rho_{\mathcal{M}}]$) functionals. Both 1-D and 3-D used the same energy functional and the differences of each case will be described in the next sections. 
+where $\rho_{\mathcal{M}}(\mathbf{x})$ is already defined [here](https://github.com/RodrigoAVargasHdz/ofdft_normflows/blob/ml4phys2023/README.md). The total energy functional ($E[\rho_{\mathcal{M}}]$) is composed with the total kinetic energy ($T[\rho_{\mathcal{M}}]$), the Hartree potential ($V_{\text{H}}[\rho_{\mathcal{M}}]$), the external potential ($V_{\text{e-N}}[\rho_{\mathcal{M}}]$) and the Dirac exchange ($E_{X}[\rho_{\mathcal{M}}]$) functionals. The differences between the 1-D and 3-D case will be presented in the next sections. 
 
 ## 1-D Case
     
@@ -15,12 +15,12 @@ Considering a one-dimensional model for diatomic molecules, the total kinetic en
 
 $$T_{\text{TF}}[\rho_{\mathcal{M}}] = \frac{\pi^2}{24} \int \left(\\rho_{\mathcal{M}}(x) \right)^{3} \mathrm{d}x.$$
 
-$$T_{\text{W}}[\rho_{\mathcal{M}}] = \frac{\lambda}{8} \int \frac{(\nabla \rho_{\mathcal{M}}(x))^2}{\rho_{\mathcal{M}}} \mathrm{d}x, $$
+$$T_{\text{W}}[\rho_{\mathcal{M}}] = \frac{\lambda_0}{8} \int \frac{(\nabla \rho_{\mathcal{M}}(x))^2}{\rho_{\mathcal{M}}} \mathrm{d}x, $$
 
-where the phenomenological parameter $\lambda$ was set to 0.2. 
+where the phenomenological parameter $\lambda_0$ was set to 0.2. 
 
 We rewrite the Weizsäcker functional in terms of the score function, 
-    $$T_{\text{W}}[\rho_{\mathcal{M}}] = \frac{\lambda}{8} \int  \left(\nabla \log \\rho_{\mathcal{M}}(x) \right)^2  \rho_{\mathcal{M}}(x) \mathrm{d}x,$$
+    $$T_{\text{W}}[\rho_{\mathcal{M}}] = \frac{\lambda_0}{8} \int  \left(\nabla \log \\rho_{\mathcal{M}}(x) \right)^2  \rho_{\mathcal{M}}(x) \mathrm{d}x,$$
 in order to use use memory-efficient gradients for optmizing $T_{\text{W}}[\rho_{\mathcal{M}}]$. 
 
 The Hartree ($V_{\text{H}}[\rho_{\mathcal{M}}]$) potential and the external potential ($V_{\text{e-N}}[\rho_{\mathcal{M}}]$) functionals both are defined by a soft version,
@@ -31,15 +31,32 @@ The Hartree ($V_{\text{H}}[\rho_{\mathcal{M}}]$) potential and the external pote
 
 where the atomic numbers $Z_\alpha$ and $Z_\beta$ are the atomic numbers, $N_e$ is the number of valence electrons and $R$ is the interatomic distance. 
 
-And we the Dirac exchange functional ($E_{X}[\rho_{\mathcal{M}}]$), 
-    $$E_{\text{X}}[\rho_{\mathcal{M}}] = -\frac{3}{4} \left(\frac{3}{\pi} \right)^{1/3} \int \rho_{\mathcal{M}}(x)^{4/3} \mathrm{d}x.$$
+And the exchange-correlation functional $E_{\text{XC}}[\rho_{\mathcal{M}}]$ is given by, 
+    $$E_{\text{XC}}[\rho_{\mathcal{M}}] = \int  \epsilon_{\text{XC}} \rho_{\mathcal{M}}$$,
+where $\epsilon_{\text{XC}}$ is, 
+
+$$\epsilon_{\text{XC}} (r_{\text{s}},\zeta) = \frac{a_{\zeta} + b_{\zeta}r_{\text{s}} + c_{\zeta} r_{\text{s}}^{2}}{1 + d_{\zeta} r_{\text{s}} + e_{\zeta} r_{\text{s}}^2 + f_{\zeta} r_{\text{s}}^3} + \frac{g_{\zeta} r_{\text{s}} \ln[{r_{\text{s}} + \alpha_{\zeta} r_{\text{s}}^{\beta_{\zeta}} }]}{1 + h_{\zeta} r_{\text{s}}^2}.$$
+
+For all one-dimensional simulations, we used, the Wigner-Seitz radius $r_{\text{s}} = \tfrac{1}{2\rho_{\mathcal{M}}}$ and $\zeta = 0$ (unpolarized density). 
 
 ## 3-D Case 
 
-To demonstrate the capability to use CNF for real-space simulations, we considered the optimization of $H_2$ and $H_{2}O$. For both chemical systems, we considered the same total energy functional where the differences are in the Hartree-Potential, $$v_{\text{e-N}}(\mathcal{x}) = -\sum_i \frac{Z_i}{\|\mathcal{x} - \mathbf{R}_i\|},$$ where no soft form approximation was used and $Z_i$ is the atomic number of the $i$ atom.
+To demonstrate the capability to use CNF for real-space simulations, we considered the optimization of $H_2$ and $H_{2}O$. For both chemical systems, we considered the same total energy functional where the differences are in the Hartree-Potential, $$v_{\text{e-N}}(\mathcal{x}) = -\sum_i \frac{Z_i}{\|(\boldsymbol{x}) - \mathbf{R}_i\|},$$ where no soft form approximation was used and $Z_i$ is the atomic number of the $i$ atom.
+
+ The Thomas-Fermi ($T_{\text{TF}}$) functional, 
+
+$$T_{\text{TF}}[\rho_{\mathcal{M}}] = \frac{3}{10}(3\pi^2)^{2/3} \int \left(\rho_{\mathcal{M}}(\boldsymbol{x}) \right)^{5/3} \mathrm{d}x.$$
+
+And in the exchange-correlation functional that is composed of the sum of the exchange (X) and correlation (C) terms,
+$$E_{\text{XC}}[\rho_{\mathcal{M}}] = \int  \epsilon_{\text{XC}} \rho_{\mathcal{M}}(\boldsymbol{x}) \mathrm{d}\boldsymbol{x} = \int  \epsilon_{\text{X}} \rho_{\mathcal{M}}(\boldsymbol{x}) \mathrm{d}\boldsymbol{x} + \int  \epsilon_{\text{C}} \rho_{\mathcal{M}}(\boldsymbol{x})  \mathrm{d}\boldsymbol{x}.$$
 
-And the Thomas-Fermi ($T_{\text{TF}}$) functional, 
+We report all different $\epsilon_{\text{X}}$ and $\epsilon_{\text{C}}$ used in the simulations,
+$$\epsilon_{\text{X}}^{\text{LDA}} = -\frac{3}{4} \left(\frac{3}{\pi}\right)^{1/3} \rho_{\mathcal{M}}(\boldsymbol{x})^{1/3}$$
+$$\epsilon_{\text{X}}^{\text{B88}} = -\beta \frac{X^2}{\left(1 + 6 \beta X \sinh^{-1}(X) \right)} (\boldsymbol{x})^{1/3}$$
+$$\epsilon_{\text{C}}^{\text{VWN}} = \frac{A}{2} \left[ \ln\left(\frac{y^2}{Y(y)}\right) + \frac{2b}{Q} \tan^{-1} \left(\frac{Q}{2y + b}\right) \frac{b y_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]$$
+$$\epsilon_{\text{C}}^{\text{PW92}} = -2A(1+\alpha_1 r_{\text{s}}) \ln \left[{1 + \frac{1}{2A(\beta_1 r_{\text{s}}^{1/2} + \beta_2 r_{\text{s}} + \beta_3 r_{\text{s}}^{3/2} + \beta_4 r_{\text{s}}^{2}}}\right]$$
 
-$$T_{\text{TF}}[\rho_{\mathcal{M}}] = \frac{3}{10}(3\pi^2)^{2/3} \int \left(\rho_{\mathcal{M}}(x) \right)^{5/3} \mathrm{d}x.$$
+where $r_{\text{s}} = \left ( \frac{3}{4\pi \rho_{\mathcal{M}}}  \right)^{\frac{1}{3}}$. 
+For $\epsilon_{\text{C}}^{\text{VWN}}$, $y = 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 SI of the paper.