Coverage for calorine/nep/training_factory.py: 100%
98 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-10-08 12:52 +0000
1import warnings
2from os import makedirs
3from pathlib import Path
4from os.path import exists, join as join_path, relpath
5from typing import Any
7import numpy as np
8from ase import Atoms
9from sklearn.model_selection import KFold
11from .io import write_nepfile, write_structures
14def setup_training(parameters: dict[str, Any],
15 structures: list[Atoms] = None,
16 train_structures: list[Atoms] = None,
17 test_structures: list[Atoms] = None,
18 enforced_structures: list[int] = [],
19 rootdir: str = '.',
20 mode: str = 'kfold',
21 n_splits: int = None,
22 train_fraction: float = None,
23 seed: int = 42,
24 overwrite: bool = False,
25 ) -> None:
26 """Sets up the input files for training a NEP via the ``nep``
27 executable of the GPUMD package.
29 Parameters
30 ----------
31 parameters
32 Dictionary containing the parameters to be set in the nep.in file.
33 See `here <https://gpumd.org/nep/input_parameters/index.html>`__
34 for an overview of these parameters.
35 structures
36 List of structures to be included. Required for modes ``'kfold'`` and
37 ``'bagging'``, and must not be set when mode ``'fixed'`` is used.
38 train_structures
39 Pre-defined list of training structures. Only used (and required) when mode
40 ``'fixed'`` is used.
41 test_structures
42 Pre-defined list of test structures. Only used (and required) when mode
43 ``'fixed'`` is used.
44 enforced_structures
45 Structures that _must_ be included in the training set, provided in the form
46 of a list of indices that refer to the content of the ``structures`` parameter.
47 Must not be set when mode ``'fixed'`` is used.
48 rootdir
49 Root directory in which to create the input files.
50 mode
51 How the test-train split is performed. Options: ``'kfold'``, ``'bagging'``,
52 and ``'fixed'``. ``'fixed'`` bypasses the split logic entirely and writes
53 ``train_structures``/``test_structures`` as given directly to a single
54 ``nepmodel`` directory under ``rootdir``, rather than the
55 ``nepmodel_full``/``nepmodel_split*`` directories written by ``'kfold'``/
56 ``'bagging'``. ``nepmodel_full`` holds every structure in ``train.xyz`` and
57 receives no ``test.xyz``.
58 n_splits
59 Number of splits of the input structures in training and test sets that ought to be
60 performed. By default no split will be done and all input structures will be used
61 for training. Must not be set when mode ``'fixed'`` is used.
62 train_fraction
63 Fraction of structures to use for training when mode ``'bagging'`` is used.
64 Must not be set when mode ``'fixed'`` is used.
65 seed
66 Random number generator seed to be used. This ensures reproducability.
67 overwrite
68 If True overwrite the content of ``rootdir`` if it exists.
69 """
70 if exists(rootdir) and not overwrite:
71 raise FileExistsError('Output directory exists.'
72 ' Set overwrite=True in order to override this behavior.')
74 if mode == 'fixed':
75 if train_structures is None or test_structures is None:
76 raise ValueError('Both train_structures and test_structures must be'
77 " provided when mode='fixed'.")
78 if structures is not None:
79 raise ValueError("structures cannot be set when mode='fixed'.")
80 if n_splits is not None:
81 raise ValueError("n_splits cannot be set when mode='fixed'.")
82 if train_fraction is not None:
83 raise ValueError("train_fraction cannot be set when mode='fixed'.")
84 if enforced_structures:
85 raise ValueError("enforced_structures cannot be set when mode='fixed'.")
86 elif mode in ('kfold', 'bagging'):
87 if structures is None:
88 raise ValueError(f"structures must be provided when mode='{mode}'.")
89 if train_structures is not None or test_structures is not None:
90 raise ValueError('train_structures/test_structures cannot be set'
91 f" when mode='{mode}'.")
93 if n_splits is not None and (n_splits <= 0 or n_splits > len(structures)):
94 raise ValueError(f'n_splits ({n_splits}) must be positive and'
95 f' must not exceed {len(structures)}.')
97 if mode == 'kfold' and train_fraction is not None:
98 raise ValueError(f'train_fraction cannot be set when mode {mode} is used')
99 elif mode == 'bagging' and (train_fraction <= 0 or train_fraction > 1):
100 raise ValueError(f'train_fraction ({train_fraction}) must be in (0,1]')
102 rs = np.random.RandomState(seed)
103 _prepare_training(parameters, structures, enforced_structures,
104 rootdir, mode, n_splits, train_fraction, rs,
105 train_structures=train_structures, test_structures=test_structures)
108def _prepare_training(parameters: dict[str, Any],
109 structures: list[Atoms],
110 enforced_structures: list[int],
111 rootdir: str,
112 mode: str,
113 n_splits: int | None,
114 train_fraction: float | None,
115 rs: np.random.RandomState,
116 train_structures: list[Atoms] | None = None,
117 test_structures: list[Atoms] | None = None) -> None:
118 """Prepares training and test sets and writes structural data as well as parameters files.
120 See docstring for `setup_training` for documentation of parameters.
121 """
122 if mode == 'fixed':
123 overlap = set(id(s) for s in train_structures) & set(id(s) for s in test_structures)
124 if overlap:
125 warnings.warn(f'{len(overlap)} structure(s) appear in both train_structures'
126 ' and test_structures.')
128 dirname = join_path(rootdir, 'nepmodel')
129 makedirs(dirname, exist_ok=True)
130 write_structures(join_path(dirname, 'train.xyz'), train_structures)
131 write_structures(join_path(dirname, 'test.xyz'), test_structures)
132 write_nepfile(join_path(dirname, 'nep.in'), parameters)
133 return
135 dirname = join_path(rootdir, 'nepmodel_full')
136 makedirs(dirname, exist_ok=True)
138 # Only write train.xyz for full model
139 _write_structures(structures, dirname, list(set(range(len(structures)))), None)
140 write_nepfile(join_path(dirname, 'nep.in'), parameters)
142 if n_splits is None:
143 return
145 n_structures = len(structures)
146 remaining_structures = list(set(range(n_structures)) - set(enforced_structures))
148 if mode == 'kfold':
149 kf = KFold(n_splits=n_splits, shuffle=True, random_state=rs)
150 for k, (train_indices, test_indices) in enumerate(kf.split(remaining_structures)):
151 # append enforced structures at the end of the training set
152 train_selection = [remaining_structures[x] for x in list(train_indices)]
153 test_selection = [remaining_structures[x] for x in list(test_indices)]
155 # sanity check: make sure there is no overlap between train and test
156 assert set(train_selection).intersection(set(test_selection)) == set(), \
157 'Train and test set should not overlap'
159 subdir = f'nepmodel_split{k+1}'
160 dirname = join_path(rootdir, subdir)
161 makedirs(dirname, exist_ok=True)
162 _write_structures(structures, dirname, train_selection, test_selection)
163 write_nepfile(join_path(dirname, 'nep.in'), parameters)
165 elif mode == 'bagging':
166 for k in range(n_splits):
167 train_selection = rs.choice(
168 remaining_structures,
169 size=int(train_fraction * n_structures) - len(enforced_structures),
170 replace=False)
172 # append enforced structures at the end of the training set
173 train_selection = list(train_selection)
174 train_selection.extend(enforced_structures)
176 # add the remaining structures to the test set
177 test_selection = list(set(range(n_structures)) - set(train_selection))
179 # sanity check: make sure there is no overlap between train and test
180 assert set(train_selection).intersection(set(test_selection)) == set(), \
181 'Train and test set should not overlap'
183 dirname = join_path(rootdir, f'nepmodel_split{k+1}')
184 makedirs(dirname, exist_ok=True)
185 _write_structures(structures, dirname, train_selection, test_selection)
186 write_nepfile(join_path(dirname, 'nep.in'), parameters)
188 else:
189 raise ValueError(f'Unknown value for mode: {mode}.')
192def _write_structures(structures: list[Atoms],
193 dirname: str,
194 train_selection: list[int],
195 test_selection: list[int] | None):
196 """Writes structures in format readable by nep executable.
198 See docstring for `setup_training` for documentation of parameters.
199 """
200 write_structures(
201 join_path(dirname, 'train.xyz'),
202 [s for k, s in enumerate(structures) if k in train_selection])
203 if test_selection is not None:
204 write_structures(
205 join_path(dirname, 'test.xyz'),
206 [s for k, s in enumerate(structures) if k in test_selection])
209def setup_fine_tuning_nep89(parameters: dict[str, Any],
210 nep: Path,
211 restart: Path,
212 **kwargs_to_setup_training) -> None:
213 """
214 Sets up a fine-tuning of the NEP89 foundation model.
216 Note that only the types, the number of generations, the batch,
217 the population, and the regularization parameters are allowed
218 to be changed.
220 The types must be a subset of the 89 types atomic species supported by
221 the NEP89 foundation model.
223 This function wraps :func:`setup_training`.
225 Parameters
226 ----------
227 parameters
228 Dictionary containing the parameters to be set in the `nep.in` file;
229 see `here <https://gpumd.org/nep/input_parameters/index.html>`__
230 for an overview of these parameters.
231 Note that only `lambda_1`, `lambda_2`, `lambda_e`, `lambda_f`, `lambda_v`,
232 `generation`, `population`, `type`, and `batch` are allowed parameters when fine-tuning.
233 nep:
234 Path to the `nep.txt` file for NEP89.
235 restart:
236 Path to the `nep.restart` file for NEP89.
237 kwargs_to_setup_training:
238 See the dosctring for `setup_training` for the rest of the parameters.
239 """
240 # Default parameters that need to be set for NEP89.
241 nep89_parameters = dict(version=4,
242 zbl=2,
243 cutoff=[6, 5],
244 n_max=[4, 4],
245 basis_size=[8, 8],
246 l_max=[4, 2, 1],
247 neuron=80)
249 for param in parameters.keys():
250 if param in nep89_parameters.keys():
251 raise ValueError(f'Parameter {param} not allowed when fine-tuning.')
253 if not Path(nep).is_file():
254 raise FileNotFoundError(f'{nep} does not exist.')
255 if not Path(restart).is_file():
256 raise FileNotFoundError(f'{restart} does not exist.')
258 # wrap nep89 and restart paths such that they match the subfolders written
259 # by setup_training
260 rootdir = kwargs_to_setup_training['rootdir']
261 if rootdir is None:
262 raise ValueError('The keyword `rootdir` must be set for setup_training.')
263 directory = Path(f'{rootdir}/nepmodel_full')
264 fine_tune_paths = [relpath(file, directory) for file in [nep, restart]]
266 fine_tuning = (dict(fine_tune=fine_tune_paths) | parameters | nep89_parameters)
267 setup_training(fine_tuning, **kwargs_to_setup_training)