PEANUT model architecture
=========================

Overview of the PEANUT model architecture and components. Let :math:`N_{atoms}` denote the number of atoms in the system,  :math:`N_{elem}` be the number of different chemical elements in the dataset, :math:`D` the dimension of the atom embedding vectors, :math:`B` the number of radial basis functions, :math:`L` the maximum degree of spherical harmonics, and :math:`T` the number of message-passing iterations.
For each atom :math:`i`, let :math:`\mathcal{N}(i)` denote the set of neighboring atoms within the cutoff radius.

Input for the model
-------------------

Cartesian coordinates

.. math::
    r = (r_i)_{i=1}^{N_{\mathit{atoms}}} \in \mathbb{R}^{3 \times N_{\mathit{atoms}}}

Atom types

.. math::
    Z = (Z_i)_{i=1}^{N_{\mathit{atoms}}} \in \{1, \dots, N_{\mathit{elem}}\}^{N_{\mathit{atoms}}}

Construct an embedding vector for each atom:

.. math::
    h_i^{(0)} = e(Z_i) \in \mathbb{R}^D

where :math:`e` is a function mapping the element type :math:`Z_i` of atom :math:`i` to an initial embedding vector :math:`h_i^{(0)}` of length :math:`D`. Distances between atom pairs are encoded as Euclidean distance:

.. math::
    d_{ij} = r_j - r_i, \quad
    r_{ij} = \lVert d_{ij} \rVert_2

Radial features are encoded using :math:`B` Behler-Parinello symmetry functions with centers :math:`R_B \in \mathbb{R}_+` and width parameters :math:`\eta_B \in \mathbb{R}_+`.
Additionally, a cutoff function :math:`f_c` is applied to ensure smoothness at the cutoff radius :math:`r_c`.

.. math::
    R_d(r_{ij}) = e^{-\eta_B (r_{ij} - R_B)^2} f_c(r_{ij})

Each distance :math:`r_{ij}` is mapped to a :math:`B`-dimensional radial feature vector:

.. math::
    R_{ij} = \begin{pmatrix}
        R_1(r_{ij}) \\
        R_2(r_{ij}) \\
        \vdots \\
        R_B(r_{ij})
    \end{pmatrix} \in \mathbb{R}^B



Spherical features are encoded using spherical harmonics:

.. math::
    \hat{d}_{ij} = \frac{d_{ij}}{\lVert d_{ij} \rVert_2}

.. math::
    Y_l^m(\hat{d}_{ij}), \quad m=-l,\dots,l, \quad l=0,\dots,L

Total number of spherical harmonics per pair:

.. math::
    |Y| = \sum_{l=0}^{L} (2l+1) = (L+1)^2

.. math::
    Y_{ij} = \bigl(Y_l^m(\hat{d}_{ij})\bigr)_{l=0,\dots,L}^{m=-l,\dots,l} \in \mathbb{R}^{(L+1)^2}

Radial and spherical features are combined:

.. math::
    C_{ij} = R_{ij} \otimes Y_{ij} \in \mathbb{R}^{B \times (L+1)^2}

.. math::
    F_{ij} = \mathit{flatten}(C_{ij})


Steps per atom :math:`i` and message-passing iteration :math:`t = 1, \dots, T`
------------------------------------------------------------------------------

Message passing: Construct message for each neighbor :math:`j` of atom :math:`i` using edge features :math:`F_{ij}` and sender node features :math:`h_j^{(t-1)}`:

.. math::
    m_{ij}^{(t)} = f_{\mathit{message}}(F_{ij}, h_j^{(t-1)}), \quad j \in \mathcal{N}(i)

Aggregation: Combine all incoming messages from all atoms :math:`j` in the neighborhood :math:`\mathcal{N}(i)` of atom :math:`i` to update the node embedding :math:`h_i^{(t)}` at iteration :math:`t`.

.. math::
    a_i^{(t)} = \frac{1}{|\mathcal{N}(i)|} \sum_{j \in \mathcal{N}(i)} m_{ij}^{(t)}

Node update:

.. math::
    h_i^{(t)} = MLP(h_i^{(t-1)}, a_i^{(t)}) + h_i^{(t-1)}

Energy prediction
-----------------

.. math::
    E_i = MLP_{\mathit{energy}}(h_i^{(T)}) \in \mathbb{R}

Total energy per batch:

.. math::
    E^{(b)} = \sum_{i \in b} E_i