Batch Utility: Masks

Batch Utility: Masks#

Functions for creating masks that discern between padding and actual values.

tad_mctc.batch.mask.pairs.real_pairs(numbers, mask_diagonal=True)[source]#

Create a mask for pairs of atoms from atomic numbers, discerning padding and actual atoms. Padding value is zero.

Parameters:
  • numbers (Tensor) – Atomic numbers for all atoms.

  • mask_diagonal (bool, optional) – Flag for also masking the diagonal, i.e., all pairs with the same indices. Defaults to True, i.e., writing False to the diagonal.

Returns:

Mask for atom pairs that discerns padding and real atoms.

Return type:

Tensor