MMTK.ConfigIO — MMTK User Guide 2.7.7 documentation

### Navigation

- index
- modules |
- MMTK User Guide 2.7.7 documentation »
- Module code »

# Source code for MMTK.ConfigIO

    # This module deals with input and output of configurations.
    #
    # Written by Konrad Hinsen
    #

    """
    I/O of molecular configurations

    Input: Z-Matrix and Cartesian
    Output: VRML
    PDB files are handled in :class:`~MMTK.PDB`.
    """

    __docformat__ = 'restructuredtext'

    from MMTK import PDB, Units, Utility
    from Scientific.Geometry.Objects3D import Sphere, Cone, Plane, Line, \
                                              rotatePoint
    from Scientific.Geometry import Vector
    from Scientific.Visualization import VRML
    from Scientific import N
    import os

    #
    # This class represents a Z-Matrix. Z-Matrix data consists of a list
    # with one element for each atom being defined. Each entry is a
    # list containing the data defining the atom.
    #
    class ZMatrix(object):

        """
        Z-Matrix specification of a molecule conformation

        ZMatrix objects can be used in chemical database entries
        to specify molecule conformations by internal coordinates.
        With the exception of the three first atoms, each atom is
        defined relative to three previously atoms by a distance,
        an angle, and a dihedral angle.
        """

        def __init__(self, data):
            """
            :param data: a list of atom definitions. Each atom definition (except
                         for the first three ones) is a list containing seven
                         elements:

                           - the atom to be defined
                           - a previously defined atom and the distance to it
                           - another previously defined atom and the angle to it
                           - a third previously defined atom and the dihedral
                             angle to it

                         The definition of the first atom contains only the first
                         element, the second atom needs the first three elements,
                         and the third atom is defined by the first five elements.
            """
            self.data = data
            self.coordinates = {}

        _substitute = True

        def findPositions(self):
            # First atom at origin
            self.coordinates[self.data[0][0]] = Vector(0,0,0)
            # Second atom along x-axis
            self.coordinates[self.data[1][0]] = Vector(self.data[1][2],0,0)
            # Third atom in xy-plane
            try:
                pos1 = self.coordinates[self.data[2][1]]
            except KeyError:
                raise ValueError("atom %d has no defined position"
                                  % self.data[2][1].number)
            try:
                pos2 = self.coordinates[self.data[2][3]]
            except KeyError:
                raise ValueError("atom %d has no defined position"
                                  % self.data[2][3].number)
            sphere = Sphere(pos1, self.data[2][2])
            cone = Cone(pos1, pos2-pos1, self.data[2][4])
            plane = Plane(Vector(0,0,0), Vector(0,0,1))
            points = sphere.intersectWith(cone).intersectWith(plane)
            self.coordinates[self.data[2][0]] = points[0]
            # All following atoms defined by distance + angle + dihedral
            for entry in self.data[3:]:
                try:
                    pos1 = self.coordinates[entry[1]]
                except KeyError:
                    raise ValueError("atom %d has no defined position"
                                      % entry[1].number)
                try:
                    pos2 = self.coordinates[entry[3]]
                except KeyError:
                    raise ValueError("atom %d has no defined position"
                                      % entry[3].number)
                try:
                    pos3 = self.coordinates[entry[5]]
                except KeyError:
                    raise ValueError("atom %d has no defined position"
                                      % entry[5].number)
                distance = entry[2]
                angle = entry[4]
                dihedral = entry[6]
                sphere = Sphere(pos1, distance)
                cone = Cone(pos1, pos2-pos1, angle)
                plane123 = Plane(pos3, pos2, pos1)
                points = sphere.intersectWith(cone).intersectWith(plane123)
                for p in points:
                    if Plane(pos2, pos1, p).normal * plane123.normal > 0:
                        break
                p = rotatePoint(p, Line(pos1, pos2-pos1), dihedral)
                self.coordinates[entry[0]] = p

        def applyTo(self, object):
            """
            Define the positions of the atoms in a chemical object by the
            internal coordinates of the Z-Matrix.

            :param object: the object to which the Z-Matrix is applied
            """
            if not len(self.coordinates):
                self.findPositions()
            for entry in self.coordinates.items():
                object.setPosition(entry[0], entry[1])
            object.normalizePosition()

    #
    # This class represents a dictionary of Cartesian positions
    #
    class Cartesian(object):

        """
        Cartesian specification of a molecule conformation

        Cartesian objects can be used in chemical database entries
        to specify molecule conformations by Cartesian coordinates.
        """

        def __init__(self, data):
            """
            :param data: a dictionary mapping atoms to tuples of length three
                         that define its Cartesian coordinates
            """
            self.dict = data

        _substitute = True

        def applyTo(self, object):
            """
            Define the positions of the atoms in a chemical object by the
            stored coordinates.

            :param object: the object to which the coordinates are applied
            """
            for a, r in self.dict.items():
                object.setPosition(a, Vector(r[0], r[1], r[2]))

    #
    # VRML output
    #
    class VRMLWireframeFile(VRML.VRMLFile):

        def __init__(self, filename, color_values = None):
            VRML.VRMLFile.__init__(self, filename, 'w')
            self.warning = 0
            self.color_values = color_values
            if self.color_values is not None:
                lower = N.minimum.reduce(color_values.array)
                upper = N.maximum.reduce(color_values.array)
                self.color_scale = VRML.ColorScale((lower, upper))

        def write(self, object, configuration = None, distance = None):
            from MMTK.ChemicalObjects import isChemicalObject
            from MMTK.Universe import InfiniteUniverse
            if distance is None:
                try:
                    distance = object.universe().distanceVector
                except AttributeError:
                    distance = InfiniteUniverse().distanceVector
            if not isChemicalObject(object):
                for o in object:
                    self.write(o, configuration, distance)
            else:
                for bu in object.bondedUnits():
                    for a in bu.atomList():
                        self.writeAtom(a, configuration)
                    if hasattr(bu, 'bonds'):
                        for b in bu.bonds:
                            self.writeBond(b, configuration, distance)

        def close(self):
            VRML.VRMLFile.close(self)
            if self.warning:
                Utility.warning('Some atoms are missing in the output file ' + \
                                'because their positions are undefined.')
                self.warning = 0

        def atomColor(self, atom):
            if self.color_values is None:
                return atom.color
            else:
                return self.color_scale(self.color_values[atom])

        def writeAtom(self, atom, configuration):
            pass

        def writeBond(self, bond, configuration, distance):
            p1 = bond.a1.position(configuration)
            p2 = bond.a2.position(configuration)
            if p1 is not None and p2 is not None:
                bond_vector = 0.5*distance(bond.a1, bond.a2, configuration)
                cut = bond_vector != 0.5*(p2-p1)
                color1 = self.atomColor(bond.a1)
                color2 = self.atomColor(bond.a2)
                material1 = VRML.EmissiveMaterial(color1)
                material2 = VRML.EmissiveMaterial(color2)
                if color1 == color2 and not cut:
                    c = VRML.Line(p1, p2, material = material1)
                    c.writeToFile(self)
                else:
                    c = VRML.Line(p1, p1+bond_vector, material = material1)
                    c.writeToFile(self)
                    c = VRML.Line(p2, p2-bond_vector, material = material2)
                    c.writeToFile(self)

    class VRMLHighlight(VRMLWireframeFile):

        def writeAtom(self, atom, configuration):
            try:
                highlight = atom.highlight
            except AttributeError:
                highlight = 0
            if highlight:
                p = atom.position(configuration)
                if p is None:
                    self.warning = 1
                else:
                    s = VRML.Sphere(p, 0.1*Units.Ang,
                                    material = VRML.DiffuseMaterial(atom.color),
                                    reuse = 1)
                    s.writeToFile(self)

    class VRMLBallAndStickFile(VRMLWireframeFile):

        def writeAtom(self, atom, configuration):
            p = atom.position(configuration)
            if p is None:
                self.warning = 1
            else:
                color = self.atomColor(atom)
                s = VRML.Sphere(p, 0.1*Units.Ang,
                                material = VRML.DiffuseMaterial(color),
                                reuse = 1)
                s.writeToFile(self)

        def writeBond(self, bond, configuration, distance):
            p1 = bond.a1.position(configuration)
            p2 = bond.a2.position(configuration)
            if p1 is not None and p2 is not None:
                bond_vector = 0.5*distance(bond.a1, bond.a2, configuration)
                cut = bond_vector != 0.5*(p2-p1)
                color1 = self.atomColor(bond.a1)
                color2 = self.atomColor(bond.a2)
                material1 = VRML.EmissiveMaterial(color1)
                material2 = VRML.EmissiveMaterial(color2)
                if color1 == color2 and not cut:
                    c = VRML.Cylinder(p1, p2, 0.03*Units.Ang,
                                      material = material1)
                    c.writeToFile(self)
                else:
                    c = VRML.Cylinder(p1, p1+bond_vector, 0.03*Units.Ang,
                                      material = material1)
                    c.writeToFile(self)
                    c = VRML.Cylinder(p2, p2-bond_vector, 0.03*Units.Ang,
                                      material = material2)
                    c.writeToFile(self)

    class VRMLChargeFile(VRMLWireframeFile):

        color_scale = VRML.SymmetricColorScale(1.)

        def writeAtom(self, atom, configuration):
            p = atom.position(configuration)
            c = atom.charge()
            c = max(min(c, 1.), -1.)
            if p is None:
                self.warning = 1
            else:
                s = VRML.Sphere(p, 0.1*Units.Ang,
                                material = VRML.Material(diffuse_color =
                                                         self.color_scale(c)))
                s.writeToFile(self)

        bond_material = VRML.DiffuseMaterial('black')

        def writeBond(self, bond, configuration, distance):
            p1 = bond.a1.position(configuration)
            p2 = bond.a2.position(configuration)
            if p1 is not None and p2 is not None:
                bond_vector = 0.5*distance(bond.a1, bond.a2, configuration)
                cut = bond_vector != 0.5*(p2-p1)
                if not cut:
                    c = VRML.Line(p1, p2, material = self.bond_material)
                    c.writeToFile(self)
                else:
                    c = VRML.Line(p1, p1+bond_vector, material = self.bond_material)
                    c.writeToFile(self)
                    c = VRML.Line(p2, p2-bond_vector, material = self.bond_material)
                    c.writeToFile(self)

    VRMLFile = VRMLWireframeFile

    #
    # Recognize some standard file types by their extensions
    #
    def fileFormatFromExtension(filename):
        filename, ext = os.path.splitext(filename)
        if ext in _file_compressions:
            filename, ext = os.path.splitext(filename)
        try:
            return _file_formats[ext]
        except KeyError:
            raise IOError('Unknown file format')

    _file_formats = {'.pdb': 'pdb', '.wrl': 'vrml'}
    _file_compressions = ['.gz', '.Z']

    #
    # Output file for a specified format
    #
    def OutputFile(filename, format = None):
        if format is None:
            format = fileFormatFromExtension(filename)
        format = tuple(format.split('.'))
        try:
            return _file_types[format](filename)
        except KeyError:
            if len(format) == 1:
                return _file_types[format[0]](filename)
            else:
                _file_types[format[0]](filename, format[1])

    _file_types = {'pdb': PDB.PDBOutputFile,
                   ('vrml',): VRMLFile,
                   ('vrml', 'wireframe'): VRMLWireframeFile,
                   ('vrml', 'highlight'): VRMLHighlight,
                   ('vrml', 'ball_and_stick'): VRMLBallAndStickFile,
                   ('vrml', 'charge'): VRMLChargeFile}

### Navigation

- index
- modules |
- MMTK User Guide 2.7.7 documentation »
- Module code »