Function main

examples/struc2vec_flight.py:51–83

(smoke=False, show=True)

def main(smoke=False, show=True):
    graph_path = SMOKE_GRAPH_PATH if smoke else FLIGHT_GRAPH_PATH
    graph = nx.read_edgelist(
        str(graph_path),
        create_using=nx.DiGraph(),
        nodetype=None,
        data=[("weight", int)],
    )

    with tempfile.TemporaryDirectory(prefix="struc2vec-") as temp_dir:
        model = Struc2Vec(
            graph,
            walk_length=3 if smoke else 10,
            num_walks=1 if smoke else 80,
            workers=1 if smoke else 4,
            verbose=0 if smoke else 40,
            temp_path=temp_dir + "/",
        )
        model.train(
            embed_size=8 if smoke else 128,
            window_size=2 if smoke else 5,
            workers=1,
            iter=1 if smoke else 3,
        )
        embeddings = model.get_embeddings()

    assert len(embeddings) > 0

    if not smoke:
        evaluate_embeddings(embeddings, FLIGHT_LABEL_PATH)
        plot_embeddings(embeddings, FLIGHT_LABEL_PATH, show=show)

    return embeddings


if __name__ == "__main__":