diff options
Diffstat (limited to 'examples')
-rw-r--r-- | examples/plot_OT_2D_samples.py | 26 | ||||
-rw-r--r-- | examples/plot_UOT_1D.py | 76 | ||||
-rw-r--r-- | examples/plot_UOT_barycenter_1D.py | 164 | ||||
-rw-r--r-- | examples/plot_barycenter_fgw.py | 184 | ||||
-rw-r--r-- | examples/plot_barycenter_lp_vs_entropic.py | 7 | ||||
-rw-r--r-- | examples/plot_fgw.py | 173 | ||||
-rw-r--r-- | examples/plot_free_support_barycenter.py | 2 | ||||
-rw-r--r-- | examples/plot_otda_color_images.py | 8 | ||||
-rw-r--r-- | examples/plot_otda_mapping_colors_images.py | 2 | ||||
-rw-r--r-- | examples/plot_screenkhorn_1D.py | 68 |
10 files changed, 703 insertions, 7 deletions
diff --git a/examples/plot_OT_2D_samples.py b/examples/plot_OT_2D_samples.py index bb952a0..63126ba 100644 --- a/examples/plot_OT_2D_samples.py +++ b/examples/plot_OT_2D_samples.py @@ -10,6 +10,7 @@ sum of diracs. The OT matrix is plotted with the samples. """ # Author: Remi Flamary <remi.flamary@unice.fr> +# Kilian Fatras <kilian.fatras@irisa.fr> # # License: MIT License @@ -100,3 +101,28 @@ pl.legend(loc=0) pl.title('OT matrix Sinkhorn with samples') pl.show() + + +############################################################################## +# Emprirical Sinkhorn +# ---------------- + +#%% sinkhorn + +# reg term +lambd = 1e-3 + +Ges = ot.bregman.empirical_sinkhorn(xs, xt, lambd) + +pl.figure(7) +pl.imshow(Ges, interpolation='nearest') +pl.title('OT matrix empirical sinkhorn') + +pl.figure(8) +ot.plot.plot2D_samples_mat(xs, xt, Ges, color=[.5, .5, 1]) +pl.plot(xs[:, 0], xs[:, 1], '+b', label='Source samples') +pl.plot(xt[:, 0], xt[:, 1], 'xr', label='Target samples') +pl.legend(loc=0) +pl.title('OT matrix Sinkhorn from samples') + +pl.show() diff --git a/examples/plot_UOT_1D.py b/examples/plot_UOT_1D.py new file mode 100644 index 0000000..2ea8b05 --- /dev/null +++ b/examples/plot_UOT_1D.py @@ -0,0 +1,76 @@ +# -*- coding: utf-8 -*- +""" +=============================== +1D Unbalanced optimal transport +=============================== + +This example illustrates the computation of Unbalanced Optimal transport +using a Kullback-Leibler relaxation. +""" + +# Author: Hicham Janati <hicham.janati@inria.fr> +# +# License: MIT License + +import numpy as np +import matplotlib.pylab as pl +import ot +import ot.plot +from ot.datasets import make_1D_gauss as gauss + +############################################################################## +# Generate data +# ------------- + + +#%% parameters + +n = 100 # nb bins + +# bin positions +x = np.arange(n, dtype=np.float64) + +# Gaussian distributions +a = gauss(n, m=20, s=5) # m= mean, s= std +b = gauss(n, m=60, s=10) + +# make distributions unbalanced +b *= 5. + +# loss matrix +M = ot.dist(x.reshape((n, 1)), x.reshape((n, 1))) +M /= M.max() + + +############################################################################## +# Plot distributions and loss matrix +# ---------------------------------- + +#%% plot the distributions + +pl.figure(1, figsize=(6.4, 3)) +pl.plot(x, a, 'b', label='Source distribution') +pl.plot(x, b, 'r', label='Target distribution') +pl.legend() + +# plot distributions and loss matrix + +pl.figure(2, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, M, 'Cost matrix M') + + +############################################################################## +# Solve Unbalanced Sinkhorn +# -------------- + + +# Sinkhorn + +epsilon = 0.1 # entropy parameter +alpha = 1. # Unbalanced KL relaxation parameter +Gs = ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha, verbose=True) + +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, Gs, 'UOT matrix Sinkhorn') + +pl.show() diff --git a/examples/plot_UOT_barycenter_1D.py b/examples/plot_UOT_barycenter_1D.py new file mode 100644 index 0000000..c8d9d3b --- /dev/null +++ b/examples/plot_UOT_barycenter_1D.py @@ -0,0 +1,164 @@ +# -*- coding: utf-8 -*- +""" +=========================================================== +1D Wasserstein barycenter demo for Unbalanced distributions +=========================================================== + +This example illustrates the computation of regularized Wassersyein Barycenter +as proposed in [10] for Unbalanced inputs. + + +[10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816. + +""" + +# Author: Hicham Janati <hicham.janati@inria.fr> +# +# License: MIT License + +import numpy as np +import matplotlib.pylab as pl +import ot +# necessary for 3d plot even if not used +from mpl_toolkits.mplot3d import Axes3D # noqa +from matplotlib.collections import PolyCollection + +############################################################################## +# Generate data +# ------------- + +# parameters + +n = 100 # nb bins + +# bin positions +x = np.arange(n, dtype=np.float64) + +# Gaussian distributions +a1 = ot.datasets.make_1D_gauss(n, m=20, s=5) # m= mean, s= std +a2 = ot.datasets.make_1D_gauss(n, m=60, s=8) + +# make unbalanced dists +a2 *= 3. + +# creating matrix A containing all distributions +A = np.vstack((a1, a2)).T +n_distributions = A.shape[1] + +# loss matrix + normalization +M = ot.utils.dist0(n) +M /= M.max() + +############################################################################## +# Plot data +# --------- + +# plot the distributions + +pl.figure(1, figsize=(6.4, 3)) +for i in range(n_distributions): + pl.plot(x, A[:, i]) +pl.title('Distributions') +pl.tight_layout() + +############################################################################## +# Barycenter computation +# ---------------------- + +# non weighted barycenter computation + +weight = 0.5 # 0<=weight<=1 +weights = np.array([1 - weight, weight]) + +# l2bary +bary_l2 = A.dot(weights) + +# wasserstein +reg = 1e-3 +alpha = 1. + +bary_wass = ot.unbalanced.barycenter_unbalanced(A, M, reg, alpha, weights) + +pl.figure(2) +pl.clf() +pl.subplot(2, 1, 1) +for i in range(n_distributions): + pl.plot(x, A[:, i]) +pl.title('Distributions') + +pl.subplot(2, 1, 2) +pl.plot(x, bary_l2, 'r', label='l2') +pl.plot(x, bary_wass, 'g', label='Wasserstein') +pl.legend() +pl.title('Barycenters') +pl.tight_layout() + +############################################################################## +# Barycentric interpolation +# ------------------------- + +# barycenter interpolation + +n_weight = 11 +weight_list = np.linspace(0, 1, n_weight) + + +B_l2 = np.zeros((n, n_weight)) + +B_wass = np.copy(B_l2) + +for i in range(0, n_weight): + weight = weight_list[i] + weights = np.array([1 - weight, weight]) + B_l2[:, i] = A.dot(weights) + B_wass[:, i] = ot.unbalanced.barycenter_unbalanced(A, M, reg, alpha, weights) + + +# plot interpolation + +pl.figure(3) + +cmap = pl.cm.get_cmap('viridis') +verts = [] +zs = weight_list +for i, z in enumerate(zs): + ys = B_l2[:, i] + verts.append(list(zip(x, ys))) + +ax = pl.gcf().gca(projection='3d') + +poly = PolyCollection(verts, facecolors=[cmap(a) for a in weight_list]) +poly.set_alpha(0.7) +ax.add_collection3d(poly, zs=zs, zdir='y') +ax.set_xlabel('x') +ax.set_xlim3d(0, n) +ax.set_ylabel(r'$\alpha$') +ax.set_ylim3d(0, 1) +ax.set_zlabel('') +ax.set_zlim3d(0, B_l2.max() * 1.01) +pl.title('Barycenter interpolation with l2') +pl.tight_layout() + +pl.figure(4) +cmap = pl.cm.get_cmap('viridis') +verts = [] +zs = weight_list +for i, z in enumerate(zs): + ys = B_wass[:, i] + verts.append(list(zip(x, ys))) + +ax = pl.gcf().gca(projection='3d') + +poly = PolyCollection(verts, facecolors=[cmap(a) for a in weight_list]) +poly.set_alpha(0.7) +ax.add_collection3d(poly, zs=zs, zdir='y') +ax.set_xlabel('x') +ax.set_xlim3d(0, n) +ax.set_ylabel(r'$\alpha$') +ax.set_ylim3d(0, 1) +ax.set_zlabel('') +ax.set_zlim3d(0, B_l2.max() * 1.01) +pl.title('Barycenter interpolation with Wasserstein') +pl.tight_layout() + +pl.show() diff --git a/examples/plot_barycenter_fgw.py b/examples/plot_barycenter_fgw.py new file mode 100644 index 0000000..77b0370 --- /dev/null +++ b/examples/plot_barycenter_fgw.py @@ -0,0 +1,184 @@ +# -*- coding: utf-8 -*- +""" +================================= +Plot graphs' barycenter using FGW +================================= + +This example illustrates the computation barycenter of labeled graphs using FGW + +Requires networkx >=2 + +.. [18] Vayer Titouan, Chapel Laetitia, Flamary R{\'e}mi, Tavenard Romain + and Courty Nicolas + "Optimal Transport for structured data with application on graphs" + International Conference on Machine Learning (ICML). 2019. + +""" + +# Author: Titouan Vayer <titouan.vayer@irisa.fr> +# +# License: MIT License + +#%% load libraries +import numpy as np +import matplotlib.pyplot as plt +import networkx as nx +import math +from scipy.sparse.csgraph import shortest_path +import matplotlib.colors as mcol +from matplotlib import cm +from ot.gromov import fgw_barycenters +#%% Graph functions + + +def find_thresh(C, inf=0.5, sup=3, step=10): + """ Trick to find the adequate thresholds from where value of the C matrix are considered close enough to say that nodes are connected + Tthe threshold is found by a linesearch between values "inf" and "sup" with "step" thresholds tested. + The optimal threshold is the one which minimizes the reconstruction error between the shortest_path matrix coming from the thresholded adjency matrix + and the original matrix. + Parameters + ---------- + C : ndarray, shape (n_nodes,n_nodes) + The structure matrix to threshold + inf : float + The beginning of the linesearch + sup : float + The end of the linesearch + step : integer + Number of thresholds tested + """ + dist = [] + search = np.linspace(inf, sup, step) + for thresh in search: + Cprime = sp_to_adjency(C, 0, thresh) + SC = shortest_path(Cprime, method='D') + SC[SC == float('inf')] = 100 + dist.append(np.linalg.norm(SC - C)) + return search[np.argmin(dist)], dist + + +def sp_to_adjency(C, threshinf=0.2, threshsup=1.8): + """ Thresholds the structure matrix in order to compute an adjency matrix. + All values between threshinf and threshsup are considered representing connected nodes and set to 1. Else are set to 0 + Parameters + ---------- + C : ndarray, shape (n_nodes,n_nodes) + The structure matrix to threshold + threshinf : float + The minimum value of distance from which the new value is set to 1 + threshsup : float + The maximum value of distance from which the new value is set to 1 + Returns + ------- + C : ndarray, shape (n_nodes,n_nodes) + The threshold matrix. Each element is in {0,1} + """ + H = np.zeros_like(C) + np.fill_diagonal(H, np.diagonal(C)) + C = C - H + C = np.minimum(np.maximum(C, threshinf), threshsup) + C[C == threshsup] = 0 + C[C != 0] = 1 + + return C + + +def build_noisy_circular_graph(N=20, mu=0, sigma=0.3, with_noise=False, structure_noise=False, p=None): + """ Create a noisy circular graph + """ + g = nx.Graph() + g.add_nodes_from(list(range(N))) + for i in range(N): + noise = float(np.random.normal(mu, sigma, 1)) + if with_noise: + g.add_node(i, attr_name=math.sin((2 * i * math.pi / N)) + noise) + else: + g.add_node(i, attr_name=math.sin(2 * i * math.pi / N)) + g.add_edge(i, i + 1) + if structure_noise: + randomint = np.random.randint(0, p) + if randomint == 0: + if i <= N - 3: + g.add_edge(i, i + 2) + if i == N - 2: + g.add_edge(i, 0) + if i == N - 1: + g.add_edge(i, 1) + g.add_edge(N, 0) + noise = float(np.random.normal(mu, sigma, 1)) + if with_noise: + g.add_node(N, attr_name=math.sin((2 * N * math.pi / N)) + noise) + else: + g.add_node(N, attr_name=math.sin(2 * N * math.pi / N)) + return g + + +def graph_colors(nx_graph, vmin=0, vmax=7): + cnorm = mcol.Normalize(vmin=vmin, vmax=vmax) + cpick = cm.ScalarMappable(norm=cnorm, cmap='viridis') + cpick.set_array([]) + val_map = {} + for k, v in nx.get_node_attributes(nx_graph, 'attr_name').items(): + val_map[k] = cpick.to_rgba(v) + colors = [] + for node in nx_graph.nodes(): + colors.append(val_map[node]) + return colors + +############################################################################## +# Generate data +# ------------- + +#%% circular dataset +# We build a dataset of noisy circular graphs. +# Noise is added on the structures by random connections and on the features by gaussian noise. + + +np.random.seed(30) +X0 = [] +for k in range(9): + X0.append(build_noisy_circular_graph(np.random.randint(15, 25), with_noise=True, structure_noise=True, p=3)) + +############################################################################## +# Plot data +# --------- + +#%% Plot graphs + +plt.figure(figsize=(8, 10)) +for i in range(len(X0)): + plt.subplot(3, 3, i + 1) + g = X0[i] + pos = nx.kamada_kawai_layout(g) + nx.draw(g, pos=pos, node_color=graph_colors(g, vmin=-1, vmax=1), with_labels=False, node_size=100) +plt.suptitle('Dataset of noisy graphs. Color indicates the label', fontsize=20) +plt.show() + +############################################################################## +# Barycenter computation +# ---------------------- + +#%% We compute the barycenter using FGW. Structure matrices are computed using the shortest_path distance in the graph +# Features distances are the euclidean distances +Cs = [shortest_path(nx.adjacency_matrix(x)) for x in X0] +ps = [np.ones(len(x.nodes())) / len(x.nodes()) for x in X0] +Ys = [np.array([v for (k, v) in nx.get_node_attributes(x, 'attr_name').items()]).reshape(-1, 1) for x in X0] +lambdas = np.array([np.ones(len(Ys)) / len(Ys)]).ravel() +sizebary = 15 # we choose a barycenter with 15 nodes + +A, C, log = fgw_barycenters(sizebary, Ys, Cs, ps, lambdas, alpha=0.95, log=True) + +############################################################################## +# Plot Barycenter +# ------------------------- + +#%% Create the barycenter +bary = nx.from_numpy_matrix(sp_to_adjency(C, threshinf=0, threshsup=find_thresh(C, sup=100, step=100)[0])) +for i, v in enumerate(A.ravel()): + bary.add_node(i, attr_name=v) + +#%% +pos = nx.kamada_kawai_layout(bary) +nx.draw(bary, pos=pos, node_color=graph_colors(bary, vmin=-1, vmax=1), with_labels=False) +plt.suptitle('Barycenter', fontsize=20) +plt.show() diff --git a/examples/plot_barycenter_lp_vs_entropic.py b/examples/plot_barycenter_lp_vs_entropic.py index b82765e..d7c72d0 100644 --- a/examples/plot_barycenter_lp_vs_entropic.py +++ b/examples/plot_barycenter_lp_vs_entropic.py @@ -102,7 +102,7 @@ pl.tight_layout() problems.append([A, [bary_l2, bary_wass, bary_wass2]]) ############################################################################## -# Dirac Data +# Stair Data # ---------- #%% parameters @@ -168,6 +168,11 @@ pl.legend() pl.title('Barycenters') pl.tight_layout() + +############################################################################## +# Dirac Data +# ---------- + #%% parameters a1 = np.zeros(n) diff --git a/examples/plot_fgw.py b/examples/plot_fgw.py new file mode 100644 index 0000000..43efc94 --- /dev/null +++ b/examples/plot_fgw.py @@ -0,0 +1,173 @@ +# -*- coding: utf-8 -*- +""" +============================== +Plot Fused-gromov-Wasserstein +============================== + +This example illustrates the computation of FGW for 1D measures[18]. + +.. [18] Vayer Titouan, Chapel Laetitia, Flamary R{\'e}mi, Tavenard Romain + and Courty Nicolas + "Optimal Transport for structured data with application on graphs" + International Conference on Machine Learning (ICML). 2019. + +""" + +# Author: Titouan Vayer <titouan.vayer@irisa.fr> +# +# License: MIT License + +import matplotlib.pyplot as pl +import numpy as np +import ot +from ot.gromov import gromov_wasserstein, fused_gromov_wasserstein + +############################################################################## +# Generate data +# --------- + +#%% parameters +# We create two 1D random measures +n = 20 # number of points in the first distribution +n2 = 30 # number of points in the second distribution +sig = 1 # std of first distribution +sig2 = 0.1 # std of second distribution + +np.random.seed(0) + +phi = np.arange(n)[:, None] +xs = phi + sig * np.random.randn(n, 1) +ys = np.vstack((np.ones((n // 2, 1)), 0 * np.ones((n // 2, 1)))) + sig2 * np.random.randn(n, 1) + +phi2 = np.arange(n2)[:, None] +xt = phi2 + sig * np.random.randn(n2, 1) +yt = np.vstack((np.ones((n2 // 2, 1)), 0 * np.ones((n2 // 2, 1)))) + sig2 * np.random.randn(n2, 1) +yt = yt[::-1, :] + +p = ot.unif(n) +q = ot.unif(n2) + +############################################################################## +# Plot data +# --------- + +#%% plot the distributions + +pl.close(10) +pl.figure(10, (7, 7)) + +pl.subplot(2, 1, 1) + +pl.scatter(ys, xs, c=phi, s=70) +pl.ylabel('Feature value a', fontsize=20) +pl.title('$\mu=\sum_i \delta_{x_i,a_i}$', fontsize=25, usetex=True, y=1) +pl.xticks(()) +pl.yticks(()) +pl.subplot(2, 1, 2) +pl.scatter(yt, xt, c=phi2, s=70) +pl.xlabel('coordinates x/y', fontsize=25) +pl.ylabel('Feature value b', fontsize=20) +pl.title('$\\nu=\sum_j \delta_{y_j,b_j}$', fontsize=25, usetex=True, y=1) +pl.yticks(()) +pl.tight_layout() +pl.show() + +############################################################################## +# Create structure matrices and across-feature distance matrix +# --------- + +#%% Structure matrices and across-features distance matrix +C1 = ot.dist(xs) +C2 = ot.dist(xt) +M = ot.dist(ys, yt) +w1 = ot.unif(C1.shape[0]) +w2 = ot.unif(C2.shape[0]) +Got = ot.emd([], [], M) + +############################################################################## +# Plot matrices +# --------- + +#%% +cmap = 'Reds' +pl.close(10) +pl.figure(10, (5, 5)) +fs = 15 +l_x = [0, 5, 10, 15] +l_y = [0, 5, 10, 15, 20, 25] +gs = pl.GridSpec(5, 5) + +ax1 = pl.subplot(gs[3:, :2]) + +pl.imshow(C1, cmap=cmap, interpolation='nearest') +pl.title("$C_1$", fontsize=fs) +pl.xlabel("$k$", fontsize=fs) +pl.ylabel("$i$", fontsize=fs) +pl.xticks(l_x) +pl.yticks(l_x) + +ax2 = pl.subplot(gs[:3, 2:]) + +pl.imshow(C2, cmap=cmap, interpolation='nearest') +pl.title("$C_2$", fontsize=fs) +pl.ylabel("$l$", fontsize=fs) +#pl.ylabel("$l$",fontsize=fs) +pl.xticks(()) +pl.yticks(l_y) +ax2.set_aspect('auto') + +ax3 = pl.subplot(gs[3:, 2:], sharex=ax2, sharey=ax1) +pl.imshow(M, cmap=cmap, interpolation='nearest') +pl.yticks(l_x) +pl.xticks(l_y) +pl.ylabel("$i$", fontsize=fs) +pl.title("$M_{AB}$", fontsize=fs) +pl.xlabel("$j$", fontsize=fs) +pl.tight_layout() +ax3.set_aspect('auto') +pl.show() + +############################################################################## +# Compute FGW/GW +# --------- + +#%% Computing FGW and GW +alpha = 1e-3 + +ot.tic() +Gwg, logw = fused_gromov_wasserstein(M, C1, C2, p, q, loss_fun='square_loss', alpha=alpha, verbose=True, log=True) +ot.toc() + +#%reload_ext WGW +Gg, log = gromov_wasserstein(C1, C2, p, q, loss_fun='square_loss', verbose=True, log=True) + +############################################################################## +# Visualize transport matrices +# --------- + +#%% visu OT matrix +cmap = 'Blues' +fs = 15 +pl.figure(2, (13, 5)) +pl.clf() +pl.subplot(1, 3, 1) +pl.imshow(Got, cmap=cmap, interpolation='nearest') +#pl.xlabel("$y$",fontsize=fs) +pl.ylabel("$i$", fontsize=fs) +pl.xticks(()) + +pl.title('Wasserstein ($M$ only)') + +pl.subplot(1, 3, 2) +pl.imshow(Gg, cmap=cmap, interpolation='nearest') +pl.title('Gromov ($C_1,C_2$ only)') +pl.xticks(()) +pl.subplot(1, 3, 3) +pl.imshow(Gwg, cmap=cmap, interpolation='nearest') +pl.title('FGW ($M+C_1,C_2$)') + +pl.xlabel("$j$", fontsize=fs) +pl.ylabel("$i$", fontsize=fs) + +pl.tight_layout() +pl.show() diff --git a/examples/plot_free_support_barycenter.py b/examples/plot_free_support_barycenter.py index b6efc59..64b89e4 100644 --- a/examples/plot_free_support_barycenter.py +++ b/examples/plot_free_support_barycenter.py @@ -62,7 +62,7 @@ X = ot.lp.free_support_barycenter(measures_locations, measures_weights, X_init, pl.figure(1) for (x_i, b_i) in zip(measures_locations, measures_weights): color = np.random.randint(low=1, high=10 * N) - pl.scatter(x_i[:, 0], x_i[:, 1], s=b * 1000, label='input measure') + pl.scatter(x_i[:, 0], x_i[:, 1], s=b_i * 1000, label='input measure') pl.scatter(X[:, 0], X[:, 1], s=b * 1000, c='black', marker='^', label='2-Wasserstein barycenter') pl.title('Data measures and their barycenter') pl.legend(loc=0) diff --git a/examples/plot_otda_color_images.py b/examples/plot_otda_color_images.py index e77aec0..62383a2 100644 --- a/examples/plot_otda_color_images.py +++ b/examples/plot_otda_color_images.py @@ -4,7 +4,7 @@ OT for image color adaptation ============================= -This example presents a way of transferring colors between two image +This example presents a way of transferring colors between two images with Optimal Transport as introduced in [6] [6] Ferradans, S., Papadakis, N., Peyre, G., & Aujol, J. F. (2014). @@ -27,7 +27,7 @@ r = np.random.RandomState(42) def im2mat(I): - """Converts and image to matrix (one pixel per line)""" + """Converts an image to matrix (one pixel per line)""" return I.reshape((I.shape[0] * I.shape[1], I.shape[2])) @@ -115,8 +115,8 @@ ot_sinkhorn.fit(Xs=Xs, Xt=Xt) transp_Xs_emd = ot_emd.transform(Xs=X1) transp_Xt_emd = ot_emd.inverse_transform(Xt=X2) -transp_Xs_sinkhorn = ot_emd.transform(Xs=X1) -transp_Xt_sinkhorn = ot_emd.inverse_transform(Xt=X2) +transp_Xs_sinkhorn = ot_sinkhorn.transform(Xs=X1) +transp_Xt_sinkhorn = ot_sinkhorn.inverse_transform(Xt=X2) I1t = minmax(mat2im(transp_Xs_emd, I1.shape)) I2t = minmax(mat2im(transp_Xt_emd, I2.shape)) diff --git a/examples/plot_otda_mapping_colors_images.py b/examples/plot_otda_mapping_colors_images.py index 5f1e844..a20eca8 100644 --- a/examples/plot_otda_mapping_colors_images.py +++ b/examples/plot_otda_mapping_colors_images.py @@ -77,7 +77,7 @@ Image_emd = minmax(mat2im(transp_Xs_emd, I1.shape)) # SinkhornTransport ot_sinkhorn = ot.da.SinkhornTransport(reg_e=1e-1) ot_sinkhorn.fit(Xs=Xs, Xt=Xt) -transp_Xs_sinkhorn = ot_emd.transform(Xs=X1) +transp_Xs_sinkhorn = ot_sinkhorn.transform(Xs=X1) Image_sinkhorn = minmax(mat2im(transp_Xs_sinkhorn, I1.shape)) ot_mapping_linear = ot.da.MappingTransport( diff --git a/examples/plot_screenkhorn_1D.py b/examples/plot_screenkhorn_1D.py new file mode 100644 index 0000000..840ead8 --- /dev/null +++ b/examples/plot_screenkhorn_1D.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +""" +=============================== +1D Screened optimal transport +=============================== + +This example illustrates the computation of Screenkhorn: +Screening Sinkhorn Algorithm for Optimal transport. +""" + +# Author: Mokhtar Z. Alaya <mokhtarzahdi.alaya@gmail.com> +# +# License: MIT License + +import numpy as np +import matplotlib.pylab as pl +import ot.plot +from ot.datasets import make_1D_gauss as gauss +from ot.bregman import screenkhorn + +############################################################################## +# Generate data +# ------------- + +#%% parameters + +n = 100 # nb bins + +# bin positions +x = np.arange(n, dtype=np.float64) + +# Gaussian distributions +a = gauss(n, m=20, s=5) # m= mean, s= std +b = gauss(n, m=60, s=10) + +# loss matrix +M = ot.dist(x.reshape((n, 1)), x.reshape((n, 1))) +M /= M.max() + +############################################################################## +# Plot distributions and loss matrix +# ---------------------------------- + +#%% plot the distributions + +pl.figure(1, figsize=(6.4, 3)) +pl.plot(x, a, 'b', label='Source distribution') +pl.plot(x, b, 'r', label='Target distribution') +pl.legend() + +# plot distributions and loss matrix + +pl.figure(2, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, M, 'Cost matrix M') + +############################################################################## +# Solve Screenkhorn +# ----------------------- + +# Screenkhorn +lambd = 2e-03 # entropy parameter +ns_budget = 30 # budget number of points to be keeped in the source distribution +nt_budget = 30 # budget number of points to be keeped in the target distribution + +G_screen = screenkhorn(a, b, M, lambd, ns_budget, nt_budget, uniform=False, restricted=True, verbose=True) +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, G_screen, 'OT matrix Screenkhorn') +pl.show() |