1import numpy as np 

2from ase.neighborlist import NeighborList 

3from ase.data import atomic_masses, chemical_symbols 

4from ase import Atoms 

5 

6 

7def find_nearest_index(array, value): 

8 array = np.asarray(array) 

9 idx = (np.abs(array - value)).argmin() 

10 return idx 

11 

12 

13def find_nearest_value(array, value): 

14 array = np.asarray(array) 

15 idx = (np.abs(array - value)).argmin() 

16 return array[idx] 

17 

18 

19def write_gpumd(fd, atoms, maximum_neighbors=None, cutoff=None, 

20 groupings=None, use_triclinic=False, species=None): 

21 """ 

22 Writes atoms into GPUMD input format. 

23 

24 Parameters 

25 ---------- 

26 fd : file 

27 File like object to which the atoms object should be written 

28 atoms : Atoms 

29 Input structure 

30 maximum_neighbors: int 

31 Maximum number of neighbors any atom can ever have (not relevant when 

32 using force constant potentials) 

33 cutoff: float 

34 Initial cutoff distance used for building the neighbor list (not 

35 relevant when using force constant potentials) 

36 groupings : list[list[list[int]]] 

37 Groups into which the individual atoms should be divided in the form of 

38 a list of list of lists. Specifically, the outer list corresponds to 

39 the grouping methods, of which there can be three at the most, which 

40 contains a list of groups in the form of lists of site indices. The 

41 sum of the lengths of the latter must be the same as the total number 

42 of atoms. 

43 use_triclinic: bool 

44 Use format for triclinic cells 

45 species : List[str] 

46 GPUMD uses integers to define atom types. This list allows customized 

47 such definitions (e.g, ['Pd', 'H'] means Pd is type 0 and H type 1). 

48 If None, this list is built by assigning each distinct species to 

49 an integer in the order of appearance in `atoms`. 

50 

51 Raises 

52 ------ 

53 ValueError 

54 Raised if parameters are incompatible 

55 """ 

56 

57 # Check velocties parameter 

58 if atoms.get_velocities() is None: 

59 has_velocity = 0 

60 else: 

61 has_velocity = 1 

62 velocities = atoms.get_velocities() 

63 

64 # Check groupings parameter 

65 if groupings is None: 

66 number_of_grouping_methods = 0 

67 else: 

68 number_of_grouping_methods = len(groupings) 

69 if number_of_grouping_methods > 3: 

70 raise ValueError('There can be no more than 3 grouping methods!') 

71 for g, grouping in enumerate(groupings): 

72 all_indices = [i for group in grouping for i in group] 

73 if len(all_indices) != len(atoms) or\ 

74 set(all_indices) != set(range(len(atoms))): 

75 raise ValueError('The indices listed in grouping method {} are' 

76 ' not compatible with the input' 

77 ' structure!'.format(g)) 

78 

79 # If not specified, estimate the maximum_neighbors 

80 if maximum_neighbors is None: 

81 if cutoff is None: 

82 cutoff = 0.1 

83 maximum_neighbors = 1 

84 else: 

85 nl = NeighborList([cutoff / 2] * len(atoms), skin=2, bothways=True) 

86 nl.update(atoms) 

87 maximum_neighbors = 0 

88 for atom in atoms: 

89 maximum_neighbors = max(maximum_neighbors, 

90 len(nl.get_neighbors(atom.index)[0])) 

91 maximum_neighbors *= 2 

92 maximum_neighbors = min(maximum_neighbors, 1024) 

93 

94 # Add header and cell parameters 

95 lines = [] 

96 if atoms.cell.orthorhombic and not use_triclinic: 

97 triclinic = 0 

98 else: 

99 triclinic = 1 

100 lines.append('{} {} {} {} {} {}'.format(len(atoms), maximum_neighbors, 

101 cutoff, triclinic, has_velocity, 

102 number_of_grouping_methods)) 

103 if triclinic: 

104 lines.append((' {}' * 12)[1:].format(*atoms.pbc.astype(int), 

105 *atoms.cell[:].flatten())) 

106 else: 

107 lines.append((' {}' * 6)[1:].format(*atoms.pbc.astype(int), 

108 *atoms.cell.lengths())) 

109 

110 # Create symbols-to-type map, i.e. integers starting at 0 

111 if not species: 

112 symbol_type_map = {} 

113 for symbol in atoms.get_chemical_symbols(): 

114 if symbol not in symbol_type_map: 

115 symbol_type_map[symbol] = len(symbol_type_map) 

116 else: 

117 if any([sym not in species 

118 for sym in set(atoms.get_chemical_symbols())]): 

119 raise ValueError('The species list does not contain all chemical ' 

120 'species that are present in the atoms object.') 

121 else: 

122 symbol_type_map = {symbol: i for i, symbol in enumerate(species)} 

123 

124 # Add lines for all atoms 

125 for a, atm in enumerate(atoms): 

126 t = symbol_type_map[atm.symbol] 

127 line = (' {}' * 5)[1:].format(t, *atm.position, atm.mass) 

128 if has_velocity: 

129 line += (' {}' * 3).format(*velocities[a]) 

130 if groupings is not None: 

131 for grouping in groupings: 

132 for i, group in enumerate(grouping): 

133 if a in group: 

134 line += ' {}'.format(i) 

135 break 

136 lines.append(line) 

137 

138 # Write file 

139 fd.write('\n'.join(lines)) 

140 

141 

142def load_xyz_input_gpumd(fd, species=None, isotope_masses=None): 

143 """ 

144 Read the structure input file for GPUMD and return an ase Atoms object 

145 togehter with a dictionary with parameters and a types-to-symbols map 

146 

147 Parameters 

148 ---------- 

149 fd : file | str 

150 File object or name of file from which to read the Atoms object 

151 species : List[str] 

152 List with the chemical symbols that correspond to each type, will take 

153 precedence over isotope_masses 

154 isotope_masses: Dict[str, List[float]] 

155 Dictionary with chemical symbols and lists of the associated atomic 

156 masses, which is used to identify the chemical symbols that correspond 

157 to the types not found in species_types. The default is to find the 

158 closest match :data:`ase.data.atomic_masses`. 

159 

160 Returns 

161 ------- 

162 atoms : Atoms 

163 Atoms object 

164 input_parameters : Dict[str, int] 

165 Dictionary with parameters from the first row of the input file, namely 

166 'N', 'M', 'cutoff', 'triclinic', 'has_velocity' and 'num_of_groups' 

167 species : List[str] 

168 List with the chemical symbols that correspond to each type 

169 

170 Raises 

171 ------ 

172 ValueError 

173 Raised if the list of species is incompatible with the input file 

174 """ 

175 # Parse first line 

176 first_line = next(fd) 

177 input_parameters = {} 

178 keys = ['N', 'M', 'cutoff', 'triclinic', 'has_velocity', 

179 'num_of_groups'] 

180 types = [float if key == 'cutoff' else int for key in keys] 

181 for k, (key, typ) in enumerate(zip(keys, types)): 

182 input_parameters[key] = typ(first_line.split()[k]) 

183 

184 # Parse second line 

185 second_line = next(fd) 

186 second_arr = np.array(second_line.split()) 

187 pbc = second_arr[:3].astype(bool) 

188 if input_parameters['triclinic']: 

189 cell = second_arr[3:].astype(float).reshape((3, 3)) 

190 else: 

191 cell = np.diag(second_arr[3:].astype(float)) 

192 

193 # Parse all remaining rows 

194 n_rows = input_parameters['N'] 

195 n_columns = 5 + input_parameters['has_velocity'] * 3 +\ 

196 input_parameters['num_of_groups'] 

197 rest_lines = [next(fd) for _ in range(n_rows)] 

198 rest_arr = np.array([line.split() for line in rest_lines]) 

199 assert rest_arr.shape == (n_rows, n_columns) 

200 

201 # Extract atom types, positions and masses 

202 atom_types = rest_arr[:, 0].astype(int) 

203 positions = rest_arr[:, 1:4].astype(float) 

204 masses = rest_arr[:, 4].astype(float) 

205 

206 # Determine the atomic species 

207 if species is None: 

208 type_symbol_map = {} 

209 if isotope_masses is not None: 

210 mass_symbols = {mass: symbol for symbol, masses in 

211 isotope_masses.items() for mass in masses} 

212 symbols = [] 

213 for atom_type, mass in zip(atom_types, masses): 

214 if species is None: 

215 if atom_type not in type_symbol_map: 

216 if isotope_masses is not None: 

217 nearest_value = find_nearest_value( 

218 list(mass_symbols.keys()), mass) 

219 symbol = mass_symbols[nearest_value] 

220 else: 

221 symbol = chemical_symbols[ 

222 find_nearest_index(atomic_masses, mass)] 

223 type_symbol_map[atom_type] = symbol 

224 else: 

225 symbol = type_symbol_map[atom_type] 

226 else: 

227 if atom_type > len(species): 

228 raise Exception('There is no entry for atom type {} in the ' 

229 'species list!'.format(atom_type)) 

230 symbol = species[atom_type] 

231 symbols.append(symbol) 

232 

233 if species is None: 

234 species = [type_symbol_map[i] for i in sorted(type_symbol_map.keys())] 

235 

236 # Create the Atoms object 

237 atoms = Atoms(symbols=symbols, positions=positions, masses=masses, pbc=pbc, 

238 cell=cell) 

239 if input_parameters['has_velocity']: 

240 velocities = rest_arr[:, 5:8].astype(float) 

241 atoms.set_velocities(velocities) 

242 if input_parameters['num_of_groups']: 

243 start_col = 5 + 3 * input_parameters['has_velocity'] 

244 groups = rest_arr[:, start_col:].astype(int) 

245 atoms.info = {i: {'groups': groups[i, :]} for i in range(n_rows)} 

246 

247 return atoms, input_parameters, species 

248 

249 

250def read_gpumd(fd, species=None, isotope_masses=None): 

251 """ 

252 Read Atoms object from a GPUMD structure input file 

253 

254 Parameters 

255 ---------- 

256 fd : file | str 

257 File object or name of file from which to read the Atoms object 

258 species : List[str] 

259 List with the chemical symbols that correspond to each type, will take 

260 precedence over isotope_masses 

261 isotope_masses: Dict[str, List[float]] 

262 Dictionary with chemical symbols and lists of the associated atomic 

263 masses, which is used to identify the chemical symbols that correspond 

264 to the types not found in species_types. The default is to find the 

265 closest match :data:`ase.data.atomic_masses`. 

266 

267 Returns 

268 ------- 

269 atoms : Atoms 

270 Atoms object 

271 

272 Raises 

273 ------ 

274 ValueError 

275 Raised if the list of species is incompatible with the input file 

276 """ 

277 

278 return load_xyz_input_gpumd(fd, species, isotope_masses)[0]