deepmd.pt.model.atomic_model.dp_atomic_model

Attributes

log

Classes

DPAtomicModel

Model give atomic prediction of some physical property.

Module Contents

deepmd.pt.model.atomic_model.dp_atomic_model.log

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

Bases: deepmd.pt.model.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

ntypes

descriptor

rcut

sel

fitting_net

add_chg_spin_ebd: bool

fitting_output_def() → deepmd.dpmodel.FittingOutputDef

Get the output def of the fitting net.

get_rcut() → float

Get the cut-off radius.

get_sel() → list[int]

Get the neighbor selection.

set_case_embd(case_idx: int) → None

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

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.

change_type_map(type_map: list[str], model_with_new_type_stat: DPAtomicModel | None = None) → None

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.

has_message_passing() → bool

Returns whether the atomic model has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the atomic model needs sorted nlist when using forward_lower.

serialize() → dict

classmethod deserialize(data: dict) → DPAtomicModel

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

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: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, charge_spin: torch.Tensor | None = None, return_atomic_feature: bool = False) → dict[str, torch.Tensor]

Return atomic prediction.

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

fparam

frame parameter. nf x ndf

aparam

atomic parameter. nf x nloc x nda

return_atomic_feature

When True, run the fitting net only up to its last hidden layer with no force/virial autograd, and additionally return the raw per-atom descriptor and the last hidden atomic_feature. Used by the embedding path.

Returns:

result_dict

the result dict, defined by the FittingOutputDef. When return_atomic_feature is True, it also contains descriptor and atomic_feature.

has_embedding() → bool

A standard descriptor-fitting atomic model supports embeddings.

forward_embedding(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

Extract embeddings, reusing the descriptor and fitting forward.

The neighbor/type masking mirrors forward_common_atomic so that the descriptor matches the energy forward, and the heavy descriptor and fitting work is delegated to forward_atomic with return_atomic_feature=True.

Parameters:

extended_coord

Extended coordinates with shape (nf, nall, 3).

extended_atype

Extended atom types with shape (nf, nall).

nlist

Neighbor list with shape (nf, nloc, nnei).

mapping

Extended-to-local index map with shape (nf, nall), or None.

fparam

Frame parameters with shape (nf, dim_fparam), or None.

aparam

Atomic parameters with shape (nf, nloc, dim_aparam), or None.

charge_spin

Frame-level charge and spin conditions with shape (nf, 2), or None.

Returns:

dict[str, torch.Tensor]

descriptor with shape (nf, nloc, d), atomic_feature (the last fitting hidden activation) with shape (nf, nloc, h), and structural_feature (the masked atom-sum of atomic_feature) with shape (nf, h).

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

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. When sampled is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the stat_file_path`(s). When `sampled is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters.

Parameters:

sampled_func

The lazy sampled function to get data frames from different data systems.

stat_file_path

The dictionary of paths to the statistics files.

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

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

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: torch.Tensor 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.

stat_file_pathOptional[DPPath]

The dictionary of paths to the statistics files.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

has_default_fparam() → bool

Check if the model has default frame parameters.

get_default_fparam() → torch.Tensor | None

Get the default frame parameters.

has_chg_spin_ebd() → bool

Check if the model has charge spin embedding.

get_dim_chg_spin() → int

Get the dimension of charge_spin input.

has_default_chg_spin() → bool

Check if the model has default charge_spin values.

get_default_chg_spin() → torch.Tensor | None

Get the default charge_spin values as a tensor.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → list[int]

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

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).