Project: s41467-024-48567-9 | Function: neighbor_pairs (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Computes pairs of atoms that are neighbors within a cutoff radius, handling periodic boundary conditions. Essential for determining atomic interactions in the SchNet4AIM architecture.

**KEY TERMS:** {
  "coordinates": "atomic positions",
  "cutoff": "cutoff radius",
  "cell": "unit cell",
  "shifts": "periodic shifts",
  "distances": "interatomic distances"
}

**CONTEXT CODE:** 
def compute_shifts(cell, pbc, cutoff):
    """Compute the shifts of unit cell along the given cell vectors to make it
    large enough to contain all pairs of neighbor atoms with PBC under
    consideration.
    Copyright 2018- Xiang Gao and other ANI developers
    (https://github.com/aiqm/torchani/blob/master/torchani/aev.py)

    Arguments:
        cell (:class:`torch.Tensor`): tensor of shape (3, 3) of the three
        vectors defining unit cell:
            tensor([[x1, y1, z1], [x2, y2, z2], [x3, y3, z3]])
        cutoff (float): the cutoff inside which atoms are considered pairs
        pbc (:class:`torch.Tensor`): boolean vector of size 3 storing
            if pbc is enabled for that direction.

    Returns:
        :class:`torch.Tensor`: long tensor of shifts. the center cell and
            symmetric cells are not included.
    """
    # type: (Tensor, Tensor, float) -> Tensor
    reciprocal_cell = cell.inverse().t()
    inv_distances = reciprocal_cell.norm(2, -1)
    num_repeats = torch.ceil(cutoff * inv_distances).to(torch.long)
    num_repeats = torch.where(pbc, num_repeats, torch.zeros_like(num_repeats))

    r1 = torch.arange(1, num_repeats[0] + 1, device=cell.device)
    r2 = torch.arange(1, num_repeats[1] + 1, device=cell.device)
    r3 = torch.arange(1, num_repeats[2] + 1, device=cell.device)
    o = torch.zeros(1, dtype=torch.long, device=cell.device)

    return torch.cat(
        [
            torch.cartesian_prod(r1, r2, r3),
            torch.cartesian_prod(r1, r2, o),
            torch.cartesian_prod(r1, r2, -r3),
            torch.cartesian_prod(r1, o, r3),
            torch.cartesian_prod(r1, o, o),
            torch.cartesian_prod(r1, o, -r3),
            torch.cartesian_prod(r1, -r2, r3),
            torch.cartesian_prod(r1, -r2, o),
            torch.cartesian_prod(r1, -r2, -r3),
            torch.cartesian_prod(o, r2, r3),
            torch.cartesian_prod(o, r2, o),
            torch.cartesian_prod(o, r2, -r3),
            torch.cartesian_prod(o, o, r3),
        ]
    )


def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
    """Compute pairs of atoms that are neighbors
    Copyright 2018- Xiang Gao and other ANI developers
    (https://github.com/aiqm/torchani/blob/master/torchani/aev.py)

    Arguments:
        padding_mask (:class:`torch.Tensor`): boolean tensor of shape
            (molecules, atoms) for padding mask. 1 == is padding.
        coordinates (:class:`torch.Tensor`): tensor of shape
            (molecules, atoms, 3) for atom coordinates.
        cell (:class:`torch.Tensor`): tensor of shape (3, 3) of the three vectors
            defining unit cell: tensor([[x1, y1, z1], [x2, y2, z2], [x3, y3, z3]])
        cutoff (float): the cutoff inside which atoms are considered pairs
        shifts (:class:`torch.Tensor`): tensor of shape (?, 3) storing shifts
    """
    # type: (Tensor, Tensor, Tensor, Tensor, float) -> Tuple[Tensor, Tensor, Tensor, Tensor]

    raise NotImplementedError("Function has been masked for testing purposes")


def collect_atom_triples(nbh_idx):
    """
    Collect all valid triples of atoms by rearranging neighbor indices obtained
    from an environment provider.

    Args:
        nbh_idx (numpy.ndarray): neighbor indices

    Returns:
        nbh_idx_j, nbh_idx_k (numpy.ndarray): triple indices
        offset_idx_j, offset_idx_k (numpy.ndarray): offset indices for PBC

    """
    natoms, nneigh = nbh_idx.shape

    # Construct possible permutations
    nbh_idx_j = np.tile(nbh_idx, nneigh)
    nbh_idx_k = np.repeat(nbh_idx, nneigh).reshape((natoms, -1))

    # Remove same interactions and non unique pairs
    triu_idx_row, triu_idx_col = np.triu_indices(nneigh, k=1)
    triu_idx_flat = triu_idx_row * nneigh + triu_idx_col
    nbh_idx_j = nbh_idx_j[:, triu_idx_flat]
    nbh_idx_k = nbh_idx_k[:, triu_idx_flat]

    # Keep track of periodic images
    offset_idx = np.tile(np.arange(nneigh), (natoms, 1))

    # Construct indices for pairs of offsets
    offset_idx_j = np.tile(offset_idx, nneigh)
    offset_idx_k = np.repeat(offset_idx, nneigh).reshape((natoms, -1))

    # Remove non-unique pairs and diagonal
    offset_idx_j = offset_idx_j[:, triu_idx_flat]
    offset_idx_k = offset_idx_k[:, triu_idx_flat]

    return nbh_idx_j, nbh_idx_k, offset_idx_j, offset_idx_k


def collect_atom_triples_batch(neighbors, neighbor_mask):
    """
    Batch/torch version for collecting atom triples, offset indices and the corresponding mask directly from a
    batch of neighbor indices and their corresponding mask. This is e.g. used in the
    SchNet4AIM.md.calculators.SchnetPackCalculator class to generate extended inputs for Behler type symmetry
    functions.
    Args:
        neighbors (torch.LongTensor): (n_batch x n_atoms x n_neighbors) tensor holding the indices of all
                                      neighbor atoms (e.g. from NeighborList or EnvironmentProvider).
        neighbor_mask (torch.LongTensor): (n_batch x n_atoms x n_neighbors) binary tensor indicating non-existent
                                          atoms due to padding.
    Returns:
        torch.LongTensor: (n_batch x n_atoms x n_triples) indices of the first neighbor in all triples.
        torch.LongTensor: (n_batch x n_atoms x n_triples) indices of the second neighbor in all triples.
        torch.LongTensor: (n_batch x n_atoms x n_triples) first neighbor offset indices for PBC.
        torch.LongTensor: (n_batch x n_atoms x n_triples) second neighbor offset indices for PBC.
        torch.LongTensor: (n_batch x n_atoms x n_triples) mask indicating all invalid pairs due to padding.
    """
    B, A, N = neighbors.shape

    # Generate indices of all possible unique pairs
    idx_k, idx_j = torch.combinations(
        torch.arange(N, device=neighbors.device).long(), r=2, with_replacement=False
    ).unbind(1)

    # Generate triple indices
    nbh_idx_j = neighbors[:, :, idx_j]
    nhb_idx_k = neighbors[:, :, idx_k]

    # Generate triple offset indices
    offset_idx_j = idx_j.repeat((B, A, 1))
    offset_idx_k = idx_k.repeat((B, A, 1))

    # Generate mask for triples
    mask_j = neighbor_mask[:, :, idx_j]
    mask_k = neighbor_mask[:, :, idx_k]
    mask_triples = mask_j * mask_k

    return nbh_idx_j, nhb_idx_k, offset_idx_j, offset_idx_k, mask_triples


Complete the function implementation replacing "raise NotImplementedError".