rJKOtt.DistributionOnGrid module

class rJKOtt.DistributionOnGrid.DenseArrayDistribution(grid, rho)

Bases: DistributionOnGrid

Stores the density on the full grid. Run for dim <= 4, or face memory issues

Parameters:

• grid (Grid) – self-explanatory

• rho (Callable) – function proportional to the density

dim

dimension of the distribution

density(x)

Returns a value, proportional to the pdf of the distribution.

Parameters:

x (np.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of density (up to normalization constant) at these points

Return type:

np.ndarray

classmethod get_double_moon(grid, a)

Get a test bimodal distribution of the form

\[\rho_{ \text{DM}}(x) \propto \exp \left( -2( \|x \|_2 - a)^2 \right) \left( \exp(-2(x_1 -a)^2) + \exp(-2(x_1 +a)^2) \right)\]

Parameters:

• grid (Grid) – grid

• a (numpy.float64) – the shift of the distribution (int the first dimension)

Returns:

the distribution

Return type:

DenseArrayDistribution

get_marginal_on_grid(marginals)

Returns a density of the marginal distribution of (x_i1 … x_ik), discretized on the associated grid

Note

Only 1 or 2 dimensional marginals supported. Order matters:

>>> get_marginal_on_grid(i, j) == get_marginal_on_grid(j, i).T

Parameters:

marginals – indices of marginals to compute

Returns:

array of shape (grid.N[i]) or (grid.N[i], grid.N[j]) — density of the marginal distribution of x_i, discretized on the associated grid

Return type:

np.ndarray

classmethod get_nonconvex(grid, a)

Get a test distribution with nonconvex potential

\[ \begin{align}\begin{aligned}\rho_{\text{NC}} \propto \exp(-V_{\text{NC}})\\V_{ \text{NC}}(x) = \left( \sum \limits_{i = 1}^3 \sqrt{|x_i - a_i|} \right)^2\end{aligned}\end{align} \]

Parameters:

• grid (Grid) – grid

• a (numpy.ndarray) – the shift of the distribution

Returns:

the distribution

Return type:

DenseArrayDistribution

log_density(x)

Returns a value, equal to the log pdf of the distribution (up to an additive constant).

Parameters:

x (numpy.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of log density (up to an additive constant) at these points

Return type:

np.ndarray

class rJKOtt.DistributionOnGrid.DistributionOnGrid(grid)

Bases: object

Convenience class for storing distributions and plotting.

Parameters:

grid (Grid) – self-explanatory

Attributes :

dim : dimension of the distribution

abstract density(x)

Returns a value, proportional to the pdf of the distribution.

Parameters:

x (np.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of density (up to normalization constant) at these points

Return type:

np.ndarray

property dim: int

abstract get_marginal_on_grid(marginals)

Returns a density of the marginal distribution of (x_i1 … x_ik), discretized on the associated grid

Note

Only 1 or 2 dimensional marginals supported. Order matters:

>>> get_marginal_on_grid(i, j) == get_marginal_on_grid(j, i).T

Parameters:

marginals (int | Tuple[int] | Tuple[int, int]) – indices of marginals to compute

Returns:

array of shape (grid.N[i]) or (grid.N[i], grid.N[j]) — density of the marginal distribution of x_i, discretized on the associated grid

Return type:

np.ndarray

abstract log_density(x)

Returns a value, equal to the log pdf of the distribution (up to an additive constant).

Parameters:

x (numpy.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of log density (up to an additive constant) at these points

Return type:

np.ndarray

plot_1d_marginal(n_marginal, *plot_args, axs=None, **plot_kwargs)

Plots marginal distributions of x_i for i = 1..dim

Parameters:

• n_marginal (int) – index of marginal to plot

• fig_and_axs (Tuple[plt.Figure, np.ndarray], optional) – figure and axis to plot

• axs (matplotlib.pyplot.Axes | None)

Returns:

figure and axes with the plot

Return type:

Tuple[plt.Figure, plt.Axes]

plot_2d_marginal(marginals=typing.Tuple[int, int], *contour_args, axs=None, fill=False, **contour_kwargs)

Plots distribution of 2-d marginals.

Parameters:

• marginals – indices of marginals to plot

• axs – plt.Axes to plot; if None, will be initialized

• fill – if True, uses plt.contourf; default plt.contour

plot_matrix_marginals(param_subset=None, sample=None, axs=None, sym=False, w_ticks=False, w_labels=True, plot_args=[], plot_kwargs={}, contour_args=[], contour_kwargs={}, scatter_args=[], scatter_kwargs={})

Plots 1- and 2-dimensional marginals on a grid of dim x dim plots.

By default, the contorus are plotted on the upper-right triangle of the matrix.

Parameters:

• param_subset (List[int] | None) – indices of parameters to plot. if None, will plot the marginals of all the parameters

• sample (numpy.ndarray | None) – an array of shape (N_sample, dim) to be plotted on the lower left submatrix

• axs (numpy.ndarray | None) – (dim, dim) array of plt.Axes to plot on; if None, will be initialized

• sym (bool) – if True, additionally plot a transpose contours on the bottom left matrix.

• w_ticks (bool) – whether to draw ticks on the axes; False generally better for higher dimension

• w_labels (bool) – whether to draw axis labels

• plot_args (List) – will be passed to plt.plot for 1-d marginals

• plot_kwargs (Dict) – will be passed to plt.plot for 1-d marginals

• contour_args (List) – will be passed to plt.contour for 2-d marginals

• contour_kwargs (Dict) – will be passed to plt.contour for 2-d marginals

• scatter_args (List) – will be passed to plt.scatter for 2-d sample marginals

• scatter_kwargs (Dict) – will be passed to plt.scatter for 2-d sample marginals

class rJKOtt.DistributionOnGrid.GaussianMixture(grid, means, covariances, weights=None)

Bases: DistributionOnGrid

Gaussian Mixture distribution.

Parameters:

• grid (Grid) – self-explanatory

• means (np.ndarray or List[np.ndarray]) – means of each Gaussian components; either np.ndarray for one component or List[np.ndarray] of length equal to number of components

• covariances (float or List[float] or List[np.ndarray]) –

  covariances of each Gaussian components; can be

  • A single scalar c. Then all the components have the same covariance matrix c*np.eye(dim)

  • A list of scalars cs. Then i-th component has covariance matrix cs[i]*np.eye(dim)

  • A list of np.ndarray. Then each of them should be an SPD matrix.

• weights (np.ndarray or None) – positive weights for each component.

dim

dimension of the distribution

property covariances: List[numpy.ndarray]

get_marginal_on_grid(marginals)

Returns a density of the marginal distribution of (x_i1 … x_ik), discretized on the associated grid

Note

Only 1 or 2 dimensional marginals supported. Order matters:

>>> get_marginal_on_grid(i, j) == get_marginal_on_grid(j, i).T

Parameters:

marginals – indices of marginals to compute

Returns:

array of shape (grid.N[i]) or (grid.N[i], grid.N[j]) — density of the marginal distribution of x_i, discretized on the associated grid

Return type:

np.ndarray

log_density(x)

Returns a value, equal to the log pdf of the distribution (up to an additive constant).

Parameters:

x (numpy.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of log density (up to an additive constant) at these points

Return type:

np.ndarray

property means: List[numpy.ndarray]

sample(n_samples)

Sample from the Gaussian Mixture distribution

Parameters:

n_samples (int) – number of samples

property weights: numpy.ndarray

class rJKOtt.DistributionOnGrid.Grid(left, right, N_nodes, dim=1)

Bases: object

Represents uniform rectangular grid.

Parameters:

• left (Union[float, List[float]]) – Left extent for the grid in each dimension

• right (Union[float, List[float]]) – Right extent for the grid in each dimension

• N_nodes (Union[int, List[int]]) – number of nodes in each dimension

• dim (int) – dimension of the distribution

Attributes :

dim : dimension of the grid

clip_sample(xs)

Given an array of points, delete those of them that are out of the grid’s span

Parameters:

xs (numpy.ndarray) – the points; should have shape (N_points, dim)

Returns:

fitting points

Return type:

np.ndarray

get_1d_grid(i)

Return grid points in i-th dimension

Parameters:

i (int) – index of the dimension; 0 <= i < dim

Returns:

array of shape (grid.N[i]) — grid points in i-th direction

Return type:

numpy.ndarray

get_2d_grid(i, j)

Return grid points in i-th dimension

Parameters:

• i (int) – index of the first dimension; 0 <= i < dim

• j (int) – index of the second dimension; 0 <= j < dim

Returns:

arrays of shape (N_nodes[i], N_nodes[j]) (suitable for scatter plot)

Return type:

Tuple[np.ndarray]

class rJKOtt.DistributionOnGrid.TensorTrainDistribution(grid, rho_tt)

Bases: DistributionOnGrid

Distribution with density on grid stored in the compressed TT format

Parameters:

• grid (Grid) – self-explanatory

• rho_tt (tt_vector) – Tensor Train decomposition of the density on grid

dim

dimension of the distribution

density(x)

Note

Returns a nearest interpolation; Will implement linear interpolation in the future…

Parameters:

x (np.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of density (up to normalization constant) at these points

Return type:

np.ndarray

classmethod gaussian(grid, ms=0.0, sigmas=1.0)

A TT approximation of the density of the distribution with each parameter being independent and distributed as

\[x_i \sim \mathcal{N}(m_i, \sigma_i),\ i = \overline{1,\ d}\]

Parameters:

• grid (Grid) – the grid to discretize on

• ms (float | List[float] | numpy.ndarray) – means \(m_i\) of each parameter

• sigmas (float | List[float] | numpy.ndarray) – standard deviations \(sigma_i\) of each parameter

Returns:

the distribution

Return type:

TensorTrainDistribution

get_credible_interval(i, prob)

Parameters:

• i (int)

• prob (float)

get_marginal_on_grid(marginals)

Returns a density of the marginal distribution of (x_i1 … x_ik), discretized on the associated grid

Note

Only 1 or 2 dimensional marginals supported. Order matters:

>>> get_marginal_on_grid(i, j) == get_marginal_on_grid(j, i).T

Parameters:

marginals – indices of marginals to compute

Returns:

array of shape (grid.N[i]) or (grid.N[i], grid.N[j]) — density of the marginal distribution of x_i, discretized on the associated grid

Return type:

np.ndarray

log_density(x)

Returns a value, equal to the log pdf of the distribution (up to an additive constant).

Parameters:

x (numpy.ndarray) – array of shape (N_x, dim) — coordinates of N_x points

Returns:

array of shape (N_x) — value of log density (up to an additive constant) at these points

Return type:

np.ndarray

classmethod rank1_fx(grid, fns)

Convenience function to create a rank-1 TT with components

\[A^i_{1k1} = f^i(x_{i,k}),\ i = \overline{1,\ d}\]

where \(x_{i,k},\ k = \overline{1, N_i}\) is the unidimensional grid in i-th direction

Parameters:

• grid (Grid) – the grid to discretize on

• fns (List[Callable] | Callable) – functions \(f^i\)

Returns:

the distribution

Return type:

TensorTrainDistribution