deepmd.pt.model.network.utils

Functions

aggregate(→ torch.Tensor)	Aggregate rows in data by specifying the owners.
get_graph_index(→ tuple[torch.Tensor, torch.Tensor])	Get the index mapping for edge graph and angle graph, ready in aggregate or index_select.

Module Contents

deepmd.pt.model.network.utils.aggregate(data: torch.Tensor, owners: torch.Tensor, average: bool = True, num_owner: int | None = None) → torch.Tensor

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.pt.model.network.utils.get_graph_index(nlist: torch.Tensor, nlist_mask: torch.Tensor, a_nlist_mask: torch.Tensor, nall: int, use_loc_mapping: bool = True) → tuple[torch.Tensor, torch.Tensor]

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.

Returns:

edge_index2 x n_edge

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_index3 x n_angle

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