Coverage for calorine/nep/io.py: 100%
281 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-14 16:36 +0000
1from os.path import exists
2from os.path import join as join_path
3from typing import Any, Iterable, TextIO
4from warnings import warn
6import numpy as np
7from ase import Atoms
8from ase.io import read, write
9from ase.stress import voigt_6_to_full_3x3_stress
10from pandas import DataFrame
12from calorine.nep.tensor_conventions import ASE_VOIGT6_ORDER, BEC_FULL9_ORDER, \
13 nep_reduced6_to_ase_voigt6
16def read_loss(filename: str) -> DataFrame:
17 """Parses a file in `loss.out` format from GPUMD and returns the
18 content as a data frame. More information concerning file format,
19 content and units can be found `here
20 <https://gpumd.org/nep/output_files/loss_out.html>`__.
22 Parameters
23 ----------
24 filename
25 input file name
27 """
28 data = np.loadtxt(filename)
29 if isinstance(data[0], np.float64):
30 # If only a single row in loss.out, append a dimension
31 data = data.reshape(1, -1)
32 if len(data[0]) == 6:
33 tags = 'total_loss L1 L2'
34 tags += ' RMSE_P_train'
35 tags += ' RMSE_P_test'
36 elif len(data[0]) == 10:
37 tags = 'total_loss L1 L2'
38 tags += ' RMSE_E_train RMSE_F_train RMSE_V_train'
39 tags += ' RMSE_E_test RMSE_F_test RMSE_V_test'
40 elif len(data[0]) == 14:
41 tags = 'total_loss L1 L2'
42 tags += ' RMSE_E_train RMSE_F_train RMSE_V_train RMSE_Q_train RMSE_Z_train'
43 tags += ' RMSE_E_test RMSE_F_test RMSE_V_test RMSE_Q_test RMSE_Z_test'
44 else:
45 raise ValueError(
46 f'Input file contains {len(data[0])} data columns. Expected 6 or 10 columns.'
47 )
48 generations = range(100, len(data) * 100 + 1, 100)
49 df = DataFrame(data=data[:, 1:], columns=tags.split(), index=generations)
50 return df
53def _write_structure_in_nep_format(structure: Atoms, f: TextIO) -> None:
54 """Write structure block into a file-like object in format readable by nep executable.
56 Parameters
57 ----------
58 structure
59 input structure; must hold information regarding energy and forces
60 f
61 file-like object to which to write
62 """
64 # Allowed keyword=value pairs. Use ASEs extyz write functionality.:
65 # lattice="ax ay az bx by bz cx cy cz" (mandatory)
66 # energy=energy_value (mandatory)
67 # virial="vxx vxy vxz vyx vyy vyz vzx vzy vzz" (optional)
68 # weight=relative_weight (optional)
69 # properties=property_name:data_type:number_of_columns
70 # species:S:1 (mandatory)
71 # pos:R:3 (mandatory)
72 # force:R:3 or forces:R:3 (mandatory)
74 # If a structure is to be used for training, it needs to either have target
75 # 1. energies and forces,
76 # 2. dipole, denoted `dipole="dx dy dz"` in the info string, or
77 # 3. polarizability/susceptibility, denoted `pol="pxx pxy pxz pyx pyy pyz pzx pzy pzz"`
78 # in the info string.
79 has_energies_and_forces = True
80 try:
81 structure.get_potential_energy()
82 structure.get_forces() # calculate forces to have them on the Atoms object
83 except RuntimeError:
84 has_energies_and_forces = False
86 has_dipole = 'dipole' in structure.info.keys()
87 has_pol = 'pol' in structure.info.keys()
89 if not has_energies_and_forces and not has_dipole and not has_pol:
90 raise RuntimeError('Failed to retrieve target energies/forces,'
91 ' dipoles, or polarizabilities for structure')
92 if np.isclose(structure.get_volume(), 0):
93 raise ValueError('Structure cell must have a non-zero volume!')
94 try:
95 structure.get_stress()
96 except RuntimeError:
97 warn('Failed to retrieve stresses for structure')
98 write(filename=f, images=structure, write_info=True, format='extxyz')
101def write_structures(outfile: str, structures: list[Atoms]) -> None:
102 """Writes structures for training/testing in format readable by nep executable.
104 Parameters
105 ----------
106 outfile
107 output filename
108 structures
109 list of structures with energy, forces, and (possibly) stresses
110 """
111 with open(outfile, 'w') as f:
112 for structure in structures:
113 _write_structure_in_nep_format(structure, f)
116def write_nepfile(filename: str, parameters: dict[str, Any]) -> None:
117 """Writes a `nep.in` configuration to the given file name.
119 Keys are written in the order in which they appear in ``parameters``, which matters
120 because the `nep` executable rejects a ``cutoff`` line that precedes the ``type`` line.
122 Parameters
123 ----------
124 filename
125 Name of the file to write.
126 The parent directory must exist.
127 parameters
128 input parameters; see `here <https://gpumd.org/nep/input_parameters/index.html>`__
129 """
130 with open(filename, 'w') as f:
131 for key, val in parameters.items():
132 f.write(f'{key} ')
133 if isinstance(val, Iterable) and not isinstance(val, str):
134 f.write(' '.join([f'{v}' for v in val]))
135 else:
136 f.write(f'{val}')
137 f.write('\n')
140def read_nepfile(filename: str) -> dict[str, Any]:
141 """Returns the content of a configuration file (`nep.in`) as a dictionary.
143 Parameters
144 ----------
145 filename
146 input file name
147 """
148 int_vals = ['version', 'neuron', 'generation', 'batch', 'population',
149 'mode', 'model_type']
150 float_vals = ['lambda_1', 'lambda_2', 'lambda_e', 'lambda_f', 'lambda_v',
151 'lambda_q', 'lambda_shear', 'force_delta', 'atomic_v', 'zbl',
152 'use_typewise_cutoff_zbl']
153 settings = {}
154 with open(filename) as f:
155 for line in f.readlines():
156 # remove comments - throw away everything after a '#'
157 cleaned = line.split('#', 1)[0].strip()
158 flds = cleaned.split()
159 if len(flds) == 0:
160 continue
161 settings[flds[0]] = ' '.join(flds[1:])
162 for key, val in settings.items():
163 if val == '':
164 # a keyword that carries no value, such as a bare use_typewise_cutoff_zbl
165 continue
166 if key in int_vals:
167 settings[key] = int(val)
168 elif key in float_vals:
169 settings[key] = float(val)
170 elif key in ['n_max', 'l_max', 'basis_size', 'charge_mode']:
171 settings[key] = [int(v) for v in val.split()]
172 elif key in ['cutoff', 'type_weight']:
173 settings[key] = [float(v) for v in val.split()]
174 elif key == 'type':
175 types = val.split()
176 types[0] = int(types[0])
177 settings[key] = types
178 return settings
181def read_structures(dirname: str) -> tuple[list[Atoms], list[Atoms]]:
182 """Parses the output files with training and test data from a nep run and returns their
183 content as two lists of structures, representing training and test data, respectively.
184 Target and predicted data are included in the :attr:`info` dict of the :class:`Atoms`
185 objects.
187 Parameters
188 ----------
189 dirname
190 Directory from which to read output files.
192 """
193 path = join_path(dirname)
194 if not exists(path):
195 raise FileNotFoundError(f'Directory {path} does not exist')
197 # fetch model type from nep input file
198 nep_info = read_nepfile(f'{path}/nep.in')
199 model_type = nep_info.get('model_type', 0)
201 # set up which files to parse, what dimensions to expect etc
202 # depending on the type of model that is parsed
203 #
204 # The `nep` executable's virial_*.out/stress_*.out/polarizability_*.out
205 # output files use a reduced-6 component order (xx,yy,zz,xy,yz,xz) that
206 # differs from ASE's own Voigt-6 convention (xx,yy,zz,yz,xz,xy) -- see
207 # GPUMD's main_nep/structure.cu `reduced_index` array (shared by the
208 # `virial=` and `pol=` extxyz parsers). Below, virial/stress/
209 # polarizability columns are converted to ASE order immediately upon
210 # parsing, so `structure.info`/`.arrays` always hold ASE-Voigt-ordered
211 # data, matching what `atoms.get_stress()` returns elsewhere in
212 # calorine. BEC's 9 columns are left as is (row-major, see
213 # BEC_FULL9_ORDER) -- BEC is not symmetric, so no Voigt form applies.
214 if model_type == 0:
215 charge_mode = nep_info.get('charge_mode', [0])[0]
216 if charge_mode not in [0, 1, 2]:
217 raise ValueError(f'Unknown charge_mode: {charge_mode}')
218 # files to parse: (sname, size, mandatory, includes_target, per_atom)
219 files_to_parse = [
220 ('energy', 1, True, True, False),
221 ('force', 3, True, True, True),
222 ('virial', 6, True, True, False),
223 ('stress', 6, True, True, False),
224 ]
225 if charge_mode in [1, 2]:
226 # files to parse: (sname, size, includes_target, per_atom)
227 files_to_parse += [
228 ('charge', 1, True, False, True),
229 ('bec', 9, False, True, True),
230 ]
231 elif model_type == 1:
232 # files to parse: (sname, size, includes_target, per_atom)
233 files_to_parse = [('dipole', 3, True, True, False)]
234 if nep_info.get('atomic_v', 0) == 1:
235 files_to_parse = [('dipole', 3, True, True, True)]
236 elif model_type == 2:
237 # files to parse: (sname, size, includes_target, per_atom)
238 files_to_parse = [('polarizability', 6, True, True, False)]
239 if nep_info.get('atomic_v', 0) == 1:
240 files_to_parse = [('polarizability', 6, True, True, True)]
241 else:
242 raise ValueError(f'Unknown model_type: {model_type}')
244 # read training and test data
245 structures = {}
246 for stype in ['train', 'test']:
247 filename = join_path(dirname, f'{stype}.xyz')
248 try:
249 structures[stype] = read(filename, format='extxyz', index=':')
250 except FileNotFoundError:
251 warn(f'File {filename} not found.')
252 structures[stype] = []
253 continue
255 n_structures = len(structures[stype])
257 # loop over files from which to read target data and predictions
258 for sname, size, mandatory, includes_target, per_atom in files_to_parse:
259 infile = f'{sname}_{stype}.out'
260 path = join_path(dirname, infile)
261 if not exists(path):
262 if mandatory:
263 raise FileNotFoundError(f'File {path} does not exist')
264 else:
265 continue
266 ts, ps = _read_data_file(path, includes_target=includes_target)
268 if ts is not None:
269 if ts.shape[1] != size:
270 raise ValueError(f'Target data in {infile} has unexpected shape:'
271 f' {ts.shape} (expected: (-1, {size}))')
272 if ps.shape[1] != size:
273 raise ValueError(f'Predicted data in {infile} has unexpected shape:'
274 f' {ps.shape} (expected: (-1, {size}))')
276 if sname in ('virial', 'stress', 'polarizability'):
277 ts = nep_reduced6_to_ase_voigt6(ts)
278 ps = nep_reduced6_to_ase_voigt6(ps)
280 if per_atom:
281 # data per-atom, e.g., forces, per-atom-virials, Born effective charges ...
282 n_atoms_total = sum([len(s) for s in structures[stype]])
283 if len(ps) != n_atoms_total:
284 raise ValueError(f'Number of atoms in {infile} ({len(ps)})'
285 f' and {stype}.xyz ({n_atoms_total}) inconsistent.')
286 n = 0
287 for structure in structures[stype]:
288 nat = len(structure)
289 if ts is not None:
290 t = np.array(ts[n: n + nat]).reshape(nat, size)
291 structure.new_array(f'{sname}_target', t)
292 p = np.array(ps[n: n + nat]).reshape(nat, size)
293 structure.new_array(f'{sname}_predicted', p)
294 n += nat
295 else:
296 # data per structure, e.g., energy, virials, stress
297 if len(ps) != n_structures:
298 raise ValueError(f'Number of structures in {infile} ({len(ps)})'
299 f' and {stype}.xyz ({n_structures}) inconsistent.')
300 for k, structure in enumerate(structures[stype]):
301 assert ts is not None, 'This should not occur. Please report.'
302 t = ts[k]
303 assert np.shape(t) == (size,)
304 structure.info[f'{sname}_target'] = t
305 p = ps[k]
306 assert np.shape(p) == (size,)
307 structure.info[f'{sname}_predicted'] = p
309 # special handling of target data for BECs
310 # If a structure has no 'bec' array in the xyz file, no target BEC data was provided.
311 # In that case nep writes zeros for both predicted and target columns. Replace both
312 # with NaN so callers can easily identify and filter out structures without BEC targets.
313 for s in structures[stype]:
314 if 'bec_target' in s.arrays and 'bec' not in s.arrays:
315 nat = len(s)
316 s.arrays['bec_target'] = np.full((nat, 9), np.nan)
317 s.arrays['bec_predicted'] = np.full((nat, 9), np.nan)
319 # special handling of per-atom TNEP
320 # Data has to be loaded as dipole/polarizability since NEP saves them in dipole_*.out
321 # Dipole/polarizability arrays are therefore moved to atomic_v here
322 if nep_info.get('atomic_v', 0) == 1:
323 if model_type == 1:
324 s.new_array('atomic_v_target', s.arrays['dipole_target'])
325 s.new_array('atomic_v_predicted', s.arrays['dipole_predicted'])
326 del s.arrays['dipole_target']
327 del s.arrays['dipole_predicted']
328 if model_type == 2:
329 s.new_array('atomic_v_target', s.arrays['polarizability_target'])
330 s.new_array('atomic_v_predicted', s.arrays['polarizability_predicted'])
331 del s.arrays['polarizability_target']
332 del s.arrays['polarizability_predicted']
334 return structures['train'], structures['test']
337def _read_data_file(
338 path: str,
339 includes_target: bool = True,
340):
341 """Private function that parses *.out files and
342 returns their content for further processing.
343 """
344 with open(path, 'r') as f:
345 lines = f.readlines()
346 target, predicted = [], []
347 for line in lines:
348 flds = line.split()
349 if includes_target:
350 if len(flds) % 2 != 0:
351 raise ValueError(f'Incorrect number of columns in {path} ({len(flds)}).')
352 n = len(flds) // 2
353 predicted.append([float(s) for s in flds[:n]])
354 target.append([float(s) for s in flds[n:]])
355 else:
356 predicted.append([float(s) for s in flds])
357 target = None
358 if target is not None:
359 target = np.array(target)
360 predicted = np.array(predicted)
361 return target, predicted
364# Maps x/y/z (vectors) and reduced-6 symmetric-tensor component names to
365# their index. virial/stress/polarizability (and per-atom atomic_v when it
366# holds a reduced-6 polarizability) are converted to ASE-Voigt order by
367# read_structures() above, so this mapping uses ASE_VOIGT6_ORDER directly.
368_REDUCED6_MAPPING = {
369 'x': 0, 'y': 1, 'z': 2,
370 **{label: i for i, label in enumerate(ASE_VOIGT6_ORDER)},
371}
372# BEC is a raw, unreduced, generally-asymmetric 9-component per-atom tensor
373# (row-major, see BEC_FULL9_ORDER) -- not a Voigt-reduced 6-component form,
374# so it needs its own, different index scheme.
375_BEC_MAPPING = {label: i for i, label in enumerate(BEC_FULL9_ORDER)}
378def get_parity_data(
379 structures: list[Atoms],
380 property: str,
381 selection: list[str] = None,
382 flatten: bool = True,
383) -> DataFrame:
384 """Returns the predicted and target energies, forces, virials or stresses
385 from a list of structures in a format suitable for generating parity plots.
387 The structures should have been read using :func:`read_structures
388 <calorine.nep.read_structures>`, such that the `info` object is
389 populated with keys of the form `<property>_<type>` where `<property>`
390 is, e.g., `energy` or `force` and `<type>` is one of `predicted` or `target`.
392 The resulting parity data is returned as a tuple of dicts, where each entry
393 corresponds to a list.
395 Parameters
396 ----------
397 structures
398 List of structures as read with :func:`read_structures <calorine.nep.read_structures>`.
399 property
400 One of `energy`, `force`, `virial`, `stress`, `bec`, `dipole`,
401 `polarizability`, or `atomic_v`.
402 selection
403 A list containing which components to return, and/or the norm.
404 For `force`, `atomic_v` (dipole), `virial`, `stress`, `polarizability`,
405 and `dipole`, possible values are `x`, `y`, `z`, `xx`, `yy`, `zz`,
406 `yz`, `xz`, `xy`, `norm`, `pressure` (the latter only for `stress`).
407 For `bec`, all nine of `xx`, `xy`, `xz`, `yx`, `yy`, `yz`, `zx`,
408 `zy`, `zz` are available (BEC is not symmetric, so unlike the other
409 properties it has no reduced/Voigt form).
410 flatten
411 if True return flattened lists; this is useful for flattening
412 the components of force or virials into a simple list
413 """
414 global_properties = ['energy', 'virial', 'stress', 'polarizability', 'dipole']
415 per_atom_properties = ['force', 'bec', 'atomic_v']
416 if property not in global_properties + per_atom_properties:
417 raise ValueError(
418 '`property` must be one of the following:'
419 + ', '.join(global_properties + per_atom_properties)
420 )
421 if property in ['energy'] and selection:
422 raise ValueError('Selection cannot be applied to scalars.')
423 if property != 'stress' and selection and 'pressure' in selection:
424 raise ValueError(f'Cannot calculate pressure for `{property}`.')
426 data = {'predicted': [], 'target': []}
427 if property in ['force', 'bec'] and flatten:
428 size = 3 if property == 'force' else 9
429 data['species'] = []
430 for structure in structures:
431 if 'species' in data:
432 data['species'].extend(np.repeat(structure.symbols, size).tolist())
433 for stype in ['predicted', 'target']:
434 property_with_stype = f'{property}_{stype}'
435 if property in global_properties:
436 if property_with_stype not in structure.info.keys():
437 raise KeyError(f'{property_with_stype} not'
438 ' available in info field of structure')
439 extracted_property = np.array(structure.info[property_with_stype])
440 else:
441 if property_with_stype not in structure.arrays:
442 raise KeyError(f'{property_with_stype} not available in arrays of structure')
443 extracted_property = np.array(structure.arrays[property_with_stype])
445 if selection is None or len(selection) == 0:
446 data[stype].append(extracted_property)
447 continue
449 if property in ['force', 'bec', 'atomic_v']:
450 extracted_property = extracted_property.T
451 selected_values = []
452 for select in selection:
453 if select == 'norm':
454 if property == 'force':
455 selected_values.append(np.linalg.norm(extracted_property, axis=0))
456 elif property == 'atomic_v':
457 if extracted_property.shape[0] == 3:
458 # per-atom dipole-like vector
459 selected_values.append(
460 np.linalg.norm(extracted_property, axis=0))
461 elif extracted_property.shape[0] == 6:
462 # per-atom polarizability-like tensor in reduced Voigt notation
463 full_tensor = voigt_6_to_full_3x3_stress(extracted_property.T)
464 selected_values.append(
465 np.linalg.norm(full_tensor, axis=(1, 2)))
466 else:
467 raise ValueError(
468 'Cannot handle selection=`norm` with property=`atomic_v`'
469 f' of size {extracted_property.shape[0]}.'
470 )
471 elif property in ['virial', 'stress']:
472 full_tensor = voigt_6_to_full_3x3_stress(extracted_property)
473 selected_values.append(np.linalg.norm(full_tensor))
474 elif property in ['dipole']:
475 selected_values.append(np.linalg.norm(extracted_property))
476 else:
477 raise ValueError(
478 f'Cannot handle selection=`norm` with property=`{property}`.')
479 continue
481 if select == 'pressure' and property == 'stress':
482 total_stress = extracted_property
483 selected_values.append(-np.sum(total_stress[:3]) / 3)
484 continue
486 mapping = _BEC_MAPPING if property == 'bec' else _REDUCED6_MAPPING
487 if select not in mapping:
488 raise ValueError(f'Selection `{select}` is not allowed.')
489 index = mapping[select]
490 if index >= extracted_property.shape[0]:
491 raise ValueError(
492 f'Selection `{select}` is not compatible with property `{property}`.'
493 )
494 selected_values.append(extracted_property[index])
496 data[stype].append(selected_values)
497 if flatten:
498 for stype in ['target', 'predicted']:
499 value = data[stype]
500 if len(np.shape(value[0])) > 0:
501 data[stype] = np.concatenate(value).ravel().tolist()
502 if property in ['force', 'atomic_v']:
503 default_labels = ['x', 'y', 'z']
504 labels = selection if selection else default_labels
505 n = len(data['target']) // len(labels)
506 data['component'] = list(labels) * n
507 elif property in ['virial', 'stress', 'polarizability']:
508 default_labels = list(ASE_VOIGT6_ORDER)
509 labels = selection if selection else default_labels
510 n = len(data['target']) // len(labels)
511 data['component'] = list(labels) * n
512 elif property in ['bec']:
513 default_labels = list(BEC_FULL9_ORDER)
514 labels = selection if selection else default_labels
515 n = len(data['target']) // len(labels)
516 data['component'] = list(labels) * n
517 df = DataFrame(data)
518 # In case of flatten, cast to float64 for compatibility
519 # with e.g. seaborn.
520 # Casting in this way breaks tensorial properties though,
521 # so skip it there.
522 if flatten:
523 df['target'] = df.target.astype('float64')
524 df['predicted'] = df.predicted.astype('float64')
525 return df