1import numpy as np 

2 

3import ase 

4from ase.data import chemical_symbols 

5from ase.utils import reader, writer 

6 

7 

8cfg_default_fields = np.array(['positions', 'momenta', 'numbers', 'magmoms']) 

9 

10 

11@writer 

12def write_cfg(fd, atoms): 

13 """Write atomic configuration to a CFG-file (native AtomEye format). 

14 See: http://mt.seas.upenn.edu/Archive/Graphics/A/ 

15 """ 

16 

17 fd.write('Number of particles = %i\n' % len(atoms)) 

18 fd.write('A = 1.0 Angstrom\n') 

19 cell = atoms.get_cell(complete=True) 

20 for i in range(3): 

21 for j in range(3): 

22 fd.write('H0(%1.1i,%1.1i) = %f A\n' % (i + 1, j + 1, cell[i, j])) 

23 

24 entry_count = 3 

25 for x in atoms.arrays.keys(): 

26 if x not in cfg_default_fields: 

27 if len(atoms.get_array(x).shape) == 1: 

28 entry_count += 1 

29 else: 

30 entry_count += atoms.get_array(x).shape[1] 

31 

32 vels = atoms.get_velocities() 

33 if isinstance(vels, np.ndarray): 

34 entry_count += 3 

35 else: 

36 fd.write('.NO_VELOCITY.\n') 

37 

38 fd.write('entry_count = %i\n' % entry_count) 

39 

40 i = 0 

41 for name, aux in atoms.arrays.items(): 

42 if name not in cfg_default_fields: 

43 if len(aux.shape) == 1: 

44 fd.write('auxiliary[%i] = %s [a.u.]\n' % (i, name)) 

45 i += 1 

46 else: 

47 if aux.shape[1] == 3: 

48 for j in range(3): 

49 fd.write('auxiliary[%i] = %s_%s [a.u.]\n' % 

50 (i, name, chr(ord('x') + j))) 

51 i += 1 

52 

53 else: 

54 for j in range(aux.shape[1]): 

55 fd.write('auxiliary[%i] = %s_%1.1i [a.u.]\n' % 

56 (i, name, j)) 

57 i += 1 

58 

59 # Distinct elements 

60 spos = atoms.get_scaled_positions() 

61 for i in atoms: 

62 el = i.symbol 

63 

64 fd.write('%f\n' % ase.data.atomic_masses[chemical_symbols.index(el)]) 

65 fd.write('%s\n' % el) 

66 

67 x, y, z = spos[i.index, :] 

68 s = '%e %e %e ' % (x, y, z) 

69 

70 if isinstance(vels, np.ndarray): 

71 vx, vy, vz = vels[i.index, :] 

72 s = s + ' %e %e %e ' % (vx, vy, vz) 

73 

74 for name, aux in atoms.arrays.items(): 

75 if name not in cfg_default_fields: 

76 if len(aux.shape) == 1: 

77 s += ' %e' % aux[i.index] 

78 else: 

79 s += (aux.shape[1] * ' %e') % tuple(aux[i.index].tolist()) 

80 

81 fd.write('%s\n' % s) 

82 

83 

84default_color = { 

85 'H': [0.800, 0.800, 0.800], 

86 'C': [0.350, 0.350, 0.350], 

87 'O': [0.800, 0.200, 0.200]} 

88 

89default_radius = {'H': 0.435, 'C': 0.655, 'O': 0.730} 

90 

91 

92def write_clr(fd, atoms): 

93 """Write extra color and radius code to a CLR-file (for use with AtomEye). 

94 Hit F12 in AtomEye to use. 

95 See: http://mt.seas.upenn.edu/Archive/Graphics/A/ 

96 """ 

97 color = None 

98 radius = None 

99 if atoms.has('color'): 

100 color = atoms.get_array('color') 

101 if atoms.has('radius'): 

102 radius = atoms.get_array('radius') 

103 

104 if color is None: 

105 color = np.zeros([len(atoms), 3], dtype=float) 

106 for a in atoms: 

107 color[a.index, :] = default_color[a.symbol] 

108 

109 if radius is None: 

110 radius = np.zeros(len(atoms), dtype=float) 

111 for a in atoms: 

112 radius[a.index] = default_radius[a.symbol] 

113 

114 radius.shape = (-1, 1) 

115 

116 for c1, c2, c3, r in np.append(color, radius, axis=1): 

117 fd.write('%f %f %f %f\n' % (c1, c2, c3, r)) 

118 

119 

120@reader 

121def read_cfg(fd): 

122 """Read atomic configuration from a CFG-file (native AtomEye format). 

123 See: http://mt.seas.upenn.edu/Archive/Graphics/A/ 

124 """ 

125 nat = None 

126 naux = 0 

127 aux = None 

128 auxstrs = None 

129 

130 cell = np.zeros([3, 3]) 

131 transform = np.eye(3) 

132 eta = np.zeros([3, 3]) 

133 

134 current_atom = 0 

135 current_symbol = None 

136 current_mass = None 

137 

138 L = fd.readline() 

139 while L: 

140 L = L.strip() 

141 if len(L) != 0 and not L.startswith('#'): 

142 if L == '.NO_VELOCITY.': 

143 vels = None 

144 naux += 3 

145 else: 

146 s = L.split('=') 

147 if len(s) == 2: 

148 key, value = s 

149 key = key.strip() 

150 value = [x.strip() for x in value.split()] 

151 if key == 'Number of particles': 

152 nat = int(value[0]) 

153 spos = np.zeros([nat, 3]) 

154 masses = np.zeros(nat) 

155 syms = [''] * nat 

156 vels = np.zeros([nat, 3]) 

157 if naux > 0: 

158 aux = np.zeros([nat, naux]) 

159 elif key == 'A': 

160 pass # unit = float(value[0]) 

161 elif key == 'entry_count': 

162 naux += int(value[0]) - 6 

163 auxstrs = [''] * naux 

164 if nat is not None: 

165 aux = np.zeros([nat, naux]) 

166 elif key.startswith('H0('): 

167 i, j = [int(x) for x in key[3:-1].split(',')] 

168 cell[i - 1, j - 1] = float(value[0]) 

169 elif key.startswith('Transform('): 

170 i, j = [int(x) for x in key[10:-1].split(',')] 

171 transform[i - 1, j - 1] = float(value[0]) 

172 elif key.startswith('eta('): 

173 i, j = [int(x) for x in key[4:-1].split(',')] 

174 eta[i - 1, j - 1] = float(value[0]) 

175 elif key.startswith('auxiliary['): 

176 i = int(key[10:-1]) 

177 auxstrs[i] = value[0] 

178 else: 

179 # Everything else must be particle data. 

180 # First check if current line contains an element mass or 

181 # name. Then we have an extended XYZ format. 

182 s = [x.strip() for x in L.split()] 

183 if len(s) == 1: 

184 if L in chemical_symbols: 

185 current_symbol = L 

186 else: 

187 current_mass = float(L) 

188 elif current_symbol is None and current_mass is None: 

189 # Standard CFG format 

190 masses[current_atom] = float(s[0]) 

191 syms[current_atom] = s[1] 

192 spos[current_atom, :] = [float(x) for x in s[2:5]] 

193 vels[current_atom, :] = [float(x) for x in s[5:8]] 

194 current_atom += 1 

195 elif (current_symbol is not None and 

196 current_mass is not None): 

197 # Extended CFG format 

198 masses[current_atom] = current_mass 

199 syms[current_atom] = current_symbol 

200 props = [float(x) for x in s] 

201 spos[current_atom, :] = props[0:3] 

202 off = 3 

203 if vels is not None: 

204 off = 6 

205 vels[current_atom, :] = props[3:6] 

206 aux[current_atom, :] = props[off:] 

207 current_atom += 1 

208 L = fd.readline() 

209 

210 # Sanity check 

211 if current_atom != nat: 

212 raise RuntimeError('Number of atoms reported for CFG file (={0}) and ' 

213 'number of atoms actually read (={1}) differ.' 

214 .format(nat, current_atom)) 

215 

216 if np.any(eta != 0): 

217 raise NotImplementedError('eta != 0 not yet implemented for CFG ' 

218 'reader.') 

219 cell = np.dot(cell, transform) 

220 

221 if vels is None: 

222 a = ase.Atoms( 

223 symbols=syms, 

224 masses=masses, 

225 scaled_positions=spos, 

226 cell=cell, 

227 pbc=True) 

228 else: 

229 a = ase.Atoms( 

230 symbols=syms, 

231 masses=masses, 

232 scaled_positions=spos, 

233 momenta=masses.reshape(-1, 1) * vels, 

234 cell=cell, 

235 pbc=True) 

236 

237 i = 0 

238 while i < naux: 

239 auxstr = auxstrs[i] 

240 if auxstr[-2:] == '_x': 

241 a.set_array(auxstr[:-2], aux[:, i:i + 3]) 

242 i += 3 

243 else: 

244 a.set_array(auxstr, aux[:, i]) 

245 i += 1 

246 

247 return a