diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index d1ef347..c38cc5e 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -5,17 +5,16 @@ name: Python package on: push: - branches: [ master ] + branches: [master] pull_request: - branches: [ master ] + branches: [master] jobs: build: - runs-on: ubuntu-latest strategy: matrix: - python-version: [ "3.9", "3.10", "3.11", "3.12", "3.13", "3.14", "3.15" ] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"] steps: - uses: actions/checkout@v3 @@ -39,6 +38,13 @@ jobs: flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude env/ # exit-zero treats all errors as warnings flake8 . --count --exit-zero --max-complexity=15 --max-line-length=120 --statistics --exclude env/ + - name: Type checking with mypy + # numpy no longer supports <3.12, so typing is unreliable + if: ${{ !contains(fromJSON('["3.10", "3.11"]'), matrix.python-version) }} + run: | + source env/bin/activate + python -m pip install mypy + python -m mypy utilities/ --install-types --non-interactive - name: Test with pytest run: | source env/bin/activate diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000..22f0816 --- /dev/null +++ b/mypy.ini @@ -0,0 +1,12 @@ +[mypy] +strict = True +disallow_any_explicit = True +files = utilities/*.py, tests/*.py + +[mypy-sklearn.*,igraph.*,leidenalg.*] +ignore_missing_imports = True + +# tests are untyped, but we want to check their usage is correct +[mypy-tests.*] +allow_untyped_defs = True +allow_untyped_calls = True \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index be4c6d2..9cefb8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "modularitypruning" version = "1.5.1" description = "Pruning tool to identify small subsets of network partitions that are significant from the perspective of stochastic block model inference." readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.10" license = "MIT" authors = [ { name = "Ryan Gibson", email = "ryan.alex.gibson@gmail.com" } @@ -17,14 +17,14 @@ classifiers = [ "Development Status :: 5 - Production/Stable", "Topic :: Scientific/Engineering :: Information Analysis", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", "Programming Language :: Python :: 3.15", - "Programming Language :: Python :: 3 :: Only" + "Programming Language :: Python :: 3 :: Only", + "Typing :: Typed" ] dependencies = [ "leidenalg", @@ -44,5 +44,8 @@ packages = ["modularitypruning"] [tool.setuptools.package-dir] modularitypruning = "utilities" +[tool.setuptools.package-data] +modularitypruning = ["py.typed"] + [tool.autopep8] max_line_length = 120 \ No newline at end of file diff --git a/tests/shared_testing_functions.py b/tests/shared_testing_functions.py index ab84e58..81538d5 100644 --- a/tests/shared_testing_functions.py +++ b/tests/shared_testing_functions.py @@ -86,7 +86,7 @@ def generate_multilayer_intralayer_SBM(copying_probability, p_in, p_out, first_l if random() < p_out: intralayer_edges.append((u, v)) - G_intralayer = ig.Graph(intralayer_edges, directed=False) + G_intralayer = ig.Graph(n=len(combined_community_labels), edges=intralayer_edges, directed=False) return G_intralayer, layer_membership diff --git a/tests/test_documentation_examples.py b/tests/test_documentation_examples.py index a1e3d86..9ec6632 100644 --- a/tests/test_documentation_examples.py +++ b/tests/test_documentation_examples.py @@ -256,7 +256,7 @@ def test_plot_multiplex_community(self): ax = plot_multiplex_community(np.array(membership), np.array(layer_vec)) ax.set_xticks(np.linspace(0, num_layers, 2 * num_layers + 1)) ax.set_xticklabels(["", "Advice", "", "Coworker", "", "Friend", ""], fontsize=14) - plt.title(f"Multiplex Communities", fontsize=14) + plt.title("Multiplex Communities", fontsize=14) plt.ylabel("Node ID", fontsize=14) plt.close() # closing this these figures instead of showing diff --git a/tests/test_parallel_leiden_performance.py b/tests/test_parallel_leiden_performance.py index 21b5b62..d4a3fa6 100644 --- a/tests/test_parallel_leiden_performance.py +++ b/tests/test_parallel_leiden_performance.py @@ -44,8 +44,8 @@ def determine_target_parallelization_speedup(num_calculations=32): cpu_utilization = psutil.cpu_percent(interval=5) if cpu_utilization > 20: PERFORMANCE_TARGET_RELATIVE_TO_PERFECT_SCALING = 0.5 - warnings.warn(f"System CPU utilization is non-negligible during parallel performance test! " - f"Dropping performance scaling target to 50%.") + warnings.warn("System CPU utilization is non-negligible during parallel performance test! " + "Dropping performance scaling target to 50%.") start_time = time() _ = [mock_calculation(i) for i in range(num_calculations)] diff --git a/tests/test_temporal_multilevel_parameter_estimation.py b/tests/test_temporal_multilevel_parameter_estimation.py index bfce87e..e88f36c 100644 --- a/tests/test_temporal_multilevel_parameter_estimation.py +++ b/tests/test_temporal_multilevel_parameter_estimation.py @@ -77,6 +77,7 @@ def retry_test_across_graph_sizes(self, test_fn, sizes=[50, 100, 250, 500], retr if not isinstance(e, AssertionError): # test failed in unusual way (perhaps keyboard interrupt?), raising immediately raise e + assert last_error is not None raise last_error def test_temporal_SBM_correct_convergence_varying_copying_probabilty(self): diff --git a/utilities/__init__.py b/utilities/__init__.py index 4a60dff..7405899 100644 --- a/utilities/__init__.py +++ b/utilities/__init__.py @@ -1,2 +1,7 @@ # this makes prune_to_stable_partitions and prune_to_multilayer_stable_partitions importable from the top-level package -from .parameter_estimation_utilities import prune_to_stable_partitions, prune_to_multilayer_stable_partitions # noqa +from .parameter_estimation_utilities import ( + prune_to_multilayer_stable_partitions as prune_to_multilayer_stable_partitions, + prune_to_stable_partitions as prune_to_stable_partitions, +) + +__all__ = ["prune_to_stable_partitions", "prune_to_multilayer_stable_partitions"] diff --git a/utilities/champ_utilities.py b/utilities/champ_utilities.py index 7dc169c..37b79d9 100644 --- a/utilities/champ_utilities.py +++ b/utilities/champ_utilities.py @@ -1,3 +1,7 @@ +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import igraph + import multiprocessing.pool import warnings from collections import defaultdict from math import floor @@ -11,9 +15,13 @@ from .partition_utilities import all_degrees, in_degrees, out_degrees, membership_to_communities, \ membership_to_layered_communities +from .types import (Domain, FloatArray, GammaRange, IntArray, LayerVector, Membership, + Partitions, PartitionCoefficients2D, PartitionCoefficients3D) -def get_interior_point(halfspaces, initial_num_sampled=50, full_retry_limit=10): +def get_interior_point( + halfspaces: FloatArray, initial_num_sampled: int = 50, full_retry_limit: int = 10 +) -> FloatArray: """ Find interior point of halfspaces (needed to perform halfspace intersection) @@ -41,19 +49,21 @@ def get_interior_point(halfspaces, initial_num_sampled=50, full_retry_limit=10): else: raise ValueError(f"get_interior_point received unhandled dimension {normals.shape[1]}!") - num_retries = 0 + num_retries: int = 0 while True: # retry until success # randomly sample some of the halfspaces - sample_len = min(initial_num_sampled, len(interior_hs)) # len(interior_hs) - sampled_hs = np.vstack((interior_hs[choice(interior_hs.shape[0], sample_len, replace=False)], boundaries)) + sample_len: int = min(initial_num_sampled, len(interior_hs)) # len(interior_hs) + sampled_hs: FloatArray = np.vstack(( + interior_hs[choice(interior_hs.shape[0], sample_len, replace=False)], boundaries)) # compute the Chebyshev center of the sampled halfspaces' intersection - norm_vector = np.reshape(np.linalg.norm(sampled_hs[:, :-1], axis=1), (sampled_hs.shape[0], 1)) - c = np.zeros((sampled_hs.shape[1],)) + norm_vector: FloatArray = np.reshape( + np.linalg.norm(sampled_hs[:, :-1], axis=1), (sampled_hs.shape[0], 1)) + c: FloatArray = np.zeros((sampled_hs.shape[1],)) c[-1] = -1 - A = np.hstack((sampled_hs[:, :-1], norm_vector)) - b = -sampled_hs[:, -1:] + A: FloatArray = np.hstack((sampled_hs[:, :-1], norm_vector)) + b: FloatArray = -sampled_hs[:, -1:] res = linprog(c, A_ub=A, b_ub=b, bounds=(-np.inf, np.inf), method='highs-ipm') @@ -65,7 +75,10 @@ def get_interior_point(halfspaces, initial_num_sampled=50, full_retry_limit=10): # 3 : Problem appears to be unbounded. # 4 : Numerical difficulties encountered. if res.success or res.status in {0, 1, 4}: - intpt = res.x[:-1] # res.x contains [interior_point, distance to enclosing polyhedron] + if res.x is None: + # satisfying mypy, this should never occur + raise ValueError(f"Scipy's linprog claimed success but returned no solution! {res}") + intpt: FloatArray = res.x[:-1] # res.x contains [interior_point, distance to enclosing polyhedron] # ensure that the computed point is actually interior to all halfspaces if (np.dot(normals, intpt) + np.transpose(offsets) < 0).all(): @@ -86,7 +99,10 @@ def get_interior_point(halfspaces, initial_num_sampled=50, full_retry_limit=10): return intpt -def CHAMP_2D(G, all_parts, gamma_0, gamma_f, single_threaded=False): +def CHAMP_2D( + G: "igraph.Graph", all_parts: Partitions, gamma_0: float, gamma_f: float, + single_threaded: bool = False +) -> list[GammaRange]: r"""Calculates the pruned set of partitions from CHAMP on ``gamma_0`` :math:`\leq \gamma \leq` ``gamma_f``. See https://doi.org/10.3390/a10030093 for more details. @@ -94,7 +110,7 @@ def CHAMP_2D(G, all_parts, gamma_0, gamma_f, single_threaded=False): :param G: graph of interest :type G: igraph.Graph :param all_parts: partitions to prune - :type all_parts: list[tuple] or set[tuple] + :type all_parts: Sequence[tuple[int, ...]] | set[tuple[int, ...]] :param gamma_0: starting gamma value for CHAMP :type gamma_0: float :param gamma_f: ending gamma value for CHAMP @@ -106,41 +122,42 @@ def CHAMP_2D(G, all_parts, gamma_0, gamma_f, single_threaded=False): - starting gamma value for the partition's domain of optimality - ending gamma value for the partition's domain of optimality - community membership tuple for the partition - :rtype: list of tuple[float, float, tuple[int]] + :rtype: list of tuple[float, float, tuple[int, ...]] """ if len(all_parts) == 0: return [] all_parts = list(all_parts) - num_partitions = len(all_parts) + num_partitions: int = len(all_parts) - partition_coefficients = partition_coefficients_2D(G, all_parts, single_threaded=single_threaded) + partition_coefficients: PartitionCoefficients2D = partition_coefficients_2D( + G, all_parts, single_threaded=single_threaded) A_hats, P_hats = partition_coefficients # add on boundaries: bottom < Q < top, gamma_0 < gamma < gamma_f - top = max(A_hats - P_hats * gamma_0) # Could potentially be optimized - bottom = min(A_hats - P_hats * gamma_f) # Could potentially be optimized - halfspaces = np.vstack((halfspaces_from_coefficients_2D(*partition_coefficients), - np.array([[0, 1, -top], [0, -1, bottom], - [1, 0, -gamma_f], [-1, 0, gamma_0]]))) + top: np.float64 = max(A_hats - P_hats * gamma_0) # Could potentially be optimized + bottom: np.float64 = min(A_hats - P_hats * gamma_f) # Could potentially be optimized + halfspaces: FloatArray = np.vstack((halfspaces_from_coefficients_2D(*partition_coefficients), + np.array([[0, 1, -top], [0, -1, bottom], + [1, 0, -gamma_f], [-1, 0, gamma_0]]))) # normalize halfspaces to try to improve ill-conditioned linear programs halfspaces = halfspaces / np.linalg.norm(halfspaces, axis=1, keepdims=True) # Could potentially scale axes so Chebyshev center is better for problem - interior_point = get_interior_point(halfspaces) - hs = HalfspaceIntersection(halfspaces, interior_point) + interior_point: FloatArray = get_interior_point(halfspaces) + hs: HalfspaceIntersection = HalfspaceIntersection(halfspaces, interior_point) # scipy does not support facets by halfspace directly, so we must compute them - facets_by_halfspace = defaultdict(list) + facets_by_halfspace: defaultdict[int, list[FloatArray]] = defaultdict(list) for v, idx in zip(hs.intersections, hs.dual_facets): assert np.isfinite(v).all() for i in idx: if i < num_partitions: facets_by_halfspace[i].append(v) - ranges = [] + ranges: list[GammaRange] = [] for i, intersections in facets_by_halfspace.items(): x1, x2 = intersections[0][0], intersections[1][0] if x1 > x2: @@ -150,7 +167,10 @@ def CHAMP_2D(G, all_parts, gamma_0, gamma_f, single_threaded=False): return sorted(ranges, key=lambda x: x[0]) -def CHAMP_3D(G_intralayer, G_interlayer, layer_vec, all_parts, gamma_0, gamma_f, omega_0, omega_f): +def CHAMP_3D( + G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: LayerVector, + all_parts: Partitions, gamma_0: float, gamma_f: float, omega_0: float, omega_f: float +) -> list[Domain]: r"""Calculates the pruned set of partitions from CHAMP on ``gamma_0`` :math:`\leq \gamma \leq` ``gamma_f`` and ``omega_0`` :math:`\leq \omega \leq` ``omega_f``. @@ -163,7 +183,7 @@ def CHAMP_3D(G_intralayer, G_interlayer, layer_vec, all_parts, gamma_0, gamma_f, :param layer_vec: list of each vertex's layer membership :type layer_vec: list[int] :param all_parts: partitions to prune - :type all_parts: iterable[tuple] + :type all_parts: Sequence[tuple[int, ...]] | set[tuple[int, ...]] :param gamma_0: starting gamma value for CHAMP :type gamma_0: float :param gamma_f: ending gamma value for CHAMP @@ -176,54 +196,56 @@ def CHAMP_3D(G_intralayer, G_interlayer, layer_vec, all_parts, gamma_0, gamma_f, - list of polygon vertices in (gamma, omega) plane for the partition's domain of optimality - community membership tuple for the partition - :rtype: list of tuple[list[float], tuple[int]] + :rtype: list of tuple[list[numpy.ndarray], tuple[int, ...]] """ if len(all_parts) == 0: return [] all_parts = list(all_parts) - num_partitions = len(all_parts) + num_partitions: int = len(all_parts) - partition_coefficients = partition_coefficients_3D(G_intralayer, G_interlayer, layer_vec, all_parts) + partition_coefficients: PartitionCoefficients3D = partition_coefficients_3D( + G_intralayer, G_interlayer, layer_vec, all_parts) A_hats, P_hats, C_hats = partition_coefficients # add on boundaries: bottom < Q < top, gamma_0 < gamma < gamma_f, omega_0 < omega < omega_f - top = max(A_hats - P_hats * gamma_0 + C_hats * omega_f) # Could potentially be optimized - bottom = min(A_hats - P_hats * gamma_f + C_hats * omega_0) # Could potentially be optimized - halfspaces = np.vstack((halfspaces_from_coefficients_3D(*partition_coefficients), - (np.array([[0, 0, 1, -top], [0, 0, -1, bottom], - [1, 0, 0, -gamma_f], [-1, 0, 0, gamma_0], - [0, 1, 0, -omega_f], [0, -1, 0, omega_0]])))) + top: np.float64 = max(A_hats - P_hats * gamma_0 + C_hats * omega_f) # Could potentially be optimized + bottom: np.float64 = min(A_hats - P_hats * gamma_f + C_hats * omega_0) # Could potentially be optimized + halfspaces: FloatArray = np.vstack((halfspaces_from_coefficients_3D(*partition_coefficients), + (np.array([[0, 0, 1, -top], [0, 0, -1, bottom], + [1, 0, 0, -gamma_f], [-1, 0, 0, gamma_0], + [0, 1, 0, -omega_f], [0, -1, 0, omega_0]])))) # normalize halfspaces to try to improve ill-conditioned linear programs halfspaces = halfspaces / np.linalg.norm(halfspaces, axis=1, keepdims=True) # Could potentially scale axes so Chebyshev center is better for problem - interior_point = get_interior_point(halfspaces) - hs = HalfspaceIntersection(halfspaces, interior_point) + interior_point: FloatArray = get_interior_point(halfspaces) + hs: HalfspaceIntersection = HalfspaceIntersection(halfspaces, interior_point) # scipy does not support facets by halfspace directly, so we must compute them - facets_by_halfspace = defaultdict(list) + facets_by_halfspace: defaultdict[int, list[FloatArray]] = defaultdict(list) for v, idx in zip(hs.intersections, hs.dual_facets): assert np.isfinite(v).all() for i in idx: if i < num_partitions: facets_by_halfspace[i].append(v) - domains = [] + domains: list[Domain] = [] for i, intersections in facets_by_halfspace.items(): # drop quality dimension / project into (gamma, omega) plane intersections = [x[:2] for x in intersections] # sort the domain's points counter-clockwise around the centroid - centroid = np.mean(intersections, axis=0) + centroid: FloatArray = np.mean(intersections, axis=0) intersections.sort(key=lambda x: np.arctan2(x[0] - centroid[0], x[1] - centroid[1])) domains.append((intersections, all_parts[i])) return domains -def partition_coefficients_2D_serial(G, partitions): +def partition_coefficients_2D_serial(G: "igraph.Graph", partitions: Partitions + ) -> PartitionCoefficients2D: """Computes A_hat and P_hat for partitions of :G: TODO: support edge weights""" @@ -255,21 +277,24 @@ def partition_coefficients_2D_serial(G, partitions): return A_hats, P_hats -def partition_coefficients_2D(G, partitions, single_threaded=False): +def partition_coefficients_2D( + G: "igraph.Graph", partitions: Partitions, single_threaded: bool = False +) -> PartitionCoefficients2D: """Computes partitions coefficients in parallel by calling partition_coefficients_2D_serial""" partitions = list(partitions) if single_threaded: A_hats, P_hats = partition_coefficients_2D_serial(G, partitions) else: - partition_chunks = [ + partition_chunks: list[list[Membership]] = [ partitions[floor(i * len(partitions) / cpu_count()):floor((i + 1) * len(partitions) / cpu_count())] for i in range(cpu_count()) ] - pool = Pool(processes=cpu_count()) - results = pool.starmap(partition_coefficients_2D_serial, - [(G, partition_chunk) for partition_chunk in partition_chunks]) + pool: "multiprocessing.pool.Pool" = Pool(processes=cpu_count()) + results: list[PartitionCoefficients2D] = \ + pool.starmap(partition_coefficients_2D_serial, + [(G, partition_chunk) for partition_chunk in partition_chunks]) pool.close() A_hats = np.array([v for A_hats, P_hats in results for v in A_hats]) @@ -280,27 +305,33 @@ def partition_coefficients_2D(G, partitions, single_threaded=False): return A_hats, P_hats -def halfspaces_from_coefficients_2D(A_hats, P_hats): +def halfspaces_from_coefficients_2D( + A_hats: IntArray, P_hats: FloatArray +) -> FloatArray: """Converts partitions' coefficients to halfspace normal, offset Q >= -P_hat*gamma + A_hat -Q - P_hat*gamma + A_hat <= 0 (-P_hat, -1) * (gamma, Q) + A_hat <= 0 """ - return np.vstack((-P_hats, -np.ones_like(P_hats), A_hats)).T + return np.vstack((-P_hats, -np.ones_like(P_hats), A_hats)).T.astype(np.float64) -def halfspaces_from_coefficients_3D(A_hats, P_hats, C_hats): +def halfspaces_from_coefficients_3D( + A_hats: IntArray, P_hats: FloatArray, C_hats: IntArray +) -> FloatArray: """Converts partitions' coefficients to halfspace normal, offset Q >= C_hat*omega - P_hat*gamma + A_hat -Q + C_hat*omega - P_hat*gamma + A_hat <= 0 (-P_hat, C_hat, -1) * (gamma, omega, Q) + A_hat <= 0 """ - return np.vstack((-P_hats, C_hats, -np.ones_like(P_hats), A_hats)).T + return np.vstack((-P_hats, C_hats, -np.ones_like(P_hats), A_hats)).T.astype(np.float64) -def partition_coefficients_3D_serial(G_intralayer, G_interlayer, layer_vec, partitions): +def partition_coefficients_3D_serial(G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", + layer_vec: LayerVector, + partitions: list[Membership]) -> PartitionCoefficients3D: """Computes A_hat, P_hat, C_hat for partitions of a graph with intralayer edges given in :G_intralayer:, interlayer edges given in :G_interlayer:, and layer membership :layer_vec: @@ -349,23 +380,26 @@ def partition_coefficients_3D_serial(G_intralayer, G_interlayer, layer_vec, part return A_hats, P_hats, C_hats -def partition_coefficients_3D(G_intralayer, G_interlayer, layer_vec, partitions): +def partition_coefficients_3D( + G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: LayerVector, + partitions: Partitions +) -> PartitionCoefficients3D: """Computes partitions coefficients in parallel by calling partition_coefficients_3D_serial""" partitions = list(partitions) - partition_chunks = [ + partition_chunks: list[list[Membership]] = [ partitions[floor(i * len(partitions) / cpu_count()):floor((i + 1) * len(partitions) / cpu_count())] for i in range(cpu_count()) ] - pool = Pool(processes=cpu_count()) - results = pool.starmap(partition_coefficients_3D_serial, - [(G_intralayer, G_interlayer, layer_vec, partition_chunk) - for partition_chunk in partition_chunks]) + pool: "multiprocessing.pool.Pool" = Pool(processes=cpu_count()) + results: list[PartitionCoefficients3D] = pool.starmap( + partition_coefficients_3D_serial, [(G_intralayer, G_interlayer, layer_vec, partition_chunk) + for partition_chunk in partition_chunks]) pool.close() - A_hats = np.array([v for A_hats, P_hats, C_hats in results for v in A_hats]) - P_hats = np.array([v for A_hats, P_hats, C_hats in results for v in P_hats]) - C_hats = np.array([v for A_hats, P_hats, C_hats in results for v in C_hats]) + A_hats: IntArray = np.array([v for A_hats, P_hats, C_hats in results for v in A_hats]) + P_hats: FloatArray = np.array([v for A_hats, P_hats, C_hats in results for v in P_hats]) + C_hats: IntArray = np.array([v for A_hats, P_hats, C_hats in results for v in C_hats]) assert len(A_hats) == len(P_hats) == len(C_hats) == len(partitions) diff --git a/utilities/leiden_utilities.py b/utilities/leiden_utilities.py index 097008b..d060884 100644 --- a/utilities/leiden_utilities.py +++ b/utilities/leiden_utilities.py @@ -1,3 +1,5 @@ +from collections.abc import Iterable +from typing import overload, Literal import functools from multiprocessing import Pool, cpu_count @@ -6,25 +8,47 @@ import numpy as np from tqdm import tqdm -LOW_MEMORY_THRESHOLD = 1e9 # 1 GB +from .types import (FloatSequence, IntSequence, LayerTuple, LayerVector, LeidenPartition, + Membership, MultilayerLeidenParts, MultilayerLeidenResult, SingleLayerLeidenResult) + +LOW_MEMORY_THRESHOLD: float = 1e9 # 1 GB @functools.lru_cache(maxsize=1000) -def sorted_tuple(t): +def sorted_tuple(t: Membership) -> Membership: """Converts a tuple to a canonical form in which the labels' first occurrences are sorted (e.g. label 0 will always occur before label 1 in the tuple). :param t: community membership of a partition - :type t: tuple[int] + :type t: tuple[int, ...] :return: a canonical representation of the membership tuple - :rtype: tuple[int] + :rtype: tuple[int, ...] """ - sort_map = {x[0]: i for i, x in enumerate(sorted(zip(*np.unique(t, return_index=True)), key=lambda x: x[1]))} + sort_map: dict[int, int] = ( + {x[0]: i for i, x in enumerate(sorted(zip(*np.unique(t, return_index=True)), key=lambda x: x[1]))} + ) return tuple(sort_map[x] for x in t) -def singlelayer_leiden(G, gamma, return_partition=False): +# explain to mypy that the return type depends on the return_partition argument +@overload +def singlelayer_leiden( + G: ig.Graph, gamma: float, *, return_partition: Literal[True], +) -> LeidenPartition: + ... + + +@overload +def singlelayer_leiden( + G: ig.Graph, gamma: float, *, return_partition: Literal[False] = False, +) -> Membership: + ... + + +def singlelayer_leiden( + G: ig.Graph, gamma: float, return_partition: bool = False, +) -> SingleLayerLeidenResult: r"""Run the Leiden modularity maximization algorithm at a single :math:`\gamma` value. :param G: graph of interest @@ -34,13 +58,14 @@ def singlelayer_leiden(G, gamma, return_partition=False): :param return_partition: if True, return a leidenalg partition. Otherwise, return a community membership tuple :type return_partition: bool :return: partition from leidenalg - :rtype: tuple[int] or leidenalg.RBConfigurationVertexPartition + :rtype: tuple[int, ...] or leidenalg.RBConfigurationVertexPartition """ if 'weight' not in G.es: G.es['weight'] = [1.0] * G.ecount() - partition = leidenalg.find_partition(G, leidenalg.RBConfigurationVertexPartition, weights='weight', - resolution_parameter=gamma) + partition: LeidenPartition = leidenalg.find_partition( + G, leidenalg.RBConfigurationVertexPartition, weights='weight', resolution_parameter=gamma + ) if return_partition: return partition @@ -48,24 +73,26 @@ def singlelayer_leiden(G, gamma, return_partition=False): return tuple(partition.membership) -def _wrapped_singlelayer_leiden(args): +def _wrapped_singlelayer_leiden(args: tuple[ig.Graph, float]) -> Membership: """Wrapped singlelayer_leiden() for use in multiprocessing.Pool.imap_unordered.""" return singlelayer_leiden(*args) -def leiden_part(G): +def leiden_part(G: ig.Graph) -> LeidenPartition: return leidenalg.RBConfigurationVertexPartition(G) -def leiden_part_with_membership(G, membership): +def leiden_part_with_membership( + G: ig.Graph, membership: IntSequence +) -> LeidenPartition: if isinstance(membership, np.ndarray): membership = membership.tolist() - part = leiden_part(G) + part: LeidenPartition = leiden_part(G) part.set_membership(membership) return part -def split_intralayer_leiden_graph(G_intralayer, layer_membership): +def split_intralayer_leiden_graph(G_intralayer: ig.Graph, layer_membership: LayerVector) -> list[ig.Graph]: """ Split intralayer network into a separate network for each layer. @@ -73,8 +100,8 @@ def split_intralayer_leiden_graph(G_intralayer, layer_membership): :param G_intralayer: intralayer graph of interest :type G_intralayer: igraph.Graph - :param layer_vec: list of each vertex's layer membership - :type layer_vec: list[int] + :param layer_membership: list of each vertex's layer membership + :type layer_membership: list[int] :return: list of intralayer networks :rtype: list[igraph.Graph] """ @@ -85,25 +112,47 @@ def split_intralayer_leiden_graph(G_intralayer, layer_membership): @functools.lru_cache(maxsize=1) -def _split_leiden_graph_layers_cached(n, G_es, is_directed, layer_membership): - T = max(layer_membership) + 1 +def _split_leiden_graph_layers_cached( + n: int, G_es: tuple[ig.Edge, ...], is_directed: bool, layer_membership: LayerTuple +) -> list[ig.Graph]: + T: int = max(layer_membership) + 1 - edges_by_layer = [[] for _ in range(T)] - weights_by_layer = [[] for _ in range(T)] + edges_by_layer: list[list[tuple[int, int]]] = [[] for _ in range(T)] + weights_by_layer: list[list[float]] = [[] for _ in range(T)] for e in G_es: - e_layer = layer_membership[e.source] + e_layer: int = layer_membership[e.source] edges_by_layer[e_layer].append((e.source, e.target)) weights_by_layer[e_layer].append(e['weight']) - G_split_layers_list = [] + G_split_layers_list: list[ig.Graph] = [] for layer_idx in range(T): - G_layer = ig.Graph(n=n, edges=edges_by_layer[layer_idx], directed=is_directed) + G_layer: ig.Graph = ig.Graph(n=n, edges=edges_by_layer[layer_idx], directed=is_directed) G_layer.es['weight'] = weights_by_layer[layer_idx] G_split_layers_list.append(G_layer) return G_split_layers_list -def multilayer_leiden(G_intralayer, G_interlayer, layer_vec, gamma, omega, optimiser=None, return_partition=False): +# explain to mypy that the return type depends on the return_partition argument +@overload +def multilayer_leiden( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector, gamma: float, + omega: float, optimiser: leidenalg.Optimiser | None = None, *, return_partition: Literal[True], +) -> list[LeidenPartition]: + ... + + +@overload +def multilayer_leiden( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector, gamma: float, + omega: float, optimiser: leidenalg.Optimiser | None = None, *, return_partition: Literal[False] = False, +) -> Membership: + ... + + +def multilayer_leiden( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector, gamma: float, + omega: float, optimiser: leidenalg.Optimiser | None = None, return_partition: bool = False, +) -> MultilayerLeidenResult: r"""Run the Leiden modularity maximization algorithm at a single (:math:`\gamma, \omega`) value. WARNING: Optimization can be EXTREMELY slow for large numbers of layers! Leidenalg does not properly implement @@ -120,11 +169,11 @@ def multilayer_leiden(G_intralayer, G_interlayer, layer_vec, gamma, omega, optim :param omega: omega (interlayer resolution parameter) to run Leiden at :type omega: float :param optimiser: if not None, use passed-in (potentially custom) leidenalg optimiser - :type optimiser: leidenalg.Optimiser + :type optimiser: leidenalg.Optimiser | None :param return_partition: if True, return a leidenalg partition. Otherwise, return a community membership tuple :type return_partition: bool :return: partition from leidenalg - :rtype: tuple[int] or leidenalg.RBConfigurationVertexPartitionWeightedLayers + :rtype: tuple[int, ...] or list[leidenalg.RBConfigurationVertexPartition] """ if 'weight' not in G_intralayer.es: G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount() @@ -132,14 +181,18 @@ def multilayer_leiden(G_intralayer, G_interlayer, layer_vec, gamma, omega, optim if 'weight' not in G_interlayer.es: G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount() - G_split_layers = split_intralayer_leiden_graph(G_intralayer, layer_vec) + G_split_layers: list[ig.Graph] = split_intralayer_leiden_graph(G_intralayer, layer_vec) if optimiser is None: optimiser = leidenalg.Optimiser() - intralayer_parts = [leidenalg.RBConfigurationVertexPartition(G_layer, weights='weight', resolution_parameter=gamma) - for G_layer in G_split_layers] - interlayer_part = leidenalg.CPMVertexPartition(G_interlayer, resolution_parameter=0.0, weights='weight') + intralayer_parts: list[LeidenPartition] = ( + [leidenalg.RBConfigurationVertexPartition(G_layer, weights='weight', resolution_parameter=gamma) + for G_layer in G_split_layers] + ) + interlayer_part: leidenalg.CPMVertexPartition = leidenalg.CPMVertexPartition( + G_interlayer, resolution_parameter=0.0, weights='weight' + ) optimiser.optimise_partition_multiplex(intralayer_parts + [interlayer_part], layer_weights=[1] * len(intralayer_parts) + [omega]) @@ -149,26 +202,35 @@ def multilayer_leiden(G_intralayer, G_interlayer, layer_vec, gamma, omega, optim return tuple(intralayer_parts[0].membership) -def _wrapped_multilayer_leiden(args): +def _wrapped_multilayer_leiden(args: tuple[ig.Graph, ig.Graph, LayerVector, float, float] + ) -> Membership: """Wrapped multilayer_leiden() for use in multiprocessing.Pool.imap_unordered.""" return multilayer_leiden(*args) -def multilayer_leiden_part(G_intralayer, G_interlayer, layer_membership): +def multilayer_leiden_part( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_membership: LayerVector +) -> MultilayerLeidenParts: if 'weight' not in G_intralayer.es: G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount() if 'weight' not in G_interlayer.es: G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount() - G_split_layers = split_intralayer_leiden_graph(G_intralayer, layer_membership) - intralayer_parts = [leidenalg.RBConfigurationVertexPartition(G_layer, weights='weight') - for G_layer in G_split_layers] - interlayer_part = leidenalg.CPMVertexPartition(G_interlayer, resolution_parameter=0.0, weights='weight') + G_split_layers: list[ig.Graph] = split_intralayer_leiden_graph(G_intralayer, layer_membership) + intralayer_parts: list[LeidenPartition] = ( + [leidenalg.RBConfigurationVertexPartition(G_layer, weights='weight') for G_layer in G_split_layers] + ) + interlayer_part: leidenalg.CPMVertexPartition = leidenalg.CPMVertexPartition( + G_interlayer, resolution_parameter=0.0, weights='weight' + ) return intralayer_parts, interlayer_part -def multilayer_leiden_part_with_membership(G_intralayer, G_interlayer, layer_membership, community_membership): +def multilayer_leiden_part_with_membership( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_membership: LayerVector, + community_membership: IntSequence +) -> MultilayerLeidenParts: if isinstance(community_membership, np.ndarray): community_membership = community_membership.tolist() intralayer_parts, interlayer_part = multilayer_leiden_part(G_intralayer, G_interlayer, layer_membership) @@ -178,27 +240,31 @@ def multilayer_leiden_part_with_membership(G_intralayer, G_interlayer, layer_mem return intralayer_parts, interlayer_part -def repeated_leiden_from_gammas(G, gammas): +def repeated_leiden_from_gammas(G: ig.Graph, gammas: FloatSequence) -> set[Membership]: return {sorted_tuple(singlelayer_leiden(G, gamma)) for gamma in gammas} -def repeated_parallel_leiden_from_gammas(G, gammas, show_progress=True): +def repeated_parallel_leiden_from_gammas( + G: ig.Graph, gammas: FloatSequence, show_progress: bool = True +) -> set[Membership]: r"""Runs the Leiden modularity maximization algorithm at each provided :math:`\gamma` value, using all CPU cores. :param G: graph of interest :type G: igraph.Graph :param gammas: list of gammas (resolution parameters) to run Leiden at - :type gammas: list[float] + :type gammas: Sequence[float] | numpy.ndarray :param show_progress: if True, render a progress bar :type show_progress: bool :return: a set of all unique partitions returned by the Leiden algorithm - :rtype: set of tuple[int] + :rtype: set of tuple[int, ...] """ - total = set() - pool_chunk_size = max(1, len(gammas) // (cpu_count() * 100)) + total: set[Membership] = set() + pool_chunk_size: int = max(1, len(gammas) // (cpu_count() * 100)) with Pool(processes=cpu_count()) as pool: - pool_iterator = pool.imap_unordered(_wrapped_singlelayer_leiden, [(G, g) for g in gammas], - chunksize=pool_chunk_size) + pool_iterator: Iterable[Membership] = ( + pool.imap_unordered(_wrapped_singlelayer_leiden, [(G, g) for g in gammas], + chunksize=pool_chunk_size) + ) if show_progress: pool_iterator = tqdm(pool_iterator, total=len(gammas)) @@ -208,13 +274,19 @@ def repeated_parallel_leiden_from_gammas(G, gammas, show_progress=True): return total -def repeated_leiden_from_gammas_omegas(G_intralayer, G_interlayer, layer_vec, gammas, omegas): - return {sorted_tuple(multilayer_leiden(G_intralayer, G_interlayer, layer_vec, gamma, omega)) +def repeated_leiden_from_gammas_omegas( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector, + gammas: FloatSequence, omegas: FloatSequence +) -> set[Membership]: + return {sorted_tuple(tuple(multilayer_leiden(G_intralayer, G_interlayer, layer_vec, gamma, omega))) for gamma in gammas for omega in omegas} -def repeated_parallel_leiden_from_gammas_omegas(G_intralayer, G_interlayer, layer_vec, gammas, omegas, - show_progress=True): +def repeated_parallel_leiden_from_gammas_omegas( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector, + gammas: FloatSequence, omegas: FloatSequence, + show_progress: bool = True +) -> set[Membership]: """ Runs leidenalg at each gamma and omega in ``gammas`` and ``omegas``, using all CPU cores available. @@ -228,23 +300,28 @@ def repeated_parallel_leiden_from_gammas_omegas(G_intralayer, G_interlayer, laye :param layer_vec: list of each vertex's layer membership :type layer_vec: list[int] :param gammas: list of gammas to run leidenalg at - :type gammas: list[float] + :type gammas: Sequence[float] | numpy.ndarray :param omegas: list of omegas to run leidenalg at - :type omegas: list[float] + :type omegas: Sequence[float] | numpy.ndarray :param show_progress: if True, render a progress bar :type show_progress: bool :return: a set of all unique partitions encountered - :rtype: set of tuple[int] + :rtype: set of tuple[int, ...] """ - resolution_parameter_points = [(gamma, omega) for gamma in gammas for omega in omegas] + resolution_parameter_points: list[tuple[float, float]] = ( + [(gamma, omega) for gamma in gammas for omega in omegas] + ) - total = set() - pool_chunk_size = max(1, len(resolution_parameter_points) // (cpu_count() * 100)) + total: set[Membership] = set() + pool_chunk_size: int = max(1, len(resolution_parameter_points) // (cpu_count() * 100)) with Pool(processes=cpu_count()) as pool: - pool_iterator = pool.imap_unordered( - _wrapped_multilayer_leiden, - [(G_intralayer, G_interlayer, layer_vec, gamma, omega) for gamma, omega in resolution_parameter_points], - chunksize=pool_chunk_size + pool_iterator: Iterable[Membership] = ( + pool.imap_unordered( + _wrapped_multilayer_leiden, + [(G_intralayer, G_interlayer, layer_vec, gamma, omega) + for gamma, omega in resolution_parameter_points], + chunksize=pool_chunk_size + ) ) if show_progress: pool_iterator = tqdm(pool_iterator, total=len(resolution_parameter_points)) diff --git a/utilities/louvain_utilities.py b/utilities/louvain_utilities.py index f5cceb5..0ff1db3 100644 --- a/utilities/louvain_utilities.py +++ b/utilities/louvain_utilities.py @@ -2,184 +2,24 @@ This is a deprecated version of leiden_utilities. Usage will be shimmed to use Leiden instead of Louvain when possible. Prior to version v1.3.0, we used the Louvain algorithm for modularity maximization instead of Leiden. The deprecated -module ``modularitypruning.louvain_utilities`` now shims single-layer functions to their corresponding Leiden versions -in ``modularitypruning.leiden_utilities`` (though it still contains the legacy multi-layer functions since they can be -faster in general -- leidenalg does not efficiently implement multilayer optimization). +module ``modularitypruning.louvain_utilities`` now shims Louvain function names to their corresponding Leiden versions +in ``modularitypruning.leiden_utilities``. """ import warnings -from math import ceil -from multiprocessing import Pool, cpu_count - -import numpy as np -import psutil from . import leiden_utilities -from .leiden_utilities import sorted_tuple, LOW_MEMORY_THRESHOLD -from .progress import Progress - -try: - import louvain # import louvain if possible -except ModuleNotFoundError: - pass warnings.simplefilter('always', DeprecationWarning) warnings.warn("modularitypruning.louvain_utilities has been replaced. " "Please use leiden_utilities in the future when possible.", DeprecationWarning) -def __getattr__(name): +def __getattr__(name: str) -> object: # there is no meaningful type for this style of universal shim + leiden_name: str = name if "louvain" in name: leiden_name = name.replace("louvain", "leiden") warnings.simplefilter('always', DeprecationWarning) warnings.warn(f"The Louvain functions have been deprecated. Replacing {repr(name)} with {repr(leiden_name)}.", DeprecationWarning) - return getattr(leiden_utilities, leiden_name) - - -def check_multilayer_louvain_capabilities(fatal=True): - """Check if we are using the version of louvain with fast multilayer optimization. - - :param fatal: if True, raise an error when we are not using the desired louvain version - """ - try: - louvain.RBConfigurationVertexPartitionWeightedLayers - except AttributeError as e: - if not fatal: - return False - - raise AttributeError("Your installation of louvain does not support fast multilayer optimization. See " - "https://github.com/wweir827/louvain-igraph and " - "https://github.com/vtraag/louvain-igraph/pull/34") from e - - return True - - -def multilayer_louvain(G_intralayer, G_interlayer, layer_vec, gamma, omega, optimiser=None, return_partition=False): - r"""Run the Louvain modularity maximization algorithm at a single (:math:`\gamma, \omega`) value. - - :param G_intralayer: intralayer graph of interest - :type G_intralayer: igraph.Graph - :param G_interlayer: interlayer graph of interest - :type G_interlayer: igraph.Graph - :param layer_vec: list of each vertex's layer membership - :type layer_vec: list[int] - :param gamma: gamma (intralayer resolution parameter) to run Louvain at - :type gamma: float - :param omega: omega (interlayer resolution parameter) to run Louvain at - :type omega: float - :param optimiser: if not None, use passed-in (potentially custom) louvain optimiser - :type optimiser: louvain.Optimiser - :param return_partition: if True, return a louvain partition. Otherwise, return a community membership tuple - :type return_partition: bool - :return: partition from louvain - :rtype: tuple[int] or louvain.RBConfigurationVertexPartitionWeightedLayers - """ - # RBConfigurationVertexPartitionWeightedLayers implements a multilayer version of "standard" modularity (i.e. - # the Reichardt and Bornholdt's Potts model with configuration null model). - check_multilayer_louvain_capabilities() - - if 'weight' not in G_intralayer.es: - G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount() - - if 'weight' not in G_interlayer.es: - G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount() - - if optimiser is None: - optimiser = louvain.Optimiser() - - intralayer_part = louvain.RBConfigurationVertexPartitionWeightedLayers(G_intralayer, layer_vec=layer_vec, - weights='weight', resolution_parameter=gamma) - interlayer_part = louvain.CPMVertexPartition(G_interlayer, resolution_parameter=0.0, weights='weight') - optimiser.optimise_partition_multiplex([intralayer_part, interlayer_part], layer_weights=[1, omega]) - - if return_partition: - return intralayer_part - else: - return tuple(intralayer_part.membership) - - -def multilayer_louvain_part(G_intralayer, G_interlayer, layer_membership): - if 'weight' not in G_intralayer.es: - G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount() - - if 'weight' not in G_interlayer.es: - G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount() - - intralayer_part = louvain.RBConfigurationVertexPartitionWeightedLayers(G_intralayer, layer_vec=layer_membership, - weights='weight') - interlayer_part = louvain.CPMVertexPartition(G_interlayer, resolution_parameter=0.0, weights='weight') - return intralayer_part, interlayer_part - - -def multilayer_louvain_part_with_membership(G_intralayer, G_interlayer, layer_membership, community_membership): - if isinstance(community_membership, np.ndarray): - community_membership = community_membership.tolist() - intralayer_part, interlayer_part = multilayer_louvain_part(G_intralayer, G_interlayer, layer_membership) - intralayer_part.set_membership(community_membership) - interlayer_part.set_membership(community_membership) - return intralayer_part, interlayer_part - - -def repeated_louvain_from_gammas_omegas(G_intralayer, G_interlayer, layer_vec, gammas, omegas): - return {sorted_tuple(multilayer_louvain(G_intralayer, G_interlayer, layer_vec, gamma, omega)) - for gamma in gammas for omega in omegas} - - -def repeated_parallel_louvain_from_gammas_omegas(G_intralayer, G_interlayer, layer_vec, gammas, omegas, - show_progress=True, chunk_dispatch=True): - """ - Runs louvain at each gamma and omega in ``gammas`` and ``omegas``, using all CPU cores available. - - :param G_intralayer: intralayer graph of interest - :type G_intralayer: igraph.Graph - :param G_interlayer: interlayer graph of interest - :type G_interlayer: igraph.Graph - :param layer_vec: list of each vertex's layer membership - :type layer_vec: list[int] - :param gammas: list of gammas to run louvain at - :type gammas: list[float] - :param omegas: list of omegas to run louvain at - :type omegas: list[float] - :param show_progress: if True, render a progress bar - :type show_progress: bool - :param chunk_dispatch: if True, dispatch parallel work in chunks. Setting this to False may increase performance, - but can lead to out-of-memory issues - :type chunk_dispatch: bool - :return: a set of all unique partitions encountered - :rtype: set of tuple[int] - """ - resolution_parameter_points = [(gamma, omega) for gamma in gammas for omega in omegas] - - pool = Pool(processes=cpu_count()) - total = set() - - chunk_size = len(resolution_parameter_points) // 99 - if chunk_size > 0 and chunk_dispatch: - chunk_params = ([(G_intralayer, G_interlayer, layer_vec, gamma, omega) - for gamma, omega in resolution_parameter_points[i:i + chunk_size]] - for i in range(0, len(resolution_parameter_points), chunk_size)) - else: - chunk_params = [[(G_intralayer, G_interlayer, layer_vec, gamma, omega) - for gamma, omega in resolution_parameter_points]] - chunk_size = len(gammas) - - if show_progress: - progress = Progress(ceil(len(resolution_parameter_points) / chunk_size)) - - for chunk in chunk_params: - for partition in pool.starmap(multilayer_louvain, chunk): - total.add(sorted_tuple(partition)) - - if show_progress: - progress.increment() - - if psutil.virtual_memory().available < LOW_MEMORY_THRESHOLD: - # Reinitialize pool to get around an apparent memory leak in multiprocessing - pool.close() - pool = Pool(processes=cpu_count()) - - if show_progress: - progress.done() - - pool.close() - return total + attr: object = getattr(leiden_utilities, leiden_name) + return attr diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py index 466c959..321b324 100644 --- a/utilities/parameter_estimation.py +++ b/utilities/parameter_estimation.py @@ -1,13 +1,20 @@ +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import igraph import leidenalg -from .leiden_utilities import singlelayer_leiden, multilayer_leiden -from .parameter_estimation_utilities import leiden_part_with_membership, estimate_singlelayer_SBM_parameters, \ - gamma_estimate_from_parameters, omega_function_from_model, estimate_multilayer_SBM_parameters +from .leiden_utilities import singlelayer_leiden, multilayer_leiden, leiden_part_with_membership +from .parameter_estimation_utilities import estimate_singlelayer_SBM_parameters, gamma_estimate_from_parameters, \ + omega_function_from_model, estimate_multilayer_SBM_parameters from .partition_utilities import in_degrees +from .types import (GammaUpdateFunc, LayerVector, LeidenPartition, MultilayerSBMParameters, + OmegaUpdateFunc, SinglelayerSBMParameters) -def iterative_monolayer_resolution_parameter_estimation(G, gamma=1.0, tol=1e-2, max_iter=25, verbose=False, - method="leiden"): +def iterative_monolayer_resolution_parameter_estimation( + G: "igraph.Graph", gamma: float = 1.0, tol: float = 1e-2, max_iter: int = 25, + verbose: bool = False, method: str = "leiden" +) -> tuple[float, LeidenPartition]: """Monolayer variant of ALG. 1 from "Relating modularity maximization and stochastic block models in multilayer networks." This is intended to determine an "optimal" value for gamma by repeatedly maximizing modularity and estimating new values for the resolution parameter. @@ -34,37 +41,39 @@ def iterative_monolayer_resolution_parameter_estimation(G, gamma=1.0, tol=1e-2, if 'weight' not in G.es: G.es['weight'] = [1.0] * G.ecount() - m = sum(G.es['weight']) + m: float = sum(G.es['weight']) if method == "louvain" or method == "leiden": - def maximize_modularity(resolution_param): + def maximize_modularity(resolution_param: float) -> LeidenPartition: return singlelayer_leiden(G, resolution_param, return_partition=True) elif method == "2-spinglass": - def maximize_modularity(resolution_param): + def maximize_modularity(resolution_param: float) -> LeidenPartition: membership = G.community_spinglass(spins=2, gamma=resolution_param).membership return leiden_part_with_membership(G, membership) else: raise ValueError(f"Community detection method {method} not supported") - def estimate_SBM_parameters(partition): + def estimate_SBM_parameters(partition: LeidenPartition) -> SinglelayerSBMParameters: return estimate_singlelayer_SBM_parameters(G, partition, m=m) - def update_gamma(omega_in, omega_out): - return gamma_estimate_from_parameters(omega_in, omega_out) + update_gamma: GammaUpdateFunc = gamma_estimate_from_parameters - part, last_gamma = None, None + part = None + last_gamma = gamma for iteration in range(max_iter): part = maximize_modularity(gamma) omega_in, omega_out = estimate_SBM_parameters(part) last_gamma = gamma - gamma = update_gamma(omega_in, omega_out) + new_gamma = update_gamma(omega_in, omega_out) - if gamma is None: + if new_gamma is None: raise ValueError(f"gamma={last_gamma:.3f} resulted in degenerate partition") + gamma = float(new_gamma) + if verbose: - print(f"Iter {iteration:>2}: {len(part)} communities with Q={part.q:.3f} and " + print(f"Iter {iteration:>2}: {len(part) if part else 0} communities with Q={part.q:.3f} and " f"gamma={last_gamma:.3f}->{gamma:.3f}") if abs(gamma - last_gamma) < tol: @@ -75,12 +84,16 @@ def update_gamma(omega_in, omega_out): f"Final move of {abs(gamma - last_gamma):.3f} was not within tolerance {tol}") if verbose: - print(f"Returned {len(part)} communities with Q={part.q:.3f} and gamma={gamma:.3f}") + print(f"Returned {len(part) if part else 0} communities with Q={part.q if part else -1:.3f} and " + f"gamma={gamma:.3f}") return gamma, part -def check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, model, m_t, T, N=None, Nt=None): +def check_multilayer_graph_consistency( + G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: LayerVector, + model: str, m_t: list[float], T: int, N: int | None = None, Nt: list[int] | None = None +) -> None: """ Checks that the structures of the intralayer and interlayer graphs are consistent and match the given model. @@ -93,43 +106,52 @@ def check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, mo :param N: number of nodes per layer :param Nt: vector of nodes per layer """ - - rules = [T > 1, - "Graph must have multiple layers", - G_interlayer.is_directed(), - "Interlayer graph should be directed", - G_interlayer.vcount() == G_intralayer.vcount(), - "Inter-layer and Intra-layer graphs must be of the same size", - len(layer_vec) == G_intralayer.vcount(), - "Layer membership vector must have length matching graph size", - all(m > 0 for m in m_t), - "All layers of graph must contain edges", - all(layer_vec[e.source] == layer_vec[e.target] for e in G_intralayer.es), - "Intralayer graph should not contain edges across layers", - model != 'temporal' or G_interlayer.ecount() == N * (T - 1), - "Interlayer temporal graph must contain (nodes per layer) * (number of layers - 1) edges", - model != 'temporal' or (G_interlayer.vcount() % T == 0 and G_intralayer.vcount() % T == 0), - "Vertex count of a temporal graph should be a multiple of the number of layers", - model != 'temporal' or all(nt == N for nt in Nt), - "Temporal networks must have the same number of nodes in every layer", - model != 'multilevel' or all(nt > 0 for nt in Nt), - "All layers of a multilevel graph must be consecutive and nonempty", - model != 'multilevel' or all(in_degree <= 1 for in_degree in in_degrees(G_interlayer)), - "Multilevel networks should have at most one interlayer in-edge per node", - model != 'multiplex' or all(nt == N for nt in Nt), - "Multiplex networks must have the same number of nodes in every layer", - model != 'multiplex' or G_interlayer.ecount() == N * T * (T - 1), - "Multiplex interlayer networks must contain edges between all pairs of layers"] + # force omitted values to fail the checks below if they were actually needed + if not N: + N = -1 + if not Nt: + Nt = [0] * T + + rules: list[bool | str] \ + = [T > 1, + "Graph must have multiple layers", + G_interlayer.is_directed(), + "Interlayer graph should be directed", + G_interlayer.vcount() == G_intralayer.vcount(), + "Inter-layer and Intra-layer graphs must be of the same size", + len(layer_vec) == G_intralayer.vcount(), + "Layer membership vector must have length matching graph size", + all(m > 0 for m in m_t), + "All layers of graph must contain edges", + all(layer_vec[e.source] == layer_vec[e.target] for e in G_intralayer.es), + "Intralayer graph should not contain edges across layers", + model != 'temporal' or G_interlayer.ecount() == N * (T - 1), + "Interlayer temporal graph must contain (nodes per layer) * (number of layers - 1) edges", + model != 'temporal' or (G_interlayer.vcount() % T == 0 + and G_intralayer.vcount() % T == 0), + "Vertex count of a temporal graph should be a multiple of the number of layers", + model != 'temporal' or all(nt == N for nt in Nt), + "Temporal networks must have the same number of nodes in every layer", + model != 'multilevel' or all(nt > 0 for nt in Nt), + "All layers of a multilevel graph must be consecutive and nonempty", + model != 'multilevel' or all(in_degree <= 1 for in_degree in in_degrees(G_interlayer)), + "Multilevel networks should have at most one interlayer in-edge per node", + model != 'multiplex' or all(nt == N for nt in Nt), + "Multiplex networks must have the same number of nodes in every layer", + model != 'multiplex' or G_interlayer.ecount() == N * T * (T - 1), + "Multiplex interlayer networks must contain edges between all pairs of layers"] checks, messages = rules[::2], rules[1::2] if not all(checks): - raise ValueError("Input graph is malformed\n" + "\n".join(m for c, m in zip(checks, messages) if not c)) + raise ValueError("Input graph is malformed\n" + "\n".join(str(m) for c, m in zip(checks, messages) if not c)) -def iterative_multilayer_resolution_parameter_estimation(G_intralayer, G_interlayer, layer_vec, gamma=1.0, omega=1.0, - gamma_tol=1e-2, omega_tol=5e-2, omega_max=1000, max_iter=25, - model='temporal', verbose=False): +def iterative_multilayer_resolution_parameter_estimation( + G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: LayerVector, + gamma: float = 1.0, omega: float = 1.0, gamma_tol: float = 1e-2, omega_tol: float = 5e-2, + omega_max: float = 1000, max_iter: int = 25, model: str = 'temporal', verbose: bool = False +) -> tuple[float, float, list[LeidenPartition]]: """ Multilayer variant of ALG. 1 from "Relating modularity maximization and stochastic block models in multilayer networks." The nested functions here are just used to match the pseudocode in the paper. @@ -160,43 +182,46 @@ def iterative_multilayer_resolution_parameter_estimation(G_intralayer, G_interla - gamma to which the iteration converged - omega to which the iteration converged - the resulting partition - :rtype: tuple[float, float, tuple[int]] + :rtype: tuple[float, float, list[leidenalg.RBConfigurationVertexPartition]] """ - global parts if 'weight' not in G_intralayer.es: G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount() if 'weight' not in G_interlayer.es: G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount() - T = max(layer_vec) + 1 # layer count - optimiser = leidenalg.Optimiser() + T: int = max(layer_vec) + 1 # layer count + optimiser: leidenalg.Optimiser = leidenalg.Optimiser() # compute total edge weights per layer - m_t = [0] * T + m_t: list[float] = [0] * T for e in G_intralayer.es: m_t[layer_vec[e.source]] += e['weight'] # compute total node counts per layer - N = G_intralayer.vcount() // T - Nt = [0] * T + N: int = G_intralayer.vcount() // T + Nt: list[int] = [0] * T for layer in layer_vec: Nt[layer] += 1 check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, model, m_t, T, N, Nt) - update_omega = omega_function_from_model(model, omega_max, T=T) - update_gamma = gamma_estimate_from_parameters + update_omega: OmegaUpdateFunc = omega_function_from_model(model, omega_max, T=T) + update_gamma: GammaUpdateFunc = gamma_estimate_from_parameters - def maximize_modularity(intralayer_resolution, interlayer_resolution): + def maximize_modularity( + intralayer_resolution: float, interlayer_resolution: float + ) -> list[LeidenPartition]: return multilayer_leiden(G_intralayer, G_interlayer, layer_vec, intralayer_resolution, interlayer_resolution, optimiser=optimiser, return_partition=True) - def estimate_SBM_parameters(partition): + def estimate_SBM_parameters(partition: LeidenPartition) -> MultilayerSBMParameters: return estimate_multilayer_SBM_parameters(G_intralayer, G_interlayer, layer_vec, partition, model, N=N, T=T, Nt=Nt, m_t=m_t) - part, K, last_gamma, last_omega = (None,) * 4 + parts = [] + K = None + last_gamma, last_omega = gamma, omega for iteration in range(max_iter): parts = maximize_modularity(gamma, omega) theta_in, theta_out, p, K = estimate_SBM_parameters(parts[0]) @@ -205,11 +230,12 @@ def estimate_SBM_parameters(partition): raise ValueError(f"gamma={gamma:.3f}, omega={omega:.3f} resulted in impossible estimate p={p:.3f}") last_gamma, last_omega = gamma, omega - gamma = update_gamma(theta_in, theta_out) + new_gamma = update_gamma(theta_in, theta_out) - if gamma is None: + if new_gamma is None: raise ValueError(f"gamma={last_gamma:.3f}, omega={last_omega:.3f} resulted in degenerate partition") + gamma = float(new_gamma) omega = update_omega(theta_in, theta_out, p, K) if verbose: @@ -226,7 +252,7 @@ def estimate_SBM_parameters(partition): f"was not within tolerance ({gamma_tol}, {omega_tol})") if verbose: - part_Q = sum(p.q for p in parts) + part_Q = sum(p.q for p in parts) if parts else -1 print(f"Returned {K} communities with Q={part_Q:.3f}, gamma={gamma:.3f}, and omega={omega:.3f}") return gamma, omega, parts diff --git a/utilities/parameter_estimation_utilities.py b/utilities/parameter_estimation_utilities.py index d9321c5..f2b5ee2 100644 --- a/utilities/parameter_estimation_utilities.py +++ b/utilities/parameter_estimation_utilities.py @@ -1,3 +1,4 @@ +from typing import overload, Literal import warnings from math import log @@ -9,9 +10,16 @@ from .champ_utilities import CHAMP_2D, CHAMP_3D from .leiden_utilities import leiden_part_with_membership, sorted_tuple from .partition_utilities import num_communities +from .types import (Domain, DomainWithEstimates, FloatArray, GammaEstimate, + GammaEstimateValue, GammaOmegaEstimate, GammaRange, GammaUpdateFunc, LayerVector, + LeidenPartition, Membership, MultilayerSBMParameters, OmegaUpdateFunc, + PersistenceFunc, SinglelayerSBMParameters) -def estimate_singlelayer_SBM_parameters(G, partition, m=None): +def estimate_singlelayer_SBM_parameters( + G: ig.Graph, partition: LeidenPartition, + m: float | None = None +) -> SinglelayerSBMParameters: r"""Estimate singlelayer SBM parameters from a graph and a partition. See https://doi.org/10.1103/PhysRevE.94.052315 for more details. @@ -21,7 +29,7 @@ def estimate_singlelayer_SBM_parameters(G, partition, m=None): :param partition: partition of interest :type partition: leidenalg.RBConfigurationVertexPartition :param m: total edge weight of graph (if None, will be computed) - :type m: float + :type m: float | None :return: SBM parameter estimates :math:`(\omega_{in}, \omega_{out})` :rtype: tuple[float, float] """ @@ -30,25 +38,29 @@ def estimate_singlelayer_SBM_parameters(G, partition, m=None): m = sum(G.es['weight']) assert isinstance(partition, leidenalg.RBConfigurationVertexPartition) - community = partition.membership + community: list[int] = partition.membership - m_in = sum(e['weight'] * (community[e.source] == community[e.target]) for e in G.es) - kappa_r_list = [0] * len(partition) + m_in: float = sum(e['weight'] * (community[e.source] == community[e.target]) for e in G.es) + kappa_r_list: list[float] = [0] * len(partition) for e in G.es: kappa_r_list[community[e.source]] += e['weight'] kappa_r_list[community[e.target]] += e['weight'] - sum_kappa_sqr = sum(x ** 2 for x in kappa_r_list) + sum_kappa_sqr: float = sum(x ** 2 for x in kappa_r_list) - omega_in = (2 * m_in) / (sum_kappa_sqr / (2 * m)) + omega_in: float = (2 * m_in) / (sum_kappa_sqr / (2 * m)) # guard for div by zero with single community partition - omega_out = (2 * m - 2 * m_in) / (2 * m - sum_kappa_sqr / (2 * m)) if len(partition) > 1 else 0 + omega_out: float = (2 * m - 2 * m_in) / (2 * m - sum_kappa_sqr / (2 * m)) if len(partition) > 1 else 0 # return estimates for omega_in, omega_out return omega_in, omega_out -def estimate_multilayer_SBM_parameters(G_intralayer, G_interlayer, layer_vec, partition, model, N=None, T=None, - Nt=None, m_t=None): +def estimate_multilayer_SBM_parameters( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, + layer_vec: LayerVector, partition: LeidenPartition, + model: str, N: int | None = None, T: int | None = None, + Nt: list[int] | None = None, m_t: list[float] | None = None +) -> MultilayerSBMParameters: """Estimates multilayer SBM parameters from a graph and a partition :param G_intralayer: intralayer graph of interest @@ -62,13 +74,13 @@ def estimate_multilayer_SBM_parameters(G_intralayer, G_interlayer, layer_vec, pa :param model: network layer topology (temporal, multilevel, multiplex) :type model: str :param N: number of nodes per layer (automatically computed if None) - :type N: int + :type N: int | None :param T: number of layers in input graph (automatically computed if None) - :type T: int + :type T: int | None :param Nt: vector of nodes per layer (automatically computed if None) - :type Nt: int + :type Nt: list[int] | None :param m_t: vector of total edge weights per layer (automatically computed if None) - :type m_t: int + :type m_t: list[float] | None :return: theta_in, theta_out, p, K :rtype: float, float, float, int """ @@ -93,42 +105,44 @@ def estimate_multilayer_SBM_parameters(G_intralayer, G_interlayer, layer_vec, pa for layer in layer_vec: Nt[layer] += 1 - K = len(partition) + K: int = len(partition) - community = partition.membership - m_t_in = [0] * T + community: list[int] = partition.membership + m_t_in: list[float] = [0] * T for e in G_intralayer.es: if community[e.source] == community[e.target] and layer_vec[e.source] == layer_vec[e.target]: m_t_in[layer_vec[e.source]] += e['weight'] - kappa_t_r_list = [[0] * K for _ in range(T)] + kappa_t_r_list: list[list[int]] = [[0] * K for _ in range(T)] for e in G_intralayer.es: layer = layer_vec[e.source] kappa_t_r_list[layer][community[e.source]] += e['weight'] kappa_t_r_list[layer][community[e.target]] += e['weight'] - sum_kappa_t_sqr = [sum(x ** 2 for x in kappa_t_r_list[t]) for t in range(T)] + sum_kappa_t_sqr: list[float] = [sum(x ** 2 for x in kappa_t_r_list[t]) for t in range(T)] - theta_in = sum(2 * m_t_in[t] for t in range(T)) / sum(sum_kappa_t_sqr[t] / (2 * m_t[t]) for t in range(T)) + theta_in: float = sum(2 * m_t_in[t] for t in range(T)) / sum(sum_kappa_t_sqr[t] / (2 * m_t[t]) for t in range(T)) # guard for div by zero with e.g. a single community partition - theta_out_numerator = sum(2 * m_t[t] - 2 * m_t_in[t] for t in range(T)) - theta_out_denominator = sum(2 * m_t[t] - sum_kappa_t_sqr[t] / (2 * m_t[t]) for t in range(T)) + theta_out_numerator: float = sum(2 * m_t[t] - 2 * m_t_in[t] for t in range(T)) + theta_out_denominator: float = sum(2 * m_t[t] - sum_kappa_t_sqr[t] / (2 * m_t[t]) for t in range(T)) if theta_out_denominator == 0: - theta_out = 0 + theta_out: float = 0 else: theta_out = theta_out_numerator / theta_out_denominator - calculate_persistence = persistence_function_from_model(model, G_interlayer, layer_vec=layer_vec, N=N, T=T, Nt=Nt) - pers = calculate_persistence(community) + calculate_persistence: PersistenceFunc = ( + persistence_function_from_model(model, G_interlayer, layer_vec=layer_vec, N=N, T=T, Nt=Nt) + ) + pers: float = calculate_persistence(community) if model == 'multiplex': # estimate p by solving polynomial root-finding problem with starting estimate p=0.5 - def f(x): - coeff = 2 * (1 - 1 / K) / (T * (T - 1)) - return coeff * sum((T - n) * x ** n for n in range(1, T)) + 1 / K - pers + def f(x: FloatArray) -> FloatArray: + coeff: float = 2 * (1 - 1 / K) / (T * (T - 1)) + return coeff * np.sum([(T - n) * x ** n for n in range(1, T)], axis=0) + 1 / K - pers # guard for div by zero with single community partition # (in this case, all community assignments persist across layers) - p = fsolve(f, np.array([0.5]))[0] if pers < 1.0 and K > 1 else 1.0 + p: float = fsolve(f, np.array([0.5]))[0] if pers < 1.0 and K > 1 else 1.0 if p < 0: p = 0 else: @@ -139,7 +153,7 @@ def f(x): return theta_in, theta_out, p, K -def gamma_estimate(G, partition): +def gamma_estimate(G: ig.Graph, partition: LeidenPartition | Membership) -> GammaEstimateValue: """Compute the "correct" value of gamma where modularity maximization becomes equivalent to maximum likelihood methods on a degree-corrected, planted partition stochastic block model. @@ -148,9 +162,9 @@ def gamma_estimate(G, partition): :param G: graph of interest :type G: igraph.Graph :param partition: partition of interest - :type partition: tuple[int] or leidenalg.RBConfigurationVertexPartition + :type partition: tuple[int, ...] or leidenalg.RBConfigurationVertexPartition :return: gamma estimate - :rtype: float + :rtype: float | None """ if 'weight' not in G.es: @@ -163,7 +177,7 @@ def gamma_estimate(G, partition): return gamma_estimate_from_parameters(omega_in, omega_out) -def gamma_estimate_from_parameters(omega_in, omega_out): +def gamma_estimate_from_parameters(omega_in: float, omega_out: float) -> GammaEstimateValue: """ Compute the "correct" value of gamma as in :meth:`~modularitypruning.parameter_estimation_utilities.gamma_estimate` from SBM parameters. @@ -173,16 +187,19 @@ def gamma_estimate_from_parameters(omega_in, omega_out): :param omega_out: within-community edge propensity of a degree-corrected, planted partition SBM :type omega_out: float :return: gamma estimate - :rtype: float + :rtype: float | None """ if omega_in == 0 or omega_out == 0: return None # degenerate partition, this could reasonably be taken to be 0 - return (omega_in - omega_out) / (np.log(omega_in) - np.log(omega_out)) + return np.float64((omega_in - omega_out) / (np.log(omega_in) - np.log(omega_out))) -def multiplex_omega_estimate_from_parameters(theta_in, theta_out, p, K, T, omega_max=1000): +def multiplex_omega_estimate_from_parameters( + theta_in: float, theta_out: float, + p: float, K: int, T: int, omega_max: float = 1000 +) -> float: """Returns the omega estimate for a multiplex multilayer model :param theta_in: SBM parameter @@ -213,7 +230,10 @@ def multiplex_omega_estimate_from_parameters(theta_in, theta_out, p, K, T, omega return log(1 + p * K / (1 - p)) / (T * (log(theta_in) - log(theta_out))) -def temporal_multilevel_omega_estimate_from_parameters(theta_in, theta_out, p, K, omega_max=1000): +def temporal_multilevel_omega_estimate_from_parameters( + theta_in: float, theta_out: float, + p: float, K: int, omega_max: float = 1000.0 +) -> float: """Returns the omega estimate for a temporal or multilevel multilayer model :param theta_in: SBM parameter @@ -239,13 +259,16 @@ def temporal_multilevel_omega_estimate_from_parameters(theta_in, theta_out, p, K return log(1 + p * K / (1 - p)) / (2 * (log(theta_in) - log(theta_out))) if p < 1.0 else omega_max -def ordinal_persistence(G_interlayer, community, N, T): +def ordinal_persistence(G_interlayer: ig.Graph, community: list[int], N: int, T: int) -> float: # ordinal persistence (temporal model) - return sum(community[e.source] == community[e.target] for e in G_interlayer.es) / (N * (T - 1)) + return float(sum(community[e.source] == community[e.target] for e in G_interlayer.es) / (N * (T - 1))) -def multilevel_persistence(G_interlayer, community, layer_vec, Nt, T): - pers_per_layer = [0] * T +def multilevel_persistence( + G_interlayer: ig.Graph, community: list[int], + layer_vec: LayerVector, Nt: list[int], T: int +) -> float: + pers_per_layer: list[float] = [0] * T for e in G_interlayer.es: pers_per_layer[layer_vec[e.target]] += (community[e.source] == community[e.target]) @@ -253,12 +276,14 @@ def multilevel_persistence(G_interlayer, community, layer_vec, Nt, T): return sum(pers_per_layer) / (T - 1) -def categorical_persistence(G_interlayer, community, N, T): +def categorical_persistence(G_interlayer: ig.Graph, community: list[int], N: int, T: int) -> float: # categorical persistence (multiplex model) - return sum(community[e.source] == community[e.target] for e in G_interlayer.es) / (N * T * (T - 1)) + return float(sum(community[e.source] == community[e.target] for e in G_interlayer.es) / (N * T * (T - 1))) -def omega_function_from_model(model, omega_max, T): +def omega_function_from_model( + model: str, omega_max: float, T: int +) -> OmegaUpdateFunc: """Returns an appropriate function (depending on the model) for computing omega from multilayer SBM parameters Specifically, it will return versions of @@ -274,10 +299,10 @@ def omega_function_from_model(model, omega_max, T): :rtype: function """ if model == 'multiplex': - def update_omega(theta_in, theta_out, p, K): + def update_omega(theta_in: float, theta_out: float, p: float, K: int) -> float: return multiplex_omega_estimate_from_parameters(theta_in, theta_out, p, K, T, omega_max=omega_max) elif model == 'temporal' or model == 'multilevel': - def update_omega(theta_in, theta_out, p, K): + def update_omega(theta_in: float, theta_out: float, p: float, K: int) -> float: return temporal_multilevel_omega_estimate_from_parameters(theta_in, theta_out, p, K, omega_max=omega_max) else: raise ValueError(f"Model {model} is not temporal, multilevel, or multiplex") @@ -285,7 +310,11 @@ def update_omega(theta_in, theta_out, p, K): return update_omega -def persistence_function_from_model(model, G_interlayer, layer_vec=None, N=None, T=None, Nt=None): +def persistence_function_from_model( + model: str, G_interlayer: ig.Graph, + layer_vec: LayerVector | None = None, N: int | None = None, + T: int | None = None, Nt: list[int] | None = None +) -> PersistenceFunc: """ Returns a function to calculate persistence according to a given multilayer model @@ -303,19 +332,19 @@ def persistence_function_from_model(model, G_interlayer, layer_vec=None, N=None, if N is None or T is None: raise ValueError("Parameters N and T cannot be None for temporal persistence calculation") - def calculate_persistence(community): + def calculate_persistence(community: list[int]) -> float: return ordinal_persistence(G_interlayer, community, N, T) elif model == 'multilevel': if Nt is None or T is None or layer_vec is None: raise ValueError("Parameters layer_vec, Nt, T cannot be None for multilevel persistence calculation") - def calculate_persistence(community): + def calculate_persistence(community: list[int]) -> float: return multilevel_persistence(G_interlayer, community, layer_vec, Nt, T) elif model == 'multiplex': if N is None or T is None: raise ValueError("Parameters N and T cannot be None for multiplex persistence calculation") - def calculate_persistence(community): + def calculate_persistence(community: list[int]) -> float: return categorical_persistence(G_interlayer, community, N, T) else: raise ValueError(f"Model {model} is not temporal, multilevel, or multiplex") @@ -323,8 +352,13 @@ def calculate_persistence(community): return calculate_persistence -def gamma_omega_estimate(G_intralayer, G_interlayer, layer_vec, membership, omega_max=1000, model='temporal', - N=None, T=None, Nt=None, m_t=None): +def gamma_omega_estimate( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, + layer_vec: LayerVector, membership: Membership, + omega_max: float = 1000, model: str = 'temporal', + N: int | None = None, T: int | None = None, Nt: list[int] | None = None, + m_t: list[float] | None = None +) -> GammaOmegaEstimate: """Returns the (gamma, omega) estimate for a multilayer network and a partition :param G_intralayer: intralayer graph of interest @@ -334,37 +368,39 @@ def gamma_omega_estimate(G_intralayer, G_interlayer, layer_vec, membership, omeg :param layer_vec: list of each vertex's layer membership :type layer_vec: list[int] :param membership: partition membership vector - :type membership: tuple[int] + :type membership: tuple[int, ...] :param omega_max: maximum allowed value for omega :type omega_max: float :param model: network layer topology (temporal, multilevel, multiplex) :type model: str :param N: number of nodes per layer (automatically computed if None) - :type N: int + :type N: int | None :param T: number of layers in input graph (automatically computed if None) - :type T: int + :type T: int | None :param Nt: vector of nodes per layer (automatically computed if None) - :type Nt: int + :type Nt: list[int] | None :param m_t: vector of total edge weights per layer (automatically computed if None) - :type m_t: int + :type m_t: list[float] | None :return: gamma estimate, omega estimate - :rtype: float, float + :rtype: float | None, float """ if T is None: T = max(layer_vec) + 1 # layer count - partition = leiden_part_with_membership(G_intralayer, membership) + partition: LeidenPartition = leiden_part_with_membership(G_intralayer, membership) theta_in, theta_out, p, K = estimate_multilayer_SBM_parameters(G_intralayer, G_interlayer, layer_vec, partition, model, N=N, T=T, Nt=Nt, m_t=m_t) - update_omega = omega_function_from_model(model, omega_max, T=T) - update_gamma = gamma_estimate_from_parameters + update_omega: OmegaUpdateFunc = ( + omega_function_from_model(model, omega_max, T=T) + ) + update_gamma: GammaUpdateFunc = gamma_estimate_from_parameters - gamma = update_gamma(theta_in, theta_out) - omega = update_omega(theta_in, theta_out, p, K) + gamma: GammaEstimateValue = update_gamma(theta_in, theta_out) + omega: float = update_omega(theta_in, theta_out, p, K) return gamma, omega -def ranges_to_gamma_estimates(G, ranges): +def ranges_to_gamma_estimates(G: ig.Graph, ranges: list[GammaRange]) -> list[GammaEstimate]: """ Compute gamma estimates as in :meth:`~modularitypruning.parameter_estimation_utilities.gamma_estimate`, given domains of optimality from :meth:`~modularitypruning.champ_utilities.CHAMP_2D`. @@ -373,16 +409,16 @@ def ranges_to_gamma_estimates(G, ranges): :type G: igraph.Graph :param ranges: list of ``(gamma_start, gamma_end, membership)`` tuples as returned from :meth:`~modularitypruning.champ_utilities.CHAMP_2D` - :type ranges: list of tuple[float, float, tuple[int]] + :type ranges: list of tuple[float, float, tuple[int, ...]] :return: a copy of input ranges with the corresponding gamma estimate appended to each tuple - :rtype: list of tuple[float, float, tuple[int], float] + :rtype: list of tuple[float, float, tuple[int, ...], float | None] """ return [(gamma_start, gamma_end, part, gamma_estimate(G, part)) for gamma_start, gamma_end, part in ranges] -def gamma_estimates_to_stable_partitions(gamma_estimates): +def gamma_estimates_to_stable_partitions(gamma_estimates: list[GammaEstimate]) -> list[Membership]: """Computes the stable partitions (i.e. those whose gamma estimates are within their domains of optimality), given domains of optimality and gamma estimates from :meth:`ranges_to_gamma_estimates`. @@ -390,16 +426,20 @@ def gamma_estimates_to_stable_partitions(gamma_estimates): :param gamma_estimates: list of ``(gamma_start, gamma_end, membership, gamma_estimate)`` tuples as returned from :meth:`~modularitypruning.champ_utilities.CHAMP_2D` - :type gamma_estimates: list[tuple] + :type gamma_estimates: list of tuple[float, float, tuple[int, ...], float | None] :return: list of community membership tuples of the stable partitions - :rtype: list of tuple[int] + :rtype: list of tuple[int, ...] """ return [membership for gamma_start, gamma_end, membership, gamma_estimate in gamma_estimates if gamma_estimate is not None and gamma_start <= gamma_estimate <= gamma_end] -def domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, domains, model='temporal'): +def domains_to_gamma_omega_estimates( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, + layer_vec: LayerVector, domains: list[Domain], + model: str = 'temporal' +) -> list[DomainWithEstimates]: """ Compute (gamma, omega) estimates as in :meth:`~modularitypruning.parameter_estimation_utilities.gamma_omega_estimate`, given domains of optimality from @@ -413,14 +453,14 @@ def domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, doma :type layer_vec: list[int] :param domains: list of ``(domain_vertices, membership)`` tuples as returned from :meth:`~modularitypruning.champ_utilities.CHAMP_3D` - :type domains: list of tuple[list[float], tuple[int]] + :type domains: list of tuple[list[numpy.ndarray], tuple[int, ...]] :param model: network layer topology (temporal, multilevel, multiplex) :type model: str :return: a copy of input domains with the corresponding gamma and omega estimates appended to each tuple - :rtype: list of tuple[list[float], tuple[int], float, float] + :rtype: list of tuple[list[numpy.ndarray], tuple[int, ...], float | None, float] """ - domains_with_estimates = [] + domains_with_estimates: list[DomainWithEstimates] = [] for polyverts, membership in domains: gamma_est, omega_est = gamma_omega_estimate(G_intralayer, G_interlayer, layer_vec, membership, model=model) @@ -428,7 +468,27 @@ def domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, doma return domains_with_estimates -def gamma_omega_estimates_to_stable_partitions(domains_with_estimates, return_membership_only=False): +# explain to mypy that the return type depends on the return_membership_only argument +@overload +def gamma_omega_estimates_to_stable_partitions( + domains_with_estimates: list[DomainWithEstimates], + return_membership_only: Literal[False] = False +) -> list[DomainWithEstimates]: + ... + + +@overload +def gamma_omega_estimates_to_stable_partitions( + domains_with_estimates: list[DomainWithEstimates], + return_membership_only: Literal[True] +) -> list[Membership]: + ... + + +def gamma_omega_estimates_to_stable_partitions( + domains_with_estimates: list[DomainWithEstimates], + return_membership_only: bool = False +) -> list[DomainWithEstimates] | list[Membership]: """Computes the stable partitions (i.e. those whose resolution parameter estimates are within their domains of optimality), given domains of optimality and (gamma, omega) estimates from :meth:`domains_to_gamma_omega_estimates`. @@ -437,47 +497,58 @@ def gamma_omega_estimates_to_stable_partitions(domains_with_estimates, return_me :param domains_with_estimates: list of ``(domain_vertices, membership, gamma_estimate, omega_estimate)`` tuples as returned from :meth:`~modularitypruning.champ_utilities.CHAMP_3D` - :type domains_with_estimates: list[tuple] - :return: list of community membership tuples of the stable partitions - :rtype: list of tuple[int] + :param return_membership_only: if True, return only community membership tuples. Otherwise, return list of domains + with estimates + :type return_membership_only: bool + :return: list of community membership tuples or domains with estimates """ - def left_or_right(x1, y1, x2, y2, x, y): + def left_or_right( + x1: float, y1: float, x2: float, + y2: float, x: float, y: float + ) -> bool: """Returns whether the point (x,y) is to the left or right of the line between (x1, y1) and (x2, y2).""" return (x - x1) * (y2 - y1) - (y - y1) * (x2 - x1) >= 0 - stable_partitions = [] + stable_memberships: list[Membership] = [] + stable_domains: list[DomainWithEstimates] = [] for polyverts, membership, gamma_est, omega_est in domains_with_estimates: if gamma_est is None or omega_est is None: print(gamma_est, omega_est) continue - centroid_x = np.mean([x[0] for x in polyverts]) - centroid_y = np.mean([x[1] for x in polyverts]) - polygon_edges = [] + centroid_x: float = np.mean([x[0] for x in polyverts]) + centroid_y: float = np.mean([x[1] for x in polyverts]) + polygon_edges: list[tuple[FloatArray, FloatArray]] = [] for i in range(len(polyverts)): p1, p2 = polyverts[i], polyverts[(i + 1) % len(polyverts)] if left_or_right(p1[0], p1[1], p2[0], p2[1], centroid_x, centroid_y): p1, p2 = p2, p1 polygon_edges.append((p1, p2)) - left_or_rights = [] + left_or_rights: list[bool] = [] for p1, p2 in polygon_edges: left_or_rights.append(left_or_right(p1[0], p1[1], p2[0], p2[1], gamma_est, omega_est)) if all(x for x in left_or_rights) or all(not x for x in left_or_rights): # if the (gamma, omega) estimate is on the same side of all polygon edges, it lies within the domain if return_membership_only: - stable_partitions.append(membership) + stable_memberships.append(membership) else: - stable_partitions.append((polyverts, membership, gamma_est, omega_est)) + stable_domains.append((polyverts, membership, gamma_est, omega_est)) - return stable_partitions + if return_membership_only: + return stable_memberships + return stable_domains -def prune_to_stable_partitions(G, parts, gamma_start, gamma_end, restrict_num_communities=None, - single_threaded=False): +def prune_to_stable_partitions( + G: ig.Graph, parts: set[Membership], + gamma_start: float, gamma_end: float, + restrict_num_communities: int | None = None, + single_threaded: bool = False +) -> list[Membership]: """ Runs our full pruning pipeline on a singlelayer network. Returns the pruned list of stable partitions whose parameter estimates are within the provided ``gamma_start`` and ``gamma_end`` bounds. @@ -487,17 +558,17 @@ def prune_to_stable_partitions(G, parts, gamma_start, gamma_end, restrict_num_co :param G: graph of interest :type G: igraph.Graph :param parts: partitions to prune - :type parts: iterable[tuple] + :type parts: set[tuple[int, ...]] :param gamma_start: starting gamma value for CHAMP :type gamma_start: float :param gamma_end: ending gamma value for CHAMP :type gamma_end: float :param restrict_num_communities: if not None, only use input partitions of this many communities - :type restrict_num_communities: int or None + :type restrict_num_communities: int | None :param single_threaded: if True, run the CHAMP step in serial :type single_threaded: bool :return: list of community membership tuples - :rtype: list[tuple[int]] + :rtype: list[tuple[int, ...]] """ if not G.is_connected(): @@ -521,16 +592,22 @@ def prune_to_stable_partitions(G, parts, gamma_start, gamma_end, restrict_num_co if len(parts) == 0: return [] - ranges = CHAMP_2D(G, parts, gamma_start, gamma_end, single_threaded=single_threaded) - gamma_estimates = ranges_to_gamma_estimates(G, ranges) - stable_parts = gamma_estimates_to_stable_partitions(gamma_estimates) + ranges: list[GammaRange] = CHAMP_2D( + G, parts, gamma_start, gamma_end, single_threaded=single_threaded + ) + gamma_estimates: list[GammaEstimate] = ranges_to_gamma_estimates(G, ranges) + stable_parts: list[Membership] = gamma_estimates_to_stable_partitions(gamma_estimates) return stable_parts -def prune_to_multilayer_stable_partitions(G_intralayer, G_interlayer, layer_vec, model, - parts, gamma_start, gamma_end, omega_start, omega_end, - restrict_num_communities=None, single_threaded=False): +def prune_to_multilayer_stable_partitions( + G_intralayer: ig.Graph, G_interlayer: ig.Graph, + layer_vec: LayerVector, model: str, + parts: set[Membership], gamma_start: float, gamma_end: float, + omega_start: float, omega_end: float, + restrict_num_communities: int | None = None, single_threaded: bool = False +) -> list[Membership]: """ Runs our full pruning pipeline on a multilayer network. Returns the pruned list of stable partitions whose parameter estimates are within the provided ``gamma_start``, ``gamma_end``, ``omega_start``, and ``omega_end`` @@ -565,7 +642,7 @@ def prune_to_multilayer_stable_partitions(G_intralayer, G_interlayer, layer_vec, :param model: network layer topology (temporal, multilevel, multiplex) :type model: str :param parts: partitions to prune - :type parts: iterable[tuple] + :type parts: set[tuple[int, ...]] :param gamma_start: starting gamma value for CHAMP :type gamma_start: float :param gamma_end: ending gamma value for CHAMP @@ -575,19 +652,19 @@ def prune_to_multilayer_stable_partitions(G_intralayer, G_interlayer, layer_vec, :param omega_end: ending omega value for CHAMP :type omega_end: float :param restrict_num_communities: if not None, only use input partitions of this many communities - :type restrict_num_communities: int or None + :type restrict_num_communities: int | None :param single_threaded: if True, run the CHAMP step in serial :type single_threaded: bool :return: list of community membership tuples - :rtype: list[tuple[int]] + :rtype: list[tuple[int, ...]] """ if single_threaded: raise NotImplementedError("Single-threaded multilayer CHAMP was never implemented. This would be fairly easy" "to add, so please raise an issue if this feature is desired.") # check graph is connected (cannot just check intra/interlayer graphs since they will both be disconnected) - G_combined = ig.Graph(edges=[(e.source, e.target) for e in G_intralayer.es] + - [(e.source, e.target) for e in G_interlayer.es]) + G_combined: ig.Graph = ig.Graph(edges=[(e.source, e.target) for e in G_intralayer.es] + + [(e.source, e.target) for e in G_interlayer.es]) if not G_combined.is_connected() or G_combined.vcount() != G_interlayer.vcount(): warnings.warn("The pruning pipeline has not been thoroughly tested on disconnected graphs. If you run into " "problems, consider using the largest connected component of your graph.") @@ -610,9 +687,13 @@ def prune_to_multilayer_stable_partitions(G_intralayer, G_interlayer, layer_vec, if len(parts) == 0: return [] - domains = CHAMP_3D(G_intralayer, G_interlayer, layer_vec, parts, gamma_start, gamma_end, - omega_start, omega_end) - domains_with_estimates = domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, domains, model) + domains: list[Domain] = CHAMP_3D( + G_intralayer, G_interlayer, layer_vec, parts, gamma_start, gamma_end, + omega_start, omega_end + ) + domains_with_estimates: list[DomainWithEstimates] = ( + domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, domains, model) + ) if any(o_est >= omega_end for _, _, g_est, o_est in domains_with_estimates if g_est is not None): warnings.warn(f"We are truncating some omega estimates to your choice of omega_end={omega_end}. You should " @@ -624,6 +705,8 @@ def prune_to_multilayer_stable_partitions(G_intralayer, G_interlayer, layer_vec, for polyverts, membership, g_est, o_est in domains_with_estimates if g_est is not None] - stable_parts = gamma_omega_estimates_to_stable_partitions(domains_with_estimates, return_membership_only=True) + stable_parts: list[Membership] = gamma_omega_estimates_to_stable_partitions( + domains_with_estimates, return_membership_only=True + ) return stable_parts diff --git a/utilities/partition_utilities.py b/utilities/partition_utilities.py index 9a603e3..685094b 100644 --- a/utilities/partition_utilities.py +++ b/utilities/partition_utilities.py @@ -1,43 +1,49 @@ +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import igraph from collections import defaultdict from sklearn.metrics import adjusted_mutual_info_score, normalized_mutual_info_score +from .types import IntSequence -def ami(p1, p2): - return adjusted_mutual_info_score(p1, p2) +def ami(p1: IntSequence, p2: IntSequence) -> float: + return float(adjusted_mutual_info_score(p1, p2)) -def nmi(p1, p2): - return normalized_mutual_info_score(p1, p2, average_method='arithmetic') +def nmi(p1: IntSequence, p2: IntSequence) -> float: + return float(normalized_mutual_info_score(p1, p2, average_method='arithmetic')) -def all_degrees(G): - return G.degree() +def all_degrees(G: "igraph.Graph") -> list[int]: + return list(G.degree()) -def in_degrees(G): - return G.indegree() +def in_degrees(G: "igraph.Graph") -> list[int]: + return list(G.indegree()) -def out_degrees(G): - return G.outdegree() +def out_degrees(G: "igraph.Graph") -> list[int]: + return list(G.outdegree()) -def membership_to_communities(membership): - communities = defaultdict(list) + +def membership_to_communities(membership: IntSequence) -> dict[int, list[int]]: + communities: defaultdict[int, list[int]] = defaultdict(list) for v, c in enumerate(membership): - communities[c].append(v) + communities[int(c)].append(v) return communities -def membership_to_layered_communities(membership, layer_membership): - layered_communities = defaultdict(list) +def membership_to_layered_communities(membership: IntSequence, layer_membership: IntSequence + ) -> defaultdict[tuple[int, int], list[int]]: + layered_communities: defaultdict[tuple[int, int], list[int]] = defaultdict(list) for v, c in enumerate(membership): - layered_communities[(c, layer_membership[v])].append(v) + layered_communities[(int(c), int(layer_membership[v]))].append(v) return layered_communities -def num_communities(membership): - n = len(set(membership)) +def num_communities(membership: IntSequence) -> int: + n: int = len(set(membership)) assert n == max(membership) + 1 return n diff --git a/utilities/plotting.py b/utilities/plotting.py index 4b2cfad..580e70a 100644 --- a/utilities/plotting.py +++ b/utilities/plotting.py @@ -1,5 +1,12 @@ +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import matplotlib.axes + import matplotlib.colorbar + import matplotlib.colors from collections import defaultdict +from collections.abc import Sequence from random import sample, shuffle +from numpy.typing import ArrayLike import matplotlib import matplotlib.pyplot as plt @@ -9,26 +16,30 @@ from matplotlib.patches import Polygon from .partition_utilities import num_communities, ami +from .types import (AxisLimits, DomainPoint, FlexibleDomain, FlexibleDomainWithEstimates, FloatArray, + FlexibleDomainVertices, GammaEstimate, IntArray, PlotColor) -def plot_adjacency(adj): +def plot_adjacency(adj: ArrayLike) -> None: plt.imshow(adj, cmap='Greys', interpolation='nearest') plt.title("Adjacency matrix") -def plot_estimates(gamma_estimates): +def plot_estimates( + gamma_estimates: Sequence[GammaEstimate] +) -> None: """Plot partition dominance ranges with gamma estimates. :param gamma_estimates: gamma estimates as returned from :meth:`~modularitypruning.parameter_estimation_utilities.ranges_to_gamma_estimates` """ - ax = plt.gca() - num_parts = len(gamma_estimates) + ax: "matplotlib.axes.Axes" = plt.gca() + num_parts: int = len(gamma_estimates) for i in range(num_parts): g_0, g_f, part, g_est = gamma_estimates[i] - k = num_communities(part) + k: int = num_communities(part) plt.plot([g_0, g_f], [k, k], zorder=1, linewidth=2) plt.scatter([g_0, g_f], [k, k], s=30, marker='x', zorder=1) @@ -39,8 +50,8 @@ def plot_estimates(gamma_estimates): g_0, g_f, part, g_est = gamma_estimates[i] k = num_communities(part) if g_est is not None: - middle = (g_0 + g_f) / 2 - k_new = 0 + middle: float = (g_0 + g_f) / 2 + k_new: int = 0 for gamma_start, gamma_end, part, _ in gamma_estimates: if gamma_start <= g_est <= gamma_end: k_new = num_communities(part) @@ -53,7 +64,8 @@ def plot_estimates(gamma_estimates): arrowprops=dict(alpha=0.5, facecolor="black", linewidth=1.5, arrowstyle="->, head_width=0.5, head_length=0.5")) else: - truncated_x, truncated_y = g_est, k_new + truncated_x: float = g_est + truncated_y: float = k_new if g_est > x_high: truncated_x = x_high elif g_est < x_low: @@ -71,7 +83,10 @@ def plot_estimates(gamma_estimates): # length_includes_head=True, alpha=0.5, zorder=2, **{"overhang": 0.5}) -def plot_2d_domains(domains, xlim, ylim, flip_axes=True, use_current_axes=False): +def plot_2d_domains( + domains: Sequence[FlexibleDomain], + xlim: AxisLimits, ylim: AxisLimits, flip_axes: bool = True, use_current_axes: bool = False +) -> None: """Plot partition dominance ranges in the (gamma, omega) plane, using the domains from CHAMP_3D. Limits output to xlim and ylim dimensions. Note that the plotting here has x=gamma and y=omega. @@ -81,30 +96,32 @@ def plot_2d_domains(domains, xlim, ylim, flip_axes=True, use_current_axes=False) :param ylim: plotting y limits """ if use_current_axes: - ax = plt.gca() + ax: "matplotlib.axes.Axes" = plt.gca() else: fig, ax = plt.subplots() - patches = [] + patches: list[Polygon] = [] for polyverts, membership in domains: if flip_axes: polyverts = [(x[1], x[0]) for x in polyverts] - polygon = Polygon(polyverts, closed=True) + polygon: Polygon = Polygon(polyverts, closed=True) patches.append(polygon) - cnorm = matplotlib.colors.Normalize(vmin=0, vmax=len(domains)) - cmap = plt.get_cmap("Set1") - available_colors = {cmap(cnorm(i)) for i in range(len(domains))} + cnorm: "matplotlib.colors.Normalize" = matplotlib.colors.Normalize(vmin=0, vmax=len(domains)) + cmap: "matplotlib.colors.Colormap" = plt.get_cmap("Set1") + available_colors: set[PlotColor] = {cmap(cnorm(i)) for i in range(len(domains))} + colors: list[PlotColor | None] if len(available_colors) == len(domains): colors = list(available_colors) shuffle(colors) else: # ensure that no adjacent domains have the same color colors = [None] * len(domains) - domain_polygons = [polyverts for polyverts, membership in domains] + domain_polygons: list[FlexibleDomainVertices] = ( + [polyverts for polyverts, membership in domains]) - vertex_to_neighbor_domain_indices = defaultdict(set) + vertex_to_neighbor_domain_indices: defaultdict[DomainPoint, set[int]] = defaultdict(set) for i, polyverts in enumerate(domain_polygons): for v in polyverts: # it turns out that scipy's halfspace intersection (using Qhull) maintains @@ -112,19 +129,23 @@ def plot_2d_domains(domains, xlim, ylim, flip_axes=True, use_current_axes=False) vertex_to_neighbor_domain_indices[tuple(v)].add(i) for i, this_polyverts in enumerate(domain_polygons): - neighboring_domains = set() + neighboring_domains: set[int] = set() for v in this_polyverts: neighboring_domains.update(vertex_to_neighbor_domain_indices[tuple(v)]) - neighbor_colors = {colors[i] for i in neighboring_domains} + neighbor_colors: set[PlotColor | None] = {colors[i] for i in neighboring_domains} colors[i] = sample(sorted(available_colors.difference(neighbor_colors)), 1)[0] - p = PatchCollection(patches, facecolors=colors, alpha=1.0, edgecolors='black', linewidths=1.5) + p: PatchCollection = PatchCollection(patches, facecolors=colors, alpha=1.0, edgecolors='black', linewidths=1.5) ax.add_collection(p) plt.xlim(xlim) plt.ylim(ylim) -def plot_2d_domains_with_estimates(domains_with_estimates, xlim, ylim, plot_estimate_points=True, flip_axes=True): +def plot_2d_domains_with_estimates( + domains_with_estimates: Sequence[FlexibleDomainWithEstimates], + xlim: AxisLimits, ylim: AxisLimits, + plot_estimate_points: bool = True, flip_axes: bool = True +) -> None: """Plot partition dominance ranges in the (gamma, omega) plane, using the domains from CHAMP_3D with their gamma and omega estimates overlaid. @@ -137,20 +158,20 @@ def plot_2d_domains_with_estimates(domains_with_estimates, xlim, ylim, plot_esti """ assert flip_axes fig, ax = plt.subplots() - patches = [] + patches: list[Polygon] = [] for polyverts, membership, gamma_est, omega_est in domains_with_estimates: if flip_axes: polyverts = [(x[1], x[0]) for x in polyverts] - polygon = Polygon(polyverts, closed=True) + polygon: Polygon = Polygon(polyverts, closed=True) patches.append(polygon) - centroid_x = np.mean([x[0] for x in polyverts]) - centroid_y = np.mean([x[1] for x in polyverts]) + centroid_x: float = np.mean([x[0] for x in polyverts]) + centroid_y: float = np.mean([x[1] for x in polyverts]) if gamma_est is not None and omega_est is not None: - width = xlim[1] - xlim[0] + width: float = xlim[1] - xlim[0] # prune arrows too far outside plotting range to avoid excess visual noise if centroid_x < xlim[0] - 0.05 * width or centroid_x > xlim[1] + 0.05 * width: @@ -172,23 +193,27 @@ def plot_2d_domains_with_estimates(domains_with_estimates, xlim, ylim, plot_esti flip_axes=flip_axes, use_current_axes=True) -def plot_2d_domains_with_num_communities(domains_with_estimates, xlim, ylim, flip_axes=True, K_max=None, tick_step=2): +def plot_2d_domains_with_num_communities( + domains_with_estimates: Sequence[FlexibleDomainWithEstimates], + xlim: AxisLimits, ylim: AxisLimits, + flip_axes: bool = True, K_max: int | None = None, tick_step: int = 2 +) -> None: """Plot partition dominance ranges in the (gamma, omega) plane, using the domains from CHAMP_3D and coloring by the number of communities. Limits output to xlim and ylim dimensions. Note that the plotting here has x=omega and y=gamma.""" assert flip_axes fig, ax = plt.subplots() - patches = [] - cm = matplotlib.cm.viridis - Ks = [] + patches: list[Polygon] = [] + cm: "matplotlib.colors.Colormap" = matplotlib.cm.viridis + Ks: list[int] = [] for polyverts, membership, gamma_est, omega_est in domains_with_estimates: if flip_axes: polyverts = [(x[1], x[0]) for x in polyverts] if any(xlim[0] <= x[0] <= xlim[1] and ylim[0] <= x[1] <= ylim[1] for x in polyverts): - polygon = Polygon(polyverts, closed=True) + polygon: Polygon = Polygon(polyverts, closed=True) patches.append(polygon) Ks.append(num_communities(membership)) @@ -197,11 +222,11 @@ def plot_2d_domains_with_num_communities(domains_with_estimates, xlim, ylim, fli else: K_max = max(Ks) - p = PatchCollection(patches, cmap=cm, alpha=1.0, edgecolors='black', linewidths=2) + p: PatchCollection = PatchCollection(patches, cmap=cm, alpha=1.0, edgecolors='black', linewidths=2) p.set_array(np.array(Ks)) ax.add_collection(p) - cbar = plt.colorbar(p, ticks=range(2, K_max + 1, tick_step)) + cbar: "matplotlib.colorbar.Colorbar" = plt.colorbar(p, ticks=range(2, K_max + 1, tick_step)) cbar.ax.tick_params(labelsize=12) cbar.set_label("Number of Communities", fontsize=16, labelpad=10) @@ -209,53 +234,58 @@ def plot_2d_domains_with_num_communities(domains_with_estimates, xlim, ylim, fli plt.ylim(ylim) -def plot_2d_domains_with_ami(domains_with_estimates, ground_truth, xlim, ylim, flip_axes=True): +def plot_2d_domains_with_ami( + domains_with_estimates: Sequence[FlexibleDomainWithEstimates], + ground_truth: list[int], xlim: AxisLimits, ylim: AxisLimits, + flip_axes: bool = True +) -> None: """Plot partition dominance ranges in the (gamma, omega) plane, using the domains from CHAMP_3D and coloring by the AMI between the partitions and ground truth. Limits output to xlim and ylim dimensions. Note that the plotting here has x=omega and y=gamma.""" assert flip_axes fig, ax = plt.subplots() - patches = [] - cm = matplotlib.cm.copper - amis = [] + patches: list[Polygon] = [] + cm: "matplotlib.colors.Colormap" = matplotlib.cm.copper + amis: list[float] = [] for polyverts, membership, gamma_est, omega_est in domains_with_estimates: if flip_axes: polyverts = [(x[1], x[0]) for x in polyverts] if any(xlim[0] <= x[0] <= xlim[1] and ylim[0] <= x[1] <= ylim[1] for x in polyverts): - polygon = Polygon(polyverts, closed=True) + polygon: Polygon = Polygon(polyverts, closed=True) patches.append(polygon) amis.append(ami(membership, ground_truth)) - p = PatchCollection(patches, cmap=cm, alpha=1.0, edgecolors='black', linewidths=2) + p: PatchCollection = PatchCollection(patches, cmap=cm, alpha=1.0, edgecolors='black', linewidths=2) p.set_array(np.array(amis + [1.0])) # this extends the colorbar to include 1.0 ax.add_collection(p) - cbar = plt.colorbar(p) + cbar: "matplotlib.colorbar.Colorbar" = plt.colorbar(p) cbar.set_label('AMI', fontsize=14, labelpad=15) plt.xlim(xlim) plt.ylim(ylim) -def plot_multiplex_community(membership, layer_vec): +def plot_multiplex_community(membership: IntArray, + layer_vec: IntArray) -> "matplotlib.axes.Axes": """Plot a visualization of multiplex community membership :param membership: partition membership vector - :type membership: tuple[int] + :type membership: numpy.ndarray :param layer_vec: list of each vertex's layer membership - :type layer_vec: list[int] + :type layer_vec: numpy.ndarray """ - T = max(layer_vec) + 1 - N = len(membership) // T - communities_per_layer = np.zeros((N, T)) + T: int = max(layer_vec) + 1 + N: int = len(membership) // T + communities_per_layer: FloatArray = np.zeros((N, T)) for i, layer in enumerate(range(T)): communities_per_layer[:, i] = membership[layer_vec == layer] - ax = plt.axes() - cmap = sbn.cubehelix_palette(as_cmap=True) + ax: "matplotlib.axes.Axes" = plt.axes() + cmap: "matplotlib.colors.ListedColormap" = sbn.cubehelix_palette(as_cmap=True) ax.pcolormesh(communities_per_layer, cmap=cmap, linewidth=0, rasterized=True) return ax diff --git a/utilities/progress.py b/utilities/progress.py deleted file mode 100644 index 1f169b2..0000000 --- a/utilities/progress.py +++ /dev/null @@ -1,31 +0,0 @@ -import math -from time import time - - -class Progress: - """Progress bar with length :length: and :total: number of iterations.""" - - def __init__(self, total, length=80, name="Progress:"): - self.total = total - self.length = length - self.start = time() - self.ipc = total / length - self.i = 0 - self.name = name - - def update(self, i): - filled_count = math.floor(self.length * i / self.total) - filled_portion = "#" * filled_count - remaining_portion = "." * (self.length - filled_count) - time_elapsed = time() - self.start - time_estimate = time_elapsed * (self.total / i) if i > 0 else math.inf - print(f"\r{self.name} [{filled_portion}{remaining_portion}] Time: {time_elapsed:.1f} / {time_estimate:.1f} s", - end='', flush=True) - self.i = i - - def increment(self): - self.update(self.i + 1) - - def done(self): - self.update(self.total) - print() diff --git a/utilities/py.typed b/utilities/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/utilities/types.py b/utilities/types.py new file mode 100644 index 0000000..7f6f287 --- /dev/null +++ b/utilities/types.py @@ -0,0 +1,62 @@ +"""Type aliases for mypy to make the code more readable.""" + +from collections.abc import Sequence +from typing import Callable, TypeAlias + +import leidenalg +import numpy as np + + +# numpy arrays +FloatArray = np.ndarray[tuple[int, ...], np.dtype[np.float64]] +IntArray = np.ndarray[tuple[int, ...], np.dtype[np.int64]] + +# mypy doesn't consider numpy arrays to be a sequence +IntSequence = Sequence[int] | IntArray +FloatSequence = Sequence[float] | FloatArray + +# community membership vector and layer membership vector +Membership = tuple[int, ...] +LayerVector = list[int] +LayerTuple = tuple[int, ...] # needed for hashing + +# collection of partitions for CHAMP +Partitions = Sequence[Membership] | set[Membership] + +# Leiden partitions +LeidenPartition: TypeAlias = leidenalg.RBConfigurationVertexPartition +MultilayerLeidenParts: TypeAlias = tuple[list[LeidenPartition], leidenalg.CPMVertexPartition] +SingleLayerLeidenResult: TypeAlias = LeidenPartition | Membership +MultilayerLeidenResult: TypeAlias = list[LeidenPartition] | Membership + +# CHAMP 2D domains and parameter estimates +GammaEstimateValue = np.float64 | None +GammaEstimate = tuple[np.float64, np.float64, Membership, GammaEstimateValue] +GammaRange = tuple[np.float64, np.float64, Membership] + +# CHAMP 3D domains and parameter estimates +DomainVertices = list[FloatArray] +Domain = tuple[DomainVertices, Membership] # plus membership +GammaOmegaEstimate = tuple[GammaEstimateValue, float] +DomainWithEstimates = tuple[DomainVertices, Membership, GammaEstimateValue, float] + +# (A_hats, P_hats) and (A_hats, P_hats, C_hats) coefficients from CHAMP +PartitionCoefficients2D = tuple[IntArray, FloatArray] +PartitionCoefficients3D = tuple[IntArray, FloatArray, IntArray] + +# SBM parameter estimates +SinglelayerSBMParameters = tuple[float, float] +MultilayerSBMParameters = tuple[float, float, float, int] + +# parameter estimation functions +OmegaUpdateFunc = Callable[[float, float, float, int], float] +PersistenceFunc = Callable[[list[int]], float] +GammaUpdateFunc = Callable[[float, float], GammaEstimateValue] + +# more flexible plotting equivalents (allows tuples, etc.) +FlexibleDomainVertices = Sequence[FloatArray] | Sequence[tuple[float, ...]] +FlexibleDomain = tuple[FlexibleDomainVertices, Membership] +FlexibleDomainWithEstimates = tuple[FlexibleDomainVertices, Membership, GammaEstimateValue, float] +PlotColor = tuple[float, ...] +AxisLimits = tuple[float, float] +DomainPoint = tuple[float, ...]