script to score frame of waters

.gitignore

@@ -112,6 +112,7 @@ autodocksuite-4.2.6-src/autogrid/x86_Linux2/
 *.xyz
 *.txt
 *.pdb
+*.pdbqt
 *.pse
 *.csv
 *.dx

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()
+

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/*",

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()
 

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."""