# Nearest Neighbours Graph Variational AutoEncoder
## 2 Materials and Methods
### Datasets  
#### 2.2.1 Energy Deposition Datasets  
The architecture presented in this work was developed for a specific application in Medical Physics, which is the generation of the distribution of the dose absorbed by a  
Figure 3: **Full scheme. ** Schematic representation of our Re-NN Graph VAE with two encoding blocks and two decoding blocks. Each block is made up of a graph convolution and a pooling (un-pooling) operation. In the lower part of the picture, the VAE sampling and the encoding (decoding) linear layers are represented. medium interacting with a given particle beam, conditioned to beam parameters and medium density. Specifically, the approach was developed for a frontier application in radiation oncology therapy, which makes use of high-energy electron beams (Flash Therapy [31]). The datasets for this task are built simulating an electron beam interacting with a material using Geant4 [32], a toolkit in C++ for Monte Carlo simulations in particle physics. The two datasets differ on the material in which electrons deposit their energy. In the first case, this material is a cubic volume filled with water. We will refer to this dataset as \"Water Dataset\". In the second case, to increase the complexity of the task, we inserted a slab of variable density in the water volume. This slab is orthogonal to the electrons' beam and has a fixed position and thickness. The density of the slab is uniformly sampled at each run of the simulation between 0 and 5 g/cm\({}^{3}\) (for reference, water density is 1 g/cm\({}^{3}\)). We will refer to this dataset as \"Water + Slab Dataset\". In both cases, the particles' energies are sampled uniformly between 50 and 100 MeV, which is the typical range of energies for the FLASH radiotherapy with high-energy electrons. Energy deposition data are collected in a cylindrical scorer, aligned with the electron beam, and divided in 28 \(\times\) 28 \(\times\) 28 voxels along \(z\), \(\theta\) and \(r\) axes. The cylindrical shape is particularly useful in our application because it allows for higher precision near the beamline. Each example in the dataset is therefore a set of 28\({}^{3}\) voxels arranged in a cylindrical shape. Voxels have only one feature, and correspond to the amount of energy deposited in them by the simulated particle beam. Each example of the Water Dataset is labelled by the initial energy of the electron beam, while in the other dataset examples are labelled by both the particles' initial energy and the slab's density. An illustration of a typical example from these datasets is shown in Figure 4. Besides the representation of the original data in the left panel, we also show how the ReNN-Pool operates on nodes. As it is possible to see, the nodes' pooling is conducted in a spatial and uniform way. Datasets, respectively, consist of 4662 and 6239 examples and are divided in train, validation and test sets on the basis of particle energy and slab density. In particular, in the Water Dataset, the test set is made up of examples with a particle's energy ranging between 70 and 80 MeV. In the Water + Slab Dataset test set, the examples have the same range of initial energies and slab density values, ranging between 2 and 3 g/cm\({}^{3}\). In both cases, the rest of the dataset is used for validation and trains with a ratio of 1/10. Test sets have been chosen in this way in order to test the network ability to interpolate between samples and generalise to unseen configurations. For both datasets, we imposed a graph structure on the data. Each voxels was associated with a node and nodes were linked within each other with a nearest neighbours  
Figure 4: **Dataset and ReNN-Pool**. From the left, the panels show the representation of a typical example from the energy deposition datasets and two pooled representations of the same example. The nodes in light grey have null features, while all others show an energy distribution within the considered range. connectivity.