plot_network.py

import sys
import click

import networkx as nx
import numpy as np

import matplotlib.pyplot as plt
import seaborn as sns

from matplotlib.cm import get_cmap
from networkx.algorithms.community.label_propagation import (
    label_propagation_communities)
from etudes.datasets.networks import load_dataset
from pathlib import Path

GOLDEN_RATIO = 0.5 * (1 + np.sqrt(5))


def pt_to_in(x):

    pt_per_in = 72.27
    return x / pt_per_in


def size(width, aspect=GOLDEN_RATIO):

    width_in = pt_to_in(width)
    return (width_in, width_in / aspect)


WIDTH = 397.48499
OUTPUT_DIR = "figures/"


@click.command()
@click.argument("name")
@click.argument("num_subgraphs", type=int)
@click.option("--input-dir", default="datasets",
              type=click.Path(file_okay=False, dir_okay=True))
@click.option("--output-dir", default=OUTPUT_DIR,
              type=click.Path(file_okay=False, dir_okay=True),
              help="Output directory.")
@click.option('--edge-width-factor', default=1.0)
@click.option('--node-cmap', default="Dark2")
@click.option('--node-size', default=300)
@click.option("--node-alpha", default=0.7)
@click.option("--edge-alpha", default=0.4)
@click.option("--label-alpha", default=0.6)
@click.option('--labels/--no-labels', default=False)
@click.option("--attacked-edge-color", default="maroon")
@click.option("--attacked-edge-style", default="dashed")
@click.option("--font-size", default="xx-small")
@click.option("--font-color", default='k')
@click.option("--font-family", default='sans-serif')
@click.option("--font-weight", default='normal')
@click.option('--context', default="paper")
@click.option('--style', default="ticks")
@click.option('--palette', default="muted")
@click.option('--width', '-w', type=float, default=WIDTH)
@click.option('--aspect', '-a', type=float, default=GOLDEN_RATIO)
@click.option('--extension', '-e', multiple=True, default=["png"])
@click.option("--seed", default=8888)
def main(name, num_subgraphs, input_dir, output_dir, edge_width_factor,
         node_cmap, node_size, node_alpha, edge_alpha, label_alpha, labels,
         attacked_edge_color, attacked_edge_style, font_size, font_color,
         font_family, font_weight, context, style, palette, width, aspect,
         extension, seed):

    figsize = size(width, aspect)
    suffix = f"{width:.0f}x{width/aspect:.0f}"

    rc = {
        "figure.figsize": figsize,
        "font.serif": ['Times New Roman'],
        "text.usetex": True,
    }
    sns.set(context=context, style=style, palette=palette, font="serif", rc=rc)

    input_path = Path(input_dir)
    output_path = Path(output_dir).joinpath(name)
    output_path.mkdir(parents=True, exist_ok=True)

    random_state = np.random.RandomState(seed)

    click.secho("Loading original graph...")
    X, y, A = load_dataset("citeseer", data_home=input_path)
    G = nx.from_scipy_sparse_matrix(A)
    classes = np.argmax(y, axis=1)

    click.secho("Loading compromised graph...")
    H = nx.read_gpickle(input_path.joinpath("citeseer_attack_add_5000_v_6.gpickle"))

    click.secho("Loading edge posterior probabilities...")
    W = np.load(input_path.joinpath("A.npy"))
    P = nx.from_numpy_array(W)

    # Find an "interesting" subgraph to analyze
    # Here we enumerate communities using label propagation and take the
    # largest one found.
    # C = max(nx.algorithms.community.label_propagation.label_propagation_communities(H), key=len)
    # C = random_state.choice(G.nodes(), replace=False, size=100)

    for i, C in enumerate(sorted(label_propagation_communities(H),
                                 key=len, reverse=True)):

        if i >= num_subgraphs:
            click.secho(f"Finished drawing graphs for {num_subgraphs:d} communities!")
            break

        click.secho(f"Drawing graphs for {i:d}th largest community "
                    f"({len(C)} nodes)")

        classes_sub = classes[list(C)]

        H_sub = H.subgraph(C)
        G_sub = G.subgraph(C)
        P_sub = P.subgraph(C)

        probs = nx.get_edge_attributes(P_sub, name="weight")
        edge_width = edge_width_factor * np.fromiter(probs.values(), dtype="float64")

        # Compute layout
        click.secho("Computing graph drawing layout...")
        pos = nx.spring_layout(H_sub, seed=random_state)  # use the compromised graph

        # Compute graph difference
        D = nx.difference(H_sub, G_sub)
        click.secho(f"Compromised subgraph contains {D.number_of_edges()} "
                    "added edges")

        fig, ax = plt.subplots()

        nx.draw_networkx_nodes(G_sub, pos=pos,
                               node_size=node_size,
                               node_color=classes_sub,
                               alpha=node_alpha, cmap=node_cmap, ax=ax)
        if labels:
            nx.draw_networkx_labels(G_sub, pos=pos,
                                    font_size=font_size,
                                    font_color=font_color,
                                    font_family=font_family,
                                    font_weight=font_weight,
                                    alpha=label_alpha, ax=ax)
        nx.draw_networkx_edges(G_sub, pos=pos, alpha=edge_alpha, ax=ax)
        nx.draw_networkx_edges(D, pos=pos, edge_color=attacked_edge_color,
                               style=attacked_edge_style, alpha=edge_alpha, ax=ax)

        for ext in extension:
            fig.savefig(output_path.joinpath(f"{i:03d}_left_{context}_{suffix}.{ext}"),
                        bbox_inches="tight")

        plt.show()

        fig, ax = plt.subplots()

        nx.draw_networkx_nodes(P_sub, pos=pos,
                               node_size=node_size,
                               node_color=classes_sub,
                               alpha=node_alpha, cmap=node_cmap, ax=ax)
        if labels:
            nx.draw_networkx_labels(P_sub, pos=pos,
                                    font_size=font_size,
                                    font_color=font_color,
                                    font_family=font_family,
                                    font_weight=font_weight,
                                    alpha=label_alpha, ax=ax)

        # Probabilities as Edge Width
        # nx.draw_networkx_edges(P_sub, pos=pos, width=edge_width, alpha=edge_alpha,
        #                        ax=ax)

        # Probabilities as Edge Color
        # nx.draw_networkx_edges(P_sub, pos=pos, edge_color=edge_width,
        #                        edge_cmap=get_cmap("Blues"),
        #                        edge_vmin=0., edge_vmax=1.,
        #                        alpha=edge_alpha, ax=ax)

        # Probabilities as Edge Color Opacity
        for *e, w in P_sub.edges.data("weight"):
            nx.draw_networkx_edges(P_sub, pos=pos, edgelist=[e], alpha=w, ax=ax)

        for ext in extension:
            fig.savefig(output_path.joinpath(f"{i:03d}_right_{context}_{suffix}.{ext}"),
                        bbox_inches="tight")

        plt.show()

    return 0


if __name__ == "__main__":
    sys.exit(main())  # pragma: no cover