Samplers

Samplers represent probability distributions with known mutual information.

Simple distributions

bmi.samplers._splitmultinormal.SplitMultinormal (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__(self, *, dim_x, dim_y, covariance, mean=None) special

Parameters:

Name	Type	Description	Default
dim_x	int	dimension of the X space	required
dim_y	int	dimension of the Y space	required
mean	Union[numpy.__array_like._SupportsArray[numpy.dtype[Any]], numpy.__nested_sequence._NestedSequence[numpy.__array_like._SupportsArray[numpy.dtype[Any]]], bool, int, float, complex, str, bytes, numpy.__nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]]	mean vector, shape (n,) where n = dim_x + dim_y. Default: zero vector	None
covariance	Union[numpy.__array_like._SupportsArray[numpy.dtype[Any]], numpy.__nested_sequence._NestedSequence[numpy.__array_like._SupportsArray[numpy.dtype[Any]]], bool, int, float, complex, str, bytes, numpy.__nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]]	covariance matrix, shape (n, n)	required

mutual_information(self)

Calculates the mutual information I(X; Y) using an exact formula.

Returns:

Type	Description
float	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.

sample(self, n_points, rng)

Returns a sample from the joint distribution P(X, Y).

Parameters:

Name	Type	Description	Default
n_points	int	sample size	required
rng	Union[Any, int]	pseudorandom number generator	required

Returns:

Type	Description
tuple[numpy.ndarray, numpy.ndarray]	X samples, shape (n_points, dim_x) Y samples, shape (n_points, dim_y). Note that these samples are paired with X samples.

bmi.samplers._split_student_t.SplitStudentT (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 readonly

Degrees of freedom.

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

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	Union[numpy.__array_like._SupportsArray[numpy.dtype[Any]], numpy.__nested_sequence._NestedSequence[numpy.__array_like._SupportsArray[numpy.dtype[Any]]], bool, int, float, complex, str, bytes, numpy.__nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]]	dispersion matrix, shape (dim_x + dim_y, dim_x + dim_y)	required
mean	Union[numpy.__array_like._SupportsArray[numpy.dtype[Any]], numpy.__nested_sequence._NestedSequence[numpy.__array_like._SupportsArray[numpy.dtype[Any]]], bool, int, float, complex, str, bytes, numpy.__nested_sequence._NestedSequence[Union[bool, int, float, complex, str, bytes]]]	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(self)

Calculates the covariance matrix.

Returns:

Type	Description
ndarray	array, shape (dim_x+dim_y, dim_x+dim_y)

Exceptions:

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

mi_correction(self)

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(self)

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

mutual_information(self)

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(self, n_points, rng)

Sampling from multivariate Student distribution.

Note

This function is based on SciPy's sampling.

bmi.samplers._additive_uniform.AdditiveUniformSampler (BaseSampler)

__init__(self, epsilon) special

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.

mutual_information(self)

Mutual information MI(X; Y).

sample(self, n_points, rng)

Returns a sample from the joint distribution P(X, Y).

Parameters:

Name	Type	Description	Default
n_points	int	sample size	required
rng	Union[int, Any]	pseudorandom number generator	required

Returns:

Type	Description
tuple[numpy.ndarray, numpy.ndarray]	X samples, shape (n_points, dim_x) Y samples, shape (n_points, dim_y). Note that these samples are paired with X samples.

bmi.samplers._splitmultinormal.BivariateNormalSampler (BaseSampler)

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

mutual_information(self)

Mutual information MI(X; Y).

sample(self, n_points, rng)

Returns a sample from the joint distribution P(X, Y).

Parameters:

Name	Type	Description	Default
n_points	int	sample size	required
rng	Union[Any, int]	pseudorandom number generator	required

Returns:

Type	Description
tuple[numpy.ndarray, numpy.ndarray]	X samples, shape (n_points, dim_x) Y samples, shape (n_points, dim_y). Note that these samples are paired with X samples.

Combining and transforming samplers

bmi.samplers._independent_coordinates.IndependentConcatenationSampler (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__(self, samplers) special

Parameters:

Name	Type	Description	Default
samplers	Sequence[bmi.interface.ISampler]	sequence of samplers to concatenate	required

mutual_information(self)

Mutual information MI(X; Y).

sample(self, n_points, rng)

Returns a sample from the joint distribution P(X, Y).

Parameters:

Name	Type	Description	Default
n_points	int	sample size	required
rng	Union[int, Any]	pseudorandom number generator	required

Returns:

Type	Description
tuple[jax.Array, jax.Array]	X samples, shape (n_points, dim_x) Y samples, shape (n_points, dim_y). Note that these samples are paired with X samples.

bmi.samplers._transformed.TransformedSampler (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__(self, base_sampler, transform_x=None, transform_y=None, add_dim_x=0, add_dim_y=0, vectorise=True) special

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

mutual_information(self)

Mutual information MI(X; Y).

sample(self, n_points, rng)

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

Returns:

Type	Description
tuple[jax.Array, jax.Array]	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(self, x, y)

Transforms given samples by f x g.

Parameters:

Name	Type	Description	Default
x	Union[jax.Array, numpy.ndarray]	samples, (n_points, dim(X))	required
y	Union[jax.Array, numpy.ndarray]	samples, (n_points, dim(Y))	required

Returns:

Type	Description
tuple[jax.Array, jax.Array]	f(x), shape (n_points, dim(X) + add_dim_x) g(y), shape (n_points, dim(Y) + add_dim_y)

Discrete random variables

bmi.samplers._discrete_continuous.DiscreteUniformMixtureSampler (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__(self, *, n_discrete=5, use_discrete_x=True) special

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.

mutual_information(self)

Mutual information MI(X; Y).

sample(self, n_points, rng)

Returns a sample from the joint distribution P(X, Y).

Parameters:

Name	Type	Description	Default
n_points	int	sample size	required
rng	Union[int, Any]	pseudorandom number generator	required

Returns:

Type	Description
tuple[jax.Array, jax.Array]	X samples, shape (n_points, dim_x) Y samples, shape (n_points, dim_y). Note that these samples are paired with X samples.

bmi.samplers._discrete_continuous.MultivariateDiscreteUniformMixtureSampler (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._discrete_continuous.ZeroInflatedPoissonizationSampler (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__(self, p=0.15, use_discrete_y=True) special

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(self, 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

sample(self, n_points, rng)

Returns a sample from the joint distribution P(X, Y).

Parameters:

Name	Type	Description	Default
n_points	int	sample size	required
rng	Union[int, Any]	pseudorandom number generator	required

Returns:

Type	Description
tuple[jax.Array, jax.Array]	X samples, shape (n_points, dim_x) Y samples, shape (n_points, dim_y). Note that these samples are paired with X samples.

Bend and Mix Models

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

Auxiliary

bmi.samplers.base.BaseSampler (ISampler)

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

dim_x: int property readonly

Dimension of the space in which the X variable is valued.

dim_y: int property readonly

Dimension of the space in which the Y variable is valued.

__init__(self, *, dim_x, dim_y) special

Initialize self. See help(type(self)) for accurate signature.