Coverage for /builds/kinetik161/ase/ase/io/cfg.py: 78.05%

1import numpy as np 

2 

3import ase 

4from ase.data import chemical_symbols 

5from ase.utils import reader, writer 

6 

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

8 

9 

10@writer 

11def write_cfg(fd, atoms): 

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

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

14 """ 

15 

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

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

18 cell = atoms.get_cell(complete=True) 

19 for i in range(3): 

20 for j in range(3): 

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

22 

23 entry_count = 3 

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

25 if x not in cfg_default_fields: 

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

27 entry_count += 1 

28 else: 

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

30 

31 vels = atoms.get_velocities() 

32 if isinstance(vels, np.ndarray): 

33 entry_count += 3 

34 else: 

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

36 

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

38 

39 i = 0 

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

41 if name not in cfg_default_fields: 

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

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

44 i += 1 

45 else: 

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

47 for j in range(3): 

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

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

50 i += 1 

51 

52 else: 

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

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

55 (i, name, j)) 

56 i += 1 

57 

58 # Distinct elements 

59 spos = atoms.get_scaled_positions() 

60 for i in atoms: 

61 el = i.symbol 

62 

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

64 fd.write(f'{el}\n') 

65 

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

67 s = f'{x:e} {y:e} {z:e} ' 

68 

69 if isinstance(vels, np.ndarray): 

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

71 s = s + f' {vx:e} {vy:e} {vz:e} ' 

72 

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

74 if name not in cfg_default_fields: 

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

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

77 else: 

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

79 

80 fd.write(f'{s}\n') 

81 

82 

83default_color = { 

84 'H': [0.800, 0.800, 0.800], 

85 'C': [0.350, 0.350, 0.350], 

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

87 

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

89 

90 

91def write_clr(fd, atoms): 

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

93 Hit F12 in AtomEye to use. 

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

95 """ 

96 color = None 

97 radius = None 

98 if atoms.has('color'): 

99 color = atoms.get_array('color') 

100 if atoms.has('radius'): 

101 radius = atoms.get_array('radius') 

102 

103 if color is None: 

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

105 for a in atoms: 

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

107 

108 if radius is None: 

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

110 for a in atoms: 

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

112 

113 radius.shape = (-1, 1) 

114 

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

116 fd.write(f'{c1:f} {c2:f} {c3:f} {r:f}\n') 

117 

118 

119@reader 

120def read_cfg(fd): 

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

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

123 """ 

124 nat = None 

125 naux = 0 

126 aux = None 

127 auxstrs = None 

128 

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

130 transform = np.eye(3) 

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

132 

133 current_atom = 0 

134 current_symbol = None 

135 current_mass = None 

136 

137 L = fd.readline() 

138 while L: 

139 L = L.strip() 

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

141 if L == '.NO_VELOCITY.': 

142 vels = None 

143 naux += 3 

144 else: 

145 s = L.split('=') 

146 if len(s) == 2: 

147 key, value = s 

148 key = key.strip() 

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

150 if key == 'Number of particles': 

151 nat = int(value[0]) 

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

153 masses = np.zeros(nat) 

154 syms = [''] * nat 

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

156 if naux > 0: 

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

158 elif key == 'A': 

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

160 elif key == 'entry_count': 

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

162 auxstrs = [''] * naux 

163 if nat is not None: 

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

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

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

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

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

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

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

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

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

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

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

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

176 auxstrs[i] = value[0] 

177 else: 

178 # Everything else must be particle data. 

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

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

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

182 if len(s) == 1: 

183 if L in chemical_symbols: 

184 current_symbol = L 

185 else: 

186 current_mass = float(L) 

187 elif current_symbol is None and current_mass is None: 

188 # Standard CFG format 

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

190 syms[current_atom] = s[1] 

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

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

193 current_atom += 1 

194 elif (current_symbol is not None and 

195 current_mass is not None): 

196 # Extended CFG format 

197 masses[current_atom] = current_mass 

198 syms[current_atom] = current_symbol 

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

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

201 off = 3 

202 if vels is not None: 

203 off = 6 

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

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

206 current_atom += 1 

207 L = fd.readline() 

208 

209 # Sanity check 

210 if current_atom != nat: 

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

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

213 .format(nat, current_atom)) 

214 

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

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

217 'reader.') 

218 cell = np.dot(cell, transform) 

219 

220 if vels is None: 

221 a = ase.Atoms( 

222 symbols=syms, 

223 masses=masses, 

224 scaled_positions=spos, 

225 cell=cell, 

226 pbc=True) 

227 else: 

228 a = ase.Atoms( 

229 symbols=syms, 

230 masses=masses, 

231 scaled_positions=spos, 

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

233 cell=cell, 

234 pbc=True) 

235 

236 i = 0 

237 while i < naux: 

238 auxstr = auxstrs[i] 

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

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

241 i += 3 

242 else: 

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

244 i += 1 

245 

246 return a