Coverage for calorine/nep/io.py: 100%
278 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-10-08 12:52 +0000
1from collections.abc import Iterable
2from os.path import exists
3from os.path import join as join_path
4from typing import Any, NamedTuple, TextIO
5from warnings import warn
7import numpy as np
8from ase import Atoms
9from ase.io import read, write
10from ase.stress import voigt_6_to_full_3x3_stress
11from pandas import DataFrame
13from calorine.gpumd._header import _read_gpumd_header, _read_gpumd_table
14from calorine.nep.tensor_conventions import ASE_VOIGT6_ORDER, BEC_FULL9_ORDER, \
15 nep_reduced6_to_ase_voigt6
18# GPUMD names of the loss.out columns. A dipole and a polarizability model share
19# the RMSE_P names, since the column count of a file without header cannot tell
20# the two apart.
21_LOSS_HEADER_COLUMN_TAGS = {
22 'total': 'total_loss',
23 'rmse_energy_train': 'RMSE_E_train',
24 'rmse_force_train': 'RMSE_F_train',
25 'rmse_virial_train': 'RMSE_V_train',
26 'rmse_charge_train': 'RMSE_Q_train',
27 'rmse_bec_train': 'RMSE_Z_train',
28 'rmse_dipole_train': 'RMSE_P_train',
29 'rmse_polarizability_train': 'RMSE_P_train',
30 'rmse_ediff_train': 'RMSE_Ediff_train',
31 'rmse_energy_test': 'RMSE_E_test',
32 'rmse_force_test': 'RMSE_F_test',
33 'rmse_virial_test': 'RMSE_V_test',
34 'rmse_charge_test': 'RMSE_Q_test',
35 'rmse_bec_test': 'RMSE_Z_test',
36 'rmse_dipole_test': 'RMSE_P_test',
37 'rmse_polarizability_test': 'RMSE_P_test',
38 'rmse_ediff_test': 'RMSE_Ediff_test',
39}
42def _loss_tags_from_column_count(ncols: int) -> list[str]:
43 """Returns the column names of a loss.out file written without header."""
44 if ncols == 6:
45 tags = 'total_loss L1 L2'
46 tags += ' RMSE_P_train'
47 tags += ' RMSE_P_test'
48 elif ncols == 10:
49 tags = 'total_loss L1 L2'
50 tags += ' RMSE_E_train RMSE_F_train RMSE_V_train'
51 tags += ' RMSE_E_test RMSE_F_test RMSE_V_test'
52 elif ncols == 14:
53 tags = 'total_loss L1 L2'
54 tags += ' RMSE_E_train RMSE_F_train RMSE_V_train RMSE_Q_train RMSE_Z_train'
55 tags += ' RMSE_E_test RMSE_F_test RMSE_V_test RMSE_Q_test RMSE_Z_test'
56 else:
57 raise ValueError(
58 f'Input file contains {ncols} data columns. Expected 6, 10 or 14 columns.'
59 )
60 return ['generation'] + tags.split()
63def read_loss(filename: str) -> DataFrame:
64 """Parses a file in `loss.out` format from GPUMD and returns the content as a data frame.
65 More information concerning file format, content and units can be found `here
66 <https://gpumd.org/nep/output_files/loss_out.html>`__.
68 The column names come from the ``# columns`` line of the header that GPUMD writes before
69 the rows of each run.
70 They are translated to ``total_loss``, ``L1``, ``L2``, and ``RMSE_<X>_train`` and
71 ``RMSE_<X>_test``.
72 Here ``<X>`` is ``E``, ``F``, ``V``, ``Q``, ``Z`` or ``Ediff`` for the energy, force,
73 virial, charge, Born effective charge and energy difference, and ``P`` for the dipole as
74 well as the polarizability.
75 A column name without translation is kept as written.
76 A file without header is read by its number of columns, which is 6 for a dipole or
77 polarizability model, 10 for a potential and 14 for a potential with charges.
79 The data frame is indexed by the generation number, which is read from the first column of
80 the file.
81 The spacing between rows is therefore the value of the ``output_interval`` keyword in
82 `nep.in`.
83 A restart makes `nep` count from zero again while appending its rows to the same file.
84 The index accumulates the runs so that it keeps increasing.
85 Runs with different columns raise a ``ValueError``.
87 Parameters
88 ----------
89 filename
90 input file name
92 """
93 header = _read_gpumd_header(filename, _LOSS_HEADER_COLUMN_TAGS)
94 df = _read_gpumd_table(filename, header, _loss_tags_from_column_count)
95 generations = df.pop('generation').to_numpy().astype(int)
96 # a restart counts from zero again while appending to the same file, so offset
97 # each run by the generation the preceding one reached
98 restarts = np.zeros(len(generations), dtype=bool)
99 if header.blocks:
100 # every run writes its own header, and a run without rows starts at the row
101 # of the run after it
102 starts = np.cumsum([block.n_rows for block in header.blocks], dtype=int)[:-1]
103 restarts[starts[starts < len(df)]] = True
104 else:
105 restarts[1:] = generations[1:] <= generations[:-1]
106 offsets = np.where(restarts, np.concatenate(([0], generations[:-1])), 0)
107 generations += np.cumsum(offsets)
108 df.index = generations
109 return df
112def _write_structure_in_nep_format(structure: Atoms, f: TextIO) -> None:
113 """Write structure block into a file-like object in format readable by nep executable.
115 Parameters
116 ----------
117 structure
118 input structure; must hold information regarding energy and forces
119 f
120 file-like object to which to write
121 """
123 # Allowed keyword=value pairs. Use ASEs extyz write functionality.:
124 # lattice="ax ay az bx by bz cx cy cz" (mandatory)
125 # energy=energy_value (mandatory)
126 # virial="vxx vxy vxz vyx vyy vyz vzx vzy vzz" (optional)
127 # weight=relative_weight (optional)
128 # properties=property_name:data_type:number_of_columns
129 # species:S:1 (mandatory)
130 # pos:R:3 (mandatory)
131 # force:R:3 or forces:R:3 (mandatory)
133 # If a structure is to be used for training, it needs to either have target
134 # 1. energies and forces,
135 # 2. dipole, denoted `dipole="dx dy dz"` in the info string, or
136 # 3. polarizability/susceptibility, denoted `pol="pxx pxy pxz pyx pyy pyz pzx pzy pzz"`
137 # in the info string.
138 has_energies_and_forces = True
139 try:
140 structure.get_potential_energy()
141 structure.get_forces() # calculate forces to have them on the Atoms object
142 except RuntimeError:
143 has_energies_and_forces = False
145 has_dipole = 'dipole' in structure.info.keys()
146 has_pol = 'pol' in structure.info.keys()
148 if not has_energies_and_forces and not has_dipole and not has_pol:
149 raise RuntimeError('Failed to retrieve target energies/forces,'
150 ' dipoles, or polarizabilities for structure')
151 if np.isclose(structure.get_volume(), 0):
152 raise ValueError('Structure cell must have a non-zero volume!')
153 try:
154 structure.get_stress()
155 except RuntimeError:
156 warn('Failed to retrieve stresses for structure')
157 write(filename=f, images=structure, write_info=True, format='extxyz')
160def write_structures(outfile: str, structures: list[Atoms]) -> None:
161 """Writes structures for training/testing in format readable by nep executable.
163 Parameters
164 ----------
165 outfile
166 output filename
167 structures
168 list of structures with energy, forces, and (possibly) stresses
169 """
170 with open(outfile, 'w') as f:
171 for structure in structures:
172 _write_structure_in_nep_format(structure, f)
175def write_nepfile(filename: str, parameters: dict[str, Any]) -> None:
176 """Writes a `nep.in` configuration to the given file name.
178 Keys are written in the order in which they appear in ``parameters``, which matters
179 because the `nep` executable rejects a ``cutoff`` line that precedes the ``type`` line.
181 Parameters
182 ----------
183 filename
184 Name of the file to write.
185 The parent directory must exist.
186 parameters
187 input parameters; see `here <https://gpumd.org/nep/input_parameters/index.html>`__
188 """
189 with open(filename, 'w') as f:
190 for key, val in parameters.items():
191 f.write(f'{key} ')
192 if isinstance(val, Iterable) and not isinstance(val, str):
193 f.write(' '.join([f'{v}' for v in val]))
194 else:
195 f.write(f'{val}')
196 f.write('\n')
199def read_nepfile(filename: str) -> dict[str, Any]:
200 """Returns the content of a configuration file (`nep.in`) as a dictionary.
202 Parameters
203 ----------
204 filename
205 input file name
206 """
207 int_vals = ['version', 'neuron', 'generation', 'batch', 'population',
208 'mode', 'model_type']
209 float_vals = ['lambda_1', 'lambda_2', 'lambda_e', 'lambda_f', 'lambda_v',
210 'lambda_q', 'lambda_shear', 'force_delta', 'atomic_v', 'zbl',
211 'use_typewise_cutoff_zbl']
212 settings = {}
213 with open(filename) as f:
214 for line in f.readlines():
215 # remove comments - throw away everything after a '#'
216 cleaned = line.split('#', 1)[0].strip()
217 flds = cleaned.split()
218 if len(flds) == 0:
219 continue
220 settings[flds[0]] = ' '.join(flds[1:])
221 for key, val in settings.items():
222 if val == '':
223 # a keyword that carries no value, such as a bare use_typewise_cutoff_zbl
224 continue
225 if key in int_vals:
226 settings[key] = int(val)
227 elif key in float_vals:
228 settings[key] = float(val)
229 elif key in ['n_max', 'l_max', 'basis_size', 'charge_mode']:
230 settings[key] = [int(v) for v in val.split()]
231 elif key in ['cutoff', 'type_weight']:
232 settings[key] = [float(v) for v in val.split()]
233 elif key == 'type':
234 types = val.split()
235 types[0] = int(types[0])
236 settings[key] = types
237 return settings
240def read_structures(dirname: str) -> tuple[list[Atoms], list[Atoms]]:
241 """Parses the output files with training and test data from a nep run and returns their
242 content as two lists of structures, representing training and test data, respectively.
243 Target and predicted data are included in the :attr:`info` dict of the :class:`Atoms`
244 objects.
246 A missing `train.xyz` or `test.xyz` yields an empty list in its place.
247 `setup_training` writes no `test.xyz` for a full model, which is trained on every
248 structure, so that case is expected and passes silently.
249 A missing `train.xyz` is reported as a warning.
251 Parameters
252 ----------
253 dirname
254 Directory from which to read output files.
256 """
257 path = join_path(dirname)
258 if not exists(path):
259 raise FileNotFoundError(f'Directory {path} does not exist')
261 # fetch model type from nep input file
262 nep_info = read_nepfile(f'{path}/nep.in')
263 model_type = nep_info.get('model_type', 0)
265 # set up which files to parse, what dimensions to expect etc
266 # depending on the type of model that is parsed
267 #
268 # The `nep` executable's virial_*.out/stress_*.out/polarizability_*.out
269 # output files use a reduced-6 component order (xx,yy,zz,xy,yz,xz) that
270 # differs from ASE's own Voigt-6 convention (xx,yy,zz,yz,xz,xy) -- see
271 # GPUMD's main_nep/structure.cu `reduced_index` array (shared by the
272 # `virial=` and `pol=` extxyz parsers). Below, virial/stress/
273 # polarizability columns are converted to ASE order immediately upon
274 # parsing, so `structure.info`/`.arrays` always hold ASE-Voigt-ordered
275 # data, in the order `atoms.get_stress()` returns.
276 # The stored stress is the virial divided by the volume, in GPa, as GPUMD
277 # writes it (`Fitness::output` in main_nep/fitness.cu, with virial = -V * stress
278 # in main_nep/structure.cu). It is therefore the negative of `atoms.get_stress()`,
279 # which is in eV/Å^3. BEC's 9 columns are left as is (row-major, see
280 # BEC_FULL9_ORDER) -- BEC is not symmetric, so no Voigt form applies.
281 if model_type == 0:
282 charge_mode = nep_info.get('charge_mode', [0])[0]
283 if charge_mode not in [0, 1, 2]:
284 raise ValueError(f'Unknown charge_mode: {charge_mode}')
285 # files to parse: (sname, size, mandatory, includes_target, per_atom)
286 files_to_parse = [
287 ('energy', 1, True, True, False),
288 ('force', 3, True, True, True),
289 ('virial', 6, True, True, False),
290 ('stress', 6, True, True, False),
291 ]
292 if charge_mode in [1, 2]:
293 # files to parse: (sname, size, includes_target, per_atom)
294 files_to_parse += [
295 ('charge', 1, True, False, True),
296 ('bec', 9, False, True, True),
297 ]
298 elif model_type == 1:
299 # files to parse: (sname, size, includes_target, per_atom)
300 files_to_parse = [('dipole', 3, True, True, False)]
301 if nep_info.get('atomic_v', 0) == 1:
302 files_to_parse = [('dipole', 3, True, True, True)]
303 elif model_type == 2:
304 # files to parse: (sname, size, includes_target, per_atom)
305 files_to_parse = [('polarizability', 6, True, True, False)]
306 if nep_info.get('atomic_v', 0) == 1:
307 files_to_parse = [('polarizability', 6, True, True, True)]
308 else:
309 raise ValueError(f'Unknown model_type: {model_type}')
311 # read training and test data
312 structures = {}
313 for stype in ['train', 'test']:
314 filename = join_path(dirname, f'{stype}.xyz')
315 try:
316 structures[stype] = read(filename, format='extxyz', index=':')
317 except FileNotFoundError:
318 # a model trained on every structure has no test.xyz
319 if stype == 'train':
320 warn(f'File {filename} not found.')
321 structures[stype] = []
322 continue
324 n_structures = len(structures[stype])
326 # loop over files from which to read target data and predictions
327 for sname, size, mandatory, includes_target, per_atom in files_to_parse:
328 infile = f'{sname}_{stype}.out'
329 path = join_path(dirname, infile)
330 if not exists(path):
331 if mandatory:
332 raise FileNotFoundError(f'File {path} does not exist')
333 else:
334 continue
335 ts, ps = _read_data_file(path, includes_target=includes_target)
337 if ts is not None:
338 if ts.shape[1] != size:
339 raise ValueError(f'Target data in {infile} has unexpected shape:'
340 f' {ts.shape} (expected: (-1, {size}))')
341 if ps.shape[1] != size:
342 raise ValueError(f'Predicted data in {infile} has unexpected shape:'
343 f' {ps.shape} (expected: (-1, {size}))')
345 if sname in ('virial', 'stress', 'polarizability'):
346 ts = nep_reduced6_to_ase_voigt6(ts)
347 ps = nep_reduced6_to_ase_voigt6(ps)
349 if per_atom:
350 # data per-atom, e.g., forces, per-atom-virials, Born effective charges ...
351 n_atoms_total = sum([len(s) for s in structures[stype]])
352 if len(ps) != n_atoms_total:
353 raise ValueError(f'Number of atoms in {infile} ({len(ps)})'
354 f' and {stype}.xyz ({n_atoms_total}) inconsistent.')
355 n = 0
356 for structure in structures[stype]:
357 nat = len(structure)
358 if ts is not None:
359 t = np.array(ts[n: n + nat]).reshape(nat, size)
360 structure.new_array(f'{sname}_target', t)
361 p = np.array(ps[n: n + nat]).reshape(nat, size)
362 structure.new_array(f'{sname}_predicted', p)
363 n += nat
364 else:
365 # data per structure, e.g., energy, virials, stress
366 if len(ps) != n_structures:
367 raise ValueError(f'Number of structures in {infile} ({len(ps)})'
368 f' and {stype}.xyz ({n_structures}) inconsistent.')
369 for k, structure in enumerate(structures[stype]):
370 assert ts is not None, 'This should not occur. Please report.'
371 t = ts[k]
372 assert np.shape(t) == (size,)
373 structure.info[f'{sname}_target'] = t
374 p = ps[k]
375 assert np.shape(p) == (size,)
376 structure.info[f'{sname}_predicted'] = p
378 # special handling of target data for BECs
379 # If a structure has no 'bec' array in the xyz file, no target BEC data was provided.
380 # In that case nep writes zeros for both predicted and target columns. Replace both
381 # with NaN so callers can easily identify and filter out structures without BEC targets.
382 for s in structures[stype]:
383 if 'bec_target' in s.arrays and 'bec' not in s.arrays:
384 nat = len(s)
385 s.arrays['bec_target'] = np.full((nat, 9), np.nan)
386 s.arrays['bec_predicted'] = np.full((nat, 9), np.nan)
388 # special handling of per-atom TNEP
389 # Data has to be loaded as dipole/polarizability since NEP saves them in dipole_*.out
390 # Dipole/polarizability arrays are therefore moved to atomic_v here
391 if nep_info.get('atomic_v', 0) == 1:
392 if model_type == 1:
393 s.new_array('atomic_v_target', s.arrays['dipole_target'])
394 s.new_array('atomic_v_predicted', s.arrays['dipole_predicted'])
395 del s.arrays['dipole_target']
396 del s.arrays['dipole_predicted']
397 if model_type == 2:
398 s.new_array('atomic_v_target', s.arrays['polarizability_target'])
399 s.new_array('atomic_v_predicted', s.arrays['polarizability_predicted'])
400 del s.arrays['polarizability_target']
401 del s.arrays['polarizability_predicted']
403 return structures['train'], structures['test']
406def _read_data_file(
407 path: str,
408 includes_target: bool = True,
409):
410 """Private function that parses *.out files and
411 returns their content for further processing.
412 """
413 with open(path, 'r') as f:
414 lines = f.readlines()
415 target, predicted = [], []
416 for line in lines:
417 flds = line.split()
418 if includes_target:
419 if len(flds) % 2 != 0:
420 raise ValueError(f'Incorrect number of columns in {path} ({len(flds)}).')
421 n = len(flds) // 2
422 predicted.append([float(s) for s in flds[:n]])
423 target.append([float(s) for s in flds[n:]])
424 else:
425 predicted.append([float(s) for s in flds])
426 target = None
427 if target is not None:
428 target = np.array(target)
429 predicted = np.array(predicted)
430 return target, predicted
433# Maps x/y/z (vectors) and reduced-6 symmetric-tensor component names to
434# their index. virial/stress/polarizability (and per-atom atomic_v when it
435# holds a reduced-6 polarizability) are converted to ASE-Voigt order by
436# read_structures() above, so this mapping uses ASE_VOIGT6_ORDER directly.
437_REDUCED6_MAPPING = {
438 'x': 0, 'y': 1, 'z': 2,
439 **{label: i for i, label in enumerate(ASE_VOIGT6_ORDER)},
440}
441# BEC is a raw, unreduced, generally-asymmetric 9-component per-atom tensor
442# (row-major, see BEC_FULL9_ORDER) -- not a Voigt-reduced 6-component form,
443# so it needs its own, different index scheme.
444_BEC_MAPPING = {label: i for i, label in enumerate(BEC_FULL9_ORDER)}
447# Component labels without a selection, by the number of components.
448_DEFAULT_LABELS = {
449 3: ['x', 'y', 'z'],
450 6: list(ASE_VOIGT6_ORDER),
451 9: list(BEC_FULL9_ORDER),
452}
455class _ParityProperty(NamedTuple):
456 """How get_parity_data reads and selects one property."""
457 # True for a value per atom in `arrays`, False for a value per structure in `info`
458 per_atom: bool
459 # numbers of components along the last axis, empty for a scalar
460 component_counts: tuple[int, ...]
461 # index of each component label
462 component_indices: dict[str, int]
463 # selections other than a component, out of `norm` and `pressure`
464 special_selections: tuple[str, ...]
465 # whether the flattened data carry a `species` column
466 has_species: bool
469_PARITY_PROPERTIES = {
470 'energy': _ParityProperty(False, (), {}, (), False),
471 'virial': _ParityProperty(False, (6,), _REDUCED6_MAPPING, ('norm',), False),
472 'stress': _ParityProperty(False, (6,), _REDUCED6_MAPPING, ('norm', 'pressure'), False),
473 'polarizability': _ParityProperty(False, (6,), _REDUCED6_MAPPING, (), False),
474 'dipole': _ParityProperty(False, (3,), _REDUCED6_MAPPING, ('norm',), False),
475 'force': _ParityProperty(True, (3,), _REDUCED6_MAPPING, ('norm',), True),
476 'bec': _ParityProperty(True, (9,), _BEC_MAPPING, (), True),
477 'atomic_v': _ParityProperty(True, (3, 6), _REDUCED6_MAPPING, ('norm',), False),
478}
481def _select_component(
482 values: np.ndarray, select: str, property: str, rules: _ParityProperty
483) -> np.ndarray:
484 """Returns one selected quantity of `values`, whose components run along the last axis."""
485 n_components = values.shape[-1]
486 if select == 'norm':
487 if 'norm' not in rules.special_selections:
488 raise ValueError(f'Cannot handle selection=`norm` with property=`{property}`.')
489 if n_components == 3:
490 return np.linalg.norm(values, axis=-1)
491 return np.linalg.norm(voigt_6_to_full_3x3_stress(values), axis=(-2, -1))
492 if select == 'pressure':
493 # the stored stress is the negative of the ASE stress, positive under compression
494 return np.sum(values[..., :3], axis=-1) / 3
495 if select not in rules.component_indices:
496 raise ValueError(f'Selection `{select}` is not allowed.')
497 if rules.component_indices[select] >= n_components:
498 raise ValueError(f'Selection `{select}` is not compatible with property `{property}`.')
499 return values[..., rules.component_indices[select]]
502def get_parity_data(
503 structures: list[Atoms],
504 property: str,
505 selection: list[str] = None,
506 flatten: bool = True,
507) -> DataFrame:
508 """Returns the predicted and target values of a property from a list of
509 structures in a format suitable for generating parity plots.
511 The structures should have been read using :func:`read_structures
512 <calorine.nep.read_structures>`, such that the `info` or `arrays` object
513 is populated with keys of the form `<property>_<type>`, where `<type>` is
514 one of `predicted` or `target`.
516 Parameters
517 ----------
518 structures
519 List of structures as read with :func:`read_structures <calorine.nep.read_structures>`.
520 property
521 One of `energy`, `force`, `virial`, `stress`, `bec`, `dipole`,
522 `polarizability`, or `atomic_v`.
523 `stress` is the virial divided by the volume, in GPa.
524 It is the negative of the stress returned by `atoms.get_stress()`.
525 selection
526 List of the components to return, and/or the norm.
527 For `force`, `atomic_v`, `virial`, `stress`, `polarizability`,
528 and `dipole`, possible values are `x`, `y`, `z`, `xx`, `yy`, `zz`,
529 `yz`, `xz`, `xy`, `norm`, and `pressure` (the latter only for `stress`).
530 `x` and `xx` select the same component, and likewise for `y` and `z`.
531 `pressure` is positive under compression.
532 `norm` is the Euclidean norm of a vector and the Frobenius norm of a
533 rank-2 tensor, and is not available for `polarizability`.
534 For `bec`, all nine of `xx`, `xy`, `xz`, `yx`, `yy`, `yz`, `zx`,
535 `zy`, `zz` are available (BEC is not symmetric, so unlike the other
536 properties it has no reduced/Voigt form).
537 flatten
538 If True, return one row per value.
539 If False, return one row per structure, where a per-atom property holds
540 an array with one row per atom and one column per component.
542 Returns
543 -------
544 Parity data with the columns `predicted` and `target`.
545 With `flatten=True`, a `component` column names the component of each
546 value for every property except `energy`.
547 For `force` and `bec`, a `species` column names the chemical symbol of
548 the atom.
549 """
550 if property not in _PARITY_PROPERTIES:
551 raise ValueError(
552 '`property` must be one of the following: ' + ', '.join(_PARITY_PROPERTIES))
553 rules = _PARITY_PROPERTIES[property]
554 if isinstance(selection, str):
555 raise TypeError('`selection` must be a list of strings.')
556 selection = [] if selection is None else list(selection)
557 if not rules.component_counts and selection:
558 raise ValueError('Selection cannot be applied to scalars.')
559 if 'pressure' in selection and 'pressure' not in rules.special_selections:
560 raise ValueError(f'Cannot calculate pressure for `{property}`.')
562 data = {'predicted': [], 'target': []}
563 if flatten and rules.has_species:
564 data['species'] = []
565 if flatten and rules.component_counts:
566 data['component'] = []
567 for structure in structures:
568 for stype in ['predicted', 'target']:
569 property_with_stype = f'{property}_{stype}'
570 if not rules.per_atom:
571 if property_with_stype not in structure.info.keys():
572 raise KeyError(f'{property_with_stype} not'
573 ' available in info field of structure')
574 values = np.array(structure.info[property_with_stype])
575 else:
576 if property_with_stype not in structure.arrays:
577 raise KeyError(f'{property_with_stype} not available in arrays of structure')
578 values = np.array(structure.arrays[property_with_stype])
579 if rules.component_counts:
580 n_dimensions = 2 if rules.per_atom else 1
581 if (values.ndim != n_dimensions
582 or values.shape[-1] not in rules.component_counts):
583 shapes = ' or '.join(
584 f'(n_atoms, {n})' if rules.per_atom else f'({n},)'
585 for n in rules.component_counts)
586 raise ValueError(f'`{property_with_stype}` has shape {values.shape},'
587 f' while `{property}` needs {shapes}.')
588 if selection:
589 values = np.stack(
590 [_select_component(values, select, property, rules) for select in selection],
591 axis=-1)
592 data[stype].append(values)
593 if 'component' in data:
594 labels = selection or _DEFAULT_LABELS[values.shape[-1]]
595 data['component'].extend(labels * (values.size // len(labels)))
596 if 'species' in data:
597 data['species'].extend(np.repeat(structure.symbols, values.shape[-1]).tolist())
598 if flatten:
599 for stype in ['predicted', 'target']:
600 if data[stype]:
601 data[stype] = np.concatenate([np.ravel(values) for values in data[stype]])
602 df = DataFrame(data)
603 # In case of flatten, cast to float64 for compatibility
604 # with e.g. seaborn.
605 # Casting in this way breaks tensorial properties though,
606 # so skip it there.
607 if flatten:
608 df['target'] = df.target.astype('float64')
609 df['predicted'] = df.predicted.astype('float64')
610 return df