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.Data.IUPACData import atom_weights # Allowed Elements 
  9   
 10  _ATOM_FORMAT_STRING="%s%5i %-4s%c%3s %c%4i%c   %8.3f%8.3f%8.3f%6.2f%6.2f      %4s%2s%2s\n" 
 11   
 12   
13 -class Select(object):
14 """ 15 Default selection (everything) during writing - can be used as base class 16 to implement selective output. This selects which entities will be written out. 17 """ 18
19 - def __repr__(self):
20 return "<Select all>"
21
22 - def accept_model(self, model):
23 """ 24 Overload this to reject models for output. 25 """ 26 return 1
27
28 - def accept_chain(self, chain):
29 """ 30 Overload this to reject chains for output. 31 """ 32 return 1
33
34 - def accept_residue(self, residue):
35 """ 36 Overload this to reject residues for output. 37 """ 38 return 1
39
40 - def accept_atom(self, atom):
41 """ 42 Overload this to reject atoms for output. 43 """ 44 return 1
45 46
47 -class PDBIO(object):
48 """ 49 Write a Structure object (or a subset of a Structure object) as a PDB file. 50 51 52 Example: 53 >>> p=PDBParser() 54 >>> s=p.get_structure("1fat", "1fat.pdb") 55 >>> io=PDBIO() 56 >>> io.set_structure(s) 57 >>> io.save("out.pdb") 58 """
59 - def __init__(self, use_model_flag=0):
60 """ 61 @param use_model_flag: if 1, force use of the MODEL record in output. 62 @type use_model_flag: int 63 """ 64 self.use_model_flag=use_model_flag
65 66 # private mathods 67
68 - def _get_atom_line(self, atom, hetfield, segid, atom_number, resname, 69 resseq, icode, chain_id, charge=" "):
70 """Returns an ATOM PDB string (PRIVATE).""" 71 if hetfield!=" ": 72 record_type="HETATM" 73 else: 74 record_type="ATOM " 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 name=atom.get_fullname() 83 altloc=atom.get_altloc() 84 x, y, z=atom.get_coord() 85 bfactor=atom.get_bfactor() 86 occupancy=atom.get_occupancy() 87 args=(record_type, atom_number, name, altloc, resname, chain_id, 88 resseq, icode, x, y, z, occupancy, bfactor, segid, 89 element, charge) 90 return _ATOM_FORMAT_STRING % args
91 92 # Public methods 93
94 - def set_structure(self, structure):
96
97 - def save(self, file, select=Select(), write_end=0):
98 """ 99 @param file: output file 100 @type file: string or filehandle 101 102 @param select: selects which entities will be written. 103 @type select: 104 select hould have the following methods: 105 - accept_model(model) 106 - accept_chain(chain) 107 - accept_residue(residue) 108 - accept_atom(atom) 109 These methods should return 1 if the entity 110 is to be written out, 0 otherwise. 111 112 Typically select is a subclass of L{Select}. 113 """ 114 get_atom_line=self._get_atom_line 115 if isinstance(file, basestring): 116 fp=open(file, "w") 117 close_file=1 118 else: 119 # filehandle, I hope :-) 120 fp=file 121 close_file=0 122 # multiple models? 123 if len(self.structure)>1 or self.use_model_flag: 124 model_flag=1 125 else: 126 model_flag=0 127 for model in self.structure.get_list(): 128 if not select.accept_model(model): 129 continue 130 # necessary for ENDMDL 131 # do not write ENDMDL if no residues were written 132 # for this model 133 model_residues_written=0 134 atom_number=1 135 if model_flag: 136 fp.write("MODEL %s\n" % model.serial_num) 137 for chain in model.get_list(): 138 if not select.accept_chain(chain): 139 continue 140 chain_id=chain.get_id() 141 # necessary for TER 142 # do not write TER if no residues were written 143 # for this chain 144 chain_residues_written=0 145 for residue in chain.get_unpacked_list(): 146 if not select.accept_residue(residue): 147 continue 148 hetfield, resseq, icode=residue.get_id() 149 resname=residue.get_resname() 150 segid=residue.get_segid() 151 for atom in residue.get_unpacked_list(): 152 if select.accept_atom(atom): 153 chain_residues_written=1 154 model_residues_written=1 155 s=get_atom_line(atom, hetfield, segid, atom_number, resname, 156 resseq, icode, chain_id) 157 fp.write(s) 158 atom_number=atom_number+1 159 if chain_residues_written: 160 fp.write("TER\n") 161 if model_flag and model_residues_written: 162 fp.write("ENDMDL\n") 163 if write_end: 164 fp.write('END\n') 165 if close_file: 166 fp.close()
167 168 if __name__=="__main__": 169 170 from Bio.PDB.PDBParser import PDBParser 171 172 import sys 173 174 p=PDBParser(PERMISSIVE=True) 175 176 s=p.get_structure("test", sys.argv[1]) 177 178 io=PDBIO() 179 io.set_structure(s) 180 io.save("out1.pdb") 181 182 fp=open("out2.pdb", "w") 183 s1=p.get_structure("test1", sys.argv[1]) 184 s2=p.get_structure("test2", sys.argv[2]) 185 io=PDBIO(1) 186 io.set_structure(s1) 187 io.save(fp) 188 io.set_structure(s2) 189 io.save(fp, write_end=1) 190 fp.close() 191