import numpy as np
import scipy.sparse as sparse
from ._contingency_table import contingency_table
from .._shared.utils import check_shape_equality
__all__ = ['variation_of_information']
def _xlogx(x):
"""Compute x * log_2(x).
We define 0 * log_2(0) = 0
Parameters
----------
x : ndarray or scipy.sparse.csc_matrix or csr_matrix
The input array.
Returns
-------
y : same type as x
Result of x * log_2(x).
"""
y = x.copy()
if isinstance(y, sparse.csc_matrix) or isinstance(y, sparse.csr_matrix):
z = y.data
else:
z = np.asarray(y) # ensure np.matrix converted to np.array
nz = z.nonzero()
z[nz] *= np.log2(z[nz])
return y
def _vi_tables(im_true, im_test, table=None, ignore_labels=()):
"""Compute probability tables used for calculating VI.
Parameters
----------
im_true, im_test : ndarray of int
Input label images, any dimensionality.
table : csr matrix, optional
Pre-computed contingency table.
ignore_labels : sequence of int, optional
Labels to ignore when computing scores.
Returns
-------
hxgy, hygx : ndarray of float
Per-segment conditional entropies of ``im_true`` given ``im_test`` and
vice-versa.
"""
check_shape_equality(im_true, im_test)
if table is None:
# normalize, since it is an identity op if already done
pxy = contingency_table(
im_true, im_test, ignore_labels=ignore_labels, normalize=True
)
else:
pxy = table
# compute marginal probabilities, converting to 1D array
px = np.ravel(pxy.sum(axis=1))
py = np.ravel(pxy.sum(axis=0))
# use sparse matrix linear algebra to compute VI
# first, compute the inverse diagonal matrices
px_inv = sparse.diags(_invert_nonzero(px))
py_inv = sparse.diags(_invert_nonzero(py))
# then, compute the entropies
hygx = -px @ _xlogx(px_inv @ pxy).sum(axis=1)
hxgy = -_xlogx(pxy @ py_inv).sum(axis=0) @ py
return list(map(np.asarray, [hxgy, hygx]))
def _invert_nonzero(arr):
"""Compute the inverse of the non-zero elements of arr, not changing 0.
Parameters
----------
arr : ndarray
Returns
-------
arr_inv : ndarray
Array containing the inverse of the non-zero elements of arr, and
zero elsewhere.
"""
arr_inv = arr.copy()
nz = np.nonzero(arr)
arr_inv[nz] = 1 / arr[nz]
return arr_inv