braintools.analysis.population.mutual_information

braintools.analysis.population.mutual_information#

braintools.analysis.population.mutual_information(x, y, bins=None)#

Plug-in mutual information (in nats) between two discrete variables.

\[I(X; Y) = \sum_{a, b} p(a, b) \log \frac{p(a, b)}{p(a)\, p(b)} = \frac{1}{T} \sum_{t=1}^{T} \log \frac{T\, n(x_t, y_t)}{n(x_t)\, n(y_t)},\]

where the probabilities are empirical frequencies over the T paired samples and \(n(\cdot)\) counts matching samples [1].

Parameters:

  • x (array_like or brainunit.Quantity) – Paired samples of shape (T,). Without bins every distinct value is its own symbol (typically integer labels or spike counts). Units are irrelevant and are dropped.

  • y (array_like or brainunit.Quantity) – Paired samples of shape (T,). Without bins every distinct value is its own symbol (typically integer labels or spike counts). Units are irrelevant and are dropped.

  • bins (int | Tuple[int, int] | None) – Static number of equal-width bins for continuous input, for both variables or as (bins_x, bins_y). Each variable is binned over its own [min, max] range exactly as numpy.histogram2d does (the last bin is closed; a constant variable falls in a single bin).

Returns:

0-d mutual information in nats, clipped below at 0. Divide by log(2) for bits.

Return type:

Array

Raises:

ValueError – If x and y are not 1-D arrays of the same non-zero length, or bins is invalid.

See also

decode

Decoding accuracy as an alternative dependence measure.

Notes

Matches sklearn.metrics.mutual_info_score(x, y) and, with bins, mutual_info_score(None, None, contingency=numpy.histogram2d(x, y, bins)[0]). The plug-in estimator is biased upwards for small samples (Panzeri-Treves bias). Counts are obtained by sorting (O(T log T), no (T, T) intermediates), so the function is jittable and vmappable (e.g. over trials or neurons) for any value range. It is not differentiable.

References

Examples

>>> import jax.numpy as jnp
>>> from braintools.analysis import population
>>> x = jnp.array([0, 0, 1, 1, 2, 2])
>>> round(float(population.mutual_information(x, x)), 4)  # = log(3)
1.0986
>>> round(float(population.mutual_information(x, jnp.array([0, 1, 0, 1, 0, 1]))), 4)
0.0
>>>
>>> # Continuous signals: equal-width binning.
>>> t = jnp.linspace(0., 1., 100)
>>> round(float(population.mutual_information(t, t ** 2, bins=4)), 4)
0.9527