NEP interface#

Training IO#

calorine provides a number of functions for preparing input files for training NEP models, including in particular the setup_training function. There are also several functions for analyzing the training process, including, e.g., the read_loss, read_structures, and get_parity_data functions.

calorine.nep.get_parity_data(structures, property, selection=None, flatten=True)[source]#

Returns the predicted and target energies, forces, virials or stresses from a list of structures in a format suitable for generating parity plots.

The structures should have been read using read_structures, such that the info object is populated with keys of the form <property>_<type> where <property> is, e.g., energy or force and <type> is one of predicted or target.

The resulting parity data is returned as a tuple of dicts, where each entry corresponds to a list.

Parameters:

structures (list[Atoms]) – List of structures as read with read_structures.
property (str) – One of energy, force, virial, stress, bec, dipole, polarizability, or atomic_v.
selection (list[str]) – A list containing which components to return, and/or the norm. For force, atomic_v (dipole), virial, stress, polarizability, and dipole, possible values are x, y, z, xx, yy, zz, yz, xz, xy, norm, pressure (the latter only for stress). For bec, all nine of xx, xy, xz, yx, yy, yz, zx, zy, zz are available (BEC is not symmetric, so unlike the other properties it has no reduced/Voigt form).
flatten (bool) – if True return flattened lists; this is useful for flattening the components of force or virials into a simple list

Return type:

DataFrame

calorine.nep.read_loss(filename)[source]#

Parses a file in loss.out format from GPUMD and returns the content as a data frame. More information concerning file format, content and units can be found here.

Parameters:: filename (str) – input file name
Return type:: DataFrame

calorine.nep.read_nepfile(filename)[source]#

Returns the content of a configuration file (nep.in) as a dictionary.

Parameters:: filename (str) – input file name
Return type:: dict[str, Any]

calorine.nep.read_structures(dirname)[source]#

Parses the output files with training and test data from a nep run and returns their content as two lists of structures, representing training and test data, respectively. Target and predicted data are included in the info dict of the Atoms objects.

Parameters:: dirname (str) – Directory from which to read output files.
Return type:: tuple[list[Atoms], list[Atoms]]

calorine.nep.setup_fine_tuning_nep89(parameters, nep, restart, **kwargs_to_setup_training)[source]#

Sets up a fine-tuning of the NEP89 foundation model.

Note that only the types, the number of generations, the batch, the population, and the regularization parameters are allowed to be changed.

The types must be a subset of the 89 types atomic species supported by the NEP89 foundation model.

This function wraps setup_training().

Parameters:

parameters (NamedTuple) – Dictionary containing the parameters to be set in the nep.in file; see here for an overview of these parameters. Note that only lambda_1, lambda_2, lambda_e, lambda_f, lambda_v, generation, population, type, and batch are allowed parameters when fine-tuning.
nep (Path) – Path to the nep.txt file for NEP89.
restart (Path) – Path to the nep.restart file for NEP89.
kwargs_to_setup_training – See the dosctring for setup_training for the rest of the parameters.

Return type:

None

calorine.nep.setup_training(parameters, structures=None, train_structures=None, test_structures=None, enforced_structures=[], rootdir='.', mode='kfold', n_splits=None, train_fraction=None, seed=42, overwrite=False)[source]#

Sets up the input files for training a NEP via the nep executable of the GPUMD package.

Parameters:

parameters (NamedTuple) – Dictionary containing the parameters to be set in the nep.in file. See here for an overview of these parameters.
structures (list[Atoms]) – List of structures to be included. Required for modes 'kfold' and 'bagging', and must not be set when mode 'fixed' is used.
train_structures (list[Atoms]) – Pre-defined list of training structures. Only used (and required) when mode 'fixed' is used.
test_structures (list[Atoms]) – Pre-defined list of test structures. Only used (and required) when mode 'fixed' is used.
enforced_structures (list[int]) – Structures that _must_ be included in the training set, provided in the form of a list of indices that refer to the content of the structures parameter. Must not be set when mode 'fixed' is used.
rootdir (str) – Root directory in which to create the input files.
mode (str) – How the test-train split is performed. Options: 'kfold', 'bagging', and 'fixed'. 'fixed' bypasses the split logic entirely and writes train_structures/test_structures as given directly to a single nepmodel directory under rootdir, rather than the nepmodel_full/nepmodel_split* directories written by 'kfold'/ 'bagging'.
n_splits (int) – Number of splits of the input structures in training and test sets that ought to be performed. By default no split will be done and all input structures will be used for training. Must not be set when mode 'fixed' is used.
train_fraction (float) – Fraction of structures to use for training when mode 'bagging' is used. Must not be set when mode 'fixed' is used.
seed (int) – Random number generator seed to be used. This ensures reproducability.
overwrite (bool) – If True overwrite the content of rootdir if it exists.

Return type:

None

calorine.nep.write_nepfile(parameters, dirname)[source]#

Writes parameters file for NEP construction.

Parameters:

parameters (NamedTuple) – input parameters; see here
dirname (str) – directory in which to place input file and links

Return type:

None

calorine.nep.write_structures(outfile, structures)[source]#

Writes structures for training/testing in format readable by nep executable.

Parameters:

outfile (str) – output filename
structures (list[Atoms]) – list of structures with energy, forces, and (possibly) stresses

Return type:

None

Correcting per-species energy offsets#

A per-species energy-reference shift can be determined with determine_energy_offsets and applied to the target energies of a set of structures with apply_energy_offsets, e.g., to correct an energy-scale mismatch before fine-tuning a reference model such as NEP89 on data from a different reference method, or to obtain isolated-atom-style reference energies before training a model from scratch. This can also be useful when the reference data exhibit a large spread in energies, which can lead to ill-behaved training. This is, for example, the case for all-electron DFT codes such as FHI-aims and Wien2k, for which the spread in total energies can reach thousands of eV per atom.

calorine.nep.determine_energy_offsets(structures, reference=None, command=None)[source]#

Determines a per-species energy offset (reference energy shift) mu_i for a set of structures, by fitting E_target - E_reference = sum_i n_i * mu_i via ordinary least squares, where n_i is the number of atoms of species i in a structure. The fit never includes a constant term, since a structure-size-independent constant would distort per-atom energy comparisons between structures of different sizes.

If reference is None, the offsets are fitted directly against the target energies of structures themselves, which is useful, for example, to obtain isolated-atom-style reference energies before training a model from scratch. If reference is given, either as a path to a NEP model in nep.txt format or as a Model object, the offsets are instead fitted against the residual between the target energies and the energies predicted by reference, e.g., to correct for an energy-scale mismatch when fine-tuning a reference model such as NEP89 on data from a different reference method.

Parameters:

structures (List[Atoms]) – Structures with target energies attached, e.g., via an ASE calculator.
reference (Union[str, Model, None]) – Path to a NEP model in nep.txt format, or a Model object, to fit the offsets against. If None, the offsets are fitted directly against the target energies of structures.
command (Optional[str]) – Command used to invoke the nep executable when reference is given. Default: nep, or the value of the CALORINE_NEP_COMMAND environment variable if set.

Return type:

Dict[str, float]

calorine.nep.apply_energy_offsets(structures, offsets)[source]#

Applies a per-species energy offset, as determined by determine_energy_offsets, to the target energies of a set of structures. Forces and stresses, if present, are passed through unchanged.

Parameters:

structures (List[Atoms]) – Structures with target energies attached, e.g., via an ASE calculator.
offsets (Dict[str, float]) – Per-species energy offsets in eV, keyed by chemical symbol, as returned by determine_energy_offsets.

Return type:

List[Atoms]

Returns:

A new list of Atoms objects, in the same order as
structures, with corrected target energies. The input structures
are not modified.

Evaluating models#

TNEP models allow one to represent tensorial properties such as dipole moment, susceptibility, or polarizability. To test and analyze these models calorine provides several specialized functions, which can also be used to implement extended Hamiltonians.

calorine.nep.get_dipole(structure, model_filename=None, debug=False)[source]#

Calculates the dipole for a given structure. A NEP model defined by a nep.txt file needs to be provided.

Parameters:

structure (Atoms) – Input structure
model_filename (Optional[str]) – Path to NEP model in nep.txt format. Defaults to None.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output. Defaults to False.

Return type:

ndarray

calorine.nep.get_dipole_gradient(structure, model_filename=None, backend='c++', method='central difference', displacement=0.01, charge=1.0, nep_command=None, debug=False)[source]#

Calculates the dipole gradient for a given structure using finite differences. A NEP model defined by a nep.txt file needs to be provided.

Parameters:

structure (Atoms) – Input structure
model_filename (Optional[str]) – Path to NEP model in nep.txt format. Defaults to None.
backend (str) – Backend to use for computing dipole gradient with finite differences. One of 'c++' (CPU), 'python' (CPU) and 'nep' (GPU). Defaults to 'c++'.
method (str) – Method for computing gradient with finite differences. One of ‘forward difference’ and ‘central difference’. Defaults to ‘central difference’
displacement (float) – Displacement in Å to use for finite differences. Defaults to 0.01.
charge (float) – System charge in units of the elemental charge. Used for correcting the dipoles before computing the gradient. Defaults to 1.0.
nep_command (Optional[str]) – Command for running the NEP executable. Default: nep, or the value of the CALORINE_NEP_COMMAND environment variable if set.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to False.

Return type:

ndarray

Returns:

dipole gradient with shape (N, 3, 3)

calorine.nep.get_polarizability(structure, model_filename=None, debug=False)[source]#

Calculates the polarizability tensor for a given structure. A NEP model defined by a nep.txt file needs to be provided. The model must be trained to predict the polarizability.

Parameters:

structure (Atoms) – Input structure
model_filename (Optional[str]) – Path to NEP model in nep.txt format. Defaults to None.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output. Defaults to False.

Return type:

ndarray

calorine.nep.get_polarizability_gradient(structure, model_filename=None, displacement=0.01, component='full', debug=False)[source]#

Calculates the dipole gradient for a given structure using finite differences. A NEP model defined by a nep.txt file needs to be provided. This function computes the derivatives using the second-order central difference method with a C++ backend.

Parameters:

structure (Atoms) – Input structure.
model_filename (Optional[str]) – Path to NEP model in nep.txt format. Defaults to None.
displacement (float) – Displacement in Å to use for finite differences. Defaults to 0.01.
component (Union[str, List[str]]) – Component or components of the polarizability tensor that the gradient should be computed for. The following components are available: x`, ``y, z, full. Option full computes the derivative whilst moving the atoms in each Cartesian direction, which yields a tensor of shape (N, 3, 6). Multiple components may be specified. Defaults to full.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults to False.

Return type:

ndarray

Returns:

polarizability gradient with shape (N, C, 6) where C is the number of components chosen.

calorine.nep.get_potential_forces_and_virials(structure, model_filename=None, debug=False)[source]#

Calculates the per-atom potential, forces and virials for a given structure. A NEP model defined by a nep.txt file needs to be provided.

Parameters:

structure (Atoms) – Input structure
model_filename (Optional[str]) – Path to NEP model. Defaults to None.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output. Defaults to False.

Return type:

Tuple[ndarray, ndarray, ndarray]

Returns:

potential with shape (natoms,)
forces with shape (natoms, 3)
virials with shape (natoms, 9)

Inspecting NEP models#

Once a model has been trained it can be analyzed in more detail. To this end, there are functions for accessing the descriptors, the latent space, or to load the entire model. The latter function (read_model) returns a Model object, which contains the entire information about this model. It is thereby possible not only to query but to manipulate the model and write the result back to disk.

calorine.nep.get_descriptors(structure, model_filename, debug=False)[source]#

Calculates the NEP descriptors for a given structure. A NEP model defined by a nep.txt can additionally be provided to get the NEP3 model specific descriptors.

Parameters:

structure (Atoms) – Input structure
model_filename (str) – Path to NEP model in nep.txt format.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output. Defaults to False.

Return type:

ndarray

calorine.nep.get_latent_space(structure, model_filename=None, debug=False)[source]#

Calculates the latent space representation of a structure, i.e, the activiations in the hidden layer. A NEP model defined by a nep.txt file needs to be provided.

Parameters:

structure (Atoms) – Input structure
model_filename (Optional[str]) – Path to NEP model. Defaults to None.
debug (bool) – Flag to toggle debug mode. Prints GPUMD output. Defaults to False.

Return type:

ndarray

calorine.nep.read_model(filename, restart_file=None)[source]#

Parses a file in nep.txt format and returns the content in the form of a Model object.

Parameters:

filename (str) – Input file name.
restart_file (str) – If provided, also read restart parameters from this file in nep.restart format and attach them to the returned model. Defaults to None.

Return type:

Model

Modifying NEP models#

A loaded Model can be edited in place and written back to disk. Species can be added, removed, or filtered via add_species, remove_species, and keep_species. The species order can be permuted via reorder to align two models used jointly by GPUMD. The network architecture can be extended or reduced via augment and prune. These two methods also cover two architectural special cases. A potential model can be promoted to a charge-aware potential_with_charges (qNEP) model by passing charge_head=True together with the desired charge_mode (1 or 2, selecting between two Ewald-summation variants). Higher-body descriptor terms (has_q_112, has_q_123, has_q_233, has_q_134) can be added or removed alongside the angular cutoff order (l_max_3b/4b/5b).

add_species and augment require the optimizer state of the model (restart_parameters, the SNES mean/sigma values) to already be loaded, since they must extend it with entries for the parameters they create. This state is read alongside nep.txt by passing restart_file=... to read_model (or via a separate call to Model.read_restart). It is written back out via the same restart_file argument to Model.write, or explicitly via Model.write_restart.

None of add_species, remove_species, keep_species, augment, or prune ever assign a sigma value themselves. Parameters that already existed keep their exact prior mu/sigma. Parameters newly created by add_species/augment get an initialized mu but an unset (NaN) sigma. Model.write_restart refuses to write while any sigma is unset. Model.set_restart_sigma is the single method that assigns sigma values. It supports several strategies (scale_mu, constant, scale_sigma, uniform, normal), restricted by target ('unset'/'set'/'all') and optionally by species/component. Typical recipes:

Fill in the gaps after add_species/augment: model.add_species(['X']).set_restart_sigma()
Bootstrap a restart from a plain nep.txt (restart_parameters is None): model.set_restart_sigma()
Re-open sigma for the existing parameters after remove_species/prune: model.remove_species([...]).set_restart_sigma(target='all')
Freeze the descriptor, train only the network weights: model.set_restart_sigma(strategy='constant', value=0.0, component='descriptor', target='all')

Parameter values are assigned by Model.initialize_parameters, which is the counterpart of set_restart_sigma and never touches a sigma. It draws the selected parameters from a distribution (uniform, constant, or normal) and accepts the same target/species/component selectors as Model.set_restart_sigma, writing both the values of the model itself and the corresponding mu entries of the restart. The main use is to give the parameters created by 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')

initialize_parameters has to come before the set_restart_sigma call that fills the unset sigma values, since an unset sigma is what marks a parameter as new. Initializing the values first also makes the default sigma strategy meaningful for those parameters, because scale_mu computes sigma = max(floor, factor * |mu|), which is just the floor as long as mu is zero.

sqrt_epsilon_infinity is the one parameter that is explicitly excluded from any strategy. Every other parameter may take any value, but sqrt_epsilon_infinity is the square root of the effective high-frequency dielectric constant, so a random value drawn from U(-1, 1) is usually not a good starting point. It is also the one parameter that augment does not start from zero. Choose its starting value through the sqrt_epsilon_infinity argument of augment instead, which defaults to 1 and rejects non-positive values:

model.augment(charge_head=True, sqrt_epsilon_infinity=1.0)

How many parameters each step affects can be read off the model. n_frozen_parameters counts the parameters with a sigma of zero and n_unset_parameters those with an unset sigma. Both are None when restart_parameters is not loaded, and both are included in the printed representation of the model. parameter_counts breaks the same numbers down by category, using the same categories as the component argument, which shows how many parameters a component-restricted call reaches. It also works without restart_parameters, in which case it reports the parameters of the model itself and only the totals per category, since there are then no sigma values to split on.

See the Modifying NEP models tutorial for worked examples.

Writing the training input for a modified model#

Continuing the training of a modified model requires a nep.in file that describes the same architecture as the nep.txt file written by Model.write. The nep executable checks the two against each other and refuses to run if they disagree, so a nep.in carried over unchanged from before a call to augment or add_species is an error rather than an inconvenience.

Model.write_nepfile takes every keyword that describes the model from the model itself, so the file it writes agrees with the model by construction:

model.write('nep.txt')
model.write_nepfile('nep.in')

The keywords covered this way are version, model_type, type, cutoff, n_max, basis_size, l_max and neuron, together with zbl and use_typewise_cutoff_zbl for a model with ZBL repulsion, and charge_mode for a charge-aware model. They are also available on their own via Model.training_parameters.

A file written like that carries no training parameters at all, since the model does not record them. To keep the ones from an earlier run, read that nep.in with read_nepfile and pass the result along:

parameters = read_nepfile('previous/nep.in')
model.write_nepfile('nep.in', parameters)

Keywords such as generation, batch, population, lambda_e and type_weight are then passed through unchanged. Keywords that describe the model are taken from the model instead, and a warning names each value that was discarded, since such a value belongs to whichever model that file was written for. A keyword that does not apply to the model at all is removed rather than carried over, which is the case that makes a hand-merged file inconsistent most easily: a zbl line left behind for a model without ZBL repulsion, or a charge_mode line left behind for a model without charges.

Two keywords are only partly determined by the model. The flip_charge value that may follow the mode on the charge_mode line configures the training run, so it is kept while the mode itself comes from the model. A flexible ZBL potential cannot be expressed in nep.in at all, as it is requested by placing a zbl.in file next to it, and a model that uses one is reported with a warning.

NEP model class#

class calorine.nep.model.Model(version, model_type, types, radial_cutoff, angular_cutoff, n_basis_radial, n_basis_angular, n_max_radial, n_max_angular, l_max_3b, l_max_4b, l_max_5b, has_q_112, has_q_123, has_q_233, has_q_134, n_descriptor_radial, n_descriptor_angular, n_neuron, n_parameters, n_descriptor_parameters, n_ann_parameters, ann_parameters, q_scaler, radial_descriptor_weights, angular_descriptor_weights, sqrt_epsilon_infinity=None, charge_mode=0, restart_parameters=None, zbl=None, zbl_typewise_cutoff_factor=None, max_neighbors_radial=None, max_neighbors_angular=None)[source]#

Objects of this class represent a NEP model in a form suitable for inspection and manipulation. Typically a Model object is instantiated by calling the read_model function.

version#

NEP version.

Type:: int

model_type#

One of potential, dipole or polarizability.

Type:: str

types#

Chemical species that this model represents.

Type:: tuple[str, …]

radial_cutoff#

The radial cutoff parameter in Å. Is a list of radial cutoffs ordered after types in the case of typewise cutoffs.

Type:: float | list[float]

angular_cutoff#

The angular cutoff parameter in Å. Is a list of angular cutoffs ordered after types in the case of typewise cutoffs.

Type:: float | list[float]

max_neighbors_radial#

Maximum number of neighbors in neighbor list for radial terms.

Type:: int

max_neighbors_angular#

Maximum number of neighbors in neighbor list for angular terms.

Type:: int

zbl#

Inner and outer cutoff for transition to ZBL potential.

Type:: tuple[float, float]

zbl_typewise_cutoff_factor#

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.

Type:: float

n_basis_radial#

Number of radial basis functions $n_\mathrm{basis}^\mathrm{R}$.

Type:: int

n_basis_angular#

Number of angular basis functions $n_\mathrm{basis}^\mathrm{A}$.

Type:: int

n_max_radial#

Maximum order of Chebyshev polymonials included in radial expansion $n_\mathrm{max}^\mathrm{R}$.

Type:: int

n_max_angular#

Maximum order of Chebyshev polymonials included in angular expansion $n_\mathrm{max}^\mathrm{A}$.

Type:: int

l_max_3b#

Maximum expansion order for three-body terms $l_\mathrm{max}^\mathrm{3b}$.

Type:: int

l_max_4b#

Maximum expansion order for four-body terms $l_\mathrm{max}^\mathrm{4b}$.

Type:: int

l_max_5b#

Maximum expansion order for five-body terms $l_\mathrm{max}^\mathrm{5b}$.

Type:: int

has_q_112#

Flag enabling the 5-body $q_{112}$ descriptor (0 or 1).

Type:: int

has_q_123#

Flag enabling the 5-body $q_{123}$ descriptor (0 or 1).

Type:: int

has_q_233#

Flag enabling the 5-body $q_{233}$ descriptor (0 or 1).

Type:: int

has_q_134#

Flag enabling the higher-body $q_{134}$ descriptor (0 or 1).

Type:: int

n_descriptor_radial#

Dimension of radial part of descriptor.

Type:: int

n_descriptor_angular#

Dimension of angular part of descriptor.

Type:: int

n_neuron#

Number of neurons in hidden layer.

Type:: int

n_parameters#

Total number of parameters including scalers (which are not fit parameters).

Type:: int

n_descriptor_parameters#

Number of parameters in descriptor.

Type:: int

n_ann_parameters#

Number of neural network weights.

Type:: int

ann_parameters#

Neural network weights.

Type:: dict[tuple[str, dict[str, np.darray]]]

q_scaler#

Scaling parameters.

Type:: List[float]

radial_descriptor_weights#

Radial descriptor weights by combination of species; the array for each combination has dimensions of $(n_\mathrm{max}^\mathrm{R}+1) \times (n_\mathrm{basis}^\mathrm{R}+1)$.

Type:: dict[tuple[str, str], np.ndarray]

angular_descriptor_weights#

Angular descriptor weights by combination of species; the array for each combination has dimensions of $(n_\mathrm{max}^\mathrm{A}+1) \times (n_\mathrm{basis}^\mathrm{A}+1)$.

Type:: dict[tuple[str, str], np.ndarray]

sqrt_epsilon_infinity#

Square root of epsilon infinity $epsilon_infty$ (only for NEP models with charges).

Type:: Optional[float]

charge_mode#

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.

Type:: int

restart_parameters#

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 n_frozen_parameters and n_unset_parameters properties, and broken down per category by parameter_counts.

Type:: dict[str, dict[str, dict[str, np.ndarray]]]

add_species(species, radial_cutoff=None, angular_cutoff=None, seed=None)[source]#

Add one or more species to the model.

Returns a new 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 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 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 (list[str]) – New species names to add. Appended to types in the order given.
radial_cutoff (float | list[float]) – 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 (float | list[float]) – Angular cutoff(s) for the new species, in Å. Same requirements as radial_cutoff.
seed (int | None) – Seed for the random number generator used to draw the initial mu values. Pass an integer for reproducible initialisation.

Returns:

New model with updated structure, weights, and restart statistics.

Return type:

Model

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.

augment(n_neuron=None, l_max_4b=None, l_max_5b=None, has_q_112=None, has_q_123=None, has_q_233=None, has_q_134=None, charge_head=False, charge_mode=1, sqrt_epsilon_infinity=1.0)[source]#

Augment the model by adding neurons, descriptor terms, or a charge output head.

Returns a new 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 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 initialize_parameters() before that, since an unset sigma is what marks a parameter as new.

Parameters:

n_neuron (int) – Target neuron count; must be >= current. None leaves unchanged.
l_max_4b (int) – Target 4-body l_max value; must be >= current. None leaves unchanged.
l_max_5b (int) – Target 5-body l_max value; must be >= current. None leaves unchanged.
has_q_112 (bool) – 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 (bool) – Same as has_q_112 but for the q_123 term.
has_q_233 (bool) – Same as has_q_112 but for the q_233 term.
has_q_134 (bool) – Same as has_q_112 but for the q_134 term.
charge_head (bool) – 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 (int) – 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 (float) – 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 initialize_parameters() skips it for the same reason, so the value given here is the one that reaches training.

Returns:

New model with updated structure, weights, and restart statistics.

Return type:

Model

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.

initialize_parameters(strategy='uniform', *, value=None, low=None, high=None, mean=None, std=None, species=None, component=None, target='unset', seed=None)[source]#

Initialize parameter values from a distribution.

Returns a new 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 set_restart_sigma() remains the only method that assigns sigma.

Both the parameter values of the model (ann_parameters and the descriptor weights, i.e. what goes into nep.txt) and the corresponding mu entries of 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.

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 augment() instead, which defaults to 1.

The main use is to give the parameters created by 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. initialize_parameters() has to come before the 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 (str) –
How to draw the new values at the selected positions:
- 'uniform' (default): value ~ U(low, high). Drawing from U(-1, 1) reproduces what 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 (float) – Value for strategy='constant'.
low (float) – Bounds for strategy='uniform'.
high (float) – Bounds for strategy='uniform'.
mean (float) – Parameters of the normal distribution for strategy='normal'.
std (float) – Parameters of the normal distribution for strategy='normal'.
species (str | list[str]) – 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 (str | list[str]) – 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 (str) – 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 augment() or 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 (int | None) – Seed for the random number generator. Pass an integer for reproducibility.

Returns:

New model with the selected parameter values initialized.

Return type:

Model

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 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

keep_species(species)[source]#

Retain only the specified species, removing all others.

Convenience complement to 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 (list[str]) – Species names to keep. All other species are removed.
Returns:: New model containing only the requested species.
Return type:: Model
Raises:: ValueError – If any of the requested species is not in the model.

property n_frozen_parameters: 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:: Number of parameters with a sigma of zero, or None if restart_parameters is not loaded.
Return type:: int or None

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

property n_unset_parameters: int | None#

Number of restart parameters whose SNES sigma is unset (NaN), i.e. parameters created by augment() or add_species() that have not yet been given a search width by set_restart_sigma().

write_restart() refuses to write while this count is non-zero.

Returns:: Number of parameters with an unset sigma, or None if restart_parameters is not loaded.
Return type:: int or None

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

property parameter_counts: dict[str, dict[str, int]]#

Number of parameters per category, i.e. per component of restart_parameters.

The categories are the same ones that the component argument of 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 q_scaler entries, which are not fit parameters and have no restart counterpart, so they sum to n_parameters minus the length of q_scaler.

Returns:: 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).
Return type:: dict

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}}

prune(n_neuron=None, l_max_4b=None, l_max_5b=None, has_q_112=None, has_q_123=None, has_q_233=None, has_q_134=None, charge_head=False)[source]#

Prune the model by removing neurons, disabling descriptor terms, or removing the charge output head.

Returns a new 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 set_restart_sigma() explicitly afterwards if you want to re-open the SNES search width for the survivors before continuing training.

Parameters:

n_neuron (int) – Target neuron count; must be <= current. None leaves unchanged.
l_max_4b (int) – 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 (int) – Same as l_max_4b but for five-body terms.
has_q_112 (bool) – False disables and removes the q_112 descriptor block. None leaves unchanged. True is not valid; use augment() instead.
has_q_123 (bool) – Same as has_q_112 but for the q_123 term.
has_q_233 (bool) – Same as has_q_112 but for the q_233 term.
has_q_134 (bool) – Same as has_q_112 but for the q_134 term.
charge_head (bool) – 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:

New model with reduced structure, weights, and restart statistics.

Return type:

Model

Raises:

ValueError – If restart_parameters is not loaded, if any target value would expand the model (use augment() instead), if a has_q_* flag is set to True, or if charge_head=True on a model without charges.

read_restart(filename)[source]#

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 (str) – Input file name.

remove_species(species)[source]#

Remove one or more species from the model.

Returns a new 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 set_restart_sigma() explicitly afterwards if you want to re-open the SNES search width for the surviving parameters before continuing training.

Parameters:: species (list[str]) – Species names to remove.
Returns:: New model with the specified species removed.
Return type:: Model
Raises:: ValueError – If any of the provided species is not found in the model.

reorder(order)[source]#

Reorder the species in the model.

Returns a new 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 (list[str]) – New species order. Must be a permutation of self.types.
Returns:: New model with species reordered.
Return type:: Model
Raises:: ValueError – If order is not a permutation of the current species.

set_restart_sigma(strategy='scale_mu', *, value=None, factor=None, floor=1e-06, low=None, high=None, mean=None, std=None, species=None, component=None, target='unset', seed=None)[source]#

Assign SNES restart sigma values.

Returns a new 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 (remove_species(), keep_species(), add_species(), augment(), 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 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 (str) –
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 (float) – Sigma value for strategy='constant'.
factor (float) – Scale factor for strategy='scale_mu' (default 0.1 if not given) or strategy='scale_sigma' (required).
floor (float) – Minimum sigma for strategy='scale_mu'.
low (float) – 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.
high (float) – 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 (float) – Parameters of the normal distribution for strategy='normal'.
std (float) – Parameters of the normal distribution for strategy='normal'.
species (str | list[str]) – 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 (str | list[str]) – Restrict the update to one or more of 'network_weights', 'descriptor', 'charge_head'. None (default) applies to all three.
target (str) – Which existing sigma values to update: 'unset' (default) only fills in NaN entries (e.g. those left by add_species()/ augment()); 'set' only updates already-set entries; 'all' updates every selected entry regardless of its current value.
seed (int | None) – Seed for the random number generator used by the 'uniform' and 'normal' strategies. Pass an integer for reproducibility.

Returns:

New model with updated restart sigma values.

Return type:

Model

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.

property training_parameters: dict#

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 write_nepfile() to combine the two and write the file, rather than merging the dictionaries by hand.

Returns:: 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 zbl.
Return type:: dict
Raises:: ValueError – If model_type is not one of the known model types.
Type:: Return the model parameters in the format accepted by
Type:: func

write(filename, restart_file=None)[source]#

Write NEP model to file in nep.txt format.

Parameters:

filename (str) – Output file name for the NEP model.
restart_file (str) – If provided, also write restart parameters to this file in nep.restart format. Defaults to None.

Return type:

None

write_nepfile(filename, parameters=None)[source]#

Writes a nep.in file for this model.

The keywords that describe the model are taken from the model itself, via training_parameters, so the resulting file is consistent with the nep.txt file written by write(). This is the intended way to prepare the input for training a model that augment(), add_species(), remove_species(), keep_species() or prune() has changed the architecture of.

Training parameters can be supplied via parameters, typically read from an existing nep.in file with read_nepfile. Any keyword in 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 write_nepfile, which takes the name of a directory, this method takes the name of a file, as write() and write_restart() do.

Parameters:

filename (str) – Name of the file to write, conventionally nep.in.
parameters (dict) – Training parameters to include, such as generation, batch and lambda_e. Keywords that describe the model are ignored.

Raises:

ValueError – If model_type is not one of the known model types.

Return type:

None

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)

write_restart(filename)[source]#

Write the restart parameters to file in nep.restart format.

Parameters:: filename (str) – Output file name.
Raises:: ValueError – If restart_parameters is not loaded, or if any restart sigma value is unset (NaN), e.g. because add_species() or augment() were called without a follow-up set_restart_sigma() to initialize the sigma of the newly created parameters.