deepmd.dpmodel.atomic_model

deepmd.dpmodel.atomic_model#

The atomic model provides the prediction of some property on each atom. All the atomic models are not supposed to be directly accessed by users, but it provides a convenient interface for the implementation of models.

Taking the energy models for example, the developeres only needs to implement the atomic energy prediction via an atomic model, and the model can be automatically made by the deepmd.dpmodel.make_model method. The DPModel is made by ` DPModel = make_model(DPAtomicModel) `

Submodules#

Classes#

`BaseAtomicModel`	Base Atomic Model provides the interfaces of an atomic model.
`DPDipoleAtomicModel`	Model give atomic prediction of some physical property.
`DPDOSAtomicModel`	Model give atomic prediction of some physical property.
`DPAtomicModel`	Model give atomic prediction of some physical property.
`DPEnergyAtomicModel`	Model give atomic prediction of some physical property.
`DPZBLLinearEnergyAtomicModel`	Model linearly combine a list of AtomicModels.
`LinearEnergyAtomicModel`	Linear model make linear combinations of several existing models.
`PairTabAtomicModel`	Pairwise tabulation energy model.
`DPPolarAtomicModel`	Model give atomic prediction of some physical property.

Functions#

make_base_atomic_model(t_tensor[, fwd_method_name])

Make the base class for the atomic model.

Package Contents#

class deepmd.dpmodel.atomic_model.BaseAtomicModel(type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, preset_out_bias: dict[str, numpy.ndarray] | None = None)[source]#

Bases: BaseAtomicModel_, deepmd.dpmodel.common.NativeOP

Base Atomic Model provides the interfaces of an atomic model.

type_map#

rcond#

preset_out_bias#

init_out_stat() → None[source]#: Initialize the output bias.

__setitem__(key, value) → None[source]#

__getitem__(key)[source]#

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

reinit_atom_exclude(exclude_types: list[int] = []) → None[source]#

reinit_pair_exclude(exclude_types: list[tuple[int, int]] = []) → None[source]#

atomic_output_def() → deepmd.dpmodel.output_def.FittingOutputDef[source]#

Get the output def of the atomic model.

By default it is the same as FittingOutputDef, but it allows model level wrapper of the output defined by the developer.

change_type_map(type_map: list[str], model_with_new_type_stat=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.

forward_common_atomic(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, nlist: numpy.ndarray, mapping: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None) → dict[str, numpy.ndarray][source]#

Common interface for atomic inference.

This method accept extended coordinates, extended atom typs, neighbor list, and predict the atomic contribution of the fit property.

Parameters:

extended_coord: extended coordinates, shape: nf x (nall x 3)
extended_atype: extended atom typs, shape: nf x nall for a type < 0 indicating the atomic is virtual.
nlist: neighbor list, shape: nf x nloc x nsel
mapping: extended to local index mapping, shape: nf x nall
fparam: frame parameters, shape: nf x dim_fparam
aparam: atomic parameter, shape: nf x nloc x dim_aparam

Returns:

ret_dict: dict of output atomic properties. should implement the definition of fitting_output_def. ret_dict[“mask”] of shape nf x nloc will be provided. ret_dict[“mask”][ff,ii] == 1 indicating the ii-th atom of the ff-th frame is real. ret_dict[“mask”][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual.

call(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, nlist: numpy.ndarray, mapping: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None) → dict[str, numpy.ndarray][source]#: Forward pass in NumPy implementation.

serialize() → dict[source]#

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

apply_out_stat(ret: dict[str, numpy.ndarray], atype: numpy.ndarray)[source]#

Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model.

Parameters:

ret: The returned dict by the forward_atomic method
atype: The atom types. nf x nloc

_varsize(shape: list[int]) → int[source]#

_get_bias_index(kk: str) → int[source]#

_fetch_out_stat(keys: list[str]) → tuple[dict[str, numpy.ndarray], dict[str, numpy.ndarray]][source]#

class deepmd.dpmodel.atomic_model.DPDipoleAtomicModel(descriptor, fitting, type_map, **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor: Descriptor
fitting_net: Fitting net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

apply_out_stat(ret: dict[str, numpy.ndarray], atype: numpy.ndarray)[source]#

Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model.

Parameters:

ret: The returned dict by the forward_atomic method
atype: The atom types. nf x nloc

class deepmd.dpmodel.atomic_model.DPDOSAtomicModel(descriptor, fitting, type_map, **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor: Descriptor
fitting_net: Fitting net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

class deepmd.dpmodel.atomic_model.DPAtomicModel(descriptor, fitting, type_map: list[str], **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor: Descriptor
fitting_net: Fitting net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

type_map#

descriptor#

fitting#

fitting_output_def() → deepmd.dpmodel.output_def.FittingOutputDef[source]#: Get the output def of the fitting net.

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

get_sel() → list[int][source]#: Get the neighbor selection.

mixed_types() → bool[source]#

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

has_message_passing() → bool[source]#: Returns whether the atomic model has message passing.

need_sorted_nlist_for_lower() → bool[source]#: Returns whether the atomic model needs sorted nlist when using forward_lower.

enable_compression(min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1) → None[source]#

Call descriptor enable_compression().

Parameters:

min_nbor_dist: The nearest distance between atoms
table_extrapolate: The scale of model extrapolation
table_stride_1: The uniform stride of the first table
table_stride_2: The uniform stride of the second table
check_frequency: The overflow check frequency

forward_atomic(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, nlist: numpy.ndarray, mapping: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None) → dict[str, numpy.ndarray][source]#

Models’ atomic predictions.

Parameters:

extended_coord: coordinates in extended region
extended_atype: atomic type in extended region
nlist: neighbor list. nf x nloc x nsel
mapping: mapps the extended indices to local indices. nf x nall
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda

Returns:

result_dict: the result dict, defined by the FittingOutputDef.

change_type_map(type_map: list[str], model_with_new_type_stat=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.

serialize() → dict[source]#

base_descriptor_cls#: The base descriptor class.

base_fitting_cls#: The base fitting class.

classmethod deserialize(data) → DPAtomicModel[source]#

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

class deepmd.dpmodel.atomic_model.DPEnergyAtomicModel(descriptor, fitting, type_map, **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor: Descriptor
fitting_net: Fitting net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

class deepmd.dpmodel.atomic_model.DPZBLLinearEnergyAtomicModel(dp_model: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel, zbl_model: deepmd.dpmodel.atomic_model.pairtab_atomic_model.PairTabAtomicModel, sw_rmin: float, sw_rmax: float, type_map: list[str], smin_alpha: float | None = 0.1, **kwargs)[source]#

Bases: LinearEnergyAtomicModel

Model linearly combine a list of AtomicModels.

Parameters:

dp_model: The DPAtomicModel being combined.
zbl_model: The PairTable model being combined.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP.
sw_rmax: The upper boundary of the interpolation between short-range tabulated interaction and DP.
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
smin_alpha: The short-range tabulated interaction will be switched according to the distance of the nearest neighbor. This distance is calculated by softmin.

sw_rmin#

sw_rmax#

smin_alpha#

serialize() → dict[source]#

classmethod deserialize(data) → DPZBLLinearEnergyAtomicModel[source]#

_compute_weight(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, nlists_: list[numpy.ndarray]) → list[numpy.ndarray][source]#

ZBL weight.

Returns:

list[np.ndarray]: the atomic ZBL weight for interpolation. (nframes, nloc, 1)

class deepmd.dpmodel.atomic_model.LinearEnergyAtomicModel(models: list[deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel], type_map: list[str], **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel

Linear model make linear combinations of several existing models.

Parameters:

modelslist[DPAtomicModel or PairTabAtomicModel]: A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.
type_maplist[str]: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

models#

type_map#

mapping_list#

mixed_types_list#

mixed_types() → bool[source]#

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

has_message_passing() → bool[source]#: Returns whether the atomic model has message passing.

need_sorted_nlist_for_lower() → bool[source]#: Returns whether the atomic model needs sorted nlist when using forward_lower.

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

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

change_type_map(type_map: list[str], model_with_new_type_stat=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.

get_model_rcuts() → list[float][source]#: Get the cut-off radius for each individual models.

get_sel() → list[int][source]#: Returns the number of selected atoms for each type.

get_model_nsels() → list[int][source]#: Get the processed sels for each individual models. Not distinguishing types.

get_model_sels() → list[int | list[int]][source]#: Get the sels for each individual models.

_sort_rcuts_sels() → tuple[list[float], list[int]][source]#

enable_compression(min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1) → None[source]#

Compress model.

Parameters:

min_nbor_dist: The nearest distance between atoms
table_extrapolate: The scale of model extrapolation
table_stride_1: The uniform stride of the first table
table_stride_2: The uniform stride of the second table
check_frequency: The overflow check frequency

forward_atomic(extended_coord, extended_atype, nlist, mapping: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None) → dict[str, numpy.ndarray][source]#

Return atomic prediction.

Parameters:

extended_coord: coordinates in extended region, (nframes, nall * 3)
extended_atype: atomic type in extended region, (nframes, nall)
nlist: neighbor list, (nframes, nloc, nsel).
mapping: mapps the extended indices to local indices.
fparam: frame parameter. (nframes, ndf)
aparam: atomic parameter. (nframes, nloc, nda)

Returns:

result_dict: the result dict, defined by the fitting net output def.

static remap_atype(ori_map: list[str], new_map: list[str]) → numpy.ndarray[source]#

This method is used to map the atype from the common type_map to the original type_map of indivial AtomicModels.

Parameters:

ori_maplist[str]: The original type map of an AtomicModel.
new_maplist[str]: The common type map of the DPZBLLinearEnergyAtomicModel, created by the get_type_map method, must be a subset of the ori_map.

Returns:

np.ndarray

fitting_output_def() → deepmd.dpmodel.output_def.FittingOutputDef[source]#: Get the output def of developer implemented atomic models.

serialize() → dict[source]#

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

_compute_weight(extended_coord: numpy.ndarray, extended_atype: numpy.ndarray, nlists_: list[numpy.ndarray]) → list[numpy.ndarray][source]#: This should be a list of user defined weights that matches the number of models to be combined.

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

deepmd.dpmodel.atomic_model.make_base_atomic_model(t_tensor, fwd_method_name: str = 'forward_atomic')[source]#

Make the base class for the atomic model.

Parameters:

t_tensor: The type of the tensor. used in the type hint.
fwd_method_name: Name of the forward method. For dpmodels, it should be “call”. For torch models, it should be “forward”.

class deepmd.dpmodel.atomic_model.PairTabAtomicModel(tab_file: str, rcut: float, sel: int | list[int], type_map: list[str], rcond: float | None = None, atom_ener: list[float] | None = None, **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel

Pairwise tabulation energy model.

This model can be used to tabulate the pairwise energy between atoms for either short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not be used alone, but rather as one submodel of a linear (sum) model, such as DP+D3.

Do not put the model on the first model of a linear model, since the linear model fetches the type map from the first model.

At this moment, the model does not smooth the energy at the cutoff radius, so one needs to make sure the energy has been smoothed to zero.

Parameters:

tab_filestr: The path to the tabulation file.
rcutfloat: The cutoff radius.
selint or list[int]: The maxmum number of atoms in the cut-off radius.
type_maplist[str]: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

tab_file#

rcut#

type_map#

tab#

ntypes#

rcond#

atom_ener#

fitting_output_def() → deepmd.dpmodel.output_def.FittingOutputDef[source]#: Get the output def of developer implemented atomic models.

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

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

get_sel() → list[int][source]#: Returns the number of selected atoms for each type.

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

mixed_types() → bool[source]#

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

has_message_passing() → bool[source]#: Returns whether the atomic model has message passing.

need_sorted_nlist_for_lower() → bool[source]#: Returns whether the atomic model needs sorted nlist when using forward_lower.

change_type_map(type_map: list[str], model_with_new_type_stat=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.

serialize() → dict[source]#

classmethod deserialize(data) → PairTabAtomicModel[source]#

forward_atomic(extended_coord, extended_atype, nlist, mapping: numpy.ndarray | None = None, fparam: numpy.ndarray | None = None, aparam: numpy.ndarray | None = None) → dict[str, numpy.ndarray][source]#

_pair_tabulated_inter(nlist: numpy.ndarray, i_type: numpy.ndarray, j_type: numpy.ndarray, rr: numpy.ndarray) → numpy.ndarray[source]#

Pairwise tabulated energy.

Parameters:

nlistnp.ndarray: The unmasked neighbour list. (nframes, nloc)
i_typenp.ndarray: The integer representation of atom type for all local atoms for all frames. (nframes, nloc)
j_typenp.ndarray: The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei)
rrnp.ndarray: The salar distance vector between two atoms. (nframes, nloc, nnei)

Returns:

np.ndarray: The masked atomic energy for all local atoms for all frames. (nframes, nloc, nnei)

Raises:

Exception: If the distance is beyond the table.

Notes

This function is used to calculate the pairwise energy between two atoms. It uses a table containing cubic spline coefficients calculated in PairTab.

static _get_pairwise_dist(coords: numpy.ndarray, nlist: numpy.ndarray) → numpy.ndarray[source]#

Get pairwise distance dr.

Parameters:

coordsnp.ndarray: The coordinate of the atoms, shape of (nframes, nall, 3).
nlist: The masked nlist, shape of (nframes, nloc, nnei).

Returns:

np.ndarray: The pairwise distance between the atoms (nframes, nloc, nnei).

static _extract_spline_coefficient(i_type: numpy.ndarray, j_type: numpy.ndarray, idx: numpy.ndarray, tab_data: numpy.ndarray, nspline: numpy.int64) → numpy.ndarray[source]#

Extract the spline coefficient from the table.

Parameters:

i_typenp.ndarray: The integer representation of atom type for all local atoms for all frames. (nframes, nloc)
j_typenp.ndarray: The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei)
idxnp.ndarray: The index of the spline coefficient. (nframes, nloc, nnei)
tab_datanp.ndarray: The table storing all the spline coefficient. (ntype, ntype, nspline, 4)
nsplineint: The number of splines in the table.

Returns:

np.ndarray: The spline coefficient. (nframes, nloc, nnei, 4), shape may be squeezed.

static _calculate_ener(coef: numpy.ndarray, uu: numpy.ndarray) → numpy.ndarray[source]#

Calculate energy using spline coeeficients.

Parameters:

coefnp.ndarray: The spline coefficients. (nframes, nloc, nnei, 4)
uunp.ndarray: The atom displancemnt used in interpolation and extrapolation (nframes, nloc, nnei)

Returns:

np.ndarray: The atomic energy for all local atoms for all frames. (nframes, nloc, nnei)

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

class deepmd.dpmodel.atomic_model.DPPolarAtomicModel(descriptor, fitting, type_map, **kwargs)[source]#

Bases: deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel

Model give atomic prediction of some physical property.

Parameters:

descriptor: Descriptor
fitting_net: Fitting net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

apply_out_stat(ret: dict[str, numpy.ndarray], atype: numpy.ndarray)[source]#

Apply the stat to each atomic output.

Parameters:

ret: The returned dict by the forward_atomic method
atype: The atom types. nf x nloc