Information Flow in the PEANUT models
=====================================

Overview of the PEANUT model architecture and components.

.. toctree::
   :maxdepth: 2
   :caption: Subsections

   peanut
   peanut_dual_cutoff
   peanut_mini

The models use a star graph architecture to efficiently capture interactions between atoms in a molecular system. 
In this architecture, each atom is represented as a node, and edges connect each atom to a central "star" node that aggregates information from all atoms. 

For each atom :math:`i`, a cutoff sphere of radius :math:`r_c` is defined, within which neighboring atoms contribute to the atom's representation. 
The star graph structure allows for efficient message passing and information aggregation, enabling the model to learn complex

.. image:: ../../_static/star_graphs.png
   :alt: Star Graph
   :width: 800px
   :align: center


Additionally, an attention mechanism can be employed to weigh the contributions of neighboring atoms based on their relevance to the target atom.
This has not been done yet.
Exemplary sketch for an attention layer is shown below:

Sketch for attention layer part 
-------------------------------

.. code-block:: text

   Node i neighbors: j1, j2, j3

   Edge MLP: compute m_ij1, m_ij2, m_ij3  <-- message passing
   Attention: compute α_ij1, α_ij2, α_ij3  <-- per message
   Weighted messages: α_ij * m_ij
   Aggregate: sum/mean → updated node embedding h_i