# Rearrangement of Single Atoms by Solving Assignment Problems via Convolutional Neural Network
## 3 Convolutional Neural Network Model Formulation
#### Understand the Data  
The following example demonstrates a model specifically designed for rearranging atoms within a general \(n\times n\) grid size. In this example, let us consider an \(n=21\), containing a total of \(n^{2}=441\) positions, with a specific number of atoms k, as shown in Figure 5. Figure 5 represents the target atom arrangement, where '1' indicates the presence of an atom, confined to a central square of \(m\times m\) dimensions, surrounded by zeros. Figure 5a illustrates an example of the input dataset or the initial state of loaded atoms, where '1's and '0's are randomly distributed, with k atoms and \(n^{2}-k\) empty spaces. Figure 5c depicts the output data, showing the positions where atoms need to be relocated. It consists of \(k+1\) distinct numbers, with k non-zero numbers denoting the coordinates of the central square of \(m\times m\) atom positions, and zero signifying empty spaces. Notably, positions with '0's in Figure 5a align with corresponding positions in Figure 5c, while positions with '1's in Figure 5a correlate with non-zero numbers in Figure 5c. The model aims to predict atom movements based on the input, ultimately achieving the target configuration depicted in Figure 5b. Step 1: Dataset preparation: Change the output shape (separate the output digits)  
The output data (Figure 5c) represent the positions where atoms need to move. Each number in the output data indicates a specific point within the grid. Each digit (hundreds, tens, ones) in these numbers holds the same importance, meaning if any digit is predicted incorrectly, it affects the entire predicted point. To enhance prediction accuracy, each output is transformed from n\({}^{2}\) numbers to either \(2n^{2}\) or \(3n^{2}\) numbers, depending on whether the original numbers are two-digit or three-digit numbers. For instance, a \(10\times 10\) grid size with two-digit numbers would result in \(2n^{2}\) numbers, while larger grids with three-digit numbers would result in \(3n^{2}\) numbers. Each set of two or three numbers represents a single point, with the separation of digits explicitly indicated. The assumption is made that '0' is represented as '000'. The modified output data consist of \(2n^{2}\) or \(3n^{2}\) positions accordingly. Specifically, the first \(n^{2}\) indices correspond to the hundreds (if applicable), followed by indicating the tens, and, finally, the last \(n^{2}\) positions represent the ones of the original output data. The dataset is split into training and testing sets with an 80:20 ratio to evaluate the model's performance. Figure 5: (a) Initial state of loaded atom (input of the model). (b) Target state of the atoms after movement. (c) Target coordinates of each atom that needs to move to achieve the target state (output of the model). Step 2: CNN model construction  
The Convolutional Neural Network (CNN) model utilized in this work is composed of several layers that, together, enable the network to learn and make accurate predictions from input data. The primary layer type is the 2D convolution layer, which is responsible for extracting features from input data. Each convolutional layer typically uses a kernel size of \(3\times 3\), with a stride of 1 and padding of 1. This configuration helps preserve the spatial dimensions of the input data while applying the filters to extract features. The number of input and output channels varies across the layers and is adjusted based on the complexity and size of the grid. In the early layers, the number of channels starts at a lower value and increases gradually in subsequent layers to capture more complex features. This increment continues up to a certain point, after which the number of channels is reduced in the later layers to decrease the model's complexity and computational load. Following the convolution layers are the linear layers. After the convolutional layers have extracted features from the input, the feature maps are flattened into a one-dimensional vector.