Chimie Théorique » QMX

import numpy as np 

from ase.atoms import Atom, Atoms

from ase.parallel import paropen

"""Module to read and write atoms in PDB file format"""

def read_pdb(fileobj, index=-1):

    """Read PDB files.

    The format is assumed to follow the description given in

    http://www.wwpdb.org/documentation/format32/sect9.html."""

    if isinstance(fileobj, str):

        fileobj = open(fileobj)

    atoms = Atoms()

    for line in fileobj.readlines():

        if line.startswith('ATOM') or line.startswith('HETATM'):

            try:

                symbol = line[12:16].strip()

                # we assume that the second character is a label 

                # in case that it is upper case

                if len(symbol) > 1 and symbol[1].isupper():

                    symbol = symbol[0]

                words = line[30:55].split()

                position = np.array([float(words[0]), 

                                     float(words[1]),

                                     float(words[2])])

                atoms.append(Atom(symbol, position))

            except:

                pass

    return atoms

def write_pdb(fileobj, images):

    """Write images to PDB-file."""

    if isinstance(fileobj, str):

        fileobj = paropen(fileobj, 'w')

    if not isinstance(images, (list, tuple)):

        images = [images]

    format = 'ATOM  %5d %-4s              %8.3f%8.3f%8.3f  0.00  0.00\n'

    # RasMol complains if the atom index exceeds 100000. There might

    # be a limit of 5 digit numbers in this field.

    MAXNUM = 100000

    symbols = images[0].get_chemical_symbols()

    natoms = len(symbols)

    for atoms in images:

        fileobj.write('MODEL         1\n')

        p = atoms.get_positions()

        for a in range(natoms):

            x, y, z = p[a]

            fileobj.write(format % (a % MAXNUM, symbols[a], x, y, z))

        fileobj.write('ENDMDL\n')