import numpy as np
from ot.gromov import (
 semirelaxed_gromov_wasserstein,
 semirelaxed_fused_gromov_wasserstein,
 gromov_wasserstein,
 fused_gromov_wasserstein,
)
import networkx
from networkx.generators.community import stochastic_block_model as sbm
# 1st graph
N2 = 30 # 2 communities
p2 = [[1.0, 0.1, 0.0], [0.1, 0.95, 0.1], [0.0, 0.1, 0.9]]
G2 = sbm(seed=0, sizes=[N2 // 3, N2 // 3, N2 // 3], p=p2)
C2 = networkx.to_numpy_array(G2)
# 2nd graph
C3 = np.identity(4) #matrice d'adjacence
# Uniform distributions
h2 = np.ones(C2.shape[0]) / C2.shape[0]
h3 = np.ones(C3.shape[0]) / C3.shape[0]
# 0) GW(C2, h2, C3, h3) for reference
OT, log = gromov_wasserstein(C2, C3, h2, h3, symmetric=True, log=True)
gw = log["gw_dist"]
# 1) srGW(C2, h2, C3)
OT_23, log_23 = semirelaxed_gromov_wasserstein(
 C2, C3, h2, symmetric=True, log=True, G0=OT
)
srgw_23 = log_23["srgw_dist"]
# 2) srGW(C2, h2, C3) with default
OT_23_default, log_23_default = semirelaxed_gromov_wasserstein(
 C2, C3, h2, symmetric=True, log=True, G0=None
)
srgw_23_default = log_23_default["srgw_dist"]
# 3) srGW(C2, h2, C3) with default
initial = np.array(np.matmul(np.matrix(h2).T, np.matrix(h3)))
OT_23_initial, log_23_initial = semirelaxed_gromov_wasserstein(
 C2, C3, h2, symmetric=True, log=True, G0=initial
)
srgw_23_initial = log_23_initial["srgw_dist"]
print("GW(C2, C3) = ", gw) #GW(C2, C3) = 0.41222222222222227
print("G0 = OT, srGW(C2, h2, C3) = ", srgw_23) #G0 = OT, srGW(C2, h2, C3) = 0.1377777777777782
print("G0 = default, srGW(C2, h2, C3) = ", srgw_23_default) #G0 = default, srGW(C2, h2, C3) = 0.41222222222222227
print("G0 = same as default according to doc, srGW(C2, h2, C3) = ", srgw_23_initial) #G0 = same as default according to doc, srGW(C2, h2, C3) = 0.41222222222222227