Bend and Mix Models

Core utilities

bmi.samplers._tfp.JointDistribution dataclass

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

Attributes:

Name	Type	Description
dist		\(P_{XY}\)
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(x, y)

Calculates pointwise mutual information at specified points.

Parameters:

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

Returns:

Type	Description
ndarray	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(n_points, key)

Sample from the joint distribution \(P_{XY}\).

Parameters:

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

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

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

Returns:

Type	Description
float	mutual information estimate
float	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.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
ndarray	PMI profile, shape (n,)

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

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

Parameters:

Name	Type	Description	Default
dist	JointDistribution	distribution to be transformed	required
x_transform	Optional[Bijector]	diffeomorphism \(f\) to transform \(X\). Defaults to identity.	None
y_transform	Optional[Bijector]	diffeomorphism \(g\) to transform \(Y\). Defaults to identity.	None

Returns:

Type	Description
JointDistribution	transformed distribution

bmi.samplers._tfp.ProductDistribution

Bases: JointDistribution

From distributions \(P_X\) and \(P_Y\) creates a distribution \(P_{XY} = P_X \otimes P_Y\), so that \(X\) and \(Y\) are independent.

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

__init__(dist_x, dist_y)

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

Bases: BaseSampler

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

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

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[KeyArray, int]	seed for the Monte Carlo sampling	0
mi_estimate_sample	int	number of samples for the Monte Carlo sampling	200000

Basic distributions

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

Constructs a multivariate normal distribution.

bmi.samplers._tfp.MultivariateNormalDistribution

Bases: JointDistribution

Multivariate normal distribution \(P_{XY}\), 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__(*, dim_x, dim_y, covariance, mean=None)

Parameters:

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

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

Constructs a multivariate Student distribution.

Parameters:

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

bmi.samplers._tfp.MultivariateStudentDistribution

Bases: JointDistribution

Multivariate Student distribution \(P_{XY}\), 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__(*, dim_x, dim_y, df, dispersion, mean=None)

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	Optional[ArrayLike]	mean vector, shape (n,) where n = dim_x + dim_y. Default: zero vector	None
dispersion	ArrayLike	dispersion matrix, shape (n, n)	required