Coverage for calorine/nep/model.py: 99%
1154 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 12:52 +0000
1import copy
2import re
3from dataclasses import dataclass
4from itertools import product
5from typing import Any, Iterable
6from warnings import warn
8import numpy as np
10from calorine.nep.io import _write_nepfile_to_path
12NetworkWeights = dict[str, dict[str, np.ndarray]]
13DescriptorWeights = dict[tuple[str, str], np.ndarray]
14RestartParameters = dict[str, dict[str, dict[str, np.ndarray]]]
16# the `model_type` value that each model type corresponds to in `nep.in`
17_MODEL_TYPE_TO_INT = {
18 'potential': 0,
19 'potential_with_charges': 0,
20 'dipole': 1,
21 'polarizability': 2,
22}
24# The `nep.in` keywords that describe the model rather than the training run. These are the
25# keywords that the `nep` executable checks against the `nep.txt` header (see
26# `Parameters::compare_with_nep_txt` in `src/main_nep/parameters.cu` of GPUMD), plus `mode`,
27# which it accepts as a synonym of `model_type`. `Model.write_nepfile` supplies all of them
28# from the model and discards any value given for them, so that a value belonging to another
29# model cannot reach the file.
30_MODEL_PARAMETERS = (
31 'version',
32 'model_type',
33 'mode',
34 'type',
35 'cutoff',
36 'n_max',
37 'basis_size',
38 'l_max',
39 'neuron',
40 'zbl',
41 'use_typewise_cutoff_zbl',
42 'charge_mode',
43)
45# the factor that the `nep` executable uses when `use_typewise_cutoff_zbl` carries no value
46_TYPEWISE_CUTOFF_ZBL_FACTOR_DEFAULT = 0.7
49def _nepfile_tokens(value: Any) -> list[str]:
50 """Returns the value of a ``nep.in`` keyword as the list of whitespace separated tokens
51 that it is written as, so that values read from a file can be compared with values taken
52 from a model irrespective of how either is represented in Python.
54 Parameters
55 ----------
56 value
57 Value of a single ``nep.in`` keyword.
58 """
59 if isinstance(value, str):
60 return value.split()
61 if isinstance(value, Iterable):
62 return [f'{v}' for v in value]
63 return [f'{value}']
66def _same_nepfile_value(supplied: Any, model_value: Any) -> bool:
67 """Returns whether two values of a ``nep.in`` keyword agree, comparing token by token and
68 numerically where both tokens are numbers, so that e.g. ``6`` and ``6.0`` agree.
70 Parameters
71 ----------
72 supplied
73 Value to check, typically read from an existing ``nep.in`` file.
74 model_value
75 Value taken from the model, or ``None`` if the model does not set this keyword.
76 """
77 if model_value is None:
78 return False
79 left = _nepfile_tokens(supplied)
80 right = _nepfile_tokens(model_value)
81 if len(left) != len(right): 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true
82 return False
83 for token_left, token_right in zip(left, right):
84 if token_left == token_right:
85 continue
86 try:
87 value_left, value_right = float(token_left), float(token_right)
88 except ValueError:
89 return False
90 if abs(value_left - value_right) > 1e-6 * (abs(value_left) + abs(value_right)):
91 return False
92 return True
95def _format_nepfile_value(value: Any) -> str:
96 """Returns the value of a ``nep.in`` keyword as it is written to file, for use in messages.
98 Parameters
99 ----------
100 value
101 Value of a single ``nep.in`` keyword.
102 """
103 return ' '.join(_nepfile_tokens(value))
106def _get_restart_contents(filename: str) -> tuple[list[float], list[float]]:
107 """Parses a ``nep.restart`` file, and returns an unformatted list of the
108 mean and standard deviation for all model parameters.
109 Intended to be used by the py:meth:`~Model.read_restart` function.
111 Parameters
112 ----------
113 filename
114 input file name
115 """
116 mu = [] # Mean
117 sigma = [] # Standard deviation
118 with open(filename) as f:
119 for k, line in enumerate(f.readlines()):
120 flds = line.split()
121 if len(flds) == 0:
122 raise IOError(f'Empty line number {k}')
123 if len(flds) == 2:
124 mu.append(float(flds[0]))
125 sigma.append(float(flds[1]))
126 else:
127 raise IOError(f'Failed to parse line {k} from {filename}')
128 return mu, sigma
131def _get_model_type(first_row: list[str]) -> str:
132 """Parses a the first row of a ``nep.txt`` file, and returns the
133 type of NEP model. Available types are `potential`, `potential_with_charges`,
134 `dipole`, and `polarizability`.
136 Parameters
137 ----------
138 first_row
139 First row of a NEP file, split by white space.
140 """
141 model_type = first_row[0]
142 if 'charge' in model_type:
143 return 'potential_with_charges'
144 elif 'dipole' in model_type:
145 return 'dipole'
146 elif 'polarizability' in model_type:
147 return 'polarizability'
148 return 'potential'
151def _get_charge_mode(model_type_token: str) -> int:
152 """Parses the charge_mode (0, 1, or 2) from the first token of a ``nep.txt``
153 header line, e.g. ``nep4_charge1``, ``nep4_zbl_charge2``. Returns 0 for
154 non-charge models.
156 Parameters
157 ----------
158 model_type_token
159 First token of the first row of a NEP file (``flds[0]``).
160 """
161 match = re.search(r'charge(\d+)', model_type_token)
162 return int(match.group(1)) if match else 0
165def _get_nep_contents(filename: str) -> tuple[dict, list[float]]:
166 """Parses a ``nep.txt`` file, and returns a dict describing the header
167 and an unformatted list of all model parameters.
168 Intended to be used by the :func:`read_model <calorine.nep.read_model>` function.
170 Parameters
171 ----------
172 filename
173 input file name
174 """
175 # parse file and split header and parameters
176 header = []
177 parameters = []
178 nheader = 5 # 5 rows for NEP2, 6-7 rows for NEP3 onwards
179 base_line = 3
180 with open(filename) as f:
181 for k, line in enumerate(f.readlines()):
182 flds = line.split()
183 if len(flds) == 0:
184 raise IOError(f'Empty line number {k}')
185 if k == 0 and 'zbl' in flds[0]:
186 base_line += 1
187 nheader += 1
188 if k == base_line and 'basis_size' in flds[0]:
189 # Introduced in nep.txt after GPUMD v3.2
190 nheader += 1
191 if k < nheader:
192 header.append(tuple(flds))
193 elif len(flds) == 1:
194 parameters.append(float(flds[0]))
195 else:
196 raise IOError(f'Failed to parse line {k} from {filename}')
197 # compile data from the header into a dict
198 data = {}
199 for flds in header:
200 if flds[0] in ['cutoff', 'zbl']:
201 data[flds[0]] = tuple(map(float, flds[1:]))
202 elif flds[0] in ['n_max', 'l_max', 'ANN', 'basis_size']:
203 data[flds[0]] = tuple(map(int, flds[1:]))
204 elif flds[0].startswith('nep'):
205 version = flds[0].replace('nep', '').split('_')[0]
206 version = int(version)
207 data['version'] = version
208 data['types'] = flds[2:]
209 data['model_type'] = _get_model_type(flds)
210 data['charge_mode'] = _get_charge_mode(flds[0])
211 else:
212 raise ValueError(f'Unknown field: {flds[0]}')
213 return data, parameters
216def _sort_descriptor_parameters(parameters: list[float],
217 types: list[str],
218 n_max_radial: int,
219 n_basis_radial: int,
220 n_max_angular: int,
221 n_basis_angular: int) -> tuple[DescriptorWeights,
222 DescriptorWeights]:
223 """Reads a list of descriptors parameters and sorts them into two
224 appropriately structured `dicts`, one for radial and one for angular descriptor weights.
225 Intended to be used by the :func:`read_model <calorine.nep.read_model>` function.
226 """
227 # split up descriptor by chemical species and radial/angular
228 n_types = len(types)
229 n = len(parameters) // (n_types * n_types)
231 m = (n_max_radial + 1) * (n_basis_radial + 1)
232 descriptor_weights = parameters.reshape((n, n_types * n_types)).T
233 descriptor_weights_radial = descriptor_weights[:, :m]
234 descriptor_weights_angular = descriptor_weights[:, m:]
236 # add descriptors to data dict
237 radial_descriptor_weights = {}
238 angular_descriptor_weights = {}
239 m = -1
240 for i, j in product(range(n_types), repeat=2):
241 m += 1
242 s1, s2 = types[i], types[j]
243 radial_descriptor_weights[(s1, s2)] = descriptor_weights_radial[m, :].reshape(
244 (n_max_radial + 1, n_basis_radial + 1)
245 )
246 angular_descriptor_weights[(s1, s2)] = descriptor_weights_angular[m, :].reshape(
247 (n_max_angular + 1, n_basis_angular + 1)
248 )
249 return radial_descriptor_weights, angular_descriptor_weights
252def _number_of_output_biases(version: int,
253 n_types: int,
254 is_model_with_charges: bool) -> int:
255 """Returns the number of output-layer bias values in a single network pass.
257 This is the trailing block of a network pass in ``nep.txt`` and ``nep.restart``,
258 counted per pass rather than per file. A polarizability model runs two passes and
259 therefore carries twice this many bias values in total, while every other model
260 type carries exactly this many.
262 Parameters
263 ----------
264 version
265 NEP version (3, 4, or 5).
266 n_types
267 Number of atomic species in the model.
268 is_model_with_charges
269 Whether the model has a charge output head, i.e. whether its type is
270 ``potential_with_charges``.
272 Returns
273 -------
274 int
275 Number of bias values per network pass: 2 for a model with charges, since
276 ``sqrt_epsilon_infinity`` precedes the global bias, ``1 + n_types`` for NEP5,
277 which adds one bias per species to the global one, and 1 otherwise.
279 Example
280 -------
281 >>> from calorine.nep.model import _number_of_output_biases
282 >>> _number_of_output_biases(4, 2, False)
283 1
284 >>> _number_of_output_biases(5, 2, False)
285 3
286 """
287 if is_model_with_charges:
288 return 2
289 if version == 5:
290 return 1 + n_types
291 return 1
294def _sort_ann_parameters(parameters: list[float],
295 ann_groupings: list[str],
296 n_neuron: int,
297 n_networks: int,
298 version: int,
299 n_descriptor: int,
300 is_polarizability_model: bool,
301 is_model_with_charges: bool
302 ) -> NetworkWeights:
303 """Reads a list of model parameters and sorts them into an appropriately structured `dict`.
304 Intended to be used by the :func:`read_model <calorine.nep.read_model>` function
305 and by :meth:`Model.read_restart <calorine.nep.Model.read_restart>`.
306 """
307 # ann_groupings holds one entry per species for NEP4 and NEP5, and a single shared
308 # entry for NEP3, which has no per-species biases for its length to bear on
309 n_bias = _number_of_output_biases(version, len(ann_groupings), is_model_with_charges)
310 n_ann_input_weights = (n_descriptor + 1) * n_neuron # weights + bias
311 n_ann_output_weights = 2*n_neuron if is_model_with_charges else n_neuron # only weights
312 n_ann_parameters = (
313 n_ann_input_weights + n_ann_output_weights
314 ) * n_networks + n_bias
316 # Group ANN parameters
317 pars = {}
318 n1 = 0
319 n_network_params = n_ann_input_weights + n_ann_output_weights # except last bias(es)
321 n_count = 2 if is_polarizability_model else 1
322 n_outputs = 2 if is_model_with_charges else 1
323 for count in range(n_count):
324 # if polarizability model, all parameters including bias are repeated
325 # need to offset n1 by +1 to handle bias
326 n1 += count
327 for s in ann_groupings:
328 # Get the parameters for the ANN; in the case of NEP4, there is effectively
329 # one network per atomic species.
330 ann_parameters = parameters[n1 : n1 + n_network_params]
331 ann_input_weights = ann_parameters[:n_ann_input_weights]
332 w0 = np.zeros((n_neuron, n_descriptor))
333 w0[...] = np.nan
334 b0 = np.zeros((n_neuron, 1))
335 b0[...] = np.nan
336 for n in range(n_neuron):
337 for nu in range(n_descriptor):
338 w0[n, nu] = ann_input_weights[n * n_descriptor + nu]
339 b0[:, 0] = ann_input_weights[n_neuron * n_descriptor :]
341 assert np.all(
342 w0.shape == (n_neuron, n_descriptor)
343 ), f'w0 has invalid shape for key {s}; please submit a bug report'
344 assert np.all(
345 b0.shape == (n_neuron, 1)
346 ), f'b0 has invalid shape for key {s}; please submit a bug report'
347 assert not np.any(
348 np.isnan(w0)
349 ), f'some weights in w0 are nan for key {s}; please submit a bug report'
350 assert not np.any(
351 np.isnan(b0)
352 ), f'some weights in b0 are nan for key {s}; please submit a bug report'
354 ann_output_weights = ann_parameters[
355 n_ann_input_weights : n_ann_input_weights + n_ann_output_weights
356 ]
357 w1 = np.zeros((1, n_neuron * n_outputs))
358 w1[0, :] = ann_output_weights[:]
359 assert np.all(
360 w1.shape == (1, n_neuron * n_outputs)
361 ), f'w1 has invalid shape for key {s}; please submit a bug report'
362 assert not np.any(
363 np.isnan(w1)
364 ), f'some weights in w1 are nan for key {s}; please submit a bug report'
366 if count == 0 and n_outputs == 1:
367 pars[s] = dict(w0=w0, b0=b0, w1=w1)
368 elif count == 0 and n_outputs == 2:
369 pars[s] = dict(w0=w0, b0=b0, w1=w1[0, :n_neuron], w1_charge=w1[0, n_neuron:])
370 else:
371 pars[s].update({'w0_polar': w0, 'b0_polar': b0, 'w1_polar': w1})
372 # Jump to bias
373 n1 += n_network_params
374 if version == 5 and not is_model_with_charges:
375 # NEP5 models additionally have one bias term per species, which is
376 # what _number_of_output_biases accounts for. Currently NEP5 only
377 # exists for potential models, but we'll keep it here in case it gets
378 # added down the line.
379 bias_label = 'b1' if count == 0 else 'b1_polar'
380 pars[s][bias_label] = parameters[n1]
381 n1 += 1
382 # For NEP3 and NEP4 we only have one bias.
383 # For NEP4 with charges we have two biases.
384 # For NEP5 we have one bias per species, and one global.
385 if count == 0 and n_outputs == 1:
386 pars['b1'] = parameters[n1]
387 elif count == 0 and n_outputs == 2:
388 pars['sqrt_epsilon_infinity'] = parameters[n1]
389 pars['b1'] = parameters[n1+1]
390 else:
391 pars['b1_polar'] = parameters[n1]
392 sum = 0
393 for s in pars.keys():
394 if s.startswith('b1') or s.startswith('sqrt'):
395 sum += 1
396 else:
397 sum += np.sum([np.array(p).size for p in pars[s].values()])
398 assert sum == n_ann_parameters * n_count, (
399 'Inconsistent number of parameters accounted for; please submit a bug report\n'
400 f'{sum} != {n_ann_parameters}'
401 )
402 return pars
405def _adaptive_sigma(mu_arr, sigma_factor: float, sigma_floor: float) -> np.ndarray:
406 """Return adaptive SNES sigma: ``max(sigma_floor, sigma_factor * |mu|)``."""
407 return np.maximum(sigma_floor, sigma_factor * np.abs(mu_arr))
410def _format_header_float(value: float) -> str:
411 """Format *value* the way GPUMD writes the float-valued header fields of ``nep.txt``, which
412 it does with ``%g`` (see ``write_nep_txt`` in ``src/main_nep/fitness.cu``). A cutoff of 6 is
413 therefore written as ``6`` rather than ``6.0``.
415 ``%g`` keeps six significant digits, so it is used only where it reproduces the value exactly.
416 Anything carrying more digits than GPUMD would have written is emitted in full rather than
417 truncated, which matters for a file that came from somewhere else.
418 """
419 text = f'{float(value):g}'
420 return text if float(text) == float(value) else repr(float(value))
423_RESTART_COMPONENTS = ('network_weights', 'descriptor', 'charge_head')
426def _restart_leaves(model, restart_params, component=None, species=None):
427 """Yield ``(mu, sigma_container, sigma_key)`` for every leaf entry of
428 *restart_params* that matches the requested *component*/*species* filters.
430 ``sigma_container[sigma_key]`` is either a numpy array or a scalar float;
431 together with ``mu`` (same shape/type) this is everything
432 :func:`_apply_sigma_strategy` needs to read and update one leaf.
434 *component* selects among ``'network_weights'`` (w0, b0, w1, and the global
435 b1 bias), ``'descriptor'`` (radial/angular descriptor weight pairs), and
436 ``'charge_head'`` (w1_charge and sqrt_epsilon_infinity). ``None`` means all
437 three. *species* restricts per-species entries (and descriptor pairs
438 involving that species) to the given species; global scalars (b1,
439 sqrt_epsilon_infinity) are only included when *species* is ``None``, since
440 they are not owned by a single species.
441 """
442 if component is None:
443 wanted = set(_RESTART_COMPONENTS)
444 else:
445 wanted = {component} if isinstance(component, str) else set(component)
446 unknown = wanted - set(_RESTART_COMPONENTS)
447 if unknown:
448 raise ValueError(
449 f'Unknown component(s) {sorted(unknown)}; expected any of '
450 f'{_RESTART_COMPONENTS}'
451 )
453 if species is None:
454 species_filter = None
455 else:
456 species_filter = {species} if isinstance(species, str) else set(species)
458 keys = model.types if model.version in (4, 5) else ['all_species']
459 ann_mu, ann_sigma = restart_params['ann_mu'], restart_params['ann_sigma']
461 if 'network_weights' in wanted:
462 for s in keys:
463 if species_filter is not None and s not in species_filter:
464 continue
465 # b1 appears here only for NEP5, which carries a per-species bias alongside
466 # the global one handled below. For every other version it is global only.
467 for pname in ('w0', 'b0', 'w1', 'b1', 'w0_polar', 'b0_polar', 'w1_polar', 'b1_polar'):
468 if pname in ann_mu[s]:
469 yield ann_mu[s][pname], ann_sigma[s], pname
470 if species_filter is None:
471 for pname in ('b1', 'b1_polar'):
472 if pname in ann_mu:
473 yield ann_mu[pname], ann_sigma, pname
475 if 'charge_head' in wanted:
476 for s in keys:
477 if species_filter is not None and s not in species_filter:
478 continue
479 if 'w1_charge' in ann_mu[s]:
480 yield ann_mu[s]['w1_charge'], ann_sigma[s], 'w1_charge'
481 if species_filter is None and 'sqrt_epsilon_infinity' in ann_mu:
482 yield ann_mu['sqrt_epsilon_infinity'], ann_sigma, 'sqrt_epsilon_infinity'
484 if 'descriptor' in wanted:
485 for desc_type in ('radial', 'angular'):
486 mu_dict = restart_params[f'{desc_type}_descriptor_mu']
487 sigma_dict = restart_params[f'{desc_type}_descriptor_sigma']
488 for pair, mu_val in mu_dict.items():
489 if species_filter is not None and not (species_filter & set(pair)):
490 continue
491 yield mu_val, sigma_dict, pair
494_CHARGE_HEAD_PARAMETERS = ('w1_charge', 'sqrt_epsilon_infinity')
497def _parameter_component(name: str) -> str:
498 """Return the restart component that the ANN parameter head *name* belongs to.
500 Mirrors the grouping that :func:`_restart_leaves` applies to a loaded restart, for code
501 that walks the parameters of the model instead.
502 """
503 base = name[:-len('_polar')] if name.endswith('_polar') else name
504 if base in _CHARGE_HEAD_PARAMETERS:
505 return 'charge_head'
506 return 'network_weights'
509def _leaf_get(container, key):
510 """Read a parameter leaf from *container*, which is either a dict or the model itself."""
511 return container[key] if isinstance(container, dict) else getattr(container, key)
514def _leaf_set(container, key, value):
515 """Write a parameter leaf to *container*, which is either a dict or the model itself."""
516 if isinstance(container, dict): 516 ↛ 519line 516 didn't jump to line 519 because the condition on line 516 was always true
517 container[key] = value
518 else:
519 setattr(container, key, value)
522def _model_parameter_leaves(model, component=None, species=None):
523 """Yield ``(container, key, species, name)`` for every parameter of *model*, so that the
524 leaf can be read via :func:`_leaf_get` and written via :func:`_leaf_set`.
526 This is the counterpart of :func:`_restart_leaves` for the parameters of the model
527 itself. *species* is the species that owns the leaf, or ``None`` for the global scalars
528 and the descriptor pairs, and *name* is the name of the parameter head as used in
529 ``restart_parameters``, so that both can be used to look up the matching restart entry.
531 The *component*/*species* filters work exactly as in :func:`_restart_leaves`: global
532 scalars are only included when *species* is ``None``, and a descriptor pair is included
533 when either of its two species is selected.
534 """
535 if component is None:
536 wanted = set(_RESTART_COMPONENTS)
537 else:
538 wanted = {component} if isinstance(component, str) else set(component)
539 unknown = wanted - set(_RESTART_COMPONENTS)
540 if unknown:
541 raise ValueError(
542 f'Unknown component(s) {sorted(unknown)}; expected any of '
543 f'{_RESTART_COMPONENTS}'
544 )
546 if species is None:
547 species_filter = None
548 else:
549 species_filter = {species} if isinstance(species, str) else set(species)
551 keys = model.types if model.version in (4, 5) else ['all_species']
553 for s in keys:
554 if species_filter is not None and s not in species_filter:
555 continue
556 for name in model.ann_parameters[s]:
557 if _parameter_component(name) in wanted:
558 yield model.ann_parameters[s], name, s, name
560 if species_filter is None:
561 for name in model.ann_parameters:
562 if name in keys:
563 continue # per-species dictionaries, handled above
564 if _parameter_component(name) in wanted:
565 yield model.ann_parameters, name, None, name
566 if model.sqrt_epsilon_infinity is not None and 'charge_head' in wanted:
567 yield model, 'sqrt_epsilon_infinity', None, 'sqrt_epsilon_infinity'
569 if 'descriptor' in wanted:
570 for descriptor_type in ('radial', 'angular'):
571 weights = getattr(model, f'{descriptor_type}_descriptor_weights')
572 for pair in weights:
573 if species_filter is not None and not (species_filter & set(pair)):
574 continue
575 yield weights, pair, None, f'{descriptor_type}_descriptor'
578def _apply_sigma_strategy(mu, sigma_container, sigma_key, strategy, target, rng, **kw):
579 """Update ``sigma_container[sigma_key]`` in place at the positions selected
580 by *target* (``'unset'`` -> NaN entries, ``'set'`` -> non-NaN entries,
581 ``'all'`` -> everything), using *mu* and *strategy* to compute new values.
582 """
583 sigma = sigma_container[sigma_key]
584 if np.isscalar(sigma) or isinstance(sigma, (float, int)):
585 current = float(sigma)
586 is_unset = np.isnan(current)
587 apply_here = (
588 target == 'all' or (target == 'unset' and is_unset)
589 or (target == 'set' and not is_unset)
590 )
591 if not apply_here:
592 return
593 mu_val = float(mu)
594 if strategy == 'constant':
595 new_val = kw['value']
596 elif strategy == 'scale_mu':
597 new_val = float(_adaptive_sigma(np.array(mu_val), kw['factor'], kw['floor']))
598 elif strategy == 'scale_sigma':
599 if is_unset: 599 ↛ 600line 599 didn't jump to line 600 because the condition on line 599 was never true
600 raise ValueError(
601 f"strategy='scale_sigma' requires an existing sigma value, but "
602 f'{sigma_key!r} is unset (NaN); set it first, e.g. with '
603 "target='unset'."
604 )
605 new_val = current * kw['factor']
606 elif strategy == 'uniform':
607 new_val = float(rng.uniform(kw['low'], kw['high']))
608 elif strategy == 'normal': 608 ↛ 611line 608 didn't jump to line 611 because the condition on line 608 was always true
609 new_val = float(abs(rng.normal(kw['mean'], kw['std'])))
610 else:
611 raise ValueError(f'Unknown strategy {strategy!r}')
612 sigma_container[sigma_key] = float(new_val)
613 return
615 sigma_arr = sigma_container[sigma_key]
616 mu_arr = np.asarray(mu, dtype=float)
617 if target == 'unset':
618 mask = np.isnan(sigma_arr)
619 elif target == 'set':
620 mask = ~np.isnan(sigma_arr)
621 else:
622 mask = np.ones_like(sigma_arr, dtype=bool)
623 if not np.any(mask):
624 return
626 if strategy == 'constant':
627 sigma_arr[mask] = kw['value']
628 elif strategy == 'scale_mu':
629 sigma_arr[mask] = _adaptive_sigma(mu_arr[mask], kw['factor'], kw['floor'])
630 elif strategy == 'scale_sigma':
631 if np.any(np.isnan(sigma_arr[mask])):
632 raise ValueError(
633 f"strategy='scale_sigma' requires existing sigma values, but "
634 f'{sigma_key!r} has unset (NaN) entries within the selected '
635 "target; set them first, e.g. with target='unset'."
636 )
637 sigma_arr[mask] = sigma_arr[mask] * kw['factor']
638 elif strategy == 'uniform':
639 sigma_arr[mask] = rng.uniform(kw['low'], kw['high'], size=int(np.sum(mask)))
640 elif strategy == 'normal': 640 ↛ 643line 640 didn't jump to line 643 because the condition on line 640 was always true
641 sigma_arr[mask] = np.abs(rng.normal(kw['mean'], kw['std'], size=int(np.sum(mask))))
642 else:
643 raise ValueError(f'Unknown strategy {strategy!r}')
646def _draw_values(strategy, rng, size=None, **kw):
647 """Return parameter values drawn according to *strategy*: a single float if *size* is
648 ``None``, otherwise an array of *size* values.
650 Unlike the sigma strategies of :func:`_apply_sigma_strategy`, the drawn values are not
651 forced to be positive, since a parameter value may have either sign.
652 """
653 if strategy == 'constant':
654 return float(kw['value']) if size is None else np.full(size, float(kw['value']))
655 if strategy == 'uniform':
656 return (float(rng.uniform(kw['low'], kw['high'])) if size is None
657 else rng.uniform(kw['low'], kw['high'], size=size))
658 if strategy == 'normal': 658 ↛ 661line 658 didn't jump to line 661 because the condition on line 658 was always true
659 return (float(rng.normal(kw['mean'], kw['std'])) if size is None
660 else rng.normal(kw['mean'], kw['std'], size=size))
661 raise ValueError(f'Unknown strategy {strategy!r}')
664def _restart_leaf_for(restart_parameters, species, name, key):
665 """Return ``(mu_container, sigma_container, restart_key)`` for the restart entry that
666 corresponds to the model parameter leaf described by *species*, *name*, and *key*, or
667 ``None`` if the restart has no counterpart for it.
668 """
669 if name.endswith('_descriptor'):
670 mu_container = restart_parameters[f'{name}_mu']
671 sigma_container = restart_parameters[f'{name}_sigma']
672 restart_key = key
673 elif species is None:
674 mu_container = restart_parameters['ann_mu']
675 sigma_container = restart_parameters['ann_sigma']
676 restart_key = name
677 else:
678 mu_container = restart_parameters['ann_mu'][species]
679 sigma_container = restart_parameters['ann_sigma'][species]
680 restart_key = name
681 if restart_key not in sigma_container or restart_key not in mu_container: 681 ↛ 686line 681 didn't jump to line 686 because the condition on line 681 was never true
682 # A parameter head that the restart tree in hand does not carry. Whether such a
683 # parameter is new cannot be told from a sigma, so it is left alone. Every head of
684 # every shipped model has a counterpart, so this guards against a tree assembled
685 # by other means rather than against any model type in particular.
686 return None
687 return mu_container, sigma_container, restart_key
690def _restart_sigma_counts(model, component=None) -> dict[str, int]:
691 """Return the number of restart sigma entries of *model* as a dict with the keys
692 ``'total'``, ``'frozen'`` (sigma of exactly zero) and ``'unset'`` (sigma of ``NaN``),
693 restricted to *component* if given.
695 The traversal goes through :func:`_restart_leaves`, so every parameter head
696 registered there is covered automatically. Scalar leaves (``b1``,
697 ``sqrt_epsilon_infinity``) count as one entry each.
698 """
699 total, frozen, unset = 0, 0, 0
700 for _, sigma_container, sigma_key in _restart_leaves(
701 model, model.restart_parameters, component
702 ):
703 sigma = np.asarray(sigma_container[sigma_key], dtype=float)
704 total += int(sigma.size)
705 frozen += int(np.count_nonzero(sigma == 0.0))
706 unset += int(np.count_nonzero(np.isnan(sigma)))
707 return {'total': total, 'frozen': frozen, 'unset': unset}
710def _model_parameter_counts(model) -> dict[str, dict[str, int]]:
711 """Return ``{component: {'total': n}}`` for the current parameters of *model*.
713 Used when no restart is loaded, so there are no sigma values to split into frozen and
714 unset entries. The traversal goes through :func:`_model_parameter_leaves`, which also
715 drives :meth:`Model.initialize_parameters`, so the two cannot disagree on which
716 component a parameter head belongs to. The ``q_scaler`` entries are excluded, since they
717 are not fit parameters and have no restart counterpart.
718 """
719 counts = {}
720 for component in _RESTART_COMPONENTS:
721 total = sum(int(np.asarray(_leaf_get(container, key)).size)
722 for container, key, _, _ in _model_parameter_leaves(model, component))
723 if total > 0:
724 counts[component] = {'total': total}
725 return counts
728def _new_restart_parameters_from_model(model) -> RestartParameters:
729 """Build a fresh restart-parameters dict from a model's current (trained)
730 parameters: ``mu`` is copied from the model, ``sigma`` is set to ``NaN``
731 everywhere (unset), to be filled in via :meth:`Model.set_restart_sigma`.
732 """
733 keys = model.types if model.version in (4, 5) else ['all_species']
734 suffixes = ['', '_polar'] if model.model_type == 'polarizability' else ['']
736 ann_mu, ann_sigma = {}, {}
737 for s in keys:
738 params = model.ann_parameters[s]
739 mu_entry, sigma_entry = {}, {}
740 for suffix in suffixes:
741 for base in ('w0', 'b0', 'w1', 'w1_charge'):
742 pname = f'{base}{suffix}'
743 if pname in params:
744 arr = np.array(params[pname], dtype=float)
745 mu_entry[pname] = arr.copy()
746 sigma_entry[pname] = np.full(arr.shape, np.nan)
747 # NEP5 adds a per-species bias to the global one, and it is a scalar
748 # rather than an array, matching how the model itself stores it.
749 bias_name = f'b1{suffix}'
750 if bias_name in params:
751 mu_entry[bias_name] = float(params[bias_name])
752 sigma_entry[bias_name] = float('nan')
753 ann_mu[s] = mu_entry
754 ann_sigma[s] = sigma_entry
756 for suffix in suffixes:
757 b1_key = f'b1{suffix}'
758 if b1_key in model.ann_parameters: 758 ↛ 756line 758 didn't jump to line 756 because the condition on line 758 was always true
759 ann_mu[b1_key] = float(model.ann_parameters[b1_key])
760 ann_sigma[b1_key] = float('nan')
761 if model.sqrt_epsilon_infinity is not None:
762 ann_mu['sqrt_epsilon_infinity'] = float(model.sqrt_epsilon_infinity)
763 ann_sigma['sqrt_epsilon_infinity'] = float('nan')
765 radial_mu = {
766 k: np.array(v, dtype=float).copy() for k, v in model.radial_descriptor_weights.items()
767 }
768 radial_sigma = {k: np.full(v.shape, np.nan) for k, v in radial_mu.items()}
769 angular_mu = {
770 k: np.array(v, dtype=float).copy() for k, v in model.angular_descriptor_weights.items()
771 }
772 angular_sigma = {k: np.full(v.shape, np.nan) for k, v in angular_mu.items()}
774 return {
775 'ann_mu': ann_mu,
776 'ann_sigma': ann_sigma,
777 'radial_descriptor_mu': radial_mu,
778 'radial_descriptor_sigma': radial_sigma,
779 'angular_descriptor_mu': angular_mu,
780 'angular_descriptor_sigma': angular_sigma,
781 }
784def _recalculate_parameter_counts(new) -> None:
785 """Recompute n_ann_parameters, n_descriptor_parameters, and n_parameters on *new*.
787 Reads all architectural state from *new* directly, so callers must update
788 new.n_neuron, new.n_descriptor_radial/angular, new.model_type, and new.types
789 before calling this function.
790 """
791 n_types = len(new.types)
792 n_desc = new.n_descriptor_radial + new.n_descriptor_angular
793 is_charged = new.model_type == 'potential_with_charges'
794 n_networks = n_types if new.version in (4, 5) else 1
795 n_output_biases = _number_of_output_biases(new.version, n_types, is_charged)
796 n_ann_input_weights = (n_desc + 1) * new.n_neuron
797 n_ann_output_weights = 2 * new.n_neuron if is_charged else new.n_neuron
798 new.n_ann_parameters = (
799 n_ann_input_weights + n_ann_output_weights
800 ) * n_networks + n_output_biases
801 new.n_descriptor_parameters = n_types ** 2 * (
802 (new.n_max_radial + 1) * (new.n_basis_radial + 1)
803 + (new.n_max_angular + 1) * (new.n_basis_angular + 1)
804 )
805 new.n_parameters = new.n_ann_parameters + new.n_descriptor_parameters + n_desc
806 if new.model_type == 'polarizability':
807 new.n_parameters += new.n_ann_parameters
810@dataclass
811class Model:
812 r"""Objects of this class represent a NEP model in a form suitable for
813 inspection and manipulation. Typically a :class:`Model` object is instantiated
814 by calling the :func:`read_model <calorine.nep.read_model>` function.
816 Attributes
817 ----------
818 version : int
819 NEP version.
820 model_type: str
821 One of ``potential``, ``dipole`` or ``polarizability``.
822 types : tuple[str, ...]
823 Chemical species that this model represents.
824 radial_cutoff : float | list[float]
825 The radial cutoff parameter in Å.
826 Is a list of radial cutoffs ordered after ``types`` in the case of typewise cutoffs.
827 angular_cutoff : float | list[float]
828 The angular cutoff parameter in Å.
829 Is a list of angular cutoffs ordered after ``types`` in the case of typewise cutoffs.
830 max_neighbors_radial : int
831 Maximum number of neighbors in neighbor list for radial terms.
832 max_neighbors_angular : int
833 Maximum number of neighbors in neighbor list for angular terms.
834 zbl : tuple[float, float]
835 Inner and outer cutoff for transition to ZBL potential.
836 zbl_typewise_cutoff_factor : float
837 Optional typewise cutoff factor for the ZBL potential, corresponding to an
838 optional third value on the ``zbl`` line in ``nep.txt`` when
839 ``use_typewise_cutoff_zbl`` is enabled during training. ``None`` if not set.
840 n_basis_radial : int
841 Number of radial basis functions :math:`n_\mathrm{basis}^\mathrm{R}`.
842 n_basis_angular : int
843 Number of angular basis functions :math:`n_\mathrm{basis}^\mathrm{A}`.
844 n_max_radial : int
845 Maximum order of Chebyshev polymonials included in
846 radial expansion :math:`n_\mathrm{max}^\mathrm{R}`.
847 n_max_angular : int
848 Maximum order of Chebyshev polymonials included in
849 angular expansion :math:`n_\mathrm{max}^\mathrm{A}`.
850 l_max_3b : int
851 Maximum expansion order for three-body terms :math:`l_\mathrm{max}^\mathrm{3b}`.
852 l_max_4b : int
853 Maximum expansion order for four-body terms :math:`l_\mathrm{max}^\mathrm{4b}`.
854 l_max_5b : int
855 Maximum expansion order for five-body terms :math:`l_\mathrm{max}^\mathrm{5b}`.
856 has_q_112 : int
857 Flag enabling the 5-body :math:`q_{112}` descriptor (0 or 1).
858 has_q_123 : int
859 Flag enabling the 5-body :math:`q_{123}` descriptor (0 or 1).
860 has_q_233 : int
861 Flag enabling the 5-body :math:`q_{233}` descriptor (0 or 1).
862 has_q_134 : int
863 Flag enabling the higher-body :math:`q_{134}` descriptor (0 or 1).
864 n_descriptor_radial : int
865 Dimension of radial part of descriptor.
866 n_descriptor_angular : int
867 Dimension of angular part of descriptor.
868 n_neuron : int
869 Number of neurons in hidden layer.
870 n_parameters : int
871 Total number of parameters including scalers (which are not fit parameters).
872 n_descriptor_parameters : int
873 Number of parameters in descriptor.
874 n_ann_parameters : int
875 Number of neural network weights.
876 ann_parameters : dict[tuple[str, dict[str, np.darray]]]
877 Neural network weights.
878 q_scaler : List[float]
879 Scaling parameters.
880 radial_descriptor_weights : dict[tuple[str, str], np.ndarray]
881 Radial descriptor weights by combination of species; the array for each combination
882 has dimensions of
883 :math:`(n_\mathrm{max}^\mathrm{R}+1) \times (n_\mathrm{basis}^\mathrm{R}+1)`.
884 angular_descriptor_weights : dict[tuple[str, str], np.ndarray]
885 Angular descriptor weights by combination of species; the array for each combination
886 has dimensions of
887 :math:`(n_\mathrm{max}^\mathrm{A}+1) \times (n_\mathrm{basis}^\mathrm{A}+1)`.
888 sqrt_epsilon_infinity : Optional[float]
889 Square root of epsilon infinity $\epsilon_\infty$ (only for NEP models with charges).
890 charge_mode : int
891 Charge algorithm variant for ``potential_with_charges`` models; 0 for
892 non-charge-aware models. 1 corresponds to a qNEP model including both real- and
893 reciprocal-space contributions. 2 corresponds to a qNEP model, including the
894 reciprocal-space contribution only.
895 restart_parameters : dict[str, dict[str, dict[str, np.ndarray]]]
896 NEP restart parameters. A nested dictionary that contains the mean (mu) and standard
897 deviation (sigma) for the ANN and descriptor parameters. Is set using the
898 py:meth:`~Model.read_restart` method. Defaults to None.
899 The state of the sigma values is summarized by the
900 :attr:`~Model.n_frozen_parameters` and :attr:`~Model.n_unset_parameters`
901 properties, and broken down per category by :attr:`~Model.parameter_counts`.
902 """
904 version: int
905 model_type: str
906 types: tuple[str, ...]
908 radial_cutoff: float | list[float]
909 angular_cutoff: float | list[float]
911 n_basis_radial: int
912 n_basis_angular: int
913 n_max_radial: int
914 n_max_angular: int
915 l_max_3b: int
916 l_max_4b: int
917 l_max_5b: int
918 has_q_112: int
919 has_q_123: int
920 has_q_233: int
921 has_q_134: int
922 n_descriptor_radial: int
923 n_descriptor_angular: int
925 n_neuron: int
926 n_parameters: int
927 n_descriptor_parameters: int
928 n_ann_parameters: int
929 ann_parameters: NetworkWeights
930 q_scaler: list[float]
931 radial_descriptor_weights: DescriptorWeights
932 angular_descriptor_weights: DescriptorWeights
933 sqrt_epsilon_infinity: float = None
934 charge_mode: int = 0
935 restart_parameters: RestartParameters = None
937 zbl: tuple[float, float] = None
938 zbl_typewise_cutoff_factor: float = None
939 max_neighbors_radial: int = None
940 max_neighbors_angular: int = None
942 _special_fields = [
943 'ann_parameters',
944 'q_scaler',
945 'radial_descriptor_weights',
946 'angular_descriptor_weights',
947 ]
949 def __str__(self) -> str:
950 s = []
951 for fld in self.__dataclass_fields__:
952 if fld not in self._special_fields:
953 value = getattr(self, fld)
954 if fld == 'restart_parameters':
955 value = self._restart_availability()
956 s += [f'{fld:22} : {value}']
957 return '\n'.join(s)
959 def _repr_html_(self) -> str:
960 s = []
961 s += ['<table border="1" class="dataframe"']
962 s += [
963 '<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>'
964 ]
965 s += ['<tbody>']
966 for fld in self.__dataclass_fields__:
967 if fld not in self._special_fields:
968 value = getattr(self, fld)
969 if fld == 'restart_parameters':
970 value = self._restart_availability()
971 s += [
972 f'<tr><td style="text-align: left;">{fld:22}</td>'
973 f'<td>{value}</td><tr>'
974 ]
975 for fld in self._special_fields:
976 d = getattr(self, fld)
977 # print('xxx', fld, d)
978 if fld.endswith('descriptor_weights'):
979 dim = list(d.values())[0].shape
980 elif fld == 'ann_parameters' and self.version == 4:
981 dim = (len(self.types), len(list(d.values())[0]))
982 else:
983 dim = len(d)
984 s += [
985 f'<tr><td style="text-align: left;">Dimension of {fld:22}</td><td>{dim}</td><tr>'
986 ]
987 s += ['</tbody>']
988 s += ['</table>']
989 return ''.join(s)
991 @property
992 def training_parameters(self) -> dict:
993 """Return the model parameters in the format accepted by :func:`write_nepfile
994 <calorine.nep.write_nepfile>`.
996 The result covers every ``nep.in`` keyword that describes the model itself, i.e. the
997 keywords that the ``nep`` executable checks against the ``nep.txt`` header before
998 training. It carries no training parameters (``lambda_*``, ``generation``, ``batch``,
999 and the like). Use :meth:`write_nepfile` to combine the two and write the file, rather
1000 than merging the dictionaries by hand.
1002 Returns
1003 -------
1004 dict
1005 Keys ``version``, ``model_type``, ``type``, ``cutoff``, ``n_max``, ``basis_size``,
1006 ``l_max`` and ``neuron``, plus ``zbl`` and ``use_typewise_cutoff_zbl`` for a model
1007 with ZBL repulsion, and ``charge_mode`` for a charge-aware model. ``zbl`` is the
1008 single outer cutoff value that the ``nep.in`` ``zbl`` keyword expects (the inner
1009 cutoff is always half of it), not the ``(inner, outer)`` pair stored in
1010 :attr:`zbl`.
1012 Raises
1013 ------
1014 ValueError
1015 If :attr:`model_type` is not one of the known model types.
1017 """
1018 l_max = [self.l_max_3b, self.l_max_4b, self.l_max_5b,
1019 self.has_q_112, self.has_q_123, self.has_q_233, self.has_q_134]
1020 while len(l_max) > 1 and l_max[-1] == 0:
1021 l_max = l_max[:-1]
1023 if isinstance(self.radial_cutoff, list):
1024 cutoff = []
1025 for rc, ac in zip(self.radial_cutoff, self.angular_cutoff):
1026 cutoff += [rc, ac]
1027 else:
1028 cutoff = [self.radial_cutoff, self.angular_cutoff]
1030 if self.model_type not in _MODEL_TYPE_TO_INT:
1031 raise ValueError(f'Unknown model_type: {self.model_type}')
1033 # `type` must precede `cutoff`, which the `nep` executable sizes by the number of
1034 # types, so the insertion order below is part of the contract.
1035 params = {
1036 'version': self.version,
1037 'model_type': _MODEL_TYPE_TO_INT[self.model_type],
1038 'type': [len(self.types)] + list(self.types),
1039 'cutoff': cutoff,
1040 'n_max': [self.n_max_radial, self.n_max_angular],
1041 'basis_size': [self.n_basis_radial, self.n_basis_angular],
1042 'l_max': l_max,
1043 'neuron': self.n_neuron,
1044 }
1045 if self.zbl is not None:
1046 zbl_inner, zbl_outer = self.zbl
1047 if zbl_inner == 0 and zbl_outer == 0:
1048 # GPUMD writes `zbl 0 0` for a flexible ZBL potential, which is requested by
1049 # placing a `zbl.in` file next to `nep.in` rather than through a keyword.
1050 warn('This model uses a flexible ZBL potential, which cannot be expressed in '
1051 'nep.in. Place the zbl.in file of the model next to nep.in and set the '
1052 'zbl cutoff by hand.')
1053 elif abs(zbl_inner - 0.5 * zbl_outer) > 1e-6 * abs(zbl_outer):
1054 warn(f'The ZBL cutoffs of this model, {zbl_inner} and {zbl_outer} Å, are not '
1055 'expressible in nep.in, where the inner cutoff is always half of the '
1056 f'outer one. Writing zbl {zbl_outer}, which implies an inner cutoff of '
1057 f'{0.5 * zbl_outer} Å.')
1058 params['zbl'] = zbl_outer
1059 if self.zbl_typewise_cutoff_factor is not None:
1060 params['use_typewise_cutoff_zbl'] = self.zbl_typewise_cutoff_factor
1061 if self.charge_mode != 0:
1062 params['charge_mode'] = self.charge_mode
1063 return params
1065 def write_nepfile(self, filename: str, parameters: dict = None) -> None:
1066 """Writes a ``nep.in`` file for this model.
1068 The keywords that describe the model are taken from the model itself, via
1069 :attr:`training_parameters`, so the resulting file is consistent with the ``nep.txt``
1070 file written by :meth:`write`. This is the intended way to prepare the input for
1071 training a model that :meth:`augment`, :meth:`add_species`, :meth:`remove_species`,
1072 :meth:`keep_species` or :meth:`prune` has changed the architecture of.
1074 Training parameters can be supplied via :attr:`parameters`, typically read from an
1075 existing ``nep.in`` file with :func:`read_nepfile <calorine.nep.read_nepfile>`. Any
1076 keyword in :attr:`parameters` that describes the model is discarded in favor of the
1077 value of the model, with a warning naming what was dropped, since such a value refers
1078 to whichever model that file was written for and not to this one.
1080 Note that unlike :func:`write_nepfile <calorine.nep.write_nepfile>`, which takes the
1081 name of a directory, this method takes the name of a file, as :meth:`write` and
1082 :meth:`write_restart` do.
1084 Parameters
1085 ----------
1086 filename
1087 Name of the file to write, conventionally ``nep.in``.
1088 parameters
1089 Training parameters to include, such as ``generation``, ``batch`` and
1090 ``lambda_e``. Keywords that describe the model are ignored.
1092 Raises
1093 ------
1094 ValueError
1095 If :attr:`model_type` is not one of the known model types.
1097 Example
1098 -------
1099 Add a species to a model and write the input files needed to continue training it::
1101 >>> from calorine.nep import read_model, read_nepfile
1102 >>> model = read_model('nep.txt', restart_file='nep.restart')
1103 >>> extended = model.add_species(['Cl'], seed=42)
1104 >>> parameters = read_nepfile('nep.in')
1105 >>> extended.write('new/nep.txt', restart_file='new/nep.restart')
1106 >>> extended.write_nepfile('new/nep.in', parameters)
1108 """
1109 model_parameters = self.training_parameters
1110 merged = dict(model_parameters)
1111 discarded = []
1112 for key, value in (parameters or {}).items():
1113 if key not in _MODEL_PARAMETERS:
1114 merged[key] = value
1115 continue
1116 model_value = model_parameters.get(key)
1117 if key in ('model_type', 'mode'):
1118 # the `nep` executable accepts `mode` as a synonym of `model_type`
1119 model_value = model_parameters['model_type']
1120 elif key == 'use_typewise_cutoff_zbl' and _nepfile_tokens(value) == []:
1121 # the bare keyword requests the default factor
1122 value = _TYPEWISE_CUTOFF_ZBL_FACTOR_DEFAULT
1123 elif key == 'charge_mode':
1124 # `charge_mode <mode> [flip_charge]`, where flip_charge configures the
1125 # training run rather than the model and is therefore kept
1126 tokens = _nepfile_tokens(value)
1127 if len(tokens) > 1 and model_value is not None:
1128 merged['charge_mode'] = [model_value] + tokens[1:]
1129 value = tokens[:1]
1130 if not _same_nepfile_value(value, model_value):
1131 discarded.append((key, value, model_value))
1132 if discarded:
1133 lines = []
1134 for key, value, model_value in discarded:
1135 supplied = _format_nepfile_value(value)
1136 if model_value is None:
1137 lines.append(f' {key}: {supplied} -> dropped, this model has no {key}')
1138 else:
1139 lines.append(f' {key}: {supplied} -> {_format_nepfile_value(model_value)}')
1140 warn('The following nep.in parameters were overridden by the model:\n'
1141 + '\n'.join(lines))
1142 _write_nepfile_to_path(merged, filename)
1144 @property
1145 def n_frozen_parameters(self) -> int | None:
1146 """Number of frozen restart parameters, i.e. parameters whose SNES sigma is
1147 exactly zero and which are therefore excluded from the search during a restart.
1149 Returns
1150 -------
1151 int or None
1152 Number of parameters with a sigma of zero, or ``None`` if
1153 ``restart_parameters`` is not loaded.
1155 Example
1156 -------
1157 Freeze everything that has already been trained and check the result::
1159 >>> model = read_model('nep4_PbTe.txt', restart_file='nep4_PbTe.restart')
1160 >>> frozen = model.set_restart_sigma(strategy='constant', value=0.0,
1161 ... target='set')
1162 >>> frozen.n_frozen_parameters
1163 2281
1164 """
1165 if self.restart_parameters is None:
1166 return None
1167 return _restart_sigma_counts(self)['frozen']
1169 @property
1170 def n_unset_parameters(self) -> int | None:
1171 """Number of restart parameters whose SNES sigma is unset (``NaN``), i.e.
1172 parameters created by :meth:`augment` or :meth:`add_species` that have not yet
1173 been given a search width by :meth:`set_restart_sigma`.
1175 :meth:`write_restart` refuses to write while this count is non-zero.
1177 Returns
1178 -------
1179 int or None
1180 Number of parameters with an unset sigma, or ``None`` if
1181 ``restart_parameters`` is not loaded.
1183 Example
1184 -------
1185 Check how many parameters an architecture change created::
1187 >>> model = read_model('nep4_PbTe.txt', restart_file='nep4_PbTe.restart')
1188 >>> model.n_unset_parameters
1189 0
1190 >>> model.augment(n_neuron=40).n_unset_parameters
1191 640
1192 """
1193 if self.restart_parameters is None:
1194 return None
1195 return _restart_sigma_counts(self)['unset']
1197 @property
1198 def parameter_counts(self) -> dict[str, dict[str, int]]:
1199 """Number of parameters per category, i.e. per component of
1200 :attr:`restart_parameters`.
1202 The categories are the same ones that the ``component`` argument of
1203 :meth:`set_restart_sigma` accepts, so this shows how many parameters a
1204 component-restricted call reaches. Only the categories that the model actually
1205 has are included: ``'charge_head'`` is absent for models without charges.
1207 With ``restart_parameters`` loaded, the counts describe the entries of the restart
1208 and are split into frozen and unset ones. Without it, they describe the current
1209 parameters of the model and only the totals are available, since there are no sigma
1210 values to split on. The totals exclude the :attr:`q_scaler` entries, which are not
1211 fit parameters and have no restart counterpart, so they sum to
1212 :attr:`n_parameters` minus the length of :attr:`q_scaler`.
1214 Returns
1215 -------
1216 dict
1217 Dictionary keyed by category. Each value holds the number of entries in total
1218 (``'total'``) and, if ``restart_parameters`` is loaded, the number that are
1219 frozen (``'frozen'``, sigma of zero) and the number that are unset
1220 (``'unset'``, sigma of ``NaN``).
1222 Example
1223 -------
1224 For a plain model only the totals are reported::
1226 >>> read_model('nep4_PbTe.txt').parameter_counts
1227 {'network_weights': {'total': 1921}, 'descriptor': {'total': 360}}
1229 With a restart loaded, the sigma values split each category further. Freeze the
1230 descriptor and check which category the frozen parameters are in::
1232 >>> model = read_model('nep4_PbTe.txt', restart_file='nep4_PbTe.restart')
1233 >>> frozen = model.set_restart_sigma(strategy='constant', value=0.0,
1234 ... component='descriptor', target='all')
1235 >>> frozen.parameter_counts
1236 {'network_weights': {'total': 1921, 'frozen': 0, 'unset': 0},
1237 'descriptor': {'total': 360, 'frozen': 360, 'unset': 0}}
1238 """
1239 if self.restart_parameters is None:
1240 return _model_parameter_counts(self)
1241 counts = {}
1242 for component in _RESTART_COMPONENTS:
1243 component_counts = _restart_sigma_counts(self, component)
1244 if component_counts['total'] > 0:
1245 counts[component] = component_counts
1246 return counts
1248 def _restart_availability(self) -> str:
1249 """Return the one-line summary of the restart-parameter state used by
1250 :meth:`__str__` and :meth:`_repr_html_`."""
1251 if self.restart_parameters is None:
1252 return 'not available'
1253 return (f'available ({self.n_frozen_parameters} frozen, '
1254 f'{self.n_unset_parameters} unset)')
1256 def remove_species(self, species: list[str]) -> 'Model':
1257 """Remove one or more species from the model.
1259 Returns a new :class:`Model` with the specified species removed.
1260 The source model is not modified.
1262 If ``restart_parameters`` are loaded, they are pruned to match (the
1263 entries for the removed species/pairs are dropped); the surviving
1264 entries are left exactly as they were. Use :meth:`set_restart_sigma`
1265 explicitly afterwards if you want to re-open the SNES search width for
1266 the surviving parameters before continuing training.
1268 Parameters
1269 ----------
1270 species
1271 Species names to remove.
1273 Returns
1274 -------
1275 Model
1276 New model with the specified species removed.
1278 Raises
1279 ------
1280 ValueError
1281 If any of the provided species is not found in the model.
1282 """
1283 for s in species:
1284 if s not in self.types:
1285 raise ValueError(f'{s} is not a species supported by the NEP model')
1287 new = copy.deepcopy(self)
1288 types_to_keep = [t for t in self.types if t not in species]
1289 new.types = tuple(types_to_keep)
1291 # Prune ANN parameters (for NEP4 and NEP5)
1292 if self.version in [4, 5]:
1293 new.ann_parameters = {
1294 key: value for key, value in new.ann_parameters.items()
1295 if key in types_to_keep or key.startswith('b1')
1296 }
1298 # Prune descriptor weights; key is a (species1, species2) tuple
1299 new.radial_descriptor_weights = {
1300 key: value for key, value in new.radial_descriptor_weights.items()
1301 if key[0] in types_to_keep and key[1] in types_to_keep
1302 }
1303 new.angular_descriptor_weights = {
1304 key: value for key, value in new.angular_descriptor_weights.items()
1305 if key[0] in types_to_keep and key[1] in types_to_keep
1306 }
1308 # Prune typewise cutoff lists so remaining species map to correct cutoffs
1309 if isinstance(self.radial_cutoff, list):
1310 indices = [i for i, t in enumerate(self.types) if t not in species]
1311 new.radial_cutoff = [self.radial_cutoff[i] for i in indices]
1312 new.angular_cutoff = [self.angular_cutoff[i] for i in indices]
1314 # Prune restart parameters to match; survivors are left untouched
1315 if new.restart_parameters is not None:
1316 for param_type in ['mu', 'sigma']:
1317 ann_key = f'ann_{param_type}'
1318 if self.version in [4, 5]:
1319 # Keep per-species keys for survivors, global bias keys, and
1320 # sqrt_epsilon_infinity (charge models)
1321 new.restart_parameters[ann_key] = {
1322 key: value for key, value in new.restart_parameters[ann_key].items()
1323 if (key in types_to_keep or key.startswith('b1')
1324 or key == 'sqrt_epsilon_infinity')
1325 }
1327 # Prune descriptor restart parameters
1328 for desc_type in ['radial', 'angular']:
1329 key = f'{desc_type}_descriptor_{param_type}'
1330 new.restart_parameters[key] = {
1331 k: v for k, v in new.restart_parameters[key].items()
1332 if k[0] in types_to_keep and k[1] in types_to_keep
1333 }
1335 # Recalculate parameter counts
1336 _recalculate_parameter_counts(new)
1338 return new
1340 def keep_species(self, species: list[str]) -> 'Model':
1341 """Retain only the specified species, removing all others.
1343 Convenience complement to :meth:`remove_species`. Useful when the set
1344 of species to drop is large (e.g. isolating two elements from a
1345 foundation model with dozens of species).
1347 Parameters
1348 ----------
1349 species
1350 Species names to keep. All other species are removed.
1352 Returns
1353 -------
1354 Model
1355 New model containing only the requested species.
1357 Raises
1358 ------
1359 ValueError
1360 If any of the requested species is not in the model.
1361 """
1362 unknown = [s for s in species if s not in self.types]
1363 if unknown:
1364 raise ValueError(
1365 f'Species not in model: {unknown}'
1366 )
1367 to_remove = [s for s in self.types if s not in species]
1368 return self.remove_species(to_remove)
1370 def reorder(self, order: list[str]) -> 'Model':
1371 """Reorder the species in the model.
1373 Returns a new :class:`Model` with species permuted according to
1374 ``order``. This is useful for aligning the species order of two
1375 models that must share the same order when used jointly by GPUMD,
1376 e.g. a NEP potential and a TNEP dipole/polarizability model
1377 referenced together via two ``potential`` lines in ``run.in`` and
1378 ``dump_dipole`` or ``dump_polarizability``.
1380 The source model is not modified. Since ``ann_parameters``,
1381 ``radial_descriptor_weights``, ``angular_descriptor_weights``, and
1382 ``restart_parameters`` are keyed by species name (or species-pair)
1383 rather than position, reordering only requires updating ``types``
1384 and, if typewise cutoffs are in use, the positional
1385 ``radial_cutoff`` and ``angular_cutoff`` lists.
1387 Parameters
1388 ----------
1389 order
1390 New species order. Must be a permutation of ``self.types``.
1392 Returns
1393 -------
1394 Model
1395 New model with species reordered.
1397 Raises
1398 ------
1399 ValueError
1400 If ``order`` is not a permutation of the current species.
1401 """
1402 if sorted(order) != sorted(self.types):
1403 raise ValueError(
1404 f'order must be a permutation of the current species {self.types}, '
1405 f'got {list(order)}'
1406 )
1408 new = copy.deepcopy(self)
1409 new.types = tuple(order)
1411 if isinstance(self.radial_cutoff, list):
1412 indices = [self.types.index(t) for t in order]
1413 new.radial_cutoff = [self.radial_cutoff[i] for i in indices]
1414 new.angular_cutoff = [self.angular_cutoff[i] for i in indices]
1416 return new
1418 def add_species(self,
1419 species: list[str],
1420 radial_cutoff: float | list[float] = None,
1421 angular_cutoff: float | list[float] = None,
1422 seed: int | None = None) -> 'Model':
1423 """Add one or more species to the model.
1425 Returns a new :class:`Model` with the requested species added. New ANN
1426 sub-networks and descriptor weight pairs are initialised by drawing
1427 ``mu`` uniformly from [-1, 1] (matching the GPUMD fresh-model
1428 initialisation); the corresponding restart sigma entries are left
1429 unset (``NaN``) — call :meth:`set_restart_sigma` afterwards to
1430 initialize them (e.g. ``model.add_species(['X']).set_restart_sigma()``
1431 fills only the new entries by default). Charge-specific parameters
1432 (``w1_charge``) are kept at ``mu = 0`` to preserve stability, also
1433 matching GPUMD. Existing parameters (``mu`` and ``sigma``) are left
1434 untouched. Call :meth:`initialize_parameters` to draw the new values from
1435 a different distribution instead, or to give ``w1_charge`` a non-zero start.
1437 Only supported for NEP4 models. For NEP3 the ANN is shared across all
1438 species and adding a per-species sub-network is not meaningful.
1440 Parameters
1441 ----------
1442 species
1443 New species names to add. Appended to ``types`` in the order given.
1444 radial_cutoff
1445 Radial cutoff(s) for the new species, in Å. Required when the model
1446 uses typewise cutoffs (i.e. ``isinstance(model.radial_cutoff, list)``
1447 is ``True``). Pass a single float or a list with one value per new
1448 species.
1449 angular_cutoff
1450 Angular cutoff(s) for the new species, in Å. Same requirements as
1451 ``radial_cutoff``.
1452 seed
1453 Seed for the random number generator used to draw the initial ``mu``
1454 values. Pass an integer for reproducible initialisation.
1456 Returns
1457 -------
1458 Model
1459 New model with updated structure, weights, and restart statistics.
1461 Raises
1462 ------
1463 ValueError
1464 If the model version is not 4, if ``restart_parameters`` are not
1465 loaded, if any species is already in the model, or if typewise
1466 cutoffs are used and ``radial_cutoff``/``angular_cutoff`` are not
1467 provided.
1468 """
1469 if self.version != 4:
1470 raise ValueError(
1471 f'add_species() only supports NEP4 models; got version {self.version}.'
1472 )
1473 for s in species:
1474 if s in self.types:
1475 raise ValueError(f'{s!r} is already in the model.')
1476 if self.restart_parameters is None:
1477 raise ValueError(
1478 'restart_parameters must be loaded before calling add_species(). '
1479 'Pass restart_file= to read_model() or call model.read_restart() first.'
1480 )
1482 uses_typewise = isinstance(self.radial_cutoff, list)
1483 if uses_typewise:
1484 if radial_cutoff is None or angular_cutoff is None:
1485 raise ValueError(
1486 'Model uses typewise cutoffs; provide radial_cutoff and angular_cutoff '
1487 'for the new species.'
1488 )
1489 rc_list = ([radial_cutoff] * len(species)
1490 if isinstance(radial_cutoff, (int, float)) else list(radial_cutoff))
1491 ac_list = ([angular_cutoff] * len(species)
1492 if isinstance(angular_cutoff, (int, float)) else list(angular_cutoff))
1493 if len(rc_list) != len(species) or len(ac_list) != len(species):
1494 raise ValueError(
1495 'Length of radial_cutoff/angular_cutoff must match the number of new species.'
1496 )
1498 new = copy.deepcopy(self)
1500 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular
1501 n_neuron = self.n_neuron
1502 is_charged = self.model_type == 'potential_with_charges'
1503 all_types_after = list(self.types) + list(species)
1504 rng = np.random.default_rng(seed)
1506 def _rand(shape):
1507 return rng.uniform(-1.0, 1.0, size=shape)
1509 # Step 1: New ANN sub-networks
1510 w1_shape = (n_neuron,) if is_charged else (1, n_neuron)
1511 for s_new in species:
1512 w0_vals = _rand((n_neuron, n_descriptor))
1513 b0_vals = _rand((n_neuron, 1))
1514 w1_vals = _rand(w1_shape)
1515 s_params = {'w0': w0_vals.copy(), 'b0': b0_vals.copy(), 'w1': w1_vals.copy()}
1516 if is_charged:
1517 s_params['w1_charge'] = np.zeros(n_neuron)
1518 new.ann_parameters[s_new] = s_params
1520 mu_entry = {'w0': w0_vals, 'b0': b0_vals, 'w1': w1_vals}
1521 sigma_entry = {
1522 'w0': np.full((n_neuron, n_descriptor), np.nan),
1523 'b0': np.full((n_neuron, 1), np.nan),
1524 'w1': np.full(w1_shape, np.nan),
1525 }
1526 if is_charged:
1527 mu_entry['w1_charge'] = np.zeros(n_neuron)
1528 sigma_entry['w1_charge'] = np.full(n_neuron, np.nan)
1529 new.restart_parameters['ann_mu'][s_new] = mu_entry
1530 new.restart_parameters['ann_sigma'][s_new] = sigma_entry
1532 # Step 2: New descriptor weight pairs
1533 n_r = (self.n_max_radial + 1, self.n_basis_radial + 1)
1534 n_a = (self.n_max_angular + 1, self.n_basis_angular + 1)
1535 existing_pairs = set(self.radial_descriptor_weights)
1536 new_pairs = {
1537 (s1, s2)
1538 for s1 in all_types_after for s2 in all_types_after
1539 if (s1, s2) not in existing_pairs
1540 }
1541 for pair in new_pairs:
1542 r_vals = _rand(n_r)
1543 a_vals = _rand(n_a)
1544 new.radial_descriptor_weights[pair] = r_vals.copy()
1545 new.angular_descriptor_weights[pair] = a_vals.copy()
1546 new.restart_parameters['radial_descriptor_mu'][pair] = r_vals
1547 new.restart_parameters['angular_descriptor_mu'][pair] = a_vals
1548 new.restart_parameters['radial_descriptor_sigma'][pair] = np.full(n_r, np.nan)
1549 new.restart_parameters['angular_descriptor_sigma'][pair] = np.full(n_a, np.nan)
1551 # Step 3: Update types and typewise cutoffs
1552 new.types = tuple(all_types_after)
1553 if uses_typewise:
1554 new.radial_cutoff = list(self.radial_cutoff) + rc_list
1555 new.angular_cutoff = list(self.angular_cutoff) + ac_list
1557 # Step 4: Recalculate parameter counts
1558 _recalculate_parameter_counts(new)
1560 return new
1562 def write(self, filename: str, restart_file: str = None) -> None:
1563 """Write NEP model to file in `nep.txt` format.
1565 Parameters
1566 ----------
1567 filename
1568 Output file name for the NEP model.
1569 restart_file
1570 If provided, also write restart parameters to this file in
1571 `nep.restart` format. Defaults to None.
1572 """
1573 with open(filename, 'w') as f:
1574 # header
1575 version_name = f'nep{self.version}'
1576 if self.zbl is not None:
1577 version_name += '_zbl'
1578 if self.model_type == 'potential_with_charges':
1579 version_name += f'_charge{self.charge_mode}'
1580 elif self.model_type != 'potential':
1581 version_name += f'_{self.model_type}'
1582 f.write(f'{version_name} {len(self.types)} {" ".join(self.types)}\n')
1583 if self.zbl is not None:
1584 zbl_tokens = list(self.zbl)
1585 if self.zbl_typewise_cutoff_factor is not None:
1586 zbl_tokens.append(self.zbl_typewise_cutoff_factor)
1587 f.write(f'zbl {" ".join(map(_format_header_float, zbl_tokens))}\n')
1588 f.write('cutoff')
1589 if isinstance(self.radial_cutoff, float) and isinstance(self.angular_cutoff, float):
1590 f.write(f' {_format_header_float(self.radial_cutoff)}'
1591 f' {_format_header_float(self.angular_cutoff)}')
1592 else:
1593 # Typewise cutoffs: one set of cutoffs per type
1594 for i in range(len(self.types)):
1595 f.write(f' {_format_header_float(self.radial_cutoff[i])}'
1596 f' {_format_header_float(self.angular_cutoff[i])}')
1597 f.write(f' {self.max_neighbors_radial} {self.max_neighbors_angular}')
1598 f.write('\n')
1599 f.write(f'n_max {self.n_max_radial} {self.n_max_angular}\n')
1600 f.write(f'basis_size {self.n_basis_radial} {self.n_basis_angular}\n')
1601 l_max_line = f'l_max {self.l_max_3b} {self.l_max_4b} {self.l_max_5b}'
1602 if self.has_q_112 or self.has_q_123 or self.has_q_233 or self.has_q_134:
1603 l_max_line += f' {self.has_q_112}'
1604 if self.has_q_123 or self.has_q_233 or self.has_q_134:
1605 l_max_line += f' {self.has_q_123}'
1606 if self.has_q_233 or self.has_q_134:
1607 l_max_line += f' {self.has_q_233}'
1608 if self.has_q_134:
1609 l_max_line += f' {self.has_q_134}'
1610 f.write(l_max_line + '\n')
1611 f.write(f'ANN {self.n_neuron} 0\n')
1613 # neural network weights
1614 keys = self.types if self.version in (4, 5) else ['all_species']
1615 suffixes = ['', '_polar'] if self.model_type == 'polarizability' else ['']
1616 for suffix in suffixes:
1617 for s in keys:
1618 # Order: w0, b0, w1 (, b1 if NEP5)
1619 # w0 indexed as: n*N_descriptor + nu
1620 w0 = self.ann_parameters[s][f'w0{suffix}']
1621 b0 = self.ann_parameters[s][f'b0{suffix}']
1622 w1 = self.ann_parameters[s][f'w1{suffix}']
1623 for n in range(self.n_neuron):
1624 for nu in range(
1625 self.n_descriptor_radial + self.n_descriptor_angular
1626 ):
1627 f.write(f'{w0[n, nu]:15.7e}\n')
1628 for b in b0[:, 0]:
1629 f.write(f'{b:15.7e}\n')
1630 for v in (w1[0, :] if w1.ndim == 2 else w1):
1631 f.write(f'{v:15.7e}\n')
1632 if f'w1_charge{suffix}' in self.ann_parameters[s]:
1633 for v in self.ann_parameters[s][f'w1_charge{suffix}']:
1634 f.write(f'{v:15.7e}\n')
1635 if self.version == 5:
1636 b1 = self.ann_parameters[s][f'b1{suffix}']
1637 f.write(f'{b1:15.7e}\n')
1638 if self.sqrt_epsilon_infinity is not None:
1639 f.write(f'{self.sqrt_epsilon_infinity:15.7e}\n')
1640 b1 = self.ann_parameters[f'b1{suffix}']
1641 f.write(f'{b1:15.7e}\n')
1643 # descriptor weights
1644 mat = []
1645 for s1 in self.types:
1646 for s2 in self.types:
1647 mat = np.hstack(
1648 [mat, self.radial_descriptor_weights[(s1, s2)].flatten()]
1649 )
1650 mat = np.hstack(
1651 [mat, self.angular_descriptor_weights[(s1, s2)].flatten()]
1652 )
1653 n_types = len(self.types)
1654 n = int(len(mat) / (n_types * n_types))
1655 mat = mat.reshape((n_types * n_types, n)).T
1656 for v in mat.flatten():
1657 f.write(f'{v:15.7e}\n')
1659 # scaler
1660 for v in self.q_scaler:
1661 f.write(f'{v:15.7e}\n')
1663 if restart_file is not None:
1664 self.write_restart(restart_file)
1666 def read_restart(self, filename: str):
1667 """Parses a file in `nep.restart` format and saves the
1668 content in the form of mean and standard deviation for each
1669 parameter in the corresponding NEP model.
1671 Parameters
1672 ----------
1673 filename
1674 Input file name.
1675 """
1676 mu, sigma = _get_restart_contents(filename)
1677 restart_parameters = np.array([mu, sigma]).T
1679 is_polarizability_model = self.model_type == 'polarizability'
1680 is_charged_model = self.model_type == 'potential_with_charges'
1682 n1 = self.n_ann_parameters
1683 n1 *= 2 if is_polarizability_model else 1
1684 n2 = n1 + self.n_descriptor_parameters
1685 ann_parameters = restart_parameters[:n1]
1686 descriptor_parameters = np.array(restart_parameters[n1:n2])
1688 if self.version == 3:
1689 n_networks = 1
1690 elif self.version in (4, 5):
1691 # one hidden layer per atomic species
1692 n_networks = len(self.types)
1693 else:
1694 raise ValueError(f'Cannot load nep.restart for NEP model version {self.version}')
1696 ann_groups = [s for s in self.ann_parameters.keys() if not s.startswith('b1')]
1697 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular
1698 restart = {}
1700 for i, content_type in enumerate(['mu', 'sigma']):
1701 ann = _sort_ann_parameters(ann_parameters[:, i],
1702 ann_groups,
1703 self.n_neuron,
1704 n_networks,
1705 self.version,
1706 n_descriptor,
1707 is_polarizability_model,
1708 is_charged_model)
1709 radial, angular = _sort_descriptor_parameters(descriptor_parameters[:, i],
1710 self.types,
1711 self.n_max_radial,
1712 self.n_basis_radial,
1713 self.n_max_angular,
1714 self.n_basis_angular)
1716 restart[f'ann_{content_type}'] = ann
1717 restart[f'radial_descriptor_{content_type}'] = radial
1718 restart[f'angular_descriptor_{content_type}'] = angular
1719 self.restart_parameters = restart
1721 def write_restart(self, filename: str):
1722 """Write the restart parameters to file in `nep.restart` format.
1724 Parameters
1725 ----------
1726 filename
1727 Output file name.
1729 Raises
1730 ------
1731 ValueError
1732 If ``restart_parameters`` is not loaded, or if any restart sigma
1733 value is unset (``NaN``), e.g. because :meth:`add_species` or
1734 :meth:`augment` were called without a follow-up
1735 :meth:`set_restart_sigma` to initialize the sigma of the newly
1736 created parameters.
1737 """
1738 if self.restart_parameters is None:
1739 raise ValueError(
1740 'restart_parameters is not loaded; nothing to write. Pass restart_file= '
1741 'to read_model(), call Model.read_restart(), or call '
1742 'Model.set_restart_sigma() to bootstrap one from the current model '
1743 'parameters before write_restart().'
1744 )
1745 for _, sigma_container, sigma_key in _restart_leaves(self, self.restart_parameters):
1746 if np.any(np.isnan(np.asarray(sigma_container[sigma_key], dtype=float))):
1747 raise ValueError(
1748 f'restart_parameters contains an unset (NaN) sigma value for '
1749 f'{sigma_key!r}. Call Model.set_restart_sigma() to initialize it '
1750 'before write_restart().'
1751 )
1752 keys = self.types if self.version in (4, 5) else ['all_species']
1753 suffixes = ['', '_polar'] if self.model_type == 'polarizability' else ['']
1754 columns = []
1755 for i, parameter in enumerate(['mu', 'sigma']):
1756 # neural network weights
1757 ann_parameters = self.restart_parameters[f'ann_{parameter}']
1758 column = []
1759 for suffix in suffixes:
1760 for s in keys:
1761 # Order: w0, b0, w1 (, b1 if NEP5)
1762 # w0 indexed as: n*N_descriptor + nu
1763 w0 = ann_parameters[s][f'w0{suffix}']
1764 b0 = ann_parameters[s][f'b0{suffix}']
1765 w1 = ann_parameters[s][f'w1{suffix}']
1766 for n in range(self.n_neuron):
1767 for nu in range(
1768 self.n_descriptor_radial + self.n_descriptor_angular
1769 ):
1770 column.append(f'{w0[n, nu]:15.7e}')
1771 for b in b0[:, 0]:
1772 column.append(f'{b:15.7e}')
1773 for v in (w1[0, :] if w1.ndim == 2 else w1):
1774 column.append(f'{v:15.7e}')
1775 if f'w1_charge{suffix}' in ann_parameters[s]:
1776 for v in ann_parameters[s][f'w1_charge{suffix}']:
1777 column.append(f'{v:15.7e}')
1778 if f'b1{suffix}' in ann_parameters[s]:
1779 column.append(f'{ann_parameters[s][f"b1{suffix}"]:15.7e}')
1780 if f'sqrt_epsilon_infinity{suffix}' in ann_parameters:
1781 column.append(f'{ann_parameters[f"sqrt_epsilon_infinity{suffix}"]:15.7e}')
1782 b1 = ann_parameters[f'b1{suffix}']
1783 column.append(f'{b1:15.7e}')
1784 columns.append(column)
1786 # descriptor weights
1787 radial_descriptor_parameters = self.restart_parameters[f'radial_descriptor_{parameter}']
1788 angular_descriptor_parameters = self.restart_parameters[
1789 f'angular_descriptor_{parameter}']
1790 mat = []
1791 for s1 in self.types:
1792 for s2 in self.types:
1793 mat = np.hstack(
1794 [mat, radial_descriptor_parameters[(s1, s2)].flatten()]
1795 )
1796 mat = np.hstack(
1797 [mat, angular_descriptor_parameters[(s1, s2)].flatten()]
1798 )
1799 n_types = len(self.types)
1800 n = int(len(mat) / (n_types * n_types))
1801 mat = mat.reshape((n_types * n_types, n)).T
1802 for v in mat.flatten():
1803 column.append(f'{v:15.7e}')
1805 # Join the mean and standard deviation columns
1806 assert len(columns[0]) == len(columns[1]), 'Length of means must match standard deviation'
1807 joined = [f'{s1} {s2}\n' for s1, s2 in zip(*columns)]
1808 with open(filename, 'w') as f:
1809 f.writelines(joined)
1811 def initialize_parameters(self,
1812 strategy: str = 'uniform',
1813 *,
1814 value: float = None,
1815 low: float = None,
1816 high: float = None,
1817 mean: float = None,
1818 std: float = None,
1819 species: str | list[str] = None,
1820 component: str | list[str] = None,
1821 target: str = 'unset',
1822 seed: int | None = None) -> 'Model':
1823 """Initialize parameter values from a distribution.
1825 Returns a new :class:`Model` in which the parameters selected by
1826 *target*/*species*/*component* are drawn according to *strategy*. The source model is
1827 not modified. Every sigma value is left exactly as it was, so
1828 :meth:`set_restart_sigma` remains the only method that assigns sigma.
1830 Both the parameter values of the model (:attr:`ann_parameters` and the descriptor
1831 weights, i.e. what goes into ``nep.txt``) and the corresponding ``mu`` entries of
1832 :attr:`restart_parameters` (the SNES mean, i.e. what goes into ``nep.restart``) are
1833 written. Only the selected positions are touched, since for a trained model the two
1834 are not the same thing: ``nep.txt`` holds the best parameters found so far while
1835 the restart holds the mean of the search distribution.
1837 :attr:`sqrt_epsilon_infinity` is the one parameter no strategy reaches, whatever
1838 *target*/*species*/*component* select. Every other parameter may take any value, but
1839 ``epsilon_infinity`` is the square of this one, the high-frequency dielectric
1840 constant, so a drawn value that is negative or close to zero is not a starting point.
1841 Set it through the ``sqrt_epsilon_infinity`` argument of :meth:`augment` instead,
1842 which defaults to 1.
1844 The main use is to give the parameters created by :meth:`augment` a starting point
1845 other than zero::
1847 model.augment(n_neuron=40) \\
1848 .initialize_parameters(strategy='uniform', low=-1, high=1, seed=42) \\
1849 .set_restart_sigma(strategy='constant', value=0.0, target='set') \\
1850 .set_restart_sigma(target='unset')
1852 The order matters. :meth:`initialize_parameters` has to come before the
1853 :meth:`set_restart_sigma` call that fills the unset sigma values, because an unset
1854 sigma is what marks a parameter as new. Initializing the values first also makes the
1855 default sigma strategy meaningful for the new parameters, since
1856 ``strategy='scale_mu'`` computes ``sigma = max(floor, factor * |mu|)``, which is just
1857 the floor as long as ``mu`` is zero.
1859 Parameters
1860 ----------
1861 strategy
1862 How to draw the new values at the selected positions:
1864 - ``'uniform'`` (default): ``value ~ U(low, high)``. Drawing from
1865 ``U(-1, 1)`` reproduces what :meth:`add_species` does for a new sub-network.
1866 - ``'constant'``: every selected parameter is set to ``value``.
1867 - ``'normal'``: ``value ~ N(mean, std)``.
1869 The values are used as drawn, including negative ones.
1870 value
1871 Value for ``strategy='constant'``.
1872 low, high
1873 Bounds for ``strategy='uniform'``.
1874 mean, std
1875 Parameters of the normal distribution for ``strategy='normal'``.
1876 species
1877 Restrict the update to one or more species (and descriptor pairs involving them).
1878 ``None`` (default) applies to all species; the global shared bias is only
1879 included when ``species`` is ``None``.
1880 component
1881 Restrict the update to one or more of ``'network_weights'``, ``'descriptor'``,
1882 ``'charge_head'``. ``None`` (default) applies to all of them. ``'charge_head'``
1883 reaches ``w1_charge`` only, ``sqrt_epsilon_infinity`` being excluded as described
1884 above.
1885 target
1886 Which parameters to initialize, selected by the state of the corresponding
1887 *sigma*: ``'unset'`` (default) only the parameters whose sigma is unset (``NaN``),
1888 i.e. those newly created by :meth:`augment` or :meth:`add_species`; ``'set'`` only
1889 those that are already part of the search; ``'all'`` every selected parameter.
1890 A value of zero is not used as the marker, since a trained parameter may be zero.
1891 seed
1892 Seed for the random number generator. Pass an integer for reproducibility.
1894 Returns
1895 -------
1896 Model
1897 New model with the selected parameter values initialized.
1899 Raises
1900 ------
1901 ValueError
1902 If ``restart_parameters`` is not loaded, if ``strategy``/``target``/``component``
1903 is not recognized, or if a strategy-specific required argument is missing.
1905 Example
1906 -------
1907 Give the neurons that :meth:`augment` added a random starting point::
1909 >>> model = read_model('nep.txt', restart_file='nep.restart')
1910 >>> grown = model.augment(n_neuron=40)
1911 >>> grown.ann_parameters['Pb']['w0'][30:].any() # new rows are zero
1912 False
1913 >>> initialized = grown.initialize_parameters(low=-1, high=1, seed=42)
1914 >>> initialized.ann_parameters['Pb']['w0'][30:].any()
1915 True
1916 """
1917 valid_strategies = {'constant', 'uniform', 'normal'}
1918 if strategy not in valid_strategies:
1919 raise ValueError(
1920 f'strategy must be one of {sorted(valid_strategies)}; got {strategy!r}'
1921 )
1922 if target not in ('unset', 'set', 'all'):
1923 raise ValueError(f"target must be 'unset', 'set', or 'all'; got {target!r}")
1924 if strategy == 'constant' and value is None:
1925 raise ValueError("strategy='constant' requires value.")
1926 if strategy == 'uniform' and (low is None or high is None):
1927 raise ValueError("strategy='uniform' requires low and high.")
1928 if strategy == 'normal' and (mean is None or std is None):
1929 raise ValueError("strategy='normal' requires mean and std.")
1930 if self.restart_parameters is None:
1931 raise ValueError(
1932 'restart_parameters must be loaded before calling initialize_parameters(), '
1933 'since the sigma values are what mark a parameter as new. Pass restart_file= '
1934 'to read_model(), call model.read_restart(), or call '
1935 'Model.set_restart_sigma() to bootstrap one first.'
1936 )
1938 new = copy.deepcopy(self)
1939 rng = np.random.default_rng(seed)
1940 kw = dict(value=value, low=low, high=high, mean=mean, std=std)
1942 for container, key, owner, name in _model_parameter_leaves(new, component, species):
1943 if name == 'sqrt_epsilon_infinity':
1944 # Deliberately left out of every strategy. The other parameters may take any
1945 # value, but epsilon_infinity is the square of this one, the high-frequency
1946 # dielectric constant, so a drawn value that is negative or close to zero is
1947 # not a starting point. augment(charge_head=True) sets it, via its
1948 # sqrt_epsilon_infinity argument, and that value is what training starts from.
1949 continue
1950 restart_leaf = _restart_leaf_for(new.restart_parameters, owner, name, key)
1951 if restart_leaf is None: 1951 ↛ 1952line 1951 didn't jump to line 1952 because the condition on line 1951 was never true
1952 continue
1953 mu_container, sigma_container, restart_key = restart_leaf
1954 sigma = sigma_container[restart_key]
1956 if np.isscalar(sigma) or isinstance(sigma, (float, int)):
1957 is_unset = np.isnan(float(sigma))
1958 apply_here = (
1959 target == 'all' or (target == 'unset' and is_unset)
1960 or (target == 'set' and not is_unset)
1961 )
1962 if not apply_here:
1963 continue
1964 new_value = float(_draw_values(strategy, rng, **kw))
1965 _leaf_set(container, key, new_value)
1966 mu_container[restart_key] = new_value
1967 continue
1969 values = _leaf_get(container, key)
1970 mu = mu_container[restart_key]
1971 if np.shape(values) != np.shape(sigma) or np.shape(mu) != np.shape(sigma): 1971 ↛ 1972line 1971 didn't jump to line 1972 because the condition on line 1971 was never true
1972 raise ValueError(
1973 f'Shape mismatch for {name!r}: the model holds {np.shape(values)} while '
1974 f'the restart holds {np.shape(mu)}. initialize_parameters() cannot map '
1975 'the two onto each other.'
1976 )
1977 if target == 'unset':
1978 mask = np.isnan(sigma)
1979 elif target == 'set':
1980 mask = ~np.isnan(sigma)
1981 else:
1982 mask = np.ones_like(sigma, dtype=bool)
1983 if not np.any(mask):
1984 continue
1985 drawn = _draw_values(strategy, rng, size=int(np.count_nonzero(mask)), **kw)
1986 values[mask] = drawn
1987 mu[mask] = drawn
1989 return new
1991 def set_restart_sigma(self,
1992 strategy: str = 'scale_mu',
1993 *,
1994 value: float = None,
1995 factor: float = None,
1996 floor: float = 1e-6,
1997 low: float = None,
1998 high: float = None,
1999 mean: float = None,
2000 std: float = None,
2001 species: str | list[str] = None,
2002 component: str | list[str] = None,
2003 target: str = 'unset',
2004 seed: int | None = None) -> 'Model':
2005 """Assign SNES restart sigma values.
2007 Returns a new :class:`Model` with sigma values updated according to
2008 *strategy*, at the positions selected by *target*/*species*/*component*.
2009 ``mu`` and every other field are left unchanged. This is the only
2010 method that ever assigns sigma values; the structural methods
2011 (:meth:`remove_species`, :meth:`keep_species`, :meth:`add_species`,
2012 :meth:`augment`, :meth:`prune`) leave the sigma of the parameters they keep
2013 untouched and mark newly created parameters' sigma as unset (``NaN``) rather
2014 than computing a value inline. Parameter values are assigned by
2015 :meth:`initialize_parameters` instead, which in turn never touches a sigma.
2017 If ``restart_parameters`` is not loaded, it is created first: ``mu`` is
2018 copied from the model's current (trained) parameters, and every sigma
2019 is initialized as unset (``NaN``). This makes it possible to bootstrap
2020 a ``nep.restart`` file "from scratch" for a plain ``nep.txt`` model.
2022 Parameters
2023 ----------
2024 strategy
2025 How to compute new sigma values at the selected positions:
2027 - ``'constant'``: ``sigma = value``.
2028 - ``'scale_mu'`` (default): ``sigma = max(floor, factor * |mu|)``,
2029 re-opening the SNES search width in proportion to each
2030 parameter's magnitude. ``factor`` defaults to ``0.1`` for this
2031 strategy.
2032 - ``'scale_sigma'``: ``sigma = sigma * factor``. Requires the
2033 selected sigma values to already be set (not ``NaN``).
2034 - ``'uniform'``: draw ``sigma ~ U(low, high)``.
2035 - ``'normal'``: draw ``sigma = |N(mean, std)|``.
2036 value
2037 Sigma value for ``strategy='constant'``.
2038 factor
2039 Scale factor for ``strategy='scale_mu'`` (default ``0.1`` if not
2040 given) or ``strategy='scale_sigma'`` (required).
2041 floor
2042 Minimum sigma for ``strategy='scale_mu'``.
2043 low, high
2044 Bounds for ``strategy='uniform'``. Since a sigma is a standard
2045 deviation, ``low`` must be positive and smaller than ``high``, which
2046 keeps every drawn value positive. Use ``strategy='constant'`` for a
2047 single value rather than ``low == high``.
2048 mean, std
2049 Parameters of the normal distribution for ``strategy='normal'``.
2050 species
2051 Restrict the update to one or more species (and descriptor pairs
2052 involving them). ``None`` (default) applies to all species; global
2053 parameters (the shared bias, ``sqrt_epsilon_infinity``) are only
2054 included when ``species`` is ``None``.
2055 component
2056 Restrict the update to one or more of ``'network_weights'``,
2057 ``'descriptor'``, ``'charge_head'``. ``None`` (default) applies to
2058 all three.
2059 target
2060 Which existing sigma values to update: ``'unset'`` (default) only
2061 fills in ``NaN`` entries (e.g. those left by :meth:`add_species`/
2062 :meth:`augment`); ``'set'`` only updates already-set entries;
2063 ``'all'`` updates every selected entry regardless of its current
2064 value.
2065 seed
2066 Seed for the random number generator used by the ``'uniform'`` and
2067 ``'normal'`` strategies. Pass an integer for reproducibility.
2069 Returns
2070 -------
2071 Model
2072 New model with updated restart sigma values.
2074 Raises
2075 ------
2076 ValueError
2077 If ``strategy``/``target``/``component`` is not recognized, if a
2078 strategy-specific required argument is missing, if
2079 ``strategy='uniform'`` is given bounds that would admit a
2080 non-positive sigma, or if ``strategy='scale_sigma'`` is applied to a
2081 still-unset (``NaN``) sigma value.
2082 """
2083 valid_strategies = {'constant', 'scale_mu', 'scale_sigma', 'uniform', 'normal'}
2084 if strategy not in valid_strategies:
2085 raise ValueError(
2086 f'strategy must be one of {sorted(valid_strategies)}; got {strategy!r}'
2087 )
2088 if target not in ('unset', 'set', 'all'):
2089 raise ValueError(f"target must be 'unset', 'set', or 'all'; got {target!r}")
2090 if strategy == 'constant' and value is None:
2091 raise ValueError("strategy='constant' requires value.")
2092 if strategy == 'scale_mu' and factor is None:
2093 factor = 0.1
2094 if strategy == 'scale_sigma' and factor is None:
2095 raise ValueError("strategy='scale_sigma' requires factor.")
2096 if strategy == 'uniform' and (low is None or high is None):
2097 raise ValueError("strategy='uniform' requires low and high.")
2098 # A sigma is a standard deviation, so the whole interval has to be positive. Both
2099 # checks are needed: numpy samples between the two bounds whatever their order, so a
2100 # positive low on its own does not bound the draw from below.
2101 if strategy == 'uniform' and low <= 0:
2102 raise ValueError(
2103 "strategy='uniform' requires a positive low, since sigma is a standard "
2104 f'deviation; got low={low!r}.'
2105 )
2106 if strategy == 'uniform' and low >= high:
2107 raise ValueError(
2108 "strategy='uniform' requires low < high, so that every drawn sigma stays "
2109 f"positive; got low={low!r} and high={high!r}. Use strategy='constant' for a "
2110 'single value.'
2111 )
2112 if strategy == 'normal' and (mean is None or std is None):
2113 raise ValueError("strategy='normal' requires mean and std.")
2115 new = copy.deepcopy(self)
2116 if new.restart_parameters is None:
2117 new.restart_parameters = _new_restart_parameters_from_model(new)
2119 rng = np.random.default_rng(seed)
2120 kw = dict(value=value, factor=factor, floor=floor, low=low, high=high, mean=mean, std=std)
2121 for mu, sigma_container, sigma_key in _restart_leaves(
2122 new, new.restart_parameters, component, species
2123 ):
2124 _apply_sigma_strategy(mu, sigma_container, sigma_key, strategy, target, rng, **kw)
2126 return new
2128 def augment(self,
2129 n_neuron: int = None,
2130 l_max_4b: int = None,
2131 l_max_5b: int = None,
2132 has_q_112: bool = None,
2133 has_q_123: bool = None,
2134 has_q_233: bool = None,
2135 has_q_134: bool = None,
2136 charge_head: bool = False,
2137 charge_mode: int = 1,
2138 sqrt_epsilon_infinity: float = 1.0) -> 'Model':
2139 """Augment the model by adding neurons, descriptor terms, or a charge output head.
2141 Returns a new :class:`Model` with the requested structural changes applied.
2142 The source model is not modified. Existing parameter values (``mu`` and
2143 ``sigma``) are preserved exactly; new parameters are initialized to
2144 ``mu = 0``, with the corresponding restart sigma left unset (``NaN``).
2145 ``sqrt_epsilon_infinity`` is the one exception: ``epsilon_infinity`` is its square,
2146 a dielectric constant, so zero is not a value it can take. It starts at the value of
2147 the ``sqrt_epsilon_infinity`` argument (1 by default) instead.
2148 Call :meth:`set_restart_sigma` afterwards to initialize the new sigma
2149 entries (e.g. ``model.augment(n_neuron=40).set_restart_sigma()`` fills
2150 only the new entries by default). To give the new parameters a starting
2151 value other than zero, call :meth:`initialize_parameters` before that, since
2152 an unset sigma is what marks a parameter as new.
2154 Parameters
2155 ----------
2156 n_neuron
2157 Target neuron count; must be >= current. ``None`` leaves unchanged.
2158 l_max_4b
2159 Target 4-body l_max value; must be >= current. ``None`` leaves unchanged.
2160 l_max_5b
2161 Target 5-body l_max value; must be >= current. ``None`` leaves unchanged.
2162 has_q_112
2163 ``True`` enables the q_112 5-body descriptor; ``None`` or ``False`` leaves
2164 the current state unchanged (disabling an already-enabled term raises).
2165 has_q_123
2166 Same as ``has_q_112`` but for the q_123 term.
2167 has_q_233
2168 Same as ``has_q_112`` but for the q_233 term.
2169 has_q_134
2170 Same as ``has_q_112`` but for the q_134 term.
2171 charge_head
2172 If ``True``, promote a ``potential`` model to ``potential_with_charges`` by
2173 adding a charge output head (w1_charge per species and sqrt_epsilon_infinity).
2174 charge_mode
2175 Charge algorithm variant to record for the new charge head; must be 1 or 2.
2176 1 corresponds to a qNEP model, including both real- and reciprocal-space
2177 contributions. 2 corresponds to a qNEP model, including the reciprocal-space
2178 contribution only. Only meaningful when ``charge_head=True``.
2179 sqrt_epsilon_infinity
2180 Starting value for the new ``sqrt_epsilon_infinity`` parameter, used only when
2181 ``charge_head=True``. Must be positive: ``epsilon_infinity`` is its square, the
2182 high-frequency dielectric constant. The default of 1 corresponds to no dielectric
2183 screening, which is the neutral starting point for training. This is the one
2184 parameter ``augment`` does not leave at zero, and :meth:`initialize_parameters`
2185 skips it for the same reason, so the value given here is the one that reaches
2186 training.
2188 Returns
2189 -------
2190 Model
2191 New model with updated structure, weights, and restart statistics.
2193 Raises
2194 ------
2195 ValueError
2196 If ``restart_parameters`` is not loaded, if ``n_neuron`` or an ``l_max_*``
2197 target is smaller than the current value, if a ``has_q_*`` flag attempts to
2198 disable an already-enabled term, or if ``charge_head=True`` on a model that
2199 is not of type ``potential`` or comes with a non-positive
2200 ``sqrt_epsilon_infinity``.
2201 """
2202 # Structural checks (independent of restart)
2203 if self.version not in (3, 4):
2204 raise ValueError(
2205 f'augment() only supports NEP versions 3 and 4; got version {self.version}.'
2206 )
2207 if n_neuron is not None and n_neuron < self.n_neuron:
2208 raise ValueError(
2209 f'n_neuron ({n_neuron}) must be >= current n_neuron ({self.n_neuron}); '
2210 'use prune() to reduce.'
2211 )
2212 if l_max_4b is not None and l_max_4b < self.l_max_4b:
2213 raise ValueError(
2214 f'l_max_4b ({l_max_4b}) must be >= current l_max_4b ({self.l_max_4b}); '
2215 'use prune() to disable.'
2216 )
2217 if l_max_5b is not None and l_max_5b < self.l_max_5b:
2218 raise ValueError(
2219 f'l_max_5b ({l_max_5b}) must be >= current l_max_5b ({self.l_max_5b}); '
2220 'use prune() to disable.'
2221 )
2222 for flag_val, name in [
2223 (has_q_112, 'has_q_112'), (has_q_123, 'has_q_123'), (has_q_233, 'has_q_233'),
2224 (has_q_134, 'has_q_134')
2225 ]:
2226 if flag_val is False and getattr(self, name):
2227 raise ValueError(
2228 f'Cannot disable {name} via augment(); '
2229 'use prune() to disable descriptor terms.'
2230 )
2231 if charge_head and self.model_type != 'potential':
2232 raise ValueError(
2233 f'charge_head=True requires model_type="potential"; '
2234 f'got "{self.model_type}".'
2235 )
2236 if charge_head and self.version != 4:
2237 # GPUMD only accepts nep4_charge1/nep4_charge2 for the qNEP charge modes, so a
2238 # nep3_charge1 model would be written but could not be read back.
2239 raise ValueError(
2240 f'charge_head=True requires a NEP4 model; got version {self.version}.'
2241 )
2242 if charge_head and charge_mode not in (1, 2):
2243 raise ValueError(f'charge_mode must be 1 or 2; got {charge_mode}.')
2244 if charge_head and not float(sqrt_epsilon_infinity) > 0:
2245 raise ValueError(
2246 'sqrt_epsilon_infinity must be positive; got '
2247 f'{sqrt_epsilon_infinity}. epsilon_infinity is its square, the '
2248 'high-frequency dielectric constant.'
2249 )
2250 if self.restart_parameters is None:
2251 raise ValueError(
2252 'restart_parameters must be loaded before calling augment(). '
2253 'Pass restart_file= to read_model() or call model.read_restart() first.'
2254 )
2256 new = copy.deepcopy(self)
2258 # Resolve new structural parameters
2259 new_l_max_4b = l_max_4b if l_max_4b is not None else self.l_max_4b
2260 new_l_max_5b = l_max_5b if l_max_5b is not None else self.l_max_5b
2261 new_has_q_112 = int(has_q_112) if has_q_112 is not None else self.has_q_112
2262 new_has_q_123 = int(has_q_123) if has_q_123 is not None else self.has_q_123
2263 new_has_q_233 = int(has_q_233) if has_q_233 is not None else self.has_q_233
2264 new_has_q_134 = int(has_q_134) if has_q_134 is not None else self.has_q_134
2265 new_n_neuron = n_neuron if n_neuron is not None else self.n_neuron
2267 new_l_max_enh = (self.l_max_3b
2268 + (new_l_max_4b > 0) + (new_l_max_5b > 0)
2269 + (new_has_q_112 > 0) + (new_has_q_123 > 0) + (new_has_q_233 > 0)
2270 + (new_has_q_134 > 0))
2271 new_n_desc_angular = (self.n_max_angular + 1) * new_l_max_enh
2272 old_n_desc = self.n_descriptor_radial + self.n_descriptor_angular
2273 new_n_desc = self.n_descriptor_radial + new_n_desc_angular
2274 delta_desc = new_n_desc - old_n_desc
2275 delta_neuron = new_n_neuron - self.n_neuron
2277 keys = self.types if self.version in (4, 5) else ['all_species']
2279 # Step 1: Expand descriptor dimensions (new columns in w0, new q_scaler entries)
2280 if delta_desc > 0:
2281 for s in keys:
2282 old_w0 = new.ann_parameters[s]['w0'] # (n_neuron_old, old_n_desc)
2283 new.ann_parameters[s]['w0'] = np.hstack(
2284 [old_w0, np.zeros((self.n_neuron, delta_desc))]
2285 )
2286 old_mu_w0 = new.restart_parameters['ann_mu'][s]['w0']
2287 new.restart_parameters['ann_mu'][s]['w0'] = np.hstack(
2288 [old_mu_w0, np.zeros((self.n_neuron, delta_desc))]
2289 )
2290 old_sigma_w0 = new.restart_parameters['ann_sigma'][s]['w0']
2291 new.restart_parameters['ann_sigma'][s]['w0'] = np.hstack(
2292 [old_sigma_w0, np.full((self.n_neuron, delta_desc), np.nan)]
2293 )
2294 new.q_scaler = list(new.q_scaler) + [1.0] * delta_desc
2296 # Step 2: Expand neuron count (new rows in w0/b0, new columns in w1)
2297 if delta_neuron > 0:
2298 for s in keys:
2299 # w0: append new rows
2300 cur_w0 = new.ann_parameters[s]['w0'] # (n_old, new_n_desc)
2301 new.ann_parameters[s]['w0'] = np.vstack(
2302 [cur_w0, np.zeros((delta_neuron, new_n_desc))]
2303 )
2304 # b0: append new rows
2305 cur_b0 = new.ann_parameters[s]['b0']
2306 new.ann_parameters[s]['b0'] = np.vstack(
2307 [cur_b0, np.zeros((delta_neuron, 1))]
2308 )
2309 # w1: append new columns; handle both 2D (standard) and 1D (charge)
2310 cur_w1 = new.ann_parameters[s]['w1']
2311 zeros_w1 = (np.zeros(delta_neuron) if cur_w1.ndim == 1
2312 else np.zeros((1, delta_neuron)))
2313 new.ann_parameters[s]['w1'] = np.hstack([cur_w1, zeros_w1])
2314 if 'w1_charge' in new.ann_parameters[s]:
2315 cur_wc = new.ann_parameters[s]['w1_charge']
2316 new.ann_parameters[s]['w1_charge'] = np.hstack([cur_wc, np.zeros(delta_neuron)])
2318 # restart w0
2319 cur_mu_w0 = new.restart_parameters['ann_mu'][s]['w0']
2320 new.restart_parameters['ann_mu'][s]['w0'] = np.vstack(
2321 [cur_mu_w0, np.zeros((delta_neuron, new_n_desc))]
2322 )
2323 cur_sigma_w0 = new.restart_parameters['ann_sigma'][s]['w0']
2324 new.restart_parameters['ann_sigma'][s]['w0'] = np.vstack(
2325 [cur_sigma_w0, np.full((delta_neuron, new_n_desc), np.nan)]
2326 )
2327 # restart b0
2328 cur_mu_b0 = new.restart_parameters['ann_mu'][s]['b0']
2329 new.restart_parameters['ann_mu'][s]['b0'] = np.vstack(
2330 [cur_mu_b0, np.zeros((delta_neuron, 1))]
2331 )
2332 cur_sigma_b0 = new.restart_parameters['ann_sigma'][s]['b0']
2333 new.restart_parameters['ann_sigma'][s]['b0'] = np.vstack(
2334 [cur_sigma_b0, np.full((delta_neuron, 1), np.nan)]
2335 )
2336 # restart w1
2337 cur_mu_w1 = new.restart_parameters['ann_mu'][s]['w1']
2338 zeros_w1 = (np.zeros(delta_neuron) if cur_mu_w1.ndim == 1
2339 else np.zeros((1, delta_neuron)))
2340 new.restart_parameters['ann_mu'][s]['w1'] = np.hstack([cur_mu_w1, zeros_w1])
2341 cur_sigma_w1 = new.restart_parameters['ann_sigma'][s]['w1']
2342 nan_w1 = (np.full(delta_neuron, np.nan) if cur_sigma_w1.ndim == 1
2343 else np.full((1, delta_neuron), np.nan))
2344 new.restart_parameters['ann_sigma'][s]['w1'] = np.hstack([cur_sigma_w1, nan_w1])
2345 if 'w1_charge' in new.restart_parameters['ann_mu'][s]:
2346 cur = new.restart_parameters['ann_mu'][s]['w1_charge']
2347 new.restart_parameters['ann_mu'][s]['w1_charge'] = np.hstack(
2348 [cur, np.zeros(delta_neuron)]
2349 )
2350 cur = new.restart_parameters['ann_sigma'][s]['w1_charge']
2351 new.restart_parameters['ann_sigma'][s]['w1_charge'] = np.hstack(
2352 [cur, np.full(delta_neuron, np.nan)]
2353 )
2355 # Step 3: Add charge output head
2356 if charge_head:
2357 new.model_type = 'potential_with_charges'
2358 new.charge_mode = charge_mode
2359 new.sqrt_epsilon_infinity = float(sqrt_epsilon_infinity)
2360 for s in keys:
2361 cur_w1 = new.ann_parameters[s]['w1'] # (1, new_n_neuron)
2362 new.ann_parameters[s]['w1'] = cur_w1[0, :] # flatten to 1D
2363 new.ann_parameters[s]['w1_charge'] = np.zeros(new_n_neuron)
2365 cur_mu_w1 = new.restart_parameters['ann_mu'][s]['w1']
2366 new.restart_parameters['ann_mu'][s]['w1'] = cur_mu_w1[0, :]
2367 new.restart_parameters['ann_mu'][s]['w1_charge'] = np.zeros(new_n_neuron)
2369 cur_sigma_w1 = new.restart_parameters['ann_sigma'][s]['w1']
2370 new.restart_parameters['ann_sigma'][s]['w1'] = cur_sigma_w1[0, :]
2371 new.restart_parameters['ann_sigma'][s]['w1_charge'] = np.full(
2372 new_n_neuron, np.nan
2373 )
2375 new.restart_parameters['ann_mu']['sqrt_epsilon_infinity'] = float(
2376 sqrt_epsilon_infinity
2377 )
2378 new.restart_parameters['ann_sigma']['sqrt_epsilon_infinity'] = float('nan')
2380 # Step 4: Update header metadata
2381 new.l_max_4b = new_l_max_4b
2382 new.l_max_5b = new_l_max_5b
2383 new.has_q_112 = new_has_q_112
2384 new.has_q_123 = new_has_q_123
2385 new.has_q_233 = new_has_q_233
2386 new.has_q_134 = new_has_q_134
2387 new.n_descriptor_angular = new_n_desc_angular
2388 new.n_neuron = new_n_neuron
2390 # Step 5: Recalculate parameter counts
2391 _recalculate_parameter_counts(new)
2393 return new
2395 def prune(self,
2396 n_neuron: int = None,
2397 l_max_4b: int = None,
2398 l_max_5b: int = None,
2399 has_q_112: bool = None,
2400 has_q_123: bool = None,
2401 has_q_233: bool = None,
2402 has_q_134: bool = None,
2403 charge_head: bool = False) -> 'Model':
2404 """Prune the model by removing neurons, disabling descriptor terms, or removing
2405 the charge output head.
2407 Returns a new :class:`Model` with the requested structural changes applied.
2408 The source model is not modified. When reducing ``n_neuron``, neurons are
2409 selected by importance score averaged over species:
2410 ``importance[n] = mean_s(||w0_s[n,:]||_2 * |w1_s[n]|)``.
2412 Surviving parameters (``mu`` and ``sigma``) are left exactly as they
2413 were. Use :meth:`set_restart_sigma` explicitly afterwards if you want
2414 to re-open the SNES search width for the survivors before continuing
2415 training.
2417 Parameters
2418 ----------
2419 n_neuron
2420 Target neuron count; must be <= current. ``None`` leaves unchanged.
2421 l_max_4b
2422 Target 4-body l_max; must be <= current. Setting to ``0`` removes the
2423 4-body angular descriptor block. Reducing to a lower non-zero value is
2424 a header-only change (descriptor dimensions unchanged). ``None`` leaves
2425 unchanged.
2426 l_max_5b
2427 Same as ``l_max_4b`` but for five-body terms.
2428 has_q_112
2429 ``False`` disables and removes the q_112 descriptor block. ``None``
2430 leaves unchanged. ``True`` is not valid; use :meth:`augment` instead.
2431 has_q_123
2432 Same as ``has_q_112`` but for the q_123 term.
2433 has_q_233
2434 Same as ``has_q_112`` but for the q_233 term.
2435 has_q_134
2436 Same as ``has_q_112`` but for the q_134 term.
2437 charge_head
2438 If ``True``, remove the charge output head from a
2439 ``potential_with_charges`` model, converting it back to ``potential``.
2440 Removes ``w1_charge`` per species and ``sqrt_epsilon_infinity`` from
2441 the restart.
2443 Returns
2444 -------
2445 Model
2446 New model with reduced structure, weights, and restart statistics.
2448 Raises
2449 ------
2450 ValueError
2451 If ``restart_parameters`` is not loaded, if any target value would
2452 expand the model (use :meth:`augment` instead), if a ``has_q_*``
2453 flag is set to ``True``, or if ``charge_head=True`` on a model
2454 without charges.
2455 """
2456 # --- Resolve target values ---
2457 new_n_neuron = n_neuron if n_neuron is not None else self.n_neuron
2458 new_l_max_4b = l_max_4b if l_max_4b is not None else self.l_max_4b
2459 new_l_max_5b = l_max_5b if l_max_5b is not None else self.l_max_5b
2460 new_has_q_112 = 0 if has_q_112 is False else self.has_q_112
2461 new_has_q_123 = 0 if has_q_123 is False else self.has_q_123
2462 new_has_q_233 = 0 if has_q_233 is False else self.has_q_233
2463 new_has_q_134 = 0 if has_q_134 is False else self.has_q_134
2465 # --- Validate ---
2466 if self.version not in (3, 4):
2467 raise ValueError(
2468 f'prune() only supports NEP versions 3 and 4; got version {self.version}.'
2469 )
2470 if new_n_neuron > self.n_neuron:
2471 raise ValueError(
2472 f'n_neuron ({new_n_neuron}) must be <= current n_neuron ({self.n_neuron}); '
2473 'use augment() to increase.'
2474 )
2475 if new_l_max_4b > self.l_max_4b:
2476 raise ValueError(
2477 f'l_max_4b ({new_l_max_4b}) must be <= current l_max_4b ({self.l_max_4b}); '
2478 'use augment() to increase.'
2479 )
2480 if new_l_max_5b > self.l_max_5b:
2481 raise ValueError(
2482 f'l_max_5b ({new_l_max_5b}) must be <= current l_max_5b ({self.l_max_5b}); '
2483 'use augment() to increase.'
2484 )
2485 for flag_val, name in [
2486 (has_q_112, 'has_q_112'), (has_q_123, 'has_q_123'),
2487 (has_q_233, 'has_q_233'), (has_q_134, 'has_q_134')
2488 ]:
2489 if flag_val is True:
2490 raise ValueError(
2491 f'Cannot enable {name} via prune(); '
2492 'use augment() to enable descriptor terms.'
2493 )
2494 if charge_head and self.model_type != 'potential_with_charges':
2495 raise ValueError(
2496 f'charge_head=True requires model_type="potential_with_charges"; '
2497 f'got "{self.model_type}".'
2498 )
2499 if self.restart_parameters is None:
2500 raise ValueError(
2501 'restart_parameters must be loaded before calling prune(). '
2502 'Pass restart_file= to read_model() or call model.read_restart() first.'
2503 )
2505 new = copy.deepcopy(self)
2506 keys = self.types if self.version in (4, 5) else ['all_species']
2508 # Step 1: Neuron pruning — keep the most important neurons
2509 if new_n_neuron < self.n_neuron:
2510 importances = []
2511 for s in keys:
2512 w0 = self.ann_parameters[s]['w0'] # (n_neuron, n_desc)
2513 w1_flat = self.ann_parameters[s]['w1'].ravel()
2514 if 'w1_charge' in self.ann_parameters[s]:
2515 output_norm = np.abs(w1_flat) + np.abs(self.ann_parameters[s]['w1_charge'])
2516 else:
2517 output_norm = np.abs(w1_flat)
2518 importances.append(np.linalg.norm(w0, axis=1) * output_norm)
2520 keep_idx = np.sort(np.argsort(np.mean(importances, axis=0))[-new_n_neuron:])
2522 for s in keys:
2523 new.ann_parameters[s]['w0'] = new.ann_parameters[s]['w0'][keep_idx, :]
2524 new.ann_parameters[s]['b0'] = new.ann_parameters[s]['b0'][keep_idx, :]
2525 w1 = new.ann_parameters[s]['w1']
2526 new.ann_parameters[s]['w1'] = w1[:, keep_idx] if w1.ndim == 2 else w1[keep_idx]
2527 if 'w1_charge' in new.ann_parameters[s]:
2528 new.ann_parameters[s]['w1_charge'] = (
2529 new.ann_parameters[s]['w1_charge'][keep_idx]
2530 )
2531 for pk in ['ann_mu', 'ann_sigma']:
2532 rp = new.restart_parameters[pk][s]
2533 rp['w0'] = rp['w0'][keep_idx, :]
2534 rp['b0'] = rp['b0'][keep_idx, :]
2535 w1 = rp['w1']
2536 rp['w1'] = w1[:, keep_idx] if w1.ndim == 2 else w1[keep_idx]
2537 if 'w1_charge' in rp:
2538 rp['w1_charge'] = rp['w1_charge'][keep_idx]
2540 # Step 2: Descriptor column pruning (disabling higher-body terms)
2541 n_per = self.n_max_angular + 1
2542 hb_terms = [
2543 (self.l_max_4b, new_l_max_4b),
2544 (self.l_max_5b, new_l_max_5b),
2545 (self.has_q_112, new_has_q_112),
2546 (self.has_q_123, new_has_q_123),
2547 (self.has_q_233, new_has_q_233),
2548 (self.has_q_134, new_has_q_134),
2549 ]
2550 keep_cols = list(range(self.n_descriptor_radial + n_per * self.l_max_3b))
2551 col_offset = len(keep_cols)
2552 for old_val, new_val in hb_terms:
2553 if old_val > 0:
2554 if new_val > 0:
2555 keep_cols.extend(range(col_offset, col_offset + n_per))
2556 col_offset += n_per
2558 old_n_desc = self.n_descriptor_radial + self.n_descriptor_angular
2559 if len(keep_cols) < old_n_desc:
2560 keep_cols = np.array(keep_cols, dtype=int)
2561 for s in keys:
2562 new.ann_parameters[s]['w0'] = new.ann_parameters[s]['w0'][:, keep_cols]
2563 for pk in ['ann_mu', 'ann_sigma']:
2564 rp = new.restart_parameters[pk][s]
2565 rp['w0'] = rp['w0'][:, keep_cols]
2566 new.q_scaler = [new.q_scaler[i] for i in keep_cols]
2568 # Step 3: Charge head removal
2569 if charge_head:
2570 new.model_type = 'potential'
2571 new.charge_mode = 0
2572 new.sqrt_epsilon_infinity = None
2573 for s in keys:
2574 w1 = new.ann_parameters[s]['w1'] # 1D (n_neuron,)
2575 new.ann_parameters[s]['w1'] = w1.reshape(1, -1)
2576 del new.ann_parameters[s]['w1_charge']
2577 for pk in ['ann_mu', 'ann_sigma']:
2578 rp = new.restart_parameters[pk][s]
2579 rp['w1'] = rp['w1'].reshape(1, -1)
2580 del rp['w1_charge']
2581 del new.restart_parameters['ann_mu']['sqrt_epsilon_infinity']
2582 del new.restart_parameters['ann_sigma']['sqrt_epsilon_infinity']
2584 # Step 4: Update header fields
2585 new.n_neuron = new_n_neuron
2586 new.l_max_4b = new_l_max_4b
2587 new.l_max_5b = new_l_max_5b
2588 new.has_q_112 = new_has_q_112
2589 new.has_q_123 = new_has_q_123
2590 new.has_q_233 = new_has_q_233
2591 new.has_q_134 = new_has_q_134
2593 new_l_max_enh = (self.l_max_3b
2594 + (new_l_max_4b > 0) + (new_l_max_5b > 0)
2595 + (new_has_q_112 > 0) + (new_has_q_123 > 0) + (new_has_q_233 > 0)
2596 + (new_has_q_134 > 0))
2597 new.n_descriptor_angular = (self.n_max_angular + 1) * new_l_max_enh
2599 # Step 5: Recalculate parameter counts
2600 _recalculate_parameter_counts(new)
2602 return new
2605def read_model(filename: str, restart_file: str = None) -> Model:
2606 """Parses a file in ``nep.txt`` format and returns the
2607 content in the form of a :class:`Model <calorine.nep.model.Model>`
2608 object.
2610 Parameters
2611 ----------
2612 filename
2613 Input file name.
2614 restart_file
2615 If provided, also read restart parameters from this file in
2616 `nep.restart` format and attach them to the returned model.
2617 Defaults to None.
2618 """
2619 data, parameters = _get_nep_contents(filename)
2621 # sanity checks
2622 for fld in ['version', 'types', 'model_type', 'cutoff', 'basis_size', 'n_max', 'l_max', 'ANN']:
2623 if fld not in data:
2624 raise ValueError(f'Invalid model file; {fld} line is missing')
2625 if data['version'] not in [3, 4, 5]:
2626 raise ValueError('Invalid model file; only NEP versions 3, 4 and 5 are currently supported')
2628 # split up zbl tuple (optional typewise cutoff factor as a third entry)
2629 if 'zbl' in data:
2630 if len(data['zbl']) == 3:
2631 data['zbl_typewise_cutoff_factor'] = data['zbl'][2]
2632 data['zbl'] = data['zbl'][:2]
2633 elif len(data['zbl']) != 2:
2634 raise ValueError(
2635 f'Invalid model file; zbl line must have 2 or 3 entries, got {len(data["zbl"])}'
2636 )
2638 # split up cutoff tuple
2639 N_types = len(data['types'])
2640 # Either global cutoffs + max neighbirs, or typewise cutoffs + max_neighbors
2641 if len(data['cutoff']) not in [4, 2*N_types+2]:
2642 raise ValueError(
2643 'Invalid model file; cutoff line must have 4 entries (global cutoffs) or '
2644 f'{2*N_types+2} entries (typewise cutoffs for {N_types} types), '
2645 f'got {len(data["cutoff"])}'
2646 )
2647 if not all(np.isfinite(data['cutoff'])):
2648 raise ValueError('Invalid model file; cutoff values must be finite')
2649 data['max_neighbors_radial'] = int(data['cutoff'][-2])
2650 data['max_neighbors_angular'] = int(data['cutoff'][-1])
2651 if len(data['cutoff']) == 2*N_types+2:
2652 # Typewise cutoffs: radial are even, angular are odd
2653 data['radial_cutoff'] = [data['cutoff'][i*2] for i in range(N_types)]
2654 data['angular_cutoff'] = [data['cutoff'][i*2+1] for i in range(N_types)]
2655 else:
2656 data['radial_cutoff'] = data['cutoff'][0]
2657 data['angular_cutoff'] = data['cutoff'][1]
2658 del data['cutoff']
2660 # split up basis_size tuple
2661 if len(data['basis_size']) != 2:
2662 raise ValueError(
2663 f'Invalid model file; basis_size line must have 2 entries, '
2664 f'got {len(data["basis_size"])}'
2665 )
2666 data['n_basis_radial'] = data['basis_size'][0]
2667 data['n_basis_angular'] = data['basis_size'][1]
2668 del data['basis_size']
2670 # split up n_max tuple
2671 if len(data['n_max']) != 2:
2672 raise ValueError(
2673 f'Invalid model file; n_max line must have 2 entries, got {len(data["n_max"])}'
2674 )
2675 data['n_max_radial'] = data['n_max'][0]
2676 data['n_max_angular'] = data['n_max'][1]
2677 del data['n_max']
2679 # split up nl_max tuple
2680 len_l = len(data['l_max'])
2681 if len_l not in [1, 2, 3, 4, 5, 6, 7]:
2682 raise ValueError(
2683 f'Invalid model file; l_max line must have between 1 and 7 entries, got {len_l}'
2684 )
2685 data['l_max_3b'] = data['l_max'][0]
2686 data['l_max_4b'] = data['l_max'][1] if len_l > 1 else 0
2687 data['l_max_5b'] = data['l_max'][2] if len_l > 2 else 0
2688 data['has_q_112'] = data['l_max'][3] if len_l > 3 else 0
2689 data['has_q_123'] = data['l_max'][4] if len_l > 4 else 0
2690 data['has_q_233'] = data['l_max'][5] if len_l > 5 else 0
2691 data['has_q_134'] = data['l_max'][6] if len_l > 6 else 0
2692 del data['l_max']
2694 # compute dimensions of descriptor components
2695 data['n_descriptor_radial'] = data['n_max_radial'] + 1
2696 l_max_enh = (data['l_max_3b']
2697 + (data['l_max_4b'] > 0)
2698 + (data['l_max_5b'] > 0)
2699 + (data['has_q_112'] > 0)
2700 + (data['has_q_123'] > 0)
2701 + (data['has_q_233'] > 0)
2702 + (data['has_q_134'] > 0))
2703 data['n_descriptor_angular'] = (data['n_max_angular'] + 1) * l_max_enh
2704 n_descriptor = data['n_descriptor_radial'] + data['n_descriptor_angular']
2706 is_charged_model = data['model_type'] == 'potential_with_charges'
2707 # compute number of parameters
2708 data['n_neuron'] = data['ANN'][0]
2709 del data['ANN']
2710 n_types = len(data['types'])
2711 # NEP4 and NEP5 have one hidden layer per atomic species, NEP3 a single shared one
2712 n = n_types if data['version'] in (4, 5) else 1
2713 n_output_biases = _number_of_output_biases(data['version'], n_types, is_charged_model)
2715 n_ann_input_weights = (n_descriptor + 1) * data['n_neuron'] # weights + bias
2716 n_ann_output_weights = 2*data['n_neuron'] if is_charged_model else data['n_neuron'] # weights
2717 n_ann_parameters = (
2718 n_ann_input_weights + n_ann_output_weights
2719 ) * n + n_output_biases
2721 n_descriptor_weights = n_types**2 * (
2722 (data['n_max_radial'] + 1) * (data['n_basis_radial'] + 1)
2723 + (data['n_max_angular'] + 1) * (data['n_basis_angular'] + 1)
2724 )
2725 data['n_parameters'] = n_ann_parameters + n_descriptor_weights + n_descriptor
2726 is_polarizability_model = data['model_type'] == 'polarizability'
2727 if data['n_parameters'] + n_ann_parameters == len(parameters):
2728 data['n_parameters'] += n_ann_parameters
2729 if not is_polarizability_model:
2730 raise ValueError(
2731 'Model is not labelled as a polarizability model, but the number of '
2732 'parameters matches a polarizability model.\n'
2733 'If this is a polarizability model trained with GPUMD <=v3.8, please '
2734 'modify the header in the nep.txt file to enable parsing '
2735 f'`nep{data["version"]}_polarizability`.\n'
2736 )
2737 if len(parameters) < data['n_parameters']:
2738 raise ValueError(
2739 'Invalid model file; expected '
2740 f'{data["n_parameters"]} parameter values, found {len(parameters)} '
2741 '-- file may be truncated'
2742 )
2743 elif len(parameters) > data['n_parameters']:
2744 raise ValueError(
2745 'Invalid model file; expected '
2746 f'{data["n_parameters"]} parameter values, found {len(parameters)} '
2747 '-- file may contain extra or corrupted data'
2748 )
2749 data['n_ann_parameters'] = n_ann_parameters
2751 # split up parameters into the ANN weights, descriptor weights, and scaling parameters
2752 n1 = n_ann_parameters
2753 n1 *= 2 if is_polarizability_model else 1
2754 n2 = n1 + n_descriptor_weights
2755 data['ann_parameters'] = parameters[:n1]
2756 descriptor_weights = np.array(parameters[n1:n2])
2757 data['q_scaler'] = parameters[n2:]
2759 # add ann parameters to data dict
2760 ann_groups = data['types'] if data['version'] in (4, 5) else ['all_species']
2761 sorted_ann_parameters = _sort_ann_parameters(data['ann_parameters'],
2762 ann_groups,
2763 data['n_neuron'],
2764 n,
2765 data['version'],
2766 n_descriptor,
2767 is_polarizability_model,
2768 is_charged_model)
2770 data['ann_parameters'] = sorted_ann_parameters
2771 if 'sqrt_epsilon_infinity' in sorted_ann_parameters.keys():
2772 data['sqrt_epsilon_infinity'] = sorted_ann_parameters['sqrt_epsilon_infinity']
2773 sorted_ann_parameters.pop('sqrt_epsilon_infinity')
2774 data['ann_parameters'] = sorted_ann_parameters
2776 # add descriptors to data dict
2777 data['n_descriptor_parameters'] = len(descriptor_weights)
2778 radial, angular = _sort_descriptor_parameters(descriptor_weights,
2779 data['types'],
2780 data['n_max_radial'],
2781 data['n_basis_radial'],
2782 data['n_max_angular'],
2783 data['n_basis_angular'])
2784 data['radial_descriptor_weights'] = radial
2785 data['angular_descriptor_weights'] = angular
2787 model = Model(**data)
2788 if restart_file is not None:
2789 model.read_restart(restart_file)
2790 return model