ot.gaussian

Optimal transport for Gaussian distributions

Functions

ot.gaussian.bures_barycenter_fixpoint(C, weights=None, num_iter=1000, eps=1e-07, log=False, nx=None)

Return the (Bures-)Wasserstein barycenter between centered Gaussian distributions.

The function estimates the (Bures)-Wasserstein barycenter between centered Gaussian distributions \(\big(\mathcal{N}(0,\Sigma_i)\big)_{i=1}^n\) [16] by solving

\[\Sigma_b = \mathrm{argmin}_{\Sigma \in S_d^{++}(\mathbb{R})}\ \sum_{i=1}^n w_i W_2^2\big(\mathcal{N}(0,\Sigma), \mathcal{N}(0, \Sigma_i)\big).\]

The barycenter still follows a Gaussian distribution \(\mathcal{N}(0,\Sigma_b)\) where \(\Sigma_b\) is solution of the following fixed-point algorithm:

\[\Sigma_b = \sum_{i=1}^n w_i \left(\Sigma_b^{1/2}\Sigma_i^{1/2}\Sigma_b^{1/2}\right)^{1/2}.\]

Parameters:

• C (array-like (k,d,d)) – covariance of k distributions

• weights (array-like (k), optional) – weights for each distribution

• method (str) – method used for the solver, either ‘fixed_point’ or ‘gradient_descent’

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• log (bool, optional) – record log if True

• nx (module, optional) – The numerical backend module to use. If not provided, the backend will be fetched from the input matrices C.

Returns:

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[16]

M. Agueh and G. Carlier, “Barycenters in the Wasserstein space”, SIAM Journal on Mathematical Analysis, vol. 43, no. 2, pp. 904-924, 2011.

ot.gaussian.bures_barycenter_gradient_descent(C, weights=None, num_iter=1000, eps=1e-07, log=False, step_size=1, batch_size=None, averaged=False, nx=None)

Return the (Bures-)Wasserstein barycenter between centered Gaussian distributions.

The function estimates the (Bures)-Wasserstein barycenter between centered Gaussian distributions \(\big(\mathcal{N}(0,\Sigma_i)\big)_{i=1}^n\) by using a gradient descent in the Wasserstein space [74, 75] on the objective

\[\mathcal{L}(\Sigma) = \sum_{i=1}^n w_i W_2^2\big(\mathcal{N}(0,\Sigma), \mathcal{N}(0,\Sigma_i)\big).\]

Parameters:

• C (array-like (k,d,d)) – covariance of k distributions

• weights (array-like (k), optional) – weights for each distribution

• method (str) – method used for the solver, either ‘fixed_point’ or ‘gradient_descent’

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• log (bool, optional) – record log if True

• step_size (float, optional) – step size for the gradient descent, 1 by default

• batch_size (int, optional) – batch size if use a stochastic gradient descent

• averaged (bool, optional) – if True, use the averaged procedure of [74]

• nx (module, optional) – The numerical backend module to use. If not provided, the backend will be fetched from the input matrices C.

Returns:

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[74]

Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). Gradient descent algorithms for Bures-Wasserstein barycenters. In Conference on Learning Theory (pp. 1276-1304). PMLR.

[75]

Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent. Advances in Neural Information Processing Systems, 34, 22132-22145.

ot.gaussian.bures_distance(Cs, Ct, paired=False, log=False, nx=None)

Return Bures distance.

The function computes the Bures distance between \(\mu_s=\mathcal{N}(0,\Sigma_s)\) and \(\mu_t=\mathcal{N}(0,\Sigma_t)\), given by (see e.g. Remark 2.31 [15]):

\[\mathbf{B}(\Sigma_s, \Sigma_t)^{2} = \text{Tr}\left(\Sigma_s + \Sigma_t - 2 \sqrt{\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2}} \right)\]

Parameters:

• Cs (array-like (d,d) or (n,d,d)) – covariance of the source distribution

• Ct (array-like (d,d) or (m,d,d)) – covariance of the target distribution

• paired (bool, optional) – if True and n==m, return the paired distances and crossed distance otherwise

• log (bool, optional) – record log if True

• nx (module, optional) – The numerical backend module to use. If not provided, the backend will be fetched from the input matrices Cs, Ct.

Returns:

• W (float if Cs and Cd of shape (d,d), array-like (n,m) if Cs of shape (n,d,d) and Ct of shape (m,d,d), array-like (n,) if Cs and Ct of shape (n, d, d) and paired is True) – Bures Wasserstein distance

• log (dict) – log dictionary return only if log==True in parameters

References

[15]

Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6), 355-607.

ot.gaussian.bures_wasserstein_barycenter(m, C, weights=None, method='fixed_point', num_iter=1000, eps=1e-07, log=False, step_size=1, batch_size=None)

Return the (Bures-)Wasserstein barycenter between Gaussian distributions.

The function estimates the (Bures)-Wasserstein barycenter between Gaussian distributions \(\big(\mathcal{N}(\mu_i,\Sigma_i)\big)_{i=1}^n\) [16, 74, 75] by solving

\[(\mu_b, \Sigma_b) = \mathrm{argmin}_{\mu,\Sigma}\ \sum_{i=1}^n w_i W_2^2\big(\mathcal{N}(\mu,\Sigma), \mathcal{N}(\mu_i, \Sigma_i)\big)\]

The barycenter still follows a Gaussian distribution \(\mathcal{N}(\mu_b,\Sigma_b)\) where:

\[\mu_b = \sum_{i=1}^n w_i \mu_i,\]

and the barycentric covariance is the solution of the following fixed-point algorithm:

\[\Sigma_b = \sum_{i=1}^n w_i \left(\Sigma_b^{1/2}\Sigma_i^{1/2}\Sigma_b^{1/2}\right)^{1/2}\]

We propose two solvers: one based on solving the previous fixed-point problem [16]. Another based on gradient descent in the Bures-Wasserstein space [74,75].

Parameters:

• m (array-like (k,d)) – mean of k distributions

• C (array-like (k,d,d)) – covariance of k distributions

• weights (array-like (k), optional) – weights for each distribution

• method (str) – method used for the solver, either ‘fixed_point’, ‘gradient_descent’, ‘stochastic_gradient_descent’ or ‘averaged_stochastic_gradient_descent’

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• log (bool, optional) – record log if True

• step_size (float, optional) – step size for the gradient descent, 1 by default

• batch_size (int, optional) – batch size if use a stochastic gradient descent. If not None, use method=’gradient_descent’

Returns:

• mb ((d,) array-like) – mean of the barycenter

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[16]

M. Agueh and G. Carlier, “Barycenters in the Wasserstein space”, SIAM Journal on Mathematical Analysis, vol. 43, no. 2, pp. 904-924, 2011.

[74]

Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). Gradient descent algorithms for Bures-Wasserstein barycenters. In Conference on Learning Theory (pp. 1276-1304). PMLR.

[75]

Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent. Advances in Neural Information Processing Systems, 34, 22132-22145.

Examples using ot.gaussian.bures_wasserstein_barycenter

Gaussian Bures-Wasserstein barycenters

Gaussian Bures-Wasserstein barycenters

ot.gaussian.bures_wasserstein_distance(ms, mt, Cs, Ct, paired=False, log=False)

Return Bures Wasserstein distance between samples.

The function computes the Bures-Wasserstein distance between \(\mu_s=\mathcal{N}(m_s,\Sigma_s)\) and \(\mu_t=\mathcal{N}(m_t,\Sigma_t)\), as discussed in remark 2.31 [15].

\[\mathcal{W}(\mu_s, \mu_t)_2^2= \left\lVert \mathbf{m}_s - \mathbf{m}_t \right\rVert^2 + \mathcal{B}(\Sigma_s, \Sigma_t)^{2}\]

where :

\[\mathbf{B}(\Sigma_s, \Sigma_t)^{2} = \text{Tr}\left(\Sigma_s + \Sigma_t - 2 \sqrt{\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2}} \right)\]

Parameters:

• ms (array-like (d,) or (n,d)) – mean of the source distribution

• mt (array-like (d,) or (m,d)) – mean of the target distribution

• Cs (array-like (d,d) or (n,d,d)) – covariance of the source distribution

• Ct (array-like (d,d) or (m,d,d)) – covariance of the target distribution

• paired (bool, optional) – if True and n==m, return the paired distances and crossed distance otherwise

• log (bool, optional) – record log if True

Returns:

• W (float if ms and md of shape (d,), array-like (n,m) if ms of shape (n,d) and mt of shape (m,d), array-like (n,) if ms and mt of shape (n,d) and paired is True) – Bures Wasserstein distance

• log (dict) – log dictionary return only if log==True in parameters

References

[15]

Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6), 355-607.

ot.gaussian.bures_wasserstein_mapping(ms, mt, Cs, Ct, log=False)

Return OT linear operator between samples.

The function estimates the optimal linear operator that aligns the two empirical distributions. This is equivalent to estimating the closed form mapping between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) as proposed in [1] and discussed in remark 2.29 in [2].

The linear operator from source to target \(M\)

\[M(\mathbf{x})= \mathbf{A} \mathbf{x} + \mathbf{b}\]

where :

\[ \begin{align}\begin{aligned}\mathbf{A} &= \Sigma_s^{-1/2} \left(\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2} \right)^{1/2} \Sigma_s^{-1/2}\\\mathbf{b} &= \mu_t - \mathbf{A} \mu_s\end{aligned}\end{align} \]

Parameters:

• ms (array-like (d,)) – mean of the source distribution

• mt (array-like (d,)) – mean of the target distribution

• Cs (array-like (d,d)) – covariance of the source distribution

• Ct (array-like (d,d)) – covariance of the target distribution

• log (bool, optional) – record log if True

Returns:

• A ((d, d) array-like) – Linear operator

• b ((1, d) array-like) – bias

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

Knott, M. and Smith, C. S. “On the optimal mapping of distributions”, Journal of Optimization Theory and Applications Vol 43, 1984

[2]

Peyré, G., & Cuturi, M. (2017). “Computational Optimal Transport”, 2018.

ot.gaussian.empirical_bures_wasserstein_barycenter(X, reg=1e-06, weights=None, num_iter=1000, eps=1e-07, w=None, bias=True, log=False)

Return OT linear operator between samples.

The function estimates the optimal barycenter of the empirical distributions. This is equivalent to resolving the fixed point algorithm for multiple Gaussian distributions \(\left\{\mathcal{N}(\mu,\Sigma)\right\}_{i=1}^n\) [1].

The barycenter still following a Gaussian distribution \(\mathcal{N}(\mu_b,\Sigma_b)\) where :

\[\mu_b = \sum_{i=1}^n w_i \mu_i\]

And the barycentric covariance is the solution of the following fixed-point algorithm:

\[\Sigma_b = \sum_{i=1}^n w_i \left(\Sigma_b^{1/2}\Sigma_i^{1/2}\Sigma_b^{1/2}\right)^{1/2}\]

Parameters:

• X (list of array-like (n,d)) – samples in each distribution

• reg (float,optional) – regularization added to the diagonals of covariances (>0)

• weights (array-like (n,), optional) – weights for each distribution

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• w (list of array-like (n,), optional) – weights for each sample in each distribution

• bias (boolean, optional) – estimate bias \(\mathbf{b}\) else \(\mathbf{b} = 0\) (default:True)

• log (bool, optional) – record log if True

Returns:

• mb ((d,) array-like) – mean of the barycenter

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

M. Agueh and G. Carlier, “Barycenters in the Wasserstein space”, SIAM Journal on Mathematical Analysis, vol. 43, no. 2, pp. 904-924, 2011.

ot.gaussian.empirical_bures_wasserstein_distance(xs, xt, reg=1e-06, ws=None, wt=None, bias=True, log=False)

Return Bures Wasserstein distance from mean and covariance of distribution.

The function estimates the Bures-Wasserstein distance between two empirical distributions source \(\mu_s\) and target \(\mu_t\), discussed in remark 2.31 [1].

The Bures Wasserstein distance between source and target distribution \(\mathcal{W}\)

\[\mathcal{W}(\mu_s, \mu_t)_2^2= \left\lVert \mathbf{m}_s - \mathbf{m}_t \right\rVert^2 + \mathcal{B}(\Sigma_s, \Sigma_t)^{2}\]

where :

\[\mathbf{B}(\Sigma_s, \Sigma_t)^{2} = \text{Tr}\left(\Sigma_s + \Sigma_t - 2 \sqrt{\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2}} \right)\]

Parameters:

• xs (array-like (ns,d)) – samples in the source domain

• xt (array-like (nt,d)) – samples in the target domain

• reg (float,optional) – regularization added to the diagonals of covariances (>0)

• ws (array-like (ns), optional) – weights for the source samples

• wt (array-like (ns), optional) – weights for the target samples

• bias (boolean, optional) – estimate bias \(\mathbf{b}\) else \(\mathbf{b} = 0\) (default:True)

• log (bool, optional) – record log if True

Returns:

• W (float) – Bures Wasserstein distance

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

Peyré, G., & Cuturi, M. (2017). “Computational Optimal Transport”, 2018.

ot.gaussian.empirical_bures_wasserstein_mapping(xs, xt, reg=1e-06, ws=None, wt=None, bias=True, log=False)

Return OT linear operator between samples.

The function estimates the optimal linear operator that aligns the two empirical distributions. This is equivalent to estimating the closed form mapping between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) as proposed in [1] and discussed in remark 2.29 in [2].

The linear operator from source to target \(M\)

\[M(\mathbf{x})= \mathbf{A} \mathbf{x} + \mathbf{b}\]

where :

\[ \begin{align}\begin{aligned}\mathbf{A} &= \Sigma_s^{-1/2} \left(\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2} \right)^{1/2} \Sigma_s^{-1/2}\\\mathbf{b} &= \mu_t - \mathbf{A} \mu_s\end{aligned}\end{align} \]

Parameters:

• xs (array-like (ns,d)) – samples in the source domain

• xt (array-like (nt,d)) – samples in the target domain

• reg (float,optional) – regularization added to the diagonals of covariances (>0)

• ws (array-like (ns,1), optional) – weights for the source samples

• wt (array-like (ns,1), optional) – weights for the target samples

• bias (boolean, optional) – estimate bias \(\mathbf{b}\) else \(\mathbf{b} = 0\) (default:True)

• log (bool, optional) – record log if True

Returns:

• A ((d, d) array-like) – Linear operator

• b ((1, d) array-like) – bias

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

Knott, M. and Smith, C. S. “On the optimal mapping of distributions”, Journal of Optimization Theory and Applications Vol 43, 1984

[2]

Peyré, G., & Cuturi, M. (2017). “Computational Optimal Transport”, 2018.

Examples using ot.gaussian.empirical_bures_wasserstein_mapping

Linear OT mapping estimation

Linear OT mapping estimation

ot.gaussian.empirical_gaussian_gromov_wasserstein_distance(xs, xt, ws=None, wt=None, log=False)

Return Gaussian Gromov-Wasserstein distance between samples.

The function estimates the Gaussian Gromov-Wasserstein distance between two Gaussien distributions source \(\mu_s\) and target \(\mu_t\), whose parameters are estimated from the provided samples \(\mathcal{X}_s\) and \(\mathcal{X}_t\). See [57] Theorem 4.1 for more details.

Parameters:

• xs (array-like (ns,d)) – samples in the source domain

• xt (array-like (nt,d)) – samples in the target domain

• ws (array-like (ns,1), optional) – weights for the source samples

• wt (array-like (ns,1), optional) – weights for the target samples

• log (bool, optional) – record log if True

Returns:

G – Gaussian Gromov-Wasserstein distance

Return type:

float

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

ot.gaussian.empirical_gaussian_gromov_wasserstein_mapping(xs, xt, ws=None, wt=None, sign_eigs=None, log=False)

Return Gaussian Gromov-Wasserstein mapping between samples.

The function estimates the Gaussian Gromov-Wasserstein mapping between two Gaussian distributions source \(\mu_s\) and target \(\mu_t\), whose parameters are estimated from the provided samples \(\mathcal{X}_s\) and \(\mathcal{X}_t\). See [57] Theorem 4.1 for more details.

Parameters:

• xs (array-like (ns,ds)) – samples in the source domain

• xt (array-like (nt,dt)) – samples in the target domain

• ws (array-like (ns,1), optional) – weights for the source samples

• wt (array-like (ns,1), optional) – weights for the target samples

• sign_eigs (array-like (min(ds,dt),) or string, optional) – sign of the eigenvalues of the mapping matrix, by default all signs will be positive. If ‘skewness’ is provided, the sign of the eigenvalues is selected as the product of the sign of the skewness of the projected data.

• log (bool, optional) – record log if True

Returns:

• A ((dt, ds) array-like) – Linear operator

• b ((1, dt) array-like) – bias

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

Examples using ot.gaussian.empirical_gaussian_gromov_wasserstein_mapping

Linear OT mapping estimation

Linear OT mapping estimation

ot.gaussian.gaussian_gromov_wasserstein_distance(Cov_s, Cov_t, log=False)

Return the Gaussian Gromov-Wasserstein value from [57].

This function return the closed form value of the Gaussian Gromov-Wasserstein distance between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) when the OT plan is assumed to be also Gaussian. See [57] Theorem 4.1 for more details.

Parameters:

• Cov_s (array-like (ds,ds)) – covariance of the source distribution

• Cov_t (array-like (dt,dt)) – covariance of the target distribution

Returns:

G – Gaussian Gromov-Wasserstein distance

Return type:

float

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

ot.gaussian.gaussian_gromov_wasserstein_mapping(mu_s, mu_t, Cov_s, Cov_t, sign_eigs=None, log=False)

Return the Gaussian Gromov-Wasserstein mapping from [57].

This function return the closed form value of the Gaussian Gromov-Wasserstein mapping between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) when the OT plan is assumed to be also Gaussian. See [57] Theorem 4.1 for more details.

Parameters:

• mu_s (array-like (ds,)) – mean of the source distribution

• mu_t (array-like (dt,)) – mean of the target distribution

• Cov_s (array-like (ds,ds)) – covariance of the source distribution

• Cov_t (array-like (dt,dt)) – covariance of the target distribution

• log (bool, optional) – record log if True

Returns:

• A ((dt, ds) array-like) – Linear operator

• b ((1, dt) array-like) – bias

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

ot.gaussian.bures_barycenter_fixpoint(C, weights=None, num_iter=1000, eps=1e-07, log=False, nx=None)

Return the (Bures-)Wasserstein barycenter between centered Gaussian distributions.

The function estimates the (Bures)-Wasserstein barycenter between centered Gaussian distributions \(\big(\mathcal{N}(0,\Sigma_i)\big)_{i=1}^n\) [16] by solving

\[\Sigma_b = \mathrm{argmin}_{\Sigma \in S_d^{++}(\mathbb{R})}\ \sum_{i=1}^n w_i W_2^2\big(\mathcal{N}(0,\Sigma), \mathcal{N}(0, \Sigma_i)\big).\]

The barycenter still follows a Gaussian distribution \(\mathcal{N}(0,\Sigma_b)\) where \(\Sigma_b\) is solution of the following fixed-point algorithm:

\[\Sigma_b = \sum_{i=1}^n w_i \left(\Sigma_b^{1/2}\Sigma_i^{1/2}\Sigma_b^{1/2}\right)^{1/2}.\]

Parameters:

• C (array-like (k,d,d)) – covariance of k distributions

• weights (array-like (k), optional) – weights for each distribution

• method (str) – method used for the solver, either ‘fixed_point’ or ‘gradient_descent’

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• log (bool, optional) – record log if True

• nx (module, optional) – The numerical backend module to use. If not provided, the backend will be fetched from the input matrices C.

Returns:

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[16]

M. Agueh and G. Carlier, “Barycenters in the Wasserstein space”, SIAM Journal on Mathematical Analysis, vol. 43, no. 2, pp. 904-924, 2011.

ot.gaussian.bures_barycenter_gradient_descent(C, weights=None, num_iter=1000, eps=1e-07, log=False, step_size=1, batch_size=None, averaged=False, nx=None)

Return the (Bures-)Wasserstein barycenter between centered Gaussian distributions.

The function estimates the (Bures)-Wasserstein barycenter between centered Gaussian distributions \(\big(\mathcal{N}(0,\Sigma_i)\big)_{i=1}^n\) by using a gradient descent in the Wasserstein space [74, 75] on the objective

\[\mathcal{L}(\Sigma) = \sum_{i=1}^n w_i W_2^2\big(\mathcal{N}(0,\Sigma), \mathcal{N}(0,\Sigma_i)\big).\]

Parameters:

• C (array-like (k,d,d)) – covariance of k distributions

• weights (array-like (k), optional) – weights for each distribution

• method (str) – method used for the solver, either ‘fixed_point’ or ‘gradient_descent’

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• log (bool, optional) – record log if True

• step_size (float, optional) – step size for the gradient descent, 1 by default

• batch_size (int, optional) – batch size if use a stochastic gradient descent

• averaged (bool, optional) – if True, use the averaged procedure of [74]

• nx (module, optional) – The numerical backend module to use. If not provided, the backend will be fetched from the input matrices C.

Returns:

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[74]

Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). Gradient descent algorithms for Bures-Wasserstein barycenters. In Conference on Learning Theory (pp. 1276-1304). PMLR.

[75]

Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent. Advances in Neural Information Processing Systems, 34, 22132-22145.

ot.gaussian.bures_distance(Cs, Ct, paired=False, log=False, nx=None)

Return Bures distance.

The function computes the Bures distance between \(\mu_s=\mathcal{N}(0,\Sigma_s)\) and \(\mu_t=\mathcal{N}(0,\Sigma_t)\), given by (see e.g. Remark 2.31 [15]):

\[\mathbf{B}(\Sigma_s, \Sigma_t)^{2} = \text{Tr}\left(\Sigma_s + \Sigma_t - 2 \sqrt{\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2}} \right)\]

Parameters:

• Cs (array-like (d,d) or (n,d,d)) – covariance of the source distribution

• Ct (array-like (d,d) or (m,d,d)) – covariance of the target distribution

• paired (bool, optional) – if True and n==m, return the paired distances and crossed distance otherwise

• log (bool, optional) – record log if True

• nx (module, optional) – The numerical backend module to use. If not provided, the backend will be fetched from the input matrices Cs, Ct.

Returns:

• W (float if Cs and Cd of shape (d,d), array-like (n,m) if Cs of shape (n,d,d) and Ct of shape (m,d,d), array-like (n,) if Cs and Ct of shape (n, d, d) and paired is True) – Bures Wasserstein distance

• log (dict) – log dictionary return only if log==True in parameters

References

[15]

Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6), 355-607.

ot.gaussian.bures_wasserstein_barycenter(m, C, weights=None, method='fixed_point', num_iter=1000, eps=1e-07, log=False, step_size=1, batch_size=None)

Return the (Bures-)Wasserstein barycenter between Gaussian distributions.

The function estimates the (Bures)-Wasserstein barycenter between Gaussian distributions \(\big(\mathcal{N}(\mu_i,\Sigma_i)\big)_{i=1}^n\) [16, 74, 75] by solving

\[(\mu_b, \Sigma_b) = \mathrm{argmin}_{\mu,\Sigma}\ \sum_{i=1}^n w_i W_2^2\big(\mathcal{N}(\mu,\Sigma), \mathcal{N}(\mu_i, \Sigma_i)\big)\]

The barycenter still follows a Gaussian distribution \(\mathcal{N}(\mu_b,\Sigma_b)\) where:

\[\mu_b = \sum_{i=1}^n w_i \mu_i,\]

and the barycentric covariance is the solution of the following fixed-point algorithm:

\[\Sigma_b = \sum_{i=1}^n w_i \left(\Sigma_b^{1/2}\Sigma_i^{1/2}\Sigma_b^{1/2}\right)^{1/2}\]

We propose two solvers: one based on solving the previous fixed-point problem [16]. Another based on gradient descent in the Bures-Wasserstein space [74,75].

Parameters:

• m (array-like (k,d)) – mean of k distributions

• C (array-like (k,d,d)) – covariance of k distributions

• weights (array-like (k), optional) – weights for each distribution

• method (str) – method used for the solver, either ‘fixed_point’, ‘gradient_descent’, ‘stochastic_gradient_descent’ or ‘averaged_stochastic_gradient_descent’

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• log (bool, optional) – record log if True

• step_size (float, optional) – step size for the gradient descent, 1 by default

• batch_size (int, optional) – batch size if use a stochastic gradient descent. If not None, use method=’gradient_descent’

Returns:

• mb ((d,) array-like) – mean of the barycenter

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[16]

M. Agueh and G. Carlier, “Barycenters in the Wasserstein space”, SIAM Journal on Mathematical Analysis, vol. 43, no. 2, pp. 904-924, 2011.

[74]

Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). Gradient descent algorithms for Bures-Wasserstein barycenters. In Conference on Learning Theory (pp. 1276-1304). PMLR.

[75]

Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent. Advances in Neural Information Processing Systems, 34, 22132-22145.

ot.gaussian.bures_wasserstein_distance(ms, mt, Cs, Ct, paired=False, log=False)

Return Bures Wasserstein distance between samples.

The function computes the Bures-Wasserstein distance between \(\mu_s=\mathcal{N}(m_s,\Sigma_s)\) and \(\mu_t=\mathcal{N}(m_t,\Sigma_t)\), as discussed in remark 2.31 [15].

\[\mathcal{W}(\mu_s, \mu_t)_2^2= \left\lVert \mathbf{m}_s - \mathbf{m}_t \right\rVert^2 + \mathcal{B}(\Sigma_s, \Sigma_t)^{2}\]

where :

\[\mathbf{B}(\Sigma_s, \Sigma_t)^{2} = \text{Tr}\left(\Sigma_s + \Sigma_t - 2 \sqrt{\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2}} \right)\]

Parameters:

• ms (array-like (d,) or (n,d)) – mean of the source distribution

• mt (array-like (d,) or (m,d)) – mean of the target distribution

• Cs (array-like (d,d) or (n,d,d)) – covariance of the source distribution

• Ct (array-like (d,d) or (m,d,d)) – covariance of the target distribution

• paired (bool, optional) – if True and n==m, return the paired distances and crossed distance otherwise

• log (bool, optional) – record log if True

Returns:

• W (float if ms and md of shape (d,), array-like (n,m) if ms of shape (n,d) and mt of shape (m,d), array-like (n,) if ms and mt of shape (n,d) and paired is True) – Bures Wasserstein distance

• log (dict) – log dictionary return only if log==True in parameters

References

[15]

Peyré, G., & Cuturi, M. (2019). Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6), 355-607.

ot.gaussian.bures_wasserstein_mapping(ms, mt, Cs, Ct, log=False)

Return OT linear operator between samples.

The function estimates the optimal linear operator that aligns the two empirical distributions. This is equivalent to estimating the closed form mapping between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) as proposed in [1] and discussed in remark 2.29 in [2].

The linear operator from source to target \(M\)

\[M(\mathbf{x})= \mathbf{A} \mathbf{x} + \mathbf{b}\]

where :

\[ \begin{align}\begin{aligned}\mathbf{A} &= \Sigma_s^{-1/2} \left(\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2} \right)^{1/2} \Sigma_s^{-1/2}\\\mathbf{b} &= \mu_t - \mathbf{A} \mu_s\end{aligned}\end{align} \]

Parameters:

• ms (array-like (d,)) – mean of the source distribution

• mt (array-like (d,)) – mean of the target distribution

• Cs (array-like (d,d)) – covariance of the source distribution

• Ct (array-like (d,d)) – covariance of the target distribution

• log (bool, optional) – record log if True

Returns:

• A ((d, d) array-like) – Linear operator

• b ((1, d) array-like) – bias

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

Knott, M. and Smith, C. S. “On the optimal mapping of distributions”, Journal of Optimization Theory and Applications Vol 43, 1984

[2]

Peyré, G., & Cuturi, M. (2017). “Computational Optimal Transport”, 2018.

ot.gaussian.empirical_bures_wasserstein_barycenter(X, reg=1e-06, weights=None, num_iter=1000, eps=1e-07, w=None, bias=True, log=False)

Return OT linear operator between samples.

The function estimates the optimal barycenter of the empirical distributions. This is equivalent to resolving the fixed point algorithm for multiple Gaussian distributions \(\left\{\mathcal{N}(\mu,\Sigma)\right\}_{i=1}^n\) [1].

The barycenter still following a Gaussian distribution \(\mathcal{N}(\mu_b,\Sigma_b)\) where :

\[\mu_b = \sum_{i=1}^n w_i \mu_i\]

And the barycentric covariance is the solution of the following fixed-point algorithm:

\[\Sigma_b = \sum_{i=1}^n w_i \left(\Sigma_b^{1/2}\Sigma_i^{1/2}\Sigma_b^{1/2}\right)^{1/2}\]

Parameters:

• X (list of array-like (n,d)) – samples in each distribution

• reg (float,optional) – regularization added to the diagonals of covariances (>0)

• weights (array-like (n,), optional) – weights for each distribution

• num_iter (int, optional) – number of iteration for the fixed point algorithm

• eps (float, optional) – tolerance for the fixed point algorithm

• w (list of array-like (n,), optional) – weights for each sample in each distribution

• bias (boolean, optional) – estimate bias \(\mathbf{b}\) else \(\mathbf{b} = 0\) (default:True)

• log (bool, optional) – record log if True

Returns:

• mb ((d,) array-like) – mean of the barycenter

• Cb ((d, d) array-like) – covariance of the barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

M. Agueh and G. Carlier, “Barycenters in the Wasserstein space”, SIAM Journal on Mathematical Analysis, vol. 43, no. 2, pp. 904-924, 2011.

ot.gaussian.empirical_bures_wasserstein_distance(xs, xt, reg=1e-06, ws=None, wt=None, bias=True, log=False)

Return Bures Wasserstein distance from mean and covariance of distribution.

The function estimates the Bures-Wasserstein distance between two empirical distributions source \(\mu_s\) and target \(\mu_t\), discussed in remark 2.31 [1].

The Bures Wasserstein distance between source and target distribution \(\mathcal{W}\)

\[\mathcal{W}(\mu_s, \mu_t)_2^2= \left\lVert \mathbf{m}_s - \mathbf{m}_t \right\rVert^2 + \mathcal{B}(\Sigma_s, \Sigma_t)^{2}\]

where :

\[\mathbf{B}(\Sigma_s, \Sigma_t)^{2} = \text{Tr}\left(\Sigma_s + \Sigma_t - 2 \sqrt{\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2}} \right)\]

Parameters:

• xs (array-like (ns,d)) – samples in the source domain

• xt (array-like (nt,d)) – samples in the target domain

• reg (float,optional) – regularization added to the diagonals of covariances (>0)

• ws (array-like (ns), optional) – weights for the source samples

• wt (array-like (ns), optional) – weights for the target samples

• bias (boolean, optional) – estimate bias \(\mathbf{b}\) else \(\mathbf{b} = 0\) (default:True)

• log (bool, optional) – record log if True

Returns:

• W (float) – Bures Wasserstein distance

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

Peyré, G., & Cuturi, M. (2017). “Computational Optimal Transport”, 2018.

ot.gaussian.empirical_bures_wasserstein_mapping(xs, xt, reg=1e-06, ws=None, wt=None, bias=True, log=False)

Return OT linear operator between samples.

The function estimates the optimal linear operator that aligns the two empirical distributions. This is equivalent to estimating the closed form mapping between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) as proposed in [1] and discussed in remark 2.29 in [2].

The linear operator from source to target \(M\)

\[M(\mathbf{x})= \mathbf{A} \mathbf{x} + \mathbf{b}\]

where :

\[ \begin{align}\begin{aligned}\mathbf{A} &= \Sigma_s^{-1/2} \left(\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2} \right)^{1/2} \Sigma_s^{-1/2}\\\mathbf{b} &= \mu_t - \mathbf{A} \mu_s\end{aligned}\end{align} \]

Parameters:

• xs (array-like (ns,d)) – samples in the source domain

• xt (array-like (nt,d)) – samples in the target domain

• reg (float,optional) – regularization added to the diagonals of covariances (>0)

• ws (array-like (ns,1), optional) – weights for the source samples

• wt (array-like (ns,1), optional) – weights for the target samples

• bias (boolean, optional) – estimate bias \(\mathbf{b}\) else \(\mathbf{b} = 0\) (default:True)

• log (bool, optional) – record log if True

Returns:

• A ((d, d) array-like) – Linear operator

• b ((1, d) array-like) – bias

• log (dict) – log dictionary return only if log==True in parameters

References

[1]

Knott, M. and Smith, C. S. “On the optimal mapping of distributions”, Journal of Optimization Theory and Applications Vol 43, 1984

[2]

Peyré, G., & Cuturi, M. (2017). “Computational Optimal Transport”, 2018.

ot.gaussian.empirical_gaussian_gromov_wasserstein_distance(xs, xt, ws=None, wt=None, log=False)

Return Gaussian Gromov-Wasserstein distance between samples.

The function estimates the Gaussian Gromov-Wasserstein distance between two Gaussien distributions source \(\mu_s\) and target \(\mu_t\), whose parameters are estimated from the provided samples \(\mathcal{X}_s\) and \(\mathcal{X}_t\). See [57] Theorem 4.1 for more details.

Parameters:

• xs (array-like (ns,d)) – samples in the source domain

• xt (array-like (nt,d)) – samples in the target domain

• ws (array-like (ns,1), optional) – weights for the source samples

• wt (array-like (ns,1), optional) – weights for the target samples

• log (bool, optional) – record log if True

Returns:

G – Gaussian Gromov-Wasserstein distance

Return type:

float

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

ot.gaussian.empirical_gaussian_gromov_wasserstein_mapping(xs, xt, ws=None, wt=None, sign_eigs=None, log=False)

Return Gaussian Gromov-Wasserstein mapping between samples.

The function estimates the Gaussian Gromov-Wasserstein mapping between two Gaussian distributions source \(\mu_s\) and target \(\mu_t\), whose parameters are estimated from the provided samples \(\mathcal{X}_s\) and \(\mathcal{X}_t\). See [57] Theorem 4.1 for more details.

Parameters:

• xs (array-like (ns,ds)) – samples in the source domain

• xt (array-like (nt,dt)) – samples in the target domain

• ws (array-like (ns,1), optional) – weights for the source samples

• wt (array-like (ns,1), optional) – weights for the target samples

• sign_eigs (array-like (min(ds,dt),) or string, optional) – sign of the eigenvalues of the mapping matrix, by default all signs will be positive. If ‘skewness’ is provided, the sign of the eigenvalues is selected as the product of the sign of the skewness of the projected data.

• log (bool, optional) – record log if True

Returns:

• A ((dt, ds) array-like) – Linear operator

• b ((1, dt) array-like) – bias

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

ot.gaussian.gaussian_gromov_wasserstein_distance(Cov_s, Cov_t, log=False)

Return the Gaussian Gromov-Wasserstein value from [57].

This function return the closed form value of the Gaussian Gromov-Wasserstein distance between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) when the OT plan is assumed to be also Gaussian. See [57] Theorem 4.1 for more details.

Parameters:

• Cov_s (array-like (ds,ds)) – covariance of the source distribution

• Cov_t (array-like (dt,dt)) – covariance of the target distribution

Returns:

G – Gaussian Gromov-Wasserstein distance

Return type:

float

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.

ot.gaussian.gaussian_gromov_wasserstein_mapping(mu_s, mu_t, Cov_s, Cov_t, sign_eigs=None, log=False)

Return the Gaussian Gromov-Wasserstein mapping from [57].

This function return the closed form value of the Gaussian Gromov-Wasserstein mapping between two Gaussian distributions \(\mathcal{N}(\mu_s,\Sigma_s)\) and \(\mathcal{N}(\mu_t,\Sigma_t)\) when the OT plan is assumed to be also Gaussian. See [57] Theorem 4.1 for more details.

Parameters:

• mu_s (array-like (ds,)) – mean of the source distribution

• mu_t (array-like (dt,)) – mean of the target distribution

• Cov_s (array-like (ds,ds)) – covariance of the source distribution

• Cov_t (array-like (dt,dt)) – covariance of the target distribution

• log (bool, optional) – record log if True

Returns:

• A ((dt, ds) array-like) – Linear operator

• b ((1, dt) array-like) – bias

References

[57]

Delon, J., Desolneux, A., & Salmona, A. (2022). Gromov–Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4), 1178-1198.