hdiff output

r32052/pdbparser.py 2017-03-08 13:30:07.593079468 +0000 r32051/pdbparser.py 2017-03-08 13:30:07.845082829 +0000
  1: #! /usr/bin/env python  1: #! /usr/bin/env python
  2:   2: 
  3: import numpy  3: import numpy
  4: import sys  4: import sys
  5:   5: 
  6: __author__ = [  6: __author__ = [
  7:         "Tristan Cranolini <tristan.cranolini@gmail.com>",  7:                 "Tristan Cranolini <tristan.cranolini@gmail.com>",
  8:         "Jonathan Barnoud <jonathan@barnoud.net>",  8:                 "Jonathan Barnoud <jonathan@barnoud.net>",
  9: ]  9: ]
 10:  10: 
 11:  
 12: class PDB(): 11: class PDB():
 13:     def __init__(self, coords_=None, resize=2000): 12:         def __init__(self, coords_=None, resize=2000):
 14:         self.atomnames = [] 13:                 self.atomnames = []
 15:         self.resnames = [] 14:                 self.resnames = []
 16:         self.resnum = [] 15:                 self.resnum = []
 17:         self.chainname = "" 16:                 self.chainname = ""
 18:  17:                 
 19:         self.resize = resize 18:                 self.resize = resize
 20:         # coords_ allows us to resize the coordinates array less often 19:                 # coords_ allows us to resize the coordinates array less often
 21:         if coords_ is None: 20:                 if coords_ is None:
 22:             self.coords_ = numpy.zeros((resize, 3)) 21:                         self.coords_ = numpy.zeros((resize, 3))
 23:             self.Natoms = 0 22:                         self.Natoms = 0
 24:         else: 23:                 else:
 25:             self.coords_ = numpy.array(coords_, dtype=float).reshape(-1, 3) 24:                         self.coords_ = numpy.array(coords_, dtype=float).reshape(-1,3)
 26:             self.Natoms = self.coords_.shape[0] 25:                         self.Natoms = self.coords_.shape[0]
 27:         self.Bfactors = [] 26:                 self.Bfactors = []
 28:  27:         
 29:     def get_coords(self): 28:         def get_coords(self):
 30:         return self.coords_[0:self.Natoms, ] 29:                 return self.coords_[0:self.Natoms,]
 31:  30:         
 32:     coords = property(get_coords) 31:         coords = property(get_coords)
 33:  32:         
 34:     def addatom(self, line): 33:         def addatom(self, line):                
 35:         self.atomnames.append(line[12:16]) 34:                 self.atomnames.append(line[12:16])
 36:         self.resnames.append(line[17:20]) 35:                 self.resnames.append(line[17:20])
 37:         self.resnum.append(int(line[22:26])) 36:                 self.resnum.append(int(line[22:26]))
 38:  37: 
 39:         if self.Natoms >= self.coords_.shape[0]: 38:                 if self.Natoms >= self.coords_.shape[0]:
 40:             self.coords_ = numpy.row_stack([self.coords_, numpy.zeros((self.resize,3))]) 39:                         self.coords_ = numpy.row_stack([self.coords_, numpy.zeros((self.resize,3))])
 41:  40: 
 42:         self.coords_[self.Natoms,:] = [line[30:38], line[38:46], line[46:54]] 41:                 self.coords_[self.Natoms,:] = [line[30:38], line[38:46], line[46:54]]
 43:         self.Natoms += 1 42:                 self.Natoms += 1
 44:         try: 43:                 try:
 45:             self.Bfactors.append(float(line[60:66])) 44:                         self.Bfactors.append(float(line[60:66]))
 46:         except (ValueError): 45:                 except (ValueError):
 47:             self.Bfactors.append(0.0) 46:                         self.Bfactors.append(0.0)
 48:      47:         
 49:     def insertatom(self, line, idx=None): 48:         def insertatom(self, line, idx=None):
 50:         if idx is None: 49:                 if idx is None:
 51:             idx = self.Natoms 50:                         idx = self.Natoms
 52:          51:                 
 53:         self.atomnames.insert(idx, line[12:16]) 52:                 self.atomnames.insert(idx, line[12:16])
 54:         self.resnames.insert(idx, line[17:20]) 53:                 self.resnames.insert(idx, line[17:20])
 55:         self.resnum.insert(idx, int(line[22:26])) 54:                 self.resnum.insert(idx, int(line[22:26]))
 56:  55: 
 57:         if self.Natoms >= self.coords_.shape[0]: 56:                 if self.Natoms >= self.coords_.shape[0]:
 58:             self.coords_ = numpy.row_stack([self.coords_, numpy.zeros((self.resize,3))]) 57:                         self.coords_ = numpy.row_stack([self.coords_, numpy.zeros((self.resize,3))])
 59:  58: 
 60:         self.coords_[idx+1:self.Natoms+1,:] = self.coords_[idx:self.Natoms,:] 59:                 self.coords_[idx+1:self.Natoms+1,:] = self.coords_[idx:self.Natoms,:]
 61:         self.coords_[idx,:] = [float(line[30:38]), float(line[38:46]), float(line[46:54])] 60:                 self.coords_[idx,:] = [float(line[30:38]), float(line[38:46]), float(line[46:54])]
 62:         self.Natoms += 1 61:                 self.Natoms += 1
 63:         try: 62:                 try:
 64:             self.Bfactors.insert(idx, float(line[60:66])) 63:                         self.Bfactors.insert(idx, float(line[60:66]))
 65:         except (ValueError): 64:                 except (ValueError):
 66:             self.Bfactors.insert(idx, 0.0) 65:                         self.Bfactors.insert(idx, 0.0)
 67:  66: 
 68:  67: 
 69:     def center(self): 68:         def center(self):
 70:         return sum(self.coords) / self.coords.shape[0] 69:                 return sum(self.coords) / self.coords.shape[0]
 71:  70: 
 72:     def movetocenter(self): 71:         def movetocenter(self):
 73:         self.coords -= self.center() 72:                 self.coords -= self.center()
 74:      73:         
 75:     def pdbline(self, idx): 74:         def pdbline(self, idx):
 76:         #return "%-6s%5d %4s%1s%3s %1s%4d%1s        %8.3f%8.3f%8.3f%6.2f%6.2f                     %2s%2s\n" \ 75:                 return "%-6s%5d %4s%1s%3s %1s%4d%1s                %8.3f%8.3f%8.3f%6.2f%6.2f                                         %2s%2s\n" \
 77:         return        "{:6s}{:5d} {:^4s}{:1s}{:3s} {:1s}{:4d}{:1s}   {:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f}          {:>2s}{:2s}\n".format(\ 76:                         % ("ATOM", idx+1, self.atomnames[idx], " ", self.resnames[idx],
 78:             "ATOM", idx+1, self.atomnames[idx], " ", self.resnames[idx],\ 77:                         self.chainname, self.resnum[idx], " ",
 79:             self.chainname, self.resnum[idx], " ",\ 78:                         self.coords[idx][0], self.coords[idx][1], self.coords[idx][2],
 80:             self.coords[idx][0], self.coords[idx][1], self.coords[idx][2],\ 79:                         0, self.Bfactors[idx], " ", " ")
 81:             0, self.Bfactors[idx], " ", " ") 80:                         
 82:              81:         
 83:      82:         def writetofile(self, filename=None, filemode="w"):
 84:     def writetofile(self, filename=None, filemode="w"): 83:                 if filename is None:
 85:         if filename is None: 84:                         pdbfile = sys.stdout
 86:             pdbfile = sys.stdout 85:                 else:
 87:         else: 86:                         pdbfile = open(filename, filemode)
 88:             pdbfile = open(filename, filemode) 87:                 
 89:          88:                 #pdbfile.write("MODEL\n")
 90:         #pdbfile.write("MODEL\n") 89:                 for idx in range(len(self.atomnames)):
 91:         for idx in range(len(self.atomnames)): 90:                         pdbfile.write(self.pdbline(idx))
 92:             pdbfile.write(self.pdbline(idx)) 91:                 
 93:          92:                 #pdbfile.write("TER\nENDMDL\n")
 94:         #pdbfile.write("TER\nENDMDL\n") 93:                 
 95:          94:                 pdbfile.close()
 96:         pdbfile.close() 95: 
 97:  96:         def _make_copy(self, filter) :
 98:     def _make_copy(self, filter) : 97:                 """
 99:         """ 98:                 Create a filtered copy of the structure.
100:         Create a filtered copy of the structure. 99: 
101: 100:                 Parameters :
102:         Parameters :101:                 ------------
103:         ------------102:                 - filter : an 1D array of boolean which will be used as mask. The lenght of
104:         - filter : an 1D array of boolean which will be used as mask. The lenght of103:                 this array have to fit the number of atoms of the current structure.
105:         this array have to fit the number of atoms of the current structure.104: 
106: 105:                 Returns :
107:         Returns :106:                 ---------
108:         ---------107:                 - copy : a filtered copy as a PDB object.
109:         - copy : a filtered copy as a PDB object.108:                 """
110:         """109:                 structure = PDB(numpy.array(self.coords[filter, 0:3]))
111:         structure = PDB(numpy.array(self.coords[filter, 0:3]))110:                 structure.atomnames = list(numpy.array(self.atomnames)[filter])
112:         structure.atomnames = list(numpy.array(self.atomnames)[filter])111:                 structure.resnames = list(numpy.array(self.resnames)[filter])
113:         structure.resnames = list(numpy.array(self.resnames)[filter])112:                 structure.resnum = list(numpy.array(self.resnum)[filter])
114:         structure.resnum = list(numpy.array(self.resnum)[filter])113:                 structure.chainname = self.chainname
115:         structure.chainname = self.chainname114:                 
116:         115:                 structure.resize = self.resize
117:         structure.resize = self.resize116:                 structure.Natoms = len(structure.atomnames)
118:         structure.Natoms = len(structure.atomnames)117:                 structure.Bfactors = self.Bfactors
119:         structure.Bfactors = self.Bfactors118:                 return structure
120:         return structure119: 
121: 120:         def filter(self, values, exclude=False, filtprop=lambda x: x.atomnames, filtcrit=lambda x, y: x in y) :
122:     def filter(self, values, exclude=False, filtprop=lambda x: x.atomnames, filtcrit=lambda x, y: x in y) :121:                 """
123:         """122:                 Return a copy filtered by atom names.
124:         Return a copy filtered by atom names.123: 
125: 124:                 Parameters :
126:         Parameters :125:                 ------------
127:         ------------126:                 - values : values used to do the filtering
128:         - values : values used to do the filtering127:                 - filtcrit : criterion for the filter
129:         - filtcrit : criterion for the filter128:                 - exclude : if False, only the atoms with the specified atom names will be
130:         - exclude : if False, only the atoms with the specified atom names will be129:                 kept. If True they will be exclude and only the other atoms will appier in
131:         kept. If True they will be exclude and only the other atoms will appier in130:                 the copy. (False by default)
132:         the copy. (False by default)131:                 - 
133:         - 132: 
134: 133:                 Returns :
135:         Returns :134:                 ---------
136:         ---------135:                 - copy : a filtered copy of the structure.
137:         - copy : a filtered copy of the structure.136:                 """
138:         """137:                 filter = [filtcrit(i, values) for i in filtprop(self)]
139:         filter = [filtcrit(i, values) for i in filtprop(self)]138:                 filter = numpy.array(filter, dtype=bool) != exclude
140:         filter = numpy.array(filter, dtype=bool) != exclude139:                 return self._make_copy(filter)
141:         return self._make_copy(filter)140: 
142: 141:         def section(self, begin, end, exclude=False) :
143:     def section(self, begin, end, exclude=False) :142:                 """
144:         """143:                 Return a copy filtered with by resid between two values.
145:         Return a copy filtered with by resid between two values.144: 
146: 145:                 Parameters :
147:         Parameters :146:                 ------------
148:         ------------147:                 - begin, end : the first and last index of the selection.
149:         - begin, end : the first and last index of the selection.148:                 - exclude : if True the selection will be exclude from the copy, else only
150:         - exclude : if True the selection will be exclude from the copy, else only149:                 the selection will be kept in the copy. (False by delfault)
151:         the selection will be kept in the copy. (False by delfault)150: 
152: 151:                 Returns :
153:         Returns :152:                 ---------
154:         ---------153:                 - copy : a filtered copy of the structure.
155:         - copy : a filtered copy of the structure.154:                 """
156:         """155:                 filter_min = (numpy.array(self.resnum) >= begin)
157:         filter_min = (numpy.array(self.resnum) >= begin)156:                 filter_max = (numpy.array(self.resnum) <= end)
158:         filter_max = (numpy.array(self.resnum) <= end)157:                 filter = (filter_min == filter_max)
159:         filter = (filter_min == filter_max)158:                 filter = (filter != exclude)
160:         filter = (filter != exclude)159:                 return self._make_copy(filter)
161:         return self._make_copy(filter)160: 
162: 161:         def filter_res(self, residues, exclude=False) :
163:     def filter_res(self, residues, exclude=False) :162:                 """
164:         """163:                 Return a copy filtered by a resid list.
165:         Return a copy filtered by a resid list.164: 
166: 165:                 Parameters :
167:         Parameters :166:                 ------------
168:         ------------167:                 - residues : a list of resid to put in the selection
169:         - residues : a list of resid to put in the selection168:                 - exclude : if True, the selection will be excluded from the copy. (False
170:         - exclude : if True, the selection will be excluded from the copy. (False169:                 by default)
171:         by default)170: 
172: 171:                 Returns :
173:         Returns :172:                 ---------
174:         ---------173:                 -copy : a filtered copy of the structure
175:         -copy : a filtered copy of the structure174:                 """
176:         """175:                 filter = numpy.zeros(len(self.resnum))
177:         filter = numpy.zeros(len(self.resnum))176:                 for index, resid in enumerate(self.resnum) :
178:         for index, resid in enumerate(self.resnum) :177:                         if resid in residues :
179:             if resid in residues :178:                                 filter[index] = 1
180:                 filter[index] = 1179:                 filter = numpy.array(filter, dtype=bool)
181:         filter = numpy.array(filter, dtype=bool)180:                 filter = (filter != exclude)
182:         filter = (filter != exclude)181:                 return self._make_copy(filter)
183:         return self._make_copy(filter)182: 
184: 183:         def check_integrity(self) :
185:     def check_integrity(self) :184:                 """
186:         """185:                 Check if there is no residue missing and if the residues are well sorted.
187:         Check if there is no residue missing and if the residues are well sorted.186: 
188: 187:                 Returns :
189:         Returns :188:                 ---------
190:         ---------189:                 - answer : False if something is wrong.
191:         - answer : False if something is wrong.190:                 """
192:         """191:                 last_id = self.resnum[0]
193:         last_id = self.resnum[0]192:                 for id in self.resnum :
194:         for id in self.resnum :193:                         if not last_id <= id <= (last_id + 1) :
195:             if not last_id <= id <= (last_id + 1) :194:                                 return False
196:                 return False195:                         last_id = id
197:             last_id = id196:                 return True
198:         return True 
199: 197: 
200: def readpdb(filename):198: def readpdb(filename):
201:     pdbfile = open(filename, "r")199:         pdbfile = open(filename, "r")
202:     models = []200:         models = []
203:     models.append(readmodel(pdbfile))201:         models.append(readmodel(pdbfile))
204:     for line in pdbfile:202:         for line in pdbfile:
205:         if line.startswith("MODEL") or line.startswith("ATOM"):203:                 if line.startswith("MODEL") or line.startswith("ATOM"):
206:             #~ models.append(readmodel(pdbfile))204:                         #~ models.append(readmodel(pdbfile))
207:             readmodel2(pdbfile, models)205:                         readmodel2(pdbfile, models)
208:     206:         
209:     return models207:         return models
210: 208: 
211: def isatom(line):209: def isatom(line):
212:         return (line.startswith("ATOM") or line.startswith("HETATM"))# and line[17:20] in PDB.standard_aa_names:210:                 return (line.startswith("ATOM") or line.startswith("HETATM"))# and line[17:20] in PDB.standard_aa_names:
213: 211: 
214: def readmodel(pdbfile):212: def readmodel(pdbfile):
215:     model = PDB()213:         model = PDB()
216:     for line in pdbfile:214:         for line in pdbfile:
217:         if isatom(line):# and line[17:20] in PDB.standard_aa_names:215:                 if isatom(line):# and line[17:20] in PDB.standard_aa_names:
218:             model.addatom(line)216:                         model.addatom(line)
219:         elif line.startswith("ENDMDL"):217:                 elif line.startswith("ENDMDL"):
220:             break218:                         break
221:     219:         
222:     model.coords = numpy.array(model.coords)220:         model.coords = numpy.array(model.coords)
223:     return model221:         return model
224: 222: 
225: def readmodel2(pdbfile, models):223: def readmodel2(pdbfile, models):
226:     model = PDB(numpy.empty_like(models[0].coords))224:         model = PDB(numpy.empty_like(models[0].coords))
227:     i = 0225:         i = 0
228:     for line in pdbfile:226:         for line in pdbfile:
229:         if isatom(line):227:                 if isatom(line):
230:             model.insertatom(line, i)228:                         model.insertatom(line, i)
231:             i += 1229:                         i += 1
232:         elif line.startswith("ENDMDL"):230:                 elif line.startswith("ENDMDL"):
233:             break231:                         break
234:     232:         
235:     models.append(model)233:         models.append(model)
236: 234: 
237: 235: 
238: if __name__ == "__main__":236: if __name__ == "__main__":
239:     models = readpdb(sys.argv[1])237:         models = readpdb(sys.argv[1])
240:     print models[0].coords.shape238:         print models[0].coords.shape
241:     print len(models[0].coords)239:         print len(models[0].coords)


legend
Lines Added 
Lines changed
 Lines Removed

hdiff - version: 2.1.0