1#------------------------------------------------------------------------------#
2#  DFTB+: general package for performing fast atomistic simulations            #
3#  Copyright (C) 2006 - 2019  DFTB+ developers group                           #
4#                                                                              #
5#  See the LICENSE file for terms of usage and distribution.                   #
6#------------------------------------------------------------------------------#
7#
8
9'''Representation of the GEN format.'''
10
11import numpy as np
12from dptools.common import OpenFile
13from dptools.geometry import Geometry
14
15__all__ = ["Gen"]
16
17
18_TOLERANCE = 1E-10
19
20
21class Gen:
22    """Representation of a GEN file.
23
24    Attributes:
25        geometry: Geometry object with atom positions and lattice vectors.
26    """
27
28    def __init__(self, geometry, fractional=False):
29        """Initializes the instance.
30
31        Args:
32            geometry: geometry object containing the geometry information.
33            fractional: Whether fractional coordinates are preferred.
34        """
35        self.geometry = geometry
36        self.fractional = fractional
37
38
39    @classmethod
40    def fromfile(cls, fobj):
41        """Creates a Gen instance from a file.
42
43        Args:
44            fobj: filename or file like object containing geometry in
45                GEN-format.
46        """
47        with OpenFile(fobj, 'r') as fp:
48            alllines = fp.readlines()
49
50        # strip out comments starting with hashmarks
51        lines = []
52        for line in alllines:
53            li = line.partition('#')[0]
54            li = li.rstrip()
55            if li:
56                lines.append(li)
57
58        words = lines[0].split()
59        natom = int(words[0])
60        flag = words[1].lower()
61        if flag == "s":
62            periodic = True
63            relative = False
64        elif flag == "f":
65            periodic = True
66            relative = True
67        else:
68            periodic = False
69            relative = False
70        specienames = lines[1].split()
71        indexes = np.empty((natom, ), dtype=int)
72        coords = np.empty((natom, 3), dtype=float)
73        for ii, line in enumerate(lines[2:2+natom]):
74            words = line.split()
75            indexes[ii] = int(words[1]) - 1
76            coords[ii] = np.array(words[2:5], dtype=float)
77        if periodic:
78            origin = np.array(lines[natom+2].split(), dtype=float)
79            latvecs = np.empty((3, 3), dtype=float)
80            for jj in range(3):
81                latvecs[jj] = np.array(lines[natom+3+jj].split(), dtype=float)
82        else:
83            origin = None
84            latvecs = None
85        geometry = Geometry(specienames, indexes, coords, latvecs, origin,
86                            relative)
87        return cls(geometry, relative)
88
89
90    def tofile(self, fobj):
91        """Writes a GEN file.
92
93        Args:
94            fobj: File name or file object where geometry should be written.
95        """
96        lines = []
97        line = ["{0:d}".format(self.geometry.natom)]
98        geo = self.geometry
99        if geo.periodic:
100            if self.fractional:
101                line.append("F")
102                coords = geo.relcoords
103            else:
104                line.append("S")
105                coords = geo.coords
106        else:
107            line.append("C")
108            coords = geo.coords
109
110        coords = _round_to_zero(coords, _TOLERANCE)
111        lines.append(" ".join(line) + "\n")
112        lines.append(" ".join(geo.specienames) + "\n")
113        for ii in range(geo.natom):
114            lines.append("{0:6d} {1:3d} {2:18.10E} {3:18.10E} {4:18.10E}\n"\
115                         .format(ii + 1, geo.indexes[ii] + 1, *coords[ii]))
116        if geo.periodic:
117            origin = _round_to_zero(geo.origin, _TOLERANCE)
118            lines.append("{0:18.10E} {1:18.10E} {2:18.10E}\n".format(*origin))
119            latvecs = _round_to_zero(geo.latvecs, _TOLERANCE)
120            for vec in latvecs:
121                lines.append("{0:18.10E} {1:18.10E} {2:18.10E}\n".format(*vec))
122        with OpenFile(fobj, 'w') as fp:
123            fp.writelines(lines)
124
125
126    def equals(self, other, tolerance=_TOLERANCE):
127        '''Checks whether object equals to an other one.
128
129        Args:
130            other (Gen): Other Gen object.
131            tolerance (float): Maximal allowed deviation in floating point
132                numbers (e.g. coordinates).
133        '''
134        if not self.fractional == other.fractional:
135            return False
136        if not self.geometry.equals(other.geometry, tolerance):
137            return False
138        return True
139
140
141def _round_to_zero(array, tolerance):
142    '''Rounds elements of an array to zero below given tolerance.'''
143    if tolerance is None:
144        return array
145    else:
146        return np.where(abs(array) < tolerance, 0.0, array)
147