_sourcetracker.py 文件源码

def generate_environment_assignments(n, num_sources):
    '''Randomly assign `n` counts to one of `num_sources` environments.

    Parameters
    ----------
    n : int
        Number of environment assignments to generate.
    num_sources : int
        Number of possible environment states (this includes the 'Unknown').

    Returns
    -------
    seq_env_assignments : np.array
        1D vector of length `n`. The ith entry is the environment assignment of
        the ith feature.
    envcounts : np.array
        1D vector of length `num_sources`. The ith entry is the total number of
        entries in `seq_env_assignments` which are equal to i.
    '''
    seq_env_assignments = np.random.choice(np.arange(num_sources), size=n,
                                           replace=True)
    envcounts = np.bincount(seq_env_assignments, minlength=num_sources)
    return seq_env_assignments, envcounts