diff --git a/README.md b/README.md index 06be2658..1e1dcb5f 100644 --- a/README.md +++ b/README.md @@ -245,4 +245,8 @@ The library is distributed under the 3-Clause BSD license. [36] Xiao, Zhiqing, Wang, Haobo, Jin, Ying, Feng, Lei, Chen, Gang, Huang, Fei, Zhao, Junbo.[SPA: A Graph Spectral Alignment Perspective for Domain Adaptation](https://arxiv.org/pdf/2310.17594). In Neurips, 2023. -[37] Xie, Renchunzi, Odonnat, Ambroise, Feofanov, Vasilii, Deng, Weijian, Zhang, Jianfeng and An, Bo. [MaNo: Exploiting Matrix Norm for Unsupervised Accuracy Estimation Under Distribution Shifts](https://arxiv.org/pdf/2405.18979). In NeurIPS, 2024. \ No newline at end of file +[37] Xie, Renchunzi, Odonnat, Ambroise, Feofanov, Vasilii, Deng, Weijian, Zhang, Jianfeng and An, Bo. [MaNo: Exploiting Matrix Norm for Unsupervised Accuracy Estimation Under Distribution Shifts](https://arxiv.org/pdf/2405.18979). In NeurIPS, 2024. + +[38] Álvarez-Esteban, Pedro C., et al. [A fixed-point approach to barycenters in Wasserstein space.](https://arxiv.org/abs/1511.05355) Journal of Mathematical Analysis and Applications 441.2 (2016): 744-762 + +[39] Montesuma, Eduardo, Fred Maurice Ngole Mboula, and Antoine Souloumiac. [Multi-source domain adaptation through dataset dictionary learning in wasserstein space.](https://arxiv.org/pdf/2307.14953) ECAI 2023. IOS Press, 2023. 1739-1746. diff --git a/examples/methods/plot_joint_wasserstein_barycenter.py b/examples/methods/plot_joint_wasserstein_barycenter.py new file mode 100644 index 00000000..efc11e21 --- /dev/null +++ b/examples/methods/plot_joint_wasserstein_barycenter.py @@ -0,0 +1,315 @@ +""" +Computation of feature-label joint Wasserstein Barycenters +========================================================== + +This example illustrates the computation of feature-label joint Wasserstein barycenters + +""" + +# Author: Eduardo Fernandes Montesuma +# +# License: BSD 3-Clause + +# %% Imports +import matplotlib.pyplot as plt +import numpy as np +import ot +from sklearn.datasets import make_moons + +from skada._mapping import joint_wasserstein_barycenter + +# %% +# Generate labeled data from multiple distributions +# ------------------------------------------------- +# +# Here, we use as the base distribution the famous +# moons dataset. We then generate other 2 measure +# by translating the original dataset. This corresponds +# to applying a linear mapping :math:`T_{b}(x) = x + b` on +# each sample in the measures' support, i.e., applying +# :math:`P_{i} = T_{b,\sharp}P_{0}.` +X0, y0 = make_moons(n_samples=100, noise=0.1) +y1 = y0.copy() +y2 = y0.copy() + +X1 = X0 + np.array([2, 0])[None, :] +X2 = X0 + np.array([1, np.sqrt(3)])[None, :] + +Xs = [X0, X1, X2] + +# Converts labels into one-hot encoded labels +Ys = [] +for y in [y0, y1, y2]: + Y = np.zeros((y.size, y.max() + 1)) + Y[np.arange(y.size), y] = 1 + Ys.append(Y) + + +# %% +# Gaussian Modeling +# ----------------- +# +# We start our illlustration of Wasserstein barycenters +# by computing the Bures-Wasserstein barycenter. In this +# case, one assumes each measure is a Gaussian with its +# own mean vector and covariance matrix. +means = np.concatenate( + [X0.mean(axis=0)[None, :], X1.mean(axis=0)[None, :], X2.mean(axis=0)[None, :]], + axis=0, +) # shape: (k, d) + +covs = np.concatenate( + [np.cov(X0.T)[None, ...], np.cov(X1.T)[None, ...], np.cov(X2.T)[None, ...]], axis=0 +) # shape: (k, d, d) + + +# %% +# Bures-Wasserstein Barycenter +# ---------------------------- +# +# Here, we compute the Bures-Wasserstein barycenter using +# the parameters obtained in the previous section. The +# barycenter is calculated using a fixed-point algorithm. +# See, for instance (Álvarez-Esteban et al., 2016) +barycenter_mean, barycenter_cov = ot.gaussian.bures_wasserstein_barycenter( + m=means, C=covs, eps=1e-8 +) + +mappings = [ + ot.gaussian.bures_wasserstein_mapping( + ms=m, Cs=C, mt=barycenter_mean, Ct=barycenter_cov + ) + for m, C in zip(means, covs) +] + +linear_XB, linear_YB = [], [] +for _X, _Y, (A, b) in zip(Xs, Ys, mappings): + linear_XB.append(_X.dot(A) + b) + linear_YB.append(_Y) +linear_XB = np.concatenate(linear_XB, axis=0) +linear_YB = np.concatenate(linear_YB, axis=0) + +# %% +# Plots the results of Gaussian Wasserstein Barycenter +# ------------------------ +# +fig, axes = plt.subplots(1, 4, figsize=(16, 4), sharex=True, sharey=True) + +names = ["$X_{0}$", "$X_{1}$", "$X_{2}$", "$X_{B}$"] +for ax, _X, _Y, name in zip( + axes, + Xs + + [ + linear_XB, + ], + Ys + + [ + linear_YB, + ], + names, +): + ax.scatter(_X[:, 0], _X[:, 1], c=_Y.argmax(axis=1), cmap=plt.cm.coolwarm) + ax.set_title +plt.suptitle("Bures-Wasserstein Barycenter") +plt.tight_layout() +plt.show() + +# %% +# Compute the Empirical Wasserstein barycenter +# -------------------------------------------- +# +# Computes the Barycenter +XB, YB = joint_wasserstein_barycenter( + Xs, + Ys, + mus=None, + XB=None, + YB=None, + muB=None, + measure_weights=None, + n_samples=X0.shape[0], + reg_e=0.0, + verbose=True, +) + +# %% +# Plots the results of Empirical Wasserstein Barycenter +# ----------------------------------------------------- +# +fig, axes = plt.subplots(1, 4, figsize=(16, 4), sharex=True, sharey=True) + +names = ["$X_{0}$", "$X_{1}$", "$X_{2}$", "$X_{B}$"] +for ax, _X, _Y, name in zip( + axes, + Xs + + [ + XB, + ], + Ys + + [ + YB, + ], + names, +): + ax.scatter(_X[:, 0], _X[:, 1], c=_Y.argmax(axis=1), cmap=plt.cm.coolwarm) + ax.set_title +plt.suptitle("Empirical Wasserstein Barycenter") +plt.tight_layout() +plt.show() + +# %% +# Compares the obtained barycenters +# --------------------------------- +# +# +fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True) + +axes[0].scatter(XB[:, 0], XB[:, 1], c=YB.argmax(axis=1), cmap=plt.cm.coolwarm) +axes[0].set_title("Empirical Wasserstein Barycenter") +axes[1].scatter( + linear_XB[:, 0], linear_XB[:, 1], c=linear_YB.argmax(axis=1), cmap=plt.cm.coolwarm +) +axes[1].set_title("Gaussian Wasserstein Barycenter") + +plt.show() + + +# %% +# When to choose the mapping strategy +# ----------------------------------- +# +# In the previous example, you saw that there +# is no much difference between the empirical +# and the Gaussian strategies for computing +# Wasserstein barycenters. This is true because +# the mapping that generates the different measures +# is an affine transformation. More generally, if +# we expect that the mappings between all the measures +# involved is affine (e.g., $T(x) = Ax + b$), then +# we can successfully use Gaussian modeling. We now +# present an example where it fails. +def non_affine_map(points, b): + """ + Apply the non-affine map T(x, y; b) = [x^2 - y^2 + b, 2xy] to a set of points. + + Parameters + ---------- + points (np.ndarray): An array of shape (N, 2), where each row is a point (x, y). + + Returns + ------- + np.ndarray: An array of shape (N, 2), where each row is the transformed point. + """ + x = points[:, 0] # Extract x-coordinates + y = points[:, 1] # Extract y-coordinates + + # Apply the transformation + x_transformed = x**2 - y**2 + b + y_transformed = 2 * x * y + + # Stack the results into a new array of shape (N, 2) + transformed_points = np.column_stack((x_transformed, y_transformed)) + return transformed_points + + +# %% +# Transforms the samples +# ---------------------- +# +# Here we transform all the samples from the first +# measure +X1 = non_affine_map(X0, b=3) +y1 = y0.copy() +Xs = [X0, X1] + +# Converts labels into one-hot encoded labels +Ys = [] +for y in [y0, y1]: + Y = np.zeros((y.size, y.max() + 1)) + Y[np.arange(y.size), y] = 1 + Ys.append(Y) + +# %% +# Plot the measures' support +# -------------------------- +# +fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True) + +axes[0].scatter(X0[:, 0], X0[:, 1], c=y0, cmap=plt.cm.coolwarm) +axes[0].set_title("Measure 0") +axes[1].scatter(X1[:, 0], X1[:, 1], c=y1, cmap=plt.cm.coolwarm) +axes[1].set_title("Measure 1") + +plt.show() + +# %% +# Compute the Bures-Wasserstein Barycenter +# ---------------------------------------- +# +means = np.concatenate( + [X0.mean(axis=0)[None, :], X1.mean(axis=0)[None, :]], axis=0 +) # shape: (k, d) + +covs = np.concatenate( + [np.cov(X0.T)[None, ...], np.cov(X1.T)[None, ...]], axis=0 +) # shape: (k, d, d) + +mappings = [ + ot.gaussian.bures_wasserstein_mapping( + ms=m, Cs=C, mt=barycenter_mean, Ct=barycenter_cov + ) + for m, C in zip(means, covs) +] + +linear_XB, linear_YB = [], [] +for _X, _Y, (A, b) in zip(Xs, Ys, mappings): + linear_XB.append(_X.dot(A) + b) + linear_YB.append(_Y) +linear_XB = np.concatenate(linear_XB, axis=0) +linear_YB = np.concatenate(linear_YB, axis=0) + +# %% +# Compute the Empirical Wasserstein barycenter +# -------------------------------------------- +# +# Computes the Barycenter +XB, YB = joint_wasserstein_barycenter( + Xs, + Ys, + mus=None, + XB=None, + YB=None, + muB=None, + measure_weights=None, + n_samples=X0.shape[0], + reg_e=0.0, + verbose=True, +) + +# %% +# Compares the obtained barycenters +# --------------------------------- +# +# Here, as you can see, the barycenter obtained +# with the Gaussian assumption is actually just +# a translated version of the input measures. +# The empirical barycenter is actually capable +# of capturing the non-linearity of the input +# measures. +fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True) + +axes[0].scatter(XB[:, 0], XB[:, 1], c=YB.argmax(axis=1), cmap=plt.cm.coolwarm) +axes[0].set_title("Empirical Barycenter") +axes[1].scatter( + linear_XB[:, 0], linear_XB[:, 1], c=linear_YB.argmax(axis=1), cmap=plt.cm.coolwarm +) +axes[1].set_title("Bures-Wasserstein Barycenter") + +plt.show() + +# %% +# References +# ---------- +# Álvarez-Esteban, Pedro C., et al. "A fixed-point approach to barycenters in +# Wasserstein space." Journal of Mathematical Analysis and Applications 441.2 +# (2016): 744-762. diff --git a/examples/methods/plot_monge_alignment_da.py b/examples/methods/plot_monge_alignment_da.py index 7bbd534d..c2c0e549 100644 --- a/examples/methods/plot_monge_alignment_da.py +++ b/examples/methods/plot_monge_alignment_da.py @@ -120,7 +120,7 @@ clf.score(X, y, sample_domain=sample_domain, allow_source=True), ) -# %% Multisource and taregt data +# %% Multisource and target data def get_multidomain_data( diff --git a/examples/methods/plot_multi_source_da.py b/examples/methods/plot_multi_source_da.py new file mode 100644 index 00000000..afe743bc --- /dev/null +++ b/examples/methods/plot_multi_source_da.py @@ -0,0 +1,141 @@ +""" +Wasserstein Barycenter Transport +================================ + +This example illustrates the method "Wasserstein Barycenter Transport" + +""" + +# Author: Eduardo Fernandes Montesuma +# +# License: BSD 3-Clause +# sphinx_gallery_thumbnail_number = 4 + +# %% Imports +import matplotlib.pyplot as plt +import numpy as np +from sklearn.linear_model import LogisticRegression + +from skada import ( + WassersteinBarycenterTransportAdapter, +) +from skada.datasets import make_multi_source_da_example + +np.random.seed(42) + +# %% +# Generate covariate shift for multi-source DA +# -------------------------------------------- +# +# We generate a simple toy example for MSDA. This +# toy example includes a set of 4 datasets (3 sources, +# 1 target), corresponding to a rotation of a base dataset +# for the angles (0.0, 10.0, 20.0, 30.0) + +X, y, sample_domain = make_multi_source_da_example( + n_datasets=4, n_samples=500, angle_min=0.0, angle_max=30, separation=10 +) + +# Converts labels into one-hot encoded labels +Y = np.zeros((y.size, y.max() + 1)) +Y[np.arange(y.size), y] = 1 + +# %% +# +# Visualize the datasets +# ---------------------- +# +# Here we visualize the datasets + +fig, axes = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True) +for k, ax in enumerate(axes.flatten()[:-1]): + ax.scatter( + x=X[sample_domain == k, 0], + y=X[sample_domain == k, 1], + c=y[sample_domain == k], + cmap=plt.cm.coolwarm, + ) + ax.set_title(f"Source domain {k + 1}") +axes[1, 1].scatter( + x=X[sample_domain == -1, 0], + y=X[sample_domain == -1, 1], + c=y[sample_domain == -1], + cmap=plt.cm.coolwarm, +) +axes[1, 1].set_title("Target domain") +plt.tight_layout() +plt.show() + +# %% +# Fit Logistic Regression +# ----------------------- +# +# Here, we fit a classifier to the target domain +clf = LogisticRegression() +clf.fit(X[sample_domain != -1], y[sample_domain != -1]) +print( + "[Source-Only] Accuracy on the target domain:" + f" {clf.score(X[sample_domain == -1], y[sample_domain==-1])}" +) + +# %% +# Fit Wasserstein Barycenter Transport +# ------------------------------------ +# +# We fit the Wasserstein Barycenter Transport to the multi-source DA example. +# This algorithms models measures through empirical measures (mixtures of diracs) +# and is capable of handling non-linear shifts between domains. +wbt = WassersteinBarycenterTransportAdapter(n_samples=500, verbose=True) +wbt.fit(X, Y, sample_domain=sample_domain) +mapped_samples = wbt.transform(X, Y, w=None, sample_domain=sample_domain) + +# %% +# Plots loss of barycenter algorithm +# ---------------------------------- +# +# Here, we plot the loss of the empirical wasserstein barycenter +# algorithm per iteration. Note that convergence happens quite fast, +# with a few iterations. +fig, ax = plt.subplots(1, 1, figsize=(5, 5)) +ax.plot(wbt.log["barycenter_computation"]["loss_hist"]) +ax.set_xlabel("Iteration") +ax.set_ylabel("Barycenter loss") +plt.show() + +# %% +# Fit a classifier on WBT mapped data +# ----------------------------------- +# +# Here, we evaluate the performance of WBT in the target domain. +clf = LogisticRegression() +clf.fit(X=mapped_samples[-1][0], y=mapped_samples[-1][1]) +print( + "[WBT] Accuracy on the target domain:" + f" {clf.score(X[sample_domain == -1], y[sample_domain==-1])}" +) + +# %% +# +# Visualize the datasets +# ---------------------- +# +# Here we visualize the target and the transported barycenter + +XB, yB = mapped_samples[-1] +fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True) +axes[0].scatter( + x=X[sample_domain == -1, 0], + y=X[sample_domain == -1, 1], + c=y[sample_domain == -1], + cmap=plt.cm.coolwarm, +) +axes[0].set_title("Target") +axes[1].scatter( + x=XB[:, 0], + y=XB[:, 1], + c=yB, + cmap=plt.cm.coolwarm, +) +axes[1].set_title("Transported Barycenter") +plt.tight_layout() +plt.show() diff --git a/skada/__init__.py b/skada/__init__.py index 927a8b09..0e59d72c 100644 --- a/skada/__init__.py +++ b/skada/__init__.py @@ -24,7 +24,9 @@ OTMappingAdapter, OTMapping, MultiLinearMongeAlignmentAdapter, - MultiLinearMongeAlignment + MultiLinearMongeAlignment, + WassersteinBarycenterTransportAdapter, + WassersteinBarycenterTransport ) from ._reweight import ( DiscriminatorReweightAdapter, @@ -139,4 +141,8 @@ "source_target_split", "per_domain_split", + + "WassersteinBarycenterTransportAdapter", + "WassersteinBarycenterTransport" + ] diff --git a/skada/_mapping.py b/skada/_mapping.py index cb1caf90..51c1e9f0 100644 --- a/skada/_mapping.py +++ b/skada/_mapping.py @@ -8,7 +8,7 @@ from abc import abstractmethod import numpy as np -from ot import da +from ot import da, emd, sinkhorn from ot.gaussian import bures_wasserstein_barycenter, bures_wasserstein_mapping from sklearn.linear_model import LogisticRegression from sklearn.metrics.pairwise import pairwise_distances @@ -28,6 +28,168 @@ ) +def joint_wasserstein_barycenter( + Xs, + Ys, + mus=None, + XB=None, + YB=None, + muB=None, + measure_weights=None, + n_samples=None, + reg_e=0.0, + label_weight=None, + n_iter_max=100, + tol=1e-4, + verbose=False, + log=False, +): + r"""Computes the Wasserstein Barycenter [1] for a list of distributions + :math:`\mathcal{P}`, containing :math:`\hat{P}_{1}, \cdots ,\hat{P}_{K}` + and weights :math:`\alpha \in \Delta_{K}`. Each distribution is + parametrized through their support + :math:`\mathbf{X}^{( P_{k} )}, k=1, \cdots ,K`. This consists on a + implementation of the Free-Support Wasserstien Barycenter of [2]. Our + implementation relies on the fixed-point iteration of [3], + + .. math:: + \hat{B}^{(it+1)} = \psi( \hat{B}^{(it)} ), + + where :math:`\psi(\hat{P}) = T_{it,\sharp}\hat{P}`, + :math:`T_{it} = \sum_{k}\alpha_{k}T_{k,it}`, for :math:`T_{k,it}`, + the barycentric mapping between :math:`\hat{P}_{k}` and + :math:`\hat{B}^{(it)}`. + + Parameters + ---------- + Xs : List of tensors + List of tensors of shape (nk, d) with the features of the support of + each distribution Pk. + Ys : List of tensors, optional (default=None) + List of tensors of shape (nk, nc) with the labels of the support of + each distribution Pk. + XB : tensor, optional (default=None) + Tensor of shape (n, d) with the initialization for the features of + the barycenter support. + YB : tensor, optional (default=None) + Tensor of shape (n, d) with the initialization for the labels of + the barycenter support. + weights : tensor, optional (default=None) + Weight of each distribution in (XP, YP). It is a tensor of shape + (K,), whose components are all positive and it sums to one. + n_samples : int, optional (default=None) + Number of samples in the barycenter support. Only used if (XB, YB) + were not given. + reg_e : float, optional (default=0.0) + Entropic regularization. If reg_e > 0.0 uses the Sinkhorn algorithm + for computing the OT plans. + label_weight : float, optional (default=None) + Weight for the label metric. It is described as beta in the main paper. + If None is given, uses beta as the maximum pairwise distance between + samples of P and Q. + n_iter_max : int, optional (default=100) + Maximum number of iterations of the Barycenter algorithm. + n_iter_sinkhorn : int, optional (default=1000) + Maximum number of iterations of the Sinkhorn algorithm. Only used for + reg_e > 0.0. + n_iter_emd : int, optional (default=1000000) + Maximum number of iterations for Linear Programming. Only used if + reg_e = 0.0. + tol : float, optional (default=1e-4) + Tolerance for the iterations of the Wasserstein barycenter algorithm. + If a given update does not change the objective function by a value + superior to tol, the algorithm halts. + """ + assert len(Xs) == len(Ys), ( + "Expected same number of domains for" + f" features and labels, but got {len(Xs)=} and {len(Ys)=}" + ) + + n_dim = Xs[0].shape[1] + n_classes = Ys[0].shape[1] + + if mus is None: + mus = [np.ones(len(Xsk)) / len(Xsk) for Xsk in Xs] + + if n_samples is None and XB is None: + # If number of points is not provided, + # assume that the support of the barycenter + # has sum(nsi) where si is the i-th source + # domain. + n_samples = int(np.sum([len(Xs_k) for Xs_k in Xs])) + + if measure_weights is None: + measure_weights = np.ones(len(Xs)) / len(Xs) + + if XB is None: + XB = np.random.randn(n_samples, n_dim) + + if YB is None: + YB = np.random.rand(n_samples, n_classes) + YB = YB / YB.sum(axis=1)[:, None] + + if muB is None: + muB = np.ones(len(XB)) / len(XB) + + it = 0 + delta = tol + 1 + last_loss = np.inf + + if verbose: + vmessage = "|{:^25}|{:^25}|{:^25}|".format("Iteration", "Loss", "dLoss") + print("-" * len(vmessage)) + print(vmessage) + print("-" * len(vmessage)) + + if log: + extra_ret = {"loss_hist": [], "d_loss": []} + + while delta > tol and it < n_iter_max: + ground_costs, ot_plans = [], [] + + for k in range(len(Xs)): + C_k = pairwise_distances(XB, Xs[k], metric="sqeuclidean") + _lw = C_k.max() if label_weight is None else label_weight + C_k += _lw * pairwise_distances(YB, Ys[k], metric="sqeuclidean") + ground_costs.append(C_k) + if reg_e > 0.0: + plan_k = sinkhorn(muB, mus[k], C_k / C_k.max(), reg_e=reg_e) + else: + plan_k = emd(muB, mus[k], C_k) + ot_plans.append(plan_k) + + loss, _XB, _YB = 0.0, np.zeros_like(XB), np.zeros_like(YB) + for k, (Xsk, Ysk, pi_k, C_k, alpha_k) in enumerate( + zip(Xs, Ys, ot_plans, ground_costs, measure_weights) + ): + _loss_k = (C_k * plan_k).sum() + loss += alpha_k * _loss_k + _XB += alpha_k * XB.shape[0] * (pi_k @ Xsk) + _YB += alpha_k * YB.shape[0] * (pi_k @ Ysk) + XB = _XB.copy() + YB = _YB.copy() + + delta = abs(loss - last_loss) + last_loss = loss + + if verbose: + vmessage = f"|{it:^25}|{loss:^25}|{delta:^25}|" + print(vmessage) + + if log: + extra_ret["loss_hist"].append(loss) + extra_ret["d_loss"].append(delta) + + it += 1 + if verbose: + print("-" * len(vmessage)) + + if log: + extra_ret["transport_plans"] = ot_plans + return XB, YB, extra_ret + return XB, YB + + class BaseOTMappingAdapter(BaseAdapter): """Base class for all DA estimators implemented using OT mapping. @@ -554,7 +716,7 @@ class MultiLinearMongeAlignmentAdapter(BaseAdapter): The method is a simplified extension of [29] using the Bures-Wasserstein distance and mapping of [7] to align multiple source domains to a - barycenter. The sued of barycenter alignment with gaussien assumption was + barycenter. The sued of barycenter alignment with gaussian assumption was proposed in [30]. Parameters @@ -752,6 +914,284 @@ def MultiLinearMongeAlignment( ) +class WassersteinBarycenterTransportAdapter(BaseAdapter): + """Maps the source domain data to the target domain data through + the Wasserstein barycenter of source domains. This class performs + a 2-step adaptation strategy proposed in [29]_, by first computing + the Wasserstein barycenter of empirical source domain measures using + the algorithm of [39]_, then applying the Barycentric mapping of [6]_ + + Parameters + ---------- + reg_e : float, default=0.0 + Entropic regularization parameter for the Sinkhorn algorithm. + n_samples : int, optional + Number of samples to use for the barycenter computation. + label_weight : float, optional + Weight for the label regularization term in the barycenter computation. + n_iter_max : int, default=100 + Maximum number of iterations for the barycenter computation. + tol : float, default=1e-4 + Tolerance for the convergence of the barycenter computation. + verbose : bool, default=False + If True, print progress messages during computation. + use_labels_target : bool, default=False + If True, use target labels during the mapping to the target domain. + + Attributes + ---------- + source_domains : dict + A dictionary containing the source domain data, labels, and weights. + target_domains : dict + A dictionary containing the target domain data, labels, and weights. + transport_plans : dict + A dictionary containing the optimal transport plans for each source domain. + barycenter_ : dict + A dictionary containing the features and labels of the computed barycenter. + mappings : dict + A dictionary containing the EMDTransport objects for mapping source domains + to the barycenter. + mapping_target : dict + A dictionary containing the EMDTransport or SinkhornTransport objects for + mapping the barycenter to target domains. + log : dict + A dictionary containing logs from the barycenter computation and + target mappings. + + References + ---------- + .. [6] N. Courty, R. Flamary, D. Tuia and A. Rakotomamonjy, + Optimal Transport for Domain Adaptation, in IEEE + Transactions on Pattern Analysis and Machine Intelligence + + .. [29] Montesuma, Eduardo Fernandes, and Fred Maurice Ngole Mboula. + "Wasserstein barycenter for multi-source domain adaptation." In Proceedings + of the IEEE/CVF conference on computer vision and pattern recognition, pp. + 16785-16793. 2021. + + .. [39] Montesuma, Eduardo, Fred Maurice Ngole Mboula, and Antoine Souloumiac. + "Multi-source domain adaptation through dataset dictionary learning in + wasserstein space." ECAI 2023. IOS Press, 2023. 1739-1746. + """ + + def __init__( + self, + reg_e=0.0, + n_samples=None, + label_weight=None, + n_iter_max=100, + tol=1e-4, + verbose=False, + use_labels_target=False, + ): + super().__init__() + self.reg_e = reg_e + self.n_samples = n_samples + self.label_weight = label_weight + self.n_iter_max = n_iter_max + self.tol = tol + self.verbose = verbose + self.use_labels_target = use_labels_target + self.log = {} + + def fit(self, X, y=None, w=None, *, sample_domain=None): + """Fit adaptation parameters. + + Parameters + ---------- + X : array-like, shape (n_samples, n_features) + The source data. + y : array-like, shape (n_samples,) + The source labels. + w : array-like, shape (n_samples,) + The source sample importances + sample_domain : array-like, shape (n_samples,) + The domain labels (same as sample_domain). + + Returns + ------- + self : object + Returns self. + """ + X, sample_domain = check_X_domain(X, sample_domain) + self.source_domains, self.target_domains = per_domain_split( + X, y, w, sample_domain=sample_domain + ) + + Xs = [ + self.source_domains[domain_index][0] for domain_index in self.source_domains + ] + Ys = [ + self.source_domains[domain_index][1] for domain_index in self.source_domains + ] + mus = [ + self.source_domains[domain_index][2] for domain_index in self.source_domains + ] + if any([mu is None for mu in mus]): + mus = None + + XB, YB, log = joint_wasserstein_barycenter( + Xs=Xs, + Ys=Ys, + mus=mus, + measure_weights=None, + n_samples=self.n_samples, + reg_e=self.reg_e, + label_weight=self.label_weight, + n_iter_max=self.n_iter_max, + tol=self.tol, + verbose=self.verbose, + log=True, + ) + self.log["barycenter_computation"] = log + + self.transport_plans = { + domain_index: log["transport_plans"][i] + for i, domain_index in enumerate(self.source_domains) + } + + self.barycenter_ = {"features": XB, "labels": YB} + + self.mappings = {} + for i, domain_index in enumerate(self.source_domains): + self.mappings[domain_index] = da.EMDTransport() + self.mappings[domain_index].coupling_ = self.transport_plans[i] + self.mappings[domain_index].mu_s = self.source_domains[domain_index][2] + self.mappings[domain_index].xs_ = self.source_domains[domain_index][0] + self.mappings[domain_index].xt_ = self.barycenter_["features"] + + self.mapping_target = { + domain: ( + da.EMDTransport(log=True).fit( + Xs=XB, + ys=YB.argmax(axis=1), + Xt=self.target_domains[domain][0], + yt=self.target_domains[domain][1] + if self.use_labels_target + else None, + ) + if self.reg_e == 0.0 + else da.SinkhornTransport( + Xs=XB, + ys=YB.argmax(axis=1), + Xt=self.target_domains[domain][0], + yt=self.target_domains[domain][1] + if self.use_labels_target + else None, + reg_e=self.reg_e, + norm="max", + log=True, + ) + ) + for domain in self.target_domains + } + self.log["mapping_targets"] = { + domain: self.mapping_target[domain].log_ for domain in self.mapping_target + } + + return self + + def fit_transform(self, X, y=None, sample_domain=None, **params): + """Predict adaptation (weights, sample or labels). + + Parameters + ---------- + X : array-like, shape (n_samples, n_features) + The source data. + y : array-like, shape (n_samples,) + The source labels. + sample_domain : array-like, shape (n_samples,) + The domain labels (same as sample_domain). + + Returns + ------- + X_t : array-like, shape (n_samples, n_components) + The data (same as X). + """ + self.fit(X, y, sample_domain=sample_domain) + return self.transform(X, sample_domain=sample_domain, allow_source=True) + + def transform( + self, X, y=None, w=None, *, sample_domain=None, allow_source=False, **params + ) -> np.ndarray: + source_domains, target_domains = per_domain_split( + X, y, w, sample_domain=sample_domain + ) + + # Checks if the arrays on each domain are the same + new_source = not any( + np.array_equal(self.source_domains[domain][0], source_domains[domain][0]) + for domain in self.source_domains + if domain in source_domains + ) + + # NOTE: Contrary to MultiLinearMongeAlignment and GaussianMixtureMultiAlignment, + # WassersteinBarycenterTransport works on empirical measures. This means + # that the mapping is only defined on the support of the original measures + # it was trained on. We can, however, extend this mapping to new samples, + # through for instance what is called the "Ferradans mapping" in [7]. + if not new_source: + # If all arrays are the same as the ones used for training, we + # don't need to recompute OT. We simply map the barycenter to + # the target. + return { + domain: ( + self.mapping_target[domain].transform( + Xs=self.barycenter_["features"], + ys=self.barycenter_["labels"].argmax(axis=1), + Xt=target_domains[domain], + yt=target_domains[domain], + ), + self.barycenter_["labels"].argmax(axis=1), + ) + for domain in self.target_domains + } + else: + # Otherwise, we re-estimate the barycenter support using the new + # provided samples. This new support is obtained through Ferradans + # mappings, for instance. + est_XB = 0.0 + for domain in source_domains: + est_XB += self.mappings[domain].transform( + Xs=source_domains[domain][0] + ) / len(source_domains) + + # We then map the estimated barycenter support to the target domain + return { + domain: self.mapping_target[domain].transform( + Xs=est_XB, ys=None, Xt=target_domains[domain] + ) + for domain in target_domains + } + + +def WassersteinBarycenterTransport( + base_estimator=None, + reg_e=0.0, + n_samples=None, + label_weight=None, + n_iter_max=100, + tol=1e-4, + verbose=False, + use_labels_target=False, +): + if base_estimator is None: + base_estimator = LogisticRegression() + + return make_da_pipeline( + WassersteinBarycenterTransportAdapter( + reg_e=reg_e, + n_samples=n_samples, + label_weight=label_weight, + n_iter_max=n_iter_max, + tol=tol, + verbose=verbose, + use_labels_target=use_labels_target, + ), + base_estimator, + ) + + def _sqrtm(C): r"""Square root of SPD matrices. diff --git a/skada/datasets/__init__.py b/skada/datasets/__init__.py index dbfe6514..b7e5170c 100644 --- a/skada/datasets/__init__.py +++ b/skada/datasets/__init__.py @@ -37,6 +37,8 @@ make_shifted_datasets, make_dataset_from_moons_distribution, make_variable_frequency_dataset, + make_multi_source_da_example, + make_classification_dataset ) from ._mnist_usps import load_mnist_usps @@ -59,4 +61,6 @@ 'fetch_amazon_review_all', 'fetch_office_home', 'fetch_office_home_all', + 'make_multi_source_da_example', + 'make_classification_dataset' ] diff --git a/skada/datasets/_samples_generator.py b/skada/datasets/_samples_generator.py index da7493c6..05a15aec 100644 --- a/skada/datasets/_samples_generator.py +++ b/skada/datasets/_samples_generator.py @@ -848,3 +848,59 @@ def make_variable_frequency_dataset( return dataset else: return dataset.pack(as_sources=["s"], as_targets=["t"], return_X_y=return_X_y) + + +def make_classification_dataset(mean, cov, v=None, separation=1, n=200): + """Creates a 2D linearly separable dataset with two classes.""" + x1 = np.random.multivariate_normal(mean, cov, size=n) + if v is None: + v = np.random.randn( + 2, + ) + v = separation * (v / np.linalg.norm(v)).reshape(1, -1) + elif np.linalg.norm(v) != separation: + v = separation * (v / np.linalg.norm(v)).reshape(1, -1) + else: + v = v.reshape(1, -1) + x2 = x1 + v + X = np.concatenate([x1, x2], axis=0) + y = np.array([0] * len(x1) + [1] * len(x2)) + + return X, y + + +def make_multi_source_da_example( + n_datasets, n_samples=400, angle_min=0.0, angle_max=45, separation=6 +): + mu = np.array([0.0, 0.0]) + angles = np.linspace(angle_min, angle_max, n_datasets) + Xs, ys, samples_domain = [], [], [] + for i in range(n_datasets - 1): + A = np.random.randn(2, 2) + cov = 0.25 * np.dot(A.T, A) + np.eye(2) + v = np.array( + [np.cos((np.pi / 180) * angles[i]), np.sin((np.pi / 180) * angles[i])] + ) + X, y = make_classification_dataset( + mu, cov, v=v, separation=separation, n=n_samples // 2 + ) + Xs.append(X) + ys.append(y) + samples_domain.append(np.ones_like(y) * i) + + A = np.random.randn(2, 2) + mu = np.array([5.0, 5.0]) + cov = 0.1 * np.dot(A.T, A) + np.eye(2) + v = np.array( + [np.cos((np.pi / 180) * angles[-1]), np.sin((np.pi / 180) * angles[-1])] + ) + Xt, yt = make_classification_dataset( + mu, cov, v=v, separation=separation, n=n_samples + ) + samples_domain.append(np.ones_like(yt) * -1) + + return ( + np.concatenate([*Xs, Xt], axis=0), + np.concatenate([*ys, yt], axis=0), + np.concatenate(samples_domain, axis=0), + )