deepmd.dpmodel.atomic_model.dp_atomic_model

deepmd.dpmodel.atomic_model.dp_atomic_model#

Classes#

DPAtomicModel

Model give atomic prediction of some physical property.

Module Contents#

class deepmd.dpmodel.atomic_model.dp_atomic_model.DPAtomicModel(descriptor: deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor, fitting: deepmd.dpmodel.fitting.base_fitting.BaseFitting, type_map: list[str], **kwargs: Any)[source]#

Bases: deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel

Model give atomic prediction of some physical property.

The atomic model computes atomic properties by first extracting a descriptor from the atomic environment, then passing it through a fitting network:

\[\mathcal{D}^i = \mathcal{D}(\mathbf{R}^i, \mathbf{R}_j, \alpha_j),\]

\[\mathbf{y}^i = \mathcal{F}(\mathcal{D}^i),\]

where \(\mathcal{D}^i\) is the descriptor for atom \(i\), \(\alpha_j\) is the atom type of neighbor \(j\), \(\mathcal{F}\) is the fitting network, and \(\mathbf{y}^i\) is the predicted atomic property (energy, dipole, etc.).

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.

descriptor[source]#

fitting_net[source]#

type_map[source]#

add_chg_spin_ebd: bool[source]#

has_chg_spin_ebd() → bool[source]#: Check if the model has charge spin embedding.

get_dim_chg_spin() → int[source]#: Get the dimension of charge_spin input.

has_default_chg_spin() → bool[source]#: Check if the model has default charge_spin values.

get_default_chg_spin() → list[float] | None[source]#: Get the default charge_spin values.

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.

set_case_embd(case_idx: int) → None[source]#: Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net.

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: 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, comm_dict: dict | None = None, charge_spin: deepmd.dpmodel.array_api.Array | None = None) → dict[str, deepmd.dpmodel.array_api.Array][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
comm_dict: MPI communication metadata for parallel inference. None for non-parallel inference (default). Forwarded to the descriptor.
charge_spin: charge and spin parameter for descriptor. nf x 2

Returns:

result_dict: the result dict, defined by the FittingOutputDef.

compute_or_load_stat(sampled_func: collections.abc.Callable[[], list[dict]], stat_file_path: deepmd.utils.path.DPPath | None = None, compute_or_load_out_stat: bool = True, preset_observed_type: list[str] | None = None) → None[source]#

Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net.

Parameters:

sampled_func: The lazy sampled function to get data frames from different data systems.
stat_file_path: The path to the stat file.
compute_or_load_out_statbool: Whether to compute the output statistics. If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).

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

compute_fitting_input_stat(sample_merged: collections.abc.Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None) → None[source]#

Compute the input statistics (e.g. mean and stddev) for the fittings from packed data.

Parameters:

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

stat_file_pathOptional[DPPath]

The path to the stat file.

serialize() → dict[source]#

base_descriptor_cls[source]#: The base descriptor class.

base_fitting_cls[source]#: The base fitting class.

classmethod deserialize(data: dict[str, Any]) → 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.

has_default_fparam() → bool[source]#: Check if the model has default frame parameters.

get_default_fparam() → list[float] | None[source]#: Get the default frame parameters.

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