| 707 | |
| 708 | |
| 709 | def test_exp_bures(nx): |
| 710 | d = 2 |
| 711 | |
| 712 | rng = np.random.RandomState(42) |
| 713 | X = rng.randn(d, d) |
| 714 | z = rng.randn(d) |
| 715 | X, z = nx.from_numpy(X, z) |
| 716 | S = X + nx.transpose(X) |
| 717 | |
| 718 | Sigma = nx.eye(d, type_as=S) |
| 719 | |
| 720 | Lambda = ot.utils.exp_bures(Sigma, S) |
| 721 | |
| 722 | # asserst SPD |
| 723 | np.testing.assert_array_less(np.zeros(1), nx.to_numpy(z.T @ Lambda @ z)) |
| 724 | |
| 725 | # OT map from Lambda to Sigma |
| 726 | Lambda12 = nx.sqrtm(Lambda) |
| 727 | Lambda12inv = nx.inv(Lambda12) |
| 728 | M = nx.sqrtm(nx.einsum("ij, jk, kl -> il", Lambda12, Sigma, Lambda12)) |
| 729 | T = nx.einsum("ij, jk, kl -> il", Lambda12inv, M, Lambda12inv) |
| 730 | |
| 731 | # exp_\Lambda(log_\Lambda(Sigma)) = Sigma |
| 732 | Sigma_exp = ot.utils.exp_bures(Lambda, T - nx.eye(d, type_as=T)) |
| 733 | np.testing.assert_allclose(nx.to_numpy(Sigma), nx.to_numpy(Sigma_exp), atol=1e-5) |
| 734 | |
| 735 | |
| 736 | def test_fun_to_numpy(nx): |