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 

22class _ZMatrixToAtoms: 

23 known_units = dict( 

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

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

26 ) 

27 

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

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

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

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

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

33 self.set_defs(defs) 

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

35 self.symbols = [] # type: List[str] 

36 self.positions = [] # type: List[Tuple[float, float, float]] 

37 

38 @property 

39 def nrows(self): 

40 return len(self.symbols) 

41 

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

43 if isinstance(value, Real): 

44 return float(value) 

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

46 if out is None: 

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

48 .format(kind, value)) 

49 return out 

50 

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

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

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

54 if defs is None: 

55 return 

56 

57 if isinstance(defs, dict): 

58 self.defs.update(**defs) 

59 return 

60 

61 if isinstance(defs, str): 

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

63 

64 for row in defs: 

65 key, val = _re_defs.split(row) 

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

67 

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

69 try: 

70 return float(val) 

71 except ValueError as e: 

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

73 if val_out is None: 

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

75 .format(val)) from e 

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

77 

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

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

80 try: 

81 return int(name) - 1 

82 except ValueError as e: 

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

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

85 .format(name)) from e 

86 return self.name_to_index[name] 

87 

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

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

90 if self.name_to_index is None: 

91 return 

92 

93 if name in self.name_to_index: 

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

95 # longer meaningful. Destroy the map. 

96 self.name_to_index = None 

97 return 

98 

99 self.name_to_index[name] = self.nrows 

100 

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

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

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

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

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

106 "hasn't been defined yet!" 

107 .format(self.nrows, indices)) 

108 

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

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

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

112 'referred indices: {}' 

113 .format(self.nrows, indices)) 

114 

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

116 str, Union[_ZMatrixRow, Tuple[float, float, float]], 

117 ]: 

118 tokens = row.split() 

119 name = tokens[0] 

120 self.set_index(name) 

121 if len(tokens) == 1: 

122 assert self.nrows == 0 

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

124 

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

126 if ind1 == -1: 

127 assert len(tokens) == 5 

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

129 dtype=float) 

130 

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

132 

133 if len(tokens) == 3: 

134 assert self.nrows == 1 

135 self.validate_indices(ind1) 

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

137 

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

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

140 

141 if len(tokens) == 5: 

142 assert self.nrows == 2 

143 self.validate_indices(ind1, ind2) 

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

145 

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

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

148 self.validate_indices(ind1, ind2, ind3) 

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

150 a_dihedral) 

151 

152 def add_atom(self, name: str, pos: Tuple[float, float, float]) -> None: 

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

154 self.symbols.append( 

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

156 ) 

157 self.positions.append(pos) 

158 

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

160 name, zrow = self.parse_row(row) 

161 

162 if not isinstance(zrow, _ZMatrixRow): 

163 self.add_atom(name, zrow) 

164 return 

165 

166 if zrow.ind3 is None: 

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

168 # have been provided. 

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

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

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

172 self.add_atom(name, pos) 

173 return 

174 

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

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

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

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

179 

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

181 # to the dihedral axis 

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

183 ax2 -= ax1 * (ax2 @ ax1) 

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

185 

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

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

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

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

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

191 

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

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

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

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

196 self.add_atom(name, pos) 

197 

198 def to_atoms(self) -> Atoms: 

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

200 

201 

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

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

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

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

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

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

208 

209 Parameters: 

210 

211 zmat: Iterable or str 

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

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

214 into a list at newlines. 

215 distance_units: str or float, optional 

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

217 Defaults to Angstrom. 

218 angle_units: str or float, optional 

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

220 Defaults to degrees. 

221 defs: dict or str, optional 

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

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

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

225 argument. 

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

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

228 the symbol definitions from the Z-matrix. 

229 

230 Returns: 

231 

232 atoms: Atoms object 

233 """ 

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

235 

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

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

238 if isinstance(zmat, str): 

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

240 

241 for row in zmat: 

242 zmatrix.add_row(row) 

243 

244 return zmatrix.to_atoms()