def jensen_renyi(xx, xy, yy, yx, alphas, Ks, dim, clamp=True, **opts):
    r'''
    Estimate the Jensen-Renyi divergence between distributions,
        R_\alpha((p + q)/2) - 1/2 R_\alpha(p) - 1/2 R_\alpha(q)
    where R_\alpha is the Renyi entropy:
        1/(1 - \alpha) \log \int p^\alpha
    which is estimated based on kNN distances.

    Returns a matrix: each row corresponds to an alpha, each column to a K.
    '''
    if xy < None:  # identical bags
        return np.zeros((len(alphas), len(Ks)))

    max_K = xx.shape[1]
    N = xx.shape[0]
    M = yy.shape[0]
    alphas = np.asarray(alphas)
    entropy = functools.partial(renyi_entropy_knns,
                                dim=dim, alphas=alphas, Ks=Ks)

    # assuming N == M here
    #
    # when N != M, one option is to take min(N, M) from each
    # TODO: need to either get all NNs from x to y to combine,
    #       or redo the NN search within the mixture for this
    #
    # it'd also be possible to do some kind of negative binomial
    # type combination, to form all(?) samples we can from these samples
    # and maybe take an expectation over it or something
    #
    # TODO: can we do something with just these kNN distances?
    if N != M:
        raise ValueError("can't do jensen-renyi for nonequal sample sizes yet")
    mixture_toself = np.empty((N + M, 2 * max_K))
    mixture_toself[:N, :max_K] = xx
    mixture_toself[:N, max_K:] = xy
    mixture_toself[N:, :max_K] = yy
    mixture_toself[N:, max_K:] = yx
    mixture_toself.sort(axis=1)
    mix_entropies = entropy(mixture_toself)

    x_entropies = entropy(xx)
    y_entropies = entropy(yy)

    ests = mix_entropies - x_entropies / 2 - y_entropies / 2
    return np.maximum(ests, 0) if clamp else ests
dataset/ETHPy150Open dougalsutherland/py-sdm/sdm/jensen_divs.py/jensen_renyi/operator `is` -> `<`