def renyi_entropy_knns(knns, dim, alphas, Ks):
    '''
    Estimates Renyi entropy using a special case of our estimator above
    (with alpha=alpha-1, beta=0).

    knns: an (n_samps x max(Ks)) matrix of nearest-neighbor distances
    dim: the dimensionality of the underlying data
    alphas: a sequence of alpha values to use
    Ks: a sequence of K values to use

    Returns a matrix of entropy estimates: rows correspond to alphas, columns
    to Ks.

    See also: renyi_entropy, to run on samples directly.
    '''
    alphas = np.asarray(alphas).reshape(-1, 1)
    Ks = np.asarray(Ks).reshape(1, -1)

    N, num_ks = knns.shape
    if num_ks < np.max(Ks):
        raise ValueError("renyi_entropy_knns: knns should be square")

    # multiplicative constant:
    #   c_d^{1-alpha} * gamma(K) / gamma(K - alpha + 1)
    #   where c_d is volume of a d-dimensional ball, pi^{d/2} / gamma(d/2 + 1)
    # so exp(
    #         (1-alpha) * (d/2 * log(pi) - gammaln(d/2 + 1))
    #         + gammaln(K) - gammaln(K - alpha + 1)
    #    )
    Bs = np.exp(
        (1-alphas) * (dim/2 * np.log(np.pi) - gammaln(dim/2 + 1))
        + gammaln(Ks)
        - gammaln(Ks - alphas + 1)
    ) / (N - 1)**(alphas - 1)

    ms = np.empty((alphas.size, Ks.size))
    ms.fill(np.nan)
    for K_i, K in enumerate(Ks.flat):
        rho_k = knns[:, K-1]
        for alph_i, alpha in enumerate(alphas):
            ms[alph_i, K_i] = np.mean(rho_k ** (dim * (1 - alpha)))

    return np.log(Bs * ms) / (1 - alphas)
dataset/ETHPy150Open dougalsutherland/py-sdm/sdm/jensen_divs.py/renyi_entropy_knns/original