| 527 | ], |
| 528 | ) |
| 529 | def test_unbalanced_multiple_inputs(nx, method): |
| 530 | # test generalized sinkhorn for unbalanced OT |
| 531 | n = 100 |
| 532 | rng = np.random.RandomState(42) |
| 533 | |
| 534 | x = rng.randn(n, 2) |
| 535 | a = ot.utils.unif(n) |
| 536 | |
| 537 | # make dists unbalanced |
| 538 | b = rng.rand(n, 2) |
| 539 | |
| 540 | M = ot.dist(x, x) |
| 541 | epsilon = 1.0 |
| 542 | reg_m = 1.0 |
| 543 | |
| 544 | a, b, M = nx.from_numpy(a, b, M) |
| 545 | |
| 546 | G, log = ot.unbalanced.sinkhorn_unbalanced( |
| 547 | a, b, M, reg=epsilon, reg_m=reg_m, method=method, log=True, verbose=True |
| 548 | ) |
| 549 | |
| 550 | # check fixed point equations |
| 551 | # in log-domain |
| 552 | fi = reg_m / (reg_m + epsilon) |
| 553 | logb = nx.log(b + 1e-16) |
| 554 | loga = nx.log(a + 1e-16)[:, None] |
| 555 | logKtu = nx.logsumexp(log["logu"][:, None, :] - M[:, :, None] / epsilon, axis=0) |
| 556 | logKv = nx.logsumexp(log["logv"][None, :] - M[:, :, None] / epsilon, axis=1) |
| 557 | v_final = fi * (logb - logKtu) |
| 558 | u_final = fi * (loga - logKv) |
| 559 | |
| 560 | print("u_final shape = {}".format(u_final.shape)) |
| 561 | print("v_final shape = {}".format(v_final.shape)) |
| 562 | print("logu shape = {}".format(log["logu"].shape)) |
| 563 | print("logv shape = {}".format(log["logv"].shape)) |
| 564 | |
| 565 | np.testing.assert_allclose( |
| 566 | nx.to_numpy(u_final), nx.to_numpy(log["logu"]), atol=1e-05 |
| 567 | ) |
| 568 | np.testing.assert_allclose( |
| 569 | nx.to_numpy(v_final), nx.to_numpy(log["logv"]), atol=1e-05 |
| 570 | ) |
| 571 | |
| 572 | # reg_type="entropy" as multiple inputs does not work for KL yet |
| 573 | losses = ot.unbalanced.sinkhorn_unbalanced2( |
| 574 | a, b, M, reg=epsilon, reg_m=reg_m, method=method, reg_type="entropy" |
| 575 | ) |
| 576 | |
| 577 | loss1 = ot.unbalanced.sinkhorn_unbalanced2( |
| 578 | a, b[:, 0], M, reg=epsilon, reg_m=reg_m, method=method, reg_type="entropy" |
| 579 | ) |
| 580 | loss2 = ot.unbalanced.sinkhorn_unbalanced2( |
| 581 | a, b[:, 1], M, reg=epsilon, reg_m=reg_m, method=method, reg_type="entropy" |
| 582 | ) |
| 583 | |
| 584 | np.testing.assert_allclose(nx.to_numpy(losses), np.array([loss1, loss2]), atol=1e-4) |
| 585 | |
| 586 | |