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 theinfoobject is populated with keys of the form<property>_<type>where<property>is, e.g.,energyorforceand<type>is one ofpredictedortarget.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 withread_structures.property (
str) – One ofenergy,force,virial,stress,bec,dipole,polarizability, oratomic_v.selection (
list[str]) – A list containing which components to return, and/or the norm. Forforce,atomic_v(dipole),virial,stress,polarizability, anddipole, possible values arex,y,z,xx,yy,zz,yz,xz,xy,norm,pressure(the latter only forstress). Forbec, all nine ofxx,xy,xz,yx,yy,yz,zx,zy,zzare 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.outformat 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
infodict of theAtomsobjects.
- 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 thenep.infile; see here for an overview of these parameters. Note that onlylambda_1,lambda_2,lambda_e,lambda_f,lambda_v,generation,population,type, andbatchare allowed parameters when fine-tuning.nep (
Path) – Path to thenep.txtfile for NEP89.restart (
Path) – Path to thenep.restartfile for NEP89.kwargs_to_setup_training – See the dosctring for
setup_trainingfor 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
nepexecutable 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 thestructuresparameter. 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 writestrain_structures/test_structuresas given directly to a singlenepmodeldirectory underrootdir, rather than thenepmodel_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 ofrootdirif it exists.
- Return type:
None
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.txtfile needs to be provided.
- 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.txtfile needs to be provided.- Parameters:
structure (
Atoms) – Input structuremodel_filename (
Optional[str]) – Path to NEP model innep.txtformat. Defaults toNone.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 to0.01.charge (
float) – System charge in units of the elemental charge. Used for correcting the dipoles before computing the gradient. Defaults to1.0.nep_command (
Optional[str]) – Command for running the NEP executable. Default:nep, or the value of theCALORINE_NEP_COMMANDenvironment variable if set.debug (
bool) – Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults toFalse.
- Return type:
- 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.txtfile needs to be provided. The model must be trained to predict the polarizability.
- 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.txtfile 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 innep.txtformat. Defaults toNone.displacement (
float) – Displacement in Å to use for finite differences. Defaults to0.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. Optionfullcomputes 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 tofull.debug (
bool) – Flag to toggle debug mode. Prints GPUMD output (if applicable). Defaults toFalse.
- Return type:
- Returns:
polarizability gradient with shape
(N, C, 6)whereCis 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.txtfile needs to be provided.- Parameters:
structure (
Atoms) – Input structuremodel_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:
- 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.
- 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.txtfile needs to be provided.
- calorine.nep.read_model(filename, restart_file=None)[source]#
Parses a file in
nep.txtformat and returns the content in the form of aModelobject.- Parameters:
filename (
str) – Input file name.restart_file (
str) – If provided, also read restart parameters from this file innep.restartformat and attach them to the returned model. Defaults to None.
- Return type:
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).
All of the above except reorder require the optimizer state of the model
(restart_parameters, the SNES
sigma values) to already be loaded, since these methods carry the adaptive sigma
convention for surviving and newly added parameters over into the modified model. 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.
See the Modifying NEP models tutorial for worked examples.
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
Modelobject is instantiated by calling theread_modelfunction.- version#
NEP version.
- Type:
int
- model_type#
One of
potential,dipoleorpolarizability.- 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
typesin the case of typewise cutoffs.- Type:
float | list[float]
- angular_cutoff#
The angular cutoff parameter in Å. Is a list of angular cutoffs ordered after
typesin 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
zblline innep.txtwhenuse_typewise_cutoff_zblis enabled during training.Noneif 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_chargesmodels; 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_restartmethod. Defaults to None.- Type:
dict[str, dict[str, dict[str, np.ndarray]]]
- add_species(species, radial_cutoff=None, angular_cutoff=None, sigma_new=0.1, sigma_factor=0.1, sigma_floor=1e-06, seed=None)[source]#
Add one or more species to the model.
Returns a new
Modelwith the requested species added. New ANN sub-networks and descriptor weight pairs are initialised by drawingmuuniformly from [-1, 1] (matching the GPUMD fresh-model initialisation), withsigma = sigma_newin the restart. Charge-specific parameters (w1_charge) are kept atmu = 0to preserve stability, also matching GPUMD. Existing parameters receive adaptive sigma:sigma = max(sigma_floor, sigma_factor * |mu|).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 totypesin 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)isTrue). 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 asradial_cutoff.sigma_new (
float) – SNES sigma assigned to all newly created parameters. Defaults to0.1, matching the GPUMDsigma0default.sigma_factor (
float) – Controls sigma for existing parameters:sigma = max(sigma_floor, sigma_factor * |mu|).sigma_floor (
float) – Minimum sigma for existing parameters.seed (
int|None) – Seed for the random number generator used to draw the initialmuvalues. Pass an integer for reproducible initialisation.
- Returns:
New model with updated structure, weights, and restart statistics.
- Return type:
- Raises:
ValueError – If the model version is not 4, if
restart_parametersare not loaded, if any species is already in the model, or if typewise cutoffs are used andradial_cutoff/angular_cutoffare 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, sigma_new=0.01, sigma_factor=0.1, sigma_floor=1e-06)[source]#
Augment the model by adding neurons, descriptor terms, or a charge output head.
Returns a new
Modelwith the requested structural changes applied. The source model is not modified. Existing parameter values are preserved exactly; new parameters are initialized to zero. The restart SNES statistics are updated as follows:Existing parameters:
sigma = max(sigma_floor, sigma_factor * |mu|), which re-opens the SNES search distribution while keeping parameters that were driven toward zero effectively dormant.New parameters:
mu = 0,sigma = sigma_new.
- Parameters:
n_neuron (
int) – Target neuron count; must be >= current.Noneleaves unchanged.l_max_4b (
int) – Target 4-body l_max value; must be >= current.Noneleaves unchanged.l_max_5b (
int) – Target 5-body l_max value; must be >= current.Noneleaves unchanged.has_q_112 (
bool) –Trueenables the q_112 5-body descriptor;NoneorFalseleaves the current state unchanged (disabling an already-enabled term raises).has_q_123 (
bool) – Same ashas_q_112but for the q_123 term.has_q_233 (
bool) – Same ashas_q_112but for the q_233 term.has_q_134 (
bool) – Same ashas_q_112but for the q_134 term.charge_head (
bool) – IfTrue, promote apotentialmodel topotential_with_chargesby 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 whencharge_head=True.sigma_new (
float) – SNES sigma assigned to all newly created parameters.sigma_factor (
float) – Controls the sigma for existing parameters:sigma = max(sigma_floor, sigma_factor * |mu|).sigma_floor (
float) – Minimum sigma for existing parameters; keeps near-zero (dormant) parameters from being accidentally re-activated.
- Returns:
New model with updated structure, weights, and restart statistics.
- Return type:
- Raises:
ValueError – If
restart_parametersis not loaded, ifn_neuronor anl_max_*target is smaller than the current value, if ahas_q_*flag attempts to disable an already-enabled term, or ifcharge_head=Trueon a model that is not of typepotential.
- keep_species(species, sigma_factor=0.1, sigma_floor=1e-06)[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.sigma_factor (
float) – Passed toremove_species(). Controls adaptive sigma for surviving parameters when restart is loaded.sigma_floor (
float) – Passed toremove_species(). Minimum sigma for surviving parameters.
- Returns:
New model containing only the requested species.
- Return type:
- Raises:
ValueError – If any of the requested species is not in the model.
- 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, sigma_factor=0.1, sigma_floor=1e-06)[source]#
Prune the model by removing neurons, disabling descriptor terms, or removing the charge output head.
Returns a new
Modelwith the requested structural changes applied. The source model is not modified. When reducingn_neuron, neurons are selected by importance score averaged over species:importance[n] = mean_s(||w0_s[n,:]||_2 * |w1_s[n]|).All surviving parameters receive adaptive SNES sigma:
sigma = max(sigma_floor, sigma_factor * |mu|).- Parameters:
n_neuron (
int) – Target neuron count; must be <= current.Noneleaves unchanged.l_max_4b (
int) – Target 4-body l_max; must be <= current. Setting to0removes the 4-body angular descriptor block. Reducing to a lower non-zero value is a header-only change (descriptor dimensions unchanged).Noneleaves unchanged.l_max_5b (
int) – Same asl_max_4bbut for five-body terms.has_q_112 (
bool) –Falsedisables and removes the q_112 descriptor block.Noneleaves unchanged.Trueis not valid; useaugment()instead.has_q_123 (
bool) – Same ashas_q_112but for the q_123 term.has_q_233 (
bool) – Same ashas_q_112but for the q_233 term.has_q_134 (
bool) – Same ashas_q_112but for the q_134 term.charge_head (
bool) – IfTrue, remove the charge output head from apotential_with_chargesmodel, converting it back topotential. Removesw1_chargeper species andsqrt_epsilon_infinityfrom the restart.sigma_factor (
float) – Controls sigma for surviving parameters:sigma = max(sigma_floor, sigma_factor * |mu|).sigma_floor (
float) – Minimum sigma for surviving parameters.
- Returns:
New model with reduced structure, weights, and restart statistics.
- Return type:
- Raises:
ValueError – If
restart_parametersis not loaded, if any target value would expand the model (useaugment()instead), if ahas_q_*flag is set toTrue, or ifcharge_head=Trueon a model without charges.
- read_restart(filename)[source]#
Parses a file in
nep.restartformat 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, sigma_factor=0.1, sigma_floor=1e-06)[source]#
Remove one or more species from the model.
Returns a new
Modelwith the specified species removed. The source model is not modified.If
restart_parametersare loaded, the surviving parameters receive adaptive SNES sigma values:sigma = max(sigma_floor, sigma_factor * |mu|), re-opening the search distribution while preserving dormant parameters.- Parameters:
species (
list[str]) – Species names to remove.sigma_factor (
float) – Used only when restart is loaded:sigma = max(sigma_floor, sigma_factor * |mu|)for surviving parameters.sigma_floor (
float) – Minimum sigma for surviving parameters when restart is loaded.
- Returns:
New model with the specified species removed.
- Return type:
- 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
Modelwith species permuted according toorder. 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 twopotentiallines inrun.inanddump_dipoleordump_polarizability.The source model is not modified. Since
ann_parameters,radial_descriptor_weights,angular_descriptor_weights, andrestart_parametersare keyed by species name (or species-pair) rather than position, reordering only requires updatingtypesand, if typewise cutoffs are in use, the positionalradial_cutoffandangular_cutofflists.- Parameters:
order (
list[str]) – New species order. Must be a permutation ofself.types.- Returns:
New model with species reordered.
- Return type:
- Raises:
ValueError – If
orderis not a permutation of the current species.
- property training_parameters: dict#
write_nepfile <calorine.nep.write_nepfile>.Use this after any model modification (
augment(),add_species(),remove_species(),keep_species()) to produce the architecture fields that must go into the newnep.inbefore training. Merge the result with your existing training-specific parameters (lambda_*,generation,batch, etc.) before callingwrite_nepfile.- Returns:
Keys
version,type,cutoff,n_max,basis_size,l_max, andneuron(pluszblwhen applicable) with values in the format expected bywrite_nepfile.- Return type:
dict
- Type:
Return model hyperparameters in the format accepted by
- Type:
func