import copy
import re
from dataclasses import dataclass
from itertools import product
from typing import Any, Iterable
from warnings import warn
import numpy as np
from calorine.nep.io import _write_nepfile_to_path
NetworkWeights = dict[str, dict[str, np.ndarray]]
DescriptorWeights = dict[tuple[str, str], np.ndarray]
RestartParameters = dict[str, dict[str, dict[str, np.ndarray]]]
# the `model_type` value that each model type corresponds to in `nep.in`
_MODEL_TYPE_TO_INT = {
'potential': 0,
'potential_with_charges': 0,
'dipole': 1,
'polarizability': 2,
}
# The `nep.in` keywords that describe the model rather than the training run. These are the
# keywords that the `nep` executable checks against the `nep.txt` header (see
# `Parameters::compare_with_nep_txt` in `src/main_nep/parameters.cu` of GPUMD), plus `mode`,
# which it accepts as a synonym of `model_type`. `Model.write_nepfile` supplies all of them
# from the model and discards any value given for them, so that a value belonging to another
# model cannot reach the file.
_MODEL_PARAMETERS = (
'version',
'model_type',
'mode',
'type',
'cutoff',
'n_max',
'basis_size',
'l_max',
'neuron',
'zbl',
'use_typewise_cutoff_zbl',
'charge_mode',
)
# the factor that the `nep` executable uses when `use_typewise_cutoff_zbl` carries no value
_TYPEWISE_CUTOFF_ZBL_FACTOR_DEFAULT = 0.7
def _nepfile_tokens(value: Any) -> list[str]:
"""Returns the value of a ``nep.in`` keyword as the list of whitespace separated tokens
that it is written as, so that values read from a file can be compared with values taken
from a model irrespective of how either is represented in Python.
Parameters
----------
value
Value of a single ``nep.in`` keyword.
"""
if isinstance(value, str):
return value.split()
if isinstance(value, Iterable):
return [f'{v}' for v in value]
return [f'{value}']
def _same_nepfile_value(supplied: Any, model_value: Any) -> bool:
"""Returns whether two values of a ``nep.in`` keyword agree, comparing token by token and
numerically where both tokens are numbers, so that e.g. ``6`` and ``6.0`` agree.
Parameters
----------
supplied
Value to check, typically read from an existing ``nep.in`` file.
model_value
Value taken from the model, or ``None`` if the model does not set this keyword.
"""
if model_value is None:
return False
left = _nepfile_tokens(supplied)
right = _nepfile_tokens(model_value)
if len(left) != len(right):
return False
for token_left, token_right in zip(left, right):
if token_left == token_right:
continue
try:
value_left, value_right = float(token_left), float(token_right)
except ValueError:
return False
if abs(value_left - value_right) > 1e-6 * (abs(value_left) + abs(value_right)):
return False
return True
def _format_nepfile_value(value: Any) -> str:
"""Returns the value of a ``nep.in`` keyword as it is written to file, for use in messages.
Parameters
----------
value
Value of a single ``nep.in`` keyword.
"""
return ' '.join(_nepfile_tokens(value))
def _get_restart_contents(filename: str) -> tuple[list[float], list[float]]:
"""Parses a ``nep.restart`` file, and returns an unformatted list of the
mean and standard deviation for all model parameters.
Intended to be used by the py:meth:`~Model.read_restart` function.
Parameters
----------
filename
input file name
"""
mu = [] # Mean
sigma = [] # Standard deviation
with open(filename) as f:
for k, line in enumerate(f.readlines()):
flds = line.split()
if len(flds) == 0:
raise IOError(f'Empty line number {k}')
if len(flds) == 2:
mu.append(float(flds[0]))
sigma.append(float(flds[1]))
else:
raise IOError(f'Failed to parse line {k} from {filename}')
return mu, sigma
def _get_model_type(first_row: list[str]) -> str:
"""Parses a the first row of a ``nep.txt`` file, and returns the
type of NEP model. Available types are `potential`, `potential_with_charges`,
`dipole`, and `polarizability`.
Parameters
----------
first_row
First row of a NEP file, split by white space.
"""
model_type = first_row[0]
if 'charge' in model_type:
return 'potential_with_charges'
elif 'dipole' in model_type:
return 'dipole'
elif 'polarizability' in model_type:
return 'polarizability'
return 'potential'
def _get_charge_mode(model_type_token: str) -> int:
"""Parses the charge_mode (0, 1, or 2) from the first token of a ``nep.txt``
header line, e.g. ``nep4_charge1``, ``nep4_zbl_charge2``. Returns 0 for
non-charge models.
Parameters
----------
model_type_token
First token of the first row of a NEP file (``flds[0]``).
"""
match = re.search(r'charge(\d+)', model_type_token)
return int(match.group(1)) if match else 0
def _get_nep_contents(filename: str) -> tuple[dict, list[float]]:
"""Parses a ``nep.txt`` file, and returns a dict describing the header
and an unformatted list of all model parameters.
Intended to be used by the :func:`read_model <calorine.nep.read_model>` function.
Parameters
----------
filename
input file name
"""
# parse file and split header and parameters
header = []
parameters = []
nheader = 5 # 5 rows for NEP2, 6-7 rows for NEP3 onwards
base_line = 3
with open(filename) as f:
for k, line in enumerate(f.readlines()):
flds = line.split()
if len(flds) == 0:
raise IOError(f'Empty line number {k}')
if k == 0 and 'zbl' in flds[0]:
base_line += 1
nheader += 1
if k == base_line and 'basis_size' in flds[0]:
# Introduced in nep.txt after GPUMD v3.2
nheader += 1
if k < nheader:
header.append(tuple(flds))
elif len(flds) == 1:
parameters.append(float(flds[0]))
else:
raise IOError(f'Failed to parse line {k} from {filename}')
# compile data from the header into a dict
data = {}
for flds in header:
if flds[0] in ['cutoff', 'zbl']:
data[flds[0]] = tuple(map(float, flds[1:]))
elif flds[0] in ['n_max', 'l_max', 'ANN', 'basis_size']:
data[flds[0]] = tuple(map(int, flds[1:]))
elif flds[0].startswith('nep'):
version = flds[0].replace('nep', '').split('_')[0]
version = int(version)
data['version'] = version
data['types'] = flds[2:]
data['model_type'] = _get_model_type(flds)
data['charge_mode'] = _get_charge_mode(flds[0])
else:
raise ValueError(f'Unknown field: {flds[0]}')
return data, parameters
def _sort_descriptor_parameters(parameters: list[float],
types: list[str],
n_max_radial: int,
n_basis_radial: int,
n_max_angular: int,
n_basis_angular: int) -> tuple[DescriptorWeights,
DescriptorWeights]:
"""Reads a list of descriptors parameters and sorts them into two
appropriately structured `dicts`, one for radial and one for angular descriptor weights.
Intended to be used by the :func:`read_model <calorine.nep.read_model>` function.
"""
# split up descriptor by chemical species and radial/angular
n_types = len(types)
n = len(parameters) // (n_types * n_types)
m = (n_max_radial + 1) * (n_basis_radial + 1)
descriptor_weights = parameters.reshape((n, n_types * n_types)).T
descriptor_weights_radial = descriptor_weights[:, :m]
descriptor_weights_angular = descriptor_weights[:, m:]
# add descriptors to data dict
radial_descriptor_weights = {}
angular_descriptor_weights = {}
m = -1
for i, j in product(range(n_types), repeat=2):
m += 1
s1, s2 = types[i], types[j]
radial_descriptor_weights[(s1, s2)] = descriptor_weights_radial[m, :].reshape(
(n_max_radial + 1, n_basis_radial + 1)
)
angular_descriptor_weights[(s1, s2)] = descriptor_weights_angular[m, :].reshape(
(n_max_angular + 1, n_basis_angular + 1)
)
return radial_descriptor_weights, angular_descriptor_weights
def _number_of_output_biases(version: int,
n_types: int,
is_model_with_charges: bool) -> int:
"""Returns the number of output-layer bias values in a single network pass.
This is the trailing block of a network pass in ``nep.txt`` and ``nep.restart``,
counted per pass rather than per file. A polarizability model runs two passes and
therefore carries twice this many bias values in total, while every other model
type carries exactly this many.
Parameters
----------
version
NEP version (3, 4, or 5).
n_types
Number of atomic species in the model.
is_model_with_charges
Whether the model has a charge output head, i.e. whether its type is
``potential_with_charges``.
Returns
-------
int
Number of bias values per network pass: 2 for a model with charges, since
``sqrt_epsilon_infinity`` precedes the global bias, ``1 + n_types`` for NEP5,
which adds one bias per species to the global one, and 1 otherwise.
Example
-------
>>> from calorine.nep.model import _number_of_output_biases
>>> _number_of_output_biases(4, 2, False)
1
>>> _number_of_output_biases(5, 2, False)
3
"""
if is_model_with_charges:
return 2
if version == 5:
return 1 + n_types
return 1
def _sort_ann_parameters(parameters: list[float],
ann_groupings: list[str],
n_neuron: int,
n_networks: int,
version: int,
n_descriptor: int,
is_polarizability_model: bool,
is_model_with_charges: bool
) -> NetworkWeights:
"""Reads a list of model parameters and sorts them into an appropriately structured `dict`.
Intended to be used by the :func:`read_model <calorine.nep.read_model>` function
and by :meth:`Model.read_restart <calorine.nep.Model.read_restart>`.
"""
# ann_groupings holds one entry per species for NEP4 and NEP5, and a single shared
# entry for NEP3, which has no per-species biases for its length to bear on
n_bias = _number_of_output_biases(version, len(ann_groupings), is_model_with_charges)
n_ann_input_weights = (n_descriptor + 1) * n_neuron # weights + bias
n_ann_output_weights = 2*n_neuron if is_model_with_charges else n_neuron # only weights
n_ann_parameters = (
n_ann_input_weights + n_ann_output_weights
) * n_networks + n_bias
# Group ANN parameters
pars = {}
n1 = 0
n_network_params = n_ann_input_weights + n_ann_output_weights # except last bias(es)
n_count = 2 if is_polarizability_model else 1
n_outputs = 2 if is_model_with_charges else 1
for count in range(n_count):
# if polarizability model, all parameters including bias are repeated
# need to offset n1 by +1 to handle bias
n1 += count
for s in ann_groupings:
# Get the parameters for the ANN; in the case of NEP4, there is effectively
# one network per atomic species.
ann_parameters = parameters[n1 : n1 + n_network_params]
ann_input_weights = ann_parameters[:n_ann_input_weights]
w0 = np.zeros((n_neuron, n_descriptor))
w0[...] = np.nan
b0 = np.zeros((n_neuron, 1))
b0[...] = np.nan
for n in range(n_neuron):
for nu in range(n_descriptor):
w0[n, nu] = ann_input_weights[n * n_descriptor + nu]
b0[:, 0] = ann_input_weights[n_neuron * n_descriptor :]
assert np.all(
w0.shape == (n_neuron, n_descriptor)
), f'w0 has invalid shape for key {s}; please submit a bug report'
assert np.all(
b0.shape == (n_neuron, 1)
), f'b0 has invalid shape for key {s}; please submit a bug report'
assert not np.any(
np.isnan(w0)
), f'some weights in w0 are nan for key {s}; please submit a bug report'
assert not np.any(
np.isnan(b0)
), f'some weights in b0 are nan for key {s}; please submit a bug report'
ann_output_weights = ann_parameters[
n_ann_input_weights : n_ann_input_weights + n_ann_output_weights
]
w1 = np.zeros((1, n_neuron * n_outputs))
w1[0, :] = ann_output_weights[:]
assert np.all(
w1.shape == (1, n_neuron * n_outputs)
), f'w1 has invalid shape for key {s}; please submit a bug report'
assert not np.any(
np.isnan(w1)
), f'some weights in w1 are nan for key {s}; please submit a bug report'
if count == 0 and n_outputs == 1:
pars[s] = dict(w0=w0, b0=b0, w1=w1)
elif count == 0 and n_outputs == 2:
pars[s] = dict(w0=w0, b0=b0, w1=w1[0, :n_neuron], w1_charge=w1[0, n_neuron:])
else:
pars[s].update({'w0_polar': w0, 'b0_polar': b0, 'w1_polar': w1})
# Jump to bias
n1 += n_network_params
if version == 5 and not is_model_with_charges:
# NEP5 models additionally have one bias term per species, which is
# what _number_of_output_biases accounts for. Currently NEP5 only
# exists for potential models, but we'll keep it here in case it gets
# added down the line.
bias_label = 'b1' if count == 0 else 'b1_polar'
pars[s][bias_label] = parameters[n1]
n1 += 1
# For NEP3 and NEP4 we only have one bias.
# For NEP4 with charges we have two biases.
# For NEP5 we have one bias per species, and one global.
if count == 0 and n_outputs == 1:
pars['b1'] = parameters[n1]
elif count == 0 and n_outputs == 2:
pars['sqrt_epsilon_infinity'] = parameters[n1]
pars['b1'] = parameters[n1+1]
else:
pars['b1_polar'] = parameters[n1]
sum = 0
for s in pars.keys():
if s.startswith('b1') or s.startswith('sqrt'):
sum += 1
else:
sum += np.sum([np.array(p).size for p in pars[s].values()])
assert sum == n_ann_parameters * n_count, (
'Inconsistent number of parameters accounted for; please submit a bug report\n'
f'{sum} != {n_ann_parameters}'
)
return pars
def _adaptive_sigma(mu_arr, sigma_factor: float, sigma_floor: float) -> np.ndarray:
"""Return adaptive SNES sigma: ``max(sigma_floor, sigma_factor * |mu|)``."""
return np.maximum(sigma_floor, sigma_factor * np.abs(mu_arr))
def _format_header_float(value: float) -> str:
"""Format *value* the way GPUMD writes the float-valued header fields of ``nep.txt``, which
it does with ``%g`` (see ``write_nep_txt`` in ``src/main_nep/fitness.cu``). A cutoff of 6 is
therefore written as ``6`` rather than ``6.0``.
``%g`` keeps six significant digits, so it is used only where it reproduces the value exactly.
Anything carrying more digits than GPUMD would have written is emitted in full rather than
truncated, which matters for a file that came from somewhere else.
"""
text = f'{float(value):g}'
return text if float(text) == float(value) else repr(float(value))
_RESTART_COMPONENTS = ('network_weights', 'descriptor', 'charge_head')
def _restart_leaves(model, restart_params, component=None, species=None):
"""Yield ``(mu, sigma_container, sigma_key)`` for every leaf entry of
*restart_params* that matches the requested *component*/*species* filters.
``sigma_container[sigma_key]`` is either a numpy array or a scalar float;
together with ``mu`` (same shape/type) this is everything
:func:`_apply_sigma_strategy` needs to read and update one leaf.
*component* selects among ``'network_weights'`` (w0, b0, w1, and the global
b1 bias), ``'descriptor'`` (radial/angular descriptor weight pairs), and
``'charge_head'`` (w1_charge and sqrt_epsilon_infinity). ``None`` means all
three. *species* restricts per-species entries (and descriptor pairs
involving that species) to the given species; global scalars (b1,
sqrt_epsilon_infinity) are only included when *species* is ``None``, since
they are not owned by a single species.
"""
if component is None:
wanted = set(_RESTART_COMPONENTS)
else:
wanted = {component} if isinstance(component, str) else set(component)
unknown = wanted - set(_RESTART_COMPONENTS)
if unknown:
raise ValueError(
f'Unknown component(s) {sorted(unknown)}; expected any of '
f'{_RESTART_COMPONENTS}'
)
if species is None:
species_filter = None
else:
species_filter = {species} if isinstance(species, str) else set(species)
keys = model.types if model.version in (4, 5) else ['all_species']
ann_mu, ann_sigma = restart_params['ann_mu'], restart_params['ann_sigma']
if 'network_weights' in wanted:
for s in keys:
if species_filter is not None and s not in species_filter:
continue
# b1 appears here only for NEP5, which carries a per-species bias alongside
# the global one handled below. For every other version it is global only.
for pname in ('w0', 'b0', 'w1', 'b1', 'w0_polar', 'b0_polar', 'w1_polar', 'b1_polar'):
if pname in ann_mu[s]:
yield ann_mu[s][pname], ann_sigma[s], pname
if species_filter is None:
for pname in ('b1', 'b1_polar'):
if pname in ann_mu:
yield ann_mu[pname], ann_sigma, pname
if 'charge_head' in wanted:
for s in keys:
if species_filter is not None and s not in species_filter:
continue
if 'w1_charge' in ann_mu[s]:
yield ann_mu[s]['w1_charge'], ann_sigma[s], 'w1_charge'
if species_filter is None and 'sqrt_epsilon_infinity' in ann_mu:
yield ann_mu['sqrt_epsilon_infinity'], ann_sigma, 'sqrt_epsilon_infinity'
if 'descriptor' in wanted:
for desc_type in ('radial', 'angular'):
mu_dict = restart_params[f'{desc_type}_descriptor_mu']
sigma_dict = restart_params[f'{desc_type}_descriptor_sigma']
for pair, mu_val in mu_dict.items():
if species_filter is not None and not (species_filter & set(pair)):
continue
yield mu_val, sigma_dict, pair
_CHARGE_HEAD_PARAMETERS = ('w1_charge', 'sqrt_epsilon_infinity')
def _parameter_component(name: str) -> str:
"""Return the restart component that the ANN parameter head *name* belongs to.
Mirrors the grouping that :func:`_restart_leaves` applies to a loaded restart, for code
that walks the parameters of the model instead.
"""
base = name[:-len('_polar')] if name.endswith('_polar') else name
if base in _CHARGE_HEAD_PARAMETERS:
return 'charge_head'
return 'network_weights'
def _leaf_get(container, key):
"""Read a parameter leaf from *container*, which is either a dict or the model itself."""
return container[key] if isinstance(container, dict) else getattr(container, key)
def _leaf_set(container, key, value):
"""Write a parameter leaf to *container*, which is either a dict or the model itself."""
if isinstance(container, dict):
container[key] = value
else:
setattr(container, key, value)
def _model_parameter_leaves(model, component=None, species=None):
"""Yield ``(container, key, species, name)`` for every parameter of *model*, so that the
leaf can be read via :func:`_leaf_get` and written via :func:`_leaf_set`.
This is the counterpart of :func:`_restart_leaves` for the parameters of the model
itself. *species* is the species that owns the leaf, or ``None`` for the global scalars
and the descriptor pairs, and *name* is the name of the parameter head as used in
``restart_parameters``, so that both can be used to look up the matching restart entry.
The *component*/*species* filters work exactly as in :func:`_restart_leaves`: global
scalars are only included when *species* is ``None``, and a descriptor pair is included
when either of its two species is selected.
"""
if component is None:
wanted = set(_RESTART_COMPONENTS)
else:
wanted = {component} if isinstance(component, str) else set(component)
unknown = wanted - set(_RESTART_COMPONENTS)
if unknown:
raise ValueError(
f'Unknown component(s) {sorted(unknown)}; expected any of '
f'{_RESTART_COMPONENTS}'
)
if species is None:
species_filter = None
else:
species_filter = {species} if isinstance(species, str) else set(species)
keys = model.types if model.version in (4, 5) else ['all_species']
for s in keys:
if species_filter is not None and s not in species_filter:
continue
for name in model.ann_parameters[s]:
if _parameter_component(name) in wanted:
yield model.ann_parameters[s], name, s, name
if species_filter is None:
for name in model.ann_parameters:
if name in keys:
continue # per-species dictionaries, handled above
if _parameter_component(name) in wanted:
yield model.ann_parameters, name, None, name
if model.sqrt_epsilon_infinity is not None and 'charge_head' in wanted:
yield model, 'sqrt_epsilon_infinity', None, 'sqrt_epsilon_infinity'
if 'descriptor' in wanted:
for descriptor_type in ('radial', 'angular'):
weights = getattr(model, f'{descriptor_type}_descriptor_weights')
for pair in weights:
if species_filter is not None and not (species_filter & set(pair)):
continue
yield weights, pair, None, f'{descriptor_type}_descriptor'
def _apply_sigma_strategy(mu, sigma_container, sigma_key, strategy, target, rng, **kw):
"""Update ``sigma_container[sigma_key]`` in place at the positions selected
by *target* (``'unset'`` -> NaN entries, ``'set'`` -> non-NaN entries,
``'all'`` -> everything), using *mu* and *strategy* to compute new values.
"""
sigma = sigma_container[sigma_key]
if np.isscalar(sigma) or isinstance(sigma, (float, int)):
current = float(sigma)
is_unset = np.isnan(current)
apply_here = (
target == 'all' or (target == 'unset' and is_unset)
or (target == 'set' and not is_unset)
)
if not apply_here:
return
mu_val = float(mu)
if strategy == 'constant':
new_val = kw['value']
elif strategy == 'scale_mu':
new_val = float(_adaptive_sigma(np.array(mu_val), kw['factor'], kw['floor']))
elif strategy == 'scale_sigma':
if is_unset:
raise ValueError(
f"strategy='scale_sigma' requires an existing sigma value, but "
f'{sigma_key!r} is unset (NaN); set it first, e.g. with '
"target='unset'."
)
new_val = current * kw['factor']
elif strategy == 'uniform':
new_val = float(rng.uniform(kw['low'], kw['high']))
elif strategy == 'normal':
new_val = float(abs(rng.normal(kw['mean'], kw['std'])))
else:
raise ValueError(f'Unknown strategy {strategy!r}')
sigma_container[sigma_key] = float(new_val)
return
sigma_arr = sigma_container[sigma_key]
mu_arr = np.asarray(mu, dtype=float)
if target == 'unset':
mask = np.isnan(sigma_arr)
elif target == 'set':
mask = ~np.isnan(sigma_arr)
else:
mask = np.ones_like(sigma_arr, dtype=bool)
if not np.any(mask):
return
if strategy == 'constant':
sigma_arr[mask] = kw['value']
elif strategy == 'scale_mu':
sigma_arr[mask] = _adaptive_sigma(mu_arr[mask], kw['factor'], kw['floor'])
elif strategy == 'scale_sigma':
if np.any(np.isnan(sigma_arr[mask])):
raise ValueError(
f"strategy='scale_sigma' requires existing sigma values, but "
f'{sigma_key!r} has unset (NaN) entries within the selected '
"target; set them first, e.g. with target='unset'."
)
sigma_arr[mask] = sigma_arr[mask] * kw['factor']
elif strategy == 'uniform':
sigma_arr[mask] = rng.uniform(kw['low'], kw['high'], size=int(np.sum(mask)))
elif strategy == 'normal':
sigma_arr[mask] = np.abs(rng.normal(kw['mean'], kw['std'], size=int(np.sum(mask))))
else:
raise ValueError(f'Unknown strategy {strategy!r}')
def _draw_values(strategy, rng, size=None, **kw):
"""Return parameter values drawn according to *strategy*: a single float if *size* is
``None``, otherwise an array of *size* values.
Unlike the sigma strategies of :func:`_apply_sigma_strategy`, the drawn values are not
forced to be positive, since a parameter value may have either sign.
"""
if strategy == 'constant':
return float(kw['value']) if size is None else np.full(size, float(kw['value']))
if strategy == 'uniform':
return (float(rng.uniform(kw['low'], kw['high'])) if size is None
else rng.uniform(kw['low'], kw['high'], size=size))
if strategy == 'normal':
return (float(rng.normal(kw['mean'], kw['std'])) if size is None
else rng.normal(kw['mean'], kw['std'], size=size))
raise ValueError(f'Unknown strategy {strategy!r}')
def _restart_leaf_for(restart_parameters, species, name, key):
"""Return ``(mu_container, sigma_container, restart_key)`` for the restart entry that
corresponds to the model parameter leaf described by *species*, *name*, and *key*, or
``None`` if the restart has no counterpart for it.
"""
if name.endswith('_descriptor'):
mu_container = restart_parameters[f'{name}_mu']
sigma_container = restart_parameters[f'{name}_sigma']
restart_key = key
elif species is None:
mu_container = restart_parameters['ann_mu']
sigma_container = restart_parameters['ann_sigma']
restart_key = name
else:
mu_container = restart_parameters['ann_mu'][species]
sigma_container = restart_parameters['ann_sigma'][species]
restart_key = name
if restart_key not in sigma_container or restart_key not in mu_container:
# A parameter head that the restart tree in hand does not carry. Whether such a
# parameter is new cannot be told from a sigma, so it is left alone. Every head of
# every shipped model has a counterpart, so this guards against a tree assembled
# by other means rather than against any model type in particular.
return None
return mu_container, sigma_container, restart_key
def _restart_sigma_counts(model, component=None) -> dict[str, int]:
"""Return the number of restart sigma entries of *model* as a dict with the keys
``'total'``, ``'frozen'`` (sigma of exactly zero) and ``'unset'`` (sigma of ``NaN``),
restricted to *component* if given.
The traversal goes through :func:`_restart_leaves`, so every parameter head
registered there is covered automatically. Scalar leaves (``b1``,
``sqrt_epsilon_infinity``) count as one entry each.
"""
total, frozen, unset = 0, 0, 0
for _, sigma_container, sigma_key in _restart_leaves(
model, model.restart_parameters, component
):
sigma = np.asarray(sigma_container[sigma_key], dtype=float)
total += int(sigma.size)
frozen += int(np.count_nonzero(sigma == 0.0))
unset += int(np.count_nonzero(np.isnan(sigma)))
return {'total': total, 'frozen': frozen, 'unset': unset}
def _model_parameter_counts(model) -> dict[str, dict[str, int]]:
"""Return ``{component: {'total': n}}`` for the current parameters of *model*.
Used when no restart is loaded, so there are no sigma values to split into frozen and
unset entries. The traversal goes through :func:`_model_parameter_leaves`, which also
drives :meth:`Model.initialize_parameters`, so the two cannot disagree on which
component a parameter head belongs to. The ``q_scaler`` entries are excluded, since they
are not fit parameters and have no restart counterpart.
"""
counts = {}
for component in _RESTART_COMPONENTS:
total = sum(int(np.asarray(_leaf_get(container, key)).size)
for container, key, _, _ in _model_parameter_leaves(model, component))
if total > 0:
counts[component] = {'total': total}
return counts
def _new_restart_parameters_from_model(model) -> RestartParameters:
"""Build a fresh restart-parameters dict from a model's current (trained)
parameters: ``mu`` is copied from the model, ``sigma`` is set to ``NaN``
everywhere (unset), to be filled in via :meth:`Model.set_restart_sigma`.
"""
keys = model.types if model.version in (4, 5) else ['all_species']
suffixes = ['', '_polar'] if model.model_type == 'polarizability' else ['']
ann_mu, ann_sigma = {}, {}
for s in keys:
params = model.ann_parameters[s]
mu_entry, sigma_entry = {}, {}
for suffix in suffixes:
for base in ('w0', 'b0', 'w1', 'w1_charge'):
pname = f'{base}{suffix}'
if pname in params:
arr = np.array(params[pname], dtype=float)
mu_entry[pname] = arr.copy()
sigma_entry[pname] = np.full(arr.shape, np.nan)
# NEP5 adds a per-species bias to the global one, and it is a scalar
# rather than an array, matching how the model itself stores it.
bias_name = f'b1{suffix}'
if bias_name in params:
mu_entry[bias_name] = float(params[bias_name])
sigma_entry[bias_name] = float('nan')
ann_mu[s] = mu_entry
ann_sigma[s] = sigma_entry
for suffix in suffixes:
b1_key = f'b1{suffix}'
if b1_key in model.ann_parameters:
ann_mu[b1_key] = float(model.ann_parameters[b1_key])
ann_sigma[b1_key] = float('nan')
if model.sqrt_epsilon_infinity is not None:
ann_mu['sqrt_epsilon_infinity'] = float(model.sqrt_epsilon_infinity)
ann_sigma['sqrt_epsilon_infinity'] = float('nan')
radial_mu = {
k: np.array(v, dtype=float).copy() for k, v in model.radial_descriptor_weights.items()
}
radial_sigma = {k: np.full(v.shape, np.nan) for k, v in radial_mu.items()}
angular_mu = {
k: np.array(v, dtype=float).copy() for k, v in model.angular_descriptor_weights.items()
}
angular_sigma = {k: np.full(v.shape, np.nan) for k, v in angular_mu.items()}
return {
'ann_mu': ann_mu,
'ann_sigma': ann_sigma,
'radial_descriptor_mu': radial_mu,
'radial_descriptor_sigma': radial_sigma,
'angular_descriptor_mu': angular_mu,
'angular_descriptor_sigma': angular_sigma,
}
def _recalculate_parameter_counts(new) -> None:
"""Recompute n_ann_parameters, n_descriptor_parameters, and n_parameters on *new*.
Reads all architectural state from *new* directly, so callers must update
new.n_neuron, new.n_descriptor_radial/angular, new.model_type, and new.types
before calling this function.
"""
n_types = len(new.types)
n_desc = new.n_descriptor_radial + new.n_descriptor_angular
is_charged = new.model_type == 'potential_with_charges'
n_networks = n_types if new.version in (4, 5) else 1
n_output_biases = _number_of_output_biases(new.version, n_types, is_charged)
n_ann_input_weights = (n_desc + 1) * new.n_neuron
n_ann_output_weights = 2 * new.n_neuron if is_charged else new.n_neuron
new.n_ann_parameters = (
n_ann_input_weights + n_ann_output_weights
) * n_networks + n_output_biases
new.n_descriptor_parameters = n_types ** 2 * (
(new.n_max_radial + 1) * (new.n_basis_radial + 1)
+ (new.n_max_angular + 1) * (new.n_basis_angular + 1)
)
new.n_parameters = new.n_ann_parameters + new.n_descriptor_parameters + n_desc
if new.model_type == 'polarizability':
new.n_parameters += new.n_ann_parameters
[docs]
@dataclass
class Model:
r"""Objects of this class represent a NEP model in a form suitable for
inspection and manipulation. Typically a :class:`Model` object is instantiated
by calling the :func:`read_model <calorine.nep.read_model>` function.
Attributes
----------
version : int
NEP version.
model_type: str
One of ``potential``, ``dipole`` or ``polarizability``.
types : tuple[str, ...]
Chemical species that this model represents.
radial_cutoff : float | list[float]
The radial cutoff parameter in Å.
Is a list of radial cutoffs ordered after ``types`` in the case of typewise cutoffs.
angular_cutoff : float | list[float]
The angular cutoff parameter in Å.
Is a list of angular cutoffs ordered after ``types`` in the case of typewise cutoffs.
max_neighbors_radial : int
Maximum number of neighbors in neighbor list for radial terms.
max_neighbors_angular : int
Maximum number of neighbors in neighbor list for angular terms.
zbl : tuple[float, float]
Inner and outer cutoff for transition to ZBL potential.
zbl_typewise_cutoff_factor : float
Optional typewise cutoff factor for the ZBL potential, corresponding to an
optional third value on the ``zbl`` line in ``nep.txt`` when
``use_typewise_cutoff_zbl`` is enabled during training. ``None`` if not set.
n_basis_radial : int
Number of radial basis functions :math:`n_\mathrm{basis}^\mathrm{R}`.
n_basis_angular : int
Number of angular basis functions :math:`n_\mathrm{basis}^\mathrm{A}`.
n_max_radial : int
Maximum order of Chebyshev polymonials included in
radial expansion :math:`n_\mathrm{max}^\mathrm{R}`.
n_max_angular : int
Maximum order of Chebyshev polymonials included in
angular expansion :math:`n_\mathrm{max}^\mathrm{A}`.
l_max_3b : int
Maximum expansion order for three-body terms :math:`l_\mathrm{max}^\mathrm{3b}`.
l_max_4b : int
Maximum expansion order for four-body terms :math:`l_\mathrm{max}^\mathrm{4b}`.
l_max_5b : int
Maximum expansion order for five-body terms :math:`l_\mathrm{max}^\mathrm{5b}`.
has_q_112 : int
Flag enabling the 5-body :math:`q_{112}` descriptor (0 or 1).
has_q_123 : int
Flag enabling the 5-body :math:`q_{123}` descriptor (0 or 1).
has_q_233 : int
Flag enabling the 5-body :math:`q_{233}` descriptor (0 or 1).
has_q_134 : int
Flag enabling the higher-body :math:`q_{134}` descriptor (0 or 1).
n_descriptor_radial : int
Dimension of radial part of descriptor.
n_descriptor_angular : int
Dimension of angular part of descriptor.
n_neuron : int
Number of neurons in hidden layer.
n_parameters : int
Total number of parameters including scalers (which are not fit parameters).
n_descriptor_parameters : int
Number of parameters in descriptor.
n_ann_parameters : int
Number of neural network weights.
ann_parameters : dict[tuple[str, dict[str, np.darray]]]
Neural network weights.
q_scaler : List[float]
Scaling parameters.
radial_descriptor_weights : dict[tuple[str, str], np.ndarray]
Radial descriptor weights by combination of species; the array for each combination
has dimensions of
:math:`(n_\mathrm{max}^\mathrm{R}+1) \times (n_\mathrm{basis}^\mathrm{R}+1)`.
angular_descriptor_weights : dict[tuple[str, str], np.ndarray]
Angular descriptor weights by combination of species; the array for each combination
has dimensions of
:math:`(n_\mathrm{max}^\mathrm{A}+1) \times (n_\mathrm{basis}^\mathrm{A}+1)`.
sqrt_epsilon_infinity : Optional[float]
Square root of epsilon infinity $\epsilon_\infty$ (only for NEP models with charges).
charge_mode : int
Charge algorithm variant for ``potential_with_charges`` models; 0 for
non-charge-aware models. 1 corresponds to a qNEP model including both real- and
reciprocal-space contributions. 2 corresponds to a qNEP model, including the
reciprocal-space contribution only.
restart_parameters : dict[str, dict[str, dict[str, np.ndarray]]]
NEP restart parameters. A nested dictionary that contains the mean (mu) and standard
deviation (sigma) for the ANN and descriptor parameters. Is set using the
py:meth:`~Model.read_restart` method. Defaults to None.
The state of the sigma values is summarized by the
:attr:`~Model.n_frozen_parameters` and :attr:`~Model.n_unset_parameters`
properties, and broken down per category by :attr:`~Model.parameter_counts`.
"""
version: int
model_type: str
types: tuple[str, ...]
radial_cutoff: float | list[float]
angular_cutoff: float | list[float]
n_basis_radial: int
n_basis_angular: int
n_max_radial: int
n_max_angular: int
l_max_3b: int
l_max_4b: int
l_max_5b: int
has_q_112: int
has_q_123: int
has_q_233: int
has_q_134: int
n_descriptor_radial: int
n_descriptor_angular: int
n_neuron: int
n_parameters: int
n_descriptor_parameters: int
n_ann_parameters: int
ann_parameters: NetworkWeights
q_scaler: list[float]
radial_descriptor_weights: DescriptorWeights
angular_descriptor_weights: DescriptorWeights
sqrt_epsilon_infinity: float = None
charge_mode: int = 0
restart_parameters: RestartParameters = None
zbl: tuple[float, float] = None
zbl_typewise_cutoff_factor: float = None
max_neighbors_radial: int = None
max_neighbors_angular: int = None
_special_fields = [
'ann_parameters',
'q_scaler',
'radial_descriptor_weights',
'angular_descriptor_weights',
]
def __str__(self) -> str:
s = []
for fld in self.__dataclass_fields__:
if fld not in self._special_fields:
value = getattr(self, fld)
if fld == 'restart_parameters':
value = self._restart_availability()
s += [f'{fld:22} : {value}']
return '\n'.join(s)
def _repr_html_(self) -> str:
s = []
s += ['<table border="1" class="dataframe"']
s += [
'<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>'
]
s += ['<tbody>']
for fld in self.__dataclass_fields__:
if fld not in self._special_fields:
value = getattr(self, fld)
if fld == 'restart_parameters':
value = self._restart_availability()
s += [
f'<tr><td style="text-align: left;">{fld:22}</td>'
f'<td>{value}</td><tr>'
]
for fld in self._special_fields:
d = getattr(self, fld)
# print('xxx', fld, d)
if fld.endswith('descriptor_weights'):
dim = list(d.values())[0].shape
elif fld == 'ann_parameters' and self.version == 4:
dim = (len(self.types), len(list(d.values())[0]))
else:
dim = len(d)
s += [
f'<tr><td style="text-align: left;">Dimension of {fld:22}</td><td>{dim}</td><tr>'
]
s += ['</tbody>']
s += ['</table>']
return ''.join(s)
@property
def training_parameters(self) -> dict:
"""Return the model parameters in the format accepted by :func:`write_nepfile
<calorine.nep.write_nepfile>`.
The result covers every ``nep.in`` keyword that describes the model itself, i.e. the
keywords that the ``nep`` executable checks against the ``nep.txt`` header before
training. It carries no training parameters (``lambda_*``, ``generation``, ``batch``,
and the like). Use :meth:`write_nepfile` to combine the two and write the file, rather
than merging the dictionaries by hand.
Returns
-------
dict
Keys ``version``, ``model_type``, ``type``, ``cutoff``, ``n_max``, ``basis_size``,
``l_max`` and ``neuron``, plus ``zbl`` and ``use_typewise_cutoff_zbl`` for a model
with ZBL repulsion, and ``charge_mode`` for a charge-aware model. ``zbl`` is the
single outer cutoff value that the ``nep.in`` ``zbl`` keyword expects (the inner
cutoff is always half of it), not the ``(inner, outer)`` pair stored in
:attr:`zbl`.
Raises
------
ValueError
If :attr:`model_type` is not one of the known model types.
"""
l_max = [self.l_max_3b, self.l_max_4b, self.l_max_5b,
self.has_q_112, self.has_q_123, self.has_q_233, self.has_q_134]
while len(l_max) > 1 and l_max[-1] == 0:
l_max = l_max[:-1]
if isinstance(self.radial_cutoff, list):
cutoff = []
for rc, ac in zip(self.radial_cutoff, self.angular_cutoff):
cutoff += [rc, ac]
else:
cutoff = [self.radial_cutoff, self.angular_cutoff]
if self.model_type not in _MODEL_TYPE_TO_INT:
raise ValueError(f'Unknown model_type: {self.model_type}')
# `type` must precede `cutoff`, which the `nep` executable sizes by the number of
# types, so the insertion order below is part of the contract.
params = {
'version': self.version,
'model_type': _MODEL_TYPE_TO_INT[self.model_type],
'type': [len(self.types)] + list(self.types),
'cutoff': cutoff,
'n_max': [self.n_max_radial, self.n_max_angular],
'basis_size': [self.n_basis_radial, self.n_basis_angular],
'l_max': l_max,
'neuron': self.n_neuron,
}
if self.zbl is not None:
zbl_inner, zbl_outer = self.zbl
if zbl_inner == 0 and zbl_outer == 0:
# GPUMD writes `zbl 0 0` for a flexible ZBL potential, which is requested by
# placing a `zbl.in` file next to `nep.in` rather than through a keyword.
warn('This model uses a flexible ZBL potential, which cannot be expressed in '
'nep.in. Place the zbl.in file of the model next to nep.in and set the '
'zbl cutoff by hand.')
elif abs(zbl_inner - 0.5 * zbl_outer) > 1e-6 * abs(zbl_outer):
warn(f'The ZBL cutoffs of this model, {zbl_inner} and {zbl_outer} Å, are not '
'expressible in nep.in, where the inner cutoff is always half of the '
f'outer one. Writing zbl {zbl_outer}, which implies an inner cutoff of '
f'{0.5 * zbl_outer} Å.')
params['zbl'] = zbl_outer
if self.zbl_typewise_cutoff_factor is not None:
params['use_typewise_cutoff_zbl'] = self.zbl_typewise_cutoff_factor
if self.charge_mode != 0:
params['charge_mode'] = self.charge_mode
return params
[docs]
def write_nepfile(self, filename: str, parameters: dict = None) -> None:
"""Writes a ``nep.in`` file for this model.
The keywords that describe the model are taken from the model itself, via
:attr:`training_parameters`, so the resulting file is consistent with the ``nep.txt``
file written by :meth:`write`. This is the intended way to prepare the input for
training a model that :meth:`augment`, :meth:`add_species`, :meth:`remove_species`,
:meth:`keep_species` or :meth:`prune` has changed the architecture of.
Training parameters can be supplied via :attr:`parameters`, typically read from an
existing ``nep.in`` file with :func:`read_nepfile <calorine.nep.read_nepfile>`. Any
keyword in :attr:`parameters` that describes the model is discarded in favor of the
value of the model, with a warning naming what was dropped, since such a value refers
to whichever model that file was written for and not to this one.
Note that unlike :func:`write_nepfile <calorine.nep.write_nepfile>`, which takes the
name of a directory, this method takes the name of a file, as :meth:`write` and
:meth:`write_restart` do.
Parameters
----------
filename
Name of the file to write, conventionally ``nep.in``.
parameters
Training parameters to include, such as ``generation``, ``batch`` and
``lambda_e``. Keywords that describe the model are ignored.
Raises
------
ValueError
If :attr:`model_type` is not one of the known model types.
Example
-------
Add a species to a model and write the input files needed to continue training it::
>>> from calorine.nep import read_model, read_nepfile
>>> model = read_model('nep.txt', restart_file='nep.restart')
>>> extended = model.add_species(['Cl'], seed=42)
>>> parameters = read_nepfile('nep.in')
>>> extended.write('new/nep.txt', restart_file='new/nep.restart')
>>> extended.write_nepfile('new/nep.in', parameters)
"""
model_parameters = self.training_parameters
merged = dict(model_parameters)
discarded = []
for key, value in (parameters or {}).items():
if key not in _MODEL_PARAMETERS:
merged[key] = value
continue
model_value = model_parameters.get(key)
if key in ('model_type', 'mode'):
# the `nep` executable accepts `mode` as a synonym of `model_type`
model_value = model_parameters['model_type']
elif key == 'use_typewise_cutoff_zbl' and _nepfile_tokens(value) == []:
# the bare keyword requests the default factor
value = _TYPEWISE_CUTOFF_ZBL_FACTOR_DEFAULT
elif key == 'charge_mode':
# `charge_mode <mode> [flip_charge]`, where flip_charge configures the
# training run rather than the model and is therefore kept
tokens = _nepfile_tokens(value)
if len(tokens) > 1 and model_value is not None:
merged['charge_mode'] = [model_value] + tokens[1:]
value = tokens[:1]
if not _same_nepfile_value(value, model_value):
discarded.append((key, value, model_value))
if discarded:
lines = []
for key, value, model_value in discarded:
supplied = _format_nepfile_value(value)
if model_value is None:
lines.append(f' {key}: {supplied} -> dropped, this model has no {key}')
else:
lines.append(f' {key}: {supplied} -> {_format_nepfile_value(model_value)}')
warn('The following nep.in parameters were overridden by the model:\n'
+ '\n'.join(lines))
_write_nepfile_to_path(merged, filename)
@property
def n_frozen_parameters(self) -> int | None:
"""Number of frozen restart parameters, i.e. parameters whose SNES sigma is
exactly zero and which are therefore excluded from the search during a restart.
Returns
-------
int or None
Number of parameters with a sigma of zero, or ``None`` if
``restart_parameters`` is not loaded.
Example
-------
Freeze everything that has already been trained and check the result::
>>> model = read_model('nep4_PbTe.txt', restart_file='nep4_PbTe.restart')
>>> frozen = model.set_restart_sigma(strategy='constant', value=0.0,
... target='set')
>>> frozen.n_frozen_parameters
2281
"""
if self.restart_parameters is None:
return None
return _restart_sigma_counts(self)['frozen']
@property
def n_unset_parameters(self) -> int | None:
"""Number of restart parameters whose SNES sigma is unset (``NaN``), i.e.
parameters created by :meth:`augment` or :meth:`add_species` that have not yet
been given a search width by :meth:`set_restart_sigma`.
:meth:`write_restart` refuses to write while this count is non-zero.
Returns
-------
int or None
Number of parameters with an unset sigma, or ``None`` if
``restart_parameters`` is not loaded.
Example
-------
Check how many parameters an architecture change created::
>>> model = read_model('nep4_PbTe.txt', restart_file='nep4_PbTe.restart')
>>> model.n_unset_parameters
0
>>> model.augment(n_neuron=40).n_unset_parameters
640
"""
if self.restart_parameters is None:
return None
return _restart_sigma_counts(self)['unset']
@property
def parameter_counts(self) -> dict[str, dict[str, int]]:
"""Number of parameters per category, i.e. per component of
:attr:`restart_parameters`.
The categories are the same ones that the ``component`` argument of
:meth:`set_restart_sigma` accepts, so this shows how many parameters a
component-restricted call reaches. Only the categories that the model actually
has are included: ``'charge_head'`` is absent for models without charges.
With ``restart_parameters`` loaded, the counts describe the entries of the restart
and are split into frozen and unset ones. Without it, they describe the current
parameters of the model and only the totals are available, since there are no sigma
values to split on. The totals exclude the :attr:`q_scaler` entries, which are not
fit parameters and have no restart counterpart, so they sum to
:attr:`n_parameters` minus the length of :attr:`q_scaler`.
Returns
-------
dict
Dictionary keyed by category. Each value holds the number of entries in total
(``'total'``) and, if ``restart_parameters`` is loaded, the number that are
frozen (``'frozen'``, sigma of zero) and the number that are unset
(``'unset'``, sigma of ``NaN``).
Example
-------
For a plain model only the totals are reported::
>>> read_model('nep4_PbTe.txt').parameter_counts
{'network_weights': {'total': 1921}, 'descriptor': {'total': 360}}
With a restart loaded, the sigma values split each category further. Freeze the
descriptor and check which category the frozen parameters are in::
>>> model = read_model('nep4_PbTe.txt', restart_file='nep4_PbTe.restart')
>>> frozen = model.set_restart_sigma(strategy='constant', value=0.0,
... component='descriptor', target='all')
>>> frozen.parameter_counts
{'network_weights': {'total': 1921, 'frozen': 0, 'unset': 0},
'descriptor': {'total': 360, 'frozen': 360, 'unset': 0}}
"""
if self.restart_parameters is None:
return _model_parameter_counts(self)
counts = {}
for component in _RESTART_COMPONENTS:
component_counts = _restart_sigma_counts(self, component)
if component_counts['total'] > 0:
counts[component] = component_counts
return counts
def _restart_availability(self) -> str:
"""Return the one-line summary of the restart-parameter state used by
:meth:`__str__` and :meth:`_repr_html_`."""
if self.restart_parameters is None:
return 'not available'
return (f'available ({self.n_frozen_parameters} frozen, '
f'{self.n_unset_parameters} unset)')
[docs]
def remove_species(self, species: list[str]) -> 'Model':
"""Remove one or more species from the model.
Returns a new :class:`Model` with the specified species removed.
The source model is not modified.
If ``restart_parameters`` are loaded, they are pruned to match (the
entries for the removed species/pairs are dropped); the surviving
entries are left exactly as they were. Use :meth:`set_restart_sigma`
explicitly afterwards if you want to re-open the SNES search width for
the surviving parameters before continuing training.
Parameters
----------
species
Species names to remove.
Returns
-------
Model
New model with the specified species removed.
Raises
------
ValueError
If any of the provided species is not found in the model.
"""
for s in species:
if s not in self.types:
raise ValueError(f'{s} is not a species supported by the NEP model')
new = copy.deepcopy(self)
types_to_keep = [t for t in self.types if t not in species]
new.types = tuple(types_to_keep)
# Prune ANN parameters (for NEP4 and NEP5)
if self.version in [4, 5]:
new.ann_parameters = {
key: value for key, value in new.ann_parameters.items()
if key in types_to_keep or key.startswith('b1')
}
# Prune descriptor weights; key is a (species1, species2) tuple
new.radial_descriptor_weights = {
key: value for key, value in new.radial_descriptor_weights.items()
if key[0] in types_to_keep and key[1] in types_to_keep
}
new.angular_descriptor_weights = {
key: value for key, value in new.angular_descriptor_weights.items()
if key[0] in types_to_keep and key[1] in types_to_keep
}
# Prune typewise cutoff lists so remaining species map to correct cutoffs
if isinstance(self.radial_cutoff, list):
indices = [i for i, t in enumerate(self.types) if t not in species]
new.radial_cutoff = [self.radial_cutoff[i] for i in indices]
new.angular_cutoff = [self.angular_cutoff[i] for i in indices]
# Prune restart parameters to match; survivors are left untouched
if new.restart_parameters is not None:
for param_type in ['mu', 'sigma']:
ann_key = f'ann_{param_type}'
if self.version in [4, 5]:
# Keep per-species keys for survivors, global bias keys, and
# sqrt_epsilon_infinity (charge models)
new.restart_parameters[ann_key] = {
key: value for key, value in new.restart_parameters[ann_key].items()
if (key in types_to_keep or key.startswith('b1')
or key == 'sqrt_epsilon_infinity')
}
# Prune descriptor restart parameters
for desc_type in ['radial', 'angular']:
key = f'{desc_type}_descriptor_{param_type}'
new.restart_parameters[key] = {
k: v for k, v in new.restart_parameters[key].items()
if k[0] in types_to_keep and k[1] in types_to_keep
}
# Recalculate parameter counts
_recalculate_parameter_counts(new)
return new
[docs]
def keep_species(self, species: list[str]) -> 'Model':
"""Retain only the specified species, removing all others.
Convenience complement to :meth:`remove_species`. Useful when the set
of species to drop is large (e.g. isolating two elements from a
foundation model with dozens of species).
Parameters
----------
species
Species names to keep. All other species are removed.
Returns
-------
Model
New model containing only the requested species.
Raises
------
ValueError
If any of the requested species is not in the model.
"""
unknown = [s for s in species if s not in self.types]
if unknown:
raise ValueError(
f'Species not in model: {unknown}'
)
to_remove = [s for s in self.types if s not in species]
return self.remove_species(to_remove)
[docs]
def reorder(self, order: list[str]) -> 'Model':
"""Reorder the species in the model.
Returns a new :class:`Model` with species permuted according to
``order``. This is useful for aligning the species order of two
models that must share the same order when used jointly by GPUMD,
e.g. a NEP potential and a TNEP dipole/polarizability model
referenced together via two ``potential`` lines in ``run.in`` and
``dump_dipole`` or ``dump_polarizability``.
The source model is not modified. Since ``ann_parameters``,
``radial_descriptor_weights``, ``angular_descriptor_weights``, and
``restart_parameters`` are keyed by species name (or species-pair)
rather than position, reordering only requires updating ``types``
and, if typewise cutoffs are in use, the positional
``radial_cutoff`` and ``angular_cutoff`` lists.
Parameters
----------
order
New species order. Must be a permutation of ``self.types``.
Returns
-------
Model
New model with species reordered.
Raises
------
ValueError
If ``order`` is not a permutation of the current species.
"""
if sorted(order) != sorted(self.types):
raise ValueError(
f'order must be a permutation of the current species {self.types}, '
f'got {list(order)}'
)
new = copy.deepcopy(self)
new.types = tuple(order)
if isinstance(self.radial_cutoff, list):
indices = [self.types.index(t) for t in order]
new.radial_cutoff = [self.radial_cutoff[i] for i in indices]
new.angular_cutoff = [self.angular_cutoff[i] for i in indices]
return new
[docs]
def add_species(self,
species: list[str],
radial_cutoff: float | list[float] = None,
angular_cutoff: float | list[float] = None,
seed: int | None = None) -> 'Model':
"""Add one or more species to the model.
Returns a new :class:`Model` with the requested species added. New ANN
sub-networks and descriptor weight pairs are initialised by drawing
``mu`` uniformly from [-1, 1] (matching the GPUMD fresh-model
initialisation); the corresponding restart sigma entries are left
unset (``NaN``) — call :meth:`set_restart_sigma` afterwards to
initialize them (e.g. ``model.add_species(['X']).set_restart_sigma()``
fills only the new entries by default). Charge-specific parameters
(``w1_charge``) are kept at ``mu = 0`` to preserve stability, also
matching GPUMD. Existing parameters (``mu`` and ``sigma``) are left
untouched. Call :meth:`initialize_parameters` to draw the new values from
a different distribution instead, or to give ``w1_charge`` a non-zero start.
Only supported for NEP4 models. For NEP3 the ANN is shared across all
species and adding a per-species sub-network is not meaningful.
Parameters
----------
species
New species names to add. Appended to ``types`` in the order given.
radial_cutoff
Radial cutoff(s) for the new species, in Å. Required when the model
uses typewise cutoffs (i.e. ``isinstance(model.radial_cutoff, list)``
is ``True``). Pass a single float or a list with one value per new
species.
angular_cutoff
Angular cutoff(s) for the new species, in Å. Same requirements as
``radial_cutoff``.
seed
Seed for the random number generator used to draw the initial ``mu``
values. Pass an integer for reproducible initialisation.
Returns
-------
Model
New model with updated structure, weights, and restart statistics.
Raises
------
ValueError
If the model version is not 4, if ``restart_parameters`` are not
loaded, if any species is already in the model, or if typewise
cutoffs are used and ``radial_cutoff``/``angular_cutoff`` are not
provided.
"""
if self.version != 4:
raise ValueError(
f'add_species() only supports NEP4 models; got version {self.version}.'
)
for s in species:
if s in self.types:
raise ValueError(f'{s!r} is already in the model.')
if self.restart_parameters is None:
raise ValueError(
'restart_parameters must be loaded before calling add_species(). '
'Pass restart_file= to read_model() or call model.read_restart() first.'
)
uses_typewise = isinstance(self.radial_cutoff, list)
if uses_typewise:
if radial_cutoff is None or angular_cutoff is None:
raise ValueError(
'Model uses typewise cutoffs; provide radial_cutoff and angular_cutoff '
'for the new species.'
)
rc_list = ([radial_cutoff] * len(species)
if isinstance(radial_cutoff, (int, float)) else list(radial_cutoff))
ac_list = ([angular_cutoff] * len(species)
if isinstance(angular_cutoff, (int, float)) else list(angular_cutoff))
if len(rc_list) != len(species) or len(ac_list) != len(species):
raise ValueError(
'Length of radial_cutoff/angular_cutoff must match the number of new species.'
)
new = copy.deepcopy(self)
n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular
n_neuron = self.n_neuron
is_charged = self.model_type == 'potential_with_charges'
all_types_after = list(self.types) + list(species)
rng = np.random.default_rng(seed)
def _rand(shape):
return rng.uniform(-1.0, 1.0, size=shape)
# Step 1: New ANN sub-networks
w1_shape = (n_neuron,) if is_charged else (1, n_neuron)
for s_new in species:
w0_vals = _rand((n_neuron, n_descriptor))
b0_vals = _rand((n_neuron, 1))
w1_vals = _rand(w1_shape)
s_params = {'w0': w0_vals.copy(), 'b0': b0_vals.copy(), 'w1': w1_vals.copy()}
if is_charged:
s_params['w1_charge'] = np.zeros(n_neuron)
new.ann_parameters[s_new] = s_params
mu_entry = {'w0': w0_vals, 'b0': b0_vals, 'w1': w1_vals}
sigma_entry = {
'w0': np.full((n_neuron, n_descriptor), np.nan),
'b0': np.full((n_neuron, 1), np.nan),
'w1': np.full(w1_shape, np.nan),
}
if is_charged:
mu_entry['w1_charge'] = np.zeros(n_neuron)
sigma_entry['w1_charge'] = np.full(n_neuron, np.nan)
new.restart_parameters['ann_mu'][s_new] = mu_entry
new.restart_parameters['ann_sigma'][s_new] = sigma_entry
# Step 2: New descriptor weight pairs
n_r = (self.n_max_radial + 1, self.n_basis_radial + 1)
n_a = (self.n_max_angular + 1, self.n_basis_angular + 1)
existing_pairs = set(self.radial_descriptor_weights)
new_pairs = {
(s1, s2)
for s1 in all_types_after for s2 in all_types_after
if (s1, s2) not in existing_pairs
}
for pair in new_pairs:
r_vals = _rand(n_r)
a_vals = _rand(n_a)
new.radial_descriptor_weights[pair] = r_vals.copy()
new.angular_descriptor_weights[pair] = a_vals.copy()
new.restart_parameters['radial_descriptor_mu'][pair] = r_vals
new.restart_parameters['angular_descriptor_mu'][pair] = a_vals
new.restart_parameters['radial_descriptor_sigma'][pair] = np.full(n_r, np.nan)
new.restart_parameters['angular_descriptor_sigma'][pair] = np.full(n_a, np.nan)
# Step 3: Update types and typewise cutoffs
new.types = tuple(all_types_after)
if uses_typewise:
new.radial_cutoff = list(self.radial_cutoff) + rc_list
new.angular_cutoff = list(self.angular_cutoff) + ac_list
# Step 4: Recalculate parameter counts
_recalculate_parameter_counts(new)
return new
[docs]
def write(self, filename: str, restart_file: str = None) -> None:
"""Write NEP model to file in `nep.txt` format.
Parameters
----------
filename
Output file name for the NEP model.
restart_file
If provided, also write restart parameters to this file in
`nep.restart` format. Defaults to None.
"""
with open(filename, 'w') as f:
# header
version_name = f'nep{self.version}'
if self.zbl is not None:
version_name += '_zbl'
if self.model_type == 'potential_with_charges':
version_name += f'_charge{self.charge_mode}'
elif self.model_type != 'potential':
version_name += f'_{self.model_type}'
f.write(f'{version_name} {len(self.types)} {" ".join(self.types)}\n')
if self.zbl is not None:
zbl_tokens = list(self.zbl)
if self.zbl_typewise_cutoff_factor is not None:
zbl_tokens.append(self.zbl_typewise_cutoff_factor)
f.write(f'zbl {" ".join(map(_format_header_float, zbl_tokens))}\n')
f.write('cutoff')
if isinstance(self.radial_cutoff, float) and isinstance(self.angular_cutoff, float):
f.write(f' {_format_header_float(self.radial_cutoff)}'
f' {_format_header_float(self.angular_cutoff)}')
else:
# Typewise cutoffs: one set of cutoffs per type
for i in range(len(self.types)):
f.write(f' {_format_header_float(self.radial_cutoff[i])}'
f' {_format_header_float(self.angular_cutoff[i])}')
f.write(f' {self.max_neighbors_radial} {self.max_neighbors_angular}')
f.write('\n')
f.write(f'n_max {self.n_max_radial} {self.n_max_angular}\n')
f.write(f'basis_size {self.n_basis_radial} {self.n_basis_angular}\n')
l_max_line = f'l_max {self.l_max_3b} {self.l_max_4b} {self.l_max_5b}'
if self.has_q_112 or self.has_q_123 or self.has_q_233 or self.has_q_134:
l_max_line += f' {self.has_q_112}'
if self.has_q_123 or self.has_q_233 or self.has_q_134:
l_max_line += f' {self.has_q_123}'
if self.has_q_233 or self.has_q_134:
l_max_line += f' {self.has_q_233}'
if self.has_q_134:
l_max_line += f' {self.has_q_134}'
f.write(l_max_line + '\n')
f.write(f'ANN {self.n_neuron} 0\n')
# neural network weights
keys = self.types if self.version in (4, 5) else ['all_species']
suffixes = ['', '_polar'] if self.model_type == 'polarizability' else ['']
for suffix in suffixes:
for s in keys:
# Order: w0, b0, w1 (, b1 if NEP5)
# w0 indexed as: n*N_descriptor + nu
w0 = self.ann_parameters[s][f'w0{suffix}']
b0 = self.ann_parameters[s][f'b0{suffix}']
w1 = self.ann_parameters[s][f'w1{suffix}']
for n in range(self.n_neuron):
for nu in range(
self.n_descriptor_radial + self.n_descriptor_angular
):
f.write(f'{w0[n, nu]:15.7e}\n')
for b in b0[:, 0]:
f.write(f'{b:15.7e}\n')
for v in (w1[0, :] if w1.ndim == 2 else w1):
f.write(f'{v:15.7e}\n')
if f'w1_charge{suffix}' in self.ann_parameters[s]:
for v in self.ann_parameters[s][f'w1_charge{suffix}']:
f.write(f'{v:15.7e}\n')
if self.version == 5:
b1 = self.ann_parameters[s][f'b1{suffix}']
f.write(f'{b1:15.7e}\n')
if self.sqrt_epsilon_infinity is not None:
f.write(f'{self.sqrt_epsilon_infinity:15.7e}\n')
b1 = self.ann_parameters[f'b1{suffix}']
f.write(f'{b1:15.7e}\n')
# descriptor weights
mat = []
for s1 in self.types:
for s2 in self.types:
mat = np.hstack(
[mat, self.radial_descriptor_weights[(s1, s2)].flatten()]
)
mat = np.hstack(
[mat, self.angular_descriptor_weights[(s1, s2)].flatten()]
)
n_types = len(self.types)
n = int(len(mat) / (n_types * n_types))
mat = mat.reshape((n_types * n_types, n)).T
for v in mat.flatten():
f.write(f'{v:15.7e}\n')
# scaler
for v in self.q_scaler:
f.write(f'{v:15.7e}\n')
if restart_file is not None:
self.write_restart(restart_file)
[docs]
def read_restart(self, filename: str):
"""Parses a file in `nep.restart` format and saves the
content in the form of mean and standard deviation for each
parameter in the corresponding NEP model.
Parameters
----------
filename
Input file name.
"""
mu, sigma = _get_restart_contents(filename)
restart_parameters = np.array([mu, sigma]).T
is_polarizability_model = self.model_type == 'polarizability'
is_charged_model = self.model_type == 'potential_with_charges'
n1 = self.n_ann_parameters
n1 *= 2 if is_polarizability_model else 1
n2 = n1 + self.n_descriptor_parameters
ann_parameters = restart_parameters[:n1]
descriptor_parameters = np.array(restart_parameters[n1:n2])
if self.version == 3:
n_networks = 1
elif self.version in (4, 5):
# one hidden layer per atomic species
n_networks = len(self.types)
else:
raise ValueError(f'Cannot load nep.restart for NEP model version {self.version}')
ann_groups = [s for s in self.ann_parameters.keys() if not s.startswith('b1')]
n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular
restart = {}
for i, content_type in enumerate(['mu', 'sigma']):
ann = _sort_ann_parameters(ann_parameters[:, i],
ann_groups,
self.n_neuron,
n_networks,
self.version,
n_descriptor,
is_polarizability_model,
is_charged_model)
radial, angular = _sort_descriptor_parameters(descriptor_parameters[:, i],
self.types,
self.n_max_radial,
self.n_basis_radial,
self.n_max_angular,
self.n_basis_angular)
restart[f'ann_{content_type}'] = ann
restart[f'radial_descriptor_{content_type}'] = radial
restart[f'angular_descriptor_{content_type}'] = angular
self.restart_parameters = restart
[docs]
def write_restart(self, filename: str):
"""Write the restart parameters to file in `nep.restart` format.
Parameters
----------
filename
Output file name.
Raises
------
ValueError
If ``restart_parameters`` is not loaded, or if any restart sigma
value is unset (``NaN``), e.g. because :meth:`add_species` or
:meth:`augment` were called without a follow-up
:meth:`set_restart_sigma` to initialize the sigma of the newly
created parameters.
"""
if self.restart_parameters is None:
raise ValueError(
'restart_parameters is not loaded; nothing to write. Pass restart_file= '
'to read_model(), call Model.read_restart(), or call '
'Model.set_restart_sigma() to bootstrap one from the current model '
'parameters before write_restart().'
)
for _, sigma_container, sigma_key in _restart_leaves(self, self.restart_parameters):
if np.any(np.isnan(np.asarray(sigma_container[sigma_key], dtype=float))):
raise ValueError(
f'restart_parameters contains an unset (NaN) sigma value for '
f'{sigma_key!r}. Call Model.set_restart_sigma() to initialize it '
'before write_restart().'
)
keys = self.types if self.version in (4, 5) else ['all_species']
suffixes = ['', '_polar'] if self.model_type == 'polarizability' else ['']
columns = []
for i, parameter in enumerate(['mu', 'sigma']):
# neural network weights
ann_parameters = self.restart_parameters[f'ann_{parameter}']
column = []
for suffix in suffixes:
for s in keys:
# Order: w0, b0, w1 (, b1 if NEP5)
# w0 indexed as: n*N_descriptor + nu
w0 = ann_parameters[s][f'w0{suffix}']
b0 = ann_parameters[s][f'b0{suffix}']
w1 = ann_parameters[s][f'w1{suffix}']
for n in range(self.n_neuron):
for nu in range(
self.n_descriptor_radial + self.n_descriptor_angular
):
column.append(f'{w0[n, nu]:15.7e}')
for b in b0[:, 0]:
column.append(f'{b:15.7e}')
for v in (w1[0, :] if w1.ndim == 2 else w1):
column.append(f'{v:15.7e}')
if f'w1_charge{suffix}' in ann_parameters[s]:
for v in ann_parameters[s][f'w1_charge{suffix}']:
column.append(f'{v:15.7e}')
if f'b1{suffix}' in ann_parameters[s]:
column.append(f'{ann_parameters[s][f"b1{suffix}"]:15.7e}')
if f'sqrt_epsilon_infinity{suffix}' in ann_parameters:
column.append(f'{ann_parameters[f"sqrt_epsilon_infinity{suffix}"]:15.7e}')
b1 = ann_parameters[f'b1{suffix}']
column.append(f'{b1:15.7e}')
columns.append(column)
# descriptor weights
radial_descriptor_parameters = self.restart_parameters[f'radial_descriptor_{parameter}']
angular_descriptor_parameters = self.restart_parameters[
f'angular_descriptor_{parameter}']
mat = []
for s1 in self.types:
for s2 in self.types:
mat = np.hstack(
[mat, radial_descriptor_parameters[(s1, s2)].flatten()]
)
mat = np.hstack(
[mat, angular_descriptor_parameters[(s1, s2)].flatten()]
)
n_types = len(self.types)
n = int(len(mat) / (n_types * n_types))
mat = mat.reshape((n_types * n_types, n)).T
for v in mat.flatten():
column.append(f'{v:15.7e}')
# Join the mean and standard deviation columns
assert len(columns[0]) == len(columns[1]), 'Length of means must match standard deviation'
joined = [f'{s1} {s2}\n' for s1, s2 in zip(*columns)]
with open(filename, 'w') as f:
f.writelines(joined)
[docs]
def initialize_parameters(self,
strategy: str = 'uniform',
*,
value: float = None,
low: float = None,
high: float = None,
mean: float = None,
std: float = None,
species: str | list[str] = None,
component: str | list[str] = None,
target: str = 'unset',
seed: int | None = None) -> 'Model':
"""Initialize parameter values from a distribution.
Returns a new :class:`Model` in which the parameters selected by
*target*/*species*/*component* are drawn according to *strategy*. The source model is
not modified. Every sigma value is left exactly as it was, so
:meth:`set_restart_sigma` remains the only method that assigns sigma.
Both the parameter values of the model (:attr:`ann_parameters` and the descriptor
weights, i.e. what goes into ``nep.txt``) and the corresponding ``mu`` entries of
:attr:`restart_parameters` (the SNES mean, i.e. what goes into ``nep.restart``) are
written. Only the selected positions are touched, since for a trained model the two
are not the same thing: ``nep.txt`` holds the best parameters found so far while
the restart holds the mean of the search distribution.
:attr:`sqrt_epsilon_infinity` is the one parameter no strategy reaches, whatever
*target*/*species*/*component* select. Every other parameter may take any value, but
``epsilon_infinity`` is the square of this one, the high-frequency dielectric
constant, so a drawn value that is negative or close to zero is not a starting point.
Set it through the ``sqrt_epsilon_infinity`` argument of :meth:`augment` instead,
which defaults to 1.
The main use is to give the parameters created by :meth:`augment` a starting point
other than zero::
model.augment(n_neuron=40) \\
.initialize_parameters(strategy='uniform', low=-1, high=1, seed=42) \\
.set_restart_sigma(strategy='constant', value=0.0, target='set') \\
.set_restart_sigma(target='unset')
The order matters. :meth:`initialize_parameters` has to come before the
:meth:`set_restart_sigma` call that fills the unset sigma values, because an unset
sigma is what marks a parameter as new. Initializing the values first also makes the
default sigma strategy meaningful for the new parameters, since
``strategy='scale_mu'`` computes ``sigma = max(floor, factor * |mu|)``, which is just
the floor as long as ``mu`` is zero.
Parameters
----------
strategy
How to draw the new values at the selected positions:
- ``'uniform'`` (default): ``value ~ U(low, high)``. Drawing from
``U(-1, 1)`` reproduces what :meth:`add_species` does for a new sub-network.
- ``'constant'``: every selected parameter is set to ``value``.
- ``'normal'``: ``value ~ N(mean, std)``.
The values are used as drawn, including negative ones.
value
Value for ``strategy='constant'``.
low, high
Bounds for ``strategy='uniform'``.
mean, std
Parameters of the normal distribution for ``strategy='normal'``.
species
Restrict the update to one or more species (and descriptor pairs involving them).
``None`` (default) applies to all species; the global shared bias is only
included when ``species`` is ``None``.
component
Restrict the update to one or more of ``'network_weights'``, ``'descriptor'``,
``'charge_head'``. ``None`` (default) applies to all of them. ``'charge_head'``
reaches ``w1_charge`` only, ``sqrt_epsilon_infinity`` being excluded as described
above.
target
Which parameters to initialize, selected by the state of the corresponding
*sigma*: ``'unset'`` (default) only the parameters whose sigma is unset (``NaN``),
i.e. those newly created by :meth:`augment` or :meth:`add_species`; ``'set'`` only
those that are already part of the search; ``'all'`` every selected parameter.
A value of zero is not used as the marker, since a trained parameter may be zero.
seed
Seed for the random number generator. Pass an integer for reproducibility.
Returns
-------
Model
New model with the selected parameter values initialized.
Raises
------
ValueError
If ``restart_parameters`` is not loaded, if ``strategy``/``target``/``component``
is not recognized, or if a strategy-specific required argument is missing.
Example
-------
Give the neurons that :meth:`augment` added a random starting point::
>>> model = read_model('nep.txt', restart_file='nep.restart')
>>> grown = model.augment(n_neuron=40)
>>> grown.ann_parameters['Pb']['w0'][30:].any() # new rows are zero
False
>>> initialized = grown.initialize_parameters(low=-1, high=1, seed=42)
>>> initialized.ann_parameters['Pb']['w0'][30:].any()
True
"""
valid_strategies = {'constant', 'uniform', 'normal'}
if strategy not in valid_strategies:
raise ValueError(
f'strategy must be one of {sorted(valid_strategies)}; got {strategy!r}'
)
if target not in ('unset', 'set', 'all'):
raise ValueError(f"target must be 'unset', 'set', or 'all'; got {target!r}")
if strategy == 'constant' and value is None:
raise ValueError("strategy='constant' requires value.")
if strategy == 'uniform' and (low is None or high is None):
raise ValueError("strategy='uniform' requires low and high.")
if strategy == 'normal' and (mean is None or std is None):
raise ValueError("strategy='normal' requires mean and std.")
if self.restart_parameters is None:
raise ValueError(
'restart_parameters must be loaded before calling initialize_parameters(), '
'since the sigma values are what mark a parameter as new. Pass restart_file= '
'to read_model(), call model.read_restart(), or call '
'Model.set_restart_sigma() to bootstrap one first.'
)
new = copy.deepcopy(self)
rng = np.random.default_rng(seed)
kw = dict(value=value, low=low, high=high, mean=mean, std=std)
for container, key, owner, name in _model_parameter_leaves(new, component, species):
if name == 'sqrt_epsilon_infinity':
# Deliberately left out of every strategy. The other parameters may take any
# value, but epsilon_infinity is the square of this one, the high-frequency
# dielectric constant, so a drawn value that is negative or close to zero is
# not a starting point. augment(charge_head=True) sets it, via its
# sqrt_epsilon_infinity argument, and that value is what training starts from.
continue
restart_leaf = _restart_leaf_for(new.restart_parameters, owner, name, key)
if restart_leaf is None:
continue
mu_container, sigma_container, restart_key = restart_leaf
sigma = sigma_container[restart_key]
if np.isscalar(sigma) or isinstance(sigma, (float, int)):
is_unset = np.isnan(float(sigma))
apply_here = (
target == 'all' or (target == 'unset' and is_unset)
or (target == 'set' and not is_unset)
)
if not apply_here:
continue
new_value = float(_draw_values(strategy, rng, **kw))
_leaf_set(container, key, new_value)
mu_container[restart_key] = new_value
continue
values = _leaf_get(container, key)
mu = mu_container[restart_key]
if np.shape(values) != np.shape(sigma) or np.shape(mu) != np.shape(sigma):
raise ValueError(
f'Shape mismatch for {name!r}: the model holds {np.shape(values)} while '
f'the restart holds {np.shape(mu)}. initialize_parameters() cannot map '
'the two onto each other.'
)
if target == 'unset':
mask = np.isnan(sigma)
elif target == 'set':
mask = ~np.isnan(sigma)
else:
mask = np.ones_like(sigma, dtype=bool)
if not np.any(mask):
continue
drawn = _draw_values(strategy, rng, size=int(np.count_nonzero(mask)), **kw)
values[mask] = drawn
mu[mask] = drawn
return new
[docs]
def set_restart_sigma(self,
strategy: str = 'scale_mu',
*,
value: float = None,
factor: float = None,
floor: float = 1e-6,
low: float = None,
high: float = None,
mean: float = None,
std: float = None,
species: str | list[str] = None,
component: str | list[str] = None,
target: str = 'unset',
seed: int | None = None) -> 'Model':
"""Assign SNES restart sigma values.
Returns a new :class:`Model` with sigma values updated according to
*strategy*, at the positions selected by *target*/*species*/*component*.
``mu`` and every other field are left unchanged. This is the only
method that ever assigns sigma values; the structural methods
(:meth:`remove_species`, :meth:`keep_species`, :meth:`add_species`,
:meth:`augment`, :meth:`prune`) leave the sigma of the parameters they keep
untouched and mark newly created parameters' sigma as unset (``NaN``) rather
than computing a value inline. Parameter values are assigned by
:meth:`initialize_parameters` instead, which in turn never touches a sigma.
If ``restart_parameters`` is not loaded, it is created first: ``mu`` is
copied from the model's current (trained) parameters, and every sigma
is initialized as unset (``NaN``). This makes it possible to bootstrap
a ``nep.restart`` file "from scratch" for a plain ``nep.txt`` model.
Parameters
----------
strategy
How to compute new sigma values at the selected positions:
- ``'constant'``: ``sigma = value``.
- ``'scale_mu'`` (default): ``sigma = max(floor, factor * |mu|)``,
re-opening the SNES search width in proportion to each
parameter's magnitude. ``factor`` defaults to ``0.1`` for this
strategy.
- ``'scale_sigma'``: ``sigma = sigma * factor``. Requires the
selected sigma values to already be set (not ``NaN``).
- ``'uniform'``: draw ``sigma ~ U(low, high)``.
- ``'normal'``: draw ``sigma = |N(mean, std)|``.
value
Sigma value for ``strategy='constant'``.
factor
Scale factor for ``strategy='scale_mu'`` (default ``0.1`` if not
given) or ``strategy='scale_sigma'`` (required).
floor
Minimum sigma for ``strategy='scale_mu'``.
low, high
Bounds for ``strategy='uniform'``. Since a sigma is a standard
deviation, ``low`` must be positive and smaller than ``high``, which
keeps every drawn value positive. Use ``strategy='constant'`` for a
single value rather than ``low == high``.
mean, std
Parameters of the normal distribution for ``strategy='normal'``.
species
Restrict the update to one or more species (and descriptor pairs
involving them). ``None`` (default) applies to all species; global
parameters (the shared bias, ``sqrt_epsilon_infinity``) are only
included when ``species`` is ``None``.
component
Restrict the update to one or more of ``'network_weights'``,
``'descriptor'``, ``'charge_head'``. ``None`` (default) applies to
all three.
target
Which existing sigma values to update: ``'unset'`` (default) only
fills in ``NaN`` entries (e.g. those left by :meth:`add_species`/
:meth:`augment`); ``'set'`` only updates already-set entries;
``'all'`` updates every selected entry regardless of its current
value.
seed
Seed for the random number generator used by the ``'uniform'`` and
``'normal'`` strategies. Pass an integer for reproducibility.
Returns
-------
Model
New model with updated restart sigma values.
Raises
------
ValueError
If ``strategy``/``target``/``component`` is not recognized, if a
strategy-specific required argument is missing, if
``strategy='uniform'`` is given bounds that would admit a
non-positive sigma, or if ``strategy='scale_sigma'`` is applied to a
still-unset (``NaN``) sigma value.
"""
valid_strategies = {'constant', 'scale_mu', 'scale_sigma', 'uniform', 'normal'}
if strategy not in valid_strategies:
raise ValueError(
f'strategy must be one of {sorted(valid_strategies)}; got {strategy!r}'
)
if target not in ('unset', 'set', 'all'):
raise ValueError(f"target must be 'unset', 'set', or 'all'; got {target!r}")
if strategy == 'constant' and value is None:
raise ValueError("strategy='constant' requires value.")
if strategy == 'scale_mu' and factor is None:
factor = 0.1
if strategy == 'scale_sigma' and factor is None:
raise ValueError("strategy='scale_sigma' requires factor.")
if strategy == 'uniform' and (low is None or high is None):
raise ValueError("strategy='uniform' requires low and high.")
# A sigma is a standard deviation, so the whole interval has to be positive. Both
# checks are needed: numpy samples between the two bounds whatever their order, so a
# positive low on its own does not bound the draw from below.
if strategy == 'uniform' and low <= 0:
raise ValueError(
"strategy='uniform' requires a positive low, since sigma is a standard "
f'deviation; got low={low!r}.'
)
if strategy == 'uniform' and low >= high:
raise ValueError(
"strategy='uniform' requires low < high, so that every drawn sigma stays "
f"positive; got low={low!r} and high={high!r}. Use strategy='constant' for a "
'single value.'
)
if strategy == 'normal' and (mean is None or std is None):
raise ValueError("strategy='normal' requires mean and std.")
new = copy.deepcopy(self)
if new.restart_parameters is None:
new.restart_parameters = _new_restart_parameters_from_model(new)
rng = np.random.default_rng(seed)
kw = dict(value=value, factor=factor, floor=floor, low=low, high=high, mean=mean, std=std)
for mu, sigma_container, sigma_key in _restart_leaves(
new, new.restart_parameters, component, species
):
_apply_sigma_strategy(mu, sigma_container, sigma_key, strategy, target, rng, **kw)
return new
[docs]
def augment(self,
n_neuron: int = None,
l_max_4b: int = None,
l_max_5b: int = None,
has_q_112: bool = None,
has_q_123: bool = None,
has_q_233: bool = None,
has_q_134: bool = None,
charge_head: bool = False,
charge_mode: int = 1,
sqrt_epsilon_infinity: float = 1.0) -> 'Model':
"""Augment the model by adding neurons, descriptor terms, or a charge output head.
Returns a new :class:`Model` with the requested structural changes applied.
The source model is not modified. Existing parameter values (``mu`` and
``sigma``) are preserved exactly; new parameters are initialized to
``mu = 0``, with the corresponding restart sigma left unset (``NaN``).
``sqrt_epsilon_infinity`` is the one exception: ``epsilon_infinity`` is its square,
a dielectric constant, so zero is not a value it can take. It starts at the value of
the ``sqrt_epsilon_infinity`` argument (1 by default) instead.
Call :meth:`set_restart_sigma` afterwards to initialize the new sigma
entries (e.g. ``model.augment(n_neuron=40).set_restart_sigma()`` fills
only the new entries by default). To give the new parameters a starting
value other than zero, call :meth:`initialize_parameters` before that, since
an unset sigma is what marks a parameter as new.
Parameters
----------
n_neuron
Target neuron count; must be >= current. ``None`` leaves unchanged.
l_max_4b
Target 4-body l_max value; must be >= current. ``None`` leaves unchanged.
l_max_5b
Target 5-body l_max value; must be >= current. ``None`` leaves unchanged.
has_q_112
``True`` enables the q_112 5-body descriptor; ``None`` or ``False`` leaves
the current state unchanged (disabling an already-enabled term raises).
has_q_123
Same as ``has_q_112`` but for the q_123 term.
has_q_233
Same as ``has_q_112`` but for the q_233 term.
has_q_134
Same as ``has_q_112`` but for the q_134 term.
charge_head
If ``True``, promote a ``potential`` model to ``potential_with_charges`` by
adding a charge output head (w1_charge per species and sqrt_epsilon_infinity).
charge_mode
Charge algorithm variant to record for the new charge head; must be 1 or 2.
1 corresponds to a qNEP model, including both real- and reciprocal-space
contributions. 2 corresponds to a qNEP model, including the reciprocal-space
contribution only. Only meaningful when ``charge_head=True``.
sqrt_epsilon_infinity
Starting value for the new ``sqrt_epsilon_infinity`` parameter, used only when
``charge_head=True``. Must be positive: ``epsilon_infinity`` is its square, the
high-frequency dielectric constant. The default of 1 corresponds to no dielectric
screening, which is the neutral starting point for training. This is the one
parameter ``augment`` does not leave at zero, and :meth:`initialize_parameters`
skips it for the same reason, so the value given here is the one that reaches
training.
Returns
-------
Model
New model with updated structure, weights, and restart statistics.
Raises
------
ValueError
If ``restart_parameters`` is not loaded, if ``n_neuron`` or an ``l_max_*``
target is smaller than the current value, if a ``has_q_*`` flag attempts to
disable an already-enabled term, or if ``charge_head=True`` on a model that
is not of type ``potential`` or comes with a non-positive
``sqrt_epsilon_infinity``.
"""
# Structural checks (independent of restart)
if self.version not in (3, 4):
raise ValueError(
f'augment() only supports NEP versions 3 and 4; got version {self.version}.'
)
if n_neuron is not None and n_neuron < self.n_neuron:
raise ValueError(
f'n_neuron ({n_neuron}) must be >= current n_neuron ({self.n_neuron}); '
'use prune() to reduce.'
)
if l_max_4b is not None and l_max_4b < self.l_max_4b:
raise ValueError(
f'l_max_4b ({l_max_4b}) must be >= current l_max_4b ({self.l_max_4b}); '
'use prune() to disable.'
)
if l_max_5b is not None and l_max_5b < self.l_max_5b:
raise ValueError(
f'l_max_5b ({l_max_5b}) must be >= current l_max_5b ({self.l_max_5b}); '
'use prune() to disable.'
)
for flag_val, name in [
(has_q_112, 'has_q_112'), (has_q_123, 'has_q_123'), (has_q_233, 'has_q_233'),
(has_q_134, 'has_q_134')
]:
if flag_val is False and getattr(self, name):
raise ValueError(
f'Cannot disable {name} via augment(); '
'use prune() to disable descriptor terms.'
)
if charge_head and self.model_type != 'potential':
raise ValueError(
f'charge_head=True requires model_type="potential"; '
f'got "{self.model_type}".'
)
if charge_head and self.version != 4:
# GPUMD only accepts nep4_charge1/nep4_charge2 for the qNEP charge modes, so a
# nep3_charge1 model would be written but could not be read back.
raise ValueError(
f'charge_head=True requires a NEP4 model; got version {self.version}.'
)
if charge_head and charge_mode not in (1, 2):
raise ValueError(f'charge_mode must be 1 or 2; got {charge_mode}.')
if charge_head and not float(sqrt_epsilon_infinity) > 0:
raise ValueError(
'sqrt_epsilon_infinity must be positive; got '
f'{sqrt_epsilon_infinity}. epsilon_infinity is its square, the '
'high-frequency dielectric constant.'
)
if self.restart_parameters is None:
raise ValueError(
'restart_parameters must be loaded before calling augment(). '
'Pass restart_file= to read_model() or call model.read_restart() first.'
)
new = copy.deepcopy(self)
# Resolve new structural parameters
new_l_max_4b = l_max_4b if l_max_4b is not None else self.l_max_4b
new_l_max_5b = l_max_5b if l_max_5b is not None else self.l_max_5b
new_has_q_112 = int(has_q_112) if has_q_112 is not None else self.has_q_112
new_has_q_123 = int(has_q_123) if has_q_123 is not None else self.has_q_123
new_has_q_233 = int(has_q_233) if has_q_233 is not None else self.has_q_233
new_has_q_134 = int(has_q_134) if has_q_134 is not None else self.has_q_134
new_n_neuron = n_neuron if n_neuron is not None else self.n_neuron
new_l_max_enh = (self.l_max_3b
+ (new_l_max_4b > 0) + (new_l_max_5b > 0)
+ (new_has_q_112 > 0) + (new_has_q_123 > 0) + (new_has_q_233 > 0)
+ (new_has_q_134 > 0))
new_n_desc_angular = (self.n_max_angular + 1) * new_l_max_enh
old_n_desc = self.n_descriptor_radial + self.n_descriptor_angular
new_n_desc = self.n_descriptor_radial + new_n_desc_angular
delta_desc = new_n_desc - old_n_desc
delta_neuron = new_n_neuron - self.n_neuron
keys = self.types if self.version in (4, 5) else ['all_species']
# Step 1: Expand descriptor dimensions (new columns in w0, new q_scaler entries)
if delta_desc > 0:
for s in keys:
old_w0 = new.ann_parameters[s]['w0'] # (n_neuron_old, old_n_desc)
new.ann_parameters[s]['w0'] = np.hstack(
[old_w0, np.zeros((self.n_neuron, delta_desc))]
)
old_mu_w0 = new.restart_parameters['ann_mu'][s]['w0']
new.restart_parameters['ann_mu'][s]['w0'] = np.hstack(
[old_mu_w0, np.zeros((self.n_neuron, delta_desc))]
)
old_sigma_w0 = new.restart_parameters['ann_sigma'][s]['w0']
new.restart_parameters['ann_sigma'][s]['w0'] = np.hstack(
[old_sigma_w0, np.full((self.n_neuron, delta_desc), np.nan)]
)
new.q_scaler = list(new.q_scaler) + [1.0] * delta_desc
# Step 2: Expand neuron count (new rows in w0/b0, new columns in w1)
if delta_neuron > 0:
for s in keys:
# w0: append new rows
cur_w0 = new.ann_parameters[s]['w0'] # (n_old, new_n_desc)
new.ann_parameters[s]['w0'] = np.vstack(
[cur_w0, np.zeros((delta_neuron, new_n_desc))]
)
# b0: append new rows
cur_b0 = new.ann_parameters[s]['b0']
new.ann_parameters[s]['b0'] = np.vstack(
[cur_b0, np.zeros((delta_neuron, 1))]
)
# w1: append new columns; handle both 2D (standard) and 1D (charge)
cur_w1 = new.ann_parameters[s]['w1']
zeros_w1 = (np.zeros(delta_neuron) if cur_w1.ndim == 1
else np.zeros((1, delta_neuron)))
new.ann_parameters[s]['w1'] = np.hstack([cur_w1, zeros_w1])
if 'w1_charge' in new.ann_parameters[s]:
cur_wc = new.ann_parameters[s]['w1_charge']
new.ann_parameters[s]['w1_charge'] = np.hstack([cur_wc, np.zeros(delta_neuron)])
# restart w0
cur_mu_w0 = new.restart_parameters['ann_mu'][s]['w0']
new.restart_parameters['ann_mu'][s]['w0'] = np.vstack(
[cur_mu_w0, np.zeros((delta_neuron, new_n_desc))]
)
cur_sigma_w0 = new.restart_parameters['ann_sigma'][s]['w0']
new.restart_parameters['ann_sigma'][s]['w0'] = np.vstack(
[cur_sigma_w0, np.full((delta_neuron, new_n_desc), np.nan)]
)
# restart b0
cur_mu_b0 = new.restart_parameters['ann_mu'][s]['b0']
new.restart_parameters['ann_mu'][s]['b0'] = np.vstack(
[cur_mu_b0, np.zeros((delta_neuron, 1))]
)
cur_sigma_b0 = new.restart_parameters['ann_sigma'][s]['b0']
new.restart_parameters['ann_sigma'][s]['b0'] = np.vstack(
[cur_sigma_b0, np.full((delta_neuron, 1), np.nan)]
)
# restart w1
cur_mu_w1 = new.restart_parameters['ann_mu'][s]['w1']
zeros_w1 = (np.zeros(delta_neuron) if cur_mu_w1.ndim == 1
else np.zeros((1, delta_neuron)))
new.restart_parameters['ann_mu'][s]['w1'] = np.hstack([cur_mu_w1, zeros_w1])
cur_sigma_w1 = new.restart_parameters['ann_sigma'][s]['w1']
nan_w1 = (np.full(delta_neuron, np.nan) if cur_sigma_w1.ndim == 1
else np.full((1, delta_neuron), np.nan))
new.restart_parameters['ann_sigma'][s]['w1'] = np.hstack([cur_sigma_w1, nan_w1])
if 'w1_charge' in new.restart_parameters['ann_mu'][s]:
cur = new.restart_parameters['ann_mu'][s]['w1_charge']
new.restart_parameters['ann_mu'][s]['w1_charge'] = np.hstack(
[cur, np.zeros(delta_neuron)]
)
cur = new.restart_parameters['ann_sigma'][s]['w1_charge']
new.restart_parameters['ann_sigma'][s]['w1_charge'] = np.hstack(
[cur, np.full(delta_neuron, np.nan)]
)
# Step 3: Add charge output head
if charge_head:
new.model_type = 'potential_with_charges'
new.charge_mode = charge_mode
new.sqrt_epsilon_infinity = float(sqrt_epsilon_infinity)
for s in keys:
cur_w1 = new.ann_parameters[s]['w1'] # (1, new_n_neuron)
new.ann_parameters[s]['w1'] = cur_w1[0, :] # flatten to 1D
new.ann_parameters[s]['w1_charge'] = np.zeros(new_n_neuron)
cur_mu_w1 = new.restart_parameters['ann_mu'][s]['w1']
new.restart_parameters['ann_mu'][s]['w1'] = cur_mu_w1[0, :]
new.restart_parameters['ann_mu'][s]['w1_charge'] = np.zeros(new_n_neuron)
cur_sigma_w1 = new.restart_parameters['ann_sigma'][s]['w1']
new.restart_parameters['ann_sigma'][s]['w1'] = cur_sigma_w1[0, :]
new.restart_parameters['ann_sigma'][s]['w1_charge'] = np.full(
new_n_neuron, np.nan
)
new.restart_parameters['ann_mu']['sqrt_epsilon_infinity'] = float(
sqrt_epsilon_infinity
)
new.restart_parameters['ann_sigma']['sqrt_epsilon_infinity'] = float('nan')
# Step 4: Update header metadata
new.l_max_4b = new_l_max_4b
new.l_max_5b = new_l_max_5b
new.has_q_112 = new_has_q_112
new.has_q_123 = new_has_q_123
new.has_q_233 = new_has_q_233
new.has_q_134 = new_has_q_134
new.n_descriptor_angular = new_n_desc_angular
new.n_neuron = new_n_neuron
# Step 5: Recalculate parameter counts
_recalculate_parameter_counts(new)
return new
[docs]
def prune(self,
n_neuron: int = None,
l_max_4b: int = None,
l_max_5b: int = None,
has_q_112: bool = None,
has_q_123: bool = None,
has_q_233: bool = None,
has_q_134: bool = None,
charge_head: bool = False) -> 'Model':
"""Prune the model by removing neurons, disabling descriptor terms, or removing
the charge output head.
Returns a new :class:`Model` with the requested structural changes applied.
The source model is not modified. When reducing ``n_neuron``, neurons are
selected by importance score averaged over species:
``importance[n] = mean_s(||w0_s[n,:]||_2 * |w1_s[n]|)``.
Surviving parameters (``mu`` and ``sigma``) are left exactly as they
were. Use :meth:`set_restart_sigma` explicitly afterwards if you want
to re-open the SNES search width for the survivors before continuing
training.
Parameters
----------
n_neuron
Target neuron count; must be <= current. ``None`` leaves unchanged.
l_max_4b
Target 4-body l_max; must be <= current. Setting to ``0`` removes the
4-body angular descriptor block. Reducing to a lower non-zero value is
a header-only change (descriptor dimensions unchanged). ``None`` leaves
unchanged.
l_max_5b
Same as ``l_max_4b`` but for five-body terms.
has_q_112
``False`` disables and removes the q_112 descriptor block. ``None``
leaves unchanged. ``True`` is not valid; use :meth:`augment` instead.
has_q_123
Same as ``has_q_112`` but for the q_123 term.
has_q_233
Same as ``has_q_112`` but for the q_233 term.
has_q_134
Same as ``has_q_112`` but for the q_134 term.
charge_head
If ``True``, remove the charge output head from a
``potential_with_charges`` model, converting it back to ``potential``.
Removes ``w1_charge`` per species and ``sqrt_epsilon_infinity`` from
the restart.
Returns
-------
Model
New model with reduced structure, weights, and restart statistics.
Raises
------
ValueError
If ``restart_parameters`` is not loaded, if any target value would
expand the model (use :meth:`augment` instead), if a ``has_q_*``
flag is set to ``True``, or if ``charge_head=True`` on a model
without charges.
"""
# --- Resolve target values ---
new_n_neuron = n_neuron if n_neuron is not None else self.n_neuron
new_l_max_4b = l_max_4b if l_max_4b is not None else self.l_max_4b
new_l_max_5b = l_max_5b if l_max_5b is not None else self.l_max_5b
new_has_q_112 = 0 if has_q_112 is False else self.has_q_112
new_has_q_123 = 0 if has_q_123 is False else self.has_q_123
new_has_q_233 = 0 if has_q_233 is False else self.has_q_233
new_has_q_134 = 0 if has_q_134 is False else self.has_q_134
# --- Validate ---
if self.version not in (3, 4):
raise ValueError(
f'prune() only supports NEP versions 3 and 4; got version {self.version}.'
)
if new_n_neuron > self.n_neuron:
raise ValueError(
f'n_neuron ({new_n_neuron}) must be <= current n_neuron ({self.n_neuron}); '
'use augment() to increase.'
)
if new_l_max_4b > self.l_max_4b:
raise ValueError(
f'l_max_4b ({new_l_max_4b}) must be <= current l_max_4b ({self.l_max_4b}); '
'use augment() to increase.'
)
if new_l_max_5b > self.l_max_5b:
raise ValueError(
f'l_max_5b ({new_l_max_5b}) must be <= current l_max_5b ({self.l_max_5b}); '
'use augment() to increase.'
)
for flag_val, name in [
(has_q_112, 'has_q_112'), (has_q_123, 'has_q_123'),
(has_q_233, 'has_q_233'), (has_q_134, 'has_q_134')
]:
if flag_val is True:
raise ValueError(
f'Cannot enable {name} via prune(); '
'use augment() to enable descriptor terms.'
)
if charge_head and self.model_type != 'potential_with_charges':
raise ValueError(
f'charge_head=True requires model_type="potential_with_charges"; '
f'got "{self.model_type}".'
)
if self.restart_parameters is None:
raise ValueError(
'restart_parameters must be loaded before calling prune(). '
'Pass restart_file= to read_model() or call model.read_restart() first.'
)
new = copy.deepcopy(self)
keys = self.types if self.version in (4, 5) else ['all_species']
# Step 1: Neuron pruning — keep the most important neurons
if new_n_neuron < self.n_neuron:
importances = []
for s in keys:
w0 = self.ann_parameters[s]['w0'] # (n_neuron, n_desc)
w1_flat = self.ann_parameters[s]['w1'].ravel()
if 'w1_charge' in self.ann_parameters[s]:
output_norm = np.abs(w1_flat) + np.abs(self.ann_parameters[s]['w1_charge'])
else:
output_norm = np.abs(w1_flat)
importances.append(np.linalg.norm(w0, axis=1) * output_norm)
keep_idx = np.sort(np.argsort(np.mean(importances, axis=0))[-new_n_neuron:])
for s in keys:
new.ann_parameters[s]['w0'] = new.ann_parameters[s]['w0'][keep_idx, :]
new.ann_parameters[s]['b0'] = new.ann_parameters[s]['b0'][keep_idx, :]
w1 = new.ann_parameters[s]['w1']
new.ann_parameters[s]['w1'] = w1[:, keep_idx] if w1.ndim == 2 else w1[keep_idx]
if 'w1_charge' in new.ann_parameters[s]:
new.ann_parameters[s]['w1_charge'] = (
new.ann_parameters[s]['w1_charge'][keep_idx]
)
for pk in ['ann_mu', 'ann_sigma']:
rp = new.restart_parameters[pk][s]
rp['w0'] = rp['w0'][keep_idx, :]
rp['b0'] = rp['b0'][keep_idx, :]
w1 = rp['w1']
rp['w1'] = w1[:, keep_idx] if w1.ndim == 2 else w1[keep_idx]
if 'w1_charge' in rp:
rp['w1_charge'] = rp['w1_charge'][keep_idx]
# Step 2: Descriptor column pruning (disabling higher-body terms)
n_per = self.n_max_angular + 1
hb_terms = [
(self.l_max_4b, new_l_max_4b),
(self.l_max_5b, new_l_max_5b),
(self.has_q_112, new_has_q_112),
(self.has_q_123, new_has_q_123),
(self.has_q_233, new_has_q_233),
(self.has_q_134, new_has_q_134),
]
keep_cols = list(range(self.n_descriptor_radial + n_per * self.l_max_3b))
col_offset = len(keep_cols)
for old_val, new_val in hb_terms:
if old_val > 0:
if new_val > 0:
keep_cols.extend(range(col_offset, col_offset + n_per))
col_offset += n_per
old_n_desc = self.n_descriptor_radial + self.n_descriptor_angular
if len(keep_cols) < old_n_desc:
keep_cols = np.array(keep_cols, dtype=int)
for s in keys:
new.ann_parameters[s]['w0'] = new.ann_parameters[s]['w0'][:, keep_cols]
for pk in ['ann_mu', 'ann_sigma']:
rp = new.restart_parameters[pk][s]
rp['w0'] = rp['w0'][:, keep_cols]
new.q_scaler = [new.q_scaler[i] for i in keep_cols]
# Step 3: Charge head removal
if charge_head:
new.model_type = 'potential'
new.charge_mode = 0
new.sqrt_epsilon_infinity = None
for s in keys:
w1 = new.ann_parameters[s]['w1'] # 1D (n_neuron,)
new.ann_parameters[s]['w1'] = w1.reshape(1, -1)
del new.ann_parameters[s]['w1_charge']
for pk in ['ann_mu', 'ann_sigma']:
rp = new.restart_parameters[pk][s]
rp['w1'] = rp['w1'].reshape(1, -1)
del rp['w1_charge']
del new.restart_parameters['ann_mu']['sqrt_epsilon_infinity']
del new.restart_parameters['ann_sigma']['sqrt_epsilon_infinity']
# Step 4: Update header fields
new.n_neuron = new_n_neuron
new.l_max_4b = new_l_max_4b
new.l_max_5b = new_l_max_5b
new.has_q_112 = new_has_q_112
new.has_q_123 = new_has_q_123
new.has_q_233 = new_has_q_233
new.has_q_134 = new_has_q_134
new_l_max_enh = (self.l_max_3b
+ (new_l_max_4b > 0) + (new_l_max_5b > 0)
+ (new_has_q_112 > 0) + (new_has_q_123 > 0) + (new_has_q_233 > 0)
+ (new_has_q_134 > 0))
new.n_descriptor_angular = (self.n_max_angular + 1) * new_l_max_enh
# Step 5: Recalculate parameter counts
_recalculate_parameter_counts(new)
return new
[docs]
def read_model(filename: str, restart_file: str = None) -> Model:
"""Parses a file in ``nep.txt`` format and returns the
content in the form of a :class:`Model <calorine.nep.model.Model>`
object.
Parameters
----------
filename
Input file name.
restart_file
If provided, also read restart parameters from this file in
`nep.restart` format and attach them to the returned model.
Defaults to None.
"""
data, parameters = _get_nep_contents(filename)
# sanity checks
for fld in ['version', 'types', 'model_type', 'cutoff', 'basis_size', 'n_max', 'l_max', 'ANN']:
if fld not in data:
raise ValueError(f'Invalid model file; {fld} line is missing')
if data['version'] not in [3, 4, 5]:
raise ValueError('Invalid model file; only NEP versions 3, 4 and 5 are currently supported')
# split up zbl tuple (optional typewise cutoff factor as a third entry)
if 'zbl' in data:
if len(data['zbl']) == 3:
data['zbl_typewise_cutoff_factor'] = data['zbl'][2]
data['zbl'] = data['zbl'][:2]
elif len(data['zbl']) != 2:
raise ValueError(
f'Invalid model file; zbl line must have 2 or 3 entries, got {len(data["zbl"])}'
)
# split up cutoff tuple
N_types = len(data['types'])
# Either global cutoffs + max neighbirs, or typewise cutoffs + max_neighbors
if len(data['cutoff']) not in [4, 2*N_types+2]:
raise ValueError(
'Invalid model file; cutoff line must have 4 entries (global cutoffs) or '
f'{2*N_types+2} entries (typewise cutoffs for {N_types} types), '
f'got {len(data["cutoff"])}'
)
if not all(np.isfinite(data['cutoff'])):
raise ValueError('Invalid model file; cutoff values must be finite')
data['max_neighbors_radial'] = int(data['cutoff'][-2])
data['max_neighbors_angular'] = int(data['cutoff'][-1])
if len(data['cutoff']) == 2*N_types+2:
# Typewise cutoffs: radial are even, angular are odd
data['radial_cutoff'] = [data['cutoff'][i*2] for i in range(N_types)]
data['angular_cutoff'] = [data['cutoff'][i*2+1] for i in range(N_types)]
else:
data['radial_cutoff'] = data['cutoff'][0]
data['angular_cutoff'] = data['cutoff'][1]
del data['cutoff']
# split up basis_size tuple
if len(data['basis_size']) != 2:
raise ValueError(
f'Invalid model file; basis_size line must have 2 entries, '
f'got {len(data["basis_size"])}'
)
data['n_basis_radial'] = data['basis_size'][0]
data['n_basis_angular'] = data['basis_size'][1]
del data['basis_size']
# split up n_max tuple
if len(data['n_max']) != 2:
raise ValueError(
f'Invalid model file; n_max line must have 2 entries, got {len(data["n_max"])}'
)
data['n_max_radial'] = data['n_max'][0]
data['n_max_angular'] = data['n_max'][1]
del data['n_max']
# split up nl_max tuple
len_l = len(data['l_max'])
if len_l not in [1, 2, 3, 4, 5, 6, 7]:
raise ValueError(
f'Invalid model file; l_max line must have between 1 and 7 entries, got {len_l}'
)
data['l_max_3b'] = data['l_max'][0]
data['l_max_4b'] = data['l_max'][1] if len_l > 1 else 0
data['l_max_5b'] = data['l_max'][2] if len_l > 2 else 0
data['has_q_112'] = data['l_max'][3] if len_l > 3 else 0
data['has_q_123'] = data['l_max'][4] if len_l > 4 else 0
data['has_q_233'] = data['l_max'][5] if len_l > 5 else 0
data['has_q_134'] = data['l_max'][6] if len_l > 6 else 0
del data['l_max']
# compute dimensions of descriptor components
data['n_descriptor_radial'] = data['n_max_radial'] + 1
l_max_enh = (data['l_max_3b']
+ (data['l_max_4b'] > 0)
+ (data['l_max_5b'] > 0)
+ (data['has_q_112'] > 0)
+ (data['has_q_123'] > 0)
+ (data['has_q_233'] > 0)
+ (data['has_q_134'] > 0))
data['n_descriptor_angular'] = (data['n_max_angular'] + 1) * l_max_enh
n_descriptor = data['n_descriptor_radial'] + data['n_descriptor_angular']
is_charged_model = data['model_type'] == 'potential_with_charges'
# compute number of parameters
data['n_neuron'] = data['ANN'][0]
del data['ANN']
n_types = len(data['types'])
# NEP4 and NEP5 have one hidden layer per atomic species, NEP3 a single shared one
n = n_types if data['version'] in (4, 5) else 1
n_output_biases = _number_of_output_biases(data['version'], n_types, is_charged_model)
n_ann_input_weights = (n_descriptor + 1) * data['n_neuron'] # weights + bias
n_ann_output_weights = 2*data['n_neuron'] if is_charged_model else data['n_neuron'] # weights
n_ann_parameters = (
n_ann_input_weights + n_ann_output_weights
) * n + n_output_biases
n_descriptor_weights = n_types**2 * (
(data['n_max_radial'] + 1) * (data['n_basis_radial'] + 1)
+ (data['n_max_angular'] + 1) * (data['n_basis_angular'] + 1)
)
data['n_parameters'] = n_ann_parameters + n_descriptor_weights + n_descriptor
is_polarizability_model = data['model_type'] == 'polarizability'
if data['n_parameters'] + n_ann_parameters == len(parameters):
data['n_parameters'] += n_ann_parameters
if not is_polarizability_model:
raise ValueError(
'Model is not labelled as a polarizability model, but the number of '
'parameters matches a polarizability model.\n'
'If this is a polarizability model trained with GPUMD <=v3.8, please '
'modify the header in the nep.txt file to enable parsing '
f'`nep{data["version"]}_polarizability`.\n'
)
if len(parameters) < data['n_parameters']:
raise ValueError(
'Invalid model file; expected '
f'{data["n_parameters"]} parameter values, found {len(parameters)} '
'-- file may be truncated'
)
elif len(parameters) > data['n_parameters']:
raise ValueError(
'Invalid model file; expected '
f'{data["n_parameters"]} parameter values, found {len(parameters)} '
'-- file may contain extra or corrupted data'
)
data['n_ann_parameters'] = n_ann_parameters
# split up parameters into the ANN weights, descriptor weights, and scaling parameters
n1 = n_ann_parameters
n1 *= 2 if is_polarizability_model else 1
n2 = n1 + n_descriptor_weights
data['ann_parameters'] = parameters[:n1]
descriptor_weights = np.array(parameters[n1:n2])
data['q_scaler'] = parameters[n2:]
# add ann parameters to data dict
ann_groups = data['types'] if data['version'] in (4, 5) else ['all_species']
sorted_ann_parameters = _sort_ann_parameters(data['ann_parameters'],
ann_groups,
data['n_neuron'],
n,
data['version'],
n_descriptor,
is_polarizability_model,
is_charged_model)
data['ann_parameters'] = sorted_ann_parameters
if 'sqrt_epsilon_infinity' in sorted_ann_parameters.keys():
data['sqrt_epsilon_infinity'] = sorted_ann_parameters['sqrt_epsilon_infinity']
sorted_ann_parameters.pop('sqrt_epsilon_infinity')
data['ann_parameters'] = sorted_ann_parameters
# add descriptors to data dict
data['n_descriptor_parameters'] = len(descriptor_weights)
radial, angular = _sort_descriptor_parameters(descriptor_weights,
data['types'],
data['n_max_radial'],
data['n_basis_radial'],
data['n_max_angular'],
data['n_basis_angular'])
data['radial_descriptor_weights'] = radial
data['angular_descriptor_weights'] = angular
model = Model(**data)
if restart_file is not None:
model.read_restart(restart_file)
return model
