Package Bio :: Package PDB :: Module PDBIO'
[hide private]
[frames] | no frames]

Source Code for Module Bio.PDB.PDBIO'

  1  # Copyright (C) 2002, Thomas Hamelryck (thamelry@binf.ku.dk) 
  2  # This code is part of the Biopython distribution and governed by its 
  3  # license.  Please see the LICENSE file that should have been included 
  4  # as part of this package. 
  5   
  6  """Output of PDB files.""" 
  7   
  8  from Bio._py3k import basestring 
  9   
 10  from Bio.PDB.StructureBuilder import StructureBuilder  # To allow saving of chains, residues, etc.. 
 11  from Bio.Data.IUPACData import atom_weights  # Allowed Elements 
 12   
 13   
 14  _ATOM_FORMAT_STRING = "%s%5i %-4s%c%3s %c%4i%c   %8.3f%8.3f%8.3f%s%6.2f      %4s%2s%2s\n" 
 15   
 16   
17 -class Select(object):
18 """Select everything fo PDB output (for use as a bas class). 19 20 Default selection (everything) during writing - can be used as base class 21 to implement selective output. This selects which entities will be written out. 22 """ 23
24 - def __repr__(self):
25 return "<Select all>"
26
27 - def accept_model(self, model):
28 """Overload this to reject models for output.""" 29 return 1
30
31 - def accept_chain(self, chain):
32 """Overload this to reject chains for output.""" 33 return 1
34
35 - def accept_residue(self, residue):
36 """Overload this to reject residues for output.""" 37 return 1
38
39 - def accept_atom(self, atom):
40 """Overload this to reject atoms for output.""" 41 return 1
42 43
44 -class PDBIO(object):
45 """Write a Structure object (or a subset of a Structure object) as a PDB file. 46 47 Example: 48 49 >>> p=PDBParser() 50 >>> s=p.get_structure("1fat", "1fat.pdb") 51 >>> io=PDBIO() 52 >>> io.set_structure(s) 53 >>> io.save("out.pdb") 54 55 """ 56
57 - def __init__(self, use_model_flag=0):
58 """Creat the PDBIO object. 59 60 @param use_model_flag: if 1, force use of the MODEL record in output. 61 @type use_model_flag: int 62 """ 63 self.use_model_flag = use_model_flag
64 65 # private mathods 66
67 - def _get_atom_line(self, atom, hetfield, segid, atom_number, resname, 68 resseq, icode, chain_id, charge=" "):
69 """Returns an ATOM PDB string (PRIVATE).""" 70 if hetfield != " ": 71 record_type = "HETATM" 72 else: 73 record_type = "ATOM " 74 75 if atom.element: 76 element = atom.element.strip().upper() 77 if element.capitalize() not in atom_weights: 78 raise ValueError("Unrecognised element %r" % atom.element) 79 element = element.rjust(2) 80 else: 81 element = " " 82 83 name = atom.get_fullname().strip() 84 # Pad atom name if: 85 # - smaller than 4 characters 86 # AND - is not C, N, O, S, H, F, P, ..., one letter elements 87 # AND - first character is NOT numeric (funky hydrogen naming rules) 88 if len(name) < 4 and name[:1].isalpha() and len(element.strip()) < 2: 89 name = " " + name 90 91 altloc = atom.get_altloc() 92 x, y, z = atom.get_coord() 93 bfactor = atom.get_bfactor() 94 occupancy = atom.get_occupancy() 95 try: 96 occupancy_str = "%6.2f" % occupancy 97 except TypeError: 98 if occupancy is None: 99 occupancy_str = " " * 6 100 import warnings 101 from Bio import BiopythonWarning 102 warnings.warn("Missing occupancy in atom %s written as blank" % 103 repr(atom.get_full_id()), BiopythonWarning) 104 else: 105 raise TypeError("Invalid occupancy %r in atom %r" 106 % (occupancy, atom.get_full_id())) 107 108 args = (record_type, atom_number, name, altloc, resname, chain_id, 109 resseq, icode, x, y, z, occupancy_str, bfactor, segid, 110 element, charge) 111 return _ATOM_FORMAT_STRING % args
112 113 # Public methods 114
115 - def set_structure(self, pdb_object):
116 # Check what the user is providing and build a structure appropriately 117 if pdb_object.level == "S": 118 structure = pdb_object 119 else: 120 sb = StructureBuilder() 121 sb.init_structure('pdb') 122 sb.init_seg(' ') 123 # Build parts as necessary 124 if pdb_object.level == "M": 125 sb.structure.add(pdb_object) 126 self.structure = sb.structure 127 else: 128 sb.init_model(0) 129 if pdb_object.level == "C": 130 sb.structure[0].add(pdb_object) 131 else: 132 sb.init_chain('A') 133 if pdb_object.level == "R": 134 try: 135 parent_id = pdb_object.parent.id 136 sb.structure[0]['A'].id = parent_id 137 except Exception: 138 pass 139 sb.structure[0]['A'].add(pdb_object) 140 else: 141 # Atom 142 sb.init_residue('DUM', ' ', 1, ' ') 143 try: 144 parent_id = pdb_object.parent.parent.id 145 sb.structure[0]['A'].id = parent_id 146 except Exception: 147 pass 148 sb.structure[0]['A'].child_list[0].add(pdb_object) 149 150 # Return structure 151 structure = sb.structure 152 self.structure = structure
153
154 - def save(self, file, select=Select(), write_end=True, preserve_atom_numbering=False):
155 """@param file: output file 156 @type file: string or filehandle 157 158 @param select: selects which entities will be written. 159 @type select: object 160 161 Typically select is a subclass of L{Select}, it should 162 have the following methods: 163 164 - accept_model(model) 165 - accept_chain(chain) 166 - accept_residue(residue) 167 - accept_atom(atom) 168 169 These methods should return 1 if the entity is to be 170 written out, 0 otherwise. 171 172 Typically select is a subclass of L{Select}. 173 """ 174 get_atom_line = self._get_atom_line 175 if isinstance(file, basestring): 176 fp = open(file, "w") 177 close_file = 1 178 else: 179 # filehandle, I hope :-) 180 fp = file 181 close_file = 0 182 # multiple models? 183 if len(self.structure) > 1 or self.use_model_flag: 184 model_flag = 1 185 else: 186 model_flag = 0 187 for model in self.structure.get_list(): 188 if not select.accept_model(model): 189 continue 190 # necessary for ENDMDL 191 # do not write ENDMDL if no residues were written 192 # for this model 193 model_residues_written = 0 194 if not preserve_atom_numbering: 195 atom_number = 1 196 if model_flag: 197 fp.write("MODEL %s\n" % model.serial_num) 198 for chain in model.get_list(): 199 if not select.accept_chain(chain): 200 continue 201 chain_id = chain.get_id() 202 # necessary for TER 203 # do not write TER if no residues were written 204 # for this chain 205 chain_residues_written = 0 206 for residue in chain.get_unpacked_list(): 207 if not select.accept_residue(residue): 208 continue 209 hetfield, resseq, icode = residue.get_id() 210 resname = residue.get_resname() 211 segid = residue.get_segid() 212 for atom in residue.get_unpacked_list(): 213 if select.accept_atom(atom): 214 chain_residues_written = 1 215 model_residues_written = 1 216 if preserve_atom_numbering: 217 atom_number = atom.get_serial_number() 218 s = get_atom_line(atom, hetfield, segid, atom_number, resname, 219 resseq, icode, chain_id) 220 fp.write(s) 221 if not preserve_atom_numbering: 222 atom_number += 1 223 if chain_residues_written: 224 fp.write("TER %5i %3s %c%4i%c \n" 225 % (atom_number, resname, chain_id, resseq, icode)) 226 227 if model_flag and model_residues_written: 228 fp.write("ENDMDL\n") 229 if write_end: 230 fp.write('END\n') 231 if close_file: 232 fp.close()
233 234 235 if __name__ == "__main__": 236 237 from Bio.PDB.PDBParser import PDBParser 238 239 import sys 240 241 p = PDBParser(PERMISSIVE=True) 242 243 s = p.get_structure("test", sys.argv[1]) 244 245 io = PDBIO() 246 io.set_structure(s) 247 io.save("out1.pdb") 248 249 with open("out2.pdb", "w") as fp: 250 s1 = p.get_structure("test1", sys.argv[1]) 251 s2 = p.get_structure("test2", sys.argv[2]) 252 io = PDBIO(1) 253 io.set_structure(s1) 254 io.save(fp) 255 io.set_structure(s2) 256 io.save(fp, write_end=1) 257