1import os
2import oddt
3from oddt.interactions import close_contacts, hbonds
4
5test_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..')
6
7class BenchInteractions(object):
8    """Spatial functions"""
9
10    def setup(self):
11        self.mols = list(oddt.toolkit.readfile('sdf', '%s/tests/data/dude/xiap/actives_docked.sdf' % test_data_dir))[:10]
12        for mol in self.mols:
13            mol.addh(only_polar=True)
14            mol.atom_dict
15        self.protein = list(oddt.toolkit.readfile('pdb', '%s/tests/data/dude/xiap/receptor_rdkit.pdb' % test_data_dir))[0]
16        self.protein.atom_dict
17        self.protein.addh(only_polar=True)
18
19    def time_close_contacts(self):
20        for mol in self.mols:
21            close_contacts(mol.atom_dict, self.protein.atom_dict, cutoff=10.)
22
23    def peakmem_close_contacts(self):
24        self.time_close_contacts()
25
26    def time_hbonds(self):
27        for mol in self.mols:
28            hbonds(mol, self.protein)
29
30    def peakmem_hbonds(self):
31        self.time_hbonds()
32