# -*- coding: utf-8 -*-
"""
Unbalanced Co-Optimal Transport and Fused Unbalanced Gromov-Wasserstein solvers
"""
# Author: Quang Huy Tran <quang-huy.tran@univ-ubs.fr>
# Alexis Thual <alexis.thual@cea.fr>
#
# License: MIT License
import warnings
from functools import partial
import ot
from ot.backend import get_backend
from ot.utils import list_to_array, get_parameter_pair
from ._utils import (
fused_unbalanced_across_spaces_cost,
uot_cost_matrix,
uot_parameters_and_measures,
)
[docs]
def fused_unbalanced_across_spaces_divergence(
X,
Y,
wx_samp=None,
wx_feat=None,
wy_samp=None,
wy_feat=None,
reg_marginals=10,
epsilon=0,
reg_type="joint",
divergence="kl",
unbalanced_solver="sinkhorn",
alpha=0,
M_samp=None,
M_feat=None,
rescale_plan=True,
init_pi=None,
init_duals=None,
max_iter=100,
tol=1e-7,
max_iter_ot=500,
tol_ot=1e-7,
log=False,
verbose=False,
**kwargs_solver,
):
r"""Compute the fused unbalanced cross-spaces divergence between two matrices equipped
with the distributions on rows and columns. We consider two cases of matrix:
- (Squared) similarity matrix in Gromov-Wasserstein setting,
whose rows and columns represent the samples.
- Arbitrary-size matrix in Co-Optimal Transport setting,
whose rows represent samples, and columns represent corresponding features/dimensions.
More precisely, this function returns the sample and feature transport plans between
:math:`(\mathbf{X}, \mathbf{w}_{xs}, \mathbf{w}_{xf})` and
:math:`(\mathbf{Y}, \mathbf{w}_{ys}, \mathbf{w}_{yf})`,
by solving the following problem using Block Coordinate Descent algorithm:
.. math::
\mathop{\arg \min}_{\mathbf{P}, \mathbf{Q}}
&\quad \sum_{i,j,k,l}
(\mathbf{X}_{i,k} - \mathbf{Y}_{j,l})^2 \mathbf{P}_{i,j} \mathbf{Q}_{k,l} \\
&+ \rho_s \mathbf{Div}(\mathbf{P}_{\# 1} \mathbf{Q}_{\# 1}^T | \mathbf{w}_{xs} \mathbf{w}_{ys}^T)
+ \rho_f \mathbf{Div}(\mathbf{P}_{\# 2} \mathbf{Q}_{\# 2}^T | \mathbf{w}_{xf} \mathbf{w}_{yf}^T) \\
&+ \alpha_s \sum_{i,j} \mathbf{P}_{i,j} \mathbf{M^{(s)}}_{i, j}
+ \alpha_f \sum_{k, l} \mathbf{Q}_{k,l} \mathbf{M^{(f)}}_{k, l}
+ \mathbf{Reg}(\mathbf{P}, \mathbf{Q})
Where:
- :math:`\mathbf{X}`: Source input (arbitrary-size) matrix
- :math:`\mathbf{Y}`: Target input (arbitrary-size) matrix
- :math:`\mathbf{M^{(s)}}`: Additional sample matrix
- :math:`\mathbf{M^{(f)}}`: Additional feature matrix
- :math:`\mathbf{w}_{xs}`: Distribution of the samples in the source space
- :math:`\mathbf{w}_{xf}`: Distribution of the features in the source space
- :math:`\mathbf{w}_{ys}`: Distribution of the samples in the target space
- :math:`\mathbf{w}_{yf}`: Distribution of the features in the target space
- :math:`\mathbf{Div}`: Either Kullback-Leibler divergence or half-squared L2 norm.
- :math:`\mathbf{Reg}`: Regularizer for sample and feature couplings.
We consider two types of regularizer:
+ Independent regularization used in unbalanced Co-Optimal Transport
.. math::
\mathbf{Reg}(\mathbf{P}, \mathbf{Q}) =
\varepsilon_s \mathbf{Div}(\mathbf{P} | \mathbf{w}_{xs} \mathbf{w}_{ys}^T)
+ \varepsilon_f \mathbf{Div}(\mathbf{Q} | \mathbf{w}_{xf} \mathbf{w}_{yf}^T)
+ Joint regularization used in fused unbalanced Gromov-Wasserstein
.. math::
\mathbf{Reg}(\mathbf{P}, \mathbf{Q}) =
\varepsilon \mathbf{Div}(\mathbf{P} \otimes \mathbf{Q} | (\mathbf{w}_{xs} \mathbf{w}_{ys}^T) \otimes (\mathbf{w}_{xf} \mathbf{w}_{yf}^T) )
.. note:: This function allows epsilon to be zero. In that case, `unbalanced_method` must be either "mm" or "lbfgsb".
Parameters
----------
X : (n_sample_x, n_feature_x) array-like, float
Source input matrix.
Y : (n_sample_y, n_feature_y) array-like, float
Target input matrix.
wx_samp : (n_sample_x, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix X.
Uniform distribution by default.
wx_feat : (n_feature_x, ) array-like, float, optional (default = None)
Histogram assigned on columns (features) of matrix X.
Uniform distribution by default.
wy_samp : (n_sample_y, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Y.
Uniform distribution by default.
wy_feat : (n_feature_y, ) array-like, float, optional (default = None)
Histogram assigned on columns (features) of matrix Y.
Uniform distribution by default.
reg_marginals: float or indexable object of length 1 or 2
Marginal relaxation terms for sample and feature couplings.
If `reg_marginals` is a scalar or an indexable object of length 1,
then the same value is applied to both marginal relaxations.
epsilon : scalar or indexable object of length 2, float or int, optional (default = 0)
Regularization parameters for entropic approximation of sample and feature couplings.
Allow the case where `epsilon` contains 0. In that case, the MM solver is used by default
instead of Sinkhorn solver. If `epsilon` is scalar, then the same value is applied to
both regularization of sample and feature couplings.
reg_type: string, optional
- If `reg_type` = "joint": then use joint regularization for couplings.
- If `reg_type` = "independent": then use independent regularization for couplings.
divergence : string, optional (default = "kl")
- If `divergence` = "kl", then Div is the Kullback-Leibler divergence.
- If `divergence` = "l2", then Div is the half squared Euclidean norm.
unbalanced_solver : string, optional (default = "sinkhorn")
Solver for the unbalanced OT subroutine.
- If `divergence` = "kl", then `unbalanced_solver` can be: "sinkhorn", "sinkhorn_log", "mm", "lbfgsb"
- If `divergence` = "l2", then `unbalanced_solver` can be "mm", "lbfgsb"
alpha : scalar or indexable object of length 2, float or int, optional (default = 0)
Coeffficient parameter of linear terms with respect to the sample and feature couplings.
If alpha is scalar, then the same alpha is applied to both linear terms.
M_samp : (n_sample_x, n_sample_y), float, optional (default = None)
Sample matrix associated to the Wasserstein linear term on sample coupling.
M_feat : (n_feature_x, n_feature_y), float, optional (default = None)
Feature matrix associated to the Wasserstein linear term on feature coupling.
rescale_plan : boolean, optional (default = True)
If True, then rescale the sample and feature transport plans within each BCD iteration,
so that they always have equal mass.
init_pi : tuple of two matrices of size (n_sample_x, n_sample_y) and
(n_feature_x, n_feature_y), optional (default = None).
Initialization of sample and feature couplings.
Uniform distributions by default.
init_duals : tuple of two tuples ((n_sample_x, ), (n_sample_y, )) and ((n_feature_x, ), (n_feature_y, )), optional (default = None).
Initialization of sample and feature dual vectors
if using Sinkhorn algorithm. Zero vectors by default.
max_iter : int, optional (default = 100)
Number of Block Coordinate Descent (BCD) iterations.
tol : float, optional (default = 1e-7)
Tolerance of BCD scheme. If the L1-norm between the current and previous
sample couplings is under this threshold, then stop BCD scheme.
max_iter_ot : int, optional (default = 100)
Number of iterations to solve each of the
two unbalanced optimal transport problems in each BCD iteration.
tol_ot : float, optional (default = 1e-7)
Tolerance of unbalanced solver for each of the
two unbalanced optimal transport problems in each BCD iteration.
log : bool, optional (default = False)
If True then the cost and four dual vectors, including
two from sample and two from feature couplings, are recorded.
verbose : bool, optional (default = False)
If True then print the COOT cost at every multiplier of `eval_bcd`-th iteration.
Returns
-------
pi_samp : (n_sample_x, n_sample_y) array-like, float
Sample coupling matrix.
pi_feat : (n_feature_x, n_feature_y) array-like, float
Feature coupling matrix.
log : dictionary, optional
Returned if `log` is True. The keys are:
error : array-like, float
list of L1 norms between the current and previous sample coupling.
duals_sample : (n_sample_x, n_sample_y) tuple, float
Pair of dual vectors when solving OT problem w.r.t the sample coupling.
duals_feature : (n_feature_x, n_feature_y) tuple, float
Pair of dual vectors when solving OT problem w.r.t the feature coupling.
linear : float
Linear part of the cost.
ucoot : float
Total cost.
backend
The proper backend for all input arrays
"""
# MAIN FUNCTION
if reg_type not in ["joint", "independent"]:
raise (NotImplementedError('Unknown reg_type="{}"'.format(reg_type)))
if divergence not in ["kl", "l2"]:
raise (NotImplementedError('Unknown divergence="{}"'.format(divergence)))
if unbalanced_solver not in ["sinkhorn", "sinkhorn_log", "mm", "lbfgsb"]:
raise (NotImplementedError('Unknown method="{}"'.format(unbalanced_solver)))
# hyperparameters
alpha_samp, alpha_feat = get_parameter_pair(alpha)
rho_x, rho_y = get_parameter_pair(reg_marginals)
eps_samp, eps_feat = get_parameter_pair(epsilon)
if reg_type == "joint": # same regularization
eps_feat = eps_samp
if unbalanced_solver in ["sinkhorn", "sinkhorn_log"] and divergence == "l2":
warnings.warn(
"Sinkhorn algorithm does not support L2 norm. \
Divergence is set to 'kl'."
)
divergence = "kl"
if unbalanced_solver in ["sinkhorn", "sinkhorn_log"] and (
eps_samp == 0 or eps_feat == 0
):
warnings.warn(
"Sinkhorn algorithm does not support unregularized problem. \
Solver is set to 'mm'."
)
unbalanced_solver = "mm"
if init_pi is None:
pi_samp, pi_feat = None, None
else:
pi_samp, pi_feat = init_pi
if init_duals is None:
init_duals = (None, None)
duals_samp, duals_feat = init_duals
arr = [X, Y]
for tuple in [duals_samp, duals_feat]:
if tuple is not None:
d1, d2 = duals_feat
if d1 is not None:
arr.append(list_to_array(d1))
if d2 is not None:
arr.append(list_to_array(d2))
nx = get_backend(
*arr, wx_samp, wx_feat, wy_samp, wy_feat, M_samp, M_feat, pi_samp, pi_feat
)
# constant input variables
if M_samp is None:
if alpha_samp > 0:
warnings.warn(
"M_samp is None but alpha_samp = {} > 0. \
The algo will treat as if alpha_samp = 0.".format(alpha_samp)
)
else:
M_samp = alpha_samp * M_samp
if M_feat is None:
if alpha_feat > 0:
warnings.warn(
"M_feat is None but alpha_feat = {} > 0. \
The algo will treat as if alpha_feat = 0.".format(alpha_feat)
)
else:
M_feat = alpha_feat * M_feat
nx_samp, nx_feat = X.shape
ny_samp, ny_feat = Y.shape
# measures on rows and columns
if wx_samp is None:
wx_samp = nx.ones(nx_samp, type_as=X) / nx_samp
if wx_feat is None:
wx_feat = nx.ones(nx_feat, type_as=X) / nx_feat
if wy_samp is None:
wy_samp = nx.ones(ny_samp, type_as=Y) / ny_samp
if wy_feat is None:
wy_feat = nx.ones(ny_feat, type_as=Y) / ny_feat
wxy_samp = wx_samp[:, None] * wy_samp[None, :]
wxy_feat = wx_feat[:, None] * wy_feat[None, :]
# initialize coupling and dual vectors
pi_samp = wxy_samp if pi_samp is None else pi_samp
pi_feat = wxy_feat if pi_feat is None else pi_feat
if unbalanced_solver in ["sinkhorn", "sinkhorn_log"]:
if duals_samp is None:
duals_samp = (nx.zeros(nx_samp, type_as=X), nx.zeros(ny_samp, type_as=Y))
if duals_feat is None:
duals_feat = (nx.zeros(nx_feat, type_as=X), nx.zeros(ny_feat, type_as=Y))
# shortcut functions
X_sqr, Y_sqr = X**2, Y**2
local_cost_samp = partial(
uot_cost_matrix,
data=(X_sqr, Y_sqr, X, Y, M_samp),
tuple_p=(wx_feat, wy_feat),
hyperparams=(rho_x, rho_y, eps_feat),
divergence=divergence,
reg_type=reg_type,
nx=nx,
)
local_cost_feat = partial(
uot_cost_matrix,
data=(X_sqr.T, Y_sqr.T, X.T, Y.T, M_feat),
tuple_p=(wx_samp, wy_samp),
hyperparams=(rho_x, rho_y, eps_samp),
divergence=divergence,
reg_type=reg_type,
nx=nx,
)
parameters_uot_l2_samp = partial(
uot_parameters_and_measures,
tuple_weights=(wx_samp, wy_samp, wxy_samp),
hyperparams=(rho_x, rho_y, eps_samp),
reg_type=reg_type,
divergence=divergence,
nx=nx,
)
parameters_uot_l2_feat = partial(
uot_parameters_and_measures,
tuple_weights=(wx_feat, wy_feat, wxy_feat),
hyperparams=(rho_x, rho_y, eps_feat),
reg_type=reg_type,
divergence=divergence,
nx=nx,
)
solver = partial(
ot.solve,
reg_type=divergence,
unbalanced_type=divergence,
method=unbalanced_solver,
max_iter=max_iter_ot,
tol=tol_ot,
verbose=False,
)
# initialize log
if log:
dict_log = {"backend": nx, "error": []}
for idx in range(max_iter):
pi_samp_prev = nx.copy(pi_samp)
# Update feature coupling
mass = nx.sum(pi_samp)
uot_cost = local_cost_feat(pi=pi_samp)
if divergence == "kl":
new_rho = (rho_x * mass, rho_y * mass)
new_eps = mass * eps_feat if reg_type == "joint" else eps_feat
new_wx, new_wy, new_wxy = wx_feat, wy_feat, wxy_feat
else: # divergence == "l2"
new_w, new_rho, new_eps = parameters_uot_l2_feat(pi_feat)
new_wx, new_wy, new_wxy = new_w
res = solver(
M=uot_cost,
a=new_wx,
b=new_wy,
reg=new_eps,
c=new_wxy,
unbalanced=new_rho,
plan_init=pi_feat,
potentials_init=duals_feat,
)
pi_feat, duals_feat = res.plan, res.potentials
if rescale_plan:
pi_feat = nx.sqrt(mass / nx.sum(pi_feat)) * pi_feat
# Update sample coupling
mass = nx.sum(pi_feat)
uot_cost = local_cost_samp(pi=pi_feat)
if divergence == "kl":
new_rho = (rho_x * mass, rho_y * mass)
new_eps = mass * eps_feat if reg_type == "joint" else eps_feat
new_wx, new_wy, new_wxy = wx_samp, wy_samp, wxy_samp
else: # divergence == "l2"
new_w, new_rho, new_eps = parameters_uot_l2_samp(pi_samp)
new_wx, new_wy, new_wxy = new_w
res = solver(
M=uot_cost,
a=new_wx,
b=new_wy,
reg=new_eps,
c=new_wxy,
unbalanced=new_rho,
plan_init=pi_samp,
potentials_init=duals_samp,
)
pi_samp, duals_samp = res.plan, res.potentials
if rescale_plan:
pi_samp = nx.sqrt(mass / nx.sum(pi_samp)) * pi_samp # shape nx x ny
# get L1 error
err = nx.sum(nx.abs(pi_samp - pi_samp_prev))
if log:
dict_log["error"].append(err)
if verbose:
print("{:5d}|{:8e}|".format(idx + 1, err))
if err < tol:
break
# sanity check
if nx.sum(nx.isnan(pi_samp)) > 0 or nx.sum(nx.isnan(pi_feat)) > 0:
raise (
ValueError(
"There is NaN in coupling. \
Adjust the relaxation or regularization parameters."
)
)
if log:
linear_cost, ucoot_cost = fused_unbalanced_across_spaces_cost(
M_linear=(M_samp, M_feat),
data=(X_sqr, Y_sqr, X, Y),
tuple_pxy_samp=(wx_samp, wy_samp, wxy_samp),
tuple_pxy_feat=(wx_feat, wy_feat, wxy_feat),
pi_samp=pi_samp,
pi_feat=pi_feat,
hyperparams=(rho_x, rho_y, eps_samp, eps_feat),
divergence=divergence,
reg_type=reg_type,
nx=nx,
)
dict_log["duals_sample"] = duals_samp
dict_log["duals_feature"] = duals_feat
dict_log["linear_cost"] = linear_cost
dict_log["ucoot_cost"] = ucoot_cost
return pi_samp, pi_feat, dict_log
else:
return pi_samp, pi_feat
[docs]
def unbalanced_co_optimal_transport(
X,
Y,
wx_samp=None,
wx_feat=None,
wy_samp=None,
wy_feat=None,
reg_marginals=10,
epsilon=0,
divergence="kl",
unbalanced_solver="mm",
alpha=0,
M_samp=None,
M_feat=None,
rescale_plan=True,
init_pi=None,
init_duals=None,
max_iter=100,
tol=1e-7,
max_iter_ot=500,
tol_ot=1e-7,
log=False,
verbose=False,
**kwargs_solve,
):
r"""Compute the unbalanced Co-Optimal Transport between two Euclidean point clouds
(represented as matrices whose rows are samples and columns are the features/dimensions).
More precisely, this function returns the sample and feature transport plans between
:math:`(\mathbf{X}, \mathbf{w}_{xs}, \mathbf{w}_{xf})` and
:math:`(\mathbf{Y}, \mathbf{w}_{ys}, \mathbf{w}_{yf})`,
by solving the following problem using Block Coordinate Descent algorithm:
.. math::
\mathop{\arg \min}_{\mathbf{P}, \mathbf{Q}} &\quad \sum_{i,j,k,l}
(\mathbf{X}_{i,k} - \mathbf{Y}_{j,l})^2 \mathbf{P}_{i,j} \mathbf{Q}_{k,l} \\
&+ \rho_s \mathbf{Div}(\mathbf{P}_{\# 1} \mathbf{Q}_{\# 1}^T | \mathbf{w}_{xs} \mathbf{w}_{ys}^T)
+ \rho_f \mathbf{Div}(\mathbf{P}_{\# 2} \mathbf{Q}_{\# 2}^T | \mathbf{w}_{xf} \mathbf{w}_{yf}^T) \\
&+ \alpha_s \sum_{i,j} \mathbf{P}_{i,j} \mathbf{M^{(s)}}_{i, j}
+ \alpha_f \sum_{k, l} \mathbf{Q}_{k,l} \mathbf{M^{(f)}}_{k, l} \\
&+ \varepsilon_s \mathbf{Div}(\mathbf{P} | \mathbf{w}_{xs} \mathbf{w}_{ys}^T)
+ \varepsilon_f \mathbf{Div}(\mathbf{Q} | \mathbf{w}_{xf} \mathbf{w}_{yf}^T)
Where:
- :math:`\mathbf{X}`: Source input (arbitrary-size) matrix
- :math:`\mathbf{Y}`: Target input (arbitrary-size) matrix
- :math:`\mathbf{M^{(s)}}`: Additional sample matrix
- :math:`\mathbf{M^{(f)}}`: Additional feature matrix
- :math:`\mathbf{w}_{xs}`: Distribution of the samples in the source space
- :math:`\mathbf{w}_{xf}`: Distribution of the features in the source space
- :math:`\mathbf{w}_{ys}`: Distribution of the samples in the target space
- :math:`\mathbf{w}_{yf}`: Distribution of the features in the target space
- :math:`\mathbf{Div}`: Either Kullback-Leibler divergence or half-squared L2 norm.
.. note:: This function allows `epsilon` to be zero. In that case, `unbalanced_method` must be either "mm" or "lbfgsb".
Parameters
----------
X : (n_sample_x, n_feature_x) array-like, float
Source input matrix.
Y : (n_sample_y, n_feature_y) array-like, float
Target input matrix.
wx_samp : (n_sample_x, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix X.
Uniform distribution by default.
wx_feat : (n_feature_x, ) array-like, float, optional (default = None)
Histogram assigned on columns (features) of matrix X.
Uniform distribution by default.
wy_samp : (n_sample_y, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Y.
Uniform distribution by default.
wy_feat : (n_feature_y, ) array-like, float, optional (default = None)
Histogram assigned on columns (features) of matrix Y.
Uniform distribution by default.
reg_marginals: float or indexable object of length 1 or 2
Marginal relaxation terms for sample and feature couplings.
If `reg_marginals is a scalar` or an indexable object of length 1,
then the same value is applied to both marginal relaxations.
epsilon : scalar or indexable object of length 2, float or int, optional (default = 0)
Regularization parameters for entropic approximation of sample and feature couplings.
Allow the case where `epsilon` contains 0. In that case, the MM solver is used by default
instead of Sinkhorn solver. If `epsilon` is scalar, then the same value is applied to
both regularization of sample and feature couplings.
divergence : string, optional (default = "kl")
- If `divergence` = "kl", then Div is the Kullback-Leibler divergence.
- If `divergence` = "l2", then Div is the half squared Euclidean norm.
unbalanced_solver : string, optional (default = "sinkhorn")
Solver for the unbalanced OT subroutine.
- If `divergence` = "kl", then `unbalanced_solver` can be: "sinkhorn", "sinkhorn_log", "mm", "lbfgsb"
- If `divergence` = "l2", then `unbalanced_solver` can be "mm", "lbfgsb"
alpha : scalar or indexable object of length 2, float or int, optional (default = 0)
Coeffficient parameter of linear terms with respect to the sample and feature couplings.
If alpha is scalar, then the same alpha is applied to both linear terms.
M_samp : (n_sample_x, n_sample_y), float, optional (default = None)
Sample matrix associated to the Wasserstein linear term on sample coupling.
M_feat : (n_feature_x, n_feature_y), float, optional (default = None)
Feature matrix associated to the Wasserstein linear term on feature coupling.
rescale_plan : boolean, optional (default = True)
If True, then rescale the sample and feature transport plans within each BCD iteration,
so that they always have equal mass.
init_pi : tuple of two matrices of size (n_sample_x, n_sample_y) and
(n_feature_x, n_feature_y), optional (default = None).
Initialization of sample and feature couplings.
Uniform distributions by default.
init_duals : tuple of two tuples ((n_sample_x, ), (n_sample_y, )) and ((n_feature_x, ), (n_feature_y, )), optional (default = None).
Initialization of sample and feature dual vectors
if using Sinkhorn algorithm. Zero vectors by default.
max_iter : int, optional (default = 100)
Number of Block Coordinate Descent (BCD) iterations.
tol : float, optional (default = 1e-7)
Tolerance of BCD scheme. If the L1-norm between the current and previous
sample couplings is under this threshold, then stop BCD scheme.
max_iter_ot : int, optional (default = 100)
Number of iterations to solve each of the
two unbalanced optimal transport problems in each BCD iteration.
tol_ot : float, optional (default = 1e-7)
Tolerance of unbalanced solver for each of the
two unbalanced optimal transport problems in each BCD iteration.
log : bool, optional (default = False)
If True then the cost and four dual vectors, including
two from sample and two from feature couplings, are recorded.
verbose : bool, optional (default = False)
If True then print the COOT cost at every multiplier of `eval_bcd`-th iteration.
Returns
-------
pi_samp : (n_sample_x, n_sample_y) array-like, float
Sample coupling matrix.
pi_feat : (n_feature_x, n_feature_y) array-like, float
Feature coupling matrix.
log : dictionary, optional
Returned if `log` is True. The keys are:
error : array-like, float
list of L1 norms between the current and previous sample coupling.
duals_sample : (n_sample_x, n_sample_y)-tuple, float
Pair of dual vectors when solving OT problem w.r.t the sample coupling.
duals_feature : (n_feature_x, n_feature_y)-tuple, float
Pair of dual vectors when solving OT problem w.r.t the feature coupling.
linear : float
Linear part of the cost.
ucoot : float
Total cost.
References
----------
.. [71] Tran, H., Janati, H., Courty, N., Flamary, R., Redko, I., Demetci, P., & Singh, R.
Unbalanced Co-Optimal Transport. AAAI Conference on Artificial Intelligence, 2023.
"""
return fused_unbalanced_across_spaces_divergence(
X=X,
Y=Y,
wx_samp=wx_samp,
wx_feat=wx_feat,
wy_samp=wy_samp,
wy_feat=wy_feat,
reg_marginals=reg_marginals,
epsilon=epsilon,
reg_type="independent",
divergence=divergence,
unbalanced_solver=unbalanced_solver,
alpha=alpha,
M_samp=M_samp,
M_feat=M_feat,
rescale_plan=rescale_plan,
init_pi=init_pi,
init_duals=init_duals,
max_iter=max_iter,
tol=tol,
max_iter_ot=max_iter_ot,
tol_ot=tol_ot,
log=log,
verbose=verbose,
**kwargs_solve,
)
[docs]
def unbalanced_co_optimal_transport2(
X,
Y,
wx_samp=None,
wx_feat=None,
wy_samp=None,
wy_feat=None,
reg_marginals=10,
epsilon=0,
divergence="kl",
unbalanced_solver="sinkhorn",
alpha=0,
M_samp=None,
M_feat=None,
rescale_plan=True,
init_pi=None,
init_duals=None,
max_iter=100,
tol=1e-7,
max_iter_ot=500,
tol_ot=1e-7,
log=False,
verbose=False,
**kwargs_solve,
):
r"""Compute the unbalanced Co-Optimal Transport between two Euclidean point clouds
(represented as matrices whose rows are samples and columns are the features/dimensions).
More precisely, this function returns the unbalanced Co-Optimal Transport cost between
:math:`(\mathbf{X}, \mathbf{w}_{xs}, \mathbf{w}_{xf})` and
:math:`(\mathbf{Y}, \mathbf{w}_{ys}, \mathbf{w}_{yf})`,
by solving the following problem using Block Coordinate Descent algorithm:
.. math::
\mathop{\min}_{\mathbf{P}, \mathbf{Q}} &\quad \sum_{i,j,k,l}
(\mathbf{X}_{i,k} - \mathbf{Y}_{j,l})^2 \mathbf{P}_{i,j} \mathbf{Q}_{k,l} \\
&+ \rho_s \mathbf{Div}(\mathbf{P}_{\# 1} \mathbf{Q}_{\# 1}^T | \mathbf{w}_{xs} \mathbf{w}_{ys}^T)
+ \rho_f \mathbf{Div}(\mathbf{P}_{\# 2} \mathbf{Q}_{\# 2}^T | \mathbf{w}_{xf} \mathbf{w}_{yf}^T) \\
&+ \alpha_s \sum_{i,j} \mathbf{P}_{i,j} \mathbf{M^{(s)}}_{i, j}
+ \alpha_f \sum_{k, l} \mathbf{Q}_{k,l} \mathbf{M^{(f)}}_{k, l} \\
&+ \varepsilon_s \mathbf{Div}(\mathbf{P} | \mathbf{w}_{xs} \mathbf{w}_{ys}^T)
+ \varepsilon_f \mathbf{Div}(\mathbf{Q} | \mathbf{w}_{xf} \mathbf{w}_{yf}^T)
Where:
- :math:`\mathbf{X}`: Source input (arbitrary-size) matrix
- :math:`\mathbf{Y}`: Target input (arbitrary-size) matrix
- :math:`\mathbf{M^{(s)}}`: Additional sample matrix
- :math:`\mathbf{M^{(f)}}`: Additional feature matrix
- :math:`\mathbf{w}_{xs}`: Distribution of the samples in the source space
- :math:`\mathbf{w}_{xf}`: Distribution of the features in the source space
- :math:`\mathbf{w}_{ys}`: Distribution of the samples in the target space
- :math:`\mathbf{w}_{yf}`: Distribution of the features in the target space
- :math:`\mathbf{Div}`: Either Kullback-Leibler divergence or half-squared L2 norm.
.. note:: This function allows `epsilon` to be zero. In that case, `unbalanced_method` must be either "mm" or "lbfgsb".
Also the computation of gradients is only supported for KL divergence. The case of half squared-L2 norm uses those of KL divergence.
Parameters
----------
X : (n_sample_x, n_feature_x) array-like, float
Source input matrix.
Y : (n_sample_y, n_feature_y) array-like, float
Target input matrix.
wx_samp : (n_sample_x, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix X.
Uniform distribution by default.
wx_feat : (n_feature_x, ) array-like, float, optional (default = None)
Histogram assigned on columns (features) of matrix X.
Uniform distribution by default.
wy_samp : (n_sample_y, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Y.
Uniform distribution by default.
wy_feat : (n_feature_y, ) array-like, float, optional (default = None)
Histogram assigned on columns (features) of matrix Y.
Uniform distribution by default.
reg_marginals: float or indexable object of length 1 or 2
Marginal relaxation terms for sample and feature couplings.
If `reg_marginals` is a scalar or an indexable object of length 1,
then the same value is applied to both marginal relaxations.
epsilon : scalar or indexable object of length 2, float or int, optional (default = 0)
Regularization parameters for entropic approximation of sample and feature couplings.
Allow the case where `epsilon` contains 0. In that case, the MM solver is used by default
instead of Sinkhorn solver. If `epsilon` is scalar, then the same value is applied to
both regularization of sample and feature couplings.
divergence : string, optional (default = "kl")
- If `divergence` = "kl", then Div is the Kullback-Leibler divergence.
- If `divergence` = "l2", then Div is the half squared Euclidean norm.
unbalanced_solver : string, optional (default = "sinkhorn")
Solver for the unbalanced OT subroutine.
- If `divergence` = "kl", then `unbalanced_solver` can be: "sinkhorn", "sinkhorn_log", "mm", "lbfgsb"
- If `divergence` = "l2", then `unbalanced_solver` can be "mm", "lbfgsb"
alpha : scalar or indexable object of length 2, float or int, optional (default = 0)
Coeffficient parameter of linear terms with respect to the sample and feature couplings.
If alpha is scalar, then the same alpha is applied to both linear terms.
M_samp : (n_sample_x, n_sample_y), float, optional (default = None)
Sample matrix associated to the Wasserstein linear term on sample coupling.
M_feat : (n_feature_x, n_feature_y), float, optional (default = None)
Feature matrix associated to the Wasserstein linear term on feature coupling.
rescale_plan : boolean, optional (default = True)
If True, then rescale the transport plans in each BCD iteration,
so that they always have equal mass.
init_pi : tuple of two matrices of size (n_sample_x, n_sample_y) and
(n_feature_x, n_feature_y), optional (default = None).
Initialization of sample and feature couplings.
Uniform distributions by default.
init_duals : tuple of two tuples ((n_sample_x, ), (n_sample_y, )) and ((n_feature_x, ), (n_feature_y, )), optional (default = None).
Initialization of sample and feature dual vectors
if using Sinkhorn algorithm. Zero vectors by default.
max_iter : int, optional (default = 100)
Number of Block Coordinate Descent (BCD) iterations.
tol : float, optional (default = 1e-7)
Tolerance of BCD scheme. If the L1-norm between the current and previous
sample couplings is under this threshold, then stop BCD scheme.
max_iter_ot : int, optional (default = 100)
Number of iterations to solve each of the
two unbalanced optimal transport problems in each BCD iteration.
tol_ot : float, optional (default = 1e-7)
Tolerance of unbalanced solver for each of the
two unbalanced optimal transport problems in each BCD iteration.
log : bool, optional (default = False)
If True then the cost and four dual vectors, including
two from sample and two from feature couplings, are recorded.
verbose : bool, optional (default = False)
If True then print the COOT cost at every multiplier of `eval_bcd`-th iteration.
Returns
-------
ucoot : float
UCOOT cost.
log : dictionary, optional
Returned if `log` is True. The keys are:
error : array-like, float
list of L1 norms between the current and previous sample coupling.
duals_sample : (n_sample_x, n_sample_y)-tuple, float
Pair of dual vectors when solving OT problem w.r.t the sample coupling.
duals_feature : (n_feature_x, n_feature_y)-tuple, float
Pair of dual vectors when solving OT problem w.r.t the feature coupling.
linear : float
Linear part of UCOOT cost.
ucoot : float
UCOOT cost.
backend
The proper backend for all input arrays
References
----------
.. [71] Tran, H., Janati, H., Courty, N., Flamary, R., Redko, I., Demetci, P., & Singh, R.
Unbalanced Co-Optimal Transport. AAAI Conference on Artificial Intelligence, 2023.
"""
if divergence != "kl":
warnings.warn(
"The computation of gradients is only supported for KL divergence, not \
for {} divergence".format(divergence)
)
pi_samp, pi_feat, log_ucoot = unbalanced_co_optimal_transport(
X=X,
Y=Y,
wx_samp=wx_samp,
wx_feat=wx_feat,
wy_samp=wy_samp,
wy_feat=wy_feat,
reg_marginals=reg_marginals,
epsilon=epsilon,
divergence=divergence,
unbalanced_solver=unbalanced_solver,
alpha=alpha,
M_samp=M_samp,
M_feat=M_feat,
rescale_plan=rescale_plan,
init_pi=init_pi,
init_duals=init_duals,
max_iter=max_iter,
tol=tol,
max_iter_ot=max_iter_ot,
tol_ot=tol_ot,
log=True,
verbose=verbose,
**kwargs_solve,
)
nx = log_ucoot["backend"]
nx_samp, nx_feat = X.shape
ny_samp, ny_feat = Y.shape
# measures on rows and columns
if wx_samp is None:
wx_samp = nx.ones(nx_samp, type_as=X) / nx_samp
if wx_feat is None:
wx_feat = nx.ones(nx_feat, type_as=X) / nx_feat
if wy_samp is None:
wy_samp = nx.ones(ny_samp, type_as=Y) / ny_samp
if wy_feat is None:
wy_feat = nx.ones(ny_feat, type_as=Y) / ny_feat
# extract parameters
rho_x, rho_y = get_parameter_pair(reg_marginals)
eps_samp, eps_feat = get_parameter_pair(epsilon)
# calculate marginals
pi1_samp, pi2_samp = nx.sum(pi_samp, 1), nx.sum(pi_samp, 0)
pi1_feat, pi2_feat = nx.sum(pi_feat, 1), nx.sum(pi_feat, 0)
m_samp, m_feat = nx.sum(pi1_samp), nx.sum(pi1_feat)
m_wx_feat, m_wx_samp = nx.sum(wx_feat), nx.sum(wx_samp)
m_wy_feat, m_wy_samp = nx.sum(wy_feat), nx.sum(wy_samp)
# calculate subgradients
gradX = 2 * X * (pi1_samp[:, None] * pi1_feat[None, :]) - 2 * nx.dot(
nx.dot(pi_samp, Y), pi_feat.T
) # shape (nx_samp, nx_feat)
gradY = 2 * Y * (pi2_samp[:, None] * pi2_feat[None, :]) - 2 * nx.dot(
nx.dot(pi_samp.T, X), pi_feat
) # shape (ny_samp, ny_feat)
grad_wx_samp = rho_x * (m_wx_feat - m_feat * pi1_samp / wx_samp) + eps_samp * (
m_wy_samp - pi1_samp / wx_samp
)
grad_wx_feat = rho_x * (m_wx_samp - m_samp * pi1_feat / wx_feat) + eps_feat * (
m_wy_feat - pi1_feat / wx_feat
)
grad_wy_samp = rho_y * (m_wy_feat - m_feat * pi2_samp / wy_samp) + eps_samp * (
m_wx_samp - pi2_samp / wy_samp
)
grad_wy_feat = rho_y * (m_wy_samp - m_samp * pi2_feat / wy_feat) + eps_feat * (
m_wx_feat - pi2_feat / wy_feat
)
# set gradients
ucoot = log_ucoot["ucoot_cost"]
ucoot = nx.set_gradients(
ucoot,
(X, Y, wx_samp, wx_feat, wy_samp, wy_feat),
(gradX, gradY, grad_wx_samp, grad_wx_feat, grad_wy_samp, grad_wy_feat),
)
if log:
return ucoot, log_ucoot
else:
return ucoot
[docs]
def fused_unbalanced_gromov_wasserstein(
Cx,
Cy,
wx=None,
wy=None,
reg_marginals=10,
epsilon=0,
divergence="kl",
unbalanced_solver="mm",
alpha=0,
M=None,
init_duals=None,
init_pi=None,
max_iter=100,
tol=1e-7,
max_iter_ot=500,
tol_ot=1e-7,
log=False,
verbose=False,
**kwargs_solve,
):
r"""Compute the lower bound of the fused unbalanced Gromov-Wasserstein (FUGW) between two similarity matrices.
In practice, this lower bound is used interchangeably with the true FUGW.
More precisely, this function returns the transport plan between
:math:`(\mathbf{C^X}, \mathbf{w_X})` and :math:`(\mathbf{C^Y}, \mathbf{w_Y})`,
by solving the following problem using Block Coordinate Descent algorithm:
.. math::
\mathop{\arg \min}_{\substack{\mathbf{P}, \mathbf{Q}: \\ mass(P) = mass(Q)}}
&\quad \sum_{i,j,k,l} (\mathbf{C^X}_{i,k} - \mathbf{C^Y}_{j,l})^2 \mathbf{P}_{i,j} \mathbf{Q}_{k,l}
+ \frac{\alpha}{2} \sum_{i,j} (\mathbf{P}_{i,j} + \mathbf{Q}_{i,j}) \mathbf{M}_{i, j} \\
&+ \rho_1 \mathbf{Div}(\mathbf{P}_{\# 1} \mathbf{Q}_{\# 1}^T | \mathbf{w_X} \mathbf{w_X}^T)
+ \rho_2 \mathbf{Div}(\mathbf{P}_{\# 2} \mathbf{Q}_{\# 2}^T | \mathbf{w_Y} \mathbf{w_Y}^T) \\
&+ \varepsilon \mathbf{Div}(\mathbf{P} \otimes \mathbf{Q} | (\mathbf{w_X} \mathbf{w_Y}^T) \otimes (\mathbf{w_X} \mathbf{w_Y}^T))
Where:
- :math:`\mathbf{C^X}`: Source similarity matrix
- :math:`\mathbf{C^Y}`: Target similarity matrix
- :math:`\mathbf{M}`: Sample matrix corresponding to the Wasserstein term
- :math:`\mathbf{w_X}`: Distribution of the samples in the source space
- :math:`\mathbf{w_Y}`: Distribution of the samples in the target space
- :math:`\mathbf{Div}`: Either Kullback-Leibler divergence or half-squared L2 norm.
.. note:: This function allows epsilon to be zero. In that case, `unbalanced_method` must be either "mm" or "lbfgsb".
Parameters
----------
Cx : (n_sample_x, n_feature_x) array-like, float
Source similarity matrix.
Cy : (n_sample_y, n_feature_y) array-like, float
Target similarity matrix.
wx : (n_sample_x, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Cx.
Uniform distribution by default.
wy : (n_sample_y, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Cy.
Uniform distribution by default.
reg_marginals: float or indexable object of length 1 or 2
Marginal relaxation terms for sample and feature couplings.
If `reg_marginals` is a scalar or an indexable object of length 1,
then the same value is applied to both marginal relaxations.
epsilon : scalar, float or int, optional (default = 0)
Regularization parameters for entropic approximation of sample and feature couplings.
Allow the case where `epsilon` contains 0. In that case, the MM solver is used by default
instead of Sinkhorn solver. If `epsilon` is scalar, then the same value is applied to
both regularization of sample and feature couplings.
divergence : string, optional (default = "kl")
- If `divergence` = "kl", then Div is the Kullback-Leibler divergence.
- If `divergence` = "l2", then Div is the half squared Euclidean norm.
unbalanced_solver : string, optional (default = "sinkhorn")
Solver for the unbalanced OT subroutine.
- If `divergence` = "kl", then `unbalanced_solver` can be: "sinkhorn", "sinkhorn_log", "mm", "lbfgsb"
- If `divergence` = "l2", then `unbalanced_solver` can be "mm", "lbfgsb"
alpha : scalar, float or int, optional (default = 0)
Coeffficient parameter of linear terms with respect to the sample and feature couplings.
If alpha is scalar, then the same alpha is applied to both linear terms.
M : (n_sample_x, n_sample_y), float, optional (default = None)
Sample matrix associated to the Wasserstein linear term on sample coupling.
init_pi :(n_sample_x, n_sample_y) array-like, optional (default = None)
Initialization of sample coupling. By default = :math:`w_X w_Y^T`.
init_duals : tuple of vectors ((n_sample_x, ), (n_sample_y, )), optional (default = None).
Initialization of sample and feature dual vectors
if using Sinkhorn algorithm. Zero vectors by default.
max_iter : int, optional (default = 100)
Number of Block Coordinate Descent (BCD) iterations.
tol : float, optional (default = 1e-7)
Tolerance of BCD scheme. If the L1-norm between the current and previous
sample couplings is under this threshold, then stop BCD scheme.
max_iter_ot : int, optional (default = 100)
Number of iterations to solve each of the
two unbalanced optimal transport problems in each BCD iteration.
tol_ot : float, optional (default = 1e-7)
Tolerance of unbalanced solver for each of the
two unbalanced optimal transport problems in each BCD iteration.
log : bool, optional (default = False)
If True then the cost and four dual vectors, including
two from sample and two from feature couplings, are recorded.
verbose : bool, optional (default = False)
If True then print the COOT cost at every multiplier of `eval_bcd`-th iteration.
Returns
-------
pi_samp : (n_sample_x, n_sample_y) array-like, float
Sample coupling matrix.
In practice, we use this matrix as solution of FUGW.
pi_samp2 : (n_sample_x, n_sample_y) array-like, float
Second sample coupling matrix.
In practice, we usually ignore this output.
log : dictionary, optional
Returned if `log` is True. The keys are:
error : array-like, float
list of L1 norms between the current and previous sample couplings.
duals : (n_sample_x, n_sample_y)-tuple, float
Pair of dual vectors when solving OT problem w.r.t the sample coupling.
linear : float
Linear part of FUGW cost.
fugw_cost : float
Total FUGW cost.
backend
The proper backend for all input arrays
References
----------
.. [70] Thual, A., Tran, H., Zemskova, T., Courty, N., Flamary, R., Dehaene, S., & Thirion, B.
Aligning individual brains with Fused Unbalanced Gromov-Wasserstein.
Advances in Neural Information Systems, 35 (2022).
.. [72] Thibault Séjourné, François-Xavier Vialard, & Gabriel Peyré.
The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation.
Neural Information Processing Systems, 34 (2021).
"""
alpha = (alpha / 2, alpha / 2)
pi_samp, pi_feat, dict_log = fused_unbalanced_across_spaces_divergence(
X=Cx,
Y=Cy,
wx_samp=wx,
wx_feat=wx,
wy_samp=wy,
wy_feat=wy,
reg_marginals=reg_marginals,
epsilon=epsilon,
reg_type="joint",
divergence=divergence,
unbalanced_solver=unbalanced_solver,
alpha=alpha,
M_samp=M,
M_feat=M,
rescale_plan=True,
init_pi=(init_pi, init_pi),
init_duals=(init_duals, init_duals),
max_iter=max_iter,
tol=tol,
max_iter_ot=max_iter_ot,
tol_ot=tol_ot,
log=True,
verbose=verbose,
**kwargs_solve,
)
if log:
log_fugw = {
"error": dict_log["error"],
"duals": dict_log["duals_sample"],
"linear_cost": dict_log["linear_cost"],
"fugw_cost": dict_log["ucoot_cost"],
"backend": dict_log["backend"],
}
return pi_samp, pi_feat, log_fugw
else:
return pi_samp, pi_feat
[docs]
def fused_unbalanced_gromov_wasserstein2(
Cx,
Cy,
wx=None,
wy=None,
reg_marginals=10,
epsilon=0,
divergence="kl",
unbalanced_solver="mm",
alpha=0,
M=None,
init_duals=None,
init_pi=None,
max_iter=100,
tol=1e-7,
max_iter_ot=500,
tol_ot=1e-7,
log=False,
verbose=False,
**kwargs_solve,
):
r"""Compute the lower bound of the fused unbalanced Gromov-Wasserstein (FUGW) between two similarity matrices.
In practice, this lower bound is used interchangeably with the true FUGW.
More precisely, this function returns the lower bound of the fused unbalanced Gromov-Wasserstein cost between
:math:`(\mathbf{C^X}, \mathbf{w_X})` and :math:`(\mathbf{C^Y}, \mathbf{w_Y})`,
by solving the following problem using Block Coordinate Descent algorithm:
.. math::
\mathop{\min}_{\substack{\mathbf{P}, \mathbf{Q}: \\ mass(P) = mass(Q)}}
&\quad \sum_{i,j,k,l} (\mathbf{C^X}_{i,k} - \mathbf{C^Y}_{j,l})^2 \mathbf{P}_{i,j} \mathbf{Q}_{k,l}
+ \frac{\alpha}{2} \sum_{i,j} (\mathbf{P}_{i,j} + \mathbf{Q}_{i,j}) \mathbf{M}_{i, j} \\
&+ \rho_1 \mathbf{Div}(\mathbf{P}_{\# 1} \mathbf{Q}_{\# 1}^T | \mathbf{w_X} \mathbf{w_X}^T)
+ \rho_2 \mathbf{Div}(\mathbf{P}_{\# 2} \mathbf{Q}_{\# 2}^T | \mathbf{w_Y} \mathbf{w_Y}^T) \\
&+ \varepsilon \mathbf{Div}(\mathbf{P} \otimes \mathbf{Q} | (\mathbf{w_X} \mathbf{w_Y}^T) \otimes (\mathbf{w_X} \mathbf{w_Y}^T))
Where:
- :math:`\mathbf{C^X}`: Source similarity matrix
- :math:`\mathbf{C^Y}`: Target similarity matrix
- :math:`\mathbf{M}`: Sample matrix corresponding to the Wasserstein term
- :math:`\mathbf{w_X}`: Distribution of the samples in the source space
- :math:`\mathbf{w_Y}`: Distribution of the samples in the target space
- :math:`\mathbf{Div}`: Either Kullback-Leibler divergence or half-squared L2 norm.
.. note:: This function allows `epsilon` to be zero. In that case, unbalanced_method must be either "mm" or "lbfgsb".
Also the computation of gradients is only supported for KL divergence, but not for half squared-L2 norm. In case of half squared-L2 norm, the calculation of KL divergence will be used.
Parameters
----------
Cx : (n_sample_x, n_feature_x) array-like, float
Source similarity matrix.
Cy : (n_sample_y, n_feature_y) array-like, float
Target similarity matrix.
wx : (n_sample_x, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Cx.
Uniform distribution by default.
wy : (n_sample_y, ) array-like, float, optional (default = None)
Histogram assigned on rows (samples) of matrix Cy.
Uniform distribution by default.
reg_marginals: float or indexable object of length 1 or 2
Marginal relaxation terms for sample and feature couplings.
If `reg_marginals` is a scalar or an indexable object of length 1,
then the same value is applied to both marginal relaxations.
epsilon : scalar, float or int, optional (default = 0)
Regularization parameters for entropic approximation of sample and feature couplings.
Allow the case where `epsilon` contains 0. In that case, the MM solver is used by default
instead of Sinkhorn solver. If `epsilon` is scalar, then the same value is applied to
both regularization of sample and feature couplings.
divergence : string, optional (default = "kl")
- If `divergence` = "kl", then Div is the Kullback-Leibler divergence.
- If `divergence` = "l2", then Div is the half squared Euclidean norm.
unbalanced_solver : string, optional (default = "sinkhorn")
Solver for the unbalanced OT subroutine.
- If `divergence` = "kl", then `unbalanced_solver` can be: "sinkhorn", "sinkhorn_log", "mm", "lbfgsb"
- If `divergence` = "l2", then `unbalanced_solver` can be "mm", "lbfgsb"
alpha : scalar, float or int, optional (default = 0)
Coeffficient parameter of linear terms with respect to the sample and feature couplings.
If alpha is scalar, then the same alpha is applied to both linear terms.
M : (n_sample_x, n_sample_y), float, optional (default = None)
Sample matrix associated to the Wasserstein linear term on sample coupling.
init_pi :(n_sample_x, n_sample_y) array-like, optional (default = None)
Initialization of sample coupling. By default = :math:`w_X w_Y^T`.
init_duals : tuple of vectors ((n_sample_x, ), (n_sample_y, )), optional (default = None).
Initialization of sample and feature dual vectors
if using Sinkhorn algorithm. Zero vectors by default.
max_iter : int, optional (default = 100)
Number of Block Coordinate Descent (BCD) iterations.
tol : float, optional (default = 1e-7)
Tolerance of BCD scheme. If the L1-norm between the current and previous
sample couplings is under this threshold, then stop BCD scheme.
max_iter_ot : int, optional (default = 100)
Number of iterations to solve each of the
two unbalanced optimal transport problems in each BCD iteration.
tol_ot : float, optional (default = 1e-7)
Tolerance of unbalanced solver for each of the
two unbalanced optimal transport problems in each BCD iteration.
log : bool, optional (default = False)
If True then the cost and four dual vectors, including
two from sample and two from feature couplings, are recorded.
verbose : bool, optional (default = False)
If True then print the COOT cost at every multiplier of `eval_bcd`-th iteration.
Returns
-------
fugw : float
Total FUGW cost
log : dictionary, optional
Returned if `log` is True. The keys are:
error : array-like, float
list of L1 norms between the current and previous sample couplings.
duals : (n_sample_x, n_sample_y)-tuple, float
Pair of dual vectors when solving OT problem w.r.t the sample coupling.
linear : float
Linear part of FUGW cost.
fugw_cost : float
Total FUGW cost.
backend
The proper backend for all input arrays
References
----------
.. [70] Thual, A., Tran, H., Zemskova, T., Courty, N., Flamary, R., Dehaene, S., & Thirion, B.
Aligning individual brains with Fused Unbalanced Gromov-Wasserstein.
Advances in Neural Information Systems, 35 (2022).
.. [72] Thibault Séjourné, François-Xavier Vialard, & Gabriel Peyré.
The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation.
Neural Information Processing Systems, 34 (2021).
"""
if divergence != "kl":
warnings.warn(
"The computation of gradients is only supported for KL divergence, \
but not for {} divergence. The gradient of the KL case will be used.".format(
divergence
)
)
pi_samp, pi_feat, log_fugw = fused_unbalanced_gromov_wasserstein(
Cx=Cx,
Cy=Cy,
wx=wx,
wy=wy,
reg_marginals=reg_marginals,
epsilon=epsilon,
divergence=divergence,
unbalanced_solver=unbalanced_solver,
alpha=alpha,
M=M,
init_duals=init_duals,
init_pi=init_pi,
max_iter=max_iter,
tol=tol,
max_iter_ot=max_iter_ot,
tol_ot=tol_ot,
log=True,
verbose=verbose,
**kwargs_solve,
)
nx = log_fugw["backend"]
sx, sy = Cx.shape[0], Cy.shape[0]
# measures on rows and columns
if wx is None:
wx = nx.ones(sx, type_as=Cx) / sx
if wy is None:
wy = nx.ones(sy, type_as=Cy) / sy
# calculate marginals
pi1_samp, pi2_samp = nx.sum(pi_samp, 1), nx.sum(pi_samp, 0)
pi1_feat, pi2_feat = nx.sum(pi_feat, 1), nx.sum(pi_feat, 0)
m_samp, m_feat = nx.sum(pi1_samp), nx.sum(pi1_feat)
m_wx, m_wy = nx.sum(wx), nx.sum(wy)
# calculate subgradients
gradX = 2 * Cx * (pi1_samp[:, None] * pi1_feat[None, :]) - 2 * nx.dot(
nx.dot(pi_samp, Cy), pi_feat.T
) # shape (nx_samp, nx_feat)
gradY = 2 * Cy * (pi2_samp[:, None] * pi2_feat[None, :]) - 2 * nx.dot(
nx.dot(pi_samp.T, Cx), pi_feat
) # shape (ny_samp, ny_feat)
gradM = alpha / 2 * (pi_samp + pi_feat)
rho_x, rho_y = get_parameter_pair(reg_marginals)
grad_wx = (
2 * m_wx * (rho_x + epsilon * m_wy**2)
- (rho_x + epsilon) * (m_feat * pi1_samp + m_samp * pi1_feat) / wx
)
grad_wy = (
2 * m_wy * (rho_y + epsilon * m_wx**2)
- (rho_y + epsilon) * (m_feat * pi2_samp + m_samp * pi2_feat) / wy
)
# set gradients
fugw = log_fugw["fugw_cost"]
fugw = nx.set_gradients(
fugw, (Cx, Cy, M, wx, wy), (gradX, gradY, gradM, grad_wx, grad_wy)
)
if log:
return fugw, log_fugw
else:
return fugw