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
]