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