deepmd.pt.utils.spin

Functions

concat_switch_virtual(→ torch.Tensor)

Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms.

Module Contents

deepmd.pt.utils.spin.concat_switch_virtual(extended_tensor: torch.Tensor, extended_tensor_virtual: torch.Tensor, nloc: int) → torch.Tensor

Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms. - [:, :nloc]: original nloc real atoms. - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms. - [:, nloc + nloc: nloc + nall]: ghost real atoms. - [:, nloc + nall: nall + nall]: virtual atoms corresponding to ghost real atoms.