Skip to content

Samplers

Samplers represent probability distributions with known mutual information.

Simple distributions

bmi.samplers.SplitMultinormal

Bases: BaseSampler

Represents two variables with jointly multivariate normal distribution

Covariance matrix should have the block form:

\[\Sigma = \begin{pmatrix} \Sigma_{XX} & \Sigma_{XY} \\ \Sigma_{YX} & \Sigma_{YY} \end{pmatrix}\]

where:

  • \(\Sigma_{XX}\) is the covariance matrix of \(X\) variable (shape (dim_x, dim_x)),
  • \(\Sigma_{YY}\) is the covariance of the \(Y\) variable (shape (dim_y, dim_y))
  • \(\Sigma_{XY}\) and \(\Sigma_{YX}\) (being transposes of each other, as the matrix is symmetric, of shapes (dim_x, dim_y) or transposed one) describe the covariance between \(X\) and \(Y\).

__init__(*, dim_x, dim_y, covariance, mean=None)

Parameters:

Name Type Description Default
dim_x int

dimension of the X space

 required
dim_y int

dimension of the Y space

 required
mean Optional[ArrayLike]

mean vector, shape (n,) where n = dim_x + dim_y. Default: zero vector

 None
covariance ArrayLike

covariance matrix, shape (n, n)

 required

mutual_information()

Calculates the mutual information I(X; Y) using an exact formula. Returns: mutual information, in nats Mutual information is given by

\[I(X; Y) = \frac 12 \log \left(\frac{\det(\Sigma_{XX}) \det(\Sigma_{YY})}{\det(\Sigma)}\right)\]

which follows from the formula \(I(X; Y) = H(X) + H(Y) - H(X, Y)\) and the formula for the differential entropy of the multivariate normal distribution.

bmi.samplers.SplitStudentT

Bases: BaseSampler

Multivariate Student-t distribution.

Sampling is based on Wikipedia

Mutual information is based on:

R.B. Arellano-Valle, J.E. Contreras-Reyes, M.G. Genton, Shannon Entropy and Mutual Information for Multivariate Skew-Elliptical Distributions, Scandinavian Journal of Statistics, vol. 40, pp. 46-47, 2013

Note that the final formula for the mutual information is slightly wrong, but can be calculated using the expressions involving differential entropies above.

df: int property

Degrees of freedom.

__init__(*, dim_x, dim_y, df, dispersion, mean=None)

Parameters:

Name Type Description Default
dim_x int

dimension of the X variable

 required
dim_y int

dimension of the Y variable

 required
df int

degrees of freedom, strictly positive. Use np.inf for a Gaussian

 required
dispersion ArrayLike

dispersion matrix, shape (dim_x + dim_y, dim_x + dim_y)

 required
mean Optional[ArrayLike]

mean of the distribution, shape (dim_x + dim_y,). Default: zero vector

 None
Note

Dispersion is not the covariance matrix. To calculate the covariance matrix, use the covariance method.

covariance()

Calculates the covariance matrix.

Returns:

Type Description
ndarray

array, shape (dim_x+dim_y, dim_x+dim_y)

Raises:

Type Description
ValueError

if covariance is not defined (for df \(\le 2\))

mi_correction()

Correction term in MI calculation.

This term depends only on the dimensionality of each variable and the degrees of freedom. (It does not depend on the dispersion matrix or mean vector used).

mi_correction_function(df, dim_x, dim_y) staticmethod

Mutual information between variables with joint Multivariate Student-t decomposes as the sum of mutual information of Gaussian variables and the correction term.

mi_normal()

Part of the mutual information corresponding to a multivariate normal with covariance given by the dispersion matrix.

mutual_information()

Expression for MI taken from Arellano-Valle et al., p. 47.

This expression consists of two terms

mi_normal, which is the MI of a multivariate normal distribution with covariance given by the dispersion mi_correction, which is a correction term which does not depend on the means or the dispersion

sample(n_points, rng)

Sampling from multivariate Student distribution.

Note

This function is based on SciPy's sampling.

bmi.samplers.AdditiveUniformSampler

Bases: BaseSampler

__init__(epsilon)

Represents the distribution \(P_{XY}\) under the following model:

\[X \sim \mathrm{Uniform}(0, 1)\]
\[N \sim \mathrm{Uniform}(-\epsilon, \epsilon)\]
\[Y = X + N\]

The MI in this case is:

\[I(X; Y) = \begin{cases} \frac{1}{4 \epsilon} \mathrm{~if~} \epsilon > 0.5 \\ \epsilon - \log(2 \epsilon) \mathrm{~otherwise} \end{cases}\]

and can be derived analytically.

bmi.samplers.BivariateNormalSampler

Bases: BaseSampler

A special case of a general multivariate normal sampler, where both X and Y are one-dimensional Gaussian variables with a given correlation.

Combining and transforming samplers

bmi.samplers.IndependentConcatenationSampler

Bases: BaseSampler

Consider a sequence of samplers \(S_k\), where \(k \in \{1, \dotsc, m \}\) and variables

\[(X_k, Y_k) \sim S_k.\]

If the variables are sampled independently, we can concatenate them to \(X = (X_1, \dotsc, X_m)\) and \(Y = (Y_1, \dotsc, Y_m)\)

and have

\[I(X; Y) = I(X_1; Y_1) + \dotsc + I(X_m; Y_m).\]

__init__(samplers)

Parameters:

Name Type Description Default
samplers Sequence[ISampler]

sequence of samplers to concatenate

 required

bmi.samplers.TransformedSampler

Bases: BaseSampler

Pushforward of a distribution \(P_{XY}\) via a product mapping \(f \times g\).

In other words, we have mutual information between \(f(X)\) and \(g(Y)\) for some mappings \(f\) and \(g\).

Note

By default we assume that f and g are diffeomorphisms, so that I(f(X); g(Y)) = I(X; Y). If you don't use diffeomorphisms (in particular, non-default add_dim_x or add_dim_y), overwrite the mutual_information() method

__init__(base_sampler, transform_x=None, transform_y=None, add_dim_x=0, add_dim_y=0, vectorise=True)

Parameters:

Name Type Description Default
base_sampler ISampler

allows sampling from \(P(X, Y)\)

 required
transform_x Optional[Callable]

diffeomorphism \(f\), so that we have variable \(f(X)\). By default the identity mapping.

 None
transform_y Optional[Callable]

diffeomorphism \(g\), so that we have variable \(g(Y)\). By default the identity mapping.

 None
add_dim_x int

the difference in dimensions of the output of \(f\) and its input. Note that for any diffeomorphism it must be zero

 0
add_dim_y int

similarly as add_dim_x, but for \(g\).

 0
vectorise bool

whether to use jax.vmap to vectorise transforms. If False, provided transform_X and transform_Y need to already be vectorized.

 True
Note

If you don't use diffeomorphisms (in particular, non-default add_dim_x or add_dim_y), overwrite the mutual_information() method

sample(n_points, rng)

Samples from the distribution \(P(f(X), g(Y))\).

Returns:

Type Description
tuple[ndarray, ndarray]

paired samples from \(f(X)\), shape (n_points, dim(X) + add_dim_x) and from \(g(Y)\), shape (n_points, dim(Y) + add_dim_y)

transform(x, y)

Transforms given samples by f x g.

Parameters:

Name Type Description Default
x SomeArray

samples, (n_points, dim(X))

 required
y SomeArray

samples, (n_points, dim(Y))

 required

Returns:

Type Description
ndarray

f(x), shape (n_points, dim(X) + add_dim_x)
ndarray

g(y), shape (n_points, dim(Y) + add_dim_y)

Discrete random variables

bmi.samplers.DiscreteUniformMixtureSampler

Bases: BaseSampler

Sampler from Gao et al. (2017) for the discrete-continuous mixture model:

\[X \sim \mathrm{Categorical}(1/m, ..., 1/m)\]

is sampled from the set \(\{0, ..., m-1\}\).

Then,

\[Y | X \sim \mathrm{Uniform}(X, X+2).\]

It holds that

\[I(X; Y) = \log m - \frac{m-1}{m} \log 2.\]

__init__(*, n_discrete=5, use_discrete_x=True)

Parameters:

Name Type Description Default
n_discrete int

number of discrete values to sample X from

 5
use_discrete_x bool

if True, X will be an integer. If False, X will be casted to a float.

 True

cite() staticmethod

Returns the BibTeX citation.

bmi.samplers.MultivariateDiscreteUniformMixtureSampler

Bases: IndependentConcatenationSampler

Multivariate alternative for DiscreteUniformMixtureSampler, which is a concatenation of several independent samplers.

Namely, for a sequence of integers \((m_k)\), the variables \(X = (X_1, ..., X_k)\) and \(Y = (Y_1, ..., Y_k)\) are sampled coordinate-wise.

Each coordinate

\[X_k \sim \mathrm{Categorical}(1/m_k, ..., 1/m_k)\]

is from the set \(\{0, ..., m_k-1\}\).

Then,

\[Y_k | X_k \sim \mathrm{Uniform}(X_k, X_k + 2).\]

Mutual information can be calculated as

\[I(X; Y) = \sum_k I(X_k; Y_k),\]

where

\[I(X_k; Y_k) = \log m_k - \frac{m_k-1}{m_k} \log 2.\]

cite() staticmethod

Returns the BibTeX citation.

bmi.samplers.ZeroInflatedPoissonizationSampler

Bases: BaseSampler

Sampler from Gao et al. (2017) modelling zero-inflated Poissonization of the exponential distribution.

Let \(X \sim \mathrm{Exponential}(1)\). Then, \(Y\) is sampled from the mixture distribution

\[Y \mid X = p\, \delta_0 + (1-p) \, \mathrm{Poisson}(X) \]

__init__(p=0.15, use_discrete_y=True)

Parameters:

Name Type Description Default
p float

zero-inflation parameter. Must be in [0, 1).

 0.15

cite() staticmethod

Returns the BibTeX citation.

mutual_information(truncation=None)

Ground-truth mutual information is equal to

\[I(X; Y) = (1-p) \cdot (2 \log 2 - \gamma - S)\]

where

\[S = \sum_{k=1}^{\infty} \log k \cdot 2^{-k},\]

so that the approximation

\[I(X; Y) \approx (1-p) \cdot 0.3012 \]

holds.

Parameters:

Name Type Description Default
truncation Optional[int]

if set to None, the above approximation will be used. Otherwise, the sum will be truncated at the given value.

 None

Bend and Mix Models

See the Bend and Mix Models subpackage API for more information.

Auxiliary

bmi.samplers.BaseSampler

Bases: ISampler

Partial implementation of the ISampler interface, convenient to inherit from.