diff options
| -rw-r--r-- | ot/bregman.py | 11 | ||||
| -rw-r--r-- | test/test_bregman.py | 8 |
2 files changed, 13 insertions, 6 deletions
diff --git a/ot/bregman.py b/ot/bregman.py index 7acfcf1..dc43834 100644 --- a/ot/bregman.py +++ b/ot/bregman.py @@ -1587,8 +1587,13 @@ def empirical_sinkhorn_divergence(X_s, X_t, reg, a=None, b=None, metric='sqeucli log['log_sinkhorn_b'] = log_b return max(0, sinkhorn_div), log + else: - sinkhorn_div = (empirical_sinkhorn2(X_s, X_t, reg, a, b, metric=metric, numIterMax=numIterMax, stopThr=1e-9, verbose=verbose, log=log, **kwargs) - - 1 / 2 * empirical_sinkhorn2(X_s, X_s, reg, a, b, metric=metric, numIterMax=numIterMax, stopThr=1e-9, verbose=verbose, log=log, **kwargs) - - 1 / 2 * empirical_sinkhorn2(X_t, X_t, reg, a, b, metric=metric, numIterMax=numIterMax, stopThr=1e-9, verbose=verbose, log=log, **kwargs)) + sinkhorn_loss_ab = empirical_sinkhorn2(X_s, X_t, reg, a, b, metric=metric, numIterMax=numIterMax, stopThr=1e-9, verbose=verbose, log=log, **kwargs) + + sinkhorn_loss_a = empirical_sinkhorn2(X_s, X_s, reg, a, b, metric=metric, numIterMax=numIterMax, stopThr=1e-9, verbose=verbose, log=log, **kwargs) + + sinkhorn_loss_b = empirical_sinkhorn2(X_t, X_t, reg, a, b, metric=metric, numIterMax=numIterMax, stopThr=1e-9, verbose=verbose, log=log, **kwargs) + + sinkhorn_div = sinkhorn_loss_ab - 1 / 2 * (sinkhorn_loss_a + sinkhorn_loss_b) return max(0, sinkhorn_div) diff --git a/test/test_bregman.py b/test/test_bregman.py index 68d3595..58700e2 100644 --- a/test/test_bregman.py +++ b/test/test_bregman.py @@ -241,11 +241,13 @@ def test_empirical_sinkhorn_divergence(): M_t = ot.dist(X_t, X_t) emp_sinkhorn_div = ot.bregman.empirical_sinkhorn_divergence(X_s, X_t, 1) - sinkhorn_div = (2 * ot.sinkhorn2(a, b, M, 1) - ot.sinkhorn2(a, a, M_s, 1) - - ot.sinkhorn2(b, b, M_t, 1)) + 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, 0.1, log=True) - sinkhorn_div_log, log_s = ot.sinkhorn(a, b, M, 0.1, log=True) + sink_div_log, log_s = ot.sinkhorn2(a, b, M, 1) + sink_div_log_a, log_s_a = ot.sinkhorn2(a, a, M_s, 1) + sink_div_log_b, log_s_b = ot.sinkhorn2(b, b, M_t, 1) + sink_div_log = sink_div_log - 1 / 2 * (sink_div_log_a + sink_div_log_b) # check constratints np.testing.assert_allclose( |