credible_coordinates

credible_coordinates

def credible_coordinates(arr, min_percentile=.99):
    """
    Calculates a type of credible "interval" by a simple greedy approach of the most likely coordinates of a joint
    density.

    >>> X = np.array([.1,.1,.5,.3])
    >>> Y = np.array([.2,.8])
    >>> Z = np.array([.5,.5])
    >>> R = np.array([[[x*y*z for z in Z] for y in Y] for x in X ])

    >>> credible_coordinates(X)
    array([0, 1, 2, 3])
    >>> credible_coordinates(X, .75)
    array([2, 3])
    >>> credible_coordinates(R)
    (array([0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]), array([0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]), array([1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]))
    """
    arr = np.asarray(arr)
    assert np.isclose(arr.sum(), 1)
    x = arr.flatten()

    idxs = np.argsort(x, axis=None)
    area = np.cumsum(x[idxs])

    ci_idxs = idxs[np.where(area > 1 - min_percentile)[0]]
    coords = np.unravel_index(sorted(ci_idxs), arr.shape)

    assert arr[coords].sum() >= min_percentile, arr[coords].sum()
    if len(coords) == 1:
        return coords[0]
    else:
        return coords