[MRG] Domain adaptation and unbalanced solvers with backend support (#343)

* First draft

* Add matrix inverse and square root to backend

* Eigen decomposition for older versions of pytorch (1.8.1 and older)

* Corrected eigen decomposition for pytorch 1.8.1 and older

* Spectral theorem is a thing

* Optimization

* small optimization

* More functions converted

* pep8

* remove a warning and prepare torch meshgrid for future torch release (which will change default indexing)

* dots and pep8

* Meshgrid corrected for older version and prepared for future versions changes

* New backend functions

* Base transport

* LinearTransport

* All transport classes + pep8

* PR added to release file

* Jcpot barycenter test

* unbalanced with backend

* pep8

* bug solve

* test of domain adaptation with backends

* solve bug for tic toc & macos

* solving scipy deprecation warning

* solving scipy deprecation warning attempt2

* solving scipy deprecation warning attempt3

* A warning is triggered when a float->int conversion is detected

* bug solve

* docs

* release file updated

* Better handling of float->int conversion in EMD

* Corrected test for is_floating_point

* docs

* release file updated

* cupy does not allow implicit cast

* fromnumpy

* added test

* test da tf jax

* test unbalanced with no provided histogram

* using type_as argument in unif function correctly

* pep8

* transport plan cast in emd changed behaviour, now trying to cast as histogram's dtype, defaulting to cost matrix

Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>

1 files changed, 50 insertions, 97 deletions

diff --git a/test/test_gromov.py b/test/test_gromov.py
index 0dcf2da..12fd2b9 100644
--- a/test/test_gromov.py
+++ b/test/test_gromov.py
@@ -35,11 +35,7 @@ def test_gromov(nx):
     C1 /= C1.max()

     C2 /= C2.max()

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    pb = nx.from_numpy(p)

-    qb = nx.from_numpy(q)

-    G0b = nx.from_numpy(G0)

+    C1b, C2b, pb, qb, G0b = nx.from_numpy(C1, C2, p, q, G0)

 

     G = ot.gromov.gromov_wasserstein(C1, C2, p, q, 'square_loss', G0=G0, verbose=True)

     Gb = nx.to_numpy(ot.gromov.gromov_wasserstein(C1b, C2b, pb, qb, 'square_loss', G0=G0b, verbose=True))

@@ -105,11 +101,7 @@ def test_gromov_dtype_device(nx):
     for tp in nx.__type_list__:

         print(nx.dtype_device(tp))

 

-        C1b = nx.from_numpy(C1, type_as=tp)

-        C2b = nx.from_numpy(C2, type_as=tp)

-        pb = nx.from_numpy(p, type_as=tp)

-        qb = nx.from_numpy(q, type_as=tp)

-        G0b = nx.from_numpy(G0, type_as=tp)

+        C1b, C2b, pb, qb, G0b = nx.from_numpy(C1, C2, p, q, G0, type_as=tp)

 

         Gb = ot.gromov.gromov_wasserstein(C1b, C2b, pb, qb, 'square_loss', G0=G0b, verbose=True)

         gw_valb = ot.gromov.gromov_wasserstein2(C1b, C2b, pb, qb, 'kl_loss', G0=G0b, log=False)

@@ -136,11 +128,7 @@ def test_gromov_device_tf():
 

     # Check that everything stays on the CPU

     with tf.device("/CPU:0"):

-        C1b = nx.from_numpy(C1)

-        C2b = nx.from_numpy(C2)

-        pb = nx.from_numpy(p)

-        qb = nx.from_numpy(q)

-        G0b = nx.from_numpy(G0)

+        C1b, C2b, pb, qb, G0b = nx.from_numpy(C1, C2, p, q, G0)

         Gb = ot.gromov.gromov_wasserstein(C1b, C2b, pb, qb, 'square_loss', G0=G0b, verbose=True)

         gw_valb = ot.gromov.gromov_wasserstein2(C1b, C2b, pb, qb, 'kl_loss', G0=G0b, log=False)

         nx.assert_same_dtype_device(C1b, Gb)

@@ -148,11 +136,7 @@ def test_gromov_device_tf():
 

     if len(tf.config.list_physical_devices('GPU')) > 0:

         # Check that everything happens on the GPU

-        C1b = nx.from_numpy(C1)

-        C2b = nx.from_numpy(C2)

-        pb = nx.from_numpy(p)

-        qb = nx.from_numpy(q)

-        G0b = nx.from_numpy(G0b)

+        C1b, C2b, pb, qb, G0b = nx.from_numpy(C1, C2, p, q, G0)

         Gb = ot.gromov.gromov_wasserstein(C1b, C2b, pb, qb, 'square_loss', verbose=True)

         gw_valb = ot.gromov.gromov_wasserstein2(C1b, C2b, pb, qb, 'kl_loss', log=False)

         nx.assert_same_dtype_device(C1b, Gb)

@@ -222,10 +206,7 @@ def test_entropic_gromov(nx):
     C1 /= C1.max()

     C2 /= C2.max()

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    pb = nx.from_numpy(p)

-    qb = nx.from_numpy(q)

+    C1b, C2b, pb, qb = nx.from_numpy(C1, C2, p, q)

 

     G = ot.gromov.entropic_gromov_wasserstein(

         C1, C2, p, q, 'square_loss', epsilon=5e-4, verbose=True)

@@ -285,10 +266,7 @@ def test_entropic_gromov_dtype_device(nx):
     for tp in nx.__type_list__:

         print(nx.dtype_device(tp))

 

-        C1b = nx.from_numpy(C1, type_as=tp)

-        C2b = nx.from_numpy(C2, type_as=tp)

-        pb = nx.from_numpy(p, type_as=tp)

-        qb = nx.from_numpy(q, type_as=tp)

+        C1b, C2b, pb, qb = nx.from_numpy(C1, C2, p, q, type_as=tp)

 

         Gb = ot.gromov.entropic_gromov_wasserstein(

             C1b, C2b, pb, qb, 'square_loss', epsilon=5e-4, verbose=True

@@ -320,10 +298,7 @@ def test_pointwise_gromov(nx):
     C1 /= C1.max()

     C2 /= C2.max()

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    pb = nx.from_numpy(p)

-    qb = nx.from_numpy(q)

+    C1b, C2b, pb, qb = nx.from_numpy(C1, C2, p, q)

 

     def loss(x, y):

         return np.abs(x - y)

@@ -381,10 +356,7 @@ def test_sampled_gromov(nx):
     C1 /= C1.max()

     C2 /= C2.max()

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    pb = nx.from_numpy(p)

-    qb = nx.from_numpy(q)

+    C1b, C2b, pb, qb = nx.from_numpy(C1, C2, p, q)

 

     def loss(x, y):

         return np.abs(x - y)

@@ -423,19 +395,15 @@ def test_gromov_barycenter(nx):
     n_samples = 3

     p = ot.unif(n_samples)

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    p1b = nx.from_numpy(p1)

-    p2b = nx.from_numpy(p2)

-    pb = nx.from_numpy(p)

+    C1b, C2b, p1b, p2b, pb = nx.from_numpy(C1, C2, p1, p2, p)

 

     Cb = ot.gromov.gromov_barycenters(

         n_samples, [C1, C2], [p1, p2], p, [.5, .5],

-        'square_loss', max_iter=100, tol=1e-3, verbose=True, random_state=42

+        'square_loss', max_iter=100, tol=1e-3, verbose=False, random_state=42

     )

     Cbb = nx.to_numpy(ot.gromov.gromov_barycenters(

         n_samples, [C1b, C2b], [p1b, p2b], pb, [.5, .5],

-        'square_loss', max_iter=100, tol=1e-3, verbose=True, random_state=42

+        'square_loss', max_iter=100, tol=1e-3, verbose=False, random_state=42

     ))

     np.testing.assert_allclose(Cb, Cbb, atol=1e-06)

     np.testing.assert_allclose(Cbb.shape, (n_samples, n_samples))

@@ -443,15 +411,15 @@ def test_gromov_barycenter(nx):
     # test of gromov_barycenters with `log` on

     Cb_, err_ = ot.gromov.gromov_barycenters(

         n_samples, [C1, C2], [p1, p2], p, [.5, .5],

-        'square_loss', max_iter=100, tol=1e-3, verbose=True, random_state=42, log=True

+        'square_loss', max_iter=100, tol=1e-3, verbose=False, random_state=42, log=True

     )

     Cbb_, errb_ = ot.gromov.gromov_barycenters(

         n_samples, [C1b, C2b], [p1b, p2b], pb, [.5, .5],

-        'square_loss', max_iter=100, tol=1e-3, verbose=True, random_state=42, log=True

+        'square_loss', max_iter=100, tol=1e-3, verbose=False, random_state=42, log=True

     )

     Cbb_ = nx.to_numpy(Cbb_)

     np.testing.assert_allclose(Cb_, Cbb_, atol=1e-06)

-    np.testing.assert_array_almost_equal(err_['err'], errb_['err'])

+    np.testing.assert_array_almost_equal(err_['err'], nx.to_numpy(*errb_['err']))

     np.testing.assert_allclose(Cbb_.shape, (n_samples, n_samples))

 

     Cb2 = ot.gromov.gromov_barycenters(

@@ -468,15 +436,15 @@ def test_gromov_barycenter(nx):
     # test of gromov_barycenters with `log` on

     Cb2_, err2_ = ot.gromov.gromov_barycenters(

         n_samples, [C1, C2], [p1, p2], p, [.5, .5],

-        'kl_loss', max_iter=100, tol=1e-3, verbose=True, random_state=42, log=True

+        'kl_loss', max_iter=100, tol=1e-3, verbose=False, random_state=42, log=True

     )

     Cb2b_, err2b_ = ot.gromov.gromov_barycenters(

         n_samples, [C1b, C2b], [p1b, p2b], pb, [.5, .5],

-        'kl_loss', max_iter=100, tol=1e-3, verbose=True, random_state=42, log=True

+        'kl_loss', max_iter=100, tol=1e-3, verbose=False, random_state=42, log=True

     )

     Cb2b_ = nx.to_numpy(Cb2b_)

     np.testing.assert_allclose(Cb2_, Cb2b_, atol=1e-06)

-    np.testing.assert_array_almost_equal(err2_['err'], err2_['err'])

+    np.testing.assert_array_almost_equal(err2_['err'], nx.to_numpy(*err2b_['err']))

     np.testing.assert_allclose(Cb2b_.shape, (n_samples, n_samples))

 

 

@@ -495,11 +463,7 @@ def test_gromov_entropic_barycenter(nx):
     n_samples = 2

     p = ot.unif(n_samples)

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    p1b = nx.from_numpy(p1)

-    p2b = nx.from_numpy(p2)

-    pb = nx.from_numpy(p)

+    C1b, C2b, p1b, p2b, pb = nx.from_numpy(C1, C2, p1, p2, p)

 

     Cb = ot.gromov.entropic_gromov_barycenters(

         n_samples, [C1, C2], [p1, p2], p, [.5, .5],

@@ -523,7 +487,7 @@ def test_gromov_entropic_barycenter(nx):
     )

     Cbb_ = nx.to_numpy(Cbb_)

     np.testing.assert_allclose(Cb_, Cbb_, atol=1e-06)

-    np.testing.assert_array_almost_equal(err_['err'], errb_['err'])

+    np.testing.assert_array_almost_equal(err_['err'], nx.to_numpy(*errb_['err']))

     np.testing.assert_allclose(Cbb_.shape, (n_samples, n_samples))

 

     Cb2 = ot.gromov.entropic_gromov_barycenters(

@@ -548,7 +512,7 @@ def test_gromov_entropic_barycenter(nx):
     )

     Cb2b_ = nx.to_numpy(Cb2b_)

     np.testing.assert_allclose(Cb2_, Cb2b_, atol=1e-06)

-    np.testing.assert_array_almost_equal(err2_['err'], err2_['err'])

+    np.testing.assert_array_almost_equal(err2_['err'], nx.to_numpy(*err2b_['err']))

     np.testing.assert_allclose(Cb2b_.shape, (n_samples, n_samples))

 

 

@@ -578,12 +542,7 @@ def test_fgw(nx):
     M = ot.dist(ys, yt)

     M /= M.max()

 

-    Mb = nx.from_numpy(M)

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    pb = nx.from_numpy(p)

-    qb = nx.from_numpy(q)

-    G0b = nx.from_numpy(G0)

+    Mb, C1b, C2b, pb, qb, G0b = nx.from_numpy(M, C1, C2, p, q, G0)

 

     G, log = ot.gromov.fused_gromov_wasserstein(M, C1, C2, p, q, 'square_loss', alpha=0.5, G0=G0, log=True)

     Gb, logb = ot.gromov.fused_gromov_wasserstein(Mb, C1b, C2b, pb, qb, 'square_loss', alpha=0.5, G0=G0b, log=True)

@@ -681,13 +640,7 @@ def test_fgw_barycenter(nx):
     n_samples = 3

     p = ot.unif(n_samples)

 

-    ysb = nx.from_numpy(ys)

-    ytb = nx.from_numpy(yt)

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    p1b = nx.from_numpy(p1)

-    p2b = nx.from_numpy(p2)

-    pb = nx.from_numpy(p)

+    ysb, ytb, C1b, C2b, p1b, p2b, pb = nx.from_numpy(ys, yt, C1, C2, p1, p2, p)

 

     Xb, Cb = ot.gromov.fgw_barycenters(

         n_samples, [ysb, ytb], [C1b, C2b], [p1b, p2b], [.5, .5], 0.5, fixed_structure=False,

@@ -731,10 +684,8 @@ def test_gromov_wasserstein_linear_unmixing(nx):
     Cdict = np.stack([C1, C2])

     p = ot.unif(n)

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    Cdictb = nx.from_numpy(Cdict)

-    pb = nx.from_numpy(p)

+    C1b, C2b, Cdictb, pb = nx.from_numpy(C1, C2, Cdict, p)

+

     tol = 10**(-5)

     # Tests without regularization

     reg = 0.

@@ -764,8 +715,8 @@ def test_gromov_wasserstein_linear_unmixing(nx):
     np.testing.assert_allclose(unmixing2, [0., 1.], atol=1e-01)

     np.testing.assert_allclose(C1_emb, nx.to_numpy(C1b_emb), atol=1e-06)

     np.testing.assert_allclose(C2_emb, nx.to_numpy(C2b_emb), atol=1e-06)

-    np.testing.assert_allclose(reconstruction1, reconstruction1b, atol=1e-06)

-    np.testing.assert_allclose(reconstruction2, reconstruction2b, atol=1e-06)

+    np.testing.assert_allclose(reconstruction1, nx.to_numpy(reconstruction1b), atol=1e-06)

+    np.testing.assert_allclose(reconstruction2, nx.to_numpy(reconstruction2b), atol=1e-06)

     np.testing.assert_allclose(C1b_emb.shape, (n, n))

     np.testing.assert_allclose(C2b_emb.shape, (n, n))

 

@@ -798,8 +749,8 @@ def test_gromov_wasserstein_linear_unmixing(nx):
     np.testing.assert_allclose(unmixing2, [0., 1.], atol=1e-01)

     np.testing.assert_allclose(C1_emb, nx.to_numpy(C1b_emb), atol=1e-06)

     np.testing.assert_allclose(C2_emb, nx.to_numpy(C2b_emb), atol=1e-06)

-    np.testing.assert_allclose(reconstruction1, reconstruction1b, atol=1e-06)

-    np.testing.assert_allclose(reconstruction2, reconstruction2b, atol=1e-06)

+    np.testing.assert_allclose(reconstruction1, nx.to_numpy(reconstruction1b), atol=1e-06)

+    np.testing.assert_allclose(reconstruction2, nx.to_numpy(reconstruction2b), atol=1e-06)

     np.testing.assert_allclose(C1b_emb.shape, (n, n))

     np.testing.assert_allclose(C2b_emb.shape, (n, n))

 

@@ -824,13 +775,14 @@ def test_gromov_wasserstein_dictionary_learning(nx):
     dataset_means = [C.mean() for C in Cs]

     np.random.seed(0)

     Cdict_init = np.random.normal(loc=np.mean(dataset_means), scale=np.std(dataset_means), size=(n_atoms, shape, shape))

+

     if projection == 'nonnegative_symmetric':

         Cdict_init = 0.5 * (Cdict_init + Cdict_init.transpose((0, 2, 1)))

         Cdict_init[Cdict_init < 0.] = 0.

-    Csb = [nx.from_numpy(C) for C in Cs]

-    psb = [nx.from_numpy(p) for p in ps]

-    qb = nx.from_numpy(q)

-    Cdict_initb = nx.from_numpy(Cdict_init)

+

+    Csb = nx.from_numpy(*Cs)

+    psb = nx.from_numpy(*ps)

+    qb, Cdict_initb = nx.from_numpy(q, Cdict_init)

 

     # Test: compare reconstruction error using initial dictionary and dictionary learned using this initialization

     # > Compute initial reconstruction of samples on this random dictionary without backend

@@ -882,6 +834,7 @@ def test_gromov_wasserstein_dictionary_learning(nx):
         )

         total_reconstruction_b += reconstruction

 

+    total_reconstruction_b = nx.to_numpy(total_reconstruction_b)

     np.testing.assert_array_less(total_reconstruction_b, initial_total_reconstruction)

     np.testing.assert_allclose(total_reconstruction_b, total_reconstruction, atol=1e-05)

     np.testing.assert_allclose(total_reconstruction_b, total_reconstruction, atol=1e-05)

@@ -924,6 +877,7 @@ def test_gromov_wasserstein_dictionary_learning(nx):
         )

         total_reconstruction_b_bis += reconstruction

 

+    total_reconstruction_b_bis = nx.to_numpy(total_reconstruction_b_bis)

     np.testing.assert_allclose(total_reconstruction_b_bis, total_reconstruction_b, atol=1e-05)

     np.testing.assert_allclose(Cdict_bis, nx.to_numpy(Cdictb_bis), atol=1e-03)

 

@@ -969,6 +923,7 @@ def test_gromov_wasserstein_dictionary_learning(nx):
         )

         total_reconstruction_b_bis2 += reconstruction

 

+    total_reconstruction_b_bis2 = nx.to_numpy(total_reconstruction_b_bis2)

     np.testing.assert_allclose(total_reconstruction_b_bis2, total_reconstruction_bis2, atol=1e-05)

 

 

@@ -985,12 +940,8 @@ def test_fused_gromov_wasserstein_linear_unmixing(nx):
     Ydict = np.stack([F, F])

     p = ot.unif(n)

 

-    C1b = nx.from_numpy(C1)

-    C2b = nx.from_numpy(C2)

-    Fb = nx.from_numpy(F)

-    Cdictb = nx.from_numpy(Cdict)

-    Ydictb = nx.from_numpy(Ydict)

-    pb = nx.from_numpy(p)

+    C1b, C2b, Fb, Cdictb, Ydictb, pb = nx.from_numpy(C1, C2, F, Cdict, Ydict, p)

+

     # Tests without regularization

     reg = 0.

 

@@ -1022,8 +973,8 @@ def test_fused_gromov_wasserstein_linear_unmixing(nx):
     np.testing.assert_allclose(C2_emb, nx.to_numpy(C2b_emb), atol=1e-03)

     np.testing.assert_allclose(Y1_emb, nx.to_numpy(Y1b_emb), atol=1e-03)

     np.testing.assert_allclose(Y2_emb, nx.to_numpy(Y2b_emb), atol=1e-03)

-    np.testing.assert_allclose(reconstruction1, reconstruction1b, atol=1e-06)

-    np.testing.assert_allclose(reconstruction2, reconstruction2b, atol=1e-06)

+    np.testing.assert_allclose(reconstruction1, nx.to_numpy(reconstruction1b), atol=1e-06)

+    np.testing.assert_allclose(reconstruction2, nx.to_numpy(reconstruction2b), atol=1e-06)

     np.testing.assert_allclose(C1b_emb.shape, (n, n))

     np.testing.assert_allclose(C2b_emb.shape, (n, n))

 

@@ -1058,8 +1009,8 @@ def test_fused_gromov_wasserstein_linear_unmixing(nx):
     np.testing.assert_allclose(C2_emb, nx.to_numpy(C2b_emb), atol=1e-03)

     np.testing.assert_allclose(Y1_emb, nx.to_numpy(Y1b_emb), atol=1e-03)

     np.testing.assert_allclose(Y2_emb, nx.to_numpy(Y2b_emb), atol=1e-03)

-    np.testing.assert_allclose(reconstruction1, reconstruction1b, atol=1e-06)

-    np.testing.assert_allclose(reconstruction2, reconstruction2b, atol=1e-06)

+    np.testing.assert_allclose(reconstruction1, nx.to_numpy(reconstruction1b), atol=1e-06)

+    np.testing.assert_allclose(reconstruction2, nx.to_numpy(reconstruction2b), atol=1e-06)

     np.testing.assert_allclose(C1b_emb.shape, (n, n))

     np.testing.assert_allclose(C2b_emb.shape, (n, n))

 

@@ -1093,12 +1044,10 @@ def test_fused_gromov_wasserstein_dictionary_learning(nx):
     dataset_feature_means = np.stack([Y.mean(axis=0) for Y in Ys])

     Ydict_init = np.random.normal(loc=dataset_feature_means.mean(axis=0), scale=dataset_feature_means.std(axis=0), size=(n_atoms, shape, 2))

 

-    Csb = [nx.from_numpy(C) for C in Cs]

-    Ysb = [nx.from_numpy(Y) for Y in Ys]

-    psb = [nx.from_numpy(p) for p in ps]

-    qb = nx.from_numpy(q)

-    Cdict_initb = nx.from_numpy(Cdict_init)

-    Ydict_initb = nx.from_numpy(Ydict_init)

+    Csb = nx.from_numpy(*Cs)

+    Ysb = nx.from_numpy(*Ys)

+    psb = nx.from_numpy(*ps)

+    qb, Cdict_initb, Ydict_initb = nx.from_numpy(q, Cdict_init, Ydict_init)

 

     # Test: Compute initial reconstruction of samples on this random dictionary

     alpha = 0.5

@@ -1151,6 +1100,7 @@ def test_fused_gromov_wasserstein_dictionary_learning(nx):
         )

         total_reconstruction_b += reconstruction

 

+    total_reconstruction_b = nx.to_numpy(total_reconstruction_b)

     np.testing.assert_array_less(total_reconstruction_b, initial_total_reconstruction)

     np.testing.assert_allclose(total_reconstruction_b, total_reconstruction, atol=1e-05)

     np.testing.assert_allclose(Cdict, nx.to_numpy(Cdictb), atol=1e-03)

@@ -1192,6 +1142,8 @@ def test_fused_gromov_wasserstein_dictionary_learning(nx):
             tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

         )

         total_reconstruction_b_bis += reconstruction

+

+    total_reconstruction_b_bis = nx.to_numpy(total_reconstruction_b_bis)

     np.testing.assert_allclose(total_reconstruction_b_bis, total_reconstruction_b, atol=1e-05)

 

     # Test: without using adam optimizer, with log and verbose set to True

@@ -1237,4 +1189,5 @@ def test_fused_gromov_wasserstein_dictionary_learning(nx):
         total_reconstruction_b_bis2 += reconstruction

 

     # > Compare results with/without backend

+    total_reconstruction_b_bis2 = nx.to_numpy(total_reconstruction_b_bis2)

     np.testing.assert_allclose(total_reconstruction_bis2, total_reconstruction_b_bis2, atol=1e-05)