From 3b38ab4011deed5daabed21ad721396cd1f08d98 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Tue, 13 Oct 2026 12:40:07 -0400
Subject: [PATCH 01/11] Add basic mypy config and CI action

---
 .github/workflows/python-package.yml | 12 ++++++++----
 mypy.ini                             |  6 ++++++
 2 files changed, 14 insertions(+), 4 deletions(-)
 create mode 100644 mypy.ini

diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index d1ef347..50afec5 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.9", "3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
 
     steps:
       - uses: actions/checkout@v3
@@ -39,6 +38,11 @@ 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
+        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..860ec61
--- /dev/null
+++ b/mypy.ini
@@ -0,0 +1,6 @@
+[mypy]
+strict = True
+files = utilities/*.py
+
+[mypy-sklearn.*,igraph.*,leidenalg.*]
+ignore_missing_imports = True
\ No newline at end of file

From 7c3085332c9c23048551ba81e6de0c83d31c9a64 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Tue, 13 Oct 2026 21:26:32 -0400
Subject: [PATCH 02/11] Autogenerate first pass for type hints with righttyper

Used `righttyper run --root utilities -m pytest -- -m "not serial" tests/`
Current package version pulled from GitHub (a18f2b0 from Sep 9 2026) since pip is greatly outdated
---
 utilities/champ_utilities.py                |  93 ++++++-------
 utilities/leiden_utilities.py               |  75 +++++------
 utilities/louvain_utilities.py              |   8 +-
 utilities/parameter_estimation.py           |  44 ++++---
 utilities/parameter_estimation_utilities.py | 137 ++++++++++----------
 utilities/partition_utilities.py            |  13 +-
 utilities/plotting.py                       |  95 +++++++-------
 7 files changed, 244 insertions(+), 221 deletions(-)

diff --git a/utilities/champ_utilities.py b/utilities/champ_utilities.py
index 7dc169c..79a90c8 100644
--- a/utilities/champ_utilities.py
+++ b/utilities/champ_utilities.py
@@ -1,3 +1,8 @@
+from typing import TYPE_CHECKING, Any
+if TYPE_CHECKING:
+    import igraph
+    import multiprocessing.pool
+    import scipy.optimize
 import warnings
 from collections import defaultdict
 from math import floor
@@ -13,7 +18,7 @@
     membership_to_layered_communities
 
 
-def get_interior_point(halfspaces, initial_num_sampled=50, full_retry_limit=10):
+def get_interior_point(halfspaces: np.ndarray[Any, np.dtype[np.float64]], initial_num_sampled: int=50, full_retry_limit: int=10) -> np.ndarray[Any, np.dtype[np.float64]]:
     """
     Find interior point of halfspaces (needed to perform halfspace intersection)
 
@@ -41,21 +46,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: np.ndarray[Any, np.dtype[np.float64]] = 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: np.ndarray[Any, np.dtype[np.float64]] = np.reshape(np.linalg.norm(sampled_hs[:, :-1], axis=1), (sampled_hs.shape[0], 1))
+        c: np.ndarray[Any, np.dtype[np.float64]] = np.zeros((sampled_hs.shape[1],))
         c[-1] = -1
-        A = np.hstack((sampled_hs[:, :-1], norm_vector))
-        b = -sampled_hs[:, -1:]
+        A: np.ndarray[Any, np.dtype[np.float64]] = np.hstack((sampled_hs[:, :-1], norm_vector))
+        b: np.ndarray[Any, np.dtype[np.float64]] = -sampled_hs[:, -1:]
 
-        res = linprog(c, A_ub=A, b_ub=b, bounds=(-np.inf, np.inf), method='highs-ipm')
+        res: "scipy.optimize.OptimizeResult" = linprog(c, A_ub=A, b_ub=b, bounds=(-np.inf, np.inf), method='highs-ipm')
 
         # check solution if optimization succeeded, even with difficulties
         #     res.status codes
@@ -65,7 +70,7 @@ 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]
+            intpt: np.ndarray[Any, np.dtype[np.float64]] = 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 +91,7 @@ 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[T1: (float, int)](G: "igraph.Graph", all_parts: list[tuple[int, ...]]|set[tuple[int, ...]], gamma_0: T1, gamma_f: T1, single_threaded: bool=False) -> list[tuple[np.float64, np.float64, tuple[int, ...]]]:
     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.
@@ -113,15 +118,15 @@ def CHAMP_2D(G, all_parts, gamma_0, gamma_f, single_threaded=False):
         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: tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]]] = 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),
+    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: np.ndarray[Any, np.dtype[np.float64]] = np.vstack((halfspaces_from_coefficients_2D(*partition_coefficients),
                             np.array([[0, 1, -top], [0, -1, bottom],
                                       [1, 0, -gamma_f], [-1, 0, gamma_0]])))
 
@@ -129,18 +134,18 @@ def CHAMP_2D(G, all_parts, gamma_0, gamma_f, single_threaded=False):
     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: np.ndarray[Any, np.dtype[np.float64]] = 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[np.ndarray[Any, np.dtype[np.float64]]]] = 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[tuple[np.float64, np.float64, tuple[int, ...]]] = []
     for i, intersections in facets_by_halfspace.items():
         x1, x2 = intersections[0][0], intersections[1][0]
         if x1 > x2:
@@ -150,7 +155,7 @@ 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[T1: (float, int), T2: (float, int)](G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int], all_parts: list[tuple[int, ...]]|set[tuple[int, ...]], gamma_0: T1, gamma_f: T1, omega_0: T2, omega_f: T2) -> list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]]:
     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``.
 
@@ -183,15 +188,15 @@ def CHAMP_3D(G_intralayer, G_interlayer, layer_vec, all_parts, gamma_0, gamma_f,
         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: tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.int64]]] = 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),
+    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: np.ndarray[Any, np.dtype[np.float64]] = 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]]))))
@@ -200,23 +205,23 @@ def CHAMP_3D(G_intralayer, G_interlayer, layer_vec, all_parts, gamma_0, gamma_f,
     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: np.ndarray[Any, np.dtype[np.float64]] = 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[np.ndarray[Any, np.dtype[np.float64]]]] = 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[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]] = []
     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: np.ndarray[Any, np.dtype[np.float64]] = 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]))
 
@@ -255,20 +260,20 @@ 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: list[tuple[int, ...]]|set[tuple[int, ...]], single_threaded: bool=False) -> tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]]]:
     """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[tuple[int, ...]]] = [
             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,
+        pool: "multiprocessing.pool.Pool" = Pool(processes=cpu_count())
+        results: list[tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]]]] = pool.starmap(partition_coefficients_2D_serial,
                                [(G, partition_chunk) for partition_chunk in partition_chunks])
         pool.close()
 
@@ -280,7 +285,7 @@ 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: np.ndarray[Any, np.dtype[np.int64]], P_hats: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
     """Converts partitions' coefficients to halfspace normal, offset
 
     Q >= -P_hat*gamma + A_hat
@@ -290,7 +295,7 @@ def halfspaces_from_coefficients_2D(A_hats, P_hats):
     return np.vstack((-P_hats, -np.ones_like(P_hats), A_hats)).T
 
 
-def halfspaces_from_coefficients_3D(A_hats, P_hats, C_hats):
+def halfspaces_from_coefficients_3D(A_hats: np.ndarray[Any, np.dtype[np.int64]], P_hats: np.ndarray[Any, np.dtype[np.float64]], C_hats: np.ndarray[Any, np.dtype[np.int64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
     """Converts partitions' coefficients to halfspace normal, offset
 
     Q >= C_hat*omega - P_hat*gamma + A_hat
@@ -349,23 +354,23 @@ 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: list[int], partitions: list[tuple[int, ...]]) -> tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.int64]]]:
     """Computes partitions coefficients in parallel by calling partition_coefficients_3D_serial"""
     partitions = list(partitions)
-    partition_chunks = [
+    partition_chunks: list[list[tuple[int, ...]]] = [
         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,
+    pool: "multiprocessing.pool.Pool" = Pool(processes=cpu_count())
+    results: list[tuple[np.ndarray[Any, np.dtype[np.float64]], ...]|tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.int64]]]] = 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: np.ndarray[Any, np.dtype[np.int64]] = np.array([v for A_hats, P_hats, C_hats in results for v in A_hats])
+    P_hats: np.ndarray[Any, np.dtype[np.float64]] = np.array([v for A_hats, P_hats, C_hats in results for v in P_hats])
+    C_hats: np.ndarray[Any, np.dtype[np.int64]] = 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..5502c7f 100644
--- a/utilities/leiden_utilities.py
+++ b/utilities/leiden_utilities.py
@@ -1,3 +1,4 @@
+from typing import Any
 import functools
 from multiprocessing import Pool, cpu_count
 
@@ -6,11 +7,11 @@
 import numpy as np
 from tqdm import tqdm
 
-LOW_MEMORY_THRESHOLD = 1e9  # 1 GB
+LOW_MEMORY_THRESHOLD: float = 1e9  # 1 GB
 
 
 @functools.lru_cache(maxsize=1000)
-def sorted_tuple(t):
+def sorted_tuple(t: tuple[int, ...]) -> tuple[int, ...]:
     """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).
 
@@ -20,11 +21,11 @@ def sorted_tuple(t):
     :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[np.int64, 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):
+def singlelayer_leiden(G: ig.Graph, gamma: float, return_partition: bool=False) -> leidenalg.RBConfigurationVertexPartition|tuple[int, ...]:
     r"""Run the Leiden modularity maximization algorithm at a single :math:`\gamma` value.
 
     :param G: graph of interest
@@ -39,7 +40,7 @@ def singlelayer_leiden(G, gamma, return_partition=False):
     if 'weight' not in G.es:
         G.es['weight'] = [1.0] * G.ecount()
 
-    partition = leidenalg.find_partition(G, leidenalg.RBConfigurationVertexPartition, weights='weight',
+    partition: leidenalg.RBConfigurationVertexPartition = leidenalg.find_partition(G, leidenalg.RBConfigurationVertexPartition, weights='weight',
                                          resolution_parameter=gamma)
 
     if return_partition:
@@ -53,19 +54,19 @@ def _wrapped_singlelayer_leiden(args):
     return singlelayer_leiden(*args)
 
 
-def leiden_part(G):
+def leiden_part(G: ig.Graph) -> leidenalg.RBConfigurationVertexPartition:
     return leidenalg.RBConfigurationVertexPartition(G)
 
 
-def leiden_part_with_membership(G, membership):
+def leiden_part_with_membership(G: ig.Graph, membership: list[int]|tuple[int, ...]) -> leidenalg.RBConfigurationVertexPartition:
     if isinstance(membership, np.ndarray):
         membership = membership.tolist()
-    part = leiden_part(G)
+    part: leidenalg.RBConfigurationVertexPartition = 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: list[int]) -> list[ig.Graph]:
     """
     Split intralayer network into a separate network for each layer.
 
@@ -85,25 +86,25 @@ 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: tuple[int, ...]) -> 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):
+def multilayer_leiden(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gamma: float, omega: float, optimiser: leidenalg.Optimiser|None=None, return_partition: bool=False) -> list[leidenalg.RBConfigurationVertexPartition]|tuple[int, ...]:
     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
@@ -132,14 +133,14 @@ 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)
+    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = [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')
+    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])
 
@@ -154,21 +155,21 @@ def _wrapped_multilayer_leiden(args):
     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: list[int]) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
     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')
+    G_split_layers: list[ig.Graph] = split_intralayer_leiden_graph(G_intralayer, layer_membership)
+    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = [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')
+    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: list[int], community_membership: tuple[int, ...]) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
     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,11 +179,11 @@ 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: list[float]) -> set[tuple[int, ...]]:
     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: np.ndarray[Any, np.dtype[np.float64]], show_progress: bool=True) -> set[tuple[int, ...]]:
     r"""Runs the Leiden modularity maximization algorithm at each provided :math:`\gamma` value, using all CPU cores.
 
     :param G: graph of interest
@@ -194,10 +195,10 @@ def repeated_parallel_leiden_from_gammas(G, gammas, show_progress=True):
     :return: a set of all unique partitions returned by the Leiden algorithm
     :rtype: set of tuple[int]
     """
-    total = set()
-    pool_chunk_size = max(1, len(gammas) // (cpu_count() * 100))
+    total: set[tuple[int, ...]] = 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],
+        pool_iterator: tqdm = 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 +209,13 @@ 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):
+def repeated_leiden_from_gammas_omegas(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gammas: list[float], omegas: list[float]) -> set[tuple[int, ...]]:
     return {sorted_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: list[int], gammas: np.ndarray[Any, np.dtype[np.float64]], omegas: np.ndarray[Any, np.dtype[np.float64]],
+                                                show_progress: bool=True) -> set[tuple[int, ...]]:
     """
     Runs leidenalg at each gamma and omega in ``gammas`` and ``omegas``, using all CPU cores available.
 
@@ -236,12 +237,12 @@ def repeated_parallel_leiden_from_gammas_omegas(G_intralayer, G_interlayer, laye
     :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]
+    resolution_parameter_points: list[tuple[np.float64, np.float64]] = [(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[tuple[int, ...]] = 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(
+        pool_iterator: tqdm = 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
diff --git a/utilities/louvain_utilities.py b/utilities/louvain_utilities.py
index f5cceb5..321fe53 100644
--- a/utilities/louvain_utilities.py
+++ b/utilities/louvain_utilities.py
@@ -6,6 +6,10 @@
 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).
 """
+from typing import TYPE_CHECKING, Callable
+if TYPE_CHECKING:
+    import igraph
+    import leidenalg
 import warnings
 from math import ceil
 from multiprocessing import Pool, cpu_count
@@ -27,9 +31,9 @@
               "Please use leiden_utilities in the future when possible.", DeprecationWarning)
 
 
-def __getattr__(name):
+def __getattr__(name: str) -> "Callable[..., set[tuple[int, ...]]]|Callable[[igraph.Graph, list[float]], set[tuple[int, ...]]]|Callable[[igraph.Graph, list[int]|tuple[int, ...]], leidenalg.RBConfigurationVertexPartition]|Callable[[igraph.Graph], leidenalg.RBConfigurationVertexPartition]|None":
     if "louvain" in name:
-        leiden_name = name.replace("louvain", "leiden")
+        leiden_name: str = 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)
diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py
index 466c959..be91afc 100644
--- a/utilities/parameter_estimation.py
+++ b/utilities/parameter_estimation.py
@@ -1,3 +1,7 @@
+from typing import TYPE_CHECKING, Callable
+if TYPE_CHECKING:
+    import igraph
+    import numpy
 import leidenalg
 
 from .leiden_utilities import singlelayer_leiden, multilayer_leiden
@@ -6,8 +10,8 @@
 from .partition_utilities import in_degrees
 
 
-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[numpy.float64, leidenalg.RBConfigurationVertexPartition]":
     """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,10 +38,10 @@ 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) -> leidenalg.RBConfigurationVertexPartition:
             return singlelayer_leiden(G, resolution_param, return_partition=True)
     elif method == "2-spinglass":
         def maximize_modularity(resolution_param):
@@ -46,10 +50,10 @@ def maximize_modularity(resolution_param):
     else:
         raise ValueError(f"Community detection method {method} not supported")
 
-    def estimate_SBM_parameters(partition):
+    def estimate_SBM_parameters(partition: leidenalg.RBConfigurationVertexPartition) -> tuple[float, float]:
         return estimate_singlelayer_SBM_parameters(G, partition, m=m)
 
-    def update_gamma(omega_in, omega_out):
+    def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
         return gamma_estimate_from_parameters(omega_in, omega_out)
 
     part, last_gamma = None, None
@@ -80,7 +84,7 @@ def update_gamma(omega_in, omega_out):
     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: list[int], 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.
 
@@ -94,7 +98,7 @@ def check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, mo
     :param Nt: vector of nodes per layer
     """
 
-    rules = [T > 1,
+    rules: list[bool|str] = [T > 1,
              "Graph must have multiple layers",
              G_interlayer.is_directed(),
              "Interlayer graph should be directed",
@@ -127,9 +131,9 @@ def check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, mo
         raise ValueError("Input graph is malformed\n" + "\n".join(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: list[int], gamma: float=1.0, omega: float=1.0,
+                                                         gamma_tol: float=1e-2, omega_tol: float=5e-2, omega_max: int=1000, max_iter: int=25,
+                                                         model: str='temporal', verbose: bool=False) -> "tuple[numpy.float64, float, list[leidenalg.RBConfigurationVertexPartition]]":
     """
     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.
@@ -170,29 +174,29 @@ def iterative_multilayer_resolution_parameter_estimation(G_intralayer, G_interla
     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: Callable[[float, float|float|int, float, int], float|float|int] = omega_function_from_model(model, omega_max, T=T)
+    update_gamma: "Callable[[float, float], numpy.float64|None]" = gamma_estimate_from_parameters
 
-    def maximize_modularity(intralayer_resolution, interlayer_resolution):
+    def maximize_modularity(intralayer_resolution: float, interlayer_resolution: float) -> list[leidenalg.RBConfigurationVertexPartition]:
         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: leidenalg.RBConfigurationVertexPartition) -> tuple[float, float, float, int]:
         return estimate_multilayer_SBM_parameters(G_intralayer, G_interlayer, layer_vec, partition, model,
                                                   N=N, T=T, Nt=Nt, m_t=m_t)
 
diff --git a/utilities/parameter_estimation_utilities.py b/utilities/parameter_estimation_utilities.py
index d9321c5..a207815 100644
--- a/utilities/parameter_estimation_utilities.py
+++ b/utilities/parameter_estimation_utilities.py
@@ -1,3 +1,4 @@
+from typing import Any, Callable
 import warnings
 from math import log
 
@@ -11,7 +12,7 @@
 from .partition_utilities import num_communities
 
 
-def estimate_singlelayer_SBM_parameters(G, partition, m=None):
+def estimate_singlelayer_SBM_parameters(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition, m: float|None=None) -> tuple[float, float]:
     r"""Estimate singlelayer SBM parameters from a graph and a partition.
 
     See https://doi.org/10.1103/PhysRevE.94.052315 for more details.
@@ -30,25 +31,25 @@ 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: list[int], partition: leidenalg.RBConfigurationVertexPartition, model: str, N: int|None=None, T: int|None=None,
+                                       Nt: list[int]|None=None, m_t: list[float]|None=None) -> tuple[float, float, float, int]:
     """Estimates multilayer SBM parameters from a graph and a partition
 
     :param G_intralayer: intralayer graph of interest
@@ -93,42 +94,42 @@ 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[float]|list[float|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: Callable[[list[int]], float] = 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))
+        def f(x: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
+            coeff: float = 2 * (1 - 1 / K) / (T * (T - 1))
             return coeff * sum((T - n) * x ** n for n in range(1, T)) + 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 +140,7 @@ def f(x):
     return theta_in, theta_out, p, K
 
 
-def gamma_estimate(G, partition):
+def gamma_estimate(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition|tuple[int, ...]) -> np.float64|None:
     """Compute the "correct" value of gamma where modularity maximization becomes equivalent to maximum likelihood
     methods on a degree-corrected, planted partition stochastic block model.
 
@@ -163,7 +164,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) -> np.float64|None:
     """
     Compute the "correct" value of gamma as in
     :meth:`~modularitypruning.parameter_estimation_utilities.gamma_estimate` from SBM parameters.
@@ -182,7 +183,7 @@ def gamma_estimate_from_parameters(omega_in, omega_out):
     return (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: int=1000) -> float:
     """Returns the omega estimate for a multiplex multilayer model
 
     :param theta_in: SBM parameter
@@ -213,7 +214,7 @@ 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[T1: (float, int)](theta_in: float, theta_out: T1, p: float, K: int, omega_max: int=1000) -> T1:
     """Returns the omega estimate for a temporal or multilevel multilayer model
 
     :param theta_in: SBM parameter
@@ -239,13 +240,13 @@ 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))
 
 
-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: list[int], 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 +254,12 @@ 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))
 
 
-def omega_function_from_model(model, omega_max, T):
+def omega_function_from_model[T1: (float, int)](model: str, omega_max: int, T: int) -> Callable[[float, T1, float, int], T1]:
     """Returns an appropriate function (depending on the model) for computing omega from multilayer SBM parameters
 
     Specifically, it will return versions of
@@ -274,10 +275,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[T2: (float, int)](theta_in: float, theta_out: T2, p: float, K: int) -> T2:
             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 +286,7 @@ 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: list[int]|None=None, N: int|None=None, T: int|None=None, Nt: list[int]|None=None) -> Callable[[list[int]], float]:
     """
     Returns a function to calculate persistence according to a given multilayer model
 
@@ -303,19 +304,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 +324,8 @@ 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: list[int], membership: tuple[int, ...], omega_max: int=1000, model: str='temporal',
+                         N: None=None, T: int|None=None, Nt: None=None, m_t: None=None) -> tuple[np.float64|None, float]:
     """Returns the (gamma, omega) estimate for a multilayer network and a partition
 
     :param G_intralayer: intralayer graph of interest
@@ -353,18 +354,18 @@ def gamma_omega_estimate(G_intralayer, G_interlayer, layer_vec, membership, omeg
     if T is None:
         T = max(layer_vec) + 1  # layer  count
 
-    partition = leiden_part_with_membership(G_intralayer, membership)
+    partition: leidenalg.RBConfigurationVertexPartition = 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: Callable[[float, float|float|int, float, int], float|float|int] = omega_function_from_model(model, omega_max, T=T)
+    update_gamma: Callable[[float, float], np.float64|None] = gamma_estimate_from_parameters
 
-    gamma = update_gamma(theta_in, theta_out)
-    omega = update_omega(theta_in, theta_out, p, K)
+    gamma: np.float64|None = 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[tuple[np.float64, np.float64, tuple[int, ...]]]) -> list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]]:
     """
     Compute gamma estimates as in :meth:`~modularitypruning.parameter_estimation_utilities.gamma_estimate`, given
     domains of optimality from :meth:`~modularitypruning.champ_utilities.CHAMP_2D`.
@@ -382,7 +383,7 @@ def ranges_to_gamma_estimates(G, ranges):
             gamma_start, gamma_end, part in ranges]
 
 
-def gamma_estimates_to_stable_partitions(gamma_estimates):
+def gamma_estimates_to_stable_partitions(gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]]) -> list[tuple[int, ...]]:
     """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`.
 
@@ -399,7 +400,7 @@ def gamma_estimates_to_stable_partitions(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: list[int], domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]], model: str='temporal') -> list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]]:
     """
     Compute (gamma, omega) estimates as in
     :meth:`~modularitypruning.parameter_estimation_utilities.gamma_omega_estimate`, given domains of optimality from
@@ -420,7 +421,7 @@ def domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, doma
     :rtype: list of tuple[list[float], tuple[int], float, float]
     """
 
-    domains_with_estimates = []
+    domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]] = []
     for polyverts, membership in domains:
         gamma_est, omega_est = gamma_omega_estimate(G_intralayer, G_interlayer, layer_vec, membership,
                                                     model=model)
@@ -428,7 +429,7 @@ 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):
+def gamma_omega_estimates_to_stable_partitions(domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], return_membership_only: bool=False) -> list[tuple[int, ...]]:
     """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`.
@@ -442,27 +443,27 @@ def gamma_omega_estimates_to_stable_partitions(domains_with_estimates, return_me
     :rtype: list of tuple[int]
     """
 
-    def left_or_right(x1, y1, x2, y2, x, y):
+    def left_or_right(x1: np.float64, y1: np.float64, x2: np.float64, y2: np.float64, x: np.float64, y: float) -> np.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_partitions: list[tuple[int, ...]] = []
 
     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: np.float64 = np.mean([x[0] for x in polyverts])
+        centroid_y: np.float64 = np.mean([x[1] for x in polyverts])
+        polygon_edges: list[tuple[np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.float64]]]] = []
         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[np.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))
 
@@ -476,8 +477,8 @@ def left_or_right(x1, y1, x2, y2, x, y):
     return stable_partitions
 
 
-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[tuple[int, ...]], gamma_start: int, gamma_end: int, restrict_num_communities: None=None,
+                               single_threaded: bool=False) -> list[tuple[int, ...]]:
     """
     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.
@@ -521,16 +522,16 @@ 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[tuple[np.float64, np.float64, tuple[int, ...]]] = CHAMP_2D(G, parts, gamma_start, gamma_end, single_threaded=single_threaded)
+    gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]] = ranges_to_gamma_estimates(G, ranges)
+    stable_parts: list[tuple[int, ...]] = 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: list[int], model: str,
+                                          parts: set[tuple[int, ...]], gamma_start: int, gamma_end: int, omega_start: int, omega_end: int,
+                                          restrict_num_communities: None=None, single_threaded: bool=False) -> list[tuple[int, ...]]:
     """
     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``
@@ -586,7 +587,7 @@ def prune_to_multilayer_stable_partitions(G_intralayer, G_interlayer, layer_vec,
                                   "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] +
+    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 "
@@ -610,9 +611,9 @@ 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,
+    domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]] = 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_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]] = 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 +625,6 @@ 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[tuple[int, ...]] = 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..1ee9efa 100644
--- a/utilities/partition_utilities.py
+++ b/utilities/partition_utilities.py
@@ -1,9 +1,12 @@
+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
 
 
-def ami(p1, p2):
+def ami(p1: tuple[int, ...], p2: list[int]) -> float:
     return adjusted_mutual_info_score(p1, p2)
 
 
@@ -11,11 +14,11 @@ def nmi(p1, p2):
     return normalized_mutual_info_score(p1, p2, average_method='arithmetic')
 
 
-def all_degrees(G):
+def all_degrees(G: "igraph.Graph") -> list[int]:
     return G.degree()
 
 
-def in_degrees(G):
+def in_degrees(G: "igraph.Graph") -> list[int]:
     return G.indegree()
 
 
@@ -37,7 +40,7 @@ def membership_to_layered_communities(membership, layer_membership):
     return layered_communities
 
 
-def num_communities(membership):
-    n = len(set(membership))
+def num_communities(membership: list[int]|tuple[int, ...]) -> 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..2256d23 100644
--- a/utilities/plotting.py
+++ b/utilities/plotting.py
@@ -1,3 +1,8 @@
+from typing import TYPE_CHECKING, Any
+if TYPE_CHECKING:
+    import matplotlib.axes
+    import matplotlib.colorbar
+    import matplotlib.colors
 from collections import defaultdict
 from random import sample, shuffle
 
@@ -16,19 +21,19 @@ def plot_adjacency(adj):
     plt.title("Adjacency matrix")
 
 
-def plot_estimates(gamma_estimates):
+def plot_estimates(gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]]) -> 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 +44,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: np.float64 = (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)
@@ -71,7 +76,7 @@ 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: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]], xlim: tuple[int, int], ylim: tuple[float, float], 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 +86,30 @@ 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.ListedColormap" = plt.get_cmap("Set1")
+    available_colors: set[tuple[np.float64, ...]] = {cmap(cnorm(i)) for i in range(len(domains))}
 
     if len(available_colors) == len(domains):
-        colors = list(available_colors)
+        colors: list[tuple[np.float64, ...]] = 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[list[np.ndarray[Any, np.dtype[np.float64]]]] = [polyverts for polyverts, membership in domains]
 
-        vertex_to_neighbor_domain_indices = defaultdict(set)
+        vertex_to_neighbor_domain_indices: defaultdict[tuple[np.float64, np.float64], 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 +117,19 @@ 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[tuple[np.float64, ...]|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: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], xlim: tuple[int, int], ylim: tuple[float, float], 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 +142,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: np.float64 = np.mean([x[0] for x in polyverts])
+        centroid_y: np.float64 = 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: int = 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 +177,23 @@ 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: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], xlim: tuple[int, int], ylim: tuple[float, float], 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.ListedColormap" = 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 +202,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,38 +214,38 @@ 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: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], ground_truth: list[int], xlim: tuple[int, int], ylim: tuple[float, float], 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.LinearSegmentedColormap" = 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: np.ndarray[Any, np.dtype[np.int64]], layer_vec: np.ndarray[Any, np.dtype[np.int64]]) -> "matplotlib.axes.Axes":
     """Plot a visualization of multiplex community membership
 
     :param membership: partition membership vector
@@ -248,14 +253,14 @@ def plot_multiplex_community(membership, layer_vec):
     :param layer_vec: list of each vertex's layer membership
     :type layer_vec: list[int]
     """
-    T = max(layer_vec) + 1
-    N = len(membership) // T
-    communities_per_layer = np.zeros((N, T))
+    T: np.int64 = max(layer_vec) + 1
+    N: np.int64 = len(membership) // T
+    communities_per_layer: np.ndarray[Any, np.dtype[np.float64]] = 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

From 172e5d08ca1f36a38ab9d623eb3920e8fa21fba5 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Thu, 15 Oct 2026 10:52:42 -0400
Subject: [PATCH 03/11] Remove legacy Louvain and Progress implementations

---
 utilities/louvain_utilities.py | 176 +++------------------------------
 utilities/progress.py          |  31 ------
 2 files changed, 11 insertions(+), 196 deletions(-)
 delete mode 100644 utilities/progress.py

diff --git a/utilities/louvain_utilities.py b/utilities/louvain_utilities.py
index 321fe53..31b40ce 100644
--- a/utilities/louvain_utilities.py
+++ b/utilities/louvain_utilities.py
@@ -2,188 +2,34 @@
 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``.
 """
 from typing import TYPE_CHECKING, Callable
 if TYPE_CHECKING:
     import igraph
     import leidenalg
 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: str) -> "Callable[..., set[tuple[int, ...]]]|Callable[[igraph.Graph, list[float]], set[tuple[int, ...]]]|Callable[[igraph.Graph, list[int]|tuple[int, ...]], leidenalg.RBConfigurationVertexPartition]|Callable[[igraph.Graph], leidenalg.RBConfigurationVertexPartition]|None":
+# TODO: simplify this type, it's not meaningful to do this for such a shim
+def __getattr__(
+      name: str,
+) -> (
+    "Callable[..., set[tuple[int, ...]]]|"
+    "Callable[[igraph.Graph, list[float]], set[tuple[int, ...]]]|"
+    "Callable[[igraph.Graph, list[int]|tuple[int, ...]], leidenalg.RBConfigurationVertexPartition]|"
+    "Callable[[igraph.Graph], leidenalg.RBConfigurationVertexPartition]|None"
+):
     if "louvain" in name:
         leiden_name: str = 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
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()

From afc7b620ddfe7ab809d6cdd4d94f7e00caaac570 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Thu, 15 Oct 2026 10:52:45 -0400
Subject: [PATCH 04/11] Manually clean up type hints in utilities scripts

[X] __init__.py
[X] champ_utilities.py
[X] leiden_utilities.py
[X] louvain_utilities.py
[X] parameter_estimation.py
[X] parameter_estimation_utilities.py
[X] partition_utilities.py
[X] plotting.py
---
 utilities/champ_utilities.py                |  86 ++++++----
 utilities/leiden_utilities.py               | 120 ++++++++++----
 utilities/louvain_utilities.py              |  20 +--
 utilities/parameter_estimation.py           | 120 ++++++++------
 utilities/parameter_estimation_utilities.py | 165 +++++++++++++++-----
 utilities/partition_utilities.py            |  39 +++--
 utilities/plotting.py                       |  62 ++++++--
 7 files changed, 413 insertions(+), 199 deletions(-)

diff --git a/utilities/champ_utilities.py b/utilities/champ_utilities.py
index 79a90c8..4f723f2 100644
--- a/utilities/champ_utilities.py
+++ b/utilities/champ_utilities.py
@@ -5,6 +5,7 @@
     import scipy.optimize
 import warnings
 from collections import defaultdict
+from collections.abc import Sequence
 from math import floor
 from multiprocessing import Pool, cpu_count
 
@@ -17,8 +18,13 @@
 from .partition_utilities import all_degrees, in_degrees, out_degrees, membership_to_communities, \
     membership_to_layered_communities
 
+_FloatArray = np.ndarray[Any, np.dtype[np.float64]]
+_IntArray = np.ndarray[Any, np.dtype[np.int64]]
 
-def get_interior_point(halfspaces: np.ndarray[Any, np.dtype[np.float64]], initial_num_sampled: int=50, full_retry_limit: int=10) -> np.ndarray[Any, np.dtype[np.float64]]:
+
+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)
 
@@ -51,10 +57,12 @@ def get_interior_point(halfspaces: np.ndarray[Any, np.dtype[np.float64]], initia
     while True:  # retry until success
         # randomly sample some of the halfspaces
         sample_len: int = min(initial_num_sampled, len(interior_hs))  # len(interior_hs)
-        sampled_hs: np.ndarray[Any, np.dtype[np.float64]] = np.vstack((interior_hs[choice(interior_hs.shape[0], sample_len, replace=False)], boundaries))
+        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.ndarray[Any, np.dtype[np.float64]] = np.reshape(np.linalg.norm(sampled_hs[:, :-1], axis=1), (sampled_hs.shape[0], 1))
+        norm_vector: _FloatArray = np.reshape(
+            np.linalg.norm(sampled_hs[:, :-1], axis=1), (sampled_hs.shape[0], 1))
         c: np.ndarray[Any, np.dtype[np.float64]] = np.zeros((sampled_hs.shape[1],))
         c[-1] = -1
         A: np.ndarray[Any, np.dtype[np.float64]] = np.hstack((sampled_hs[:, :-1], norm_vector))
@@ -70,7 +78,7 @@ def get_interior_point(halfspaces: np.ndarray[Any, np.dtype[np.float64]], initia
         #     3 : Problem appears to be unbounded.
         #     4 : Numerical difficulties encountered.
         if res.success or res.status in {0, 1, 4}:
-            intpt: np.ndarray[Any, np.dtype[np.float64]] = res.x[:-1]  # res.x contains [interior_point, distance to enclosing polyhedron]
+            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():
@@ -91,7 +99,10 @@ def get_interior_point(halfspaces: np.ndarray[Any, np.dtype[np.float64]], initia
     return intpt
 
 
-def CHAMP_2D[T1: (float, int)](G: "igraph.Graph", all_parts: list[tuple[int, ...]]|set[tuple[int, ...]], gamma_0: T1, gamma_f: T1, single_threaded: bool=False) -> list[tuple[np.float64, np.float64, tuple[int, ...]]]:
+def CHAMP_2D[T1: (float, int)](
+    G: "igraph.Graph", all_parts: Sequence[tuple[int, ...]] | set[tuple[int, ...]], gamma_0: T1, gamma_f: T1,
+    single_threaded: bool = False
+) -> list[tuple[np.float64, np.float64, tuple[int, ...]]]:
     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.
@@ -120,15 +131,16 @@ def CHAMP_2D[T1: (float, int)](G: "igraph.Graph", all_parts: list[tuple[int, ...
     all_parts = list(all_parts)
     num_partitions: int = len(all_parts)
 
-    partition_coefficients: tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]]] = partition_coefficients_2D(G, all_parts, single_threaded=single_threaded)
+    partition_coefficients: tuple[_IntArray, _FloatArray] = 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: 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: np.ndarray[Any, np.dtype[np.float64]] = np.vstack((halfspaces_from_coefficients_2D(*partition_coefficients),
-                            np.array([[0, 1, -top], [0, -1, bottom],
-                                      [1, 0, -gamma_f], [-1, 0, gamma_0]])))
+    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)
@@ -155,7 +167,10 @@ def CHAMP_2D[T1: (float, int)](G: "igraph.Graph", all_parts: list[tuple[int, ...
     return sorted(ranges, key=lambda x: x[0])
 
 
-def CHAMP_3D[T1: (float, int), T2: (float, int)](G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int], all_parts: list[tuple[int, ...]]|set[tuple[int, ...]], gamma_0: T1, gamma_f: T1, omega_0: T2, omega_f: T2) -> list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]]:
+def CHAMP_3D[T1: (float, int), T2: (float, int)](
+    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
+    all_parts: Sequence[tuple[int, ...]] | set[tuple[int, ...]], gamma_0: T1, gamma_f: T1, omega_0: T2, omega_f: T2
+) -> list[tuple[list[_FloatArray], tuple[int, ...]]]:
     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``.
 
@@ -168,7 +183,7 @@ def CHAMP_3D[T1: (float, int), T2: (float, int)](G_intralayer: "igraph.Graph", G
     :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
@@ -190,16 +205,17 @@ def CHAMP_3D[T1: (float, int), T2: (float, int)](G_intralayer: "igraph.Graph", G
     all_parts = list(all_parts)
     num_partitions: int = len(all_parts)
 
-    partition_coefficients: tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.int64]]] = partition_coefficients_3D(G_intralayer, G_interlayer, layer_vec, all_parts)
+    partition_coefficients: tuple[_IntArray, _FloatArray, _IntArray] = 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: 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: np.ndarray[Any, np.dtype[np.float64]] = 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]]))))
+    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)
@@ -228,7 +244,8 @@ def CHAMP_3D[T1: (float, int), T2: (float, int)](G_intralayer: "igraph.Graph", G
     return domains
 
 
-def partition_coefficients_2D_serial(G, partitions):
+def partition_coefficients_2D_serial(G: "igraph.Graph", partitions: Sequence[tuple[int, ...]]
+                                     ) -> tuple[_IntArray, _FloatArray]:
     """Computes A_hat and P_hat for partitions of :G:
 
     TODO: support edge weights"""
@@ -260,7 +277,9 @@ def partition_coefficients_2D_serial(G, partitions):
     return A_hats, P_hats
 
 
-def partition_coefficients_2D(G: "igraph.Graph", partitions: list[tuple[int, ...]]|set[tuple[int, ...]], single_threaded: bool=False) -> tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]]]:
+def partition_coefficients_2D(
+    G: "igraph.Graph", partitions: Sequence[tuple[int, ...]], single_threaded: bool = False
+) -> tuple[_IntArray, _FloatArray]:
     """Computes partitions coefficients in parallel by calling partition_coefficients_2D_serial"""
     partitions = list(partitions)
 
@@ -273,8 +292,9 @@ def partition_coefficients_2D(G: "igraph.Graph", partitions: list[tuple[int, ...
         ]
 
         pool: "multiprocessing.pool.Pool" = Pool(processes=cpu_count())
-        results: list[tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]]]] = pool.starmap(partition_coefficients_2D_serial,
-                               [(G, partition_chunk) for partition_chunk in partition_chunks])
+        results: list[tuple[_IntArray, _FloatArray]] = \
+            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])
@@ -285,27 +305,32 @@ def partition_coefficients_2D(G: "igraph.Graph", partitions: list[tuple[int, ...
     return A_hats, P_hats
 
 
-def halfspaces_from_coefficients_2D(A_hats: np.ndarray[Any, np.dtype[np.int64]], P_hats: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
+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: np.ndarray[Any, np.dtype[np.int64]], P_hats: np.ndarray[Any, np.dtype[np.float64]], C_hats: np.ndarray[Any, np.dtype[np.int64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
+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: _IntArray,
+                                     partitions: list[_IntArray]) -> tuple[_IntArray, _FloatArray, _IntArray]:
     """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:
 
@@ -354,7 +379,10 @@ 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: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int], partitions: list[tuple[int, ...]]) -> tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.int64]]]:
+def partition_coefficients_3D(
+    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
+    partitions: list[tuple[int, ...]]
+) -> tuple[_IntArray, _FloatArray, _IntArray]:
     """Computes partitions coefficients in parallel by calling partition_coefficients_3D_serial"""
     partitions = list(partitions)
     partition_chunks: list[list[tuple[int, ...]]] = [
@@ -363,9 +391,9 @@ def partition_coefficients_3D(G_intralayer: "igraph.Graph", G_interlayer: "igrap
     ]
 
     pool: "multiprocessing.pool.Pool" = Pool(processes=cpu_count())
-    results: list[tuple[np.ndarray[Any, np.dtype[np.float64]], ...]|tuple[np.ndarray[Any, np.dtype[np.int64]], np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.int64]]]] = pool.starmap(partition_coefficients_3D_serial,
-                           [(G_intralayer, G_interlayer, layer_vec, partition_chunk)
-                            for partition_chunk in partition_chunks])
+    results: list[tuple[_FloatArray, ...] | tuple[_IntArray, _FloatArray, _IntArray]] = 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.ndarray[Any, np.dtype[np.int64]] = np.array([v for A_hats, P_hats, C_hats in results for v in A_hats])
diff --git a/utilities/leiden_utilities.py b/utilities/leiden_utilities.py
index 5502c7f..6d8e1c2 100644
--- a/utilities/leiden_utilities.py
+++ b/utilities/leiden_utilities.py
@@ -1,4 +1,5 @@
-from typing import Any
+from collections.abc import Iterable
+from typing import Any, cast, overload, Literal
 import functools
 from multiprocessing import Pool, cpu_count
 
@@ -21,11 +22,15 @@ def sorted_tuple(t: tuple[int, ...]) -> tuple[int, ...]:
     :rtype: tuple[int]
     """
 
-    sort_map: dict[np.int64, int] = {x[0]: i for i, x in enumerate(sorted(zip(*np.unique(t, return_index=True)), key=lambda x: x[1]))}
+    sort_map: dict[np.int64 | 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: ig.Graph, gamma: float, return_partition: bool=False) -> leidenalg.RBConfigurationVertexPartition|tuple[int, ...]:
+def singlelayer_leiden(
+    G: ig.Graph, gamma: float, return_partition: bool = False,
+) -> leidenalg.RBConfigurationVertexPartition | tuple[int, ...]:
     r"""Run the Leiden modularity maximization algorithm at a single :math:`\gamma` value.
 
     :param G: graph of interest
@@ -40,8 +45,9 @@ def singlelayer_leiden(G: ig.Graph, gamma: float, return_partition: bool=False)
     if 'weight' not in G.es:
         G.es['weight'] = [1.0] * G.ecount()
 
-    partition: leidenalg.RBConfigurationVertexPartition = leidenalg.find_partition(G, leidenalg.RBConfigurationVertexPartition, weights='weight',
-                                         resolution_parameter=gamma)
+    partition: leidenalg.RBConfigurationVertexPartition = leidenalg.find_partition(
+        G, leidenalg.RBConfigurationVertexPartition, weights='weight', resolution_parameter=gamma
+    )
 
     if return_partition:
         return partition
@@ -49,16 +55,18 @@ def singlelayer_leiden(G: ig.Graph, gamma: float, return_partition: bool=False)
         return tuple(partition.membership)
 
 
-def _wrapped_singlelayer_leiden(args):
+def _wrapped_singlelayer_leiden(args: tuple[ig.Graph, float]) -> tuple[int, ...]:
     """Wrapped singlelayer_leiden() for use in multiprocessing.Pool.imap_unordered."""
-    return singlelayer_leiden(*args)
+    return cast(tuple[int, ...], singlelayer_leiden(*args))
 
 
 def leiden_part(G: ig.Graph) -> leidenalg.RBConfigurationVertexPartition:
     return leidenalg.RBConfigurationVertexPartition(G)
 
 
-def leiden_part_with_membership(G: ig.Graph, membership: list[int]|tuple[int, ...]) -> leidenalg.RBConfigurationVertexPartition:
+def leiden_part_with_membership(
+    G: ig.Graph, membership: list[int] | tuple[int, ...]
+) -> leidenalg.RBConfigurationVertexPartition:
     if isinstance(membership, np.ndarray):
         membership = membership.tolist()
     part: leidenalg.RBConfigurationVertexPartition = leiden_part(G)
@@ -86,7 +94,9 @@ def split_intralayer_leiden_graph(G_intralayer: ig.Graph, layer_membership: list
 
 
 @functools.lru_cache(maxsize=1)
-def _split_leiden_graph_layers_cached(n: int, G_es: tuple[ig.Edge, ...], is_directed: bool, layer_membership: tuple[int, ...]) -> list[ig.Graph]:
+def _split_leiden_graph_layers_cached(
+    n: int, G_es: tuple[ig.Edge, ...], is_directed: bool, layer_membership: tuple[int, ...]
+) -> list[ig.Graph]:
     T: int = max(layer_membership) + 1
 
     edges_by_layer: list[list[tuple[int, int]]] = [[] for _ in range(T)]
@@ -104,7 +114,27 @@ def _split_leiden_graph_layers_cached(n: int, G_es: tuple[ig.Edge, ...], is_dire
     return G_split_layers_list
 
 
-def multilayer_leiden(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gamma: float, omega: float, optimiser: leidenalg.Optimiser|None=None, return_partition: bool=False) -> list[leidenalg.RBConfigurationVertexPartition]|tuple[int, ...]:
+# 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: list[int], gamma: float,
+    omega: float, optimiser: leidenalg.Optimiser | None = None, *, return_partition: Literal[True],
+) -> list[leidenalg.RBConfigurationVertexPartition]:
+    ...
+
+
+@overload
+def multilayer_leiden(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gamma: float,
+    omega: float, optimiser: leidenalg.Optimiser | None = None, *, return_partition: Literal[False] = False,
+) -> tuple[int, ...]:
+    ...
+
+
+def multilayer_leiden(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gamma: float,
+    omega: float, optimiser: leidenalg.Optimiser | None = None, return_partition: bool = False,
+) -> list[leidenalg.RBConfigurationVertexPartition] | tuple[int, ...]:
     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
@@ -138,9 +168,13 @@ def multilayer_leiden(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec:
     if optimiser is None:
         optimiser = leidenalg.Optimiser()
 
-    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = [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')
+    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = (
+        [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])
 
@@ -150,12 +184,15 @@ def multilayer_leiden(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec:
         return tuple(intralayer_parts[0].membership)
 
 
-def _wrapped_multilayer_leiden(args):
+def _wrapped_multilayer_leiden(args: tuple[ig.Graph, ig.Graph, list[int], float, float]
+                               ) -> tuple[int, ...]:
     """Wrapped multilayer_leiden() for use in multiprocessing.Pool.imap_unordered."""
     return multilayer_leiden(*args)
 
 
-def multilayer_leiden_part(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_membership: list[int]) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
+def multilayer_leiden_part(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_membership: list[int]
+) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
     if 'weight' not in G_intralayer.es:
         G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount()
 
@@ -163,13 +200,19 @@ def multilayer_leiden_part(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer
         G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount()
 
     G_split_layers: list[ig.Graph] = split_intralayer_leiden_graph(G_intralayer, layer_membership)
-    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = [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')
+    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = (
+        [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: ig.Graph, G_interlayer: ig.Graph, layer_membership: list[int], community_membership: tuple[int, ...]) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
+def multilayer_leiden_part_with_membership(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_membership: list[int],
+    community_membership: tuple[int, ...]
+) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
     if isinstance(community_membership, np.ndarray):
         community_membership = community_membership.tolist()
     intralayer_parts, interlayer_part = multilayer_leiden_part(G_intralayer, G_interlayer, layer_membership)
@@ -183,7 +226,9 @@ def repeated_leiden_from_gammas(G: ig.Graph, gammas: list[float]) -> set[tuple[i
     return {sorted_tuple(singlelayer_leiden(G, gamma)) for gamma in gammas}
 
 
-def repeated_parallel_leiden_from_gammas(G: ig.Graph, gammas: np.ndarray[Any, np.dtype[np.float64]], show_progress: bool=True) -> set[tuple[int, ...]]:
+def repeated_parallel_leiden_from_gammas(
+    G: ig.Graph, gammas: np.ndarray[Any, np.dtype[np.float64]], show_progress: bool = True
+) -> set[tuple[int, ...]]:
     r"""Runs the Leiden modularity maximization algorithm at each provided :math:`\gamma` value, using all CPU cores.
 
     :param G: graph of interest
@@ -198,8 +243,10 @@ def repeated_parallel_leiden_from_gammas(G: ig.Graph, gammas: np.ndarray[Any, np
     total: set[tuple[int, ...]] = set()
     pool_chunk_size: int = max(1, len(gammas) // (cpu_count() * 100))
     with Pool(processes=cpu_count()) as pool:
-        pool_iterator: tqdm = pool.imap_unordered(_wrapped_singlelayer_leiden, [(G, g) for g in gammas],
-                                            chunksize=pool_chunk_size)
+        pool_iterator: Iterable[tuple[int, ...]] = (
+            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))
 
@@ -209,13 +256,19 @@ def repeated_parallel_leiden_from_gammas(G: ig.Graph, gammas: np.ndarray[Any, np
     return total
 
 
-def repeated_leiden_from_gammas_omegas(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gammas: list[float], omegas: list[float]) -> set[tuple[int, ...]]:
-    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: list[int],
+    gammas: list[float], omegas: list[float]
+) -> set[tuple[int, ...]]:
+    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: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gammas: np.ndarray[Any, np.dtype[np.float64]], omegas: np.ndarray[Any, np.dtype[np.float64]],
-                                                show_progress: bool=True) -> set[tuple[int, ...]]:
+def repeated_parallel_leiden_from_gammas_omegas(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int],
+    gammas: np.ndarray[Any, np.dtype[np.float64]], omegas: np.ndarray[Any, np.dtype[np.float64]],
+    show_progress: bool = True
+) -> set[tuple[int, ...]]:
     """
     Runs leidenalg at each gamma and omega in ``gammas`` and ``omegas``, using all CPU cores available.
 
@@ -237,15 +290,20 @@ def repeated_parallel_leiden_from_gammas_omegas(G_intralayer: ig.Graph, G_interl
     :return: a set of all unique partitions encountered
     :rtype: set of tuple[int]
     """
-    resolution_parameter_points: list[tuple[np.float64, np.float64]] = [(gamma, omega) for gamma in gammas for omega in omegas]
+    resolution_parameter_points: list[tuple[np.float64, np.float64]] = (
+        [(gamma, omega) for gamma in gammas for omega in omegas]
+    )
 
     total: set[tuple[int, ...]] = set()
     pool_chunk_size: int = max(1, len(resolution_parameter_points) // (cpu_count() * 100))
     with Pool(processes=cpu_count()) as pool:
-        pool_iterator: tqdm = 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[tuple[int, ...]] = (
+            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 31b40ce..5e5fdb2 100644
--- a/utilities/louvain_utilities.py
+++ b/utilities/louvain_utilities.py
@@ -5,10 +5,7 @@
 module ``modularitypruning.louvain_utilities`` now shims Louvain function names to their corresponding Leiden versions
 in ``modularitypruning.leiden_utilities``.
 """
-from typing import TYPE_CHECKING, Callable
-if TYPE_CHECKING:
-    import igraph
-    import leidenalg
+from typing import Any
 import warnings
 
 from . import leiden_utilities
@@ -18,18 +15,11 @@
               "Please use leiden_utilities in the future when possible.", DeprecationWarning)
 
 
-# TODO: simplify this type, it's not meaningful to do this for such a shim
-def __getattr__(
-      name: str,
-) -> (
-    "Callable[..., set[tuple[int, ...]]]|"
-    "Callable[[igraph.Graph, list[float]], set[tuple[int, ...]]]|"
-    "Callable[[igraph.Graph, list[int]|tuple[int, ...]], leidenalg.RBConfigurationVertexPartition]|"
-    "Callable[[igraph.Graph], leidenalg.RBConfigurationVertexPartition]|None"
-):
+def __getattr__(name: str) -> Any:  # there is no meaningful type for this style of universal shim
+    leiden_name: str = name
     if "louvain" in name:
-        leiden_name: str = name.replace("louvain", "leiden")
+        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)
+    return getattr(leiden_utilities, leiden_name)
diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py
index be91afc..55ab23b 100644
--- a/utilities/parameter_estimation.py
+++ b/utilities/parameter_estimation.py
@@ -3,15 +3,18 @@
     import igraph
     import numpy
 import leidenalg
+import numpy as np
 
-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
 
 
-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[numpy.float64, leidenalg.RBConfigurationVertexPartition]":
+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[numpy.float64 | float | None, leidenalg.RBConfigurationVertexPartition]":
     """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.
@@ -44,7 +47,7 @@ def iterative_monolayer_resolution_parameter_estimation(G: "igraph.Graph", gamma
         def maximize_modularity(resolution_param: float) -> leidenalg.RBConfigurationVertexPartition:
             return singlelayer_leiden(G, resolution_param, return_partition=True)
     elif method == "2-spinglass":
-        def maximize_modularity(resolution_param):
+        def maximize_modularity(resolution_param: float) -> leidenalg.RBConfigurationVertexPartition:
             membership = G.community_spinglass(spins=2, gamma=resolution_param).membership
             return leiden_part_with_membership(G, membership)
     else:
@@ -54,9 +57,11 @@ def estimate_SBM_parameters(partition: leidenalg.RBConfigurationVertexPartition)
         return estimate_singlelayer_SBM_parameters(G, partition, m=m)
 
     def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
-        return gamma_estimate_from_parameters(omega_in, omega_out)
+        g = gamma_estimate_from_parameters(omega_in, omega_out)
+        return g if g else np.float64(0.0)
 
-    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)
@@ -68,7 +73,7 @@ def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
             raise ValueError(f"gamma={last_gamma:.3f} resulted in degenerate partition")
 
         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:
@@ -79,12 +84,16 @@ def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
                   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 None:.3f} and "
+              f"gamma={gamma:.3f}")
 
     return gamma, part
 
 
-def check_multilayer_graph_consistency(G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int], model: str, m_t: list[float], T: int, N: int|None=None, Nt: list[int]|None=None) -> None:
+def check_multilayer_graph_consistency(
+    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
+    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.
 
@@ -97,43 +106,51 @@ def check_multilayer_graph_consistency(G_intralayer: "igraph.Graph", G_interlaye
     :param N: number of nodes per layer
     :param Nt: vector of nodes per layer
     """
-
-    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"]
+    # 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: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int], gamma: float=1.0, omega: float=1.0,
-                                                         gamma_tol: float=1e-2, omega_tol: float=5e-2, omega_max: int=1000, max_iter: int=25,
-                                                         model: str='temporal', verbose: bool=False) -> "tuple[numpy.float64, float, list[leidenalg.RBConfigurationVertexPartition]]":
+def iterative_multilayer_resolution_parameter_estimation(
+    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
+    gamma: float = 1.0, omega: float = 1.0, gamma_tol: float = 1e-2, omega_tol: float = 5e-2,
+    omega_max: int = 1000, max_iter: int = 25, model: str = 'temporal', verbose: bool = False
+) -> "tuple[numpy.float64 | float, float, list[leidenalg.RBConfigurationVertexPartition]]":
     """
     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.
@@ -167,7 +184,6 @@ def iterative_multilayer_resolution_parameter_estimation(G_intralayer: "igraph.G
     :rtype: tuple[float, float, tuple[int]]
     """
 
-    global parts
     if 'weight' not in G_intralayer.es:
         G_intralayer.es['weight'] = [1.0] * G_intralayer.ecount()
 
@@ -189,10 +205,16 @@ def iterative_multilayer_resolution_parameter_estimation(G_intralayer: "igraph.G
         Nt[layer] += 1
 
     check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, model, m_t, T, N, Nt)
-    update_omega: Callable[[float, float|float|int, float, int], float|float|int] = omega_function_from_model(model, omega_max, T=T)
-    update_gamma: "Callable[[float, float], numpy.float64|None]" = gamma_estimate_from_parameters
+    update_omega: Callable[[float, float | float | int, float, int], float | float | int] = (
+        omega_function_from_model(model, omega_max, T=T))
+
+    def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
+        g = gamma_estimate_from_parameters(omega_in, omega_out)
+        return g if g else np.float64(0.0)
 
-    def maximize_modularity(intralayer_resolution: float, interlayer_resolution: float) -> list[leidenalg.RBConfigurationVertexPartition]:
+    def maximize_modularity(
+        intralayer_resolution: float, interlayer_resolution: float
+    ) -> list[leidenalg.RBConfigurationVertexPartition]:
         return multilayer_leiden(G_intralayer, G_interlayer, layer_vec, intralayer_resolution, interlayer_resolution,
                                  optimiser=optimiser, return_partition=True)
 
@@ -200,7 +222,9 @@ def estimate_SBM_parameters(partition: leidenalg.RBConfigurationVertexPartition)
         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])
@@ -230,7 +254,7 @@ def estimate_SBM_parameters(partition: leidenalg.RBConfigurationVertexPartition)
                   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 None
         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 a207815..4f6b026 100644
--- a/utilities/parameter_estimation_utilities.py
+++ b/utilities/parameter_estimation_utilities.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable
+from typing import Any, Callable, overload, Literal
 import warnings
 from math import log
 
@@ -12,7 +12,15 @@
 from .partition_utilities import num_communities
 
 
-def estimate_singlelayer_SBM_parameters(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition, m: float|None=None) -> tuple[float, float]:
+_Verts = list[np.ndarray[Any, np.dtype[np.float64]]]
+_GammaRange = tuple[np.float64, np.float64, tuple[int, ...]]
+_GammaEst = tuple[np.float64, np.float64, tuple[int, ...], None | np.float64]
+
+
+def estimate_singlelayer_SBM_parameters(
+    G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition,
+    m: float | None = None
+) -> tuple[float, float]:
     r"""Estimate singlelayer SBM parameters from a graph and a partition.
 
     See https://doi.org/10.1103/PhysRevE.94.052315 for more details.
@@ -48,8 +56,12 @@ def estimate_singlelayer_SBM_parameters(G: ig.Graph, partition: leidenalg.RBConf
     return omega_in, omega_out
 
 
-def estimate_multilayer_SBM_parameters(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], partition: leidenalg.RBConfigurationVertexPartition, model: str, N: int|None=None, T: int|None=None,
-                                       Nt: list[int]|None=None, m_t: list[float]|None=None) -> tuple[float, float, float, int]:
+def estimate_multilayer_SBM_parameters(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph,
+    layer_vec: list[int], partition: leidenalg.RBConfigurationVertexPartition,
+    model: str, N: int | None = None, T: int | None = None,
+    Nt: list[int] | None = None, m_t: list[float] | None = None
+) -> tuple[float, float, float, int]:
     """Estimates multilayer SBM parameters from a graph and a partition
 
     :param G_intralayer: intralayer graph of interest
@@ -102,7 +114,7 @@ def estimate_multilayer_SBM_parameters(G_intralayer: ig.Graph, G_interlayer: ig.
         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: list[list[float]|list[float|int]] = [[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']
@@ -119,13 +131,15 @@ def estimate_multilayer_SBM_parameters(G_intralayer: ig.Graph, G_interlayer: ig.
     else:
         theta_out = theta_out_numerator / theta_out_denominator
 
-    calculate_persistence: Callable[[list[int]], float] = persistence_function_from_model(model, G_interlayer, layer_vec=layer_vec, N=N, T=T, Nt=Nt)
+    calculate_persistence: Callable[[list[int]], float] = (
+        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: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
             coeff: float = 2 * (1 - 1 / K) / (T * (T - 1))
-            return coeff * sum((T - n) * x ** n for n in range(1, T)) + 1 / K - pers
+            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)
@@ -140,7 +154,7 @@ def f(x: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.f
     return theta_in, theta_out, p, K
 
 
-def gamma_estimate(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition|tuple[int, ...]) -> np.float64|None:
+def gamma_estimate(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition | tuple[int, ...]) -> np.float64 | None:
     """Compute the "correct" value of gamma where modularity maximization becomes equivalent to maximum likelihood
     methods on a degree-corrected, planted partition stochastic block model.
 
@@ -164,7 +178,7 @@ def gamma_estimate(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartit
     return gamma_estimate_from_parameters(omega_in, omega_out)
 
 
-def gamma_estimate_from_parameters(omega_in: float, omega_out: float) -> np.float64|None:
+def gamma_estimate_from_parameters(omega_in: float, omega_out: float) -> np.float64 | None:
     """
     Compute the "correct" value of gamma as in
     :meth:`~modularitypruning.parameter_estimation_utilities.gamma_estimate` from SBM parameters.
@@ -180,10 +194,13 @@ def gamma_estimate_from_parameters(omega_in: float, omega_out: float) -> np.floa
     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: float, theta_out: float, p: float, K: int, T: int, omega_max: int=1000) -> float:
+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
@@ -214,7 +231,10 @@ def multiplex_omega_estimate_from_parameters(theta_in: float, theta_out: float,
     return log(1 + p * K / (1 - p)) / (T * (log(theta_in) - log(theta_out)))
 
 
-def temporal_multilevel_omega_estimate_from_parameters[T1: (float, int)](theta_in: float, theta_out: T1, p: float, K: int, omega_max: int=1000) -> T1:
+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
@@ -242,10 +262,13 @@ def temporal_multilevel_omega_estimate_from_parameters[T1: (float, int)](theta_i
 
 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: ig.Graph, community: list[int], layer_vec: list[int], Nt: list[int], T: int) -> float:
+def multilevel_persistence(
+    G_interlayer: ig.Graph, community: list[int],
+    layer_vec: list[int], 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])
@@ -256,10 +279,12 @@ def multilevel_persistence(G_interlayer: ig.Graph, community: list[int], layer_v
 
 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[T1: (float, int)](model: str, omega_max: int, T: int) -> Callable[[float, T1, float, int], T1]:
+def omega_function_from_model(
+    model: str, omega_max: float, T: int
+) -> Callable[[float, float, float, int], float]:
     """Returns an appropriate function (depending on the model) for computing omega from multilayer SBM parameters
 
     Specifically, it will return versions of
@@ -278,7 +303,7 @@ def omega_function_from_model[T1: (float, int)](model: str, omega_max: int, T: i
         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[T2: (float, int)](theta_in: float, theta_out: T2, p: float, K: int) -> T2:
+        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")
@@ -286,7 +311,11 @@ def update_omega[T2: (float, int)](theta_in: float, theta_out: T2, p: float, K:
     return update_omega
 
 
-def persistence_function_from_model(model: str, G_interlayer: ig.Graph, layer_vec: list[int]|None=None, N: int|None=None, T: int|None=None, Nt: list[int]|None=None) -> Callable[[list[int]], float]:
+def persistence_function_from_model(
+    model: str, G_interlayer: ig.Graph,
+    layer_vec: list[int] | None = None, N: int | None = None,
+    T: int | None = None, Nt: list[int] | None = None
+) -> Callable[[list[int]], float]:
     """
     Returns a function to calculate persistence according to a given multilayer model
 
@@ -324,8 +353,12 @@ def calculate_persistence(community: list[int]) -> float:
     return calculate_persistence
 
 
-def gamma_omega_estimate(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], membership: tuple[int, ...], omega_max: int=1000, model: str='temporal',
-                         N: None=None, T: int|None=None, Nt: None=None, m_t: None=None) -> tuple[np.float64|None, float]:
+def gamma_omega_estimate(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph,
+    layer_vec: list[int], membership: tuple[int, ...],
+    omega_max: int = 1000, model: str = 'temporal',
+    N: None = None, T: int | None = None, Nt: None = None, m_t: None = None
+) -> tuple[np.float64 | None, float]:
     """Returns the (gamma, omega) estimate for a multilayer network and a partition
 
     :param G_intralayer: intralayer graph of interest
@@ -357,15 +390,17 @@ def gamma_omega_estimate(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_v
     partition: leidenalg.RBConfigurationVertexPartition = 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: Callable[[float, float|float|int, float, int], float|float|int] = omega_function_from_model(model, omega_max, T=T)
-    update_gamma: Callable[[float, float], np.float64|None] = gamma_estimate_from_parameters
+    update_omega: Callable[[float, float | float | int, float, int], float | float | int] = (
+        omega_function_from_model(model, omega_max, T=T)
+    )
+    update_gamma: Callable[[float, float], np.float64 | None] = gamma_estimate_from_parameters
 
-    gamma: np.float64|None = update_gamma(theta_in, theta_out)
+    gamma: np.float64 | None = 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: ig.Graph, ranges: list[tuple[np.float64, np.float64, tuple[int, ...]]]) -> list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]]:
+def ranges_to_gamma_estimates(G: ig.Graph, ranges: list[_GammaRange]) -> list[_GammaEst]:
     """
     Compute gamma estimates as in :meth:`~modularitypruning.parameter_estimation_utilities.gamma_estimate`, given
     domains of optimality from :meth:`~modularitypruning.champ_utilities.CHAMP_2D`.
@@ -383,7 +418,7 @@ def ranges_to_gamma_estimates(G: ig.Graph, ranges: list[tuple[np.float64, np.flo
             gamma_start, gamma_end, part in ranges]
 
 
-def gamma_estimates_to_stable_partitions(gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]]) -> list[tuple[int, ...]]:
+def gamma_estimates_to_stable_partitions(gamma_estimates: list[_GammaEst]) -> list[tuple[int, ...]]:
     """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`.
 
@@ -400,7 +435,11 @@ def gamma_estimates_to_stable_partitions(gamma_estimates: list[tuple[np.float64,
             if gamma_estimate is not None and gamma_start <= gamma_estimate <= gamma_end]
 
 
-def domains_to_gamma_omega_estimates(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]], model: str='temporal') -> list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]]:
+def domains_to_gamma_omega_estimates(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph,
+    layer_vec: list[int], domains: list[tuple[_Verts, tuple[int, ...]]],
+    model: str = 'temporal'
+) -> list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]]:
     """
     Compute (gamma, omega) estimates as in
     :meth:`~modularitypruning.parameter_estimation_utilities.gamma_omega_estimate`, given domains of optimality from
@@ -418,10 +457,10 @@ def domains_to_gamma_omega_estimates(G_intralayer: ig.Graph, G_interlayer: ig.Gr
     :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[float], tuple[int], float | None, float | None]
     """
 
-    domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]] = []
+    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]] = []
     for polyverts, membership in domains:
         gamma_est, omega_est = gamma_omega_estimate(G_intralayer, G_interlayer, layer_vec, membership,
                                                     model=model)
@@ -429,7 +468,27 @@ def domains_to_gamma_omega_estimates(G_intralayer: ig.Graph, G_interlayer: ig.Gr
     return domains_with_estimates
 
 
-def gamma_omega_estimates_to_stable_partitions(domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], return_membership_only: bool=False) -> list[tuple[int, ...]]:
+# 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[tuple[_Verts, tuple[int, ...], np.float64 | None, float]],
+    return_membership_only: Literal[False] = False
+) -> list[tuple[_Verts, tuple[int, ...], np.float64 | None, float | None]]:
+    ...
+
+
+@overload
+def gamma_omega_estimates_to_stable_partitions(
+    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]],
+    return_membership_only: Literal[True]
+) -> list[tuple[int, ...]]:
+    ...
+
+
+def gamma_omega_estimates_to_stable_partitions(
+    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]],
+    return_membership_only: bool = False
+) -> list[tuple[_Verts, tuple[int, ...], np.float64 | None, float | None]] | list[tuple[int, ...]]:
     """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`.
@@ -443,11 +502,15 @@ def gamma_omega_estimates_to_stable_partitions(domains_with_estimates: list[tupl
     :rtype: list of tuple[int]
     """
 
-    def left_or_right(x1: np.float64, y1: np.float64, x2: np.float64, y2: np.float64, x: np.float64, y: float) -> np.bool:
+    def left_or_right(
+        x1: np.float64, y1: np.float64, x2: np.float64,
+        y2: np.float64, x: np.float64, y: float
+    ) -> np.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: list[tuple[int, ...]] = []
+    # TODO: remove Anys
+    stable_partitions: list[Any] = []
 
     for polyverts, membership, gamma_est, omega_est in domains_with_estimates:
         if gamma_est is None or omega_est is None:
@@ -477,8 +540,12 @@ def left_or_right(x1: np.float64, y1: np.float64, x2: np.float64, y2: np.float64
     return stable_partitions
 
 
-def prune_to_stable_partitions(G: ig.Graph, parts: set[tuple[int, ...]], gamma_start: int, gamma_end: int, restrict_num_communities: None=None,
-                               single_threaded: bool=False) -> list[tuple[int, ...]]:
+def prune_to_stable_partitions(
+    G: ig.Graph, parts: set[tuple[int, ...]],
+    gamma_start: int, gamma_end: int,
+    restrict_num_communities: None = None,
+    single_threaded: bool = False
+) -> list[tuple[int, ...]]:
     """
     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.
@@ -522,16 +589,22 @@ def prune_to_stable_partitions(G: ig.Graph, parts: set[tuple[int, ...]], gamma_s
     if len(parts) == 0:
         return []
 
-    ranges: list[tuple[np.float64, np.float64, tuple[int, ...]]] = CHAMP_2D(G, parts, gamma_start, gamma_end, single_threaded=single_threaded)
-    gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]] = ranges_to_gamma_estimates(G, ranges)
+    ranges: list[_GammaRange] = CHAMP_2D(
+        G, parts, gamma_start, gamma_end, single_threaded=single_threaded
+    )
+    gamma_estimates: list[_GammaEst] = ranges_to_gamma_estimates(G, ranges)
     stable_parts: list[tuple[int, ...]] = gamma_estimates_to_stable_partitions(gamma_estimates)
 
     return stable_parts
 
 
-def prune_to_multilayer_stable_partitions(G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], model: str,
-                                          parts: set[tuple[int, ...]], gamma_start: int, gamma_end: int, omega_start: int, omega_end: int,
-                                          restrict_num_communities: None=None, single_threaded: bool=False) -> list[tuple[int, ...]]:
+def prune_to_multilayer_stable_partitions(
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph,
+    layer_vec: list[int], model: str,
+    parts: set[tuple[int, ...]], gamma_start: int, gamma_end: int,
+    omega_start: int, omega_end: int,
+    restrict_num_communities: None = None, single_threaded: bool = False
+) -> list[tuple[int, ...]]:
     """
     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``
@@ -588,7 +661,7 @@ def prune_to_multilayer_stable_partitions(G_intralayer: ig.Graph, G_interlayer:
 
     # check graph is connected (cannot just check intra/interlayer graphs since they will both be disconnected)
     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])
+                                    [(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.")
@@ -611,9 +684,13 @@ def prune_to_multilayer_stable_partitions(G_intralayer: ig.Graph, G_interlayer:
     if len(parts) == 0:
         return []
 
-    domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]] = CHAMP_3D(G_intralayer, G_interlayer, layer_vec, parts, gamma_start, gamma_end,
-                       omega_start, omega_end)
-    domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]] = domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, domains, model)
+    domains: list[tuple[_Verts, tuple[int, ...]]] = CHAMP_3D(
+        G_intralayer, G_interlayer, layer_vec, parts, gamma_start, gamma_end,
+        omega_start, omega_end
+    )
+    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]] = (
+        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 "
@@ -625,6 +702,8 @@ def prune_to_multilayer_stable_partitions(G_intralayer: ig.Graph, G_interlayer:
                               for polyverts, membership, g_est, o_est in domains_with_estimates
                               if g_est is not None]
 
-    stable_parts: list[tuple[int, ...]] = gamma_omega_estimates_to_stable_partitions(domains_with_estimates, return_membership_only=True)
+    stable_parts: list[tuple[int, ...]] = 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 1ee9efa..b9c23dd 100644
--- a/utilities/partition_utilities.py
+++ b/utilities/partition_utilities.py
@@ -1,46 +1,53 @@
-from typing import TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 if TYPE_CHECKING:
     import igraph
 from collections import defaultdict
+from collections.abc import Sequence
 
+import numpy as np
 from sklearn.metrics import adjusted_mutual_info_score, normalized_mutual_info_score
 
+# mypy doesn't consider numpy arrays to be a sequence
+# TODO: remove Anys
+IntSequence = Sequence[int] | np.ndarray[Any, Any]
 
-def ami(p1: tuple[int, ...], p2: list[int]) -> float:
-    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: "igraph.Graph") -> list[int]:
-    return G.degree()
+    return list(G.degree())
 
 
 def in_degrees(G: "igraph.Graph") -> list[int]:
-    return G.indegree()
+    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: list[int]|tuple[int, ...]) -> int:
+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 2256d23..47840a2 100644
--- a/utilities/plotting.py
+++ b/utilities/plotting.py
@@ -4,7 +4,9 @@
     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
@@ -16,12 +18,15 @@
 from .partition_utilities import num_communities, ami
 
 
-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: list[tuple[np.float64, np.float64, tuple[int, ...], None]|tuple[np.float64, np.float64, tuple[int, ...], np.float64]]) -> None:
+def plot_estimates(
+    gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]
+                          | tuple[np.float64, np.float64, tuple[int, ...], np.float64]]
+) -> None:
     """Plot partition dominance ranges with gamma estimates.
 
     :param gamma_estimates: gamma estimates as returned from
@@ -58,7 +63,8 @@ def plot_estimates(gamma_estimates: list[tuple[np.float64, np.float64, tuple[int
                             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:
@@ -76,7 +82,11 @@ def plot_estimates(gamma_estimates: list[tuple[np.float64, np.float64, tuple[int
                 #           length_includes_head=True, alpha=0.5, zorder=2, **{"overhang": 0.5})
 
 
-def plot_2d_domains(domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]], xlim: tuple[int, int], ylim: tuple[float, float], flip_axes: bool=True, use_current_axes: bool=False) -> None:
+# TODO: remove Anys
+def plot_2d_domains(
+    domains: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]], tuple[int, ...]]],
+        xlim: tuple[int, int], ylim: tuple[float, float], 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.
@@ -98,16 +108,18 @@ def plot_2d_domains(domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64
         patches.append(polygon)
 
     cnorm: "matplotlib.colors.Normalize" = matplotlib.colors.Normalize(vmin=0, vmax=len(domains))
-    cmap: "matplotlib.colors.ListedColormap" = plt.get_cmap("Set1")
-    available_colors: set[tuple[np.float64, ...]] = {cmap(cnorm(i)) for i in range(len(domains))}
+    cmap: "matplotlib.colors.Colormap" = plt.get_cmap("Set1")
+    available_colors: set[tuple[float, ...]] = {cmap(cnorm(i)) for i in range(len(domains))}
 
+    colors: list[tuple[float, ...] | None]
     if len(available_colors) == len(domains):
-        colors: list[tuple[np.float64, ...]] = list(available_colors)
+        colors = list(available_colors)
         shuffle(colors)
     else:
         # ensure that no adjacent domains have the same color
         colors = [None] * len(domains)
-        domain_polygons: list[list[np.ndarray[Any, np.dtype[np.float64]]]] = [polyverts for polyverts, membership in domains]
+        domain_polygons: list[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]]] = (
+            [polyverts for polyverts, membership in domains])
 
         vertex_to_neighbor_domain_indices: defaultdict[tuple[np.float64, np.float64], set[int]] = defaultdict(set)
         for i, polyverts in enumerate(domain_polygons):
@@ -120,7 +132,7 @@ def plot_2d_domains(domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64
             neighboring_domains: set[int] = set()
             for v in this_polyverts:
                 neighboring_domains.update(vertex_to_neighbor_domain_indices[tuple(v)])
-            neighbor_colors: set[tuple[np.float64, ...]|None] = {colors[i] for i in neighboring_domains}
+            neighbor_colors: set[tuple[float, ...] | None] = {colors[i] for i in neighboring_domains}
             colors[i] = sample(sorted(available_colors.difference(neighbor_colors)), 1)[0]
 
     p: PatchCollection = PatchCollection(patches, facecolors=colors, alpha=1.0, edgecolors='black', linewidths=1.5)
@@ -129,7 +141,12 @@ def plot_2d_domains(domains: list[tuple[list[np.ndarray[Any, np.dtype[np.float64
     plt.ylim(ylim)
 
 
-def plot_2d_domains_with_estimates(domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], xlim: tuple[int, int], ylim: tuple[float, float], plot_estimate_points: bool=True, flip_axes: bool=True) -> None:
+def plot_2d_domains_with_estimates(
+    domains_with_estimates: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]],
+                                       tuple[int, ...], np.float64, float]],
+    xlim: tuple[int, int], ylim: tuple[float, float],
+        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.
 
@@ -151,8 +168,8 @@ def plot_2d_domains_with_estimates(domains_with_estimates: list[tuple[list[np.nd
         polygon: Polygon = Polygon(polyverts, closed=True)
         patches.append(polygon)
 
-        centroid_x: np.float64 = np.mean([x[0] for x in polyverts])
-        centroid_y: np.float64 = 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: int = xlim[1] - xlim[0]
@@ -177,7 +194,12 @@ def plot_2d_domains_with_estimates(domains_with_estimates: list[tuple[list[np.nd
                     flip_axes=flip_axes, use_current_axes=True)
 
 
-def plot_2d_domains_with_num_communities(domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], xlim: tuple[int, int], ylim: tuple[float, float], flip_axes: bool=True, K_max: int|None=None, tick_step: int=2) -> None:
+def plot_2d_domains_with_num_communities(
+    domains_with_estimates: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]],
+                                       tuple[int, ...], np.float64, float]],
+    xlim: tuple[int, int], ylim: tuple[float, float],
+        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.
 
@@ -185,7 +207,7 @@ def plot_2d_domains_with_num_communities(domains_with_estimates: list[tuple[list
     assert flip_axes
     fig, ax = plt.subplots()
     patches: list[Polygon] = []
-    cm: "matplotlib.colors.ListedColormap" = matplotlib.cm.viridis
+    cm: "matplotlib.colors.Colormap" = matplotlib.cm.viridis
     Ks: list[int] = []
 
     for polyverts, membership, gamma_est, omega_est in domains_with_estimates:
@@ -214,7 +236,12 @@ def plot_2d_domains_with_num_communities(domains_with_estimates: list[tuple[list
     plt.ylim(ylim)
 
 
-def plot_2d_domains_with_ami(domains_with_estimates: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...], np.float64, float]], ground_truth: list[int], xlim: tuple[int, int], ylim: tuple[float, float], flip_axes: bool=True) -> None:
+def plot_2d_domains_with_ami(
+    domains_with_estimates: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]],
+                                       tuple[int, ...], np.float64, float]],
+        ground_truth: list[int], xlim: tuple[int, int], ylim: tuple[float, float],
+        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.
 
@@ -222,7 +249,7 @@ def plot_2d_domains_with_ami(domains_with_estimates: list[tuple[list[np.ndarray[
     assert flip_axes
     fig, ax = plt.subplots()
     patches: list[Polygon] = []
-    cm: "matplotlib.colors.LinearSegmentedColormap" = matplotlib.cm.copper
+    cm: "matplotlib.colors.Colormap" = matplotlib.cm.copper
     amis: list[float] = []
 
     for polyverts, membership, gamma_est, omega_est in domains_with_estimates:
@@ -245,7 +272,8 @@ def plot_2d_domains_with_ami(domains_with_estimates: list[tuple[list[np.ndarray[
     plt.ylim(ylim)
 
 
-def plot_multiplex_community(membership: np.ndarray[Any, np.dtype[np.int64]], layer_vec: np.ndarray[Any, np.dtype[np.int64]]) -> "matplotlib.axes.Axes":
+def plot_multiplex_community(membership: np.ndarray[Any, np.dtype[np.int64]],
+                             layer_vec: np.ndarray[Any, np.dtype[np.int64]]) -> "matplotlib.axes.Axes":
     """Plot a visualization of multiplex community membership
 
     :param membership: partition membership vector

From c3fb28c9fe42c34859c3aaaaced5416d4d093be1 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Wed, 14 Oct 2026 21:39:13 -0400
Subject: [PATCH 05/11] Greatly simplify typing by consolidating into one file

Also remove all use of Any in type hints
---
 mypy.ini                                    |   1 +
 utilities/champ_utilities.py                | 116 ++++++++--------
 utilities/leiden_utilities.py               | 108 ++++++++-------
 utilities/louvain_utilities.py              |   6 +-
 utilities/parameter_estimation.py           |  31 ++---
 utilities/parameter_estimation_utilities.py | 138 ++++++++++----------
 utilities/partition_utilities.py            |   7 +-
 utilities/plotting.py                       |  52 ++++----
 utilities/types.py                          |  62 +++++++++
 9 files changed, 300 insertions(+), 221 deletions(-)
 create mode 100644 utilities/types.py

diff --git a/mypy.ini b/mypy.ini
index 860ec61..6a9ba22 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -1,5 +1,6 @@
 [mypy]
 strict = True
+disallow_any_explicit = True
 files = utilities/*.py
 
 [mypy-sklearn.*,igraph.*,leidenalg.*]
diff --git a/utilities/champ_utilities.py b/utilities/champ_utilities.py
index 4f723f2..bc0df8a 100644
--- a/utilities/champ_utilities.py
+++ b/utilities/champ_utilities.py
@@ -1,11 +1,10 @@
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     import igraph
     import multiprocessing.pool
     import scipy.optimize
 import warnings
 from collections import defaultdict
-from collections.abc import Sequence
 from math import floor
 from multiprocessing import Pool, cpu_count
 
@@ -17,14 +16,13 @@
 
 from .partition_utilities import all_degrees, in_degrees, out_degrees, membership_to_communities, \
     membership_to_layered_communities
-
-_FloatArray = np.ndarray[Any, np.dtype[np.float64]]
-_IntArray = np.ndarray[Any, np.dtype[np.int64]]
+from .types import (Domain, FloatArray, GammaRange, IntArray, LayerVector, Membership,
+                    Partitions, PartitionCoefficients2D, PartitionCoefficients3D)
 
 
 def get_interior_point(
-    halfspaces: _FloatArray, initial_num_sampled: int = 50, full_retry_limit: int = 10
-) -> _FloatArray:
+    halfspaces: FloatArray, initial_num_sampled: int = 50, full_retry_limit: int = 10
+) -> FloatArray:
     """
     Find interior point of halfspaces (needed to perform halfspace intersection)
 
@@ -57,18 +55,18 @@ def get_interior_point(
     while True:  # retry until success
         # randomly sample some of the halfspaces
         sample_len: int = min(initial_num_sampled, len(interior_hs))  # len(interior_hs)
-        sampled_hs: _FloatArray = np.vstack((
+        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: _FloatArray = np.reshape(
+        norm_vector: FloatArray = np.reshape(
             np.linalg.norm(sampled_hs[:, :-1], axis=1), (sampled_hs.shape[0], 1))
-        c: np.ndarray[Any, np.dtype[np.float64]] = np.zeros((sampled_hs.shape[1],))
+        c: FloatArray = np.zeros((sampled_hs.shape[1],))
         c[-1] = -1
-        A: np.ndarray[Any, np.dtype[np.float64]] = np.hstack((sampled_hs[:, :-1], norm_vector))
-        b: np.ndarray[Any, np.dtype[np.float64]] = -sampled_hs[:, -1:]
+        A: FloatArray = np.hstack((sampled_hs[:, :-1], norm_vector))
+        b: FloatArray = -sampled_hs[:, -1:]
 
-        res: "scipy.optimize.OptimizeResult" = linprog(c, A_ub=A, b_ub=b, bounds=(-np.inf, np.inf), method='highs-ipm')
+        res = linprog(c, A_ub=A, b_ub=b, bounds=(-np.inf, np.inf), method='highs-ipm')
 
         # check solution if optimization succeeded, even with difficulties
         #     res.status codes
@@ -78,7 +76,10 @@ def get_interior_point(
         #     3 : Problem appears to be unbounded.
         #     4 : Numerical difficulties encountered.
         if res.success or res.status in {0, 1, 4}:
-            intpt: _FloatArray = 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():
@@ -99,10 +100,10 @@ def get_interior_point(
     return intpt
 
 
-def CHAMP_2D[T1: (float, int)](
-    G: "igraph.Graph", all_parts: Sequence[tuple[int, ...]] | set[tuple[int, ...]], gamma_0: T1, gamma_f: T1,
+def CHAMP_2D(
+    G: "igraph.Graph", all_parts: Partitions, gamma_0: float, gamma_f: float,
     single_threaded: bool = False
-) -> list[tuple[np.float64, np.float64, tuple[int, ...]]]:
+) -> 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.
@@ -131,33 +132,33 @@ def CHAMP_2D[T1: (float, int)](
     all_parts = list(all_parts)
     num_partitions: int = len(all_parts)
 
-    partition_coefficients: tuple[_IntArray, _FloatArray] = partition_coefficients_2D(
+    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: 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]])))
+    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: np.ndarray[Any, np.dtype[np.float64]] = get_interior_point(halfspaces)
+    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[int, list[np.ndarray[Any, np.dtype[np.float64]]]] = 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: list[tuple[np.float64, np.float64, tuple[int, ...]]] = []
+    ranges: list[GammaRange] = []
     for i, intersections in facets_by_halfspace.items():
         x1, x2 = intersections[0][0], intersections[1][0]
         if x1 > x2:
@@ -167,10 +168,10 @@ def CHAMP_2D[T1: (float, int)](
     return sorted(ranges, key=lambda x: x[0])
 
 
-def CHAMP_3D[T1: (float, int), T2: (float, int)](
-    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
-    all_parts: Sequence[tuple[int, ...]] | set[tuple[int, ...]], gamma_0: T1, gamma_f: T1, omega_0: T2, omega_f: T2
-) -> list[tuple[list[_FloatArray], tuple[int, ...]]]:
+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``.
 
@@ -205,47 +206,47 @@ def CHAMP_3D[T1: (float, int), T2: (float, int)](
     all_parts = list(all_parts)
     num_partitions: int = len(all_parts)
 
-    partition_coefficients: tuple[_IntArray, _FloatArray, _IntArray] = partition_coefficients_3D(
+    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: 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]]))))
+    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: np.ndarray[Any, np.dtype[np.float64]] = get_interior_point(halfspaces)
+    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[int, list[np.ndarray[Any, np.dtype[np.float64]]]] = 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: list[tuple[list[np.ndarray[Any, np.dtype[np.float64]]], tuple[int, ...]]] = []
+    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.ndarray[Any, np.dtype[np.float64]] = 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: "igraph.Graph", partitions: Sequence[tuple[int, ...]]
-                                     ) -> tuple[_IntArray, _FloatArray]:
+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"""
@@ -278,21 +279,21 @@ def partition_coefficients_2D_serial(G: "igraph.Graph", partitions: Sequence[tup
 
 
 def partition_coefficients_2D(
-    G: "igraph.Graph", partitions: Sequence[tuple[int, ...]], single_threaded: bool = False
-) -> tuple[_IntArray, _FloatArray]:
+    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: list[list[tuple[int, ...]]] = [
+        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: "multiprocessing.pool.Pool" = Pool(processes=cpu_count())
-        results: list[tuple[_IntArray, _FloatArray]] = \
+        results: list[PartitionCoefficients2D] = \
             pool.starmap(partition_coefficients_2D_serial,
                          [(G, partition_chunk) for partition_chunk in partition_chunks])
         pool.close()
@@ -306,8 +307,8 @@ def partition_coefficients_2D(
 
 
 def halfspaces_from_coefficients_2D(
-    A_hats: _IntArray, P_hats: _FloatArray
-) -> _FloatArray:
+    A_hats: IntArray, P_hats: FloatArray
+) -> FloatArray:
     """Converts partitions' coefficients to halfspace normal, offset
 
     Q >= -P_hat*gamma + A_hat
@@ -318,8 +319,8 @@ def halfspaces_from_coefficients_2D(
 
 
 def halfspaces_from_coefficients_3D(
-    A_hats: _IntArray, P_hats: _FloatArray, C_hats: _IntArray
-) -> _FloatArray:
+    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
@@ -329,8 +330,9 @@ def halfspaces_from_coefficients_3D(
     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: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: _IntArray,
-                                     partitions: list[_IntArray]) -> tuple[_IntArray, _FloatArray, _IntArray]:
+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:
 
@@ -380,25 +382,25 @@ def partition_coefficients_3D_serial(G_intralayer: "igraph.Graph", G_interlayer:
 
 
 def partition_coefficients_3D(
-    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
-    partitions: list[tuple[int, ...]]
-) -> tuple[_IntArray, _FloatArray, _IntArray]:
+    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: list[list[tuple[int, ...]]] = [
+    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: "multiprocessing.pool.Pool" = Pool(processes=cpu_count())
-    results: list[tuple[_FloatArray, ...] | tuple[_IntArray, _FloatArray, _IntArray]] = pool.starmap(
+    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.ndarray[Any, np.dtype[np.int64]] = np.array([v for A_hats, P_hats, C_hats in results for v in A_hats])
-    P_hats: np.ndarray[Any, np.dtype[np.float64]] = np.array([v for A_hats, P_hats, C_hats in results for v in P_hats])
-    C_hats: np.ndarray[Any, np.dtype[np.int64]] = 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 6d8e1c2..311144b 100644
--- a/utilities/leiden_utilities.py
+++ b/utilities/leiden_utilities.py
@@ -1,5 +1,5 @@
 from collections.abc import Iterable
-from typing import Any, cast, overload, Literal
+from typing import overload, Literal
 import functools
 from multiprocessing import Pool, cpu_count
 
@@ -8,11 +8,14 @@
 import numpy as np
 from tqdm import tqdm
 
+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: tuple[int, ...]) -> tuple[int, ...]:
+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).
 
@@ -22,15 +25,30 @@ def sorted_tuple(t: tuple[int, ...]) -> tuple[int, ...]:
     :rtype: tuple[int]
     """
 
-    sort_map: dict[np.int64 | int, int] = (
+    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)
 
 
+# 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,
-) -> leidenalg.RBConfigurationVertexPartition | tuple[int, ...]:
+) -> SingleLayerLeidenResult:
     r"""Run the Leiden modularity maximization algorithm at a single :math:`\gamma` value.
 
     :param G: graph of interest
@@ -45,7 +63,7 @@ def singlelayer_leiden(
     if 'weight' not in G.es:
         G.es['weight'] = [1.0] * G.ecount()
 
-    partition: leidenalg.RBConfigurationVertexPartition = leidenalg.find_partition(
+    partition: LeidenPartition = leidenalg.find_partition(
         G, leidenalg.RBConfigurationVertexPartition, weights='weight', resolution_parameter=gamma
     )
 
@@ -55,26 +73,26 @@ def singlelayer_leiden(
         return tuple(partition.membership)
 
 
-def _wrapped_singlelayer_leiden(args: tuple[ig.Graph, float]) -> tuple[int, ...]:
+def _wrapped_singlelayer_leiden(args: tuple[ig.Graph, float]) -> Membership:
     """Wrapped singlelayer_leiden() for use in multiprocessing.Pool.imap_unordered."""
-    return cast(tuple[int, ...], singlelayer_leiden(*args))
+    return singlelayer_leiden(*args)
 
 
-def leiden_part(G: ig.Graph) -> leidenalg.RBConfigurationVertexPartition:
+def leiden_part(G: ig.Graph) -> LeidenPartition:
     return leidenalg.RBConfigurationVertexPartition(G)
 
 
 def leiden_part_with_membership(
-    G: ig.Graph, membership: list[int] | tuple[int, ...]
-) -> leidenalg.RBConfigurationVertexPartition:
+    G: ig.Graph, membership: IntSequence
+) -> LeidenPartition:
     if isinstance(membership, np.ndarray):
         membership = membership.tolist()
-    part: leidenalg.RBConfigurationVertexPartition = leiden_part(G)
+    part: LeidenPartition = leiden_part(G)
     part.set_membership(membership)
     return part
 
 
-def split_intralayer_leiden_graph(G_intralayer: ig.Graph, layer_membership: list[int]) -> list[ig.Graph]:
+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.
 
@@ -83,7 +101,7 @@ def split_intralayer_leiden_graph(G_intralayer: ig.Graph, layer_membership: list
     :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]
+    :type layer_vec: LayerVector
     :return: list of intralayer networks
     :rtype: list[igraph.Graph]
     """
@@ -95,7 +113,7 @@ def split_intralayer_leiden_graph(G_intralayer: ig.Graph, layer_membership: list
 
 @functools.lru_cache(maxsize=1)
 def _split_leiden_graph_layers_cached(
-    n: int, G_es: tuple[ig.Edge, ...], is_directed: bool, layer_membership: tuple[int, ...]
+    n: int, G_es: tuple[ig.Edge, ...], is_directed: bool, layer_membership: LayerTuple
 ) -> list[ig.Graph]:
     T: int = max(layer_membership) + 1
 
@@ -117,24 +135,24 @@ def _split_leiden_graph_layers_cached(
 # 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: list[int], gamma: float,
+    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[leidenalg.RBConfigurationVertexPartition]:
+) -> list[LeidenPartition]:
     ...
 
 
 @overload
 def multilayer_leiden(
-    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gamma: float,
+    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,
-) -> tuple[int, ...]:
+) -> Membership:
     ...
 
 
 def multilayer_leiden(
-    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int], gamma: float,
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector, gamma: float,
     omega: float, optimiser: leidenalg.Optimiser | None = None, return_partition: bool = False,
-) -> list[leidenalg.RBConfigurationVertexPartition] | tuple[int, ...]:
+) -> 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
@@ -145,7 +163,7 @@ def multilayer_leiden(
     :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]
+    :type layer_vec: LayerVector
     :param gamma: gamma (intralayer resolution parameter) to run Leiden at
     :type gamma: float
     :param omega: omega (interlayer resolution parameter) to run Leiden at
@@ -168,7 +186,7 @@ def multilayer_leiden(
     if optimiser is None:
         optimiser = leidenalg.Optimiser()
 
-    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = (
+    intralayer_parts: list[LeidenPartition] = (
         [leidenalg.RBConfigurationVertexPartition(G_layer, weights='weight', resolution_parameter=gamma)
          for G_layer in G_split_layers]
     )
@@ -184,15 +202,15 @@ def multilayer_leiden(
         return tuple(intralayer_parts[0].membership)
 
 
-def _wrapped_multilayer_leiden(args: tuple[ig.Graph, ig.Graph, list[int], float, float]
-                               ) -> tuple[int, ...]:
+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: ig.Graph, G_interlayer: ig.Graph, layer_membership: list[int]
-) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
+    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()
 
@@ -200,7 +218,7 @@ def multilayer_leiden_part(
         G_interlayer.es['weight'] = [1.0] * G_interlayer.ecount()
 
     G_split_layers: list[ig.Graph] = split_intralayer_leiden_graph(G_intralayer, layer_membership)
-    intralayer_parts: list[leidenalg.RBConfigurationVertexPartition] = (
+    intralayer_parts: list[LeidenPartition] = (
         [leidenalg.RBConfigurationVertexPartition(G_layer, weights='weight') for G_layer in G_split_layers]
     )
     interlayer_part: leidenalg.CPMVertexPartition = leidenalg.CPMVertexPartition(
@@ -210,9 +228,9 @@ def multilayer_leiden_part(
 
 
 def multilayer_leiden_part_with_membership(
-    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_membership: list[int],
-    community_membership: tuple[int, ...]
-) -> tuple[list[leidenalg.RBConfigurationVertexPartition], leidenalg.CPMVertexPartition]:
+    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)
@@ -222,13 +240,13 @@ def multilayer_leiden_part_with_membership(
     return intralayer_parts, interlayer_part
 
 
-def repeated_leiden_from_gammas(G: ig.Graph, gammas: list[float]) -> set[tuple[int, ...]]:
+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: ig.Graph, gammas: np.ndarray[Any, np.dtype[np.float64]], show_progress: bool = True
-) -> set[tuple[int, ...]]:
+    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
@@ -240,10 +258,10 @@ def repeated_parallel_leiden_from_gammas(
     :return: a set of all unique partitions returned by the Leiden algorithm
     :rtype: set of tuple[int]
     """
-    total: set[tuple[int, ...]] = set()
+    total: set[Membership] = set()
     pool_chunk_size: int = max(1, len(gammas) // (cpu_count() * 100))
     with Pool(processes=cpu_count()) as pool:
-        pool_iterator: Iterable[tuple[int, ...]] = (
+        pool_iterator: Iterable[Membership] = (
             pool.imap_unordered(_wrapped_singlelayer_leiden, [(G, g) for g in gammas],
                                 chunksize=pool_chunk_size)
         )
@@ -257,18 +275,18 @@ def repeated_parallel_leiden_from_gammas(
 
 
 def repeated_leiden_from_gammas_omegas(
-    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int],
-    gammas: list[float], omegas: list[float]
-) -> set[tuple[int, ...]]:
+    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: ig.Graph, G_interlayer: ig.Graph, layer_vec: list[int],
-    gammas: np.ndarray[Any, np.dtype[np.float64]], omegas: np.ndarray[Any, np.dtype[np.float64]],
+    G_intralayer: ig.Graph, G_interlayer: ig.Graph, layer_vec: LayerVector,
+    gammas: FloatSequence, omegas: FloatSequence,
     show_progress: bool = True
-) -> set[tuple[int, ...]]:
+) -> set[Membership]:
     """
     Runs leidenalg at each gamma and omega in ``gammas`` and ``omegas``, using all CPU cores available.
 
@@ -280,7 +298,7 @@ def repeated_parallel_leiden_from_gammas_omegas(
     :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]
+    :type layer_vec: LayerVector
     :param gammas: list of gammas to run leidenalg at
     :type gammas: list[float]
     :param omegas: list of omegas to run leidenalg at
@@ -290,14 +308,14 @@ def repeated_parallel_leiden_from_gammas_omegas(
     :return: a set of all unique partitions encountered
     :rtype: set of tuple[int]
     """
-    resolution_parameter_points: list[tuple[np.float64, np.float64]] = (
+    resolution_parameter_points: list[tuple[float, float]] = (
         [(gamma, omega) for gamma in gammas for omega in omegas]
     )
 
-    total: set[tuple[int, ...]] = set()
+    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: Iterable[tuple[int, ...]] = (
+        pool_iterator: Iterable[Membership] = (
             pool.imap_unordered(
                 _wrapped_multilayer_leiden,
                 [(G_intralayer, G_interlayer, layer_vec, gamma, omega)
diff --git a/utilities/louvain_utilities.py b/utilities/louvain_utilities.py
index 5e5fdb2..0ff1db3 100644
--- a/utilities/louvain_utilities.py
+++ b/utilities/louvain_utilities.py
@@ -5,7 +5,6 @@
 module ``modularitypruning.louvain_utilities`` now shims Louvain function names to their corresponding Leiden versions
 in ``modularitypruning.leiden_utilities``.
 """
-from typing import Any
 import warnings
 
 from . import leiden_utilities
@@ -15,11 +14,12 @@
               "Please use leiden_utilities in the future when possible.", DeprecationWarning)
 
 
-def __getattr__(name: str) -> Any:  # there is no meaningful type for this style of universal shim
+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)
+    attr: object = getattr(leiden_utilities, leiden_name)
+    return attr
diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py
index 55ab23b..bf663dd 100644
--- a/utilities/parameter_estimation.py
+++ b/utilities/parameter_estimation.py
@@ -1,7 +1,6 @@
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     import igraph
-    import numpy
 import leidenalg
 import numpy as np
 
@@ -9,12 +8,14 @@
 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 (LayerVector, LeidenPartition, MultilayerSBMParameters,
+                    OmegaUpdateFunc, SinglelayerSBMParameters)
 
 
 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[numpy.float64 | float | None, leidenalg.RBConfigurationVertexPartition]":
+) -> 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.
@@ -44,19 +45,19 @@ def iterative_monolayer_resolution_parameter_estimation(
     m: float = sum(G.es['weight'])
 
     if method == "louvain" or method == "leiden":
-        def maximize_modularity(resolution_param: float) -> leidenalg.RBConfigurationVertexPartition:
+        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: float) -> leidenalg.RBConfigurationVertexPartition:
+        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: leidenalg.RBConfigurationVertexPartition) -> tuple[float, float]:
+    def estimate_SBM_parameters(partition: LeidenPartition) -> SinglelayerSBMParameters:
         return estimate_singlelayer_SBM_parameters(G, partition, m=m)
 
-    def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
+    def update_gamma(omega_in: float, omega_out: float) -> float:
         g = gamma_estimate_from_parameters(omega_in, omega_out)
         return g if g else np.float64(0.0)
 
@@ -91,7 +92,7 @@ def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
 
 
 def check_multilayer_graph_consistency(
-    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
+    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:
     """
@@ -147,10 +148,10 @@ def check_multilayer_graph_consistency(
 
 
 def iterative_multilayer_resolution_parameter_estimation(
-    G_intralayer: "igraph.Graph", G_interlayer: "igraph.Graph", layer_vec: list[int],
+    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: int = 1000, max_iter: int = 25, model: str = 'temporal', verbose: bool = False
-) -> "tuple[numpy.float64 | float, float, list[leidenalg.RBConfigurationVertexPartition]]":
+    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.
@@ -205,20 +206,20 @@ def iterative_multilayer_resolution_parameter_estimation(
         Nt[layer] += 1
 
     check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, model, m_t, T, N, Nt)
-    update_omega: Callable[[float, float | float | int, float, int], float | float | int] = (
+    update_omega: OmegaUpdateFunc = (
         omega_function_from_model(model, omega_max, T=T))
 
-    def update_gamma(omega_in: float, omega_out: float) -> "numpy.float64":
+    def update_gamma(omega_in: float, omega_out: float) -> float:
         g = gamma_estimate_from_parameters(omega_in, omega_out)
         return g if g else np.float64(0.0)
 
     def maximize_modularity(
         intralayer_resolution: float, interlayer_resolution: float
-    ) -> list[leidenalg.RBConfigurationVertexPartition]:
+    ) -> 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: leidenalg.RBConfigurationVertexPartition) -> tuple[float, float, float, int]:
+    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)
 
diff --git a/utilities/parameter_estimation_utilities.py b/utilities/parameter_estimation_utilities.py
index 4f6b026..fa6804e 100644
--- a/utilities/parameter_estimation_utilities.py
+++ b/utilities/parameter_estimation_utilities.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, overload, Literal
+from typing import overload, Literal
 import warnings
 from math import log
 
@@ -10,17 +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
-
-
-_Verts = list[np.ndarray[Any, np.dtype[np.float64]]]
-_GammaRange = tuple[np.float64, np.float64, tuple[int, ...]]
-_GammaEst = tuple[np.float64, np.float64, tuple[int, ...], None | np.float64]
+from .types import (Domain, DomainWithEstimates, FloatArray, GammaEstimate,
+                    GammaEstimateValue, GammaOmegaEstimate, GammaRange, GammaUpdateFunc, LayerVector,
+                    LeidenPartition, Membership, MultilayerSBMParameters, OmegaUpdateFunc,
+                    PersistenceFunc, SinglelayerSBMParameters)
 
 
 def estimate_singlelayer_SBM_parameters(
-    G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition,
+    G: ig.Graph, partition: LeidenPartition,
     m: float | None = None
-) -> tuple[float, float]:
+) -> SinglelayerSBMParameters:
     r"""Estimate singlelayer SBM parameters from a graph and a partition.
 
     See https://doi.org/10.1103/PhysRevE.94.052315 for more details.
@@ -58,10 +57,10 @@ def estimate_singlelayer_SBM_parameters(
 
 def estimate_multilayer_SBM_parameters(
     G_intralayer: ig.Graph, G_interlayer: ig.Graph,
-    layer_vec: list[int], partition: leidenalg.RBConfigurationVertexPartition,
+    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
-) -> tuple[float, float, float, int]:
+) -> MultilayerSBMParameters:
     """Estimates multilayer SBM parameters from a graph and a partition
 
     :param G_intralayer: intralayer graph of interest
@@ -69,7 +68,7 @@ def estimate_multilayer_SBM_parameters(
     :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]
+    :type layer_vec: LayerVector
     :param partition: partition of interest
     :type partition: leidenalg.RBConfigurationVertexPartition
     :param model: network layer topology (temporal, multilevel, multiplex)
@@ -131,13 +130,13 @@ def estimate_multilayer_SBM_parameters(
     else:
         theta_out = theta_out_numerator / theta_out_denominator
 
-    calculate_persistence: Callable[[list[int]], float] = (
+    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: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.float64]]:
+        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
 
@@ -154,7 +153,7 @@ def f(x: np.ndarray[Any, np.dtype[np.float64]]) -> np.ndarray[Any, np.dtype[np.f
     return theta_in, theta_out, p, K
 
 
-def gamma_estimate(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartition | tuple[int, ...]) -> np.float64 | None:
+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.
 
@@ -178,7 +177,7 @@ def gamma_estimate(G: ig.Graph, partition: leidenalg.RBConfigurationVertexPartit
     return gamma_estimate_from_parameters(omega_in, omega_out)
 
 
-def gamma_estimate_from_parameters(omega_in: float, omega_out: float) -> np.float64 | None:
+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.
@@ -267,7 +266,7 @@ def ordinal_persistence(G_interlayer: ig.Graph, community: list[int], N: int, T:
 
 def multilevel_persistence(
     G_interlayer: ig.Graph, community: list[int],
-    layer_vec: list[int], Nt: list[int], T: int
+    layer_vec: LayerVector, Nt: list[int], T: int
 ) -> float:
     pers_per_layer: list[float] = [0] * T
     for e in G_interlayer.es:
@@ -284,7 +283,7 @@ def categorical_persistence(G_interlayer: ig.Graph, community: list[int], N: int
 
 def omega_function_from_model(
     model: str, omega_max: float, T: int
-) -> Callable[[float, float, float, int], float]:
+) -> OmegaUpdateFunc:
     """Returns an appropriate function (depending on the model) for computing omega from multilayer SBM parameters
 
     Specifically, it will return versions of
@@ -313,9 +312,9 @@ def update_omega(theta_in: float, theta_out: float, p: float, K: int) -> float:
 
 def persistence_function_from_model(
     model: str, G_interlayer: ig.Graph,
-    layer_vec: list[int] | None = None, N: int | None = None,
+    layer_vec: LayerVector | None = None, N: int | None = None,
     T: int | None = None, Nt: list[int] | None = None
-) -> Callable[[list[int]], float]:
+) -> PersistenceFunc:
     """
     Returns a function to calculate persistence according to a given multilayer model
 
@@ -355,10 +354,11 @@ def calculate_persistence(community: list[int]) -> float:
 
 def gamma_omega_estimate(
     G_intralayer: ig.Graph, G_interlayer: ig.Graph,
-    layer_vec: list[int], membership: tuple[int, ...],
-    omega_max: int = 1000, model: str = 'temporal',
-    N: None = None, T: int | None = None, Nt: None = None, m_t: None = None
-) -> tuple[np.float64 | None, float]:
+    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
@@ -366,7 +366,7 @@ def gamma_omega_estimate(
     :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]
+    :type layer_vec: LayerVector
     :param membership: partition membership vector
     :type membership: tuple[int]
     :param omega_max: maximum allowed value for omega
@@ -387,20 +387,20 @@ def gamma_omega_estimate(
     if T is None:
         T = max(layer_vec) + 1  # layer  count
 
-    partition: leidenalg.RBConfigurationVertexPartition = 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: Callable[[float, float | float | int, float, int], float | float | int] = (
+    update_omega: OmegaUpdateFunc = (
         omega_function_from_model(model, omega_max, T=T)
     )
-    update_gamma: Callable[[float, float], np.float64 | None] = gamma_estimate_from_parameters
+    update_gamma: GammaUpdateFunc = gamma_estimate_from_parameters
 
-    gamma: np.float64 | None = update_gamma(theta_in, theta_out)
+    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: ig.Graph, ranges: list[_GammaRange]) -> list[_GammaEst]:
+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`.
@@ -418,7 +418,7 @@ def ranges_to_gamma_estimates(G: ig.Graph, ranges: list[_GammaRange]) -> list[_G
             gamma_start, gamma_end, part in ranges]
 
 
-def gamma_estimates_to_stable_partitions(gamma_estimates: list[_GammaEst]) -> list[tuple[int, ...]]:
+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`.
 
@@ -437,9 +437,9 @@ def gamma_estimates_to_stable_partitions(gamma_estimates: list[_GammaEst]) -> li
 
 def domains_to_gamma_omega_estimates(
     G_intralayer: ig.Graph, G_interlayer: ig.Graph,
-    layer_vec: list[int], domains: list[tuple[_Verts, tuple[int, ...]]],
+    layer_vec: LayerVector, domains: list[Domain],
     model: str = 'temporal'
-) -> list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]]:
+) -> list[DomainWithEstimates]:
     """
     Compute (gamma, omega) estimates as in
     :meth:`~modularitypruning.parameter_estimation_utilities.gamma_omega_estimate`, given domains of optimality from
@@ -450,7 +450,7 @@ def domains_to_gamma_omega_estimates(
     :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]
+    :type layer_vec: LayerVector
     :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]]
@@ -460,7 +460,7 @@ def domains_to_gamma_omega_estimates(
     :rtype: list of tuple[list[float], tuple[int], float | None, float | None]
     """
 
-    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]] = []
+    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)
@@ -471,24 +471,24 @@ def domains_to_gamma_omega_estimates(
 # 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[tuple[_Verts, tuple[int, ...], np.float64 | None, float]],
+    domains_with_estimates: list[DomainWithEstimates],
     return_membership_only: Literal[False] = False
-) -> list[tuple[_Verts, tuple[int, ...], np.float64 | None, float | None]]:
+) -> list[DomainWithEstimates]:
     ...
 
 
 @overload
 def gamma_omega_estimates_to_stable_partitions(
-    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]],
+    domains_with_estimates: list[DomainWithEstimates],
     return_membership_only: Literal[True]
-) -> list[tuple[int, ...]]:
+) -> list[Membership]:
     ...
 
 
 def gamma_omega_estimates_to_stable_partitions(
-    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]],
+    domains_with_estimates: list[DomainWithEstimates],
     return_membership_only: bool = False
-) -> list[tuple[_Verts, tuple[int, ...], np.float64 | None, float | None]] | list[tuple[int, ...]]:
+) -> 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`.
@@ -503,49 +503,51 @@ def gamma_omega_estimates_to_stable_partitions(
     """
 
     def left_or_right(
-        x1: np.float64, y1: np.float64, x2: np.float64,
-        y2: np.float64, x: np.float64, y: float
-    ) -> np.bool:
+        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
 
-    # TODO: remove Anys
-    stable_partitions: list[Any] = []
+    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.float64 = np.mean([x[0] for x in polyverts])
-        centroid_y: np.float64 = np.mean([x[1] for x in polyverts])
-        polygon_edges: list[tuple[np.ndarray[Any, np.dtype[np.float64]], np.ndarray[Any, np.dtype[np.float64]]]] = []
+        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: list[np.bool] = []
+        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: ig.Graph, parts: set[tuple[int, ...]],
-    gamma_start: int, gamma_end: int,
-    restrict_num_communities: None = None,
+    G: ig.Graph, parts: set[Membership],
+    gamma_start: float, gamma_end: float,
+    restrict_num_communities: int | None = None,
     single_threaded: bool = False
-) -> list[tuple[int, ...]]:
+) -> 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.
@@ -589,22 +591,22 @@ def prune_to_stable_partitions(
     if len(parts) == 0:
         return []
 
-    ranges: list[_GammaRange] = CHAMP_2D(
+    ranges: list[GammaRange] = CHAMP_2D(
         G, parts, gamma_start, gamma_end, single_threaded=single_threaded
     )
-    gamma_estimates: list[_GammaEst] = ranges_to_gamma_estimates(G, ranges)
-    stable_parts: list[tuple[int, ...]] = gamma_estimates_to_stable_partitions(gamma_estimates)
+    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: ig.Graph, G_interlayer: ig.Graph,
-    layer_vec: list[int], model: str,
-    parts: set[tuple[int, ...]], gamma_start: int, gamma_end: int,
-    omega_start: int, omega_end: int,
-    restrict_num_communities: None = None, single_threaded: bool = False
-) -> list[tuple[int, ...]]:
+    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``
@@ -635,7 +637,7 @@ def prune_to_multilayer_stable_partitions(
     :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]
+    :type layer_vec: LayerVector
     :param model: network layer topology (temporal, multilevel, multiplex)
     :type model: str
     :param parts: partitions to prune
@@ -684,11 +686,11 @@ def prune_to_multilayer_stable_partitions(
     if len(parts) == 0:
         return []
 
-    domains: list[tuple[_Verts, tuple[int, ...]]] = CHAMP_3D(
+    domains: list[Domain] = CHAMP_3D(
         G_intralayer, G_interlayer, layer_vec, parts, gamma_start, gamma_end,
         omega_start, omega_end
     )
-    domains_with_estimates: list[tuple[_Verts, tuple[int, ...], np.float64 | None, float]] = (
+    domains_with_estimates: list[DomainWithEstimates] = (
         domains_to_gamma_omega_estimates(G_intralayer, G_interlayer, layer_vec, domains, model)
     )
 
@@ -702,7 +704,7 @@ def prune_to_multilayer_stable_partitions(
                               for polyverts, membership, g_est, o_est in domains_with_estimates
                               if g_est is not None]
 
-    stable_parts: list[tuple[int, ...]] = gamma_omega_estimates_to_stable_partitions(
+    stable_parts: list[Membership] = gamma_omega_estimates_to_stable_partitions(
         domains_with_estimates, return_membership_only=True
     )
 
diff --git a/utilities/partition_utilities.py b/utilities/partition_utilities.py
index b9c23dd..0ae9c9b 100644
--- a/utilities/partition_utilities.py
+++ b/utilities/partition_utilities.py
@@ -1,15 +1,12 @@
-from typing import Any, TYPE_CHECKING
+from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     import igraph
 from collections import defaultdict
-from collections.abc import Sequence
 
 import numpy as np
 from sklearn.metrics import adjusted_mutual_info_score, normalized_mutual_info_score
 
-# mypy doesn't consider numpy arrays to be a sequence
-# TODO: remove Anys
-IntSequence = Sequence[int] | np.ndarray[Any, Any]
+from .types import IntSequence
 
 
 def ami(p1: IntSequence, p2: IntSequence) -> float:
diff --git a/utilities/plotting.py b/utilities/plotting.py
index 47840a2..1099690 100644
--- a/utilities/plotting.py
+++ b/utilities/plotting.py
@@ -1,10 +1,9 @@
-from typing import TYPE_CHECKING, Any
+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
 
@@ -16,6 +15,8 @@
 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: ArrayLike) -> None:
@@ -24,8 +25,7 @@ def plot_adjacency(adj: ArrayLike) -> None:
 
 
 def plot_estimates(
-    gamma_estimates: list[tuple[np.float64, np.float64, tuple[int, ...], None]
-                          | tuple[np.float64, np.float64, tuple[int, ...], np.float64]]
+    gamma_estimates: list[GammaEstimate]
 ) -> None:
     """Plot partition dominance ranges with gamma estimates.
 
@@ -49,7 +49,7 @@ def plot_estimates(
         g_0, g_f, part, g_est = gamma_estimates[i]
         k = num_communities(part)
         if g_est is not None:
-            middle: np.float64 = (g_0 + g_f) / 2
+            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:
@@ -82,10 +82,9 @@ def plot_estimates(
                 #           length_includes_head=True, alpha=0.5, zorder=2, **{"overhang": 0.5})
 
 
-# TODO: remove Anys
 def plot_2d_domains(
-    domains: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]], tuple[int, ...]]],
-        xlim: tuple[int, int], ylim: tuple[float, float], flip_axes: bool = True, use_current_axes: bool = False
+    domains: list[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.
 
@@ -109,19 +108,19 @@ def plot_2d_domains(
 
     cnorm: "matplotlib.colors.Normalize" = matplotlib.colors.Normalize(vmin=0, vmax=len(domains))
     cmap: "matplotlib.colors.Colormap" = plt.get_cmap("Set1")
-    available_colors: set[tuple[float, ...]] = {cmap(cnorm(i)) for i in range(len(domains))}
+    available_colors: set[PlotColor] = {cmap(cnorm(i)) for i in range(len(domains))}
 
-    colors: list[tuple[float, ...] | None]
+    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: list[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]]] = (
+        domain_polygons: list[FlexibleDomainVertices] = (
             [polyverts for polyverts, membership in domains])
 
-        vertex_to_neighbor_domain_indices: defaultdict[tuple[np.float64, np.float64], set[int]] = 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
@@ -132,7 +131,7 @@ def plot_2d_domains(
             neighboring_domains: set[int] = set()
             for v in this_polyverts:
                 neighboring_domains.update(vertex_to_neighbor_domain_indices[tuple(v)])
-            neighbor_colors: set[tuple[float, ...] | None] = {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 = PatchCollection(patches, facecolors=colors, alpha=1.0, edgecolors='black', linewidths=1.5)
@@ -142,9 +141,8 @@ def plot_2d_domains(
 
 
 def plot_2d_domains_with_estimates(
-    domains_with_estimates: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]],
-                                       tuple[int, ...], np.float64, float]],
-    xlim: tuple[int, int], ylim: tuple[float, float],
+    domains_with_estimates: list[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
@@ -172,7 +170,7 @@ def plot_2d_domains_with_estimates(
         centroid_y: float = np.mean([x[1] for x in polyverts])
 
         if gamma_est is not None and omega_est is not None:
-            width: int = 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:
@@ -195,9 +193,8 @@ def plot_2d_domains_with_estimates(
 
 
 def plot_2d_domains_with_num_communities(
-    domains_with_estimates: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]],
-                                       tuple[int, ...], np.float64, float]],
-    xlim: tuple[int, int], ylim: tuple[float, float],
+    domains_with_estimates: list[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
@@ -237,9 +234,8 @@ def plot_2d_domains_with_num_communities(
 
 
 def plot_2d_domains_with_ami(
-    domains_with_estimates: list[tuple[Sequence[np.ndarray[Any, np.dtype[np.float64]] | tuple[Any, Any]],
-                                       tuple[int, ...], np.float64, float]],
-        ground_truth: list[int], xlim: tuple[int, int], ylim: tuple[float, float],
+    domains_with_estimates: list[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
@@ -272,8 +268,8 @@ def plot_2d_domains_with_ami(
     plt.ylim(ylim)
 
 
-def plot_multiplex_community(membership: np.ndarray[Any, np.dtype[np.int64]],
-                             layer_vec: np.ndarray[Any, np.dtype[np.int64]]) -> "matplotlib.axes.Axes":
+def plot_multiplex_community(membership: IntArray,
+                             layer_vec: IntArray) -> "matplotlib.axes.Axes":
     """Plot a visualization of multiplex community membership
 
     :param membership: partition membership vector
@@ -281,9 +277,9 @@ def plot_multiplex_community(membership: np.ndarray[Any, np.dtype[np.int64]],
     :param layer_vec: list of each vertex's layer membership
     :type layer_vec: list[int]
     """
-    T: np.int64 = max(layer_vec) + 1
-    N: np.int64 = len(membership) // T
-    communities_per_layer: np.ndarray[Any, np.dtype[np.float64]] = 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]
 
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, ...]

From 13694db7ad4531f25c036b4808e43bbf0770aaed Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Thu, 15 Oct 2026 09:52:39 -0400
Subject: [PATCH 06/11] Add back degenerate gamma estimation

---
 utilities/parameter_estimation.py | 30 +++++++++++++-----------------
 utilities/partition_utilities.py  |  1 -
 2 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py
index bf663dd..dfe0914 100644
--- a/utilities/parameter_estimation.py
+++ b/utilities/parameter_estimation.py
@@ -2,13 +2,12 @@
 if TYPE_CHECKING:
     import igraph
 import leidenalg
-import numpy as np
 
 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 (LayerVector, LeidenPartition, MultilayerSBMParameters,
+from .types import (GammaUpdateFunc, LayerVector, LeidenPartition, MultilayerSBMParameters,
                     OmegaUpdateFunc, SinglelayerSBMParameters)
 
 
@@ -57,9 +56,7 @@ def maximize_modularity(resolution_param: float) -> LeidenPartition:
     def estimate_SBM_parameters(partition: LeidenPartition) -> SinglelayerSBMParameters:
         return estimate_singlelayer_SBM_parameters(G, partition, m=m)
 
-    def update_gamma(omega_in: float, omega_out: float) -> float:
-        g = gamma_estimate_from_parameters(omega_in, omega_out)
-        return g if g else np.float64(0.0)
+    update_gamma: GammaUpdateFunc = gamma_estimate_from_parameters
 
     part = None
     last_gamma = gamma
@@ -68,11 +65,13 @@ def update_gamma(omega_in: float, omega_out: float) -> float:
         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) if part else 0} communities with Q={part.q:.3f} and "
                   f"gamma={last_gamma:.3f}->{gamma:.3f}")
@@ -85,7 +84,7 @@ def update_gamma(omega_in: float, omega_out: float) -> float:
                   f"Final move of {abs(gamma - last_gamma):.3f} was not within tolerance {tol}")
 
     if verbose:
-        print(f"Returned {len(part) if part else 0} communities with Q={part.q if part else None:.3f} and "
+        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
@@ -206,12 +205,8 @@ def iterative_multilayer_resolution_parameter_estimation(
         Nt[layer] += 1
 
     check_multilayer_graph_consistency(G_intralayer, G_interlayer, layer_vec, model, m_t, T, N, Nt)
-    update_omega: OmegaUpdateFunc = (
-        omega_function_from_model(model, omega_max, T=T))
-
-    def update_gamma(omega_in: float, omega_out: float) -> float:
-        g = gamma_estimate_from_parameters(omega_in, omega_out)
-        return g if g else np.float64(0.0)
+    update_omega: OmegaUpdateFunc = omega_function_from_model(model, omega_max, T=T)
+    update_gamma: GammaUpdateFunc = gamma_estimate_from_parameters
 
     def maximize_modularity(
         intralayer_resolution: float, interlayer_resolution: float
@@ -234,11 +229,12 @@ def estimate_SBM_parameters(partition: LeidenPartition) -> MultilayerSBMParamete
             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:
@@ -255,7 +251,7 @@ def estimate_SBM_parameters(partition: LeidenPartition) -> MultilayerSBMParamete
                   f"was not within tolerance ({gamma_tol}, {omega_tol})")
 
     if verbose:
-        part_Q = sum(p.q for p in parts) if parts else None
+        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/partition_utilities.py b/utilities/partition_utilities.py
index 0ae9c9b..685094b 100644
--- a/utilities/partition_utilities.py
+++ b/utilities/partition_utilities.py
@@ -3,7 +3,6 @@
     import igraph
 from collections import defaultdict
 
-import numpy as np
 from sklearn.metrics import adjusted_mutual_info_score, normalized_mutual_info_score
 
 from .types import IntSequence

From 811899623425f54a1fa7c508a3837c8c83ed3bc7 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Wed, 14 Oct 2026 22:47:11 -0400
Subject: [PATCH 07/11] Improve type specificity in docstrings

---
 utilities/champ_utilities.py                |  6 +-
 utilities/leiden_utilities.py               | 30 ++++-----
 utilities/parameter_estimation.py           |  6 +-
 utilities/parameter_estimation_utilities.py | 67 +++++++++++----------
 utilities/plotting.py                       |  4 +-
 5 files changed, 57 insertions(+), 56 deletions(-)

diff --git a/utilities/champ_utilities.py b/utilities/champ_utilities.py
index bc0df8a..e0407b3 100644
--- a/utilities/champ_utilities.py
+++ b/utilities/champ_utilities.py
@@ -111,7 +111,7 @@ def CHAMP_2D(
     :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
@@ -123,7 +123,7 @@ def CHAMP_2D(
         - 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:
@@ -197,7 +197,7 @@ def CHAMP_3D(
 
         - 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:
diff --git a/utilities/leiden_utilities.py b/utilities/leiden_utilities.py
index 311144b..d060884 100644
--- a/utilities/leiden_utilities.py
+++ b/utilities/leiden_utilities.py
@@ -9,7 +9,7 @@
 from tqdm import tqdm
 
 from .types import (FloatSequence, IntSequence, LayerTuple, LayerVector, LeidenPartition,
-                     Membership, MultilayerLeidenParts, MultilayerLeidenResult, SingleLayerLeidenResult)
+                    Membership, MultilayerLeidenParts, MultilayerLeidenResult, SingleLayerLeidenResult)
 
 LOW_MEMORY_THRESHOLD: float = 1e9  # 1 GB
 
@@ -20,9 +20,9 @@ def sorted_tuple(t: Membership) -> Membership:
     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: dict[int, int] = (
@@ -58,7 +58,7 @@ def singlelayer_leiden(
     :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()
@@ -100,8 +100,8 @@ def split_intralayer_leiden_graph(G_intralayer: ig.Graph, layer_membership: Laye
 
     :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: LayerVector
+    :param layer_membership: list of each vertex's layer membership
+    :type layer_membership: list[int]
     :return: list of intralayer networks
     :rtype: list[igraph.Graph]
     """
@@ -163,17 +163,17 @@ def multilayer_leiden(
     :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: LayerVector
+    :type layer_vec: list[int]
     :param gamma: gamma (intralayer resolution parameter) to run Leiden at
     :type gamma: float
     :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()
@@ -252,11 +252,11 @@ def repeated_parallel_leiden_from_gammas(
     :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[Membership] = set()
     pool_chunk_size: int = max(1, len(gammas) // (cpu_count() * 100))
@@ -298,15 +298,15 @@ def repeated_parallel_leiden_from_gammas_omegas(
     :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: LayerVector
+    :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: list[tuple[float, float]] = (
         [(gamma, omega) for gamma in gammas for omega in omegas]
diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py
index dfe0914..ea7b568 100644
--- a/utilities/parameter_estimation.py
+++ b/utilities/parameter_estimation.py
@@ -126,8 +126,8 @@ def check_multilayer_graph_consistency(
                                "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),
+                               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",
@@ -181,7 +181,7 @@ def iterative_multilayer_resolution_parameter_estimation(
         - 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]]
     """
 
     if 'weight' not in G_intralayer.es:
diff --git a/utilities/parameter_estimation_utilities.py b/utilities/parameter_estimation_utilities.py
index fa6804e..f2b5ee2 100644
--- a/utilities/parameter_estimation_utilities.py
+++ b/utilities/parameter_estimation_utilities.py
@@ -29,7 +29,7 @@ def estimate_singlelayer_SBM_parameters(
     :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]
     """
@@ -68,19 +68,19 @@ def estimate_multilayer_SBM_parameters(
     :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: LayerVector
+    :type layer_vec: list[int]
     :param partition: partition of interest
     :type partition: leidenalg.RBConfigurationVertexPartition
     :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
     """
@@ -162,9 +162,9 @@ def gamma_estimate(G: ig.Graph, partition: LeidenPartition | Membership) -> Gamm
     :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:
@@ -187,7 +187,7 @@ def gamma_estimate_from_parameters(omega_in: float, omega_out: float) -> GammaEs
     :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:
@@ -366,23 +366,23 @@ def gamma_omega_estimate(
     :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: LayerVector
+    :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
@@ -409,9 +409,9 @@ def ranges_to_gamma_estimates(G: ig.Graph, ranges: list[GammaRange]) -> list[Gam
     :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
@@ -426,9 +426,9 @@ def gamma_estimates_to_stable_partitions(gamma_estimates: list[GammaEstimate]) -
 
     :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
@@ -450,14 +450,14 @@ def domains_to_gamma_omega_estimates(
     :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: LayerVector
+    :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 | None, float | None]
+    :rtype: list of tuple[list[numpy.ndarray], tuple[int, ...], float | None, float]
     """
 
     domains_with_estimates: list[DomainWithEstimates] = []
@@ -497,9 +497,10 @@ def gamma_omega_estimates_to_stable_partitions(
 
     :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(
@@ -557,17 +558,17 @@ def prune_to_stable_partitions(
     :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():
@@ -637,11 +638,11 @@ def prune_to_multilayer_stable_partitions(
     :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: LayerVector
+    :type layer_vec: list[int]
     :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
@@ -651,11 +652,11 @@ def prune_to_multilayer_stable_partitions(
     :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"
diff --git a/utilities/plotting.py b/utilities/plotting.py
index 1099690..5c6c9a2 100644
--- a/utilities/plotting.py
+++ b/utilities/plotting.py
@@ -273,9 +273,9 @@ def plot_multiplex_community(membership: IntArray,
     """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: int = max(layer_vec) + 1
     N: int = len(membership) // T

From 0e15803423df6c945200a13c537a664d66c00f19 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Thu, 15 Oct 2026 12:13:42 -0400
Subject: [PATCH 08/11] Add typing to pyproject and expand mypy checks to tests

---
 mypy.ini                                              |  9 +++++++--
 pyproject.toml                                        |  6 +++++-
 .../test_temporal_multilevel_parameter_estimation.py  |  1 +
 utilities/__init__.py                                 |  7 ++++++-
 utilities/plotting.py                                 | 11 ++++++-----
 utilities/py.typed                                    |  0
 6 files changed, 25 insertions(+), 9 deletions(-)
 create mode 100644 utilities/py.typed

diff --git a/mypy.ini b/mypy.ini
index 6a9ba22..22f0816 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -1,7 +1,12 @@
 [mypy]
 strict = True
 disallow_any_explicit = True
-files = utilities/*.py
+files = utilities/*.py, tests/*.py
 
 [mypy-sklearn.*,igraph.*,leidenalg.*]
-ignore_missing_imports = True
\ No newline at end of file
+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..0831cae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,8 @@ classifiers = [
     "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 +45,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/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/plotting.py b/utilities/plotting.py
index 5c6c9a2..580e70a 100644
--- a/utilities/plotting.py
+++ b/utilities/plotting.py
@@ -4,6 +4,7 @@
     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
 
@@ -25,7 +26,7 @@ def plot_adjacency(adj: ArrayLike) -> None:
 
 
 def plot_estimates(
-    gamma_estimates: list[GammaEstimate]
+    gamma_estimates: Sequence[GammaEstimate]
 ) -> None:
     """Plot partition dominance ranges with gamma estimates.
 
@@ -83,7 +84,7 @@ def plot_estimates(
 
 
 def plot_2d_domains(
-    domains: list[FlexibleDomain],
+    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.
@@ -141,7 +142,7 @@ def plot_2d_domains(
 
 
 def plot_2d_domains_with_estimates(
-    domains_with_estimates: list[FlexibleDomainWithEstimates],
+    domains_with_estimates: Sequence[FlexibleDomainWithEstimates],
     xlim: AxisLimits, ylim: AxisLimits,
         plot_estimate_points: bool = True, flip_axes: bool = True
 ) -> None:
@@ -193,7 +194,7 @@ def plot_2d_domains_with_estimates(
 
 
 def plot_2d_domains_with_num_communities(
-    domains_with_estimates: list[FlexibleDomainWithEstimates],
+    domains_with_estimates: Sequence[FlexibleDomainWithEstimates],
     xlim: AxisLimits, ylim: AxisLimits,
         flip_axes: bool = True, K_max: int | None = None, tick_step: int = 2
 ) -> None:
@@ -234,7 +235,7 @@ def plot_2d_domains_with_num_communities(
 
 
 def plot_2d_domains_with_ami(
-    domains_with_estimates: list[FlexibleDomainWithEstimates],
+    domains_with_estimates: Sequence[FlexibleDomainWithEstimates],
         ground_truth: list[int], xlim: AxisLimits, ylim: AxisLimits,
         flip_axes: bool = True
 ) -> None:
diff --git a/utilities/py.typed b/utilities/py.typed
new file mode 100644
index 0000000..e69de29

From 67f9e799a5a14ebbdc5e3767e3840cd94fdde2c6 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Wed, 14 Oct 2026 22:48:13 -0400
Subject: [PATCH 09/11] Update CI to run mypy on modern Python versions

Drop Python 3.9 testing for newer typing features
---
 .github/workflows/python-package.yml | 4 +++-
 pyproject.toml                       | 3 +--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index 50afec5..c38cc5e 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -14,7 +14,7 @@ jobs:
     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 +39,8 @@ jobs:
           # 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
diff --git a/pyproject.toml b/pyproject.toml
index 0831cae..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,7 +17,6 @@ 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",

From 256364dbde13e08d02efbf03c6f411fdcdcbd8d4 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Thu, 15 Oct 2026 10:59:39 -0400
Subject: [PATCH 10/11] Resolve flake8 linting complaints

---
 tests/test_documentation_examples.py      |  2 +-
 tests/test_parallel_leiden_performance.py |  4 +-
 utilities/champ_utilities.py              |  1 -
 utilities/parameter_estimation.py         | 55 ++++++++++++-----------
 4 files changed, 31 insertions(+), 31 deletions(-)

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/utilities/champ_utilities.py b/utilities/champ_utilities.py
index e0407b3..37b79d9 100644
--- a/utilities/champ_utilities.py
+++ b/utilities/champ_utilities.py
@@ -2,7 +2,6 @@
 if TYPE_CHECKING:
     import igraph
     import multiprocessing.pool
-    import scipy.optimize
 import warnings
 from collections import defaultdict
 from math import floor
diff --git a/utilities/parameter_estimation.py b/utilities/parameter_estimation.py
index ea7b568..321b324 100644
--- a/utilities/parameter_estimation.py
+++ b/utilities/parameter_estimation.py
@@ -112,33 +112,34 @@ def check_multilayer_graph_consistency(
     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"]
+    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]
 

From 9d3620391f60dc562c19a45306a014320dba5d87 Mon Sep 17 00:00:00 2001
From: Ryan Gibson <ryan.alexander.gibson@gmail.com>
Date: Thu, 15 Oct 2026 12:31:44 -0400
Subject: [PATCH 11/11] Make test's multilayer SBM function more robust

---
 tests/shared_testing_functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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