atomsengine.py 3.22 KB
Newer Older
1
import ase.io
2
3
import logging
logger = logging.getLogger(__name__)
4
5
6
7
8
9


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

10
    See the dictionary 'formats' for all the supported formats and a brief
11
12
    explanation.Reading is primarily done by ASE or MDAnalysis, but in some cases own
    implementation is used. Returns all coordinates as numpy arrays.
13
    """
14
    formats = {
15
16
17
18
        "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.",
19
    }
20

21
    def determine_tool(self, format):
22
23
24
        """Determines which tool to use for extracting trajectories in the
        given format.
        """
25
        formats = {
26
27
28
            "xyz": "ASE",
            "cif": "ASE",
            "pdb": "ASE",
29
30
31
32
33
34
        }
        result = formats.get(format)
        if result:
            return result
        else:
            logger.warning("The format '{}' is not supported by AtomsEngine.".format(format))
35

36
37
38
39
40
41
42
43
44
45
    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

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

54
    def iread(self, file_handle, format, index=0):
55
56
57
58
59
        """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.
        """

60
61
62
        if not self.check_format_support(format):
            return

63
64
65
        if file_handle is None:
            print "NONE"

66
67
        tool = self.determine_tool(format)

68
69
70
71
72
73
74
75
76
77
        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):
        """
        """
78
79
80
        # 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.
81
82
83
84
85
86
87
88
89
        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)
90

91
92
    def custom_iread(self, file_handle, format, index):
        """
93
        """
94
        pass