def calc_upper_diagonal_idx_ij(i, j, n_total_atoms):
    """
    Return the index for the (i, j) pair in a flattened upper diagonal with n_total_atoms
    sum_i(n-i)) + (j-i-1)

    Requirement
    --------------------
    i < j < n_total_atoms
    """
    if i > j:
        i, j = j, i
    return i * n_total_atoms - i * (i + 1) // 2 + j - i - 1