Source code for pygenstability.pygenstability

r"""PyGenStability code to solve generalized Markov Stability including Markov stability.

The generalized Markov Stability is of the form

.. math::

    Q_{gen}(t,H) = \mathrm{Tr} \left [H^T \left (F(t)-\sum_{k=1}^m v_{2k-1} v_{2k}^T\right)H\right]

where :math:`F(t)` is the quality matrix and :math:`v_k` are null model vectors.
The choice of the quality matrix and null model vectors are arbitrary in the generalized
Markov Stability setting, and can be parametrised via built-in constructors, or specified by
the user via the constructor module.
"""

from __future__ import annotations

import itertools
import logging
import multiprocessing
from collections import defaultdict
from functools import partial
from functools import wraps
from time import time
from typing import Any
from typing import Callable
from typing import Sequence

try:
    import igraph as ig
    import leidenalg

    _NO_LEIDEN = False
except ImportError:  # pragma: no cover
    _NO_LEIDEN = True

import numpy as np
import scipy.sparse as sp
from scipy.stats import entropy
from sklearn.metrics import mutual_info_score
from tqdm import tqdm

try:
    from pygenstability import generalized_louvain

    _NO_LOUVAIN = False
except ImportError:  # pragma: no cover
    _NO_LOUVAIN = True

from pygenstability.constructors import Constructor
from pygenstability.constructors import load_constructor
from pygenstability.io import save_results
from pygenstability.optimal_scales import identify_optimal_scales

L = logging.getLogger(__name__)
_DTYPE = np.float64
THRESHOLD = 1e-8


def _timing(f: Callable[..., Any]) -> Callable[..., Any]:  # pragma: no cover
    """Use as decorator to time a function execution if logging is in DEBUG mode."""

    @wraps(f)
    def wrap(*args: Any, **kw: Any) -> Any:
        if logging.root.level == logging.DEBUG:
            t_start = time()
            result = f(*args, **kw)
            t_end = time()
            with open("timing.csv", "a", encoding="utf-8") as file:
                print(f"{f.__name__}, {t_end - t_start}", file=file)
        else:
            result = f(*args, **kw)
        return result

    return wrap


def _get_chunksize(n_comp: int, pool: Any) -> int:
    """Split jobs accross workers for speedup."""
    return max(1, int(n_comp / pool._processes))  # pylint: disable=protected-access


def _graph_checks(graph: Any, dtype: Any = _DTYPE) -> Any:
    """Do some checks and preprocessing of the graph."""
    graph = sp.csr_matrix(graph, dtype=dtype)
    if sp.csgraph.connected_components(graph)[0] > 1:
        raise Exception(
            f"Graph not connected, with {sp.csgraph.connected_components(graph)[0]} components"
        )

    if sp.linalg.norm(graph - graph.T) > 0:
        L.warning("Your graph is directed!")

    if np.min(graph) < 0:
        L.warning("You have negative weights, consider using signed constructor.")

    graph.eliminate_zeros()
    return graph


def _get_scales(
    min_scale: float = -2.0,
    max_scale: float = 0.5,
    n_scale: int = 20,
    log_scale: bool = True,
    scales: np.ndarray | None = None,
) -> np.ndarray:
    """Get the scale vectors."""
    if scales is not None:
        return scales
    if log_scale:
        return np.logspace(min_scale, max_scale, n_scale)
    return np.linspace(min_scale, max_scale, n_scale)


def _get_params(all_locals: dict[str, Any]) -> dict[str, Any]:
    """Get run paramters from the local variables."""
    del all_locals["graph"]
    if hasattr(all_locals["constructor"], "get_data"):
        all_locals["constructor"] = "custom constructor"
    return all_locals


@_timing
def _get_constructor_data(
    constructor: Any,
    scales: np.ndarray,
    pool: Any,
    tqdm_disable: bool = False,
) -> list[dict[str, Any]]:
    return list(
        tqdm(
            pool.imap(constructor.get_data, scales),
            total=len(scales),
            disable=tqdm_disable,
        )
    )


def _check_method(method: str) -> str:  # pragma: no cover
    if _NO_LEIDEN and _NO_LOUVAIN:
        raise Exception("Without Louvain or Leiden solver, we cannot run PyGenStability")

    if method == "louvain" and _NO_LOUVAIN:
        L.warning("Louvain is not available, we fallback to leiden.")
        return "leiden"

    if method == "leiden" and _NO_LEIDEN:
        L.warning("Leiden is not available, we fallback to louvain.")
        return "louvain"

    return method


[docs] @_timing def run( graph: sp.spmatrix | None = None, constructor: str | Constructor = "linearized", min_scale: float = -2.0, max_scale: float = 0.5, n_scale: int = 20, log_scale: bool = True, scales: np.ndarray | None = None, n_tries: int = 100, with_all_tries: bool = False, with_NVI: bool = True, n_NVI: int = 20, with_postprocessing: bool = True, with_ttprime: bool = True, with_spectral_gap: bool = False, exp_comp_mode: str = "spectral", result_file: str = "results.pkl", n_workers: int = 4, tqdm_disable: bool = False, with_optimal_scales: bool = True, optimal_scales_kwargs: dict | None = None, method: str = "louvain", constructor_kwargs: dict | None = None, seed: int | None = None, ) -> dict: """This is the main function to compute graph clustering across scales with Markov Stability. This function needs a graph object as an adjacency matrix encoded with scipy.csgraph. The default settings will provide a fast and generic run with linearized Markov Stability, which corresponds to modularity with a scale parameter. Other built-in constructors are available to perform Markov Stability with matrix exponential computations. Custom constructors can be added via the constructor module. Additional parameters can be used to set the range and number of scales, number of trials for generalized Markov Stability optimisation, with Louvain or Leiden algorithm. Args: graph (scipy.csgraph): graph to cluster, if None, the constructor cannot be a str constructor (str/function): name of the generalized Markov Stability constructor, or custom constructor function. It must have two arguments, graph and scale. min_scale (float): minimum Markov scale max_scale (float): maximum Markov scale n_scale (int): number of scale steps log_scale (bool): use linear or log scales for scales scales (array): custom scale vector, if provided, it will override the other scale arguments n_tries (int): number of generalized Markov Stability optimisation evaluations with_all_tries (bools): store all partitions with stability values found in different optimisation evaluations with_NVI (bool): compute NVI(t) between generalized Markov Stability optimisations at each scale t n_NVI (int): number of randomly chosen generalized Markov Stability optimisations to estimate NVI with_postprocessing (bool): apply the final postprocessing step with_ttprime (bool): compute the NVI(t,tprime) matrix to compare scales t and tprime with_spectral_gap (bool): normalise scale by spectral gap exp_comp_mode (str): mode to compute matrix exponential, can be expm or spectral result_file (str): path to the result file n_workers (int): number of workers for multiprocessing tqdm_disable (bool): disable progress bars with_optimal_scales (bool): apply optimal scale selection algorithm optimal_scales_kwargs (dict): kwargs to pass to optimal scale selection, see optimal_scale module. method (str): optimisation method, louvain or leiden constructor_kwargs (dict): additional kwargs to pass to constructor prepare method seed (int): seed for the random number generator; pass an int for reproducible runs, or leave as None for a fresh non-deterministic seed. Returns: Results dict with the following entries - 'run_params': dict with parameters used to run the code - 'scales': scales of the scan - 'number_of_communities': number of communities at each scale - 'stability': value of stability cost function at each scale - 'community_id': community node labels at each scale - 'all_tries': all community node labels with stability values found in different optimisation evaluations at each scale (included if with_all_tries==True) - 'NVI': NVI(t) at each scale - 'ttprime': NVI(t,tprime) matrix - 'block_nvi': block NVI curve (included if with_optimal_scales==True) - 'selected_partitions': selected partitions (included if with_optimal_scales==True) """ method = _check_method(method) run_params = _get_params(locals()) rng = np.random.default_rng(seed) graph = _graph_checks(graph) scales = _get_scales( min_scale=min_scale, max_scale=max_scale, n_scale=n_scale, log_scale=log_scale, scales=scales, ) # sync n_scale to the resolved scan length, not the user kwarg n_scale = len(scales) run_params["n_scale"] = n_scale exp_comp_mode = _resolve_exp_comp_mode(exp_comp_mode, constructor) # avoid mutating a caller-supplied constructor_kwargs dict. constructor_kwargs = { **(constructor_kwargs or {}), "with_spectral_gap": with_spectral_gap, "exp_comp_mode": exp_comp_mode, } constructor_obj = load_constructor(constructor, graph, **constructor_kwargs) with multiprocessing.Pool(n_workers) as pool: L.info("Precompute constructors...") constructor_data = _get_constructor_data( constructor_obj, scales, pool, tqdm_disable=tqdm_disable ) L.info("Optimise stability...") scan_results = _scan_scales( constructor_data, scales, pool, rng, n_tries=n_tries, method=method, with_NVI=with_NVI, n_NVI=n_NVI, with_all_tries=with_all_tries, result_file=result_file, tqdm_disable=tqdm_disable, n_scale=n_scale, run_params=run_params, ) all_results = _run_post_scan_analysis( scan_results, pool, constructor_data, method=method, tqdm_disable=tqdm_disable, with_postprocessing=with_postprocessing, with_ttprime=with_ttprime, with_optimal_scales=with_optimal_scales, optimal_scales_kwargs=optimal_scales_kwargs, n_scale=n_scale, ) save_results(all_results, filename=result_file) return dict(all_results)
def _resolve_exp_comp_mode(exp_comp_mode: str, constructor: Any) -> str: """Validate exp_comp_mode and force expm for directed/signed constructors.""" assert exp_comp_mode in ["spectral", "expm"] if constructor in ("directed", "linearized_directed", "signed"): L.info("We cannot use spectral exponential computation for directed constructor") return "expm" return exp_comp_mode def _scan_scales( constructor_data: list[dict[str, Any]], scales: np.ndarray, pool: Any, rng: np.random.Generator, *, n_tries: int, method: str, with_NVI: bool, n_NVI: int, with_all_tries: bool, result_file: str, tqdm_disable: bool, n_scale: int, run_params: dict[str, Any], ) -> defaultdict: """Run the per-scale optimisation loop and aggregate results.""" all_results: defaultdict = defaultdict(list) all_results["run_params"] = run_params for i, t in tqdm(enumerate(scales), total=n_scale, disable=tqdm_disable): results = _run_optimisations(constructor_data[i], n_tries, pool, rng, method) communities = _process_runs(t, results, all_results) if with_NVI: _compute_NVI(communities, all_results, pool, n_partitions=min(n_NVI, n_tries)) if with_all_tries: all_results["all_tries"].append(results) save_results(all_results, filename=result_file) return all_results def _run_post_scan_analysis( all_results: defaultdict, pool: Any, constructor_data: list[dict[str, Any]], *, method: str, tqdm_disable: bool, with_postprocessing: bool, with_ttprime: bool, with_optimal_scales: bool, optimal_scales_kwargs: dict | None, n_scale: int, ) -> dict[str, Any]: """Apply postprocessing, ttprime, and optimal-scale selection.""" if with_postprocessing: L.info("Apply postprocessing...") _apply_postprocessing(all_results, pool, constructor_data, tqdm_disable, method=method) if with_ttprime or with_optimal_scales: L.info("Compute ttprimes...") _compute_ttprime(all_results, pool) if with_optimal_scales: L.info("Identify optimal scales...") if optimal_scales_kwargs is None: optimal_scales_kwargs = { "kernel_size": max(2, int(0.1 * n_scale)), "window_size": max(2, int(0.1 * n_scale)), "basin_radius": max(1, int(0.01 * n_scale)), } return identify_optimal_scales(all_results, **optimal_scales_kwargs) return all_results def _process_runs( scale: float, results: list[tuple[float, list[int]]], all_results: defaultdict, ) -> np.ndarray: """For each scale pick partition with highest stability among all iterations.""" # collect results from different optimisation runs stabilities = np.array([res[0] for res in results]) communities = np.array([res[1] for res in results]) # find index for highest stability best_run_id = np.argmax(stabilities) # save results for partition with highest stability all_results["scales"].append(scale) all_results["number_of_communities"].append(len(np.unique(communities[best_run_id]))) all_results["stability"].append(stabilities[best_run_id]) # we assign strictly increasing community IDs all_results["community_id"].append(_assign_increasing_ids(communities[best_run_id])) return communities def _assign_increasing_ids(community_id: np.ndarray) -> np.ndarray: """Assign strictly increasing community IDs starting from 0.""" community_id = np.asarray(community_id) unique_ids, first_ind, inverse = np.unique(community_id, return_index=True, return_inverse=True) relabel = np.empty(len(unique_ids), dtype=np.intp) relabel[np.argsort(first_ind)] = np.arange(len(unique_ids)) return relabel[inverse] @_timing def _compute_NVI( communities: np.ndarray, all_results: defaultdict, pool: Any, n_partitions: int = 10, ) -> None: """Compute NVI measure between the first n_partitions.""" # NVI is pairwise; with < 2 partitions there is nothing to compare. if n_partitions < 2: # pragma: no cover all_results["NVI"].append(0.0) return selected_partitions = communities[:n_partitions] # prepare worker to compute NVI between selected partitions worker = partial(evaluate_NVI, partitions=selected_partitions) # we compute pairwise NVI only for i != j because NVI is a metric index_pairs = list(itertools.combinations(range(n_partitions), 2)) chunksize = _get_chunksize(len(index_pairs), pool) # compute using pool of workers nvi_off_diagonal = list(pool.imap(worker, index_pairs, chunksize=chunksize)) # we compute the mean NVI, using the fact that NVI is a metric nvi_mean = 2 * np.sum(nvi_off_diagonal) / n_partitions**2 # append mean NVI to results all_results["NVI"].append(nvi_mean)
[docs] def evaluate_NVI(index_pair: Sequence[int], partitions: Any) -> float: r"""Evaluations of Normalized Variation of Information (NVI). NVI is defined for two partitions :math:`p_0` and :math:`p_1` as: .. math:: NVI = \frac{E(p_0) + E(p_1) - 2MI(p_0, p_1)}{JE(p_0, p_1)} where :math:`E` is the entropy, :math:`JE` the joint entropy and :math:`MI` the mutual information. Args: index_pair (list): list of two indices to select pairs of partitions partitions (list): list of partitions Returns: float, Normalized Variation Information """ p_0 = partitions[index_pair[0]] p_1 = partitions[index_pair[1]] MI = mutual_info_score(p_0, p_1) E_0 = entropy(np.unique(p_0, return_counts=True)[1]) E_1 = entropy(np.unique(p_1, return_counts=True)[1]) JE = E_0 + E_1 - MI if abs(JE) < 1e-8: return 0.0 return (JE - MI) / JE
def _to_indices( matrix: Any, directed: bool = False ) -> tuple[tuple[np.ndarray, np.ndarray], np.ndarray]: """Convert a sparse matrix to indices and values. The generalized stability ``Tr[H^T F H]`` depends only on the symmetric part of the quality matrix ``F``, so we symmetrise to ``(F + F^T) / 2`` first. This makes the Louvain backend (which only receives the lower triangle) and the Leiden backend (full directed graph) compute the same, correct edge term for asymmetric/directed quality matrices; for symmetric ``F`` it is a no-op. Args: matrix (sparse): sparse matrix to convert directed (bool): used for Leiden, which works if graph is full """ matrix = 0.5 * (matrix + matrix.T) if not directed: matrix = sp.tril(matrix) rows, cols, values = sp.find(matrix) return (rows, cols), values def _leiden_multiplex_layers( graph: Any, quality_values: np.ndarray, null_model: np.ndarray, initial_membership: list[int] | None = None, ) -> tuple[list[Any], list[float]]: r"""Build the Leiden multiplex layers encoding generalized Markov Stability. Each null pair ``(a, b)`` is split via the polarisation identity :math:`S_a S_b = \tfrac{1}{4}(S_{a+b}^2 - S_{a-b}^2)` into two CPM layers with ``node_sizes`` :math:`(a+b)/2` (weight ``+1``) and :math:`(a-b)/2` (weight ``-1``), which handles asymmetric and signed nulls. The edge term gets its own layer so it is counted once. For a symmetric pair (``a == b``) the second layer vanishes, recovering the single-layer trick of V. Traag (https://github.com/vtraag/leidenalg/pull/109#issuecomment-1283963065). """ n_nodes = graph.vcount() membership = {} if initial_membership is None else {"initial_membership": initial_membership} # edge/quality term in a dedicated layer with no null contribution (node_sizes = 0) partitions = [ leidenalg.CPMVertexPartition( graph, weights=quality_values, node_sizes=[0] * n_nodes, correct_self_loops=True, **membership, ) ] layer_weights = [1.0] # null model terms, with no edge contribution (weights = 0), via polarisation zero_weights = [0.0] * len(quality_values) null_model = np.asarray(null_model) for null_plus, null_minus in zip(null_model[::2], null_model[1::2]): for node_sizes, weight in ( (0.5 * (null_plus + null_minus), 1.0), (0.5 * (null_plus - null_minus), -1.0), ): # skip vanishing layers (e.g. the antisymmetric layer of a symmetric pair) if not np.any(node_sizes): continue partitions.append( leidenalg.CPMVertexPartition( graph, weights=zero_weights, node_sizes=node_sizes.tolist(), correct_self_loops=True, **membership, ) ) layer_weights.append(weight) return partitions, layer_weights @_timing def _optimise( try_idx: int, seed: int, quality_indices: tuple[np.ndarray, np.ndarray], quality_values: np.ndarray, null_model: np.ndarray, global_shift: float, method: str = "louvain", ) -> tuple[float, list[int]]: """Worker for generalized Markov Stability optimisation runs.""" if method == "louvain": # generalized_louvain is a pybind11 C++ extension; mypy can't see its symbols. stability, community_id = generalized_louvain.run_louvain( # type: ignore[attr-defined] quality_indices[0], quality_indices[1], quality_values, len(quality_values), null_model, np.shape(null_model)[0], 1.0, seed, ) elif method == "leiden": G = ig.Graph(edges=zip(*quality_indices), directed=True) partitions, layer_weights = _leiden_multiplex_layers(G, quality_values, null_model) optimiser = leidenalg.Optimiser() optimiser.set_rng_seed(int(seed)) # we initialise stability stability = sum(weight * p.quality() for weight, p in zip(layer_weights, partitions)) # we use Leiden to find optimal partition and update stability according to improvement stability += optimiser.optimise_partition_multiplex(partitions, layer_weights=layer_weights) community_id = partitions[0].membership else: # pragma: no cover raise ValueError(f"Unknown method {method!r}, expected 'louvain' or 'leiden'") return stability + global_shift, community_id def _evaluate_quality( partition_id: list[int], quality_indices: tuple[np.ndarray, np.ndarray], quality_values: np.ndarray, null_model: np.ndarray, global_shift: float, method: str = "louvain", ) -> float: """Worker for generalized Markov Stability evaluations.""" # evaluate using Louvain method if method == "louvain": quality = generalized_louvain.evaluate_quality( # type: ignore[attr-defined] quality_indices[0], quality_indices[1], quality_values, len(quality_values), null_model, np.shape(null_model)[0], 1.0, partition_id, ) # evaluate using Leiden method if method == "leiden": G = ig.Graph(edges=zip(*quality_indices), directed=True) partitions, layer_weights = _leiden_multiplex_layers( G, quality_values, null_model, initial_membership=partition_id ) quality = sum(weight * p.quality() for weight, p in zip(layer_weights, partitions)) return quality + global_shift def _run_optimisations( constructor: dict[str, Any], n_runs: int, pool: Any, rng: np.random.Generator, method: str = "louvain", ) -> list[tuple[float, list[int]]]: """Run several generalized Markov Stability optimisation on the current quality matrix.""" quality_indices, quality_values = _to_indices( constructor["quality"], directed=method == "leiden" ) worker = partial( _optimise, quality_indices=quality_indices, quality_values=quality_values, null_model=constructor["null_model"], global_shift=constructor.get("shift", 0.0), method=method, ) # seed each worker deterministically from the run-level rng seeds = rng.integers(0, int(1e8), size=n_runs).tolist() chunksize = _get_chunksize(n_runs, pool) return pool.starmap(worker, zip(range(n_runs), seeds), chunksize=chunksize) @_timing def _compute_ttprime(all_results: defaultdict, pool: Any) -> None: """Compute NVI(t,t') from the Markov stability results.""" # prepare worker to compute NVI between selected partitions worker = partial(evaluate_NVI, partitions=all_results["community_id"]) # we compute NVI only for t < t' because NVI is a metric index_pairs = list(itertools.combinations(range(len(all_results["scales"])), 2)) chunksize = _get_chunksize(len(index_pairs), pool) # compute NVI(t,t') for t < t' ttprime_list = pool.map(worker, index_pairs, chunksize=chunksize) # store NVI(t,'t) as symmetric matrix with zero diagonal all_results["ttprime"] = np.zeros([len(all_results["scales"]), len(all_results["scales"])]) for i, ttp in enumerate(ttprime_list): all_results["ttprime"][index_pairs[i][0], index_pairs[i][1]] = ttp all_results["ttprime"] += all_results["ttprime"].T @_timing def _apply_postprocessing( all_results: defaultdict, pool: Any, constructors: list[dict[str, Any]], tqdm_disable: bool = False, method: str = "louvain", ) -> None: """Apply postprocessing.""" all_results_raw = all_results.copy() # iterate through all scales for i, constructor in tqdm( enumerate(constructors), total=len(constructors), disable=tqdm_disable ): quality_indices, quality_values = _to_indices( constructor["quality"], directed=method == "leiden" ) # prepare _evaluate_quality() function for parallel processing worker = partial( _evaluate_quality, quality_indices=quality_indices, quality_values=quality_values, null_model=constructor["null_model"], global_shift=constructor.get("shift", 0.0), method=method, ) # find index in sequence of partitions that leads to highest quality score quality_scores = pool.map( worker, all_results_raw["community_id"], chunksize=_get_chunksize(len(all_results_raw["community_id"]), pool), ) best_quality_id = np.argmax(quality_scores) # only if the new found score is better update partition if np.abs(quality_scores[best_quality_id] - all_results["stability"][i]) > THRESHOLD: # replace old partition with new partition all_results["community_id"][i] = all_results_raw["community_id"][best_quality_id] # assign new quality score all_results["stability"][i] = quality_scores[best_quality_id] # update number of communities all_results["number_of_communities"][i] = all_results_raw["number_of_communities"][ best_quality_id ]