diff --git a/.gitignore b/.gitignore index c712383..fc5e790 100644 --- a/.gitignore +++ b/.gitignore @@ -112,6 +112,7 @@ autodocksuite-4.2.6-src/autogrid/x86_Linux2/ *.xyz *.txt *.pdb +*.pdbqt *.pse *.csv *.dx diff --git a/scripts/score_frame_waters.py b/scripts/score_frame_waters.py new file mode 100644 index 0000000..896156c --- /dev/null +++ b/scripts/score_frame_waters.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# WaterKit +# +# Launch waterkit +# + +import os +import argparse +import shutil + +from waterkit import AutoGrid +from waterkit import Map +from waterkit import Molecule +from waterkit import WaterBox +from waterkit import WaterKit +from waterkit import Water +from waterkit import utils +from vina import Vina + + +def cmd_lineparser(): + parser = argparse.ArgumentParser(description='waterkit') + parser.add_argument('-w', '--frame', help='input frame with waters [.pdb]', required=True) + parser.add_argument('-i', '--mol', dest='receptor_pdbqtfilename', required=True, + action='store', help='receptor file') + parser.add_argument('-c', '--center', dest='box_center', nargs=3, type=float, + action='store', help='center of the box') + parser.add_argument('-s', '--size', dest='box_size', nargs=3, type=int, required=True, + action='store', help='size of the box in Angstrom') + parser.add_argument('--autogrid_exec_path', dest='autogrid_exec_path', default='autogrid4', + action='store', help='path to the autogrid4 executable (default: autogrid4') + return parser.parse_args() + + +if __name__ == "__main__": + args = cmd_lineparser() + receptor_pdbqtfilename = args.receptor_pdbqtfilename + water_frame_filename = args.frame + box_center = args.box_center + box_size = args.box_size + autogrid_exec_path = args.autogrid_exec_path + water_model = 'tip3p' + spherical_water_maps = (None, None) + temperature = None # we don't need it here + + print(f"{box_size=}") + + # Read PDBQT/MOL2 file, Waterfield file and AutoDock grid map + receptor = Molecule.from_file(receptor_pdbqtfilename) + + with utils.temporary_directory(prefix='wk_', dir='.', clean=False) as tmp_dir: + # Generate AutoDock maps using the Amber ff14SB forcefield + receptor.to_pdbqt_file('receptor.pdbqt') + ff14sb_param_file = os.path.join(utils.path_module('waterkit'), 'data/ff14SB_parameters.dat') + ag = AutoGrid(autogrid_exec_path, ff14sb_param_file) + ad_map = ag.run('receptor.pdbqt', ['OW'], box_center, box_size, smooth=0, dielectric=1) + + if spherical_water_maps[0] is None: + # Convert amber atom types to AutoDock atom types + ad_receptor = utils.convert_amber_to_autodock_types(receptor) + ad_receptor.to_pdbqt_file('receptor_ad.pdbqt') + + # Generate Vina maps for the spherical maps + v = Vina(verbosity=0) + v.set_receptor('receptor_ad.pdbqt') + v.compute_vina_maps(box_center, box_size, force_even_voxels=False) + v.write_maps('vina') + sw_map = Map('vina.O_DA.map', 'SW') + else: + raise NotImplementedError("hard-coded to use vina's O_DA map as spherical water map") + + ad_map.add_map('SW', sw_map._maps['SW']) + + # It is more cleaner if we prepare the maps (OW, HW for tip3p, OT, HT, LP for tip5p) before + utils.prepare_water_map(ad_map, water_model) + + waters, reskeys = Water.from_file(water_frame_filename) + nr_waters = len(waters) + for index in range(nr_waters): + + # add other waters to maps + this_water = waters[index] + other_waters = waters[:index] + waters[(index + 1):] + # WaterBox init makes a copy of ad_map internally in self.map + wbox = WaterBox(receptor, ad_map, temperature, water_model, spherical_water_maps[1]) + for other_water in other_waters: + + wbox._wopt._update_maps(other_water) + + # compute energy of current water + oxygen_xyz = this_water.coordinates(1) + print(oxygen_xyz) + water_info = this_water.atom_informations() + energy = wbox.map.energy_coordinates(oxygen_xyz, water_info["t"][0])[0] + print(f"{index:3} {reskeys[index]:6} {energy:12.6f} O (rec+wat)") + energy_rec = ad_map.energy_coordinates(oxygen_xyz, water_info["t"][0])[0] + print(f"{index:3} {reskeys[index]:6} {energy_rec:12.6f} O (just rec)") + for i, atom_type in enumerate(water_info["t"][1:]): + # i + 2 because 1-indexed and because we skipped first element in enumerate + e = wbox.map.energy_coordinates(this_water.coordinates(i+2), atom_type)[0] + energy += e + print(f"{index:3} {reskeys[index]:6} {e:12.6f} H (rec+wat)") + + er = ad_map.energy_coordinates(this_water.coordinates(i+2), atom_type)[0] + energy_rec += er + print(f"{index:3} {reskeys[index]:6} {er:12.6f} H (just rec)") + + print(f"{index:3} {reskeys[index]:6} {energy:12.6f} HOH (rec+wat)") + print(f"{index:3} {reskeys[index]:6} {energy_rec:12.6f} HOH (just rec)") + print() + diff --git a/setup.py b/setup.py index d2b485d..5dc5f3c 100644 --- a/setup.py +++ b/setup.py @@ -31,6 +31,7 @@ def find_files(directory): "scripts/wk_minimize_trajectory.py", "scripts/wk_get_spherical_map.py", "scripts/wk_generate_gaff2_parameters.py"], + "scripts/score_frame_waters.py"], package_data={ "waterkit" : ["data/*", "data/water/spherical/*", diff --git a/waterkit/sampling.py b/waterkit/sampling.py index d5af605..57046c6 100644 --- a/waterkit/sampling.py +++ b/waterkit/sampling.py @@ -285,7 +285,7 @@ def _update_maps(self, water): """ spacing = self._ad_map._spacing ad_map = self._ad_map - boxsize = np.array([8, 8, 8]) + boxsize = np.array([15, 15, 15]) npts = np.round(boxsize / spacing).astype(int) // 2 * 2 + 1 water_xyz = water.coordinates() diff --git a/waterkit/water.py b/waterkit/water.py index 4a9cec4..a3528f6 100644 --- a/waterkit/water.py +++ b/waterkit/water.py @@ -56,7 +56,7 @@ def __init__(self, xyz, atom_type="W", partial_charge=0., hb_anchor=None, hb_vec self._add_atom(xyz, atom_type, partial_charge) @classmethod - def from_file(cls, fname, atom_type="W", partial_charge=0.): + def from_file(cls, fname, water_model="tip3p"): """Create list of Water objects from a PDB file. The water molecules created are spherical. @@ -70,26 +70,105 @@ def from_file(cls, fname, atom_type="W", partial_charge=0.): list: list of Water molecule objects """ - waters = [] - - # Get name and file extension - name, file_extension = os.path.splitext(fname) - - if file_extension == ".pdbqt": - file_extension = "pdb" + if water_model != "tip3p": + raise NotImplementedError("hard-coded tip3p in Water.from_file") - # Read PDB file - obconv = ob.OBConversion() - obconv.SetInFormat(file_extension) - OBMol = ob.OBMol() - obconv.ReadFile(OBMol, fname) + waters = [] - for x in ob.OBMolAtomIter(OBMol): - if x.IsOxygen(): - xyz = np.array([x.GetX(), x.GetY(), x.GetZ()]) - waters.append(cls(xyz, atom_type, partial_charge)) + with open(fname) as f: + pdb_string = f.read() - return waters + def _add_if_new(to_dict, key, value, repeat_log): + if key in to_dict: + repeat_log.add(key) + else: + to_dict[key] = value + return + + blocks_by_residue = {} + reskey_to_resname = {} + reskey = None + buffered_reskey = None + interrupted_residues = set() + pdb_block = [] + + for line in pdb_string.splitlines(True): + if line.startswith("TER") and reskey is not None: + _add_if_new(blocks_by_residue, reskey, pdb_block, interrupted_residues) + blocks_by_residue[reskey] = pdb_block + pdb_block = [] + reskey = None + buffered_reskey = None + if line.startswith("ATOM") or line.startswith("HETATM"): + atomname = line[12:16].strip() + altloc = line[16:17].strip() + resname = line[17:20].strip() + chainid = line[21:22].strip() + resnum = int(line[22:26].strip()) + icode = line[26:27].strip() + x = float(line[30:38]) + y = float(line[38:46]) + z = float(line[46:54]) + element = line[76:78].strip() + reskey = f"{chainid}:{resnum}{icode}" # e.g. ":42", "A:42B" + reskey_to_resname.setdefault(reskey, set()) + reskey_to_resname[reskey].add(resname) + atom = ( + atomname, altloc, resname, chainid, + resnum, icode, x, y, z, element, + ) + + if reskey == buffered_reskey: # this line continues existing residue + pdb_block.append(atom) + else: + if buffered_reskey is not None: + _add_if_new( + blocks_by_residue, + buffered_reskey, + pdb_block, + interrupted_residues, + ) + buffered_reskey = reskey + pdb_block = [atom] + + if pdb_block: # there was not a TER line + _add_if_new(blocks_by_residue, reskey, pdb_block, interrupted_residues) + + # verify that each identifier (e.g. "A:17" has a single resname + violations = {k: v for k, v in reskey_to_resname.items() if len(v) != 1} + if len(violations): + msg = "each residue key must have exactly 1 resname" + eol + msg += f"but got {violations=}" + raise ValueError(msg) + + if interrupted_residues: + msg = f"interrupted residues in PDB: {interrupted_residues}" + raise ValueError(msg) + + for reskey, pdb_block in blocks_by_residue.items(): + if list(reskey_to_resname[reskey])[0] not in ["HOH", "WAT"]: + continue + xyz = [] + # find oxygen index + oxygen_index = None + for i, atom in enumerate(pdb_block): + if atom[9] == "O": + oxygen_index = i + if oxygen_index is None: + raise RuntimeError("got water without element O") + oxygen = pdb_block[oxygen_index] + xyz = np.array([oxygen[6], oxygen[7], oxygen[8]]) + water = cls(xyz, "OW", -0.834) + for i, atom in enumerate(pdb_block): + if i == oxygen_index: + continue + if atom[9] != "H": + raise RuntimeError("expected hydrogen atom in water") + xyz = np.array([atom[6], atom[7], atom[8]]) + water._add_atom(xyz, "HW", +0.417) + waters.append(water) + + return waters, list(blocks_by_residue.keys()) def _add_atom(self, xyz, atom_type, partial_charge): """Add an atom to the molecule."""