--- a
+++ b/utils.py
@@ -0,0 +1,234 @@
+from typing import Union, Iterable
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from rdkit import Chem
+import networkx as nx
+from networkx.algorithms import isomorphism
+from Bio.PDB.Polypeptide import is_aa
+
+
+class Queue():
+    def __init__(self, max_len=50):
+        self.items = []
+        self.max_len = max_len
+
+    def __len__(self):
+        return len(self.items)
+
+    def add(self, item):
+        self.items.insert(0, item)
+        if len(self) > self.max_len:
+            self.items.pop()
+
+    def mean(self):
+        return np.mean(self.items)
+
+    def std(self):
+        return np.std(self.items)
+
+
+def reverse_tensor(x):
+    return x[torch.arange(x.size(0) - 1, -1, -1)]
+
+
+#####
+
+
+def get_grad_norm(
+        parameters: Union[torch.Tensor, Iterable[torch.Tensor]],
+        norm_type: float = 2.0) -> torch.Tensor:
+    """
+    Adapted from: https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html#clip_grad_norm_
+    """
+
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+
+    norm_type = float(norm_type)
+
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+
+    device = parameters[0].grad.device
+
+    total_norm = torch.norm(torch.stack(
+        [torch.norm(p.grad.detach(), norm_type).to(device) for p in
+         parameters]), norm_type)
+
+    return total_norm
+
+
+def write_xyz_file(coords, atom_types, filename):
+    out = f"{len(coords)}\n\n"
+    assert len(coords) == len(atom_types)
+    for i in range(len(coords)):
+        out += f"{atom_types[i]} {coords[i, 0]:.3f} {coords[i, 1]:.3f} {coords[i, 2]:.3f}\n"
+    with open(filename, 'w') as f:
+        f.write(out)
+
+
+def write_sdf_file(sdf_path, molecules):
+    # NOTE Changed to be compatitble with more versions of rdkit
+    #with Chem.SDWriter(str(sdf_path)) as w:
+    #    for mol in molecules:
+    #        w.write(mol)
+
+    w = Chem.SDWriter(str(sdf_path))
+    w.SetKekulize(False)
+    for m in molecules:
+        if m is not None:
+            w.write(m)
+
+    # print(f'Wrote SDF file to {sdf_path}')
+
+
+def residues_to_atoms(x_ca, atom_encoder):
+    x = x_ca
+    one_hot = F.one_hot(
+        torch.tensor(atom_encoder['C'], device=x_ca.device),
+        num_classes=len(atom_encoder)
+    ).repeat(*x_ca.shape[:-1], 1)
+    return x, one_hot
+
+
+def get_residue_with_resi(pdb_chain, resi):
+    res = [x for x in pdb_chain.get_residues() if x.id[1] == resi]
+    assert len(res) == 1
+    return res[0]
+
+
+def get_pocket_from_ligand(pdb_model, ligand, dist_cutoff=8.0):
+
+    if ligand.endswith(".sdf"):
+        # ligand as sdf file
+        rdmol = Chem.SDMolSupplier(str(ligand))[0]
+        ligand_coords = torch.from_numpy(rdmol.GetConformer().GetPositions()).float()
+        resi = None
+    else:
+        # ligand contained in PDB; given in <chain>:<resi> format
+        chain, resi = ligand.split(':')
+        ligand = get_residue_with_resi(pdb_model[chain], int(resi))
+        ligand_coords = torch.from_numpy(
+            np.array([a.get_coord() for a in ligand.get_atoms()]))
+
+    pocket_residues = []
+    for residue in pdb_model.get_residues():
+        if residue.id[1] == resi:
+            continue  # skip ligand itself
+
+        res_coords = torch.from_numpy(
+            np.array([a.get_coord() for a in residue.get_atoms()]))
+        if is_aa(residue.get_resname(), standard=True) \
+                and torch.cdist(res_coords, ligand_coords).min() < dist_cutoff:
+            pocket_residues.append(residue)
+
+    return pocket_residues
+
+
+def batch_to_list(data, batch_mask):
+    # data_list = []
+    # for i in torch.unique(batch_mask):
+    #     data_list.append(data[batch_mask == i])
+    # return data_list
+
+    # make sure batch_mask is increasing
+    idx = torch.argsort(batch_mask)
+    batch_mask = batch_mask[idx]
+    data = data[idx]
+
+    chunk_sizes = torch.unique(batch_mask, return_counts=True)[1].tolist()
+    return torch.split(data, chunk_sizes)
+
+
+def num_nodes_to_batch_mask(n_samples, num_nodes, device):
+    assert isinstance(num_nodes, int) or len(num_nodes) == n_samples
+
+    if isinstance(num_nodes, torch.Tensor):
+        num_nodes = num_nodes.to(device)
+
+    sample_inds = torch.arange(n_samples, device=device)
+
+    return torch.repeat_interleave(sample_inds, num_nodes)
+
+
+def rdmol_to_nxgraph(rdmol):
+    graph = nx.Graph()
+    for atom in rdmol.GetAtoms():
+        # Add the atoms as nodes
+        graph.add_node(atom.GetIdx(), atom_type=atom.GetAtomicNum())
+
+    # Add the bonds as edges
+    for bond in rdmol.GetBonds():
+        graph.add_edge(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx())
+
+    return graph
+
+
+def calc_rmsd(mol_a, mol_b):
+    """ Calculate RMSD of two molecules with unknown atom correspondence. """
+    graph_a = rdmol_to_nxgraph(mol_a)
+    graph_b = rdmol_to_nxgraph(mol_b)
+
+    gm = isomorphism.GraphMatcher(
+        graph_a, graph_b,
+        node_match=lambda na, nb: na['atom_type'] == nb['atom_type'])
+
+    isomorphisms = list(gm.isomorphisms_iter())
+    if len(isomorphisms) < 1:
+        return None
+
+    all_rmsds = []
+    for mapping in isomorphisms:
+        atom_types_a = [atom.GetAtomicNum() for atom in mol_a.GetAtoms()]
+        atom_types_b = [mol_b.GetAtomWithIdx(mapping[i]).GetAtomicNum()
+                        for i in range(mol_b.GetNumAtoms())]
+        assert atom_types_a == atom_types_b
+
+        conf_a = mol_a.GetConformer()
+        coords_a = np.array([conf_a.GetAtomPosition(i)
+                             for i in range(mol_a.GetNumAtoms())])
+        conf_b = mol_b.GetConformer()
+        coords_b = np.array([conf_b.GetAtomPosition(mapping[i])
+                             for i in range(mol_b.GetNumAtoms())])
+
+        diff = coords_a - coords_b
+        rmsd = np.sqrt(np.mean(np.sum(diff * diff, axis=1)))
+        all_rmsds.append(rmsd)
+
+    if len(isomorphisms) > 1:
+        print("More than one isomorphism found. Returning minimum RMSD.")
+
+    return min(all_rmsds)
+
+
+class AppendVirtualNodes:
+    def __init__(self, max_ligand_size, atom_encoder, symbol):
+        self.max_ligand_size = max_ligand_size
+        self.atom_encoder = atom_encoder
+        self.vidx = atom_encoder[symbol]
+
+    def __call__(self, data):
+
+        n_virt = self.max_ligand_size - data['num_lig_atoms']
+        mu = data['lig_coords'].mean(0, keepdim=True)
+        sigma = data['lig_coords'].std(0).max()
+        virt_coords = torch.randn(n_virt, 3) * sigma + mu
+
+        # insert virtual atom column
+        one_hot = torch.cat((data['lig_one_hot'][:, :self.vidx],
+                            torch.zeros(data['num_lig_atoms'])[:, None],
+                            data['lig_one_hot'][:, self.vidx:]), dim=1)
+        virt_one_hot = torch.zeros(n_virt, len(self.atom_encoder))
+        virt_one_hot[:, self.vidx] = 1
+        virt_mask = torch.ones(n_virt) * data['lig_mask'][0]
+
+        data['lig_coords'] = torch.cat((data['lig_coords'], virt_coords))
+        data['lig_one_hot'] = torch.cat((one_hot, virt_one_hot))
+        data['num_lig_atoms'] = self.max_ligand_size
+        data['lig_mask'] = torch.cat((data['lig_mask'], virt_mask))
+        data['num_virtual_atoms'] = n_virt
+
+        return data