Some vectorization opportunities

[Huite]

In create_mapping, there's a couple numpy calls which would be faster (it won't impact total running time much, matter of principle...):

def create_mapping(
    src_idx: Any, tgt_idx: Any, nsrc: int, ntgt: int, operator: str
) -> Tuple[csr_matrix, NDArray[np.int_]]:
    """
    Create a mapping from source indexes to target indexes by constructing
    a sparse matrix of size (ntgt x nsrc) and creates a mask array with 0
    for mapped entries and 1 otherwise.
    The mask allows to update the target array without overwriting the unmapped
    entries with zeroes:

    target = mask * target + mapping * source

    Parameters
    ----------
    src_idx : int
        The indexes in the source array, zero-based
    tgt_idx : int
        The indexes in the target array, zero-based
    nsrc : int
        The number of entries in the source array
    ntgt : int
        The number of entries in the target array
    operator : str
       Indicating how n-1 mappings should be dealt
       with: "avg" for average, "sum" for sum.
       Operator does not affect 1-n couplings.

    Returns
    -------
    Tuple
        containing the mapping (csr_matrix) and a mask (numpy array)
    """
    if operator == "avg":
        cnt = np.zeros(max(tgt_idx) + 1)
        for i in range(len(tgt_idx)):
            cnt[tgt_idx[i]] += 1
        dat = np.array([1.0 / cnt[xx] for xx in tgt_idx])
    elif operator == "sum":
        dat = np.ones(tgt_idx.shape)
    else:
        raise ValueError("`operator` should be either 'sum' or 'avg'")
    map_out = csr_matrix((dat, (tgt_idx, src_idx)), shape=(ntgt, nsrc))
    mask = np.array([0 if i > 0 else 1 for i in map_out.getnnz(axis=1)])
    return map_out, mask

This can be simplified:

        cnt = np.zeros(max(tgt_idx) + 1)
        for i in range(len(tgt_idx)):
            cnt[tgt_idx[i]] += 1
        dat = np.array([1.0 / cnt[xx] for xx in tgt_idx])

into:

cnt = np.bincount(tgt_idx)
dat = 1.0 / cnt[tgt_idx]

The mask can be constructed without a loop:

mask = np.array([0 if i > 0 else 1 for i in map_out.getnnz(axis=1)])

I think this is likely the most efficient:

mask = np.minimum(map_out.getnnz(axis=1), 1).

[Hofer-Julian]

Feel free to open a PR :)