[MRG] GW dictionary learning (#319)

* add fgw dictionary learning feature

* add fgw dictionary learning feature

* plot gromov wasserstein dictionary learning

* Update __init__.py

* fix pep8 errors exact E501 line too long

* fix last pep8 issues

* add unitary tests for (F)GW dictionary learning without using autodifferentiable functions

* correct tests for (F)GW dictionary learning without using autodiff

* correct tests for (F)GW dictionary learning without using autodiff

* fix docs and notations

* answer to review: improve tests, docs, examples + make node weights optional

* fix pep8 and examples

* improve docs + tests + thumbnail

* make example faster

* improve ex

* update README.md

* make GDL tests faster

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

1 files changed, 542 insertions, 12 deletions

diff --git a/test/test_gromov.py b/test/test_gromov.py
index 4b995d5..329f99c 100644
--- a/test/test_gromov.py
+++ b/test/test_gromov.py
@@ -3,6 +3,7 @@
 # Author: Erwan Vautier <erwan.vautier@gmail.com>

 #         Nicolas Courty <ncourty@irisa.fr>

 #         Titouan Vayer <titouan.vayer@irisa.fr>

+#         Cédric Vincent-Cuaz <cedric.vincent-cuaz@inria.fr>

 #

 # License: MIT License

 

@@ -26,6 +27,7 @@ def test_gromov(nx):
 

     p = ot.unif(n_samples)

     q = ot.unif(n_samples)

+    G0 = p[:, None] * q[None, :]

 

     C1 = ot.dist(xs, xs)

     C2 = ot.dist(xt, xt)

@@ -37,9 +39,10 @@ def test_gromov(nx):
     C2b = nx.from_numpy(C2)

     pb = nx.from_numpy(p)

     qb = nx.from_numpy(q)

+    G0b = nx.from_numpy(G0)

 

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

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

+    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))

 

     # check constraints

     np.testing.assert_allclose(G, Gb, atol=1e-06)

@@ -56,9 +59,9 @@ def test_gromov(nx):
     gwb, logb = ot.gromov.gromov_wasserstein2(C1b, C2b, pb, qb, 'kl_loss', log=True)

     gwb = nx.to_numpy(gwb)

 

-    gw_val = ot.gromov.gromov_wasserstein2(C1, C2, p, q, 'kl_loss', log=False)

+    gw_val = ot.gromov.gromov_wasserstein2(C1, C2, p, q, 'kl_loss', G0=G0, log=False)

     gw_valb = nx.to_numpy(

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

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

     )

 

     G = log['T']

@@ -91,6 +94,7 @@ def test_gromov_dtype_device(nx):
 

     p = ot.unif(n_samples)

     q = ot.unif(n_samples)

+    G0 = p[:, None] * q[None, :]

 

     C1 = ot.dist(xs, xs)

     C2 = ot.dist(xt, xt)

@@ -105,9 +109,10 @@ def test_gromov_dtype_device(nx):
         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)

 

-        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)

+        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)

         nx.assert_same_dtype_device(C1b, gw_valb)

@@ -123,6 +128,7 @@ def test_gromov_device_tf():
     xt = xs[::-1].copy()

     p = ot.unif(n_samples)

     q = ot.unif(n_samples)

+    G0 = p[:, None] * q[None, :]

     C1 = ot.dist(xs, xs)

     C2 = ot.dist(xt, xt)

     C1 /= C1.max()

@@ -134,8 +140,9 @@ def test_gromov_device_tf():
         C2b = nx.from_numpy(C2)

         pb = nx.from_numpy(p)

         qb = nx.from_numpy(q)

-        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)

+        G0b = nx.from_numpy(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)

         nx.assert_same_dtype_device(C1b, gw_valb)

 

@@ -145,6 +152,7 @@ def test_gromov_device_tf():
         C2b = nx.from_numpy(C2)

         pb = nx.from_numpy(p)

         qb = nx.from_numpy(q)

+        G0b = nx.from_numpy(G0b)

         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)

@@ -554,6 +562,7 @@ def test_fgw(nx):
 

     p = ot.unif(n_samples)

     q = ot.unif(n_samples)

+    G0 = p[:, None] * q[None, :]

 

     C1 = ot.dist(xs, xs)

     C2 = ot.dist(xt, xt)

@@ -569,9 +578,10 @@ def test_fgw(nx):
     C2b = nx.from_numpy(C2)

     pb = nx.from_numpy(p)

     qb = nx.from_numpy(q)

+    G0b = nx.from_numpy(G0)

 

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

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

+    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)

     Gb = nx.to_numpy(Gb)

 

     # check constraints

@@ -586,8 +596,8 @@ def test_fgw(nx):
     np.testing.assert_allclose(

         Gb, np.flipud(Id), atol=1e-04)  # cf convergence gromov

 

-    fgw, log = ot.gromov.fused_gromov_wasserstein2(M, C1, C2, p, q, 'square_loss', alpha=0.5, log=True)

-    fgwb, logb = ot.gromov.fused_gromov_wasserstein2(Mb, C1b, C2b, pb, qb, 'square_loss', alpha=0.5, log=True)

+    fgw, log = ot.gromov.fused_gromov_wasserstein2(M, C1, C2, p, q, 'square_loss', G0=None, alpha=0.5, log=True)

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

     fgwb = nx.to_numpy(fgwb)

 

     G = log['T']

@@ -698,3 +708,523 @@ def test_fgw_barycenter(nx):
     Xb, Cb = nx.to_numpy(Xb), nx.to_numpy(Cb)

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

     np.testing.assert_allclose(Xb.shape, (n_samples, ys.shape[1]))

+

+

+def test_gromov_wasserstein_linear_unmixing(nx):

+    n = 10

+

+    X1, y1 = ot.datasets.make_data_classif('3gauss', n, random_state=42)

+    X2, y2 = ot.datasets.make_data_classif('3gauss2', n, random_state=42)

+

+    C1 = ot.dist(X1)

+    C2 = ot.dist(X2)

+    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)

+    tol = 10**(-5)

+    # Tests without regularization

+    reg = 0.

+    unmixing1, C1_emb, OT, reconstruction1 = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C1, Cdict, reg=reg, p=p, q=p,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing1b, C1b_emb, OTb, reconstruction1b = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C1b, Cdictb, reg=reg, p=None, q=None,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2, C2_emb, OT, reconstruction2 = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C2, Cdict, reg=reg, p=None, q=None,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2b, C2b_emb, OTb, reconstruction2b = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C2b, Cdictb, reg=reg, p=pb, q=pb,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

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

+    np.testing.assert_allclose(unmixing1, [1., 0.], atol=1e-01)

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

+    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(C1b_emb.shape, (n, n))

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

+

+    # Tests with regularization

+

+    reg = 0.001

+    unmixing1, C1_emb, OT, reconstruction1 = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C1, Cdict, reg=reg, p=p, q=p,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing1b, C1b_emb, OTb, reconstruction1b = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C1b, Cdictb, reg=reg, p=None, q=None,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2, C2_emb, OT, reconstruction2 = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C2, Cdict, reg=reg, p=None, q=None,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2b, C2b_emb, OTb, reconstruction2b = ot.gromov.gromov_wasserstein_linear_unmixing(

+        C2b, Cdictb, reg=reg, p=pb, q=pb,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+    )

+

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

+    np.testing.assert_allclose(unmixing1, [1., 0.], atol=1e-01)

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

+    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(C1b_emb.shape, (n, n))

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

+

+

+def test_gromov_wasserstein_dictionary_learning(nx):

+

+    # create dataset composed from 2 structures which are repeated 5 times

+    shape = 10

+    n_samples = 2

+    n_atoms = 2

+    projection = 'nonnegative_symmetric'

+    X1, y1 = ot.datasets.make_data_classif('3gauss', shape, random_state=42)

+    X2, y2 = ot.datasets.make_data_classif('3gauss2', shape, random_state=42)

+    C1 = ot.dist(X1)

+    C2 = ot.dist(X2)

+    Cs = [C1.copy() for _ in range(n_samples // 2)] + [C2.copy() for _ in range(n_samples // 2)]

+    ps = [ot.unif(shape) for _ in range(n_samples)]

+    q = ot.unif(shape)

+

+    # Provide initialization for the graph dictionary of shape (n_atoms, shape, shape)

+    # following the same procedure than implemented in gromov_wasserstein_dictionary_learning.

+    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)

+

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

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

+    use_adam_optimizer = True

+    verbose = False

+    tol = 10**(-5)

+    epochs = 1

+

+    initial_total_reconstruction = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Cs[i], Cdict_init, p=ps[i], q=q, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        initial_total_reconstruction += reconstruction

+

+    # > Learn the dictionary using this init

+    Cdict, log = ot.gromov.gromov_wasserstein_dictionary_learning(

+        Cs, D=n_atoms, nt=shape, ps=ps, q=q, Cdict_init=Cdict_init,

+        epochs=epochs, batch_size=2 * n_samples, learning_rate=1., reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary without backend

+    total_reconstruction = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Cs[i], Cdict, p=None, q=None, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction += reconstruction

+

+    np.testing.assert_array_less(total_reconstruction, initial_total_reconstruction)

+

+    # Test: Perform same experiments after going through backend

+

+    Cdictb, log = ot.gromov.gromov_wasserstein_dictionary_learning(

+        Csb, D=n_atoms, nt=shape, ps=None, q=None, Cdict_init=Cdict_initb,

+        epochs=epochs, batch_size=n_samples, learning_rate=1., reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # Compute reconstruction of samples on learned dictionary

+    total_reconstruction_b = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Csb[i], Cdictb, p=psb[i], q=qb, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_b += reconstruction

+

+    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)

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

+

+    # Test: Perform same comparison without providing the initial dictionary being an optional input

+    #       knowing than the initialization scheme is the same than implemented to set the benchmarked initialization.

+    np.random.seed(0)

+    Cdict_bis, log = ot.gromov.gromov_wasserstein_dictionary_learning(

+        Cs, D=n_atoms, nt=shape, ps=None, q=None, Cdict_init=None,

+        epochs=epochs, batch_size=n_samples, learning_rate=1., reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_bis = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Cs[i], Cdict_bis, p=ps[i], q=q, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_bis += reconstruction

+

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

+

+    # Test: Same after going through backend

+    np.random.seed(0)

+    Cdictb_bis, log = ot.gromov.gromov_wasserstein_dictionary_learning(

+        Csb, D=n_atoms, nt=shape, ps=psb, q=qb, Cdict_init=None,

+        epochs=epochs, batch_size=n_samples, learning_rate=1., reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_b_bis = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Csb[i], Cdictb_bis, p=None, q=None, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_b_bis += reconstruction

+

+    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)

+

+    # Test: Perform same comparison without providing the initial dictionary being an optional input

+    #       and testing other optimization settings untested until now.

+    #       We pass previously estimated dictionaries to speed up the process.

+    use_adam_optimizer = False

+    verbose = True

+    use_log = True

+

+    np.random.seed(0)

+    Cdict_bis2, log = ot.gromov.gromov_wasserstein_dictionary_learning(

+        Cs, D=n_atoms, nt=shape, ps=ps, q=q, Cdict_init=Cdict,

+        epochs=epochs, batch_size=n_samples, learning_rate=10., reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=use_log, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_bis2 = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Cs[i], Cdict_bis2, p=ps[i], q=q, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_bis2 += reconstruction

+

+    np.testing.assert_array_less(total_reconstruction_bis2, total_reconstruction)

+

+    # Test: Same after going through backend

+    np.random.seed(0)

+    Cdictb_bis2, log = ot.gromov.gromov_wasserstein_dictionary_learning(

+        Csb, D=n_atoms, nt=shape, ps=psb, q=qb, Cdict_init=Cdictb,

+        epochs=epochs, batch_size=n_samples, learning_rate=10., reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=use_log, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_b_bis2 = 0

+    for i in range(n_samples):

+        _, _, _, reconstruction = ot.gromov.gromov_wasserstein_linear_unmixing(

+            Csb[i], Cdictb_bis2, p=psb[i], q=qb, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_b_bis2 += reconstruction

+

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

+

+

+def test_fused_gromov_wasserstein_linear_unmixing(nx):

+

+    n = 10

+    X1, y1 = ot.datasets.make_data_classif('3gauss', n, random_state=42)

+    X2, y2 = ot.datasets.make_data_classif('3gauss2', n, random_state=42)

+    F, y = ot.datasets.make_data_classif('3gauss', n, random_state=42)

+

+    C1 = ot.dist(X1)

+    C2 = ot.dist(X2)

+    Cdict = np.stack([C1, C2])

+    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)

+    # Tests without regularization

+    reg = 0.

+

+    unmixing1, C1_emb, Y1_emb, OT, reconstruction1 = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C1, F, Cdict, Ydict, p=p, q=p, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing1b, C1b_emb, Y1b_emb, OTb, reconstruction1b = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C1b, Fb, Cdictb, Ydictb, p=None, q=None, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2, C2_emb, Y2_emb, OT, reconstruction2 = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C2, F, Cdict, Ydict, p=None, q=None, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2b, C2b_emb, Y2b_emb, OTb, reconstruction2b = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C2b, Fb, Cdictb, Ydictb, p=pb, q=pb, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

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

+    np.testing.assert_allclose(unmixing1, [1., 0.], atol=1e-01)

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

+    np.testing.assert_allclose(unmixing2, [0., 1.], atol=1e-01)

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

+    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(C1b_emb.shape, (n, n))

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

+

+    # Tests with regularization

+    reg = 0.001

+

+    unmixing1, C1_emb, Y1_emb, OT, reconstruction1 = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C1, F, Cdict, Ydict, p=p, q=p, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing1b, C1b_emb, Y1b_emb, OTb, reconstruction1b = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C1b, Fb, Cdictb, Ydictb, p=None, q=None, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2, C2_emb, Y2_emb, OT, reconstruction2 = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C2, F, Cdict, Ydict, p=None, q=None, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

+    unmixing2b, C2b_emb, Y2b_emb, OTb, reconstruction2b = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+        C2b, Fb, Cdictb, Ydictb, p=pb, q=pb, alpha=0.5, reg=reg,

+        tol_outer=10**(-6), tol_inner=10**(-6), max_iter_outer=20, max_iter_inner=200

+    )

+

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

+    np.testing.assert_allclose(unmixing1, [1., 0.], atol=1e-01)

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

+    np.testing.assert_allclose(unmixing2, [0., 1.], atol=1e-01)

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

+    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(C1b_emb.shape, (n, n))

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

+

+

+def test_fused_gromov_wasserstein_dictionary_learning(nx):

+

+    # create dataset composed from 2 structures which are repeated 5 times

+    shape = 10

+    n_samples = 2

+    n_atoms = 2

+    projection = 'nonnegative_symmetric'

+    X1, y1 = ot.datasets.make_data_classif('3gauss', shape, random_state=42)

+    X2, y2 = ot.datasets.make_data_classif('3gauss2', shape, random_state=42)

+    F, y = ot.datasets.make_data_classif('3gauss', shape, random_state=42)

+

+    C1 = ot.dist(X1)

+    C2 = ot.dist(X2)

+    Cs = [C1.copy() for _ in range(n_samples // 2)] + [C2.copy() for _ in range(n_samples // 2)]

+    Ys = [F.copy() for _ in range(n_samples)]

+    ps = [ot.unif(shape) for _ in range(n_samples)]

+    q = ot.unif(shape)

+

+    # Provide initialization for the graph dictionary of shape (n_atoms, shape, shape)

+    # following the same procedure than implemented in gromov_wasserstein_dictionary_learning.

+    dataset_structure_means = [C.mean() for C in Cs]

+    np.random.seed(0)

+    Cdict_init = np.random.normal(loc=np.mean(dataset_structure_means), scale=np.std(dataset_structure_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.

+    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)

+

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

+    alpha = 0.5

+    use_adam_optimizer = True

+    verbose = False

+    tol = 1e-05

+    epochs = 1

+

+    initial_total_reconstruction = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Cs[i], Ys[i], Cdict_init, Ydict_init, p=ps[i], q=q,

+            alpha=alpha, reg=0., tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        initial_total_reconstruction += reconstruction

+

+    # > Learn a dictionary using this given initialization and check that the reconstruction loss

+    # on the learned dictionary is lower than the one using its initialization.

+    Cdict, Ydict, log = ot.gromov.fused_gromov_wasserstein_dictionary_learning(

+        Cs, Ys, D=n_atoms, nt=shape, ps=ps, q=q, Cdict_init=Cdict_init, Ydict_init=Ydict_init,

+        epochs=epochs, batch_size=n_samples, learning_rate_C=1., learning_rate_Y=1., alpha=alpha, reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Cs[i], Ys[i], Cdict, Ydict, p=None, q=None, alpha=alpha, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction += reconstruction

+    # Compare both

+    np.testing.assert_array_less(total_reconstruction, initial_total_reconstruction)

+

+    # Test: Perform same experiments after going through backend

+

+    Cdictb, Ydictb, log = ot.gromov.fused_gromov_wasserstein_dictionary_learning(

+        Csb, Ysb, D=n_atoms, nt=shape, ps=None, q=None, Cdict_init=Cdict_initb, Ydict_init=Ydict_initb,

+        epochs=epochs, batch_size=2 * n_samples, learning_rate_C=1., learning_rate_Y=1., alpha=alpha, reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_b = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Csb[i], Ysb[i], Cdictb, Ydictb, p=psb[i], q=qb, alpha=alpha, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_b += reconstruction

+

+    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)

+    np.testing.assert_allclose(Ydict, nx.to_numpy(Ydictb), atol=1e-03)

+

+    # Test: Perform similar experiment without providing the initial dictionary being an optional input

+    np.random.seed(0)

+    Cdict_bis, Ydict_bis, log = ot.gromov.fused_gromov_wasserstein_dictionary_learning(

+        Cs, Ys, D=n_atoms, nt=shape, ps=None, q=None, Cdict_init=None, Ydict_init=None,

+        epochs=epochs, batch_size=n_samples, learning_rate_C=1., learning_rate_Y=1., alpha=alpha, reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_bis = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Cs[i], Ys[i], Cdict_bis, Ydict_bis, p=ps[i], q=q, alpha=alpha, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_bis += reconstruction

+

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

+

+    # > Same after going through backend

+    np.random.seed(0)

+    Cdictb_bis, Ydictb_bis, log = ot.gromov.fused_gromov_wasserstein_dictionary_learning(

+        Csb, Ysb, D=n_atoms, nt=shape, ps=None, q=None, Cdict_init=None, Ydict_init=None,

+        epochs=epochs, batch_size=n_samples, learning_rate_C=1., learning_rate_Y=1., alpha=alpha, reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=False, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_b_bis = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Csb[i], Ysb[i], Cdictb_bis, Ydictb_bis, p=psb[i], q=qb, alpha=alpha, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_b_bis += reconstruction

+    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

+    use_adam_optimizer = False

+    verbose = True

+    use_log = True

+

+    # > Experiment providing previously estimated dictionary to speed up the test compared to providing initial random init.

+    np.random.seed(0)

+    Cdict_bis2, Ydict_bis2, log = ot.gromov.fused_gromov_wasserstein_dictionary_learning(

+        Cs, Ys, D=n_atoms, nt=shape, ps=ps, q=q, Cdict_init=Cdict, Ydict_init=Ydict,

+        epochs=epochs, batch_size=n_samples, learning_rate_C=10., learning_rate_Y=10., alpha=alpha, reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=use_log, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_bis2 = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Cs[i], Ys[i], Cdict_bis2, Ydict_bis2, p=ps[i], q=q, alpha=alpha, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_bis2 += reconstruction

+

+    np.testing.assert_array_less(total_reconstruction_bis2, total_reconstruction)

+

+    # > Same after going through backend

+    np.random.seed(0)

+    Cdictb_bis2, Ydictb_bis2, log = ot.gromov.fused_gromov_wasserstein_dictionary_learning(

+        Csb, Ysb, D=n_atoms, nt=shape, ps=None, q=None, Cdict_init=Cdictb, Ydict_init=Ydictb,

+        epochs=epochs, batch_size=n_samples, learning_rate_C=10., learning_rate_Y=10., alpha=alpha, reg=0.,

+        tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200,

+        projection=projection, use_log=use_log, use_adam_optimizer=use_adam_optimizer, verbose=verbose

+    )

+

+    # > Compute reconstruction of samples on learned dictionary

+    total_reconstruction_b_bis2 = 0

+    for i in range(n_samples):

+        _, _, _, _, reconstruction = ot.gromov.fused_gromov_wasserstein_linear_unmixing(

+            Csb[i], Ysb[i], Cdictb_bis2, Ydictb_bis2, p=None, q=None, alpha=alpha, reg=0.,

+            tol_outer=tol, tol_inner=tol, max_iter_outer=20, max_iter_inner=200

+        )

+        total_reconstruction_b_bis2 += reconstruction

+

+    # > Compare results with/without backend

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