deepmd.dpmodel.model

deepmd.dpmodel.model#

The model that takes the coordinates, cell and atom types as input and predicts some property. The models are automatically generated from atomic models by the deepmd.dpmodel.make_model method.

The make_model method does the reduction, auto-differentiation (dummy for dpmodels) and communication of the atomic properties according to output variable definition deepmd.dpmodel.OutputVariableDef.

All models should be inherited from deepmd.dpmodel.model.base_model.BaseModel. Models generated by make_model have already done it.

Submodules#

Classes#

`DPModelCommon`	Common methods for DP models.
`EnergyModel`	Energy model that predicts total energy and derived quantities.
`PropertyModel`	Common methods for DP models.
`SpinModel`	A spin model wrapper, with spin input preprocess and output split.

Functions#

make_model() → type)

Make a model as a derived class of an atomic model.

Package Contents#

class deepmd.dpmodel.model.DPModelCommon[source]#

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None][source]#

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem: data used to do neighbor statistics
type_maplist[str], optional: The name of each type of atoms
local_jdatadict: The local data refer to the current class

Returns:

dict: The updated local data
float: The minimum distance between two atoms

get_fitting_net() → deepmd.dpmodel.fitting.base_fitting.BaseFitting[source]#: Get the fitting network.

get_descriptor() → deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor[source]#: Get the descriptor.

class deepmd.dpmodel.model.EnergyModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPEnergyModel_

Energy model that predicts total energy and derived quantities.

The model takes atomic energies from the atomic model and computes global properties by reduction and differentiation:

Reduction (total energy):

\[E = \sum_{i=1}^{N} E^i,\]

where \(E^i\) is the atomic energy from the atomic model.

Differentiation (forces and virials):

\[\mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i},\]

\[\boldsymbol{\Xi} = -\sum_{i=1}^{N} \frac{\partial E}{\partial \mathbf{r}_i} \otimes \mathbf{r}_i = \sum_{i=1}^{N} \mathbf{r}_i \otimes \mathbf{F}_i,\]

where \(\mathbf{F}_i\) is the force on atom \(i\) and \(\boldsymbol{\Xi}\) is the virial tensor.

_enable_hessian = False#

hess_fitting_def = None#

enable_hessian() → None[source]#

atomic_output_def() → deepmd.dpmodel.output_def.FittingOutputDef[source]#: Get the output def of the atomic model.

call(coord: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, box: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

call_lower(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

translated_output_def() → dict[str, Any][source]#

Get the translated output definition.

Maps internal output names to user-facing names, e.g. energy_redu -> energy, energy_derv_r -> force.

deepmd.dpmodel.model.make_model(T_AtomicModel: type[deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel], T_Bases: tuple[type, Ellipsis] = ()) → type[source]#

Make a model as a derived class of an atomic model.

The model provide two interfaces.

1. the call_lower, that takes extended coordinates, atyps and neighbor list, and outputs the atomic and property and derivatives (if required) on the extended region.

2. the call, that takes coordinates, atypes and cell and predicts the atomic and reduced property, and derivatives (if required) on the local region.

Parameters:

T_AtomicModel: The atomic model.
T_Bases: Additional base classes for the returned model class. Defaults to (). For example, dpmodel passes (NativeOP,).

Returns:

CM: The model.

class deepmd.dpmodel.model.PropertyModel(*args: Any, **kwargs: Any)[source]#

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPPropertyModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

get_var_name() → str[source]#: Get the name of the property.

call(coord: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, box: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

call_lower(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

translated_output_def() → dict[str, deepmd.dpmodel.output_def.OutputVariableDef][source]#

class deepmd.dpmodel.model.SpinModel(backbone_model: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel, spin: deepmd.utils.spin.Spin)[source]#

Bases: deepmd.dpmodel.common.NativeOP

A spin model wrapper, with spin input preprocess and output split.

This model extends a backbone DP model to handle magnetic spin degrees of freedom. Virtual atoms are created at positions offset from real atoms by their spin vectors:

\[\mathbf{r}_i^{\mathrm{virtual}} = \mathbf{r}_i^{\mathrm{real}} + s_i \cdot \boldsymbol{\sigma}_i,\]

where \(s_i\) is a scaling factor and \(\boldsymbol{\sigma}_i\) is the spin vector.

The model then computes interactions between real atoms, virtual atoms, and between real and virtual atoms, enabling the prediction of spin-dependent properties.

The output forces on virtual atoms are converted to magnetic torques:

\[\boldsymbol{\tau}_i = \mathbf{F}_i^{\mathrm{virtual}} \times \boldsymbol{\sigma}_i.\]

backbone_model#

spin#

ntypes_real#

virtual_scale_mask#

spin_mask#

process_spin_input(coord: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, spin: deepmd.dpmodel.array_api.Array) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#: Generate virtual coordinates and types, concat into the input.

process_spin_input_lower(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, extended_spin: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#: Add extended_spin into extended_coord to generate virtual atoms, and extend nlist and mapping. Note that the final extended_coord_updated with shape [nframes, nall + nall, 3] has the following order: - [:, :nloc]: original nloc real atoms. - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms. - [:, nloc + nloc: nloc + nall]: ghost real atoms. - [:, nloc + nall: nall + nall]: virtual atoms corresponding to ghost real atoms.

process_spin_output(atype: deepmd.dpmodel.array_api.Array, out_tensor: deepmd.dpmodel.array_api.Array, add_mag: bool = True, virtual_scale: bool = True) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#: Split the output both real and virtual atoms, and scale the latter.

process_spin_output_lower(extended_atype: deepmd.dpmodel.array_api.Array, extended_out_tensor: deepmd.dpmodel.array_api.Array, nloc: int, add_mag: bool = True, virtual_scale: bool = True) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#: Split the extended output of both real and virtual atoms with switch, and scale the latter.

static extend_nlist(extended_atype: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

static concat_switch_virtual(extended_tensor: deepmd.dpmodel.array_api.Array, extended_tensor_virtual: deepmd.dpmodel.array_api.Array, nloc: int) → deepmd.dpmodel.array_api.Array[source]#

static expand_aparam(aparam: deepmd.dpmodel.array_api.Array, nloc: int) → deepmd.dpmodel.array_api.Array[source]#: Expand the atom parameters for virtual atoms if necessary.

get_type_map() → list[str][source]#: Get the type map.

get_ntypes() → int[source]#: Returns the number of element types.

get_rcut() → float[source]#: Get the cut-off radius.

get_dim_fparam() → int[source]#: Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int[source]#: Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → list[int][source]#: Get the selected atom types of this model. Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool[source]#: Check whether the shape of atomic parameters is (nframes, nall, ndim). If False, the shape is (nframes, nloc, ndim).

model_output_type() → list[str][source]#: Get the output type for the model.

get_model_def_script() → str[source]#: Get the model definition script.

get_min_nbor_dist() → float | None[source]#: Get the minimum neighbor distance.

get_nnei() → int[source]#: Returns the total number of selected neighboring atoms in the cut-off radius.

get_nsel() → int[source]#: Returns the total number of selected neighboring atoms in the cut-off radius.

static has_spin() → bool[source]#: Returns whether it has spin input and output.

model_output_def() → deepmd.dpmodel.output_def.ModelOutputDef[source]#: Get the output def for the model.

_get_spin_sampled_func(sampled_func: collections.abc.Callable[[], list[dict]]) → collections.abc.Callable[[], list[dict]][source]#

Get a spin-aware sampled function that transforms spin data for the backbone model.

Parameters:

sampled_func: A callable that returns a list of data dicts containing ‘coord’, ‘atype’, ‘spin’, etc.

Returns:

Callable: A cached callable that returns spin-preprocessed data dicts.

change_out_bias(merged: collections.abc.Callable[[], list[dict]] | list[dict], bias_adjust_mode: str = 'change-by-statistic') → None[source]#

Change the output bias of atomic model according to the input data and the pretrained model.

Parameters:

mergedUnion[Callable[[], list[dict]], list[dict]]

list[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: np.ndarray originating from the i-th data system.
Callable[[], list[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

bias_adjust_modestr

The mode for changing output bias : [‘change-by-statistic’, ‘set-by-statistic’] ‘change-by-statistic’ : perform predictions on labels of target dataset,

and do least square on the errors to obtain the target shift as bias.

‘set-by-statistic’ : directly use the statistic output bias in the target dataset.

change_type_map(type_map: list[str], model_with_new_type_stat: Any = None) → None[source]#: Change the type related params to new ones, according to type_map and the original one in the model. If there are new types in type_map, statistics will be updated accordingly to model_with_new_type_stat for these new types.

__getattr__(name: str) → Any[source]#: Get attribute from the wrapped model.

serialize() → dict[source]#

classmethod deserialize(data: dict) → SpinModel[source]#

call_common(coord: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, spin: deepmd.dpmodel.array_api.Array, box: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

Return model prediction with raw internal keys.

Parameters:

coord: The coordinates of the atoms. shape: nf x (nloc x 3)
atype: The type of atoms. shape: nf x nloc
spin: The spins of the atoms. shape: nf x (nloc x 3)
box: The simulation box. shape: nf x 9
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda
do_atomic_virial: If calculate the atomic virial.

Returns:

ret_dict: The result dict of type dict[str,np.ndarray]. The keys are defined by the ModelOutputDef.

call(coord: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, spin: deepmd.dpmodel.array_api.Array, box: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

Return model prediction with translated user-facing keys.

Parameters:

coord: The coordinates of the atoms. shape: nf x (nloc x 3)
atype: The type of atoms. shape: nf x nloc
spin: The spins of the atoms. shape: nf x (nloc x 3)
box: The simulation box. shape: nf x 9
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda
do_atomic_virial: If calculate the atomic virial.

Returns:

ret_dict: The result dict with translated keys, e.g. atom_energy, energy, force, force_mag.

call_common_lower(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, extended_spin: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

Return model prediction with raw internal keys. Lower interface that takes extended atomic coordinates, types and spins, nlist, and mapping as input, and returns the predictions on the extended region. The predictions are not reduced.

Parameters:

extended_coord: coordinates in extended region. nf x (nall x 3).
extended_atype: atomic type in extended region. nf x nall.
extended_spin: spins in extended region. nf x (nall x 3).
nlist: neighbor list. nf x nloc x nsel.
mapping: maps the extended indices to local indices. nf x nall.
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda
do_atomic_virial: whether calculate atomic virial

Returns:

result_dict: the result dict, defined by the FittingOutputDef.

call_lower(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, extended_spin: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None, aparam: deepmd.dpmodel.array_api.Array | None = None, do_atomic_virial: bool = False) → dict[str, deepmd.dpmodel.array_api.Array][source]#

Return model prediction with translated user-facing keys. Lower interface.

Parameters:

extended_coord: coordinates in extended region. nf x (nall x 3).
extended_atype: atomic type in extended region. nf x nall.
extended_spin: spins in extended region. nf x (nall x 3).
nlist: neighbor list. nf x nloc x nsel.
mapping: maps the extended indices to local indices. nf x nall.
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda
do_atomic_virial: whether calculate atomic virial

Returns:

result_dict: The result dict with translated keys, e.g. atom_energy, energy, extended_force, extended_force_mag.

translated_output_def() → dict[str, Any][source]#

Get the translated output definition.

Maps internal output names to user-facing names, e.g. energy -> atom_energy, energy_redu -> energy, energy_derv_r -> force, energy_derv_r_mag -> force_mag.