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

1# fmt: off 

2 

3import numpy as np 

4 

5import ase 

6from ase.data import chemical_symbols 

7from ase.utils import reader, writer 

8 

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

10 

11 

12@writer 

13def write_cfg(fd, atoms): 

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

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

16 """ 

17 

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

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

20 cell = atoms.get_cell(complete=True) 

21 for i in range(3): 

22 for j in range(3): 

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

24 

25 entry_count = 3 

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

27 if x not in cfg_default_fields: 

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

29 entry_count += 1 

30 else: 

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

32 

33 vels = atoms.get_velocities() 

34 if isinstance(vels, np.ndarray): 

35 entry_count += 3 

36 else: 

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

38 

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

40 

41 i = 0 

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

43 if name not in cfg_default_fields: 

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

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

46 i += 1 

47 else: 

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

49 for j in range(3): 

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

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

52 i += 1 

53 

54 else: 

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

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

57 (i, name, j)) 

58 i += 1 

59 

60 # Distinct elements 

61 spos = atoms.get_scaled_positions() 

62 for i in atoms: 

63 el = i.symbol 

64 

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

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

67 

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

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

70 

71 if isinstance(vels, np.ndarray): 

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

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

74 

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

76 if name not in cfg_default_fields: 

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

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

79 else: 

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

81 

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

83 

84 

85default_color = { 

86 'H': [0.800, 0.800, 0.800], 

87 'C': [0.350, 0.350, 0.350], 

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

89 

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

91 

92 

93def write_clr(fd, atoms): 

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

95 Hit F12 in AtomEye to use. 

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

97 """ 

98 color = None 

99 radius = None 

100 if atoms.has('color'): 

101 color = atoms.get_array('color') 

102 if atoms.has('radius'): 

103 radius = atoms.get_array('radius') 

104 

105 if color is None: 

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

107 for a in atoms: 

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

109 

110 if radius is None: 

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

112 for a in atoms: 

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

114 

115 radius.shape = (-1, 1) 

116 

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

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

119 

120 

121@reader 

122def read_cfg(fd): 

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

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

125 """ 

126 nat = None 

127 naux = 0 

128 aux = None 

129 auxstrs = None 

130 

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

132 transform = np.eye(3) 

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

134 

135 current_atom = 0 

136 current_symbol = None 

137 current_mass = None 

138 

139 L = fd.readline() 

140 while L: 

141 L = L.strip() 

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

143 if L == '.NO_VELOCITY.': 

144 vels = None 

145 naux += 3 

146 else: 

147 s = L.split('=') 

148 if len(s) == 2: 

149 key, value = s 

150 key = key.strip() 

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

152 if key == 'Number of particles': 

153 nat = int(value[0]) 

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

155 masses = np.zeros(nat) 

156 syms = [''] * nat 

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

158 if naux > 0: 

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

160 elif key == 'A': 

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

162 elif key == 'entry_count': 

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

164 auxstrs = [''] * naux 

165 if nat is not None: 

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

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

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

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

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

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

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

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

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

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

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

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

178 auxstrs[i] = value[0] 

179 else: 

180 # Everything else must be particle data. 

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

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

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

184 if len(s) == 1: 

185 if L in chemical_symbols: 

186 current_symbol = L 

187 else: 

188 current_mass = float(L) 

189 elif current_symbol is None and current_mass is None: 

190 # Standard CFG format 

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

192 syms[current_atom] = s[1] 

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

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

195 current_atom += 1 

196 elif (current_symbol is not None and 

197 current_mass is not None): 

198 # Extended CFG format 

199 masses[current_atom] = current_mass 

200 syms[current_atom] = current_symbol 

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

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

203 off = 3 

204 if vels is not None: 

205 off = 6 

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

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

208 current_atom += 1 

209 L = fd.readline() 

210 

211 # Sanity check 

212 if current_atom != nat: 

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

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

215 .format(nat, current_atom)) 

216 

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

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

219 'reader.') 

220 cell = np.dot(cell, transform) 

221 

222 if vels is None: 

223 a = ase.Atoms( 

224 symbols=syms, 

225 masses=masses, 

226 scaled_positions=spos, 

227 cell=cell, 

228 pbc=True) 

229 else: 

230 a = ase.Atoms( 

231 symbols=syms, 

232 masses=masses, 

233 scaled_positions=spos, 

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

235 cell=cell, 

236 pbc=True) 

237 

238 i = 0 

239 while i < naux: 

240 auxstr = auxstrs[i] 

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

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

243 i += 3 

244 else: 

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

246 i += 1 

247 

248 return a