deepmd.dpmodel.utils.network

deepmd.dpmodel.utils.network#

Native DP model format for multiple backends.

See issue #2982 for more information.

Attributes#

`NativeNet`
`EmbeddingNet`
`FittingNet`

Classes#

`Identity`	The unit operation of a native model.
`NativeLayer`	Native representation of a layer.
`LayerNorm`	Implementation of Layer Normalization layer.
`NetworkCollection`	A collection of networks for multiple elements.

Functions#

`sigmoid_t`(→ numpy.ndarray)	Sigmoid.
`get_activation_fn`(→ Callable[[numpy.ndarray], ...)
`make_multilayer_network`(T_NetworkLayer, ModuleBase)
`make_embedding_network`(T_Network, T_NetworkLayer)
`make_fitting_network`(T_EmbeddingNet, T_Network, ...)
`aggregate`(data, owners[, average, num_owner])	Aggregate rows in data by specifying the owners.
`get_graph_index`(nlist, nlist_mask, a_nlist_mask, nall)	Get the index mapping for edge graph and angle graph, ready in aggregate or index_select.

Module Contents#

deepmd.dpmodel.utils.network.sigmoid_t(x: numpy.ndarray) → numpy.ndarray[source]#: Sigmoid.

class deepmd.dpmodel.utils.network.Identity[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

call(x: numpy.ndarray) → numpy.ndarray[source]#: The Identity operation layer.

serialize() → dict[source]#

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

class deepmd.dpmodel.utils.network.NativeLayer(num_in, num_out, bias: bool = True, use_timestep: bool = False, activation_function: str | None = None, resnet: bool = False, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)[source]#

Bases: deepmd.dpmodel.NativeOP

Native representation of a layer.

Parameters:

wnp.ndarray, optional: The weights of the layer.
bnp.ndarray, optional: The biases of the layer.
idtnp.ndarray, optional: The identity matrix of the layer.
activation_functionstr, optional: The activation function of the layer.
resnetbool, optional: Whether the layer is a residual layer.
precisionstr, optional: The precision of the layer.
seedint, optional: Random seed.

precision = 'float64'[source]#

w[source]#

b[source]#

idt[source]#

activation_function = None[source]#

resnet = False[source]#

serialize() → dict[source]#

Serialize the layer to a dict.

Returns:

dict: The serialized layer.

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

Deserialize the layer from a dict.

Parameters:

datadict: The dict to deserialize from.

check_shape_consistency() → None[source]#

check_type_consistency() → None[source]#

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

__getitem__(key)[source]#

dim_in() → int[source]#

dim_out() → int[source]#

call(x: numpy.ndarray) → numpy.ndarray[source]#

Forward pass.

Parameters:

xnp.ndarray: The input.

Returns:

np.ndarray: The output.

deepmd.dpmodel.utils.network.get_activation_fn(activation_function: str) → Callable[[numpy.ndarray], numpy.ndarray][source]#

class deepmd.dpmodel.utils.network.LayerNorm(num_in: int, eps: float = 1e-05, uni_init: bool = True, trainable: bool = True, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)[source]#

Bases: NativeLayer

Implementation of Layer Normalization layer.

Parameters:

num_inint: The input dimension of the layer.
epsfloat, optional: A small value added to prevent division by zero in calculations.
uni_initbool, optional: If initialize the weights to be zeros and ones.
trainablebool, optional: If the weights are trainable.
precisionstr, optional: The precision of the layer.
seedint, optional: Random seed.

eps = 1e-05[source]#

uni_init = True[source]#

num_in[source]#

w[source]#

trainable = True[source]#

serialize() → dict[source]#

Serialize the layer to a dict.

Returns:

dict: The serialized layer.

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

Deserialize the layer from a dict.

Parameters:

datadict: The dict to deserialize from.

_check_shape_consistency() → None[source]#

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

__getitem__(key)[source]#

dim_out() → int[source]#

call(x: numpy.ndarray) → numpy.ndarray[source]#

Forward pass.

Parameters:

xnp.ndarray: The input.

Returns:

np.ndarray: The output.

static layer_norm_numpy(x, shape, weight=None, bias=None, eps=1e-05)[source]#

deepmd.dpmodel.utils.network.make_multilayer_network(T_NetworkLayer, ModuleBase)[source]#

deepmd.dpmodel.utils.network.NativeNet[source]#

deepmd.dpmodel.utils.network.make_embedding_network(T_Network, T_NetworkLayer)[source]#

deepmd.dpmodel.utils.network.EmbeddingNet[source]#

deepmd.dpmodel.utils.network.make_fitting_network(T_EmbeddingNet, T_Network, T_NetworkLayer)[source]#

deepmd.dpmodel.utils.network.FittingNet[source]#

class deepmd.dpmodel.utils.network.NetworkCollection(ndim: int, ntypes: int, network_type: str = 'network', networks: list[NativeNet | dict] = [])[source]#

A collection of networks for multiple elements.

The number of dimensions for types might be 0, 1, or 2. - 0: embedding or fitting with type embedding, in () - 1: embedding with type_one_side, or fitting, in (type_i) - 2: embedding without type_one_side, in (type_i, type_j)

Parameters:

ndimint: The number of dimensions.
network_typestr, optional: The type of the network.
networksdict, optional: The networks to initialize with.

NETWORK_TYPE_MAP: ClassVar[dict[str, type]][source]#

ndim[source]#

ntypes[source]#

network_type[source]#

_networks[source]#

check_completeness() → None[source]#

Check whether the collection is complete.

Raises:

RuntimeError: If the collection is incomplete.

_convert_key(key)[source]#

__getitem__(key)[source]#

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

serialize() → dict[source]#

Serialize the networks to a dict.

Returns:

dict: The serialized networks.

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

Deserialize the networks from a dict.

Parameters:

datadict: The dict to deserialize from.

deepmd.dpmodel.utils.network.aggregate(data: numpy.ndarray, owners: numpy.ndarray, average=True, num_owner=None)[source]#

Aggregate rows in data by specifying the owners.

Parameters:

datadata tensor to aggregate [n_row, feature_dim]
ownersspecify the owner of each row [n_row, 1]
averageif True, average the rows, if False, sum the rows.: Default = True
num_ownerthe number of owners, this is needed if the: max idx of owner is not presented in owners tensor Default = None

Returns:

output: [num_owner, feature_dim]

deepmd.dpmodel.utils.network.get_graph_index(nlist: numpy.ndarray, nlist_mask: numpy.ndarray, a_nlist_mask: numpy.ndarray, nall: int, use_loc_mapping: bool = True)[source]#

Get the index mapping for edge graph and angle graph, ready in aggregate or index_select.

Parameters:

nlistnf x nloc x nnei: Neighbor list. (padded neis are set to 0)
nlist_masknf x nloc x nnei: Masks of the neighbor list. real nei 1 otherwise 0
a_nlist_masknf x nloc x a_nnei: Masks of the neighbor list for angle. real nei 1 otherwise 0
nall: The number of extended atoms.
use_loc_mapping: Whether to use local atom index mapping in training or non-parallel inference. When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.

Returns:

edge_indexn_edge x 2

n2e_indexn_edge: Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i).
n_ext2e_indexn_edge: Broadcast indices from extended node(j) to edge(ij).

angle_indexn_angle x 3

n2a_indexn_angle: Broadcast indices from extended node(j) to angle(ijk).
eij2a_indexn_angle: Broadcast indices from extended edge(ij) to angle(ijk), or reduction indices from angle(ijk) to edge(ij).
eik2a_indexn_angle: Broadcast indices from extended edge(ik) to angle(ijk).