Implementing neighbor lists instead of distances

[MihailBogojeski]

WIP: I started the implementation of neighbor lists for our model, and incorporated the neighbor lists into get_distances and get_atom_edge_encodings and get_residue_edge_encodings.

Issue: After struggling for way too many hours, I still haven't figured out what the variable closest_neighbors is supposed to do, so I just assumed it lists the closest neighbor that is not in the same protein as the current atom/residue, but this seems to be false because my implementation doesn't match the old one.

What needs to be done:

1. Checking where the adjacency tensor is used (I saw multiple places in data_loader.py), and adjusting them so that they work with the neighbor lists instead of the full quadratic matrix.

2. When combining the different proteins into batches, the distances/edge features need to be stacked together so that there is no batch dimension, and the edges of all proteins are kept in the same dimension. The index lists also need to be modified with the proper offset in order to avoid overlap, here's a code snippet given a list of index lists corresponding to different proteins in a batch:

    ```
    prev_max = 0
    # each entry in list_idx_is corresponds to an index list for a protein in the batch
    for i in range(len(list_idx_is)):
        list_idx_is[i] += prev_max
        list_idx_js[i] += prev_max
        max_i = torch.max(idx_is[i])
        max_j = torch.max(idx_is[i])
        prev_max = max(max_i, max_j) + 1
   
    stacked_idx_is = torch.cat(idx_is, dim=0)
    stacked_idx_js = torch.cat(idx_js, dim=0)
    # proteins contains a list of the different proteins in the batch
    # create structure to keep track of which atoms/residues belong to which protein/batch
    protein_batch_idx = []
    for i in range(len(proteins)):
          protein_batch_idx.append(torch.ones(proteins[i].shape[0]) * i)
   
    proteins = torch.cat(proteins, dim=0)
    protein_batch_idx = torch.cat(protein_batch_idx).type(torch.LongTensor)
   

[moritzschaefer]

I'm on it
(PS: It's the wrong repository ;))