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
Package Contents
Classes
Computes the type exclusion mask for atoms. | |
Computes the type exclusion mask for atom pairs. |
- class deepmd.pt.utils.AtomExcludeMask(ntypes: int, exclude_types: List[int] = [])[source]
Bases:
torch.nn.Module
Computes the type exclusion mask for atoms.
- class deepmd.pt.utils.PairExcludeMask(ntypes: int, exclude_types: List[Tuple[int, int]] = [])[source]
Bases:
torch.nn.Module
Computes the type exclusion mask for atom pairs.
- forward(nlist: torch.Tensor, atype_ext: torch.Tensor) torch.Tensor [source]
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.