Skip to content

[WIP] Add new features to unbalanced solvers #551

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rflamary merged 20 commits into from
Nov 6, 2023
Merged

[WIP] Add new features to unbalanced solvers #551

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
e903625
add new features to unbalanced solvers
6Ulm Oct 31, 2023
ff3b057
add new features to unbalanced solvers
6Ulm Oct 31, 2023
0912026
fix bug in test
6Ulm Oct 31, 2023
4fa2195
remove stab_sinkhorn
6Ulm Oct 31, 2023
a288be0
remove kl
6Ulm Oct 31, 2023
e5c3f84
fix bug in lbfgsb_unbalanced
6Ulm Oct 31, 2023
3787cdc
fix bug in lbfgsb_unbalanced
6Ulm Oct 31, 2023
810dfe1
fix bug in KL in sinkhorn_unbalanced
6Ulm Nov 1, 2023
6883c82
fix diverge to origina/master
6Ulm Nov 1, 2023
48b6729
edit release.md
6Ulm Nov 1, 2023
2ad05ab
fix test
6Ulm Nov 1, 2023
2ccb2aa
add test and rearrange arguments
6Ulm Nov 2, 2023
51d6e24
fix test
6Ulm Nov 2, 2023
60cd0d7
fix test
6Ulm Nov 2, 2023
c0d8391
fix test
6Ulm Nov 2, 2023
4db194a
fix bug in test
6Ulm Nov 2, 2023
e8298d0
fix bug in doctest
6Ulm Nov 2, 2023
1bccd47
fix bug in doctest
6Ulm Nov 2, 2023
c2c0e96
add test for more coverage
6Ulm Nov 2, 2023
d03ff06
fix conflict
6Ulm Nov 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
fix bug in test
  • Loading branch information
@6Ulm
6Ulm committed Nov 2, 2023
commit 4db194a865ae8363240f08094a5ab7d330da4989
69 changes: 45 additions & 24 deletions ot/unbalanced.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,30 +277,55 @@ def sinkhorn_unbalanced2(a, b, M, reg, reg_m, method='sinkhorn',
ot.unbalanced.sinkhorn_reg_scaling: Unbalanced Sinkhorn with epsilon scaling :ref:`[9, 10] <references-sinkhorn-unbalanced2>`

"""
b = list_to_array(b)
M, a, b = list_to_array(M, a, b)
nx = get_backend(M, a, b)

if len(b.shape) < 2:
b = b[:, None]
if method.lower() == 'sinkhorn':
res = sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)

elif method.lower() == 'sinkhorn_stabilized':
res = sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)
elif method.lower() in ['sinkhorn_reg_scaling']:
warnings.warn('Method not implemented yet. Using classic Sinkhorn-Knopp')
res = sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)
else:
raise ValueError('Unknown method %s.' % method)

if method.lower() == 'sinkhorn':
return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)
if log:
return nx.sum(M * res[0]), res[1]
else:
return nx.sum(M * res)

elif method.lower() == 'sinkhorn_stabilized':
return sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr,
verbose=verbose,
log=log, **kwargs)
elif method.lower() in ['sinkhorn_reg_scaling']:
warnings.warn('Method not implemented yet. Using classic Sinkhorn-Knopp')
return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)
else:
raise ValueError('Unknown method %s.' % method)
if method.lower() == 'sinkhorn':
return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)

elif method.lower() == 'sinkhorn_stabilized':
return sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)
elif method.lower() in ['sinkhorn_reg_scaling']:
warnings.warn('Method not implemented yet. Using classic Sinkhorn-Knopp')
return sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type,
warmstart, numItermax=numItermax,
stopThr=stopThr, verbose=verbose,
log=log, **kwargs)
else:
raise ValueError('Unknown method %s.' % method)


def sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
Expand Down Expand Up @@ -443,8 +468,6 @@ def sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
v = nx.ones(dim_b, type_as=M)
else:
u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1])
Comment thread
6Ulm marked this conversation as resolved.
if not n_hists:
u, v = u.reshape(-1), v.reshape(-1)

if reg_type == "kl":
K = nx.exp(-M / reg) * a.reshape(-1)[:, None] * b.reshape(-1)[None, :]
Expand Down Expand Up @@ -654,8 +677,6 @@ def sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
v = nx.ones(dim_b, type_as=M)
else:
u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1])
if not n_hists:
u, v = u.reshape(-1), v.reshape(-1)

if reg_type == "kl":
log_ab = nx.log(a + 1e-16).reshape(-1)[:, None] + nx.log(b + 1e-16).reshape(-1)[None, :]
Expand Down
12 changes: 5 additions & 7 deletions test/test_unbalanced.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@


@pytest.mark.parametrize("method,reg_type", itertools.product(["sinkhorn", "sinkhorn_stabilized"], ["kl", "entropy"]))
# @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_stabilized"])
def test_unbalanced_convergence(nx, method, reg_type):
# test generalized sinkhorn for unbalanced OT
n = 100
Expand Down Expand Up @@ -80,7 +79,6 @@ def test_unbalanced_convergence(nx, method, reg_type):


@pytest.mark.parametrize("method,reg_type", itertools.product(["sinkhorn", "sinkhorn_stabilized"], ["kl", "entropy"]))
# @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_stabilized"])
def test_unbalanced_warmstart(nx, method, reg_type):
# test generalized sinkhorn for unbalanced OT
n = 100
Expand Down Expand Up @@ -115,8 +113,8 @@ def test_unbalanced_warmstart(nx, method, reg_type):
reg_type=reg_type, warmstart=warmstart, verbose=True
)

_, log = ot.lp.emd(a, b, M, log=True)
warmstart1 = (log["u"], log["v"])
_, log_emd = ot.lp.emd(a, b, M, log=True)
warmstart1 = (log_emd["u"], log_emd["v"])
G1, log1 = ot.unbalanced.sinkhorn_unbalanced(
a, b, M, reg=epsilon, reg_m=reg_m, method=method,
reg_type=reg_type, warmstart=warmstart1, log=True, verbose=True
Expand All @@ -126,9 +124,6 @@ def test_unbalanced_warmstart(nx, method, reg_type):
reg_type=reg_type, warmstart=warmstart1, verbose=True
)

np.testing.assert_allclose(nx.to_numpy(loss), nx.to_numpy(loss0), atol=1e-5)
np.testing.assert_allclose(nx.to_numpy(loss0), nx.to_numpy(loss1), atol=1e-5)

np.testing.assert_allclose(
nx.to_numpy(log["logu"]), nx.to_numpy(log0["logu"]), atol=1e-05)
np.testing.assert_allclose(
Expand All @@ -141,6 +136,9 @@ def test_unbalanced_warmstart(nx, method, reg_type):
np.testing.assert_allclose(nx.to_numpy(G), nx.to_numpy(G0), atol=1e-05)
np.testing.assert_allclose(nx.to_numpy(G0), nx.to_numpy(G1), atol=1e-05)

np.testing.assert_allclose(nx.to_numpy(loss), nx.to_numpy(loss0), atol=1e-5)
np.testing.assert_allclose(nx.to_numpy(loss0), nx.to_numpy(loss1), atol=1e-5)


@pytest.mark.parametrize("method,reg_m", itertools.product(["sinkhorn", "sinkhorn_stabilized"], [1, float("inf")]))
def test_unbalanced_relaxation_parameters(nx, method, reg_m):
Expand Down