atomsengine.py

import ase.io
import logging
logger = logging.getLogger(__name__)


#===============================================================================
class AtomsEngine(object):
    """Used to parse various different atomic coordinate files.

    See the dictionary 'formats' for all the supported formats and a brief
    explanation.Reading is primarily done by ASE or MDAnalysis, but in some cases own
    implementation is used. Returns all coordinates as numpy arrays.
    """
    formats = {
        "xyz":    "(.xyz): The XYZ file format.",
        "cif":    "(.cif): Crystallographic Information File",
        "pdb":    "(.pdb): Protein Data Bank",
        #"dcd":    "(.dcd): Binary trajectory file format used by CHARMM, NAMD, and X-PLOR.",
    }

    def determine_tool(self, format):
        """Determines which tool to use for extracting trajectories in the
        given format.
        """
        formats = {
            "xyz": "ASE",
            "cif": "ASE",
            "pdb": "ASE",
        }
        result = formats.get(format)
        if result:
            return result
        else:
            logger.warning("The format '{}' is not supported by AtomsEngine.".format(format))

    def check_format_support(self, format):
        """Check if the given format is supported.
        """

        if format not in AtomsEngine.formats:
            logger.error("The format '{}' is not supported by AtomsEngine.".format(format))
            return False
        else:
            return True

    def n_atoms(self, file_handle, format):
        """Read the first configuration of the coordinate file to extract the
        number of atoms in it.
        """
        iterator = self.iread(file_handle, format)
        pos = iterator.next()
        return pos.shape[0]

    def iread(self, file_handle, format, index=0):
        """Returns an iterator that goes through the given trajectory file one
        configuration at a time. Good for e.g. streaming the contents to disc as the
        whole file doesn't have to be loaded into memory.
        """

        if not self.check_format_support(format):
            return

        if file_handle is None:
            print "NONE"

        tool = self.determine_tool(format)

        if tool == "ASE":
            return self.ase_iread(file_handle, format, index)
        elif tool == "custom":
            return self.custom_iread(file_handle, format, index)
        elif tool == "MDAnalysis":
            return self.mdanalysis_iread(file_handle, format, index)

    def ase_iread(self, file_handle, format, index):
        """
        """
        # After reading the ASE source code, it seems that the ASE iread does
        # actually read the entire file into memory and the yields the
        # configurations from it. Should be checked at some point.
        def ase_generator(iterator):
            """Used to wrap an iterator returned by ase.io.iread so that it returns
            the positions instead of the ase.Atoms object.
            """
            for value in iterator:
                yield value.get_positions()

        iterator = ase.io.iread(file_handle, format=format)
        return ase_generator(iterator)

    def custom_iread(self, file_handle, format, index):
        """
        """
        pass