import warnings
from os import makedirs
from pathlib import Path
from os.path import exists, join as join_path
from typing import NamedTuple
import numpy as np
from ase import Atoms
from sklearn.model_selection import KFold
from .io import write_nepfile, write_structures
[docs]
def setup_training(parameters: NamedTuple,
structures: list[Atoms] = None,
train_structures: list[Atoms] = None,
test_structures: list[Atoms] = None,
enforced_structures: list[int] = [],
rootdir: str = '.',
mode: str = 'kfold',
n_splits: int = None,
train_fraction: float = None,
seed: int = 42,
overwrite: bool = False,
) -> None:
"""Sets up the input files for training a NEP via the ``nep``
executable of the GPUMD package.
Parameters
----------
parameters
Dictionary containing the parameters to be set in the nep.in file.
See `here <https://gpumd.org/nep/input_parameters/index.html>`__
for an overview of these parameters.
structures
List of structures to be included. Required for modes ``'kfold'`` and
``'bagging'``, and must not be set when mode ``'fixed'`` is used.
train_structures
Pre-defined list of training structures. Only used (and required) when mode
``'fixed'`` is used.
test_structures
Pre-defined list of test structures. Only used (and required) when mode
``'fixed'`` is used.
enforced_structures
Structures that _must_ be included in the training set, provided in the form
of a list of indices that refer to the content of the ``structures`` parameter.
Must not be set when mode ``'fixed'`` is used.
rootdir
Root directory in which to create the input files.
mode
How the test-train split is performed. Options: ``'kfold'``, ``'bagging'``,
and ``'fixed'``. ``'fixed'`` bypasses the split logic entirely and writes
``train_structures``/``test_structures`` as given directly to a single
``nepmodel`` directory under ``rootdir``, rather than the
``nepmodel_full``/``nepmodel_split*`` directories written by ``'kfold'``/
``'bagging'``.
n_splits
Number of splits of the input structures in training and test sets that ought to be
performed. By default no split will be done and all input structures will be used
for training. Must not be set when mode ``'fixed'`` is used.
train_fraction
Fraction of structures to use for training when mode ``'bagging'`` is used.
Must not be set when mode ``'fixed'`` is used.
seed
Random number generator seed to be used. This ensures reproducability.
overwrite
If True overwrite the content of ``rootdir`` if it exists.
"""
if exists(rootdir) and not overwrite:
raise FileExistsError('Output directory exists.'
' Set overwrite=True in order to override this behavior.')
if mode == 'fixed':
if train_structures is None or test_structures is None:
raise ValueError('Both train_structures and test_structures must be'
" provided when mode='fixed'.")
if structures is not None:
raise ValueError("structures cannot be set when mode='fixed'.")
if n_splits is not None:
raise ValueError("n_splits cannot be set when mode='fixed'.")
if train_fraction is not None:
raise ValueError("train_fraction cannot be set when mode='fixed'.")
if enforced_structures:
raise ValueError("enforced_structures cannot be set when mode='fixed'.")
elif mode in ('kfold', 'bagging'):
if structures is None:
raise ValueError(f"structures must be provided when mode='{mode}'.")
if train_structures is not None or test_structures is not None:
raise ValueError('train_structures/test_structures cannot be set'
f" when mode='{mode}'.")
if n_splits is not None and (n_splits <= 0 or n_splits > len(structures)):
raise ValueError(f'n_splits ({n_splits}) must be positive and'
f' must not exceed {len(structures)}.')
if mode == 'kfold' and train_fraction is not None:
raise ValueError(f'train_fraction cannot be set when mode {mode} is used')
elif mode == 'bagging' and (train_fraction <= 0 or train_fraction > 1):
raise ValueError(f'train_fraction ({train_fraction}) must be in (0,1]')
rs = np.random.RandomState(seed)
_prepare_training(parameters, structures, enforced_structures,
rootdir, mode, n_splits, train_fraction, rs,
train_structures=train_structures, test_structures=test_structures)
def _prepare_training(parameters: NamedTuple,
structures: list[Atoms],
enforced_structures: list[int],
rootdir: str,
mode: str,
n_splits: int | None,
train_fraction: float | None,
rs: np.random.RandomState,
train_structures: list[Atoms] | None = None,
test_structures: list[Atoms] | None = None) -> None:
"""Prepares training and test sets and writes structural data as well as parameters files.
See docstring for `setup_training` for documentation of parameters.
"""
if mode == 'fixed':
overlap = set(id(s) for s in train_structures) & set(id(s) for s in test_structures)
if overlap:
warnings.warn(f'{len(overlap)} structure(s) appear in both train_structures'
' and test_structures.')
dirname = join_path(rootdir, 'nepmodel')
makedirs(dirname, exist_ok=True)
write_structures(join_path(dirname, 'train.xyz'), train_structures)
write_structures(join_path(dirname, 'test.xyz'), test_structures)
write_nepfile(parameters, dirname)
return
dirname = join_path(rootdir, 'nepmodel_full')
makedirs(dirname, exist_ok=True)
_write_structures(structures, dirname, list(set(range(len(structures)))), [0])
write_nepfile(parameters, dirname)
if n_splits is None:
return
n_structures = len(structures)
remaining_structures = list(set(range(n_structures)) - set(enforced_structures))
if mode == 'kfold':
kf = KFold(n_splits=n_splits, shuffle=True, random_state=rs)
for k, (train_indices, test_indices) in enumerate(kf.split(remaining_structures)):
# append enforced structures at the end of the training set
train_selection = [remaining_structures[x] for x in list(train_indices)]
test_selection = [remaining_structures[x] for x in list(test_indices)]
# sanity check: make sure there is no overlap between train and test
assert set(train_selection).intersection(set(test_selection)) == set(), \
'Train and test set should not overlap'
subdir = f'nepmodel_split{k+1}'
dirname = join_path(rootdir, subdir)
makedirs(dirname, exist_ok=True)
_write_structures(structures, dirname, train_selection, test_selection)
write_nepfile(parameters, dirname)
elif mode == 'bagging':
for k in range(n_splits):
train_selection = rs.choice(
remaining_structures,
size=int(train_fraction * n_structures) - len(enforced_structures),
replace=False)
# append enforced structures at the end of the training set
train_selection = list(train_selection)
train_selection.extend(enforced_structures)
# add the remaining structures to the test set
test_selection = list(set(range(n_structures)) - set(train_selection))
# sanity check: make sure there is no overlap between train and test
assert set(train_selection).intersection(set(test_selection)) == set(), \
'Train and test set should not overlap'
dirname = join_path(rootdir, f'nepmodel_split{k+1}')
makedirs(dirname, exist_ok=True)
_write_structures(structures, dirname, train_selection, test_selection)
write_nepfile(parameters, dirname)
else:
raise ValueError(f'Unknown value for mode: {mode}.')
def _write_structures(structures: list[Atoms],
dirname: str,
train_selection: list[int],
test_selection: list[int]):
"""Writes structures in format readable by nep executable.
See docstring for `setup_training` for documentation of parameters.
"""
write_structures(
join_path(dirname, 'train.xyz'),
[s for k, s in enumerate(structures) if k in train_selection])
write_structures(
join_path(dirname, 'test.xyz'),
[s for k, s in enumerate(structures) if k in test_selection])
[docs]
def setup_fine_tuning_nep89(parameters: NamedTuple,
nep: Path,
restart: Path,
**kwargs_to_setup_training) -> None:
"""
Sets up a fine-tuning of the NEP89 foundation model.
Note that only the types, the number of generations, the batch,
the population, and the regularization parameters are allowed
to be changed.
The types must be a subset of the 89 types atomic species supported by
the NEP89 foundation model.
This function wraps :func:`setup_training`.
Parameters
----------
parameters
Dictionary containing the parameters to be set in the `nep.in` file;
see `here <https://gpumd.org/nep/input_parameters/index.html>`__
for an overview of these parameters.
Note that only `lambda_1`, `lambda_2`, `lambda_e`, `lambda_f`, `lambda_v`,
`generation`, `population`, `type`, and `batch` are allowed parameters when fine-tuning.
nep:
Path to the `nep.txt` file for NEP89.
restart:
Path to the `nep.restart` file for NEP89.
kwargs_to_setup_training:
See the dosctring for `setup_training` for the rest of the parameters.
"""
# Default parameters that need to be set for NEP89.
nep89_parameters = dict(version=4,
zbl=2,
cutoff=[6, 5],
n_max=[4, 4],
basis_size=[8, 8],
l_max=[4, 2, 1],
neuron=80)
for param in parameters.keys():
if param in nep89_parameters.keys():
raise ValueError(f'Parameter {param} not allowed when fine-tuning.')
if not Path(nep).is_file():
raise FileNotFoundError(f'{nep} does not exist.')
if not Path(restart).is_file():
raise FileNotFoundError(f'{restart} does not exist.')
# wrap nep89 and restart paths such that they match the subfolders written
# by setup_training
rootdir = kwargs_to_setup_training['rootdir']
if rootdir is None:
raise ValueError('The keyword `rootdir` must be set for setup_training.')
directory = Path(f'{rootdir}/nepmodel_full')
fine_tune_paths = [str(Path(file).relative_to(directory, walk_up=True))
for file in [nep, restart]]
fine_tuning = (dict(fine_tune=fine_tune_paths) | parameters | nep89_parameters)
setup_training(fine_tuning, **kwargs_to_setup_training)
