def test_sinkhorn2_variants_multi_b(nx): # test sinkhorn n = 50 rng = np.random.RandomState(0) x = rng.randn(n, 2) u = ot.utils.unif(n) b = rng.rand(n, 3) b = b / np.sum(b, 0, keepdims=True) M = ot.dist(x, x) ub, bb, M_nx = nx.from_numpy(u, b, M) G = ot.sinkhorn2(u, b, M, 1, method='sinkhorn', stopThr=1e-10) Gl = nx.to_numpy( ot.sinkhorn2(ub, bb, M_nx, 1, method='sinkhorn_log', stopThr=1e-10)) G0 = nx.to_numpy( ot.sinkhorn2(ub, bb, M_nx, 1, method='sinkhorn', stopThr=1e-10)) Gs = nx.to_numpy( ot.sinkhorn2(ub, bb, M_nx, 1, method='sinkhorn_stabilized', stopThr=1e-10)) # check values np.testing.assert_allclose(G, G0, atol=1e-05) np.testing.assert_allclose(G, Gl, atol=1e-05) np.testing.assert_allclose(G0, Gs, atol=1e-05)
def test_empirical_sinkhorn_divergence(): # Test sinkhorn divergence n = 10 a = np.linspace(1, n, n) a /= a.sum() b = ot.unif(n) X_s = np.reshape(np.arange(n), (n, 1)) X_t = np.reshape(np.arange(0, n * 2, 2), (n, 1)) M = ot.dist(X_s, X_t) M_s = ot.dist(X_s, X_s) M_t = ot.dist(X_t, X_t) emp_sinkhorn_div = ot.bregman.empirical_sinkhorn_divergence(X_s, X_t, 1, a=a, b=b) sinkhorn_div = (ot.sinkhorn2(a, b, M, 1) - 1 / 2 * ot.sinkhorn2(a, a, M_s, 1) - 1 / 2 * ot.sinkhorn2(b, b, M_t, 1)) emp_sinkhorn_div_log, log_es = ot.bregman.empirical_sinkhorn_divergence(X_s, X_t, 1, a=a, b=b, log=True) sink_div_log_ab, log_s_ab = ot.sinkhorn2(a, b, M, 1, log=True) sink_div_log_a, log_s_a = ot.sinkhorn2(a, a, M_s, 1, log=True) sink_div_log_b, log_s_b = ot.sinkhorn2(b, b, M_t, 1, log=True) sink_div_log = sink_div_log_ab - 1 / 2 * (sink_div_log_a + sink_div_log_b) # check constraints np.testing.assert_allclose( emp_sinkhorn_div, sinkhorn_div, atol=1e-05) # cf conv emp sinkhorn np.testing.assert_allclose( emp_sinkhorn_div_log, sink_div_log, atol=1e-05) # cf conv emp sinkhorn
def test_not_implemented_method(): # test sinkhorn w = 10 n = w**2 rng = np.random.RandomState(42) A_img = rng.rand(2, w, w) A_flat = A_img.reshape(n, 2) a1, a2 = A_flat.T M_flat = ot.utils.dist0(n) not_implemented = "new_method" reg = 0.01 with pytest.raises(ValueError): ot.sinkhorn(a1, a2, M_flat, reg, method=not_implemented) with pytest.raises(ValueError): ot.sinkhorn2(a1, a2, M_flat, reg, method=not_implemented) with pytest.raises(ValueError): ot.barycenter(A_flat, M_flat, reg, method=not_implemented) with pytest.raises(ValueError): ot.bregman.barycenter_debiased(A_flat, M_flat, reg, method=not_implemented) with pytest.raises(ValueError): ot.bregman.convolutional_barycenter2d(A_img, reg, method=not_implemented) with pytest.raises(ValueError): ot.bregman.convolutional_barycenter2d_debiased(A_img, reg, method=not_implemented)
def test_sinkhorn_stabilization(): # test sinkhorn n = 100 a1 = ot.datasets.make_1D_gauss(n, m=30, s=10) a2 = ot.datasets.make_1D_gauss(n, m=40, s=10) M = ot.utils.dist0(n) reg = 1e-5 loss1 = ot.sinkhorn2(a1, a2, M, reg, method="sinkhorn_log") loss2 = ot.sinkhorn2(a1, a2, M, reg, tau=1, method="sinkhorn_stabilized") np.testing.assert_allclose(loss1, loss2, atol=1e-06) # cf convergence sinkhorn
def coupling_W2(coupling_1, coupling_2, source, target, epsilon): """ Returns the entropically-regularized W2 distance between two couplings """ cost_matrix = coupling_to_coupling_cost_matrix(source, target) return ot.sinkhorn2(coupling_1.flatten(), coupling_2.flatten(), cost_matrix, epsilon)
def test_lazy_empirical_sinkhorn(): # test sinkhorn n = 10 a = ot.unif(n) b = ot.unif(n) numIterMax = 1000 X_s = np.reshape(np.arange(n), (n, 1)) X_t = np.reshape(np.arange(0, n), (n, 1)) M = ot.dist(X_s, X_t) M_m = ot.dist(X_s, X_t, metric='minkowski') f, g = ot.bregman.empirical_sinkhorn(X_s, X_t, 1, numIterMax=numIterMax, isLazy=True, batchSize=(1, 3), verbose=True) G_sqe = np.exp(f[:, None] + g[None, :] - M / 1) sinkhorn_sqe = ot.sinkhorn(a, b, M, 1) f, g, log_es = ot.bregman.empirical_sinkhorn(X_s, X_t, 0.1, numIterMax=numIterMax, isLazy=True, batchSize=1, log=True) G_log = np.exp(f[:, None] + g[None, :] - M / 0.1) sinkhorn_log, log_s = ot.sinkhorn(a, b, M, 0.1, log=True) f, g = ot.bregman.empirical_sinkhorn(X_s, X_t, 1, metric='minkowski', numIterMax=numIterMax, isLazy=True, batchSize=1) G_m = np.exp(f[:, None] + g[None, :] - M_m / 1) sinkhorn_m = ot.sinkhorn(a, b, M_m, 1) loss_emp_sinkhorn, log = ot.bregman.empirical_sinkhorn2( X_s, X_t, 1, numIterMax=numIterMax, isLazy=True, batchSize=1, log=True) loss_sinkhorn = ot.sinkhorn2(a, b, M, 1) # check constratints np.testing.assert_allclose(sinkhorn_sqe.sum(1), G_sqe.sum(1), atol=1e-05) # metric sqeuclidian np.testing.assert_allclose(sinkhorn_sqe.sum(0), G_sqe.sum(0), atol=1e-05) # metric sqeuclidian np.testing.assert_allclose(sinkhorn_log.sum(1), G_log.sum(1), atol=1e-05) # log np.testing.assert_allclose(sinkhorn_log.sum(0), G_log.sum(0), atol=1e-05) # log np.testing.assert_allclose(sinkhorn_m.sum(1), G_m.sum(1), atol=1e-05) # metric euclidian np.testing.assert_allclose(sinkhorn_m.sum(0), G_m.sum(0), atol=1e-05) # metric euclidian np.testing.assert_allclose(loss_emp_sinkhorn, loss_sinkhorn, atol=1e-05)
def test_sinkhorn_multi_b(method, verbose, warn): # test sinkhorn n = 10 rng = np.random.RandomState(0) x = rng.randn(n, 2) u = ot.utils.unif(n) b = rng.rand(n, 3) b = b / np.sum(b, 0, keepdims=True) M = ot.dist(x, x) loss0, log = ot.sinkhorn(u, b, M, .1, method=method, stopThr=1e-10, log=True) loss = [ ot.sinkhorn2(u, b[:, k], M, .1, method=method, stopThr=1e-10, verbose=verbose, warn=warn) for k in range(3) ] # check constraints np.testing.assert_allclose(loss0, loss, atol=1e-4) # cf convergence sinkhorn
def _edge_curvature( edge, measures, geodesic_distances, measure_cutoff=1e-6, sinkhorn_regularisation=0, weighted_curvature=False, ): """Compute curvature for an edge.""" node_x, node_y = edge m_x, m_y = measures[node_x], measures[node_y] Nx = np.where(m_x >= measure_cutoff * np.max(m_x))[0] Ny = np.where(m_y >= measure_cutoff * np.max(m_y))[0] m_x, m_y = m_x[Nx], m_y[Ny] m_x /= m_x.sum() m_y /= m_y.sum() distances_xy = geodesic_distances[np.ix_(Nx, Ny)] if sinkhorn_regularisation > 0: wasserstein_distance = ot.sinkhorn2(m_x, m_y, distances_xy, sinkhorn_regularisation)[0] else: wasserstein_distance = ot.emd2(m_x, m_y, distances_xy) if weighted_curvature: return geodesic_distances[node_x, node_y] - wasserstein_distance return 1.0 - wasserstein_distance / geodesic_distances[node_x, node_y]
def test_convergence_warning(method): # test sinkhorn n = 100 a1 = ot.datasets.make_1D_gauss(n, m=30, s=10) a2 = ot.datasets.make_1D_gauss(n, m=40, s=10) A = np.asarray([a1, a2]).T M = ot.utils.dist0(n) with pytest.warns(UserWarning): ot.sinkhorn(a1, a2, M, 1., method=method, stopThr=0, numItermax=1) if method in ["sinkhorn", "sinkhorn_stabilized", "sinkhorn_log"]: with pytest.warns(UserWarning): ot.barycenter(A, M, 1, method=method, stopThr=0, numItermax=1) with pytest.warns(UserWarning): ot.sinkhorn2(a1, a2, M, 1, method=method, stopThr=0, numItermax=1)
def _sinkhorn_distance(x, y, d): """Compute the approximate optimal transportation distance (Sinkhorn distance) of the given density distributions. Parameters ---------- x : (m,) numpy.ndarray Source's density distributions, includes source and source's neighbors. y : (n,) numpy.ndarray Target's density distributions, includes source and source's neighbors. d : (m, n) numpy.ndarray Shortest path matrix. Returns ------- m : float Sinkhorn distance, an approximate optimal transportation distance. """ t0 = time.time() m = ot.sinkhorn2(x, y, d, 1e-1, method='sinkhorn')[0] logger.debug( "%8f secs for Sinkhorn. dist. \t#source_nbr: %d, #target_nbr: %d" % (time.time() - t0, len(x), len(y))) return m
def ot_loss(mapped, target, device='cpu'): reg = 5 nx = mapped.shape[0] ny = target.shape[0] ab = torch.ones(nx) / nx ab = ab.to(device) M = ot.dist(mapped, target) # euclidean by default loss = ot.sinkhorn2(ab, ab, M, reg) return loss
def sinkhorn(X, Y, options): """sinkhorn distance(regularized OT)""" D = pdist2(X, Y, options.metric) N, M = np.shape(X)[0], np.shape(Y)[0] a, b = np.ones(N) / N, np.ones(M) / M dist = ot.sinkhorn2(a, b, D, options.regularize)[0] T = ot.sinkhorn(a, b, D, 0.1) return dist, T
def sinkhorn(X, Y, options): """Sinkhorn Distance Regularized OT""" device = 'cuda' if options.cuda else 'cpu' D = pdist2(X, Y, options) N, M = X.size()[0], Y.size()[0] a, b = torch.ones(N, dtype=torch.float64).to(device) / N, torch.ones( M, dtype=torch.float64).to(device) / M dist = ot.sinkhorn2(a, b, D, options.regularize)[0] T = ot.sinkhorn(a, b, D, 0.1) return dist, T
def test_empirical_sinkhorn_divergence(nx): # Test sinkhorn divergence n = 10 a = np.linspace(1, n, n) a /= a.sum() b = ot.unif(n) X_s = np.reshape(np.arange(n, dtype=np.float64), (n, 1)) X_t = np.reshape(np.arange(0, n * 2, 2, dtype=np.float64), (n, 1)) M = ot.dist(X_s, X_t) M_s = ot.dist(X_s, X_s) M_t = ot.dist(X_t, X_t) ab, bb, X_sb, X_tb, M_nx, M_sb, M_tb = nx.from_numpy( a, b, X_s, X_t, M, M_s, M_t) emp_sinkhorn_div = nx.to_numpy( ot.bregman.empirical_sinkhorn_divergence(X_sb, X_tb, 1, a=ab, b=bb)) sinkhorn_div = nx.to_numpy( ot.sinkhorn2(ab, bb, M_nx, 1) - 1 / 2 * ot.sinkhorn2(ab, ab, M_sb, 1) - 1 / 2 * ot.sinkhorn2(bb, bb, M_tb, 1)) emp_sinkhorn_div_np = ot.bregman.empirical_sinkhorn_divergence(X_s, X_t, 1, a=a, b=b) # check constraints np.testing.assert_allclose(emp_sinkhorn_div, emp_sinkhorn_div_np, atol=1e-05) np.testing.assert_allclose(emp_sinkhorn_div, sinkhorn_div, atol=1e-05) # cf conv emp sinkhorn ot.bregman.empirical_sinkhorn_divergence(X_sb, X_tb, 1, a=ab, b=bb, log=True)
def check_gradient(): (loss, log) = ot.sinkhorn2(a, b, M, lambd, log=True) subgrad = lambd * np.log(log["u"]) subgrad = subgrad.reshape((subgrad.shape[0], )) subgrad -= np.mean(subgrad) eps = 1e-4 grad = np.zeros(a.shape[0]) for i in range(a.shape[0]): direction = np.zeros(a.shape[0]) direction[i] = 1 direction -= np.mean(direction) grad[i] = (ot.sinkhorn2(a + eps * direction, b, M, lambd) - loss) / eps print(a.shape) print(subgrad.shape) loss2 = ot.sinkhorn2(a - 0.1 * subgrad, b, M, lambd) # disturbed_subgrad = subgrad + 0.01 * np.random.normal(size = (subgrad.shape[0],)) # disturbed_subgrad -= np.mean(disturbed_subgrad) # loss3 = ot.sinkhorn2(a - 0.01 * disturbed_subgrad, b, M, lambd) print(loss) print(loss2) # print(loss3) return (grad, subgrad)
def _sinkhorn_distance(self, x, y, d): """ Compute the approximate optimal transportation distance (Sinkhorn distance) of the given density distributions. :param x: Source's neighbors distributions :param y: Target's neighbors distributions :param d: Cost dictionary :return: Sinkhorn distance """ t0 = time.time() m = ot.sinkhorn2(x, y, d, 1e-1, method='sinkhorn')[0] logger.debug( "%8f secs for Sinkhorn. dist. \t#source_nbr: %d, #target_nbr: %d" % (time.time() - t0, len(x), len(y))) return m
def test_sinkhorn2_variants_device_tf(method): nx = ot.backend.TensorflowBackend() n = 100 x = np.random.randn(n, 2) u = ot.utils.unif(n) M = ot.dist(x, x) # Check that everything stays on the CPU with tf.device("/CPU:0"): ub, Mb = nx.from_numpy(u, M) Gb = ot.sinkhorn(ub, ub, Mb, 1, method=method, stopThr=1e-10) lossb = ot.sinkhorn2(ub, ub, Mb, 1, method=method, stopThr=1e-10) nx.assert_same_dtype_device(Mb, Gb) nx.assert_same_dtype_device(Mb, lossb) if len(tf.config.list_physical_devices('GPU')) > 0: # Check that everything happens on the GPU ub, Mb = nx.from_numpy(u, M) Gb = ot.sinkhorn(ub, ub, Mb, 1, method=method, stopThr=1e-10) lossb = ot.sinkhorn2(ub, ub, Mb, 1, method=method, stopThr=1e-10) nx.assert_same_dtype_device(Mb, Gb) nx.assert_same_dtype_device(Mb, lossb) assert nx.dtype_device(Gb)[1].startswith("GPU")
def main(): na = 100 nb = 150 reg = 0.5 mu_s = np.array([0, 0]) cov_s = np.array([[1, 0], [0, 1]]) mu_t = np.array([4, 4]) cov_t = np.array([[1, -.8], [-.8, 1]]) x_tf = tf.placeholder(dtype=tf.float32, shape=[na, 2]) y_tf = tf.placeholder(dtype=tf.float32, shape=[nb, 2]) M_tf = dmat_tf(x_tf, y_tf) tf_sinkhorn_loss = sink_tf(M_tf, (na, nb), reg) print("I can compute the gradient for a", tf.gradients(tf_sinkhorn_loss, x_tf)) print("I can compute the gradient for b", tf.gradients(tf_sinkhorn_loss, y_tf)) sess = tf.InteractiveSession() sess.run(tf.global_variables_initializer()) xs = ot.datasets.make_2D_samples_gauss(na, mu_s, cov_s) xt = ot.datasets.make_2D_samples_gauss(nb, mu_t, cov_t) # Visualization plt.figure(1) plt.plot(xs[:, 0], xs[:, 1], '+b', label='Source samples') plt.plot(xt[:, 0], xt[:, 1], 'xr', label='Target samples') plt.legend(loc=0) plt.title('Source and target distributions') plt.show() # TF - sinkhorn tf_sinkhorn_loss_val = sess.run(tf_sinkhorn_loss, feed_dict={ x_tf: xs, y_tf: xt }) print(' tf_sinkhorn_loss', tf_sinkhorn_loss_val) # POT - sinkhorn M = ot.dist(xs.copy(), xt.copy(), metric='euclidean') a = np.ones((na, )) / na b = np.ones((nb, )) / nb # uniform distribution on samples pot_sinkhorn_loss = ot.sinkhorn2(a, b, M, reg)[0] print('pot_sinkhorn_loss', pot_sinkhorn_loss)
def min_a(b, M, lambd): n = M.shape[0] a_tilde = np.ones(n) / n a_hat = a_tilde.copy() a_hat_old = a_hat.copy() converged = False t = 0.5 tol = 1e-5 its = 0 gamma = 1e-2 while not converged: # its += 1 # print("Iteration: {}".format(its)) # beta = (t+1)/2 # a = (1 - 1/beta)*a_hat + 1/beta * a_tilde # (_, log) = ot.sinkhorn2(a, b, M, lambd, log = True) # u = log["u"].reshape(n) # alpha = lambd * np.log(u) # alpha -= np.mean(alpha) # # a_tilde *= u**(-t * beta * lambd) # a_tilde *= np.exp(-t * beta * alpha) # a_tilde /= np.sum(a_tilde) # a_hat_old = a_hat.copy() # a_hat = (1 - 1/beta) * a_hat + 1/beta * a_tilde # t += 1 its += 1 print("Iteration: {}".format(its)) beta = (t + 1) / 2 a = (1 - 1 / beta) * a_hat + 1 / beta * a_tilde (_, log) = ot.sinkhorn2(a, b, M, lambd, log=True) u = log["u"].reshape(n) alpha = lambd * np.log(u) alpha -= np.mean(alpha) # a_tilde *= u**(-t * beta * lambd) a_hat_old = a_hat.copy() a_hat -= gamma * alpha a_hat[a_hat < 0] = 0 a_hat /= np.sum(a_hat) print(a_hat_old) print(a_hat) if np.linalg.norm(a_hat - a_hat_old) < tol: converged = True return a_hat
def test_sinkhorn2_variants_dtype_device(nx, method): n = 100 x = np.random.randn(n, 2) u = ot.utils.unif(n) M = ot.dist(x, x) for tp in nx.__type_list__: print(nx.dtype_device(tp)) ub, Mb = nx.from_numpy(u, M, type_as=tp) lossb = ot.sinkhorn2(ub, ub, Mb, 1, method=method, stopThr=1e-10) nx.assert_same_dtype_device(Mb, lossb)
def quick_e_score(self, n1, n2): """Pass in an AvgLink and return negative average distance.""" if n1.needs_update: n1._update() if n2.needs_update: n2._update() # rows i want 1 is num_samples by dim rows_i_want_1 = n1.mat # rows i want 2 is num_samples by dim rows_i_want_2 = n2.mat # compute the point cloud wasserstein distance between the normalized # distributions. M = cdist(rows_i_want_1, rows_i_want_2) a = np.ones(rows_i_want_1.shape[0]) / rows_i_want_1.shape[0] b = np.ones(rows_i_want_2.shape[0]) / rows_i_want_2.shape[0] dist = ot.sinkhorn2(a,b,M,self.sinkhorn_reg,method='sinkhorn_stabilized',numItermax=self.max_sinkhorn_iter) return -dist[0]
def _wmd(self, i, row, X_train): union_idx = np.union1d(X_train[i].indices, row.indices) W_minimal = self.W_embed[union_idx] W_dist = euclidean_distances(W_minimal) bow_i = X_train[i, union_idx].A.ravel() bow_j = row[:, union_idx].A.ravel() if self.sinkhorn: return ot.sinkhorn2( bow_i, bow_j, W_dist, self.sinkhorn_reg, numItermax=50, method="sinkhorn_stabilized", )[0] else: return ot.emd2(bow_i, bow_j, W_dist)
def forward(self, x, y): self.save_for_backward(x, y) loss = torch.zeros(1).cuda() for i in range(x.shape[0]): a = x[i].cpu().numpy() a[a <= 0] = 1e-9 a = a / np.sum(a) b = y[i].cpu() b = nn.Softmax()(b.view(1, -1)).data.numpy().reshape((-1, )) b[b <= 0] = 1e-9 b = b / np.sum(b) dis, log_i = ot.sinkhorn2(a, b, self.m, self.reg, log=True) self.log[i] = torch.FloatTensor(log_i['u']) loss += dis[0] return loss
def test_sinkhorn2_backends(nx): n_samples = 100 n_features = 2 rng = np.random.RandomState(0) x = rng.randn(n_samples, n_features) y = rng.randn(n_samples, n_features) a = ot.utils.unif(n_samples) M = ot.dist(x, y) G = ot.sinkhorn(a, a, M, 1) ab, M_nx = nx.from_numpy(a, M) Gb = ot.sinkhorn2(ab, ab, M_nx, 1) np.allclose(G, nx.to_numpy(Gb))
def test_empirical_sinkhorn(nx): # test sinkhorn n = 10 a = ot.unif(n) b = ot.unif(n) X_s = np.reshape(1.0 * np.arange(n), (n, 1)) X_t = np.reshape(1.0 * np.arange(0, n), (n, 1)) M = ot.dist(X_s, X_t) M_m = ot.dist(X_s, X_t, metric='euclidean') ab, bb, X_sb, X_tb, M_nx, M_mb = nx.from_numpy(a, b, X_s, X_t, M, M_m) G_sqe = nx.to_numpy(ot.bregman.empirical_sinkhorn(X_sb, X_tb, 1)) sinkhorn_sqe = nx.to_numpy(ot.sinkhorn(ab, bb, M_nx, 1)) G_log, log_es = ot.bregman.empirical_sinkhorn(X_sb, X_tb, 0.1, log=True) G_log = nx.to_numpy(G_log) sinkhorn_log, log_s = ot.sinkhorn(ab, bb, M_nx, 0.1, log=True) sinkhorn_log = nx.to_numpy(sinkhorn_log) G_m = nx.to_numpy( ot.bregman.empirical_sinkhorn(X_sb, X_tb, 1, metric='euclidean')) sinkhorn_m = nx.to_numpy(ot.sinkhorn(ab, bb, M_mb, 1)) loss_emp_sinkhorn = nx.to_numpy( ot.bregman.empirical_sinkhorn2(X_sb, X_tb, 1)) loss_sinkhorn = nx.to_numpy(ot.sinkhorn2(ab, bb, M_nx, 1)) # check constraints np.testing.assert_allclose(sinkhorn_sqe.sum(1), G_sqe.sum(1), atol=1e-05) # metric sqeuclidian np.testing.assert_allclose(sinkhorn_sqe.sum(0), G_sqe.sum(0), atol=1e-05) # metric sqeuclidian np.testing.assert_allclose(sinkhorn_log.sum(1), G_log.sum(1), atol=1e-05) # log np.testing.assert_allclose(sinkhorn_log.sum(0), G_log.sum(0), atol=1e-05) # log np.testing.assert_allclose(sinkhorn_m.sum(1), G_m.sum(1), atol=1e-05) # metric euclidian np.testing.assert_allclose(sinkhorn_m.sum(0), G_m.sum(0), atol=1e-05) # metric euclidian np.testing.assert_allclose(loss_emp_sinkhorn, loss_sinkhorn, atol=1e-05)
def ws(p, q, x, dist='euclidean'): """ Wasserstein distance btwn p and q Parameters ----- p : distribution q : distribution x : support of p and q p, q, and x must be of shape (n_samples, dimension) Returns ----- Wasserstein distance (scalar) """ reg = 1e-2 M = ot.dist(x, x, metric=dist) M /= M.max() ws = ot.sinkhorn2(p, q, M, reg) return ws[0]
def forward(self, x, y): self.save_for_backward(x, y) loss = torch.zeros(1).cuda() for i in range(x.shape[0]): a = x[i].cpu().numpy().reshape((-1, )) a[a <= 0] = 1e-9 a = a / np.sum(a) #print(y[i]) b = y[i].cpu().numpy().reshape((-1, )) b[b <= 0] = 1e-9 b = b / np.sum(b) dis, log_i = ot.sinkhorn2(a, b, self.m, self.reg, log=True, method='sinkhorn_stabilized') self.log[i] = torch.FloatTensor(log_i['logu'][:, 0]) loss += dis[0] return loss
def test_empirical_sinkhorn(): # test sinkhorn n = 100 a = ot.unif(n) b = ot.unif(n) X_s = np.reshape(np.arange(n), (n, 1)) X_t = np.reshape(np.arange(0, n), (n, 1)) M = ot.dist(X_s, X_t) M_m = ot.dist(X_s, X_t, metric='minkowski') G_sqe = ot.bregman.empirical_sinkhorn(X_s, X_t, 1) sinkhorn_sqe = ot.sinkhorn(a, b, M, 1) G_log, log_es = ot.bregman.empirical_sinkhorn(X_s, X_t, 0.1, log=True) sinkhorn_log, log_s = ot.sinkhorn(a, b, M, 0.1, log=True) G_m = ot.bregman.empirical_sinkhorn(X_s, X_t, 1, metric='minkowski') sinkhorn_m = ot.sinkhorn(a, b, M_m, 1) loss_emp_sinkhorn = ot.bregman.empirical_sinkhorn2(X_s, X_t, 1) loss_sinkhorn = ot.sinkhorn2(a, b, M, 1) # check constratints np.testing.assert_allclose(sinkhorn_sqe.sum(1), G_sqe.sum(1), atol=1e-05) # metric sqeuclidian np.testing.assert_allclose(sinkhorn_sqe.sum(0), G_sqe.sum(0), atol=1e-05) # metric sqeuclidian np.testing.assert_allclose(sinkhorn_log.sum(1), G_log.sum(1), atol=1e-05) # log np.testing.assert_allclose(sinkhorn_log.sum(0), G_log.sum(0), atol=1e-05) # log np.testing.assert_allclose(sinkhorn_m.sum(1), G_m.sum(1), atol=1e-05) # metric euclidian np.testing.assert_allclose(sinkhorn_m.sum(0), G_m.sum(0), atol=1e-05) # metric euclidian np.testing.assert_allclose(loss_emp_sinkhorn, loss_sinkhorn, atol=1e-05)
def reg_wasserstein(root_1: Root, root_2: Root, reg=0.5, method="sinkhorn_stabilized", numItermax=1000): """Regularized Wasserstein distance, computed with Sinkhorn algorithm.""" sq_dist_matrix = distance_matrix(root_1.nodes, root_2.nodes)**2 rescaled_nodes_1 = root_1.nodes_weights / root_1.nodes_weights.sum() rescaled_nodes_2 = root_2.nodes_weights / root_2.nodes_weights.sum() coupling = ot.sinkhorn(a=rescaled_nodes_1, b=rescaled_nodes_2, M=sq_dist_matrix, reg=reg, method=method, numItermax=numItermax) cost = ot.sinkhorn2(a=rescaled_nodes_1, b=rescaled_nodes_2, M=sq_dist_matrix, reg=reg, method=method, numItermax=numItermax) return np.sqrt(cost), coupling
def test_sinkhorn2_gradients(): n_samples = 100 n_features = 2 rng = np.random.RandomState(0) x = rng.randn(n_samples, n_features) y = rng.randn(n_samples, n_features) a = ot.utils.unif(n_samples) M = ot.dist(x, y) if torch: a1 = torch.tensor(a, requires_grad=True) b1 = torch.tensor(a, requires_grad=True) M1 = torch.tensor(M, requires_grad=True) val = ot.sinkhorn2(a1, b1, M1, 1) val.backward() assert a1.shape == a1.grad.shape assert b1.shape == b1.grad.shape assert M1.shape == M1.grad.shape