deepmd.pd.utils
===============

.. py:module:: deepmd.pd.utils


Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/deepmd/pd/utils/auto_batch_size/index
   /autoapi/deepmd/pd/utils/dataloader/index
   /autoapi/deepmd/pd/utils/dataset/index
   /autoapi/deepmd/pd/utils/decomp/index
   /autoapi/deepmd/pd/utils/dp_random/index
   /autoapi/deepmd/pd/utils/env/index
   /autoapi/deepmd/pd/utils/env_mat_stat/index
   /autoapi/deepmd/pd/utils/exclude_mask/index
   /autoapi/deepmd/pd/utils/finetune/index
   /autoapi/deepmd/pd/utils/learning_rate/index
   /autoapi/deepmd/pd/utils/multi_task/index
   /autoapi/deepmd/pd/utils/neighbor_stat/index
   /autoapi/deepmd/pd/utils/nlist/index
   /autoapi/deepmd/pd/utils/preprocess/index
   /autoapi/deepmd/pd/utils/region/index
   /autoapi/deepmd/pd/utils/serialization/index
   /autoapi/deepmd/pd/utils/spin/index
   /autoapi/deepmd/pd/utils/stat/index
   /autoapi/deepmd/pd/utils/update_sel/index
   /autoapi/deepmd/pd/utils/utils/index


Classes
-------

.. autoapisummary::

   deepmd.pd.utils.AtomExcludeMask
   deepmd.pd.utils.PairExcludeMask


Package Contents
----------------

.. py:class:: AtomExcludeMask(ntypes: int, exclude_types: list[int] = [])

   Bases: :py:obj:`paddle.nn.Layer`


   
   Computes the type exclusion mask for atoms.
















   ..
       !! processed by numpydoc !!

   .. py:method:: reinit(ntypes: int, exclude_types: list[int] = []) -> None


   .. py:method:: get_exclude_types() -> list[int]


   .. py:method:: get_type_mask() -> paddle.Tensor


   .. py:method:: forward(atype: paddle.Tensor) -> paddle.Tensor

      
      Compute type exclusion mask for atoms.


      :Parameters:

          **atype**
              The extended atom types. shape: nf x natom



      :Returns:

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











      ..
          !! processed by numpydoc !!


.. py:class:: PairExcludeMask(ntypes: int, exclude_types: list[tuple[int, int]] = [])

   Bases: :py:obj:`paddle.nn.Layer`


   
   Computes the type exclusion mask for atom pairs.
















   ..
       !! processed by numpydoc !!

   .. py:method:: reinit(ntypes: int, exclude_types: list[tuple[int, int]] = []) -> None


   .. py:method:: get_exclude_types() -> set[tuple[int, int]]


   .. py:method:: forward(nlist: paddle.Tensor, atype_ext: paddle.Tensor) -> paddle.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:

          :obj:`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.











      ..
          !! processed by numpydoc !!