laplace emd+sinkhorn

1 files changed, 106 insertions, 1 deletions

diff --git a/test/test_da.py b/test/test_da.py
index f700df9..15f4308 100644
--- a/test/test_da.py
+++ b/test/test_da.py
@@ -8,7 +8,6 @@ import numpy as np
 from numpy.testing import assert_allclose, assert_equal
 
 import ot
-from ot.bregman import jcpot_barycenter
 from ot.datasets import make_data_classif
 from ot.utils import unif
 
@@ -602,3 +601,109 @@ def test_jcpot_transport_class():
 
     # check that the oos method is working
     assert_equal(transp_Xs_new.shape, Xs_new.shape)
+
+def test_emd_laplace_class():
+    """test_emd_laplace_transport
+    """
+    ns = 150
+    nt = 200
+
+    Xs, ys = make_data_classif('3gauss', ns)
+    Xt, yt = make_data_classif('3gauss2', nt)
+
+    otda = ot.da.EMDLaplaceTransport(reg_lap=0.01, max_iter=1000, tol=1e-9, verbose=False, log=True)
+
+    # test its computed
+    otda.fit(Xs=Xs, ys=ys, Xt=Xt)
+
+    assert hasattr(otda, "coupling_")
+    assert hasattr(otda, "log_")
+
+    # test dimensions of coupling
+    assert_equal(otda.coupling_.shape, ((Xs.shape[0], Xt.shape[0])))
+
+    # test all margin constraints
+    mu_s = unif(ns)
+    mu_t = unif(nt)
+
+    assert_allclose(
+        np.sum(otda.coupling_, axis=0), mu_t, rtol=1e-3, atol=1e-3)
+    assert_allclose(
+        np.sum(otda.coupling_, axis=1), mu_s, rtol=1e-3, atol=1e-3)
+
+    # test transform
+    transp_Xs = otda.transform(Xs=Xs)
+    [assert_equal(x.shape, y.shape) for x, y in zip(transp_Xs, Xs)]
+
+    Xs_new, _ = make_data_classif('3gauss', ns + 1)
+    transp_Xs_new = otda.transform(Xs_new)
+
+    # check that the oos method is working
+    assert_equal(transp_Xs_new.shape, Xs_new.shape)
+
+    # test inverse transform
+    transp_Xt = otda.inverse_transform(Xt=Xt)
+    assert_equal(transp_Xt.shape, Xt.shape)
+
+    Xt_new, _ = make_data_classif('3gauss2', nt + 1)
+    transp_Xt_new = otda.inverse_transform(Xt=Xt_new)
+
+    # check that the oos method is working
+    assert_equal(transp_Xt_new.shape, Xt_new.shape)
+
+    # test fit_transform
+    transp_Xs = otda.fit_transform(Xs=Xs, Xt=Xt)
+    assert_equal(transp_Xs.shape, Xs.shape)
+
+def test_sinkhorn_laplace_class():
+    """test_sinkhorn_laplace_transport
+    """
+    ns = 150
+    nt = 200
+
+    Xs, ys = make_data_classif('3gauss', ns)
+    Xt, yt = make_data_classif('3gauss2', nt)
+
+    otda = ot.da.SinkhornLaplaceTransport(reg_e = 1, reg_lap=0.01, max_iter=1000, tol=1e-9, verbose=False, log=True)
+
+    # test its computed
+    otda.fit(Xs=Xs, ys=ys, Xt=Xt)
+
+    assert hasattr(otda, "coupling_")
+    assert hasattr(otda, "log_")
+
+    # test dimensions of coupling
+    assert_equal(otda.coupling_.shape, ((Xs.shape[0], Xt.shape[0])))
+
+    # test all margin constraints
+    mu_s = unif(ns)
+    mu_t = unif(nt)
+
+    assert_allclose(
+        np.sum(otda.coupling_, axis=0), mu_t, rtol=1e-3, atol=1e-3)
+    assert_allclose(
+        np.sum(otda.coupling_, axis=1), mu_s, rtol=1e-3, atol=1e-3)
+
+    # test transform
+    transp_Xs = otda.transform(Xs=Xs)
+    [assert_equal(x.shape, y.shape) for x, y in zip(transp_Xs, Xs)]
+
+    Xs_new, _ = make_data_classif('3gauss', ns + 1)
+    transp_Xs_new = otda.transform(Xs_new)
+
+    # check that the oos method is working
+    assert_equal(transp_Xs_new.shape, Xs_new.shape)
+
+    # test inverse transform
+    transp_Xt = otda.inverse_transform(Xt=Xt)
+    assert_equal(transp_Xt.shape, Xt.shape)
+
+    Xt_new, _ = make_data_classif('3gauss2', nt + 1)
+    transp_Xt_new = otda.inverse_transform(Xt=Xt_new)
+
+    # check that the oos method is working
+    assert_equal(transp_Xt_new.shape, Xt_new.shape)
+
+    # test fit_transform
+    transp_Xs = otda.fit_transform(Xs=Xs, Xt=Xt)
+    assert_equal(transp_Xs.shape, Xs.shape)
\ No newline at end of file