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