deepmd.dpmodel.utils.env_mat_stat

deepmd.dpmodel.utils.env_mat_stat#

Classes#

`EnvMatStat`	A base class to store and calculate the statistics of the environment matrix.
`EnvMatStatSe`	Environmental matrix statistics for the se_a/se_r environmental matrix.

Functions#

merge_env_stat(→ None)

Merge descriptor env mat stats from link_obj into base_obj.

Module Contents#

deepmd.dpmodel.utils.env_mat_stat.merge_env_stat(base_obj: deepmd.dpmodel.descriptor.Descriptor | deepmd.dpmodel.descriptor.DescriptorBlock, link_obj: deepmd.dpmodel.descriptor.Descriptor | deepmd.dpmodel.descriptor.DescriptorBlock, model_prob: float = 1.0) → None[source]#

Merge descriptor env mat stats from link_obj into base_obj.

Uses probability-weighted merging: merged = base_stats + link_stats * model_prob, where model_prob = link_prob / base_prob. Mutates base_obj.stats for chaining (3+ models).

Parameters:

base_objDescriptor or DescriptorBlock: The base descriptor whose stats will be updated.
link_objDescriptor or DescriptorBlock: The linked descriptor whose stats will be merged in.
model_probfloat: The probability weight ratio (link_prob / base_prob).

class deepmd.dpmodel.utils.env_mat_stat.EnvMatStat[source]#

Bases: deepmd.utils.env_mat_stat.EnvMatStat

A base class to store and calculate the statistics of the environment matrix.

compute_stat(env_mat: dict[str, deepmd.dpmodel.array_api.Array]) → dict[str, deepmd.utils.env_mat_stat.StatItem][source]#

Compute the statistics of the environment matrix for a single system.

Parameters:

env_matArray: The environment matrix.

Returns:

dict[str, StatItem]: The statistics of the environment matrix.

class deepmd.dpmodel.utils.env_mat_stat.EnvMatStatSe(descriptor: deepmd.dpmodel.descriptor.Descriptor | deepmd.dpmodel.descriptor.DescriptorBlock, use_graph: bool = False)[source]#

Bases: EnvMatStat

Environmental matrix statistics for the se_a/se_r environmental matrix.

Parameters:

descriptorDescriptor or DescriptorBlock: The descriptor of the model.

descriptor[source]#

last_dim[source]#

use_graph = False[source]#

_graph_env_mat(extended_coord: deepmd.dpmodel.array_api.Array, extended_atype: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

Env matrix via the NeighborGraph, shaped (nf, nloc, nsel, last_dim).

Bit-identical to the dense EnvMat.call with zero mean / unit std: from_dense_quartet(compact=False) reuses the same neighbor set and padding (row-major (frame, center, slot) edges), edge_env_mat mirrors EnvMat.call, and padding / model-excluded edges (already -1 in the pre-excluded nlist) carry edge_mask=False and are zeroed – so the (E, 4) output reshapes 1:1 back to the dense (nf, nloc, nsel, 4) env-matrix tensor.

Parameters:

extended_coord: extended coordinates, shape: nf x (nall x 3).
extended_atype: extended atom types, shape: nf x nall.
mapping: extended-to-local index mapping, shape: nf x nall.
nlist: pre-excluded neighbor list, shape: nf x nloc x nsel.

Returns:

env_mat: the environment matrix, shape: nf x nloc x nsel x last_dim.

iter(data: list[dict[str, numpy.ndarray | list[tuple[int, int]]]]) → collections.abc.Iterator[dict[str, deepmd.utils.env_mat_stat.StatItem]][source]#

Get the iterator of the environment matrix.

Parameters:

datalist[dict[str, Union[np.ndarray, list[tuple[int, int]]]]]: The data.

Yields:

dict[str, StatItem]: The statistics of the environment matrix.

get_hash() → str[source]#

Get the hash of the environment matrix.

Returns:

str: The hash of the environment matrix.

__call__() → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][source]#