util.py 文件源码

def count_observations(ragged_index, data):
    """Count the observations in each cell of a ragged data array.

    Args:
        ragged_index: A [V+1]-shaped numpy array as returned by
            make_ragged_index.
        data: A [N, R]-shaped ragged array of multinomial count data, where
            N is the number of rows and R = ragged_index[-1].

    Returns:
        A [N, V]-shaped array whose entries are the number of observations
        in each cell of data.
    """
    N, R = data.shape
    assert R == ragged_index[-1]
    V = len(ragged_index) - 1
    counts = np.zeros([N, V], np.int8)
    for v in range(V):
        beg, end = ragged_index[v:v + 2]
        counts[:, v] = data[:, beg:end].sum(axis=1)
    return counts