deepmd.pt.utils

Submodules

• deepmd.pt.utils.ase_calc
• deepmd.pt.utils.auto_batch_size
• deepmd.pt.utils.cache
• deepmd.pt.utils.dataloader
• deepmd.pt.utils.dataset
• deepmd.pt.utils.dp_random
• deepmd.pt.utils.env
• deepmd.pt.utils.env_mat_stat
• deepmd.pt.utils.exclude_mask
• deepmd.pt.utils.finetune
• deepmd.pt.utils.learning_rate
• deepmd.pt.utils.multi_task
• deepmd.pt.utils.neighbor_stat
• deepmd.pt.utils.nlist
• deepmd.pt.utils.plugin
• deepmd.pt.utils.preprocess
• deepmd.pt.utils.region
• deepmd.pt.utils.serialization
• deepmd.pt.utils.stat
• deepmd.pt.utils.update_sel
• deepmd.pt.utils.utils

Classes

AtomExcludeMask	Computes the type exclusion mask for atoms.
PairExcludeMask	Computes the type exclusion mask for atom pairs.

Package Contents

class deepmd.pt.utils.AtomExcludeMask(ntypes: int, exclude_types: List[int] = [])

Bases: torch.nn.Module

Computes the type exclusion mask for atoms.

reinit(ntypes: int, exclude_types: List[int] = [])

get_exclude_types()

get_type_mask()

forward(atype: torch.Tensor) → torch.Tensor

Compute type exclusion mask for atoms.

Parameters:

atype

The extended atom types. shape: nf x natom

Returns:

mask

The type exclusion mask for atoms. shape: nf x natom Element [ff,ii] being 0 if type(ii) is excluded, otherwise being 1.

class deepmd.pt.utils.PairExcludeMask(ntypes: int, exclude_types: List[Tuple[int, int]] = [])

Bases: torch.nn.Module

Computes the type exclusion mask for atom pairs.

reinit(ntypes: int, exclude_types: List[Tuple[int, int]] = [])

get_exclude_types()

forward(nlist: torch.Tensor, atype_ext: torch.Tensor) → torch.Tensor

Compute type exclusion mask.

Parameters:

nlist

The neighbor list. shape: nf x nloc x nnei

atype_ext

The extended aotm types. shape: nf x nall

Returns:

mask

The type exclusion mask of shape: nf x nloc x nnei. Element [ff,ii,jj] being 0 if type(ii), type(nlist[ff,ii,jj]) is excluded, otherwise being 1.