Coverage for calorine/nep/training_factory.py: 100%
97 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 12:50 +0000
1import warnings
2from os import makedirs
3from pathlib import Path
4from os.path import exists, join as join_path, relpath
5from typing import NamedTuple
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: NamedTuple,
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'``.
57 n_splits
58 Number of splits of the input structures in training and test sets that ought to be
59 performed. By default no split will be done and all input structures will be used
60 for training. Must not be set when mode ``'fixed'`` is used.
61 train_fraction
62 Fraction of structures to use for training when mode ``'bagging'`` is used.
63 Must not be set when mode ``'fixed'`` is used.
64 seed
65 Random number generator seed to be used. This ensures reproducability.
66 overwrite
67 If True overwrite the content of ``rootdir`` if it exists.
68 """
69 if exists(rootdir) and not overwrite:
70 raise FileExistsError('Output directory exists.'
71 ' Set overwrite=True in order to override this behavior.')
73 if mode == 'fixed':
74 if train_structures is None or test_structures is None:
75 raise ValueError('Both train_structures and test_structures must be'
76 " provided when mode='fixed'.")
77 if structures is not None:
78 raise ValueError("structures cannot be set when mode='fixed'.")
79 if n_splits is not None:
80 raise ValueError("n_splits cannot be set when mode='fixed'.")
81 if train_fraction is not None:
82 raise ValueError("train_fraction cannot be set when mode='fixed'.")
83 if enforced_structures:
84 raise ValueError("enforced_structures cannot be set when mode='fixed'.")
85 elif mode in ('kfold', 'bagging'):
86 if structures is None:
87 raise ValueError(f"structures must be provided when mode='{mode}'.")
88 if train_structures is not None or test_structures is not None:
89 raise ValueError('train_structures/test_structures cannot be set'
90 f" when mode='{mode}'.")
92 if n_splits is not None and (n_splits <= 0 or n_splits > len(structures)):
93 raise ValueError(f'n_splits ({n_splits}) must be positive and'
94 f' must not exceed {len(structures)}.')
96 if mode == 'kfold' and train_fraction is not None:
97 raise ValueError(f'train_fraction cannot be set when mode {mode} is used')
98 elif mode == 'bagging' and (train_fraction <= 0 or train_fraction > 1):
99 raise ValueError(f'train_fraction ({train_fraction}) must be in (0,1]')
101 rs = np.random.RandomState(seed)
102 _prepare_training(parameters, structures, enforced_structures,
103 rootdir, mode, n_splits, train_fraction, rs,
104 train_structures=train_structures, test_structures=test_structures)
107def _prepare_training(parameters: NamedTuple,
108 structures: list[Atoms],
109 enforced_structures: list[int],
110 rootdir: str,
111 mode: str,
112 n_splits: int | None,
113 train_fraction: float | None,
114 rs: np.random.RandomState,
115 train_structures: list[Atoms] | None = None,
116 test_structures: list[Atoms] | None = None) -> None:
117 """Prepares training and test sets and writes structural data as well as parameters files.
119 See docstring for `setup_training` for documentation of parameters.
120 """
121 if mode == 'fixed':
122 overlap = set(id(s) for s in train_structures) & set(id(s) for s in test_structures)
123 if overlap:
124 warnings.warn(f'{len(overlap)} structure(s) appear in both train_structures'
125 ' and test_structures.')
127 dirname = join_path(rootdir, 'nepmodel')
128 makedirs(dirname, exist_ok=True)
129 write_structures(join_path(dirname, 'train.xyz'), train_structures)
130 write_structures(join_path(dirname, 'test.xyz'), test_structures)
131 write_nepfile(parameters, dirname)
132 return
134 dirname = join_path(rootdir, 'nepmodel_full')
135 makedirs(dirname, exist_ok=True)
136 _write_structures(structures, dirname, list(set(range(len(structures)))), [0])
137 write_nepfile(parameters, dirname)
139 if n_splits is None:
140 return
142 n_structures = len(structures)
143 remaining_structures = list(set(range(n_structures)) - set(enforced_structures))
145 if mode == 'kfold':
146 kf = KFold(n_splits=n_splits, shuffle=True, random_state=rs)
147 for k, (train_indices, test_indices) in enumerate(kf.split(remaining_structures)):
148 # append enforced structures at the end of the training set
149 train_selection = [remaining_structures[x] for x in list(train_indices)]
150 test_selection = [remaining_structures[x] for x in list(test_indices)]
152 # sanity check: make sure there is no overlap between train and test
153 assert set(train_selection).intersection(set(test_selection)) == set(), \
154 'Train and test set should not overlap'
156 subdir = f'nepmodel_split{k+1}'
157 dirname = join_path(rootdir, subdir)
158 makedirs(dirname, exist_ok=True)
159 _write_structures(structures, dirname, train_selection, test_selection)
160 write_nepfile(parameters, dirname)
162 elif mode == 'bagging':
163 for k in range(n_splits):
164 train_selection = rs.choice(
165 remaining_structures,
166 size=int(train_fraction * n_structures) - len(enforced_structures),
167 replace=False)
169 # append enforced structures at the end of the training set
170 train_selection = list(train_selection)
171 train_selection.extend(enforced_structures)
173 # add the remaining structures to the test set
174 test_selection = list(set(range(n_structures)) - set(train_selection))
176 # sanity check: make sure there is no overlap between train and test
177 assert set(train_selection).intersection(set(test_selection)) == set(), \
178 'Train and test set should not overlap'
180 dirname = join_path(rootdir, f'nepmodel_split{k+1}')
181 makedirs(dirname, exist_ok=True)
182 _write_structures(structures, dirname, train_selection, test_selection)
183 write_nepfile(parameters, dirname)
185 else:
186 raise ValueError(f'Unknown value for mode: {mode}.')
189def _write_structures(structures: list[Atoms],
190 dirname: str,
191 train_selection: list[int],
192 test_selection: list[int]):
193 """Writes structures in format readable by nep executable.
195 See docstring for `setup_training` for documentation of parameters.
196 """
197 write_structures(
198 join_path(dirname, 'train.xyz'),
199 [s for k, s in enumerate(structures) if k in train_selection])
200 write_structures(
201 join_path(dirname, 'test.xyz'),
202 [s for k, s in enumerate(structures) if k in test_selection])
205def setup_fine_tuning_nep89(parameters: NamedTuple,
206 nep: Path,
207 restart: Path,
208 **kwargs_to_setup_training) -> None:
209 """
210 Sets up a fine-tuning of the NEP89 foundation model.
212 Note that only the types, the number of generations, the batch,
213 the population, and the regularization parameters are allowed
214 to be changed.
216 The types must be a subset of the 89 types atomic species supported by
217 the NEP89 foundation model.
219 This function wraps :func:`setup_training`.
221 Parameters
222 ----------
223 parameters
224 Dictionary containing the parameters to be set in the `nep.in` file;
225 see `here <https://gpumd.org/nep/input_parameters/index.html>`__
226 for an overview of these parameters.
227 Note that only `lambda_1`, `lambda_2`, `lambda_e`, `lambda_f`, `lambda_v`,
228 `generation`, `population`, `type`, and `batch` are allowed parameters when fine-tuning.
229 nep:
230 Path to the `nep.txt` file for NEP89.
231 restart:
232 Path to the `nep.restart` file for NEP89.
233 kwargs_to_setup_training:
234 See the dosctring for `setup_training` for the rest of the parameters.
235 """
236 # Default parameters that need to be set for NEP89.
237 nep89_parameters = dict(version=4,
238 zbl=2,
239 cutoff=[6, 5],
240 n_max=[4, 4],
241 basis_size=[8, 8],
242 l_max=[4, 2, 1],
243 neuron=80)
245 for param in parameters.keys():
246 if param in nep89_parameters.keys():
247 raise ValueError(f'Parameter {param} not allowed when fine-tuning.')
249 if not Path(nep).is_file():
250 raise FileNotFoundError(f'{nep} does not exist.')
251 if not Path(restart).is_file():
252 raise FileNotFoundError(f'{restart} does not exist.')
254 # wrap nep89 and restart paths such that they match the subfolders written
255 # by setup_training
256 rootdir = kwargs_to_setup_training['rootdir']
257 if rootdir is None:
258 raise ValueError('The keyword `rootdir` must be set for setup_training.')
259 directory = Path(f'{rootdir}/nepmodel_full')
260 fine_tune_paths = [relpath(file, directory) for file in [nep, restart]]
262 fine_tuning = (dict(fine_tune=fine_tune_paths) | parameters | nep89_parameters)
263 setup_training(fine_tuning, **kwargs_to_setup_training)