def test_load_mols(testdir, system, n_ligand, n_receptor, ext):
ligname = os.path.join(system, f"{system}_ligand.{ext}")
recname = os.path.join(system, f"{system}_protein.pdb")
ulig = mda.Universe(os.path.join(testdir, ligname))
urec = mda.Universe(os.path.join(testdir, recname))
assert len(ulig.atoms) == n_ligand
assert len(urec.atoms) == n_receptor
systems = loaders.load_mols(ligname, recname, testdir)
assert len(systems) == 1
system = systems[0]
assert len(system.atoms) == n_ligand + n_receptor
lig = system.select_atoms("resname LIG")
assert len(lig.atoms) == n_ligand
assert set(lig.resnames) == set(["LIG"])
assert set(lig.resnums) == set([1])
assert set(lig.resids) == set([1])
assert set(lig.record_types) == set(["HETATM"])
assert set(lig.segids) == set([""])
def test_load_mols_fail_rec(testdir):
system = "1a4r"
ligname = os.path.join(system, f"{system}_ligand.pdb")
recname = os.path.join(system, f"{system}_FAIL.pdb")
with pytest.raises(Exception):
system = loaders.load_mols(ligname, recname, testdir)
def test_select(testdir, system, distance, n_ligand, n_receptor, ext):
"""
Selection compared with PyMol selection:
sele byres PROTEIN within DISTANCE of LIGAND
"""
ligname = os.path.join(system, f"{system}_ligand.{ext}")
recname = os.path.join(system, f"{system}_protein.pdb")
systems = loaders.load_mols(ligname, recname, testdir)
assert len(systems) == 1
system = systems[0]
atoms, coordinates = loaders.select(system, distance)
assert len(atoms) == n_ligand + n_receptor
assert coordinates.shape == (n_ligand + n_receptor, 3)
def test_select_removeHs(testdir, system, distance, n_ligand, n_receptor, ext):
"""
Selection compared with PyMol selection:
sele byres PROTEIN within DISTANCE of LIGAND
Hydrogen atoms were counted by hand and removed
(H* does not select 1HD1 while *H* also selects CH2).
"""
ligname = os.path.join(system, f"{system}_ligand.{ext}")
recname = os.path.join(system, f"{system}_protein.pdb")
systems = loaders.load_mols(ligname, recname, testdir)
assert len(systems) == 1
system = systems[0]
atoms, coordinates = loaders.select(system, distance, removeHs=True)
# assert len(atoms) == n_ligand + n_receptor
assert coordinates.shape == (n_ligand + n_receptor, 3)
def test_load_mols(testdir, system, n_ligand, n_receptor, ext):
ligname = os.path.join(system, f"{system}_docking.{ext}")
ligfile = os.path.join(testdir, ligname)
recname = os.path.join(system, f"{system}_protein.pdb")
obmols = [obmol for obmol in pybel.readfile(ext, ligfile)]
systems = loaders.load_mols(ligname, recname, testdir)
assert len(systems) == 9
for system, obmol in zip(systems, obmols):
assert len(system.atoms) == n_ligand + n_receptor
lig = system.select_atoms("resname LIG")
assert len(lig.atoms) == n_ligand
# Check ligand coordinates
for idx, atom in enumerate(obmol):
assert np.allclose(lig.atoms.positions[idx], atom.coords)
# Load and apply atom to index mapping with amap
amap = utils.load_amap(args.amap)
n_species = len(amap)
# Load AEVComputer and model
AEVC = utils.loadAEVC(args.aev)
model = utils.loadmodel(args.model)
AEVC.to(device)
model.to(device)
with open(args.gradfile, "r") as f:
for line in tqdm.tqdm(f):
label, recfile, ligfile = line.split()
systems = loaders.load_mols(ligfile, recfile, args.datapaths)
# TODO: Allow multiple systems?
assert len(systems) == 1
system = systems[0]
# Select ligand and residues
# TODO: Unify selections
selection = system.select_atoms(
f"(byres (around {args.distance} (resname LIG))) or (resname LIG)"
)
if args.removeHs:
mask = selection.elements != "H"
sel_idxs = selection.ix[mask]