node2vec.py

"""This script implements the node2vec algorithm for learning node embeddings.

Most of the code is taken from the reference implementation. The code has been
updated to Python 3 (thanks also to Steve Corman for this issue
https://github.com/aditya-grover/node2vec/issues/124#issuecomment-1372624686).
"""

# pylint: disable=line-too-long,import-error

# Copyright (c) 2025 Emanuele Petriglia <inbox@emanuelepetriglia.com>
# Copyright (c) 2016 Aditya Grover
# All rights reserved. This file is licensed under the MIT license.


import argparse
import random
from pathlib import Path

import numpy as np
import networkx as nx
from gensim.models import Word2Vec


class Graph:
    """A graph representation for node2vec."""

    def __init__(self, nx_graph, is_directed, p, q):
        """Initialize the graph.

        Args:
            nx_graph: A networkx graph.
            is_directed: A boolean indicating whether the graph is directed.
            p: The return hyperparameter.
            q: The in-out hyperparameter.
        """
        self.graph = nx_graph
        self.is_directed = is_directed
        self.p = p
        self.q = q

        self.alias_nodes = None
        self.alias_edges = None

    def node2vec_walk(self, walk_length, start_node):
        """Simulate a random walk starting from start node.

        Args:
            walk_length: The length of the walk.
            start_node: The starting node for the walk.

        Returns:
            A list of nodes in the walk.
        """
        graph = self.graph
        alias_nodes = self.alias_nodes
        alias_edges = self.alias_edges

        walk = [start_node]

        while len(walk) < walk_length:
            cur = walk[-1]
            cur_nbrs = sorted(graph.neighbors(cur))
            if len(cur_nbrs) > 0:
                if len(walk) == 1:
                    walk.append(
                        cur_nbrs[alias_draw(alias_nodes[cur][0], alias_nodes[cur][1])]
                    )
                else:
                    prev = walk[-2]
                    next_elem = cur_nbrs[
                        alias_draw(
                            alias_edges[(prev, cur)][0], alias_edges[(prev, cur)][1]
                        )
                    ]
                    walk.append(next_elem)
            else:
                break

        return walk

    def simulate_walks(self, num_walks, walk_length):
        """Repeatedly simulate random walks from each node.

        Args:
            num_walks: The number of walks to simulate from each node.
            walk_length: The length of each walk.

        Returns:
            A list of walks.
        """
        graph = self.graph
        walks = []
        nodes = list(graph.nodes())
        print("Walk iteration:")
        for walk_iter in range(num_walks):
            print(str(walk_iter + 1), "/", str(num_walks))
            random.shuffle(nodes)
            for node in nodes:
                walks.append(
                    self.node2vec_walk(walk_length=walk_length, start_node=node)
                )

        return walks

    def get_alias_edge(self, src, dst):
        """Get the alias edge setup lists for a given edge.

        Args:
            src: The source node of the edge.
            dst: The destination node of the edge.

        Returns:
            A tuple of (alias, q) for the edge.
        """
        graph = self.graph
        p = self.p
        q = self.q

        unnormalized_probs = []
        for dst_nbr in sorted(graph.neighbors(dst)):
            if dst_nbr == src:
                unnormalized_probs.append(graph[dst][dst_nbr]["weight"] / p)
            elif graph.has_edge(dst_nbr, src):
                unnormalized_probs.append(graph[dst][dst_nbr]["weight"])
            else:
                unnormalized_probs.append(graph[dst][dst_nbr]["weight"] / q)
        norm_const = sum(unnormalized_probs)
        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

        return alias_setup(normalized_probs)

    def preprocess_transition_probs(self):
        """Preprocessing of transition probabilities for guiding the random walks."""
        graph = self.graph
        is_directed = self.is_directed

        alias_nodes = {}
        for node in graph.nodes():
            unnormalized_probs = [
                graph[node][nbr]["weight"] for nbr in sorted(graph.neighbors(node))
            ]
            norm_const = sum(unnormalized_probs)
            normalized_probs = [
                float(u_prob) / norm_const for u_prob in unnormalized_probs
            ]
            alias_nodes[node] = alias_setup(normalized_probs)

        alias_edges = {}

        if is_directed:
            for edge in graph.edges():
                alias_edges[edge] = self.get_alias_edge(edge[0], edge[1])
        else:
            for edge in graph.edges():
                alias_edges[edge] = self.get_alias_edge(edge[0], edge[1])
                alias_edges[(edge[1], edge[0])] = self.get_alias_edge(edge[1], edge[0])

        self.alias_nodes = alias_nodes
        self.alias_edges = alias_edges


def alias_setup(probs):
    """Compute utility lists for non-uniform sampling from discrete distributions.

    Refer to https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/
    for details.

    Args:
        probs: A list of probabilities.

    Returns:
        A tuple of (j, q) for alias sampling.
    """
    k = len(probs)
    q = np.zeros(k)
    j = np.zeros(k, dtype=np.int32)

    smaller = []
    larger = []
    for kk, prob in enumerate(probs):
        q[kk] = k * prob
        if q[kk] < 1.0:
            smaller.append(kk)
        else:
            larger.append(kk)

    while len(smaller) > 0 and len(larger) > 0:
        small = smaller.pop()
        large = larger.pop()

        j[small] = large
        q[large] = q[large] + q[small] - 1.0
        if q[large] < 1.0:
            smaller.append(large)
        else:
            larger.append(large)

    return j, q


def alias_draw(j, q):
    """Draw sample from a non-uniform discrete distribution using alias sampling.

    Args:
        j: The J table from alias_setup.
        q: The q table from alias_setup.

    Returns:
        A random sample from the distribution.
    """
    k = len(j)

    kk = int(np.floor(np.random.rand() * k))
    if np.random.rand() < q[kk]:
        return kk

    return j[kk]


def read_graph(graph_path, weighted, directed):
    """Reads the input network in networkx.

    Args:
        graph_path: The path to the graph file.
        weighted: A boolean indicating whether the graph is weighted.
        directed: A boolean indicating whether the graph is directed.

    Returns:
        A networkx graph.
    """
    if weighted:
        # The "data" parameter is not needed for read_weighted_edgelist.
        graph = nx.read_weighted_edgelist(
            graph_path, nodetype=int, create_using=nx.DiGraph()
        )
    else:
        graph = nx.read_edgelist(graph_path, nodetype=int, create_using=nx.DiGraph())
        for edge in graph.edges():
            graph[edge[0]][edge[1]]["weight"] = 1

    if not directed:
        graph = graph.to_undirected()

    return graph


# pylint: disable=too-many-arguments, too-many-positional-arguments
def learn_embeddings(walks, dimensions, window_size, workers, iters, output):
    """Learn embeddings by optimizing the Skipgram objective using SGD.

    Args:
        walks: A list of walks.
        dimensions: The number of dimensions for the embeddings.
        window_size: The context window size.
        workers: The number of parallel workers.
        iters: The number of epochs in SGD.
        output: The path to the output file.
    """
    walks = [list(map(str, walk)) for walk in walks]
    model = Word2Vec(
        walks,
        vector_size=dimensions,
        window=window_size,
        min_count=0,
        sg=1,
        negative=10,
        workers=workers,
        epochs=iters,
    )
    model.wv.save_word2vec_format(output)


# pylint: disable=redefined-outer-name
def main(args):
    """Pipeline for representational learning for all nodes in a graph.

    Args:
        args: The command line arguments.
    """
    nx_graph = read_graph(args.input, args.weighted, args.directed)
    graph = Graph(nx_graph, args.directed, args.p, args.q)
    graph.preprocess_transition_probs()
    walks = graph.simulate_walks(args.num_walks, args.walk_length)
    learn_embeddings(
        walks, args.dimensions, args.window_size, args.workers, args.iters, args.output
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run node2vec.")

    parser.add_argument(
        "--input", type=Path, required=True, help="Input graph path (.edgelist)."
    )

    parser.add_argument(
        "--output", type=Path, required=True, help="Embeddings path (.emb)."
    )

    parser.add_argument(
        "--dimensions",
        type=int,
        default=128,
        help="Number of dimensions. Default is 128.",
    )

    parser.add_argument(
        "--walk-length",
        type=int,
        default=80,
        help="Length of walk per source. Default is 80.",
    )

    parser.add_argument(
        "--num-walks",
        type=int,
        default=10,
        help="Number of walks per source. Default is 10.",
    )

    parser.add_argument(
        "--window-size",
        type=int,
        default=10,
        help="Context size for optimization. Default is 10.",
    )

    parser.add_argument("--iter", default=1, type=int, help="Number of epochs in SGD")

    parser.add_argument(
        "--workers",
        type=int,
        default=8,
        help="Number of parallel workers. Default is 8.",
    )

    parser.add_argument(
        "--p", type=float, default=1, help="Return hyperparameter. Default is 1."
    )

    parser.add_argument(
        "--q", type=float, default=1, help="Inout hyperparameter. Default is 1."
    )

    parser.add_argument(
        "--weighted",
        dest="weighted",
        action="store_true",
        help="Boolean specifying (un)weighted. Default is unweighted.",
    )
    parser.add_argument("--unweighted", dest="unweighted", action="store_false")
    parser.set_defaults(weighted=False)

    parser.add_argument(
        "--directed",
        dest="directed",
        action="store_true",
        help="Graph is (un)directed. Default is undirected.",
    )
    parser.add_argument("--undirected", dest="undirected", action="store_false")
    parser.set_defaults(directed=False)

    args = parser.parse_args()

    main(args)