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 - def __init__(self, use_model_flag=0):
56 """Creat the PDBIO object. 57 58 @param use_model_flag: if 1, force use of the MODEL record in output. 59 @type use_model_flag: int 60 """ 61 self.use_model_flag = use_model_flag
62 63 # private mathods 64
65 - def _get_atom_line(self, atom, hetfield, segid, atom_number, resname, 66 resseq, icode, chain_id, charge=" "):
67 """Returns an ATOM PDB string (PRIVATE).""" 68 if hetfield != " ": 69 record_type = "HETATM" 70 else: 71 record_type = "ATOM " 72 if atom.element: 73 element = atom.element.strip().upper() 74 if element.capitalize() not in atom_weights: 75 raise ValueError("Unrecognised element %r" % atom.element) 76 element = element.rjust(2) 77 else: 78 element = " " 79 name = atom.get_fullname() 80 altloc = atom.get_altloc() 81 x, y, z = atom.get_coord() 82 bfactor = atom.get_bfactor() 83 occupancy = atom.get_occupancy() 84 try: 85 occupancy_str = "%6.2f" % occupancy 86 except TypeError: 87 if occupancy is None: 88 occupancy_str = " " * 6 89 import warnings 90 from Bio import BiopythonWarning 91 warnings.warn("Missing occupancy in atom %s written as blank" % 92 repr(atom.get_full_id()), BiopythonWarning) 93 else: 94 raise TypeError("Invalid occupancy %r in atom %r" 95 % (occupancy, atom.get_full_id())) 96 97 args = (record_type, atom_number, name, altloc, resname, chain_id, 98 resseq, icode, x, y, z, occupancy_str, bfactor, segid, 99 element, charge) 100 return _ATOM_FORMAT_STRING % args
101 102 # Public methods 103
104 - def set_structure(self, pdb_object):
105 # Check what the user is providing and build a structure appropriately 106 if pdb_object.level == "S": 107 structure = pdb_object 108 else: 109 sb = StructureBuilder() 110 sb.init_structure('pdb') 111 sb.init_seg(' ') 112 # Build parts as necessary 113 if pdb_object.level == "M": 114 sb.structure.add(pdb_object) 115 self.structure = sb.structure 116 else: 117 sb.init_model(0) 118 if pdb_object.level == "C": 119 sb.structure[0].add(pdb_object) 120 else: 121 sb.init_chain('A') 122 if pdb_object.level == "R": 123 try: 124 parent_id = pdb_object.parent.id 125 sb.structure[0]['A'].id = parent_id 126 except Exception: 127 pass 128 sb.structure[0]['A'].add(pdb_object) 129 else: 130 # Atom 131 sb.init_residue('DUM', ' ', 1, ' ') 132 try: 133 parent_id = pdb_object.parent.parent.id 134 sb.structure[0]['A'].id = parent_id 135 except Exception: 136 pass 137 sb.structure[0]['A'].child_list[0].add(pdb_object) 138 139 # Return structure 140 structure = sb.structure 141 self.structure = structure
142
143 - def save(self, file, select=Select(), write_end=True):
144 """ 145 @param file: output file 146 @type file: string or filehandle 147 148 @param select: selects which entities will be written. 149 @type select: object 150 151 Typically select is a subclass of L{Select}, it should 152 have the following methods: 153 154 - accept_model(model) 155 - accept_chain(chain) 156 - accept_residue(residue) 157 - accept_atom(atom) 158 159 These methods should return 1 if the entity is to be 160 written out, 0 otherwise. 161 162 Typically select is a subclass of L{Select}. 163 """ 164 get_atom_line = self._get_atom_line 165 if isinstance(file, basestring): 166 fp = open(file, "w") 167 close_file = 1 168 else: 169 # filehandle, I hope :-) 170 fp = file 171 close_file = 0 172 # multiple models? 173 if len(self.structure) > 1 or self.use_model_flag: 174 model_flag = 1 175 else: 176 model_flag = 0 177 for model in self.structure.get_list(): 178 if not select.accept_model(model): 179 continue 180 # necessary for ENDMDL 181 # do not write ENDMDL if no residues were written 182 # for this model 183 model_residues_written = 0 184 atom_number = 1 185 if model_flag: 186 fp.write("MODEL %s\n" % model.serial_num) 187 for chain in model.get_list(): 188 if not select.accept_chain(chain): 189 continue 190 chain_id = chain.get_id() 191 # necessary for TER 192 # do not write TER if no residues were written 193 # for this chain 194 chain_residues_written = 0 195 for residue in chain.get_unpacked_list(): 196 if not select.accept_residue(residue): 197 continue 198 hetfield, resseq, icode = residue.get_id() 199 resname = residue.get_resname() 200 segid = residue.get_segid() 201 for atom in residue.get_unpacked_list(): 202 if select.accept_atom(atom): 203 chain_residues_written = 1 204 model_residues_written = 1 205 s = get_atom_line(atom, hetfield, segid, atom_number, resname, 206 resseq, icode, chain_id) 207 fp.write(s) 208 atom_number = atom_number + 1 209 if chain_residues_written: 210 fp.write("TER\n") 211 if model_flag and model_residues_written: 212 fp.write("ENDMDL\n") 213 if write_end: 214 fp.write('END\n') 215 if close_file: 216 fp.close()
217 218 if __name__ == "__main__": 219 220 from Bio.PDB.PDBParser import PDBParser 221 222 import sys 223 224 p = PDBParser(PERMISSIVE=True) 225 226 s = p.get_structure("test", sys.argv[1]) 227 228 io = PDBIO() 229 io.set_structure(s) 230 io.save("out1.pdb") 231 232 with open("out2.pdb", "w") as fp: 233 s1 = p.get_structure("test1", sys.argv[1]) 234 s2 = p.get_structure("test2", sys.argv[2]) 235 io = PDBIO(1) 236 io.set_structure(s1) 237 io.save(fp) 238 io.set_structure(s2) 239 io.save(fp, write_end=1) 240