Source code for gerrychain.meta.diversity

"""
Simple tooling to collect diversity stats on chain runs
"""

from collections.abc import Hashable, Iterable, Iterator
from dataclasses import dataclass

from gerrychain.partition import Partition




@dataclass
class DiversityStats:
    """
    Lightweight stats object that reports the diversity of a given chain.

    Attributes:
        unique_plans (int): The number of unique plans seen so far.
        unique_districts (int): The number of unique districts seen so far.
        steps_taken (int): The number of steps taken so far.
    Example usage::

        DiversityStats(unique_plans=44162, unique_districts=82992, steps_taken=100000)
    """

    unique_plans: int
    unique_districts: int
    steps_taken: int





def collect_diversity_stats(
    chain: Iterable[Partition],
) -> Iterator[tuple[Partition, DiversityStats]]:
    """Report the diversity of the chain being run, live, as a drop-in wrapper.

    Requires the cut_edges updater on each `Partition` object. Plans/districts are considered
    distinct if they are not isomorphic. That is, relabled plans and districts are considered
    non-unique and counted as duplicate.

    Example usage::
        for partition, stats in collect_diversity_stats( Replay( graph, "sample-run.chain" ) ):
            print(stats) # normal chain stuff here

    Args:
        chain (Iterable[Partition]): A chain object to collect stats on.

    Returns:
        Iterator[tuple[Partition, DiversityStats]]: Pairs of partitions and cumulative statistics.
    """
    seen_plans: set[frozenset[frozenset[Hashable]]] = set()
    seen_districts: set[frozenset[Hashable]] = set()

    unique_plans = 0
    unique_districts = 0
    steps_taken = 0

    for partition in chain:
        steps_taken += 1

        for _, nodes in partition.assignment.parts.items():
            hashable_nodes = frozenset(nodes)
            if hashable_nodes not in seen_districts:
                unique_districts += 1
                seen_districts.add(hashable_nodes)

        cut_edges = partition["cut_edges"]
        hashable_cut_edges = frozenset(frozenset(edge) for edge in cut_edges)
        if hashable_cut_edges not in seen_plans:
            unique_plans += 1
            seen_plans.add(hashable_cut_edges)

        stats = DiversityStats(
            unique_plans=unique_plans,
            unique_districts=unique_districts,
            steps_taken=steps_taken,
        )

        yield partition, stats