+ "# Add weights on the edges for visualization later on\nweight_intra_G2 = 5\nweight_inter_G2 = 0.5\nweight_intra_G3 = 1.\nweight_inter_G3 = 1.5\n\nweightedG2 = networkx.Graph()\npart_G2 = [G2.nodes[i]['block'] for i in range(N2)]\n\nfor node in G2.nodes():\n weightedG2.add_node(node)\nfor i, j in G2.edges():\n if part_G2[i] == part_G2[j]:\n weightedG2.add_edge(i, j, weight=weight_intra_G2)\n else:\n weightedG2.add_edge(i, j, weight=weight_inter_G2)\n\nweightedG3 = networkx.Graph()\npart_G3 = [G3.nodes[i]['block'] for i in range(N3)]\n\nfor node in G3.nodes():\n weightedG3.add_node(node)\nfor i, j in G3.edges():\n if part_G3[i] == part_G3[j]:\n weightedG3.add_edge(i, j, weight=weight_intra_G3)\n else:\n weightedG3.add_edge(i, j, weight=weight_inter_G3)\n\n\ndef draw_graph(G, C, nodes_color_part, Gweights=None,\n pos=None, edge_color='black', node_size=None,\n shiftx=0, seed=0):\n\n if (pos is None):\n pos = networkx.spring_layout(G, scale=1., seed=seed)\n\n if shiftx != 0:\n for k, v in pos.items():\n v[0] = v[0] + shiftx\n\n alpha_edge = 0.7\n width_edge = 1.8\n if Gweights is None:\n networkx.draw_networkx_edges(G, pos, width=width_edge, alpha=alpha_edge, edge_color=edge_color)\n else:\n # We make more visible connections between activated nodes\n n = len(Gweights)\n edgelist_activated = []\n edgelist_deactivated = []\n for i in range(n):\n for j in range(n):\n if Gweights[i] * Gweights[j] * C[i, j] > 0:\n edgelist_activated.append((i, j))\n elif C[i, j] > 0:\n edgelist_deactivated.append((i, j))\n\n networkx.draw_networkx_edges(G, pos, edgelist=edgelist_activated,\n width=width_edge, alpha=alpha_edge,\n edge_color=edge_color)\n networkx.draw_networkx_edges(G, pos, edgelist=edgelist_deactivated,\n width=width_edge, alpha=0.1,\n edge_color=edge_color)\n\n if Gweights is None:\n for node, node_color in enumerate(nodes_color_part):\n networkx.draw_networkx_nodes(G, pos, nodelist=[node],\n node_size=node_size, alpha=1,\n node_color=node_color)\n else:\n scaled_Gweights = Gweights / (0.5 * Gweights.max())\n nodes_size = node_size * scaled_Gweights\n for node, node_color in enumerate(nodes_color_part):\n networkx.draw_networkx_nodes(G, pos, nodelist=[node],\n node_size=nodes_size[node], alpha=1,\n node_color=node_color)\n return pos\n\n\ndef draw_transp_colored_GW(G1, C1, G2, C2, part_G1, p1, p2, T,\n pos1=None, pos2=None, shiftx=4, switchx=False,\n node_size=70, seed_G1=0, seed_G2=0):\n starting_color = 0\n # get graphs partition and their coloring\n part1 = part_G1.copy()\n unique_colors = ['C%s' % (starting_color + i) for i in np.unique(part1)]\n nodes_color_part1 = []\n for cluster in part1:\n nodes_color_part1.append(unique_colors[cluster])\n\n nodes_color_part2 = []\n # T: getting colors assignment from argmin of columns\n for i in range(len(G2.nodes())):\n j = np.argmax(T[:, i])\n nodes_color_part2.append(nodes_color_part1[j])\n pos1 = draw_graph(G1, C1, nodes_color_part1, Gweights=p1,\n pos=pos1, node_size=node_size, shiftx=0, seed=seed_G1)\n pos2 = draw_graph(G2, C2, nodes_color_part2, Gweights=p2, pos=pos2,\n node_size=node_size, shiftx=shiftx, seed=seed_G2)\n\n for k1, v1 in pos1.items():\n max_Tk1 = np.max(T[k1, :])\n for k2, v2 in pos2.items():\n if (T[k1, k2] > 0):\n pl.plot([pos1[k1][0], pos2[k2][0]],\n [pos1[k1][1], pos2[k2][1]],\n '-', lw=0.7, alpha=min(T[k1, k2] / max_Tk1 + 0.1, 1.),\n color=nodes_color_part1[k1])\n return pos1, pos2\n\n\nnode_size = 40\nfontsize = 13\nseed_G2 = 0\nseed_G3 = 4\n\npl.figure(2, figsize=(12, 3.5))\npl.clf()\npl.subplot(131)\npl.axis('off')\npl.axis\npl.title('(CG algo) FGW=%s \\n \\n OT sparsity = %s \\n feasibility error = %s' % (\n np.round(log0['fgw_dist'], 3), str(np.round(fgw0_sparsity, 2)) + ' %',\n np.round(err0, 4)), fontsize=fontsize)\n\np0, q0 = fgw0.sum(1), fgw0.sum(0) # check marginals\n\npos1, pos2 = draw_transp_colored_GW(\n weightedG2, C2, weightedG3, C3, part_G2, p1=p0, p2=q0, T=fgw0,\n shiftx=1.5, node_size=node_size, seed_G1=seed_G2, seed_G2=seed_G3)\n\npl.subplot(132)\npl.axis('off')\n\np, q = fgw.sum(1), fgw.sum(0) # check marginals\n\npl.title('(PP algo) FGW=%s\\n \\n OT sparsity = %s \\n feasibility error = %s' % (\n np.round(log['fgw_dist'], 3), str(np.round(fgw_sparsity, 2)) + ' %',\n np.round(err, 4)), fontsize=fontsize)\n\npos1, pos2 = draw_transp_colored_GW(\n weightedG2, C2, weightedG3, C3, part_G2, p1=p, p2=q, T=fgw,\n pos1=pos1, pos2=pos2, shiftx=0., node_size=node_size, seed_G1=0, seed_G2=0)\n\npl.subplot(133)\npl.axis('off')\n\npe, qe = fgwe.sum(1), fgwe.sum(0) # check marginals\n\npl.title('Entropic FGW=%s\\n \\n OT sparsity = %s \\n feasibility error = %s' % (\n np.round(loge['fgw_dist'], 3), str(np.round(fgwe_sparsity, 2)) + ' %',\n np.round(erre, 4)), fontsize=fontsize)\n\npos1, pos2 = draw_transp_colored_GW(\n weightedG2, C2, weightedG3, C3, part_G2, p1=pe, p2=qe, T=fgwe,\n pos1=pos1, pos2=pos2, shiftx=0., node_size=node_size, seed_G1=0, seed_G2=0)\n\npl.tight_layout()\n\npl.show()"
0 commit comments