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): 

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 

46 Raises 

47 ------ 

48 ValueError 

49 Raised if parameters are incompatible 

50 """ 

51 

52 # Check velocties parameter 

53 if atoms.get_velocities() is None: 

54 has_velocity = 0 

55 else: 

56 has_velocity = 1 

57 velocities = atoms.get_velocities() 

58 

59 # Check groupings parameter 

60 if groupings is None: 

61 number_of_grouping_methods = 0 

62 else: 

63 number_of_grouping_methods = len(groupings) 

64 if number_of_grouping_methods > 3: 

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

66 for g, grouping in enumerate(groupings): 

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

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

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

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

71 ' not compatible with the input' 

72 ' structure!'.format(g)) 

73 

74 # If not specified, estimate the maximum_neighbors 

75 if maximum_neighbors is None: 

76 if cutoff is None: 

77 cutoff = 0.1 

78 maximum_neighbors = 1 

79 else: 

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

81 nl.update(atoms) 

82 maximum_neighbors = 0 

83 for atom in atoms: 

84 maximum_neighbors = max(maximum_neighbors, 

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

86 maximum_neighbors *= 2 

87 

88 # Add header and cell parameters 

89 lines = [] 

90 if atoms.cell.orthorhombic and not use_triclinic: 

91 triclinic = 0 

92 else: 

93 triclinic = 1 

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

95 cutoff, triclinic, has_velocity, 

96 number_of_grouping_methods)) 

97 if triclinic: 

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

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

100 else: 

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

102 *atoms.cell.lengths())) 

103 

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

105 symbol_type_map = {} 

106 for symbol in atoms.get_chemical_symbols(): 

107 if symbol not in symbol_type_map: 

108 symbol_type_map[symbol] = len(symbol_type_map) 

109 

110 # Add lines for all atoms 

111 for a, atm in enumerate(atoms): 

112 t = symbol_type_map[atm.symbol] 

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

114 if has_velocity: 

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

116 if groupings is not None: 

117 for grouping in groupings: 

118 for i, group in enumerate(grouping): 

119 if a in group: 

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

121 break 

122 lines.append(line) 

123 

124 # Write file 

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

126 

127 

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

129 

130 """ 

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

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

133 

134 Parameters 

135 ---------- 

136 fd : file | str 

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

138 species : List[str] 

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

140 precedence over isotope_masses 

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

142 Dictionary with chemical symbols and lists of the associated atomic 

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

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

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

146 

147 Returns 

148 ------- 

149 atoms : Atoms 

150 Atoms object 

151 input_parameters : Dict[str, int] 

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

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

154 species : List[str] 

155 List with the chemical symbols that correspond to each type 

156 

157 Raises 

158 ------ 

159 ValueError 

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

161 """ 

162 # Parse first line 

163 first_line = next(fd) 

164 print(first_line) 

165 input_parameters = {} 

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

167 'num_of_groups'] 

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

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

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

171 

172 # Parse second line 

173 second_line = next(fd) 

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

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

176 if input_parameters['triclinic']: 

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

178 else: 

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

180 

181 # Parse all remaining rows 

182 n_rows = input_parameters['N'] 

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

184 input_parameters['num_of_groups'] 

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

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

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

188 

189 # Extract atom types, positions and masses 

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

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

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

193 

194 # Determine the atomic species 

195 if species is None: 

196 type_symbol_map = {} 

197 if isotope_masses is not None: 

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

199 isotope_masses.items() for mass in masses} 

200 symbols = [] 

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

202 if species is None: 

203 if atom_type not in type_symbol_map: 

204 if isotope_masses is not None: 

205 nearest_value = find_nearest_value( 

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

207 symbol = mass_symbols[nearest_value] 

208 else: 

209 symbol = chemical_symbols[ 

210 find_nearest_index(atomic_masses, mass)] 

211 type_symbol_map[atom_type] = symbol 

212 else: 

213 symbol = type_symbol_map[atom_type] 

214 else: 

215 if atom_type > len(species): 

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

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

218 symbol = species[atom_type] 

219 symbols.append(symbol) 

220 

221 if species is None: 

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

223 

224 # Create the Atoms object 

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

226 cell=cell) 

227 if input_parameters['has_velocity']: 

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

229 atoms.set_velocities(velocities) 

230 if input_parameters['num_of_groups']: 

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

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

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

234 

235 return atoms, input_parameters, species 

236 

237 

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

239 """ 

240 Read Atoms object from a GPUMD structure input file 

241 

242 Parameters 

243 ---------- 

244 fd : file | str 

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

246 species : List[str] 

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

248 precedence over isotope_masses 

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

250 Dictionary with chemical symbols and lists of the associated atomic 

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

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

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

254 

255 Returns 

256 ------- 

257 atoms : Atoms 

258 Atoms object 

259 

260 Raises 

261 ------ 

262 ValueError 

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

264 """ 

265 

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