1from typing import Dict, List, Tuple, Union, Optional 

2from numbers import Real 

3from collections import namedtuple 

4import re 

5from string import digits 

6import numpy as np 

7from ase import Atoms 

8from ase.units import Angstrom, Bohr, nm 

9 

10 

11# split on newlines or semicolons 

12_re_linesplit = re.compile(r'\n|;') 

13# split definitions on whitespace or on "=" (possibly also with whitespace) 

14_re_defs = re.compile(r'\s*=\s*|\s+') 

15 

16 

17_ZMatrixRow = namedtuple( 

18 '_ZMatrixRow', 'ind1 dist ind2 a_bend ind3 a_dihedral', 

19) 

20 

21 

22ThreeFloats = Union[Tuple[float, float, float], np.ndarray] 

23 

24 

25class _ZMatrixToAtoms: 

26 known_units = dict( 

27 distance={'angstrom': Angstrom, 'bohr': Bohr, 'au': Bohr, 'nm': nm}, 

28 angle={'radians': 1., 'degrees': np.pi / 180}, 

29 ) 

30 

31 def __init__(self, dconv: Union[str, Real], aconv: Union[str, Real], 

32 defs: Optional[Union[Dict[str, float], 

33 str, List[str]]] = None) -> None: 

34 self.dconv = self.get_units('distance', dconv) # type: float 

35 self.aconv = self.get_units('angle', aconv) # type: float 

36 self.set_defs(defs) 

37 self.name_to_index: Optional[Dict[str, int]] = dict() 

38 self.symbols: List[str] = [] 

39 self.positions: List[ThreeFloats] = [] 

40 

41 @property 

42 def nrows(self): 

43 return len(self.symbols) 

44 

45 def get_units(self, kind: str, value: Union[str, Real]) -> float: 

46 if isinstance(value, Real): 

47 return float(value) 

48 out = self.known_units[kind].get(value.lower()) 

49 if out is None: 

50 raise ValueError("Unknown {} units: {}" 

51 .format(kind, value)) 

52 return out 

53 

54 def set_defs(self, defs: Union[Dict[str, float], str, 

55 List[str], None]) -> None: 

56 self.defs = dict() # type: Dict[str, float] 

57 if defs is None: 

58 return 

59 

60 if isinstance(defs, dict): 

61 self.defs.update(**defs) 

62 return 

63 

64 if isinstance(defs, str): 

65 defs = _re_linesplit.split(defs.strip()) 

66 

67 for row in defs: 

68 key, val = _re_defs.split(row) 

69 self.defs[key] = self.get_var(val) 

70 

71 def get_var(self, val: str) -> float: 

72 try: 

73 return float(val) 

74 except ValueError as e: 

75 val_out = self.defs.get(val.lstrip('+-')) 

76 if val_out is None: 

77 raise ValueError('Invalid value encountered in Z-matrix: {}' 

78 .format(val)) from e 

79 return val_out * (-1 if val.startswith('-') else 1) 

80 

81 def get_index(self, name: str) -> int: 

82 """Find index for a given atom name""" 

83 try: 

84 return int(name) - 1 

85 except ValueError as e: 

86 if self.name_to_index is None or name not in self.name_to_index: 

87 raise ValueError('Failed to determine index for name "{}"' 

88 .format(name)) from e 

89 return self.name_to_index[name] 

90 

91 def set_index(self, name: str) -> None: 

92 """Assign index to a given atom name for name -> index lookup""" 

93 if self.name_to_index is None: 

94 return 

95 

96 if name in self.name_to_index: 

97 # "name" has been encountered before, so name_to_index is no 

98 # longer meaningful. Destroy the map. 

99 self.name_to_index = None 

100 return 

101 

102 self.name_to_index[name] = self.nrows 

103 

104 def validate_indices(self, *indices: int) -> None: 

105 """Raises an error if indices in a Z-matrix row are invalid.""" 

106 if any(np.array(indices) >= self.nrows): 

107 raise ValueError('An invalid Z-matrix was provided! Row {} refers ' 

108 'to atom indices {}, at least one of which ' 

109 "hasn't been defined yet!" 

110 .format(self.nrows, indices)) 

111 

112 if len(indices) != len(set(indices)): 

113 raise ValueError('An atom index has been used more than once a ' 

114 'row of the Z-matrix! Row numbers {}, ' 

115 'referred indices: {}' 

116 .format(self.nrows, indices)) 

117 

118 def parse_row(self, row: str) -> Tuple[ 

119 str, Union[_ZMatrixRow, ThreeFloats], 

120 ]: 

121 tokens = row.split() 

122 name = tokens[0] 

123 self.set_index(name) 

124 if len(tokens) == 1: 

125 assert self.nrows == 0 

126 return name, np.zeros(3, dtype=float) 

127 

128 ind1 = self.get_index(tokens[1]) 

129 if ind1 == -1: 

130 assert len(tokens) == 5 

131 return name, np.array(list(map(self.get_var, tokens[2:])), 

132 dtype=float) 

133 

134 dist = self.dconv * self.get_var(tokens[2]) 

135 

136 if len(tokens) == 3: 

137 assert self.nrows == 1 

138 self.validate_indices(ind1) 

139 return name, np.array([dist, 0, 0], dtype=float) 

140 

141 ind2 = self.get_index(tokens[3]) 

142 a_bend = self.aconv * self.get_var(tokens[4]) 

143 

144 if len(tokens) == 5: 

145 assert self.nrows == 2 

146 self.validate_indices(ind1, ind2) 

147 return name, _ZMatrixRow(ind1, dist, ind2, a_bend, None, None) 

148 

149 ind3 = self.get_index(tokens[5]) 

150 a_dihedral = self.aconv * self.get_var(tokens[6]) 

151 self.validate_indices(ind1, ind2, ind3) 

152 return name, _ZMatrixRow(ind1, dist, ind2, a_bend, ind3, 

153 a_dihedral) 

154 

155 def add_atom(self, name: str, pos: ThreeFloats) -> None: 

156 """Sets the symbol and position of an atom.""" 

157 self.symbols.append( 

158 ''.join([c for c in name if c not in digits]).capitalize() 

159 ) 

160 self.positions.append(pos) 

161 

162 def add_row(self, row: str) -> None: 

163 name, zrow = self.parse_row(row) 

164 

165 if not isinstance(zrow, _ZMatrixRow): 

166 self.add_atom(name, zrow) 

167 return 

168 

169 if zrow.ind3 is None: 

170 # This is the third atom, so only a bond distance and an angle 

171 # have been provided. 

172 pos = self.positions[zrow.ind1].copy() 

173 pos[0] += zrow.dist * np.cos(zrow.a_bend) * (zrow.ind2 - zrow.ind1) 

174 pos[1] += zrow.dist * np.sin(zrow.a_bend) 

175 self.add_atom(name, pos) 

176 return 

177 

178 # ax1 is the dihedral axis, which is defined by the bond vector 

179 # between the two inner atoms in the dihedral, ind1 and ind2 

180 ax1 = self.positions[zrow.ind2] - self.positions[zrow.ind1] 

181 ax1 /= np.linalg.norm(ax1) 

182 

183 # ax2 lies within the 1-2-3 plane, and it is perpendicular 

184 # to the dihedral axis 

185 ax2 = self.positions[zrow.ind2] - self.positions[zrow.ind3] 

186 ax2 -= ax1 * (ax2 @ ax1) 

187 ax2 /= np.linalg.norm(ax2) 

188 

189 # ax3 is a vector that forms the appropriate dihedral angle, though 

190 # the bending angle is 90 degrees, rather than a_bend. It is formed 

191 # from a linear combination of ax2 and (ax2 x ax1) 

192 ax3 = (ax2 * np.cos(zrow.a_dihedral) 

193 + np.cross(ax2, ax1) * np.sin(zrow.a_dihedral)) 

194 

195 # The final position vector is a linear combination of ax1 and ax3. 

196 pos = ax1 * np.cos(zrow.a_bend) - ax3 * np.sin(zrow.a_bend) 

197 pos *= zrow.dist / np.linalg.norm(pos) 

198 pos += self.positions[zrow.ind1] 

199 self.add_atom(name, pos) 

200 

201 def to_atoms(self) -> Atoms: 

202 return Atoms(self.symbols, self.positions) 

203 

204 

205def parse_zmatrix(zmat: Union[str, List[str]], 

206 distance_units: Union[str, Real] = 'angstrom', 

207 angle_units: Union[str, Real] = 'degrees', 

208 defs: Optional[Union[Dict[str, float], str, 

209 List[str]]] = None) -> Atoms: 

210 """Converts a Z-matrix into an Atoms object. 

211 

212 Parameters: 

213 

214 zmat: Iterable or str 

215 The Z-matrix to be parsed. Iteration over `zmat` should yield the rows 

216 of the Z-matrix. If `zmat` is a str, it will be automatically split 

217 into a list at newlines. 

218 distance_units: str or float, optional 

219 The units of distance in the provided Z-matrix. 

220 Defaults to Angstrom. 

221 angle_units: str or float, optional 

222 The units for angles in the provided Z-matrix. 

223 Defaults to degrees. 

224 defs: dict or str, optional 

225 If `zmat` contains symbols for bond distances, bending angles, and/or 

226 dihedral angles instead of numeric values, then the definition of 

227 those symbols should be passed to this function using this keyword 

228 argument. 

229 Note: The symbol definitions are typically printed adjacent to the 

230 Z-matrix itself, but this function will not automatically separate 

231 the symbol definitions from the Z-matrix. 

232 

233 Returns: 

234 

235 atoms: Atoms object 

236 """ 

237 zmatrix = _ZMatrixToAtoms(distance_units, angle_units, defs=defs) 

238 

239 # zmat should be a list containing the rows of the z-matrix. 

240 # for convenience, allow block strings and split at newlines. 

241 if isinstance(zmat, str): 

242 zmat = _re_linesplit.split(zmat.strip()) 

243 

244 for row in zmat: 

245 zmatrix.add_row(row) 

246 

247 return zmatrix.to_atoms()