Back to index

python-biopython  1.60
PDBIO.py
Go to the documentation of this file.
00001 # Copyright (C) 2002, Thomas Hamelryck (thamelry@binf.ku.dk)
00002 # This code is part of the Biopython distribution and governed by its
00003 # license.  Please see the LICENSE file that should have been included
00004 # as part of this package.
00005 
00006 """Output of PDB files."""
00007 
00008 from Bio.Data.IUPACData import atom_weights # Allowed Elements
00009 
00010 _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"
00011 
00012 
00013 class Select(object):
00014     """
00015     Default selection (everything) during writing - can be used as base class
00016     to implement selective output. This selects which entities will be written out.
00017     """
00018 
00019     def __repr__(self):
00020         return "<Select all>"
00021 
00022     def accept_model(self, model):
00023         """
00024         Overload this to reject models for output.
00025         """
00026         return 1
00027 
00028     def accept_chain(self, chain):
00029         """
00030         Overload this to reject chains for output.
00031         """
00032         return 1
00033 
00034     def accept_residue(self, residue):
00035         """
00036         Overload this to reject residues for output.
00037         """
00038         return 1
00039 
00040     def accept_atom(self, atom):
00041         """
00042         Overload this to reject atoms for output.
00043         """
00044         return 1
00045 
00046 
00047 class PDBIO(object):
00048     """
00049     Write a Structure object (or a subset of a Structure object) as a PDB file.
00050 
00051 
00052     Example:
00053         >>> p=PDBParser()
00054         >>> s=p.get_structure("1fat", "1fat.pdb")
00055         >>> io=PDBIO()
00056         >>> io.set_structure(s)
00057         >>> io.save("out.pdb")
00058     """
00059     def __init__(self, use_model_flag=0):
00060         """
00061         @param use_model_flag: if 1, force use of the MODEL record in output.
00062         @type use_model_flag: int
00063         """
00064         self.use_model_flag=use_model_flag
00065     
00066     # private mathods
00067 
00068     def _get_atom_line(self, atom, hetfield, segid, atom_number, resname, 
00069         resseq, icode, chain_id, charge="  "):
00070         """Returns an ATOM PDB string (PRIVATE)."""
00071         if hetfield!=" ":
00072             record_type="HETATM"
00073         else:
00074             record_type="ATOM  "
00075         if atom.element:
00076             element = atom.element.strip().upper()
00077             if element.capitalize() not in atom_weights:
00078                 raise ValueError("Unrecognised element %r" % atom.element)
00079             element = element.rjust(2)
00080         else:
00081             element = "  "
00082         name=atom.get_fullname()
00083         altloc=atom.get_altloc()
00084         x, y, z=atom.get_coord()
00085         bfactor=atom.get_bfactor()
00086         occupancy=atom.get_occupancy()
00087         args=(record_type, atom_number, name, altloc, resname, chain_id,
00088             resseq, icode, x, y, z, occupancy, bfactor, segid,
00089             element, charge)
00090         return _ATOM_FORMAT_STRING % args
00091 
00092     # Public methods
00093 
00094     def set_structure(self, structure):
00095         self.structure=structure
00096 
00097     def save(self, file, select=Select(), write_end=0):
00098         """
00099         @param file: output file
00100         @type file: string or filehandle 
00101 
00102         @param select: selects which entities will be written.
00103         @type select: 
00104             select hould have the following methods:
00105                 - accept_model(model)
00106                 - accept_chain(chain)
00107                 - accept_residue(residue)
00108                 - accept_atom(atom)
00109             These methods should return 1 if the entity
00110             is to be written out, 0 otherwise.
00111 
00112             Typically select is a subclass of L{Select}.
00113         """
00114         get_atom_line=self._get_atom_line
00115         if isinstance(file, basestring):
00116             fp=open(file, "w")
00117             close_file=1
00118         else:
00119             # filehandle, I hope :-)
00120             fp=file
00121             close_file=0
00122         # multiple models?
00123         if len(self.structure)>1 or self.use_model_flag:
00124             model_flag=1
00125         else:
00126             model_flag=0
00127         for model in self.structure.get_list():
00128             if not select.accept_model(model):
00129                 continue
00130             # necessary for ENDMDL 
00131             # do not write ENDMDL if no residues were written
00132             # for this model
00133             model_residues_written=0
00134             atom_number=1
00135             if model_flag:
00136                 fp.write("MODEL      %s\n" % model.serial_num)
00137             for chain in model.get_list():
00138                 if not select.accept_chain(chain):
00139                     continue
00140                 chain_id=chain.get_id()
00141                 # necessary for TER 
00142                 # do not write TER if no residues were written
00143                 # for this chain
00144                 chain_residues_written=0
00145                 for residue in chain.get_unpacked_list():
00146                     if not select.accept_residue(residue):
00147                         continue
00148                     hetfield, resseq, icode=residue.get_id()
00149                     resname=residue.get_resname()  
00150                     segid=residue.get_segid()
00151                     for atom in residue.get_unpacked_list():
00152                         if select.accept_atom(atom):
00153                             chain_residues_written=1
00154                             model_residues_written=1
00155                             s=get_atom_line(atom, hetfield, segid, atom_number, resname,
00156                                 resseq, icode, chain_id)
00157                             fp.write(s)
00158                             atom_number=atom_number+1
00159                 if chain_residues_written:
00160                     fp.write("TER\n")
00161             if model_flag and model_residues_written:
00162                 fp.write("ENDMDL\n")
00163             if write_end:
00164                 fp.write('END\n')
00165         if close_file:
00166             fp.close()
00167 
00168 if __name__=="__main__":
00169     
00170     from Bio.PDB.PDBParser import PDBParser
00171 
00172     import sys
00173 
00174     p=PDBParser(PERMISSIVE=True)
00175 
00176     s=p.get_structure("test", sys.argv[1])
00177 
00178     io=PDBIO()
00179     io.set_structure(s)
00180     io.save("out1.pdb")
00181 
00182     fp=open("out2.pdb", "w")
00183     s1=p.get_structure("test1", sys.argv[1])
00184     s2=p.get_structure("test2", sys.argv[2])
00185     io=PDBIO(1)
00186     io.set_structure(s1)
00187     io.save(fp)
00188     io.set_structure(s2)
00189     io.save(fp, write_end=1)
00190     fp.close()
00191 
00192 
00193 
00194