Update utils.py

utils.py

@@ -1,17 +1,3 @@
-#   Copyright 2021 Gabriele Orlando
-#
-#   Licensed under the Apache License, Version 2.0 (the "License");
-#   you may not use this file except in compliance with the License.
-#   You may obtain a copy of the License at
-#
-#       http://www.apache.org/licenses/LICENSE-2.0
-#
-#   Unless required by applicable law or agreed to in writing, software
-#   distributed under the License is distributed on an "AS IS" BASIS,
-#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#   See the License for the specific language governing permissions and
-#   limitations under the License.
-
 import os,torch
 from pyuul.sources.globalVariables import *
 from pyuul.sources import  hashings
@@ -30,20 +16,16 @@ setup_seed(100)
 def parseSDF(SDFFile):
 	"""
 	function to parse pdb files. It can be used to parse a single file or all the pdb files in a folder. In case a folder is given, the coordinates are gonna be padded
-
 	Parameters
 	----------
 	SDFFile : str
 		path of the PDB file or of the folder containing multiple PDB files
-
 	Returns
 	-------
 	coords : torch.Tensor
 		coordinates of the atoms in the pdb file(s). Shape ( batch, numberOfAtoms, 3)
-
 	atomNames : list
 		a list of the atom identifier. It encodes atom type, residue type, residue position and chain
-
 	"""
 	if not os.path.isdir(SDFFile):
 		fil = SDFFile
@@ -98,7 +80,6 @@ def parsePDB(PDBFile,keep_only_chains=None,keep_hetatm=True,bb_only=False):
 
 	"""
 	function to parse pdb files. It can be used to parse a single file or all the pdb files in a folder. In case a folder is given, the coordinates are gonna be padded
-
 	Parameters
 	----------
 	PDBFile : str
@@ -113,10 +94,8 @@ def parsePDB(PDBFile,keep_only_chains=None,keep_hetatm=True,bb_only=False):
 	-------
 	coords : torch.Tensor
 		coordinates of the atoms in the pdb file(s). Shape ( batch, numberOfAtoms, 3)
-
 	atomNames : list
 		a list of the atom identifier. It encodes atom type, residue type, residue position and chain
-
 	"""
 
 	bbatoms = ["N", "CA", "C"]
@@ -200,31 +179,26 @@ def parsePDB(PDBFile,keep_only_chains=None,keep_hetatm=True,bb_only=False):
 					atomNamesTMP += ["HET_"+str(resnum)+"_"+atnameHet+"_"+line[21]]
 			coords+=[torch.tensor(coordsTMP)]
 			atomNames += [atomNamesTMP]
-        print(atomNames)
-        print(pdbname)
-        print(pdb_num)
+		print(atomNames)
+		print(pdbname)
+		print(pdb_num)
 		return torch.torch.nn.utils.rnn.pad_sequence(coords, batch_first=True, padding_value=PADDING_INDEX), atomNames, pdbname, pdb_num
 
 
 def atomlistToChannels(atomNames, hashing="Element_Hashing", device="cpu"):
 	"""
 	function to get channels from atom names (obtained parsing the pdb files with the parsePDB function)
-
 	Parameters
 	----------
 	atomNames : list
 		atom names obtained parsing the pdb files with the parsePDB function
-
 	hashing : "TPL_Hashing" or "Element_Hashing" or dict
 		define which atoms are grouped together. You can use two default hashings or build your own hashing:
-
 		TPL_Hashing: uses the hashing of torch protein library (https://github.com/lupoglaz/TorchProteinLibrary)
 		Element_Hashing: groups atoms in accordnce with the element only: C -> 0, N -> 1, O ->2, P ->3, S- >4, H ->5, everything else ->6
-
 		Alternatively, if you are not happy with the default hashings, you can build a dictionary of dictionaries that defines the channel of every atom type in the pdb.
 		the first dictionary has the residue tag (three letters amino acid code) as key (3 letters compound name for hetero atoms, as written in the PDB file)
 		every residue key is associated to a dictionary, which the atom tags (as written in the PDB files) as keys and the channel (int) as value
-
 		for example, you can define the channels just based on the atom element as following:
 		{
 		'CYS': {'N': 1, 'O': 2, 'C': 0, 'SG': 3, 'CB': 0, 'CA': 0}, # channels for cysteine atoms
@@ -233,21 +207,16 @@ def atomlistToChannels(atomNames, hashing="Element_Hashing", device="cpu"):
 		'GOL': {'O1':2,'O2':2,'O3':2,'C1':0,'C2':0,'C3':0}, # channels for glycerol atom
 		...
 		}
-
 		The default encoding is the one that assigns a different channel to each element
-
 		other encodings can be found in sources/hashings.py
-
 	device : torch.device
 			The device on which the model should run. E.g. torch.device("cuda") or torch.device("cpu:0")
 	Returns
 	-------
 	coords : torch.Tensor
 		coordinates of the atoms in the pdb file(s). Shape ( batch, numberOfAtoms, 3)
-
 	channels : torch.tensor
 		the channel of every atom. Shape (batch,numberOfAtoms)
-
 	"""
 	if hashing == "TPL_Hashing":
 		hashing = hashings.TPLatom_hash
@@ -295,20 +264,15 @@ def atomlistToChannels(atomNames, hashing="Element_Hashing", device="cpu"):
 def atomlistToRadius(atomList, hashing="FoldX_radius", device="cpu"):
 	"""
 	function to get radius from atom names (obtained parsing the pdb files with the parsePDB function)
-
-
-
 	Parameters
 	----------
 	atomNames : list
 		atom names obtained parsing the pdb files with the parsePDB function
 	hashing : FoldX_radius or dict
 		"FoldX_radius" provides the radius used by the FoldX force field
-
 		Alternatively, if you are not happy with the foldX radius, you can build a dictionary of dictionaries that defines the radius of every atom type in the pdb.
 		The first dictionary has the residue tag (three letters amino acid code) as key (3 letters compound name for hetero atoms, as written in the PDB file)
 		every residue key is associated to a dictionary, which the atom tags (as written in the PDB files) as keys and the radius (float) as value
-
 		for example, you can define the radius as following:
 		{
 		'CYS': {'N': 1.45, 'O': 1.37, 'C': 1.7, 'SG': 1.7, 'CB': 1.7, 'CA': 1.7}, # radius for cysteine atoms
@@ -317,21 +281,16 @@ def atomlistToRadius(atomList, hashing="FoldX_radius", device="cpu"):
 		'GOL': {'O1':1.37,'O2':1.37,'O3':1.37,'C1':1.7,'C2':1.7,'C3':1.7}, # radius for glycerol atoms
 		...
 		}
-
 		The default radius are the ones defined in FoldX
-
 		Radius default dictionary can be found in sources/hashings.py
-
 	device : torch.device
 			The device on which the model should run. E.g. torch.device("cuda") or torch.device("cpu:0")
 	Returns
 	-------
 	coords : torch.Tensor
 		coordinates of the atoms in the pdb file(s). Shape ( batch, numberOfAtoms, 3)
-
 	radius : torch.tensor
 		The radius of every atom. Shape (batch,numberOfAtoms)
-
 	"""
 	if hashing == "FoldX_radius":
 		hashing = hashings.radius
@@ -359,45 +318,36 @@ def atomlistToRadius(atomList, hashing="FoldX_radius", device="cpu"):
 
 '''
 def write_pdb(batchedCoords, atomNames , name=None, output_folder="outpdb/"): #I need to add the chain id
-
 	if name is None:
 		name = range(len(batchedCoords))
-
 	for struct in range(len(name)):
 		f = open(output_folder + str(name[struct]) + ".pdb", "w")
-
 		coords=batchedCoords[struct].data.numpy()
 		atname=atomNames[struct]
 		for i in range(len(coords)):
-
 			rnName = atname[i].split("_")[0]#hashings.resi_hash_inverse[resi_list[i]]
 			atName = atname[i].split("_")[2]#hashings.atom_hash_inverse[resi_list[i]][atom_list[i]]
 			pos = atname[i].split("_")[1]
 			chain = "A"
-
 			num = " " * (5 - len(str(i))) + str(i)
 			a_name = atName + " " * (4 - len(atName))
 			numres = " " * (4 - len(str(pos))) + str(pos)
-
 			x = round(float(coords[i][0]), 3)
 			sx = str(x)
 			while len(sx.split(".")[1]) < 3:
 				sx += "0"
 			x = " " * (8 - len(sx)) + sx
-
 			y = round(float(coords[i][1]), 3)
 			sy = str(y)
 			while len(sy.split(".")[1]) < 3:
 				sy += "0"
 			y = " " * (8 - len(sy)) + sy
-
 			z = round(float(coords[i][2]), 3)
 			sz = str(z)
 			while len(sz.split(".")[1]) < 3:
 				sz += "0"
 			z = " " * (8 - len(sz)) + sz
 			chain = " " * (2 - len(chain)) + chain
-
 			if rnName !="HET":
 				f.write("ATOM  " + num + "  " + a_name + "" + rnName + chain + numres + "    " + x + y + z + "  1.00 64.10           " + atName[0] + "\n")
 			else: