Bend and Mix Models

Core utilities

`bmi.samplers._tfp._core.JointDistribution` `dataclass`

The main object of this package, representing a Bend and Mix Model (BMM), i.e., a joint distribution $P_{X Y}$ together with the marginal distributions $P_{X}$ and $P_{Y}$ .

Attributes:

Name	Type	Description
`dist`		$P_{X Y}$
`dist_x`	`Distribution`	$P_{X}$
`dist_y`	`Distribution`	$P_{Y}$
`dim_x`	`int`	dimension of the support of $X$
`dim_y`	`int`	dimension of the support of $Y$
`analytic_mi`	`Optional[float]`	analytical mutual information. Use `None` if unknown (in most cases)

`pmi(self, x, y)`

Calculates pointwise mutual information at specified points.

Parameters:

Name	Type	Description	Default
`x`	`Array`	points in the X space, shape `(n_points, dim_x)`	required
`y`	`Array`	points in the Y space, shape `(n_points, dim_y)`	required

Returns:

Type	Description
`Array`	pointwise mutual information evaluated at (x, y) points, shape `(n_points,)`

Note

This function is vectorized, i.e. it can calculate PMI for multiple points at once.

`sample(self, n_points, key)`

Sample from the joint distribution $P_{X Y}$ .

Parameters:

Name	Type	Description	Default
`n_points`	`int`	number of samples to draw	required
`key`	`Array`	JAX random key	required

`bmi.samplers._tfp._core.monte_carlo_mi_estimate(key, dist, n)`

Estimates the mutual information $I (X; Y)$ using Monte Carlo sampling.

Returns:

Type	Description
`tuple[float, float]`	mutual information estimate standard error estimate

Note

It is worth to run this procedure multiple times and see whether the standard error estimate is accurate.

`bmi.samplers._tfp._core.pmi_profile(key, dist, n)`

Monte Carlo draws a sample of size n from the PMI distribution.

Parameters:

Name	Type	Description	Default
`key`	`Array`	JAX random key, used to generate the sample	required
`dist`	`JointDistribution`	distribution	required
`n`	`int`	number of points to sample	required

Returns:

Type	Description
`Array`	PMI profile, shape `(n,)`

`bmi.samplers._tfp._core.transform(dist, x_transform=None, y_transform=None)`

For given diffeomorphisms $f$ and $g$ transforms the joint distribution $P_{X Y}$ into $P_{f (X) g (Y)}$ .

Parameters:

Name	Type	Description	Default
`dist`	`JointDistribution`	distribution to be transformed	required
`x_transform`	`Optional[tensorflow_probability.substrates.jax.bijectors.bijector.Bijector]`	diffeomorphism $f$ to transform $X$ . Defaults to identity.	`None`
`y_transform`	`Optional[tensorflow_probability.substrates.jax.bijectors.bijector.Bijector]`	diffeomorphism $g$ to transform $Y$ . Defaults to identity.	`None`

Returns:

Type	Description
`JointDistribution`	transformed distribution

`bmi.samplers._tfp._product.ProductDistribution (JointDistribution)`

From distributions $P_{X}$ and $P_{Y}$ creates a distribution $P_{X Y} = P_{X} \otimes P_{Y}$ , so that $X$ and $Y$ are independent.

In particular, $I (X; Y) = 0$ .

`init(self, dist_x, dist_y)` `special`

Creates a product distribution.

Parameters:

Name	Type	Description	Default
`dist_x`	`Distribution`	distribution $P_{X}$	required
`dist_y`	`Distribution`	distribution $P_{Y}$	required

`bmi.samplers._tfp._wrapper.BMMSampler (BaseSampler)`

Wraps a given Bend and Mix Model (BMM) into a sampler.

`init(self, dist, mi=None, mi_estimate_seed=0, mi_estimate_sample=200000)` `special`

Parameters:

Name	Type	Description	Default
`dist`	`JointDistribution`	distribution represented by a BMM to be wrapped	required
`mi`	`Optional[float]`	mutual information of the distribution, if already calculated. If not provided, it will be estimated via Monte Carlo sampling.	`None`
`mi_estimate_seed`	`Union[Any, int]`	seed for the Monte Carlo sampling	`0`
`mi_estimate_sample`	`int`	number of samples for the Monte Carlo sampling	`200000`

`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.

Basic distributions

`bmi.samplers._tfp._normal.construct_multivariate_normal_distribution(mean, covariance)`

Constructs a multivariate normal distribution.

`bmi.samplers._tfp._normal.MultivariateNormalDistribution (JointDistribution)`

Multivariate normal distribution $P_{X Y}$ , such that $P_{X}$ is a multivariate normal distribution on the space of dimension dim_x and $P_{Y}$ is a multivariate normal distribution on the space of dimension dim_y.

`init(self, *, dim_x, dim_y, covariance, mean=None)` `special`

Parameters:

Name	Type	Description	Default
`dim_x`	`int`	dimension of the $X$ support	required
`dim_y`	`int`	dimension of the $Y$ support	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

`bmi.samplers._tfp._student.construct_multivariate_student_distribution(mean, dispersion, df)`

Constructs a multivariate Student distribution.

Parameters:

Name	Type	Description	Default
`mean`	`Array`	location vector, shape `(dim,)`	required
`dispersion`	`Array`	dispersion matrix, shape `(dim, dim)`	required
`df`	`Union[int, float]`	degrees of freedom	required

`bmi.samplers._tfp._student.MultivariateStudentDistribution (JointDistribution)`

Multivariate Student distribution $P_{X Y}$ , such that $P_{X}$ is a multivariate Student distribution on the space of dimension dim_x and $P_{Y}$ is a multivariate Student distribution on the space of dimension dim_y.

Note that the degrees of freedom df are the same for all distributions.

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

Parameters:

Name	Type	Description	Default
`dim_x`	`int`	dimension of the $X$ support	required
`dim_y`	`int`	dimension of the $Y$ support	required
`df`	`int`	degrees of freedom	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`
`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 `(n, n)`	required

Bend and Mix Models

Core utilities

bmi.samplers._tfp._core.JointDistribution dataclass

pmi(self, x, y)

sample(self, n_points, key)

bmi.samplers._tfp._core.monte_carlo_mi_estimate(key, dist, n)

bmi.samplers._tfp._core.pmi_profile(key, dist, n)

bmi.samplers._tfp._core.transform(dist, x_transform=None, y_transform=None)

bmi.samplers._tfp._product.ProductDistribution (JointDistribution)

__init__(self, dist_x, dist_y) special

bmi.samplers._tfp._wrapper.BMMSampler (BaseSampler)

__init__(self, dist, mi=None, mi_estimate_seed=0, mi_estimate_sample=200000) special

mutual_information(self)

sample(self, n_points, rng)

Basic distributions

bmi.samplers._tfp._normal.construct_multivariate_normal_distribution(mean, covariance)

bmi.samplers._tfp._normal.MultivariateNormalDistribution (JointDistribution)

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

bmi.samplers._tfp._student.construct_multivariate_student_distribution(mean, dispersion, df)

bmi.samplers._tfp._student.MultivariateStudentDistribution (JointDistribution)

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

`bmi.samplers._tfp._core.JointDistribution` `dataclass`

`pmi(self, x, y)`

`sample(self, n_points, key)`

`bmi.samplers._tfp._core.monte_carlo_mi_estimate(key, dist, n)`

`bmi.samplers._tfp._core.pmi_profile(key, dist, n)`

`bmi.samplers._tfp._core.transform(dist, x_transform=None, y_transform=None)`

`bmi.samplers._tfp._product.ProductDistribution (JointDistribution)`

`init(self, dist_x, dist_y)` `special`

`bmi.samplers._tfp._wrapper.BMMSampler (BaseSampler)`

`init(self, dist, mi=None, mi_estimate_seed=0, mi_estimate_sample=200000)` `special`

`mutual_information(self)`

`sample(self, n_points, rng)`

`bmi.samplers._tfp._normal.construct_multivariate_normal_distribution(mean, covariance)`

`bmi.samplers._tfp._normal.MultivariateNormalDistribution (JointDistribution)`

`init(self, *, dim_x, dim_y, covariance, mean=None)` `special`

`bmi.samplers._tfp._student.construct_multivariate_student_distribution(mean, dispersion, df)`

`bmi.samplers._tfp._student.MultivariateStudentDistribution (JointDistribution)`

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