from collections.abc import Iterable
from os.path import exists
from os.path import join as join_path
from typing import Any, NamedTuple, TextIO
from warnings import warn
import numpy as np
from ase import Atoms
from ase.io import read, write
from ase.stress import voigt_6_to_full_3x3_stress
from pandas import DataFrame
from calorine.gpumd._header import _read_gpumd_header, _read_gpumd_table
from calorine.nep.tensor_conventions import ASE_VOIGT6_ORDER, BEC_FULL9_ORDER, \
nep_reduced6_to_ase_voigt6
# GPUMD names of the loss.out columns. A dipole and a polarizability model share
# the RMSE_P names, since the column count of a file without header cannot tell
# the two apart.
_LOSS_HEADER_COLUMN_TAGS = {
'total': 'total_loss',
'rmse_energy_train': 'RMSE_E_train',
'rmse_force_train': 'RMSE_F_train',
'rmse_virial_train': 'RMSE_V_train',
'rmse_charge_train': 'RMSE_Q_train',
'rmse_bec_train': 'RMSE_Z_train',
'rmse_dipole_train': 'RMSE_P_train',
'rmse_polarizability_train': 'RMSE_P_train',
'rmse_ediff_train': 'RMSE_Ediff_train',
'rmse_energy_test': 'RMSE_E_test',
'rmse_force_test': 'RMSE_F_test',
'rmse_virial_test': 'RMSE_V_test',
'rmse_charge_test': 'RMSE_Q_test',
'rmse_bec_test': 'RMSE_Z_test',
'rmse_dipole_test': 'RMSE_P_test',
'rmse_polarizability_test': 'RMSE_P_test',
'rmse_ediff_test': 'RMSE_Ediff_test',
}
def _loss_tags_from_column_count(ncols: int) -> list[str]:
"""Returns the column names of a loss.out file written without header."""
if ncols == 6:
tags = 'total_loss L1 L2'
tags += ' RMSE_P_train'
tags += ' RMSE_P_test'
elif ncols == 10:
tags = 'total_loss L1 L2'
tags += ' RMSE_E_train RMSE_F_train RMSE_V_train'
tags += ' RMSE_E_test RMSE_F_test RMSE_V_test'
elif ncols == 14:
tags = 'total_loss L1 L2'
tags += ' RMSE_E_train RMSE_F_train RMSE_V_train RMSE_Q_train RMSE_Z_train'
tags += ' RMSE_E_test RMSE_F_test RMSE_V_test RMSE_Q_test RMSE_Z_test'
else:
raise ValueError(
f'Input file contains {ncols} data columns. Expected 6, 10 or 14 columns.'
)
return ['generation'] + tags.split()
[docs]
def read_loss(filename: str) -> DataFrame:
"""Parses a file in `loss.out` format from GPUMD and returns the content as a data frame.
More information concerning file format, content and units can be found `here
<https://gpumd.org/nep/output_files/loss_out.html>`__.
The column names come from the ``# columns`` line of the header that GPUMD writes before
the rows of each run.
They are translated to ``total_loss``, ``L1``, ``L2``, and ``RMSE_<X>_train`` and
``RMSE_<X>_test``.
Here ``<X>`` is ``E``, ``F``, ``V``, ``Q``, ``Z`` or ``Ediff`` for the energy, force,
virial, charge, Born effective charge and energy difference, and ``P`` for the dipole as
well as the polarizability.
A column name without translation is kept as written.
A file without header is read by its number of columns, which is 6 for a dipole or
polarizability model, 10 for a potential and 14 for a potential with charges.
The data frame is indexed by the generation number, which is read from the first column of
the file.
The spacing between rows is therefore the value of the ``output_interval`` keyword in
`nep.in`.
A restart makes `nep` count from zero again while appending its rows to the same file.
The index accumulates the runs so that it keeps increasing.
Runs with different columns raise a ``ValueError``.
Parameters
----------
filename
input file name
"""
header = _read_gpumd_header(filename, _LOSS_HEADER_COLUMN_TAGS)
df = _read_gpumd_table(filename, header, _loss_tags_from_column_count)
generations = df.pop('generation').to_numpy().astype(int)
# a restart counts from zero again while appending to the same file, so offset
# each run by the generation the preceding one reached
restarts = np.zeros(len(generations), dtype=bool)
if header.blocks:
# every run writes its own header, and a run without rows starts at the row
# of the run after it
starts = np.cumsum([block.n_rows for block in header.blocks], dtype=int)[:-1]
restarts[starts[starts < len(df)]] = True
else:
restarts[1:] = generations[1:] <= generations[:-1]
offsets = np.where(restarts, np.concatenate(([0], generations[:-1])), 0)
generations += np.cumsum(offsets)
df.index = generations
return df
def _write_structure_in_nep_format(structure: Atoms, f: TextIO) -> None:
"""Write structure block into a file-like object in format readable by nep executable.
Parameters
----------
structure
input structure; must hold information regarding energy and forces
f
file-like object to which to write
"""
# Allowed keyword=value pairs. Use ASEs extyz write functionality.:
# lattice="ax ay az bx by bz cx cy cz" (mandatory)
# energy=energy_value (mandatory)
# virial="vxx vxy vxz vyx vyy vyz vzx vzy vzz" (optional)
# weight=relative_weight (optional)
# properties=property_name:data_type:number_of_columns
# species:S:1 (mandatory)
# pos:R:3 (mandatory)
# force:R:3 or forces:R:3 (mandatory)
# If a structure is to be used for training, it needs to either have target
# 1. energies and forces,
# 2. dipole, denoted `dipole="dx dy dz"` in the info string, or
# 3. polarizability/susceptibility, denoted `pol="pxx pxy pxz pyx pyy pyz pzx pzy pzz"`
# in the info string.
has_energies_and_forces = True
try:
structure.get_potential_energy()
structure.get_forces() # calculate forces to have them on the Atoms object
except RuntimeError:
has_energies_and_forces = False
has_dipole = 'dipole' in structure.info.keys()
has_pol = 'pol' in structure.info.keys()
if not has_energies_and_forces and not has_dipole and not has_pol:
raise RuntimeError('Failed to retrieve target energies/forces,'
' dipoles, or polarizabilities for structure')
if np.isclose(structure.get_volume(), 0):
raise ValueError('Structure cell must have a non-zero volume!')
try:
structure.get_stress()
except RuntimeError:
warn('Failed to retrieve stresses for structure')
write(filename=f, images=structure, write_info=True, format='extxyz')
[docs]
def write_structures(outfile: str, structures: list[Atoms]) -> None:
"""Writes structures for training/testing in format readable by nep executable.
Parameters
----------
outfile
output filename
structures
list of structures with energy, forces, and (possibly) stresses
"""
with open(outfile, 'w') as f:
for structure in structures:
_write_structure_in_nep_format(structure, f)
[docs]
def write_nepfile(filename: str, parameters: dict[str, Any]) -> None:
"""Writes a `nep.in` configuration to the given file name.
Keys are written in the order in which they appear in ``parameters``, which matters
because the `nep` executable rejects a ``cutoff`` line that precedes the ``type`` line.
Parameters
----------
filename
Name of the file to write.
The parent directory must exist.
parameters
input parameters; see `here <https://gpumd.org/nep/input_parameters/index.html>`__
"""
with open(filename, 'w') as f:
for key, val in parameters.items():
f.write(f'{key} ')
if isinstance(val, Iterable) and not isinstance(val, str):
f.write(' '.join([f'{v}' for v in val]))
else:
f.write(f'{val}')
f.write('\n')
[docs]
def read_nepfile(filename: str) -> dict[str, Any]:
"""Returns the content of a configuration file (`nep.in`) as a dictionary.
Parameters
----------
filename
input file name
"""
int_vals = ['version', 'neuron', 'generation', 'batch', 'population',
'mode', 'model_type']
float_vals = ['lambda_1', 'lambda_2', 'lambda_e', 'lambda_f', 'lambda_v',
'lambda_q', 'lambda_shear', 'force_delta', 'atomic_v', 'zbl',
'use_typewise_cutoff_zbl']
settings = {}
with open(filename) as f:
for line in f.readlines():
# remove comments - throw away everything after a '#'
cleaned = line.split('#', 1)[0].strip()
flds = cleaned.split()
if len(flds) == 0:
continue
settings[flds[0]] = ' '.join(flds[1:])
for key, val in settings.items():
if val == '':
# a keyword that carries no value, such as a bare use_typewise_cutoff_zbl
continue
if key in int_vals:
settings[key] = int(val)
elif key in float_vals:
settings[key] = float(val)
elif key in ['n_max', 'l_max', 'basis_size', 'charge_mode']:
settings[key] = [int(v) for v in val.split()]
elif key in ['cutoff', 'type_weight']:
settings[key] = [float(v) for v in val.split()]
elif key == 'type':
types = val.split()
types[0] = int(types[0])
settings[key] = types
return settings
[docs]
def read_structures(dirname: str) -> tuple[list[Atoms], list[Atoms]]:
"""Parses the output files with training and test data from a nep run and returns their
content as two lists of structures, representing training and test data, respectively.
Target and predicted data are included in the :attr:`info` dict of the :class:`Atoms`
objects.
A missing `train.xyz` or `test.xyz` yields an empty list in its place.
`setup_training` writes no `test.xyz` for a full model, which is trained on every
structure, so that case is expected and passes silently.
A missing `train.xyz` is reported as a warning.
Parameters
----------
dirname
Directory from which to read output files.
"""
path = join_path(dirname)
if not exists(path):
raise FileNotFoundError(f'Directory {path} does not exist')
# fetch model type from nep input file
nep_info = read_nepfile(f'{path}/nep.in')
model_type = nep_info.get('model_type', 0)
# set up which files to parse, what dimensions to expect etc
# depending on the type of model that is parsed
#
# The `nep` executable's virial_*.out/stress_*.out/polarizability_*.out
# output files use a reduced-6 component order (xx,yy,zz,xy,yz,xz) that
# differs from ASE's own Voigt-6 convention (xx,yy,zz,yz,xz,xy) -- see
# GPUMD's main_nep/structure.cu `reduced_index` array (shared by the
# `virial=` and `pol=` extxyz parsers). Below, virial/stress/
# polarizability columns are converted to ASE order immediately upon
# parsing, so `structure.info`/`.arrays` always hold ASE-Voigt-ordered
# data, in the order `atoms.get_stress()` returns.
# The stored stress is the virial divided by the volume, in GPa, as GPUMD
# writes it (`Fitness::output` in main_nep/fitness.cu, with virial = -V * stress
# in main_nep/structure.cu). It is therefore the negative of `atoms.get_stress()`,
# which is in eV/Å^3. BEC's 9 columns are left as is (row-major, see
# BEC_FULL9_ORDER) -- BEC is not symmetric, so no Voigt form applies.
if model_type == 0:
charge_mode = nep_info.get('charge_mode', [0])[0]
if charge_mode not in [0, 1, 2]:
raise ValueError(f'Unknown charge_mode: {charge_mode}')
# files to parse: (sname, size, mandatory, includes_target, per_atom)
files_to_parse = [
('energy', 1, True, True, False),
('force', 3, True, True, True),
('virial', 6, True, True, False),
('stress', 6, True, True, False),
]
if charge_mode in [1, 2]:
# files to parse: (sname, size, includes_target, per_atom)
files_to_parse += [
('charge', 1, True, False, True),
('bec', 9, False, True, True),
]
elif model_type == 1:
# files to parse: (sname, size, includes_target, per_atom)
files_to_parse = [('dipole', 3, True, True, False)]
if nep_info.get('atomic_v', 0) == 1:
files_to_parse = [('dipole', 3, True, True, True)]
elif model_type == 2:
# files to parse: (sname, size, includes_target, per_atom)
files_to_parse = [('polarizability', 6, True, True, False)]
if nep_info.get('atomic_v', 0) == 1:
files_to_parse = [('polarizability', 6, True, True, True)]
else:
raise ValueError(f'Unknown model_type: {model_type}')
# read training and test data
structures = {}
for stype in ['train', 'test']:
filename = join_path(dirname, f'{stype}.xyz')
try:
structures[stype] = read(filename, format='extxyz', index=':')
except FileNotFoundError:
# a model trained on every structure has no test.xyz
if stype == 'train':
warn(f'File {filename} not found.')
structures[stype] = []
continue
n_structures = len(structures[stype])
# loop over files from which to read target data and predictions
for sname, size, mandatory, includes_target, per_atom in files_to_parse:
infile = f'{sname}_{stype}.out'
path = join_path(dirname, infile)
if not exists(path):
if mandatory:
raise FileNotFoundError(f'File {path} does not exist')
else:
continue
ts, ps = _read_data_file(path, includes_target=includes_target)
if ts is not None:
if ts.shape[1] != size:
raise ValueError(f'Target data in {infile} has unexpected shape:'
f' {ts.shape} (expected: (-1, {size}))')
if ps.shape[1] != size:
raise ValueError(f'Predicted data in {infile} has unexpected shape:'
f' {ps.shape} (expected: (-1, {size}))')
if sname in ('virial', 'stress', 'polarizability'):
ts = nep_reduced6_to_ase_voigt6(ts)
ps = nep_reduced6_to_ase_voigt6(ps)
if per_atom:
# data per-atom, e.g., forces, per-atom-virials, Born effective charges ...
n_atoms_total = sum([len(s) for s in structures[stype]])
if len(ps) != n_atoms_total:
raise ValueError(f'Number of atoms in {infile} ({len(ps)})'
f' and {stype}.xyz ({n_atoms_total}) inconsistent.')
n = 0
for structure in structures[stype]:
nat = len(structure)
if ts is not None:
t = np.array(ts[n: n + nat]).reshape(nat, size)
structure.new_array(f'{sname}_target', t)
p = np.array(ps[n: n + nat]).reshape(nat, size)
structure.new_array(f'{sname}_predicted', p)
n += nat
else:
# data per structure, e.g., energy, virials, stress
if len(ps) != n_structures:
raise ValueError(f'Number of structures in {infile} ({len(ps)})'
f' and {stype}.xyz ({n_structures}) inconsistent.')
for k, structure in enumerate(structures[stype]):
assert ts is not None, 'This should not occur. Please report.'
t = ts[k]
assert np.shape(t) == (size,)
structure.info[f'{sname}_target'] = t
p = ps[k]
assert np.shape(p) == (size,)
structure.info[f'{sname}_predicted'] = p
# special handling of target data for BECs
# If a structure has no 'bec' array in the xyz file, no target BEC data was provided.
# In that case nep writes zeros for both predicted and target columns. Replace both
# with NaN so callers can easily identify and filter out structures without BEC targets.
for s in structures[stype]:
if 'bec_target' in s.arrays and 'bec' not in s.arrays:
nat = len(s)
s.arrays['bec_target'] = np.full((nat, 9), np.nan)
s.arrays['bec_predicted'] = np.full((nat, 9), np.nan)
# special handling of per-atom TNEP
# Data has to be loaded as dipole/polarizability since NEP saves them in dipole_*.out
# Dipole/polarizability arrays are therefore moved to atomic_v here
if nep_info.get('atomic_v', 0) == 1:
if model_type == 1:
s.new_array('atomic_v_target', s.arrays['dipole_target'])
s.new_array('atomic_v_predicted', s.arrays['dipole_predicted'])
del s.arrays['dipole_target']
del s.arrays['dipole_predicted']
if model_type == 2:
s.new_array('atomic_v_target', s.arrays['polarizability_target'])
s.new_array('atomic_v_predicted', s.arrays['polarizability_predicted'])
del s.arrays['polarizability_target']
del s.arrays['polarizability_predicted']
return structures['train'], structures['test']
def _read_data_file(
path: str,
includes_target: bool = True,
):
"""Private function that parses *.out files and
returns their content for further processing.
"""
with open(path, 'r') as f:
lines = f.readlines()
target, predicted = [], []
for line in lines:
flds = line.split()
if includes_target:
if len(flds) % 2 != 0:
raise ValueError(f'Incorrect number of columns in {path} ({len(flds)}).')
n = len(flds) // 2
predicted.append([float(s) for s in flds[:n]])
target.append([float(s) for s in flds[n:]])
else:
predicted.append([float(s) for s in flds])
target = None
if target is not None:
target = np.array(target)
predicted = np.array(predicted)
return target, predicted
# Maps x/y/z (vectors) and reduced-6 symmetric-tensor component names to
# their index. virial/stress/polarizability (and per-atom atomic_v when it
# holds a reduced-6 polarizability) are converted to ASE-Voigt order by
# read_structures() above, so this mapping uses ASE_VOIGT6_ORDER directly.
_REDUCED6_MAPPING = {
'x': 0, 'y': 1, 'z': 2,
**{label: i for i, label in enumerate(ASE_VOIGT6_ORDER)},
}
# BEC is a raw, unreduced, generally-asymmetric 9-component per-atom tensor
# (row-major, see BEC_FULL9_ORDER) -- not a Voigt-reduced 6-component form,
# so it needs its own, different index scheme.
_BEC_MAPPING = {label: i for i, label in enumerate(BEC_FULL9_ORDER)}
# Component labels without a selection, by the number of components.
_DEFAULT_LABELS = {
3: ['x', 'y', 'z'],
6: list(ASE_VOIGT6_ORDER),
9: list(BEC_FULL9_ORDER),
}
class _ParityProperty(NamedTuple):
"""How get_parity_data reads and selects one property."""
# True for a value per atom in `arrays`, False for a value per structure in `info`
per_atom: bool
# numbers of components along the last axis, empty for a scalar
component_counts: tuple[int, ...]
# index of each component label
component_indices: dict[str, int]
# selections other than a component, out of `norm` and `pressure`
special_selections: tuple[str, ...]
# whether the flattened data carry a `species` column
has_species: bool
_PARITY_PROPERTIES = {
'energy': _ParityProperty(False, (), {}, (), False),
'virial': _ParityProperty(False, (6,), _REDUCED6_MAPPING, ('norm',), False),
'stress': _ParityProperty(False, (6,), _REDUCED6_MAPPING, ('norm', 'pressure'), False),
'polarizability': _ParityProperty(False, (6,), _REDUCED6_MAPPING, (), False),
'dipole': _ParityProperty(False, (3,), _REDUCED6_MAPPING, ('norm',), False),
'force': _ParityProperty(True, (3,), _REDUCED6_MAPPING, ('norm',), True),
'bec': _ParityProperty(True, (9,), _BEC_MAPPING, (), True),
'atomic_v': _ParityProperty(True, (3, 6), _REDUCED6_MAPPING, ('norm',), False),
}
def _select_component(
values: np.ndarray, select: str, property: str, rules: _ParityProperty
) -> np.ndarray:
"""Returns one selected quantity of `values`, whose components run along the last axis."""
n_components = values.shape[-1]
if select == 'norm':
if 'norm' not in rules.special_selections:
raise ValueError(f'Cannot handle selection=`norm` with property=`{property}`.')
if n_components == 3:
return np.linalg.norm(values, axis=-1)
return np.linalg.norm(voigt_6_to_full_3x3_stress(values), axis=(-2, -1))
if select == 'pressure':
# the stored stress is the negative of the ASE stress, positive under compression
return np.sum(values[..., :3], axis=-1) / 3
if select not in rules.component_indices:
raise ValueError(f'Selection `{select}` is not allowed.')
if rules.component_indices[select] >= n_components:
raise ValueError(f'Selection `{select}` is not compatible with property `{property}`.')
return values[..., rules.component_indices[select]]
[docs]
def get_parity_data(
structures: list[Atoms],
property: str,
selection: list[str] = None,
flatten: bool = True,
) -> DataFrame:
"""Returns the predicted and target values of a property from a list of
structures in a format suitable for generating parity plots.
The structures should have been read using :func:`read_structures
<calorine.nep.read_structures>`, such that the `info` or `arrays` object
is populated with keys of the form `<property>_<type>`, where `<type>` is
one of `predicted` or `target`.
Parameters
----------
structures
List of structures as read with :func:`read_structures <calorine.nep.read_structures>`.
property
One of `energy`, `force`, `virial`, `stress`, `bec`, `dipole`,
`polarizability`, or `atomic_v`.
`stress` is the virial divided by the volume, in GPa.
It is the negative of the stress returned by `atoms.get_stress()`.
selection
List of the components to return, and/or the norm.
For `force`, `atomic_v`, `virial`, `stress`, `polarizability`,
and `dipole`, possible values are `x`, `y`, `z`, `xx`, `yy`, `zz`,
`yz`, `xz`, `xy`, `norm`, and `pressure` (the latter only for `stress`).
`x` and `xx` select the same component, and likewise for `y` and `z`.
`pressure` is positive under compression.
`norm` is the Euclidean norm of a vector and the Frobenius norm of a
rank-2 tensor, and is not available for `polarizability`.
For `bec`, all nine of `xx`, `xy`, `xz`, `yx`, `yy`, `yz`, `zx`,
`zy`, `zz` are available (BEC is not symmetric, so unlike the other
properties it has no reduced/Voigt form).
flatten
If True, return one row per value.
If False, return one row per structure, where a per-atom property holds
an array with one row per atom and one column per component.
Returns
-------
Parity data with the columns `predicted` and `target`.
With `flatten=True`, a `component` column names the component of each
value for every property except `energy`.
For `force` and `bec`, a `species` column names the chemical symbol of
the atom.
"""
if property not in _PARITY_PROPERTIES:
raise ValueError(
'`property` must be one of the following: ' + ', '.join(_PARITY_PROPERTIES))
rules = _PARITY_PROPERTIES[property]
if isinstance(selection, str):
raise TypeError('`selection` must be a list of strings.')
selection = [] if selection is None else list(selection)
if not rules.component_counts and selection:
raise ValueError('Selection cannot be applied to scalars.')
if 'pressure' in selection and 'pressure' not in rules.special_selections:
raise ValueError(f'Cannot calculate pressure for `{property}`.')
data = {'predicted': [], 'target': []}
if flatten and rules.has_species:
data['species'] = []
if flatten and rules.component_counts:
data['component'] = []
for structure in structures:
for stype in ['predicted', 'target']:
property_with_stype = f'{property}_{stype}'
if not rules.per_atom:
if property_with_stype not in structure.info.keys():
raise KeyError(f'{property_with_stype} not'
' available in info field of structure')
values = np.array(structure.info[property_with_stype])
else:
if property_with_stype not in structure.arrays:
raise KeyError(f'{property_with_stype} not available in arrays of structure')
values = np.array(structure.arrays[property_with_stype])
if rules.component_counts:
n_dimensions = 2 if rules.per_atom else 1
if (values.ndim != n_dimensions
or values.shape[-1] not in rules.component_counts):
shapes = ' or '.join(
f'(n_atoms, {n})' if rules.per_atom else f'({n},)'
for n in rules.component_counts)
raise ValueError(f'`{property_with_stype}` has shape {values.shape},'
f' while `{property}` needs {shapes}.')
if selection:
values = np.stack(
[_select_component(values, select, property, rules) for select in selection],
axis=-1)
data[stype].append(values)
if 'component' in data:
labels = selection or _DEFAULT_LABELS[values.shape[-1]]
data['component'].extend(labels * (values.size // len(labels)))
if 'species' in data:
data['species'].extend(np.repeat(structure.symbols, values.shape[-1]).tolist())
if flatten:
for stype in ['predicted', 'target']:
if data[stype]:
data[stype] = np.concatenate([np.ravel(values) for values in data[stype]])
df = DataFrame(data)
# In case of flatten, cast to float64 for compatibility
# with e.g. seaborn.
# Casting in this way breaks tensorial properties though,
# so skip it there.
if flatten:
df['target'] = df.target.astype('float64')
df['predicted'] = df.predicted.astype('float64')
return df
