Coverage for calorine/nep/io.py: 100%
283 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 12:52 +0000
1from os.path import exists
2from os.path import join as join_path
3from typing import Any, Iterable, NamedTuple, 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_to_path(parameters: dict[str, Any], filename: str) -> None:
117 """Writes a `nep.in` configuration to the given file name.
119 Keys are written in the order in which they appear in :attr:`parameters`, which matters
120 because the `nep` executable rejects a ``cutoff`` line that precedes the ``type`` line.
122 Parameters
123 ----------
124 parameters
125 input parameters; see `here <https://gpumd.org/nep/input_parameters/index.html>`__
126 filename
127 name of the file to write
128 """
129 with open(filename, 'w') as f:
130 for key, val in parameters.items():
131 f.write(f'{key} ')
132 if isinstance(val, Iterable) and not isinstance(val, str):
133 f.write(' '.join([f'{v}' for v in val]))
134 else:
135 f.write(f'{val}')
136 f.write('\n')
139def write_nepfile(parameters: NamedTuple, dirname: str) -> None:
140 """Writes parameters file for NEP construction.
142 Parameters
143 ----------
144 parameters
145 input parameters; see `here <https://gpumd.org/nep/input_parameters/index.html>`__
146 dirname
147 directory in which to place input file and links
148 """
149 _write_nepfile_to_path(parameters, join_path(dirname, 'nep.in'))
152def read_nepfile(filename: str) -> dict[str, Any]:
153 """Returns the content of a configuration file (`nep.in`) as a dictionary.
155 Parameters
156 ----------
157 filename
158 input file name
159 """
160 int_vals = ['version', 'neuron', 'generation', 'batch', 'population',
161 'mode', 'model_type']
162 float_vals = ['lambda_1', 'lambda_2', 'lambda_e', 'lambda_f', 'lambda_v',
163 'lambda_q', 'lambda_shear', 'force_delta', 'atomic_v', 'zbl',
164 'use_typewise_cutoff_zbl']
165 settings = {}
166 with open(filename) as f:
167 for line in f.readlines():
168 # remove comments - throw away everything after a '#'
169 cleaned = line.split('#', 1)[0].strip()
170 flds = cleaned.split()
171 if len(flds) == 0:
172 continue
173 settings[flds[0]] = ' '.join(flds[1:])
174 for key, val in settings.items():
175 if val == '':
176 # a keyword that carries no value, such as a bare use_typewise_cutoff_zbl
177 continue
178 if key in int_vals:
179 settings[key] = int(val)
180 elif key in float_vals:
181 settings[key] = float(val)
182 elif key in ['n_max', 'l_max', 'basis_size', 'charge_mode']:
183 settings[key] = [int(v) for v in val.split()]
184 elif key in ['cutoff', 'type_weight']:
185 settings[key] = [float(v) for v in val.split()]
186 elif key == 'type':
187 types = val.split()
188 types[0] = int(types[0])
189 settings[key] = types
190 return settings
193def read_structures(dirname: str) -> tuple[list[Atoms], list[Atoms]]:
194 """Parses the output files with training and test data from a nep run and returns their
195 content as two lists of structures, representing training and test data, respectively.
196 Target and predicted data are included in the :attr:`info` dict of the :class:`Atoms`
197 objects.
199 Parameters
200 ----------
201 dirname
202 Directory from which to read output files.
204 """
205 path = join_path(dirname)
206 if not exists(path):
207 raise FileNotFoundError(f'Directory {path} does not exist')
209 # fetch model type from nep input file
210 nep_info = read_nepfile(f'{path}/nep.in')
211 model_type = nep_info.get('model_type', 0)
213 # set up which files to parse, what dimensions to expect etc
214 # depending on the type of model that is parsed
215 #
216 # The `nep` executable's virial_*.out/stress_*.out/polarizability_*.out
217 # output files use a reduced-6 component order (xx,yy,zz,xy,yz,xz) that
218 # differs from ASE's own Voigt-6 convention (xx,yy,zz,yz,xz,xy) -- see
219 # GPUMD's main_nep/structure.cu `reduced_index` array (shared by the
220 # `virial=` and `pol=` extxyz parsers). Below, virial/stress/
221 # polarizability columns are converted to ASE order immediately upon
222 # parsing, so `structure.info`/`.arrays` always hold ASE-Voigt-ordered
223 # data, matching what `atoms.get_stress()` returns elsewhere in
224 # calorine. BEC's 9 columns are left as is (row-major, see
225 # BEC_FULL9_ORDER) -- BEC is not symmetric, so no Voigt form applies.
226 if model_type == 0:
227 charge_mode = nep_info.get('charge_mode', [0])[0]
228 if charge_mode not in [0, 1, 2]:
229 raise ValueError(f'Unknown charge_mode: {charge_mode}')
230 # files to parse: (sname, size, mandatory, includes_target, per_atom)
231 files_to_parse = [
232 ('energy', 1, True, True, False),
233 ('force', 3, True, True, True),
234 ('virial', 6, True, True, False),
235 ('stress', 6, True, True, False),
236 ]
237 if charge_mode in [1, 2]:
238 # files to parse: (sname, size, includes_target, per_atom)
239 files_to_parse += [
240 ('charge', 1, True, False, True),
241 ('bec', 9, False, True, True),
242 ]
243 elif model_type == 1:
244 # files to parse: (sname, size, includes_target, per_atom)
245 files_to_parse = [('dipole', 3, True, True, False)]
246 if nep_info.get('atomic_v', 0) == 1:
247 files_to_parse = [('dipole', 3, True, True, True)]
248 elif model_type == 2:
249 # files to parse: (sname, size, includes_target, per_atom)
250 files_to_parse = [('polarizability', 6, True, True, False)]
251 if nep_info.get('atomic_v', 0) == 1:
252 files_to_parse = [('polarizability', 6, True, True, True)]
253 else:
254 raise ValueError(f'Unknown model_type: {model_type}')
256 # read training and test data
257 structures = {}
258 for stype in ['train', 'test']:
259 filename = join_path(dirname, f'{stype}.xyz')
260 try:
261 structures[stype] = read(filename, format='extxyz', index=':')
262 except FileNotFoundError:
263 warn(f'File {filename} not found.')
264 structures[stype] = []
265 continue
267 n_structures = len(structures[stype])
269 # loop over files from which to read target data and predictions
270 for sname, size, mandatory, includes_target, per_atom in files_to_parse:
271 infile = f'{sname}_{stype}.out'
272 path = join_path(dirname, infile)
273 if not exists(path):
274 if mandatory:
275 raise FileNotFoundError(f'File {path} does not exist')
276 else:
277 continue
278 ts, ps = _read_data_file(path, includes_target=includes_target)
280 if ts is not None:
281 if ts.shape[1] != size:
282 raise ValueError(f'Target data in {infile} has unexpected shape:'
283 f' {ts.shape} (expected: (-1, {size}))')
284 if ps.shape[1] != size:
285 raise ValueError(f'Predicted data in {infile} has unexpected shape:'
286 f' {ps.shape} (expected: (-1, {size}))')
288 if sname in ('virial', 'stress', 'polarizability'):
289 ts = nep_reduced6_to_ase_voigt6(ts)
290 ps = nep_reduced6_to_ase_voigt6(ps)
292 if per_atom:
293 # data per-atom, e.g., forces, per-atom-virials, Born effective charges ...
294 n_atoms_total = sum([len(s) for s in structures[stype]])
295 if len(ps) != n_atoms_total:
296 raise ValueError(f'Number of atoms in {infile} ({len(ps)})'
297 f' and {stype}.xyz ({n_atoms_total}) inconsistent.')
298 n = 0
299 for structure in structures[stype]:
300 nat = len(structure)
301 if ts is not None:
302 t = np.array(ts[n: n + nat]).reshape(nat, size)
303 structure.new_array(f'{sname}_target', t)
304 p = np.array(ps[n: n + nat]).reshape(nat, size)
305 structure.new_array(f'{sname}_predicted', p)
306 n += nat
307 else:
308 # data per structure, e.g., energy, virials, stress
309 if len(ps) != n_structures:
310 raise ValueError(f'Number of structures in {infile} ({len(ps)})'
311 f' and {stype}.xyz ({n_structures}) inconsistent.')
312 for k, structure in enumerate(structures[stype]):
313 assert ts is not None, 'This should not occur. Please report.'
314 t = ts[k]
315 assert np.shape(t) == (size,)
316 structure.info[f'{sname}_target'] = t
317 p = ps[k]
318 assert np.shape(p) == (size,)
319 structure.info[f'{sname}_predicted'] = p
321 # special handling of target data for BECs
322 # If a structure has no 'bec' array in the xyz file, no target BEC data was provided.
323 # In that case nep writes zeros for both predicted and target columns. Replace both
324 # with NaN so callers can easily identify and filter out structures without BEC targets.
325 for s in structures[stype]:
326 if 'bec_target' in s.arrays and 'bec' not in s.arrays:
327 nat = len(s)
328 s.arrays['bec_target'] = np.full((nat, 9), np.nan)
329 s.arrays['bec_predicted'] = np.full((nat, 9), np.nan)
331 # special handling of per-atom TNEP
332 # Data has to be loaded as dipole/polarizability since NEP saves them in dipole_*.out
333 # Dipole/polarizability arrays are therefore moved to atomic_v here
334 if nep_info.get('atomic_v', 0) == 1:
335 if model_type == 1:
336 s.new_array('atomic_v_target', s.arrays['dipole_target'])
337 s.new_array('atomic_v_predicted', s.arrays['dipole_predicted'])
338 del s.arrays['dipole_target']
339 del s.arrays['dipole_predicted']
340 if model_type == 2:
341 s.new_array('atomic_v_target', s.arrays['polarizability_target'])
342 s.new_array('atomic_v_predicted', s.arrays['polarizability_predicted'])
343 del s.arrays['polarizability_target']
344 del s.arrays['polarizability_predicted']
346 return structures['train'], structures['test']
349def _read_data_file(
350 path: str,
351 includes_target: bool = True,
352):
353 """Private function that parses *.out files and
354 returns their content for further processing.
355 """
356 with open(path, 'r') as f:
357 lines = f.readlines()
358 target, predicted = [], []
359 for line in lines:
360 flds = line.split()
361 if includes_target:
362 if len(flds) % 2 != 0:
363 raise ValueError(f'Incorrect number of columns in {path} ({len(flds)}).')
364 n = len(flds) // 2
365 predicted.append([float(s) for s in flds[:n]])
366 target.append([float(s) for s in flds[n:]])
367 else:
368 predicted.append([float(s) for s in flds])
369 target = None
370 if target is not None:
371 target = np.array(target)
372 predicted = np.array(predicted)
373 return target, predicted
376# Maps x/y/z (vectors) and reduced-6 symmetric-tensor component names to
377# their index. virial/stress/polarizability (and per-atom atomic_v when it
378# holds a reduced-6 polarizability) are converted to ASE-Voigt order by
379# read_structures() above, so this mapping uses ASE_VOIGT6_ORDER directly.
380_REDUCED6_MAPPING = {
381 'x': 0, 'y': 1, 'z': 2,
382 **{label: i for i, label in enumerate(ASE_VOIGT6_ORDER)},
383}
384# BEC is a raw, unreduced, generally-asymmetric 9-component per-atom tensor
385# (row-major, see BEC_FULL9_ORDER) -- not a Voigt-reduced 6-component form,
386# so it needs its own, different index scheme.
387_BEC_MAPPING = {label: i for i, label in enumerate(BEC_FULL9_ORDER)}
390def get_parity_data(
391 structures: list[Atoms],
392 property: str,
393 selection: list[str] = None,
394 flatten: bool = True,
395) -> DataFrame:
396 """Returns the predicted and target energies, forces, virials or stresses
397 from a list of structures in a format suitable for generating parity plots.
399 The structures should have been read using :func:`read_structures
400 <calorine.nep.read_structures>`, such that the `info` object is
401 populated with keys of the form `<property>_<type>` where `<property>`
402 is, e.g., `energy` or `force` and `<type>` is one of `predicted` or `target`.
404 The resulting parity data is returned as a tuple of dicts, where each entry
405 corresponds to a list.
407 Parameters
408 ----------
409 structures
410 List of structures as read with :func:`read_structures <calorine.nep.read_structures>`.
411 property
412 One of `energy`, `force`, `virial`, `stress`, `bec`, `dipole`,
413 `polarizability`, or `atomic_v`.
414 selection
415 A list containing which components to return, and/or the norm.
416 For `force`, `atomic_v` (dipole), `virial`, `stress`, `polarizability`,
417 and `dipole`, possible values are `x`, `y`, `z`, `xx`, `yy`, `zz`,
418 `yz`, `xz`, `xy`, `norm`, `pressure` (the latter only for `stress`).
419 For `bec`, all nine of `xx`, `xy`, `xz`, `yx`, `yy`, `yz`, `zx`,
420 `zy`, `zz` are available (BEC is not symmetric, so unlike the other
421 properties it has no reduced/Voigt form).
422 flatten
423 if True return flattened lists; this is useful for flattening
424 the components of force or virials into a simple list
425 """
426 global_properties = ['energy', 'virial', 'stress', 'polarizability', 'dipole']
427 per_atom_properties = ['force', 'bec', 'atomic_v']
428 if property not in global_properties + per_atom_properties:
429 raise ValueError(
430 '`property` must be one of the following:'
431 + ', '.join(global_properties + per_atom_properties)
432 )
433 if property in ['energy'] and selection:
434 raise ValueError('Selection cannot be applied to scalars.')
435 if property != 'stress' and selection and 'pressure' in selection:
436 raise ValueError(f'Cannot calculate pressure for `{property}`.')
438 data = {'predicted': [], 'target': []}
439 if property in ['force', 'bec'] and flatten:
440 size = 3 if property == 'force' else 9
441 data['species'] = []
442 for structure in structures:
443 if 'species' in data:
444 data['species'].extend(np.repeat(structure.symbols, size).tolist())
445 for stype in ['predicted', 'target']:
446 property_with_stype = f'{property}_{stype}'
447 if property in global_properties:
448 if property_with_stype not in structure.info.keys():
449 raise KeyError(f'{property_with_stype} not'
450 ' available in info field of structure')
451 extracted_property = np.array(structure.info[property_with_stype])
452 else:
453 if property_with_stype not in structure.arrays:
454 raise KeyError(f'{property_with_stype} not available in arrays of structure')
455 extracted_property = np.array(structure.arrays[property_with_stype])
457 if selection is None or len(selection) == 0:
458 data[stype].append(extracted_property)
459 continue
461 if property in ['force', 'bec', 'atomic_v']:
462 extracted_property = extracted_property.T
463 selected_values = []
464 for select in selection:
465 if select == 'norm':
466 if property == 'force':
467 selected_values.append(np.linalg.norm(extracted_property, axis=0))
468 elif property == 'atomic_v':
469 if extracted_property.shape[0] == 3:
470 # per-atom dipole-like vector
471 selected_values.append(
472 np.linalg.norm(extracted_property, axis=0))
473 elif extracted_property.shape[0] == 6:
474 # per-atom polarizability-like tensor in reduced Voigt notation
475 full_tensor = voigt_6_to_full_3x3_stress(extracted_property.T)
476 selected_values.append(
477 np.linalg.norm(full_tensor, axis=(1, 2)))
478 else:
479 raise ValueError(
480 'Cannot handle selection=`norm` with property=`atomic_v`'
481 f' of size {extracted_property.shape[0]}.'
482 )
483 elif property in ['virial', 'stress']:
484 full_tensor = voigt_6_to_full_3x3_stress(extracted_property)
485 selected_values.append(np.linalg.norm(full_tensor))
486 elif property in ['dipole']:
487 selected_values.append(np.linalg.norm(extracted_property))
488 else:
489 raise ValueError(
490 f'Cannot handle selection=`norm` with property=`{property}`.')
491 continue
493 if select == 'pressure' and property == 'stress':
494 total_stress = extracted_property
495 selected_values.append(-np.sum(total_stress[:3]) / 3)
496 continue
498 mapping = _BEC_MAPPING if property == 'bec' else _REDUCED6_MAPPING
499 if select not in mapping:
500 raise ValueError(f'Selection `{select}` is not allowed.')
501 index = mapping[select]
502 if index >= extracted_property.shape[0]:
503 raise ValueError(
504 f'Selection `{select}` is not compatible with property `{property}`.'
505 )
506 selected_values.append(extracted_property[index])
508 data[stype].append(selected_values)
509 if flatten:
510 for stype in ['target', 'predicted']:
511 value = data[stype]
512 if len(np.shape(value[0])) > 0:
513 data[stype] = np.concatenate(value).ravel().tolist()
514 if property in ['force', 'atomic_v']:
515 default_labels = ['x', 'y', 'z']
516 labels = selection if selection else default_labels
517 n = len(data['target']) // len(labels)
518 data['component'] = list(labels) * n
519 elif property in ['virial', 'stress', 'polarizability']:
520 default_labels = list(ASE_VOIGT6_ORDER)
521 labels = selection if selection else default_labels
522 n = len(data['target']) // len(labels)
523 data['component'] = list(labels) * n
524 elif property in ['bec']:
525 default_labels = list(BEC_FULL9_ORDER)
526 labels = selection if selection else default_labels
527 n = len(data['target']) // len(labels)
528 data['component'] = list(labels) * n
529 df = DataFrame(data)
530 # In case of flatten, cast to float64 for compatibility
531 # with e.g. seaborn.
532 # Casting in this way breaks tensorial properties though,
533 # so skip it there.
534 if flatten:
535 df['target'] = df.target.astype('float64')
536 df['predicted'] = df.predicted.astype('float64')
537 return df