Coverage for /builds/kinetik161/ase/ase/io/gpumd.py: 90.68%
118 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1import numpy as np
3from ase import Atoms
4from ase.data import atomic_masses, chemical_symbols
5from ase.neighborlist import NeighborList
8def find_nearest_index(array, value):
9 array = np.asarray(array)
10 idx = (np.abs(array - value)).argmin()
11 return idx
14def find_nearest_value(array, value):
15 array = np.asarray(array)
16 idx = (np.abs(array - value)).argmin()
17 return array[idx]
20def write_gpumd(fd, atoms, maximum_neighbors=None, cutoff=None,
21 groupings=None, use_triclinic=False, species=None):
22 """
23 Writes atoms into GPUMD input format.
25 Parameters
26 ----------
27 fd : file
28 File like object to which the atoms object should be written
29 atoms : Atoms
30 Input structure
31 maximum_neighbors: int
32 Maximum number of neighbors any atom can ever have (not relevant when
33 using force constant potentials)
34 cutoff: float
35 Initial cutoff distance used for building the neighbor list (not
36 relevant when using force constant potentials)
37 groupings : list[list[list[int]]]
38 Groups into which the individual atoms should be divided in the form of
39 a list of list of lists. Specifically, the outer list corresponds to
40 the grouping methods, of which there can be three at the most, which
41 contains a list of groups in the form of lists of site indices. The
42 sum of the lengths of the latter must be the same as the total number
43 of atoms.
44 use_triclinic: bool
45 Use format for triclinic cells
46 species : List[str]
47 GPUMD uses integers to define atom types. This list allows customized
48 such definitions (e.g, ['Pd', 'H'] means Pd is type 0 and H type 1).
49 If None, this list is built by assigning each distinct species to
50 an integer in the order of appearance in `atoms`.
52 Raises
53 ------
54 ValueError
55 Raised if parameters are incompatible
56 """
58 # Check velocties parameter
59 if atoms.get_velocities() is None:
60 has_velocity = 0
61 else:
62 has_velocity = 1
63 velocities = atoms.get_velocities()
65 # Check groupings parameter
66 if groupings is None:
67 number_of_grouping_methods = 0
68 else:
69 number_of_grouping_methods = len(groupings)
70 if number_of_grouping_methods > 3:
71 raise ValueError('There can be no more than 3 grouping methods!')
72 for g, grouping in enumerate(groupings):
73 all_indices = [i for group in grouping for i in group]
74 if len(all_indices) != len(atoms) or\
75 set(all_indices) != set(range(len(atoms))):
76 raise ValueError('The indices listed in grouping method {} are'
77 ' not compatible with the input'
78 ' structure!'.format(g))
80 # If not specified, estimate the maximum_neighbors
81 if maximum_neighbors is None:
82 if cutoff is None:
83 cutoff = 0.1
84 maximum_neighbors = 1
85 else:
86 nl = NeighborList([cutoff / 2] * len(atoms), skin=2, bothways=True)
87 nl.update(atoms)
88 maximum_neighbors = 0
89 for atom in atoms:
90 maximum_neighbors = max(maximum_neighbors,
91 len(nl.get_neighbors(atom.index)[0]))
92 maximum_neighbors *= 2
93 maximum_neighbors = min(maximum_neighbors, 1024)
95 # Add header and cell parameters
96 lines = []
97 if atoms.cell.orthorhombic and not use_triclinic:
98 triclinic = 0
99 else:
100 triclinic = 1
101 lines.append('{} {} {} {} {} {}'.format(len(atoms), maximum_neighbors,
102 cutoff, triclinic, has_velocity,
103 number_of_grouping_methods))
104 if triclinic:
105 lines.append((' {}' * 12)[1:].format(*atoms.pbc.astype(int),
106 *atoms.cell[:].flatten()))
107 else:
108 lines.append((' {}' * 6)[1:].format(*atoms.pbc.astype(int),
109 *atoms.cell.lengths()))
111 # Create symbols-to-type map, i.e. integers starting at 0
112 if not species:
113 symbol_type_map = {}
114 for symbol in atoms.get_chemical_symbols():
115 if symbol not in symbol_type_map:
116 symbol_type_map[symbol] = len(symbol_type_map)
117 else:
118 if any(sym not in species
119 for sym in set(atoms.get_chemical_symbols())):
120 raise ValueError('The species list does not contain all chemical '
121 'species that are present in the atoms object.')
122 else:
123 symbol_type_map = {symbol: i for i, symbol in enumerate(species)}
125 # Add lines for all atoms
126 for a, atm in enumerate(atoms):
127 t = symbol_type_map[atm.symbol]
128 line = (' {}' * 5)[1:].format(t, *atm.position, atm.mass)
129 if has_velocity:
130 line += (' {}' * 3).format(*velocities[a])
131 if groupings is not None:
132 for grouping in groupings:
133 for i, group in enumerate(grouping):
134 if a in group:
135 line += f' {i}'
136 break
137 lines.append(line)
139 # Write file
140 fd.write('\n'.join(lines))
143def load_xyz_input_gpumd(fd, species=None, isotope_masses=None):
144 """
145 Read the structure input file for GPUMD and return an ase Atoms object
146 togehter with a dictionary with parameters and a types-to-symbols map
148 Parameters
149 ----------
150 fd : file | str
151 File object or name of file from which to read the Atoms object
152 species : List[str]
153 List with the chemical symbols that correspond to each type, will take
154 precedence over isotope_masses
155 isotope_masses: Dict[str, List[float]]
156 Dictionary with chemical symbols and lists of the associated atomic
157 masses, which is used to identify the chemical symbols that correspond
158 to the types not found in species_types. The default is to find the
159 closest match :data:`ase.data.atomic_masses`.
161 Returns
162 -------
163 atoms : Atoms
164 Atoms object
165 input_parameters : Dict[str, int]
166 Dictionary with parameters from the first row of the input file, namely
167 'N', 'M', 'cutoff', 'triclinic', 'has_velocity' and 'num_of_groups'
168 species : List[str]
169 List with the chemical symbols that correspond to each type
171 Raises
172 ------
173 ValueError
174 Raised if the list of species is incompatible with the input file
175 """
176 # Parse first line
177 first_line = next(fd)
178 input_parameters = {}
179 keys = ['N', 'M', 'cutoff', 'triclinic', 'has_velocity',
180 'num_of_groups']
181 types = [float if key == 'cutoff' else int for key in keys]
182 for k, (key, typ) in enumerate(zip(keys, types)):
183 input_parameters[key] = typ(first_line.split()[k])
185 # Parse second line
186 second_line = next(fd)
187 second_arr = np.array(second_line.split())
188 pbc = second_arr[:3].astype(bool)
189 if input_parameters['triclinic']:
190 cell = second_arr[3:].astype(float).reshape((3, 3))
191 else:
192 cell = np.diag(second_arr[3:].astype(float))
194 # Parse all remaining rows
195 n_rows = input_parameters['N']
196 n_columns = 5 + input_parameters['has_velocity'] * 3 +\
197 input_parameters['num_of_groups']
198 rest_lines = [next(fd) for _ in range(n_rows)]
199 rest_arr = np.array([line.split() for line in rest_lines])
200 assert rest_arr.shape == (n_rows, n_columns)
202 # Extract atom types, positions and masses
203 atom_types = rest_arr[:, 0].astype(int)
204 positions = rest_arr[:, 1:4].astype(float)
205 masses = rest_arr[:, 4].astype(float)
207 # Determine the atomic species
208 if species is None:
209 type_symbol_map = {}
210 if isotope_masses is not None:
211 mass_symbols = {mass: symbol for symbol, masses in
212 isotope_masses.items() for mass in masses}
213 symbols = []
214 for atom_type, mass in zip(atom_types, masses):
215 if species is None:
216 if atom_type not in type_symbol_map:
217 if isotope_masses is not None:
218 nearest_value = find_nearest_value(
219 list(mass_symbols.keys()), mass)
220 symbol = mass_symbols[nearest_value]
221 else:
222 symbol = chemical_symbols[
223 find_nearest_index(atomic_masses, mass)]
224 type_symbol_map[atom_type] = symbol
225 else:
226 symbol = type_symbol_map[atom_type]
227 else:
228 if atom_type > len(species):
229 raise Exception('There is no entry for atom type {} in the '
230 'species list!'.format(atom_type))
231 symbol = species[atom_type]
232 symbols.append(symbol)
234 if species is None:
235 species = [type_symbol_map[i] for i in sorted(type_symbol_map.keys())]
237 # Create the Atoms object
238 atoms = Atoms(symbols=symbols, positions=positions, masses=masses, pbc=pbc,
239 cell=cell)
240 if input_parameters['has_velocity']:
241 velocities = rest_arr[:, 5:8].astype(float)
242 atoms.set_velocities(velocities)
243 if input_parameters['num_of_groups']:
244 start_col = 5 + 3 * input_parameters['has_velocity']
245 groups = rest_arr[:, start_col:].astype(int)
246 atoms.info = {i: {'groups': groups[i, :]} for i in range(n_rows)}
248 return atoms, input_parameters, species
251def read_gpumd(fd, species=None, isotope_masses=None):
252 """
253 Read Atoms object from a GPUMD structure input file
255 Parameters
256 ----------
257 fd : file | str
258 File object or name of file from which to read the Atoms object
259 species : List[str]
260 List with the chemical symbols that correspond to each type, will take
261 precedence over isotope_masses
262 isotope_masses: Dict[str, List[float]]
263 Dictionary with chemical symbols and lists of the associated atomic
264 masses, which is used to identify the chemical symbols that correspond
265 to the types not found in species_types. The default is to find the
266 closest match :data:`ase.data.atomic_masses`.
268 Returns
269 -------
270 atoms : Atoms
271 Atoms object
273 Raises
274 ------
275 ValueError
276 Raised if the list of species is incompatible with the input file
277 """
279 return load_xyz_input_gpumd(fd, species, isotope_masses)[0]