"""
This module implements algorithms for finding balanced subsets of nodes in a graph.
By "balanced" we mean that the total population in the subset is appropriate if the graph
were divided into N subsets of approximately the same size - where N is the number of districts
desired.
Stated differently, if you want to partition a graph into N equally sized subsets (by population)
of connected nodes, then this is the module that has those algorithms.
It leverages the GerryChain Graph object to handle graph structures.
Key functionalities include:
- The `_PopulatedGraph` class, which represents a graph with additional population data,
and methods for assessing and modifying this data.
- Functions for finding balanced edge cuts in a populated graph, either through
contraction or memoization techniques.
- The routines to find balanced subsets of nodes, such as bipartition_tree()
Dependencies:
- random: Provides random number generation for probabilistic approaches.
- typing: Used for type hints.
"""
import itertools
import random
import warnings
from collections import deque
from collections.abc import Callable, Hashable, Mapping, Sequence, Set as AbstractSet
from functools import partial
from inspect import signature
from typing import NamedTuple, Protocol, cast
from weakref import WeakSet
from .._deprecated import adapt_legacy_callable, adapt_legacy_cut_choice, deprecated_parameters
from .._rng import make_rng
from ..graph import FrozenGraph, Graph
from .spanning_tree import random_spanning_tree
# An owning operation shares one RNG across all of its bipartition attempts.
_node_repeats_warning_rngs: WeakSet[random.Random] = WeakSet()
"""
RustworkX Issues:
Note: This module has been modified in order to be able to operate on
RustworkX.PyGraph objects in addition to NetworkX.Graph objects. The
new GerryChain Graph object embeds a graph object that can be based
either on NetworkX or RustworkX (the old GerryChain graph object was
a subclass of NetworkX.Graph). The reason for supporting
RustworkX.PyGrapy is performance - it is much faster than NetworkX.
The default usage model is for users to create graphs using Networkx,
both because of legacy concerns but also because it is just convenient
to create graphs in NetworkX. When the user creates a Partition
object, the embedded NetworkX.Graph object will be automatically
converted to be a RustworkX.PyGraph object. And since most of the
routines in this module are used after the creation of a Partition
object, the underlying graph operations will most often be performed by
RustworkX code.
There is a way to override this behavior by setting the value of a
variable in the code (in partition.py). The variable in
partition.py that controls this (and its default setting) is:
test_performance_using_NX_graph = False
As the name of the variable suggests, setting the variable to True
allows comparison of performance between NX and RX. If set to True,
the code that creates a partition does NOT convert the underlying
graph object from NX to RX. So, the code is designed to work on
both kinds of embedded graph objects - NX and RX.
Many of the functions in this file operate on subgraphs, and subgraphs
behave differently from NX subgraphs. In particular, RustworkX subgraphs
typically change the IDs for the nodes in the subgraph, so that a node with
an ID of say, 5, in a parent graph might have an ID of say, 2, in the subgraph.
It is the same node, with the same data, but its ID has changed. Note that
this affects edges too - edges are also reassigned new edge_ids.
To deal with subgraphs having different node_ids from their parent graph
the code has implemented a mapping (dictionary) for subgraph node_ids to
parent graph node_ids. This makes it possible for routines to convert any
results obtained using a subgraph to the appropriate node_ids for the
parent graph. Stated differently, if a routine computes node_id or edge_id
related results using subgraphs, then those results need to be translated
back (in the case of RX) into the context of the parent graph, and this
mapping makes that possible.
One way that the code in this module tries to reduce the risk of using
the wrong node_ids and edge_ids is to have all calls on subgraph()
happen as actual parameters to a function call. This ensures that
the node_ids and edge_ids of the subgraph are only visible in the
context of the called function - they cannot "leak" into the code
that uses parent graph node_ids and edge_ids. It is the responsibility
of the called function to translate any node_id and edge_id results
back to the context of the parent graph.
So - if you decide to write custom code that involves subgraphs, please
spend a little time reviewing how the code in this module is implemented
so that you can avoid subtle nasty bugs...
Note: Predecessor and successor functions have been moved to the new
GerryChain Graph object. The reason for moving them was to remove dependencies
on NetworkX (and RustworkX) from this module.
"""
GraphLike = Graph | FrozenGraph
SpanningTreeFn = Callable[..., Graph]
class _PopulatedGraph:
"""
A class representing a graph with population information. It is used by
the code that finds districts (subsets of nodes) that each have
the appropriate population (within epsilon of the ideal district population).
Attributes:
graph (Graph): The underlying graph structure.
subsets (dict): A dictionary mapping nodes to their subsets. This is used
as a way to accumulate nodes that "belong together", so it is kind of
a scratchpad for nodes that hopefully will become a district.
population (dict): A dictionary mapping nodes to the population of the node.
tot_pop (int | float): The total population of the graph.
ideal_pop (float): The ideal population for each district.
epsilon (float): The tolerance for population deviation from the ideal population within
each
district.
"""
def __init__(
self,
graph: GraphLike,
populations: dict[Hashable, int | float],
ideal_pop: int | float,
epsilon: float,
) -> None:
"""Initialize a _PopulatedGraph instance.
Args:
graph (Graph): The underlying graph structure.
populations (dict): A dictionary mapping nodes to their populations.
ideal_pop (int | float): The ideal population for each district.
epsilon (float): The tolerance for population deviation as a percentage of the ideal
population within each district.
"""
self.graph = graph
self.subsets = {node_id: {node_id} for node_id in graph.node_indices}
self.population = populations.copy()
self.tot_pop = sum(self.population.values())
self.ideal_pop = ideal_pop
self.epsilon = epsilon
# Note: _degrees maintains the number of edges from a node which can change
# when a node is "contracted" - see _contract_node() below.
self._degrees = {node_id: graph.degree(node_id) for node_id in graph.node_indices}
def degree(self, node: Hashable) -> int:
return self._degrees[node]
def _contract_node(self, node: Hashable, parent: Hashable) -> None:
# Merge the population and the subset of nodes from "self"
# into the parent node, and reduce the degrees of the parent
# indicating that the "self" node is no longer in the tree.
self.population[parent] += self.population[node]
self.subsets[parent] |= self.subsets[node]
self._degrees[parent] -= 1
def is_within_epsilon(self, population: int | float) -> bool:
"""Return whether a population is within epsilon of the ideal district population.
This is the single definition of "balanced" used by both cut finders. Keeping it in one
place matters: the bound is inclusive, and when the two finders spelled the comparison out
separately they disagreed about whether an exactly-on-target cut counted.
Args:
population (int | float): The population to test.
Returns:
bool: True if the population is acceptable.
"""
return abs(population - self.ideal_pop) <= self.epsilon * self.ideal_pop
# frm: only ever used inside this file
# But maybe this is intended to be used externally...
def has_ideal_population(self, node: Hashable, single_district_cut: bool = False) -> bool:
"""Checks if a merged node is within epsilon of the ideal population.
Args:
node (Hashable): The node to check.
single_district_cut (bool, optional): Whether or not we are cutting off a single
district. When set to False, we check if the node we are cutting and the remaining
graph are both within epsilon of the ideal population. When set to True, we only
check if the node we are cutting is within epsilon of the ideal population.
Defaults to False.
Returns:
bool: True if the node has an ideal population within the graph up to epsilon.
"""
node_population = self.population[node]
if single_district_cut:
return self.is_within_epsilon(node_population)
return self.is_within_epsilon(node_population) and self.is_within_epsilon(
self.tot_pop - node_population
)
def __repr__(self) -> str:
graph_info = f"Graph(nodes={len(self.graph.node_indices)}, edges={len(self.graph.edges)})"
return (
f"{self.__class__.__name__}("
f"graph={graph_info}, "
f"total_population={self.tot_pop}, "
f"ideal_population={self.ideal_pop}, "
f"epsilon={self.epsilon})"
)
# Define a _Cut object to describe a set of nodes that could form a district.
#
# A little background:
#
# GerryChain uses spanning trees to find parts / districts that have the
# appropriate population. The general approach is to create a spanning
# tree from a graph of nodes and to then walk the tree looking for an
# edge that has on one side or the other a set of nodes that has the
# appropriate population for a part / district. When such an edge is
# found, the code creates a "_Cut" object to record: 1) the edge in
# the spanning tree, 2) the "weight" of the edge, and 3) the set of
# nodes that could form a part / district.
#
# The "weight" is used to choose which _Cut (if more than one is found)
# to use to create a part / district.
#
# Note that a _Cut object is only ever used inside this file.
#
class _Cut(NamedTuple):
"""A candidate edge cut and the nodes on one side of it."""
edge: tuple[Hashable, Hashable]
weight: float
subset: frozenset[Hashable]
class FindBalancedEdgeCutsFn(Protocol):
"""Enumerates the population-balanced cuts of an already-drawn spanning tree.
A bipartition function draws a spanning tree, calls one of these to find every edge whose
removal leaves acceptable populations, and then hands the result to a :class:`CutChoiceFn` to
pick one. :func:`find_balanced_edge_cuts_memoization` and
:func:`find_balanced_edge_cuts_contraction` both match this shape.
Args:
h (_PopulatedGraph): The spanning tree annotated with node populations, the ideal
population, and the allowed epsilon.
single_district_cut (bool, optional): When False, both sides of the cut must be within
epsilon of the ideal population. When True, only the returned side must be, which is
what lets callers peel off a single district of a given size. Default is False.
rootnode_choice_fn (Callable | None, optional): Chooses the root to hang the tree from,
given the candidate nodes. Defaults to a uniform random choice. Finders that examine
every edge regardless of root ignore this.
rng (random.Random): The RNG supplied by the owning operation.
Returns:
list[_Cut]: Every balanced cut found, which may be empty if the tree admits none.
"""
def __call__(
self,
h: _PopulatedGraph,
/,
*,
single_district_cut: bool = False,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
rng: random.Random,
) -> list[_Cut]: ...
class _BoundFindBalancedEdgeCutsFn(Protocol):
"""A :class:`FindBalancedEdgeCutsFn` with ``single_district_cut`` and ``rng`` already bound.
:func:`bipartition_tree` binds those two before handing the finder to the retry loop, which
then varies only the root choice between attempts.
"""
def __call__(
self,
h: _PopulatedGraph,
/,
*,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
) -> list[_Cut]: ...
[docs]
class BipartitionTreeFn(Protocol):
"""The bipartitioning step used by seed-plan generators.
Called with the subgraph to split and keyword configuration; returns the node set of one
side of the bipartition. :func:`bipartition_tree` (or a ``functools.partial`` of it, to
bind extra options like ``max_attempts``) matches this shape.
"""
def __call__(
self,
subgraph_to_split: GraphLike,
/,
*,
pop_col: str,
pop_target: float,
epsilon: float,
node_repeats: int = 0,
single_district_cut: bool = False,
rng: random.Random,
) -> AbstractSet[Hashable]: ...
[docs]
class ReComBipartitionTreeFn(Protocol):
"""The bipartitioning step used by ReCom proposals.
This extends the seed-plan call shape with ``region_surcharge`` and returns the node set on one
side of the bipartition.
"""
def __call__(
self,
subgraph_to_split: GraphLike,
/,
*,
pop_col: str,
pop_target: float,
epsilon: float,
node_repeats: int = 0,
single_district_cut: bool = False,
region_surcharge: dict[str, float] | None = None,
rng: random.Random,
) -> AbstractSet[Hashable]: ...
class CutChoiceFn(Protocol):
"""Selects which balanced cut to make. Called as ``cut_choice_fn(cuts, rng=rng)``."""
def __call__(self, cut_edge_list: list[_Cut], /, *, rng: random.Random) -> _Cut: ...
class RegionAwareCutChoiceFn(Protocol):
"""Region-aware variant of :class:`CutChoiceFn`.
A cut chooser whose signature declares ``populated_graph`` and ``region_surcharge``
parameters (detected via ``inspect.signature``) is additionally handed the populated
spanning tree and the region surcharges when selecting a cut.
"""
def __call__(
self,
cut_edge_list: list[_Cut],
/,
*,
populated_graph: _PopulatedGraph,
region_surcharge: dict[str, float],
rng: random.Random,
) -> _Cut: ...
def _bfs_predecessors_and_successors_for_tree(
tree: GraphLike, root: Hashable, build_successors: bool = False
) -> tuple[dict[Hashable, Hashable], dict[Hashable, list[Hashable]] | None]:
"""Return ``(pred, succ)`` parent/child maps for the tree ``graph`` rooted at ``root``.
``pred`` maps every non-root node to its parent (the neighbor on the path toward
``root``). When ``build_successors`` is True, ``succ`` maps every internal node to the
list of its children; otherwise ``succ`` is None.
This is one tight breadth-first traversal that fetches ``graph.neighbors()`` once per
node as it is visited. It replaces ``graph.predecessors(root)`` and
``graph.successors(root)``, each of which ran its own traversal through the
``Graph.generic_bfs_*`` generators. The memoization path previously paid for that traversal
twice (predecessors then successors); fusing them here computes both in one pass.
For a tree, every non-root node has a unique parent, so ``pred`` is identical to the map
the old BFS produced regardless of neighbor-visitation order.
Special thanks to Jeanne Clelland who's original work informed this section of code.
Args:
tree (Graph): The tree for which to compute predecessors and successors.
root (Hashable): The root node of the tree.
build_successors (bool, optional): Whether to build the successors map. Defaults to False.
Returns:
tuple[dict[Hashable, Hashable], dict[Hashable, list[Hashable]] | None]: A tuple containing the predecessors
map and the successors map (or None if not built).
"""
# TODO: Technically, you can run this on a graph and it should be fine, but there is an implicit
# assumption that the passed graph is a tree, and we should probably check this somewhere
# if there is a cheap way to do it.
pred: dict[Hashable, Hashable] = {}
succ: dict[Hashable, list[Hashable]] | None = {} if build_successors else None
seen = {root}
queue = deque([root])
while queue:
node = queue.popleft()
children = None
for neighbor in tree.neighbors(node):
if neighbor not in seen:
seen.add(neighbor)
pred[neighbor] = node
queue.append(neighbor)
if build_successors:
assert succ is not None
if children is None:
children = succ[node] = []
children.append(neighbor)
return pred, succ
[docs]
@deprecated_parameters(
renamed={
"one_sided_cut": "single_district_cut",
"choice": "rootnode_choice_fn",
}
)
def find_balanced_edge_cuts_contraction(
h: _PopulatedGraph,
single_district_cut: bool = False,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
*,
rng: random.Random | int | None = None,
) -> list[_Cut]:
"""Find balanced edge cuts using contraction.
Args:
h (_PopulatedGraph): The populated graph.
single_district_cut (bool, optional): Whether or not we are cutting off a single district.
When set to False, we check if the node we are cutting and the remaining graph are both
within epsilon of the ideal population. When set to True, we only check if the node we
are cutting is within epsilon of the ideal population. Defaults to False.
rootnode_choice_fn (Callable | None): The function used to select the root node_id.
Defaults to the supplied RNG's ``choice``.
rng (random.Random | int | None, optional): Source of randomness. Pass a shared
``Random`` for repeated standalone calls; an integer restarts the stream each call.
Returns:
list[_Cut]: A list of balanced edge cuts.
"""
rng = make_rng(rng)
if rootnode_choice_fn is None:
rootnode_choice_fn = rng.choice
node_ids = h.graph.nodes if h.graph.is_nx_graph() else h.graph.node_indices
# NOTE: The root must have degree > 1 here: the loop below seeds its queue from every degree-1
# node and then looks up `pred[leaf]`. A leaf root would land in that queue without having a
# parent, raising KeyError. On a two-node tree both nodes are leaves, so nothing qualifies and
# `rng.choice` raises IndexError.
root = rootnode_choice_fn([node_id for node_id in node_ids if h.degree(node_id) > 1])
# Parent map for iteratively contracting leaves. This used to call h.graph.predecessors(root)
pred, _ = _bfs_predecessors_and_successors_for_tree(h.graph, root)
cuts = []
# frm: Work up from leaf nodes to find subtrees with the "correct"
# population. The algorighm starts with real leaf nodes, but
# if a node does not have the "correct" population, then that
# node is merged (contracted) into its parent, effectively
# creating another leaf node which is then added to the end
# of the queue.
#
# In this way, we calculate the total population of subtrees
# by going bottom up, until we find a subtree that has the
# "correct" population for a cut.
# frm: ??? Note that there is at least one other routine in this file
# that does something similar (perhaps exactly the same).
# Need to figure out why there are more than one way to do this...
leaves = deque(node_id for node_id in node_ids if h.degree(node_id) == 1)
while len(leaves) > 0:
leaf = leaves.popleft()
if h.has_ideal_population(leaf, single_district_cut=single_district_cut):
# frm: If the population of the subtree rooted in this node is the correct
# size, then add it to the cut list. Note that if single_district_cut == False,
# then the cut means the cut bisects the graph of the spanning tree.
e = (leaf, pred[leaf])
cuts.append(
_Cut(
edge=e,
weight=h.graph.edge_data(h.graph.get_edge_id_from_edge(e)).get(
"random_weight", rng.random()
),
subset=frozenset(h.subsets[leaf].copy()),
)
)
# Contract the leaf: frm: merge the leaf's population into the parent and add the
# parent to "leaves"
parent = pred[leaf]
# Add child population and subsets to parent, reduce parent's degree by 1
# This effectively removes the leaf from the tree, adding all of its data
# to the parent.
h._contract_node(leaf, parent)
if h.degree(parent) == 1 and parent != root:
# frm: Only add the parent to the end of the queue when we are merging
# the last leaf - this makes sure we only add the parent node to
# the queue one time...
leaves.append(parent)
return cuts
def _calc_pops(
succ: dict[Hashable, list[Hashable]], root: Hashable, h: _PopulatedGraph
) -> dict[Hashable, int | float]:
"""Return A dictionary mapping nodes to their subtree populations.
Args:
succ (dict): The successors of the graph.
root (Hashable): The root node of the graph.
h (_PopulatedGraph): The populated graph.
Returns:
dict: A dictionary mapping nodes to their subtree populations.
"""
# frm: This took me a while to sort out what was going on.
# Conceptually it is easy - given a tree anchored in a root node,
# calculate the population in each subtree going bottom-up.
# The stack (deque) provides the mechanism for going bottom-up.
# On the way down, you just put nodes in the stack (append is like
# push() which seems odd to me, but whatever...) then on the way back
# up, you add the totals for each child to your own population and
# presto you have the total population for each subtree...
#
# For this to work, you just need to have a list of nodes with
# their successors associated with them...
#
subtree_pops: dict[Hashable, int | float] = {}
stack = deque(n for n in succ[root])
while stack:
next_node = stack.pop()
if next_node not in subtree_pops:
if next_node in succ:
children = succ[next_node]
if all(c in subtree_pops for c in children):
subtree_pops[next_node] = sum(subtree_pops[c] for c in children)
subtree_pops[next_node] += h.population[next_node]
else:
stack.append(next_node)
for c in children:
if c not in subtree_pops:
stack.append(c)
else:
subtree_pops[next_node] = h.population[next_node]
return subtree_pops
def _nodes_in_subtree(start: Hashable, succ: dict[Hashable, list[Hashable]]) -> set[Hashable]:
"""
Collects the nodes in a subtree that are rooted in the "start" node.
Args:
start (Hashable): The start node.
succ (dict): The successors of the graph.
Returns:
set: A set of nodes for a particular district (only one side of the cut).
"""
"""
frm: Compute the nodes in a subtree defined by a _Cut.
This routine computes the set of nodes in a subtree rooted in the
node identified by "start" in the tree defined by "succ".
As such it is highly dependent on context and is not generally
useful. That is, it is essentially just a way to refactor some
code used in a couple of places so that the logic in the code is
in one place instead of several.
To be specific, _Cuts are always relative to a specific tree for
a partition. This tree is a "spanning tree" that converts the
graph into a DAG. _Cuts are then computed by finding subtrees
of that DAG that have the appropriate population (this could
presumably be modified to include other factors).
When a _Cut is created, we want to collect all of the nodes that
are in the subtree, and this is what this routine does. It
merely starts at the root of the subtree (start) and goes down
the subtree, adding each node to a set.
frm: TODO: Documentation: Add the above explanation for what a _Cut is and how
we find them by converting the graph to a DAG and
then looking for subtrees to a block header at the
top of this file. It will give the reader some
idea wtf is going on... ;-)
"""
nodes = set()
queue = deque([start])
while queue:
next_node = queue.pop()
if next_node not in nodes:
nodes.add(next_node)
if next_node in succ:
for c in succ[next_node]:
if c not in nodes:
queue.append(c)
return nodes
# frm: used externally by tree_proposals.py
[docs]
@deprecated_parameters(
renamed={
"one_sided_cut": "single_district_cut",
"choice": "rootnode_choice_fn",
}
)
def find_balanced_edge_cuts_memoization(
h: _PopulatedGraph,
single_district_cut: bool = False,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
*,
rng: random.Random | int | None = None,
) -> list[_Cut]:
"""Find balanced edge cuts using memoization.
This function takes a _PopulatedGraph object and a choice function as input and returns a list
of balanced edge cuts. A balanced edge cut is defined as a cut that divides the graph into two
subsets, such that the population of each subset is close to the ideal population defined by
the _PopulatedGraph object.
Args:
h (_PopulatedGraph): The _PopulatedGraph object representing the graph.
single_district_cut (bool, optional): Whether or not we are cutting off a single district.
When set to False, we check if the node we are cutting and the remaining graph are both
within epsilon of the ideal population. When set to True, we only check if the node we
are cutting is within epsilon of the ideal population. Defaults to False.
rootnode_choice_fn (Callable | None): The choice function used to select the root
node_id. Defaults to the supplied RNG's ``choice``.
rng (random.Random | int | None, optional): Source of randomness. Pass a shared
``Random`` for repeated standalone calls; an integer restarts the stream each call.
Returns:
list[_Cut]: A list of balanced edge cuts.
"""
rng = make_rng(rng)
if rootnode_choice_fn is None:
rootnode_choice_fn = rng.choice
node_ids = h.graph.nodes if h.graph.is_nx_graph() else h.graph.node_indices
# NOTE: Any node may serve as the root. Every non-root node contributes the one edge joining it
# to its parent, and a tree on n nodes has n - 1 of each, so the candidate edges are all of them
# wherever we root. The root only decides which side of an edge counts as "below", and so which
# side a cut hands back. `find_balanced_edge_cuts_contraction` is different: it requires a root
# of degree > 1, because it seeds its worklist from every degree-1 node and then looks up
# `pred[leaf]`, which a leaf root would not have.
root = rootnode_choice_fn(list(node_ids))
# Parent and child maps for the tree rooted at `root`. For a tree each non-root node has a
# unique parent, so `pred` is identical to the old map; the order of children within
# succ[node] may differ, which does not affect the set of cuts found.
pred, succ = _bfs_predecessors_and_successors_for_tree(h.graph, root, build_successors=True)
assert succ is not None
total_pop = h.tot_pop
# Calculate the population of each subtree in the "succ" tree
subtree_pops = _calc_pops(succ, root, h)
def cut_at(node: Hashable, subset: AbstractSet[Hashable]) -> _Cut:
"""Build the cut that removes the edge joining ``node`` to its parent.
``subset`` is the side of that edge the cut hands back. The edge keeps whatever
``random_weight`` the spanning tree gave it, falling back to a fresh draw. Note that the
draw happens either way, so the RNG advances once per cut found.
"""
edge = (node, pred[node])
fallback_weight = rng.random()
return _Cut(
edge=edge,
weight=h.graph.edge_data(h.graph.get_edge_id_from_edge(edge)).get(
"random_weight", fallback_weight
),
subset=frozenset(subset),
)
def below(node: Hashable) -> AbstractSet[Hashable]:
"""The subtree hanging below ``node``, including ``node`` itself."""
return _nodes_in_subtree(node, succ)
def above(node: Hashable) -> AbstractSet[Hashable]:
"""Everything except the subtree below ``node``."""
return set(h.graph.node_indices) - _nodes_in_subtree(node, succ)
cuts = []
if single_district_cut:
# Peeling off one district: only the side we hand back has to be balanced, so a node
# qualifies if either the subtree below it or everything above it hits the target. The
# remainder is left for later cuts to divide.
for node, tree_pop in subtree_pops.items():
if h.is_within_epsilon(tree_pop):
cuts.append(cut_at(node, below(node)))
elif h.is_within_epsilon(total_pop - tree_pop):
cuts.append(cut_at(node, above(node)))
else:
# Bisecting the tree: both sides have to be balanced at once, which is a much stricter
# test and typically admits far fewer cuts.
for node, tree_pop in subtree_pops.items():
if h.is_within_epsilon(tree_pop) and h.is_within_epsilon(total_pop - tree_pop):
cuts.append(cut_at(node, above(node)))
return cuts
# frm: only used in this file and in a test
[docs]
class BipartitionWarning(UserWarning):
"""
Generally raised when it is proving difficult to find a balanced cut.
"""
pass
# frm: only used in this file and in a test
[docs]
class ReselectException(Exception):
"""
Raised when the tree-splitting algorithm is unable to find a
balanced cut after some maximum number of attempts, but the
user has allowed the algorithm to reselect the pair of
districts from parent graph to try and recombine.
"""
pass
def _max_weight_choice(cut_edge_list: list[_Cut], *, rng: random.Random) -> _Cut:
"""Selects a cut from a list of cuts based on the maximum weight.
This random weight is either assigned during the call to the minimum spanning tree algorithm
(Kruskal's) algorithm or it is generated during the selection of the balanced edges (cf.
`find_balanced_edge_cuts_memoization` and `find_balanced_edge_cuts_contraction`).
This function returns the cut with the highest weight.
In the case where a region aware chain is run, this will preferentially select for cuts that
span different regions, rather than cuts that are interior to that region (the likelihood of
this is generally controlled by the ``region_surcharge`` parameter).
In any case where the surcharges are either not set or zero, this is effectively the same as
calling random.choice() on the list of cuts. Under the above conditions, all of the weights on
the cuts are randomly generated on the interval [0,1], and there is no outside force that might
make the weight assigned to a particular type of cut higher than another.
Args:
cut_edge_list (list[_Cut]): A list of _Cut objects. Each object has an edge, a weight, and a
subset attribute.
rng (random.Random): Random number generator used for fallback selection.
Returns:
_Cut: The cut with the highest random weight.
"""
# Just in case, default to random choice
if not isinstance(cut_edge_list[0], _Cut) or cut_edge_list[0].weight is None:
return rng.choice(cut_edge_list)
# frm: ???: this strikes me as possibly expensive. Computing the
# max in a list is O(N) so not terrible, but this
# might be called lots of times (need to know more about
# how it is used). Would it make sense to have the
# cut_edge_list sorted before it is frozen? I think it
# is now a set, so it would need to be a list... Not
# urgent, but worth looking into at some point...
#
return max(cut_edge_list, key=lambda cut: cut.weight)
def _random_choice(cut_edge_list: list[_Cut], *, rng: random.Random) -> _Cut:
return rng.choice(cut_edge_list)
def _power_set_sorted_by_size_then_sum(
region_surcharge_dict: dict[str, float],
) -> list[tuple[str, ...]]:
"""Power set sorted by size then sum.
This function computes the power set of regions that are listed in the region_surcharge_dict,.
and then sorts that power set by size and then sum.
Args:
region_surcharge_dict (dict): Description
"""
num_regions = len(region_surcharge_dict)
power_set = [
s
for i in range(1, num_regions + 1)
for s in itertools.combinations(region_surcharge_dict.keys(), i)
]
# Sort the subsets in descending order based on
# the sum of their corresponding values in the dictionary
sorted_power_set = sorted(
power_set, key=lambda s: (len(s), sum(region_surcharge_dict[i] for i in s)), reverse=True
)
return sorted_power_set
# Note that the populated graph and the region surcharge are passed
# by object reference. This means that a copy is not made since we
# are not modifying the object in the function, and the speed of
# this randomized selection will not suffer for it.
def _region_preferred_max_weight_choice(
cut_edge_list: list[_Cut],
populated_graph: _PopulatedGraph,
region_surcharge: dict[str, float],
*,
rng: random.Random,
) -> _Cut:
# frm: ???: There is no NX/RX dependency in this routine, but I do
# not yet understand what it does or why...
"""Selects a cut from a list of cuts based on the maximum weight, with a preference for
cuts that span different regions.
This is similare to the `_max_weight_choice` function except that it will preferentially
select one of the cuts that has the highest surcharge. So, if we have a weight dict of the form
``{region1: wt1, region2: wt2}`` , then this function first looks for a cut that is a cut edge
for both ``region1`` and ``region2`` and then selects the one with the highest weight. If no
such cut exists, then it will then look for a cut that is a cut edge for the region with the
highest surcharge (presumably the region that we care more about not splitting).
In the case of 3 regions, it will first look for a cut that is a cut edge for all 3 regions,
then for a cut that is a cut edge for 2 regions sorted by the highest total surcharge, and then
for a cut that is a cut edge for the region with the highest surcharge.
For the case of 4 or more regions, the power set starts to get a bit large, so we default back
to the `_max_weight_choice` function and just select the cut with the highest weight,
which will still preferentially select for cuts that span the most regions that we care about.
Args:
populated_graph (_PopulatedGraph): The populated graph.
region_surcharge (dict): A dictionary of surcharges for the spanning tree algorithm.
cut_edge_list (list[_Cut]): A list of _Cut objects. Each object has an edge, a weight, and a
subset attribute.
rng (random.Random): Random number generator used for selection.
Returns:
_Cut: A random _Cut from the set of possible _Cuts with the highest surcharge.
"""
# Check to see if the list of possible edge_cuts is one
# with edge weights. If not, then just select an edge at random.
if (
not isinstance(region_surcharge, dict)
or not isinstance(cut_edge_list[0], _Cut)
or cut_edge_list[0].weight is None
):
return rng.choice(cut_edge_list)
# Early return for simple cases
if len(region_surcharge) < 1 or len(region_surcharge) > 3:
return _max_weight_choice(cut_edge_list, rng=rng)
# Prepare data for efficient access
edge_region_info = {
cut: {
# Given a cut_edge_list (whose elements have an
# attribute, "edge",) construct a dict
# that associates with each "cut" the
# values of the region_surcharge values
# for both nodes in the edge.
#
# So, if the region_surcharge dict was
# {"muni": 0.2, "water": 0.8} then for
# each cut, cut_n, there would be a
# dict value that looked like:
#
# {
# "muni": (<node_1_muni_weight>, <node_2_muni_weight>)
# "water": (<node_1_water_weight>, <node_2_water_weight>)
# }
#
key: (
populated_graph.graph.node_data(cut.edge[0]).get(key),
populated_graph.graph.node_data(cut.edge[1]).get(key),
)
for key in region_surcharge
}
for cut in cut_edge_list
}
# Generate power set sorted by surcharge, then filter cuts based
# on region matching
power_set = _power_set_sorted_by_size_then_sum(region_surcharge)
for region_combination in power_set:
suitable_cuts = [
cut
for cut in cut_edge_list
if all(
edge_region_info[cut][key][0] != edge_region_info[cut][key][1]
for key in region_combination
)
]
if suitable_cuts:
return _max_weight_choice(suitable_cuts, rng=rng)
return _max_weight_choice(cut_edge_list, rng=rng)
def _internal_bipartition_tree(
subgraph_to_split: GraphLike,
pop_col: str,
pop_target: int | float,
epsilon: float,
node_repeats: int = 0,
spanning_tree: GraphLike | None = None,
spanning_tree_fn: SpanningTreeFn = random_spanning_tree,
region_surcharge: dict[str, float] | None = None,
spanning_tree_fn_kwargs: Mapping[str, object] | None = None,
find_balanced_edge_cuts_fn: FindBalancedEdgeCutsFn = find_balanced_edge_cuts_memoization,
single_district_cut: bool = False,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
max_attempts: int = 100000,
warn_attempts: int = 1000,
allow_pair_reselection: bool = False,
repeat_until_valid: bool = True, # frm: TODO: Have this NOT default...
cut_choice_fn: CutChoiceFn | RegionAwareCutChoiceFn = _region_preferred_max_weight_choice,
*,
rng: random.Random,
) -> tuple[int, AbstractSet[Hashable]]:
"""Find a population-balanced connected subset of nodes.
The returned subset induces a connected subgraph, and its complement forms the other part of
the partition.
This function finds a balanced 2 partition of a graph by drawing a spanning tree and finding an
edge to cut that leaves at most an epsilon imbalance between the populations of the parts.
Each spanning tree is searched once plus ``node_repeats`` additional times before a new tree
is drawn.
Builds up a connected subgraph with a connected complement whose population is ``epsilon *
pop_target`` away from ``pop_target``.
Args:
subgraph_to_split (Graph | FrozenGraph): The graph to partition.
pop_col (str): The node attribute holding the population of each node.
pop_target (int | float): The target population for the returned subset of nodes.
epsilon (float): The allowable deviation from ``pop_target`` (as a percentage of
``pop_target``) for the subgraph's population.
node_repeats (int, optional): Additional roots to try on each spanning tree before
drawing a new tree. Defaults to 0. Positive values are useful with contraction or
custom cut-edge finders, but not with the default memoized finder.
spanning_tree (Graph | None, optional): The spanning tree for the algorithm to use (used
when the algorithm chooses a new root and for testing).
spanning_tree_fn (Callable, optional): The random spanning tree algorithm to use if a
spanning tree is not provided. Defaults to `random_spanning_tree`.
region_surcharge (dict | None, optional): A dictionary of surcharges for the spanning
tree algorithm. Defaults to None.
spanning_tree_fn_kwargs (dict | None, optional): Extra keyword arguments forwarded
verbatim to ``spanning_tree_fn``. Use this to set function-specific options that are
not named here, e.g. ``{"treat_unassigned_as_single_region": True}`` for
``random_spanning_tree``. Defaults to None.
find_balanced_edge_cuts_fn (FindBalancedEdgeCutsFn, optional): The function to find
balanced edge cuts. Defaults to `find_balanced_edge_cuts_memoization`. A custom
finder must accept ``single_district_cut``, ``rootnode_choice_fn``, and ``rng`` as
keyword arguments; see :class:`FindBalancedEdgeCutsFn`.
single_district_cut (bool, optional): Passed to the ``find_balanced_edge_cuts_fn``.
Determines whether we are cutting off a single district when partitioning the tree. When
set to False, we check if the node we are cutting and the remaining graph are both
within epsilon of the ideal population. When set to True, we only check if the node we
are cutting is within epsilon of the ideal population. Defaults to False.
repeat_until_valid (bool, optional): Whether to keep drawing trees until a valid cut is
found. Defaults to ``True``.
rootnode_choice_fn (Callable | None): The function to make a random choice of root node
for the population tree. Passed to ``find_balanced_edge_cuts_fn``. Can be substituted
for testing. Defaults to the supplied RNG's ``choice``.
max_attempts (int): The maximum number of attempts that should be made
to bipartition. Defaults to 100,000.
warn_attempts (int, optional): The number of attempts after which a warning is issued if a
balanced cut cannot be found. Defaults to 1000.
allow_pair_reselection (bool, optional): Whether we would like to return an error to the
calling function to ask it to reselect the pair of nodes to try and recombine. Defaults
to False.
cut_choice_fn (CutChoiceFn | RegionAwareCutChoiceFn, optional): Function used to
select the cut edge. Defaults to `_region_preferred_max_weight_choice`.
Note that this function should gracefully handle the case when the edges in the list of
possible balanced cuts do not have edge weights - in this case, it should default to
a uniform random choice.
rng (random.Random): The RNG supplied by the owning operation.
Returns:
tuple[int, AbstractSet[Hashable]]: The number of possible cuts and the selected connected
subset of nodes.
Raises:
BipartitionWarning: If a possible cut cannot be found after 1000 attempts.
RuntimeError: If a possible cut cannot be found after the maximum number of attempts
given by ``max_attempts``.
"""
# Create a new empty dict now, instead of having a default value for the
# paramter which would instantiate once instead of for every function invocation.
if region_surcharge is None:
region_surcharge = {}
spanning_tree_fn = adapt_legacy_callable(
spanning_tree_fn,
"Spanning-tree function",
dropped={"region_surcharge"},
)
find_balanced_edge_cuts_fn = cast(
FindBalancedEdgeCutsFn,
adapt_legacy_callable(
find_balanced_edge_cuts_fn,
"Balanced-edge-cut function",
renamed={
"single_district_cut": "one_sided_cut",
"rootnode_choice_fn": "choice",
},
),
)
cut_choice_fn = cast(
CutChoiceFn | RegionAwareCutChoiceFn,
adapt_legacy_cut_choice(cut_choice_fn),
)
if spanning_tree_fn_kwargs and "region_surcharge" in spanning_tree_fn_kwargs:
raise ValueError(
"Pass region_surcharge via the region_surcharge parameter, not inside "
"spanning_tree_fn_kwargs."
)
unwrapped_cut_finder = find_balanced_edge_cuts_fn
while isinstance(unwrapped_cut_finder, partial):
unwrapped_cut_finder = unwrapped_cut_finder.func
# Warn the user if they are using node_repeats with the default cut finder. The memoization
# cut finder exhaustively searches each spanning tree, so node_repeats is not beneficial.
# NOTE: The existence of the default value of 2 was a hold-over from a very early version
# of GerryChain where the contraction method was the default. Lots of legacy code
# still uses the default value of 2, so we don't want to change it. However, we do
# want to warn the user that it is not beneficial with the default cut finder which
# is now memoization.
if (
node_repeats > 0
and unwrapped_cut_finder is find_balanced_edge_cuts_memoization
and rng not in _node_repeats_warning_rngs
):
warnings.warn(
"node_repeats is not beneficial with `find_balanced_edge_cuts_memoization`, which "
"exhaustively searches each spanning tree. Set node_repeats=0 to redraw the tree "
"after an unsuccessful search. Positive values are useful only with "
"`find_balanced_edge_cuts_contraction` or a custom cut-edge finder.",
UserWarning,
stacklevel=2,
)
_node_repeats_warning_rngs.add(rng)
spanning_tree_fn = partial(
spanning_tree_fn,
region_surcharge=region_surcharge,
rng=rng,
**(spanning_tree_fn_kwargs or {}),
)
bound_find_balanced_edge_cuts_fn: _BoundFindBalancedEdgeCutsFn = partial(
find_balanced_edge_cuts_fn, single_district_cut=single_district_cut, rng=rng
)
# Find possible edge cuts if they exist.
#
# Note that _get_possible_edge_cuts_and_populated_graph() raises an
# exception if max_attempts is exceeded.
#
possible_cuts, populated_graph = _get_possible_edge_cuts_and_populated_graph(
graph_to_split=subgraph_to_split,
pop_col=pop_col,
pop_target=pop_target,
epsilon=epsilon,
node_repeats=node_repeats,
spanning_tree=spanning_tree,
spanning_tree_fn=spanning_tree_fn,
find_balanced_edge_cuts_fn=bound_find_balanced_edge_cuts_fn,
rootnode_choice_fn=rootnode_choice_fn,
max_attempts=max_attempts,
warn_attempts=warn_attempts,
repeat_until_valid=repeat_until_valid,
allow_pair_reselection=allow_pair_reselection,
)
num_cuts = len(possible_cuts)
if num_cuts != 0:
# Note: There are two different function signatures for cut_choice_fn().
#
# One just takes the set of possible cuts, but the other takes
# two additional parameters, region_surcharge, and populated_graph,
# so we need to see which cut_choice_fn() we are dealing with.
#
if (
"region_surcharge" in signature(cut_choice_fn).parameters
and "populated_graph" in signature(cut_choice_fn).parameters
):
chosen_cut = cast(RegionAwareCutChoiceFn, cut_choice_fn)(
possible_cuts,
populated_graph=populated_graph,
region_surcharge=region_surcharge,
rng=rng,
)
else:
chosen_cut = cast(CutChoiceFn, cut_choice_fn)(possible_cuts, rng=rng)
parent_node_ids = subgraph_to_split.translate_subgraph_node_ids_for_set_of_nodes(
chosen_cut.subset
)
# frm: Not sure if it is important that the returned set be a frozenset...
return num_cuts, frozenset(parent_node_ids)
else:
# No balanced edge_cuts were found.
#
# This is OK if the repeat_until_valid flag is False. In that case (which is
# related to reversible_recom()) we should return 0 for num_cuts and an
# empty list of node_ids. This will signal to the caller that we failed to
# find any balanced edge_cuts.
#
# However, if repeat_until_valid is True, then an exception will be raised by
# lower level code, so we should never end up in the else-part of this if-stmt.
# One might argue that using exceptions this way is not quite cool, but that
# is the way this code works...
#
if not repeat_until_valid:
empty_set = set()
return 0, empty_set
else:
raise Exception("This should never happen...")
[docs]
@deprecated_parameters(
renamed={
"graph": "subgraph_to_split",
"balance_edge_fn": "find_balanced_edge_cuts_fn",
"one_sided_cut": "single_district_cut",
"choice": "rootnode_choice_fn",
"cut_choice": "cut_choice_fn",
}
)
def bipartition_tree(
subgraph_to_split: GraphLike,
pop_col: str,
pop_target: int | float,
epsilon: float,
node_repeats: int = 0,
spanning_tree: GraphLike | None = None,
spanning_tree_fn: SpanningTreeFn = random_spanning_tree,
region_surcharge: dict[str, float] | None = None,
spanning_tree_fn_kwargs: Mapping[str, object] | None = None,
find_balanced_edge_cuts_fn: FindBalancedEdgeCutsFn = find_balanced_edge_cuts_memoization,
single_district_cut: bool = False,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
repeat_until_valid: bool = True,
max_attempts: int = 100000,
warn_attempts: int = 1000,
allow_pair_reselection: bool = False,
cut_choice_fn: CutChoiceFn | RegionAwareCutChoiceFn = _region_preferred_max_weight_choice,
*,
rng: random.Random | int | None = None,
) -> AbstractSet[Hashable]:
"""Find a population-balanced connected subset of nodes.
The returned subset induces a connected subgraph, and its complement forms the other part of
the partition.
This function finds a balanced 2 partition of a graph by drawing a spanning tree and finding an
edge to cut that leaves at most an epsilon imbalance between the populations of the parts.
Each spanning tree is searched once plus ``node_repeats`` additional times before a new tree
is drawn.
Builds up a connected subgraph with a connected complement whose population is ``epsilon *
pop_target`` away from ``pop_target``.
Args:
subgraph_to_split (Graph | FrozenGraph): The graph to partition.
pop_col (str): The node attribute holding the population of each node.
pop_target (int | float): The target population for the returned subset of nodes.
epsilon (float): The allowable deviation from ``pop_target`` (as a percentage of
``pop_target``) for the subgraph's population.
node_repeats (int, optional): Additional roots to try on each spanning tree before
drawing a new tree. Defaults to 0. Positive values are useful with contraction or
custom cut-edge finders, but not with the default memoized finder.
spanning_tree (Graph | None, optional): The spanning tree for the algorithm to use (used
when the algorithm chooses a new root and for testing).
spanning_tree_fn (Callable, optional): The random spanning tree algorithm to use if a
spanning tree is not provided. Defaults to `random_spanning_tree`.
region_surcharge (dict | None, optional): A dictionary of surcharges for the spanning
tree algorithm. Defaults to None.
spanning_tree_fn_kwargs (dict | None, optional): Extra keyword arguments forwarded verbatim
to ``spanning_tree_fn``. Use this to set function-specific options that are not named
here, e.g. ``{"treat_unassigned_as_single_region": True}`` for ``random_spanning_tree``.
Defaults to None.
find_balanced_edge_cuts_fn (FindBalancedEdgeCutsFn, optional): The function to find
balanced edge cuts. Defaults to `find_balanced_edge_cuts_memoization`. A custom
finder must accept ``single_district_cut``, ``rootnode_choice_fn``, and ``rng`` as
keyword arguments; see :class:`FindBalancedEdgeCutsFn`.
single_district_cut (bool, optional): Passed to the ``find_balanced_edge_cuts_fn``.
Determines whether we are cutting off a single district when partitioning the tree. When
set to False, we check if the node we are cutting and the remaining graph are both
within epsilon of the ideal population. When set to True, we only check if the node we
are cutting is within epsilon of the ideal population. Defaults to False.
repeat_until_valid (bool, optional): Whether to keep drawing trees until a valid cut is
found. Defaults to ``True``.
rootnode_choice_fn (Callable | None): The function to make a random choice of root node
for the population tree. Passed to ``find_balanced_edge_cuts_fn``. Can be substituted
for testing. Defaults to the supplied RNG's ``choice``.
max_attempts (int, optional): The maximum number of attempts that should be made
to bipartition. Defaults to 100000.
warn_attempts (int, optional): The number of attempts after which a warning is issued if a
balanced cut cannot be found. Defaults to 1000.
allow_pair_reselection (bool, optional): Whether we would like to return an error to the
calling function to ask it to reselect the pair of nodes to try and recombine. Defaults
to False.
cut_choice_fn (CutChoiceFn | RegionAwareCutChoiceFn, optional): Function used to
select the cut edge. Defaults to `_region_preferred_max_weight_choice`.
Note that this function should gracefully handle the case when the edges in the list of
possible balanced cuts do not have edge weights - in this case, it should default to
a uniform random choice.
rng (random.Random | int | None, optional): Source of randomness. Pass a shared
``Random`` for repeated standalone calls; an integer restarts the stream each call.
Returns:
set[Hashable] | frozenset[Hashable]: A subset of nodes whose induced subgraph is connected.
The other part of the partition is the complement of this subset.
Raises:
BipartitionWarning: If a possible cut cannot be found after 1000 attempts.
RuntimeError: If a possible cut cannot be found after the maximum number of attempts
given by ``max_attempts``.
"""
rng = make_rng(rng)
num_cuts, node_ids = _internal_bipartition_tree(
subgraph_to_split=subgraph_to_split,
pop_col=pop_col,
pop_target=pop_target,
epsilon=epsilon,
node_repeats=node_repeats,
spanning_tree=spanning_tree,
spanning_tree_fn=spanning_tree_fn,
region_surcharge=region_surcharge,
spanning_tree_fn_kwargs=spanning_tree_fn_kwargs,
find_balanced_edge_cuts_fn=find_balanced_edge_cuts_fn,
single_district_cut=single_district_cut,
rootnode_choice_fn=rootnode_choice_fn,
max_attempts=max_attempts,
warn_attempts=warn_attempts,
allow_pair_reselection=allow_pair_reselection,
repeat_until_valid=repeat_until_valid,
cut_choice_fn=cut_choice_fn,
rng=rng,
)
# Note that the value of repeat_until_valid in the above call controls
# whether node_ids is allowed to be the empty_set.
#
# Typically, repeat_until_valid is set to True (usually by default) which
# tells the _internal_bipartition_tree() function to keep trying to find
# balanced edge cuts until it reaches max_attempt, at which point it
# raises an exception (hence no node_ids returned at all).
#
# However, if repeat_until_valid is set to False (which is done for
# reversible_recom(), then _internal_bipartition_tree() will NOT loop
# and it will just return whatever it finds for balanced edge cuts
# after the first attempt, in which case node_ids can be the empty set).
return node_ids
def _get_possible_edge_cuts_and_populated_graph(
#
# Note: Complexity Alert... _get_possible_edge_cuts_and_populated_graph does NOT translate
# node_ids to parent graph context. The reason is that if there are several possible
# edge_cuts found, it is not yet clear which of the edge_cuts to actually cut, and hence
# which node_ids to translate, so we leave both the selection of which edge_cut to cut and
# the translation of node_ids to the caller.
#
graph_to_split: GraphLike,
pop_col: str,
pop_target: int | float,
epsilon: float,
node_repeats: int = 0,
spanning_tree: GraphLike | None = None,
spanning_tree_fn: SpanningTreeFn = random_spanning_tree,
find_balanced_edge_cuts_fn: _BoundFindBalancedEdgeCutsFn = find_balanced_edge_cuts_memoization,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
repeat_until_valid: bool = True,
warn_attempts: int = 1000,
max_attempts: int = 100000,
allow_pair_reselection: bool = False,
) -> tuple[list[_Cut], _PopulatedGraph]:
"""Randomly bipartitions a tree into two subgraphs until a valid bipartition is found.
frm: TODO: Documentation: Peter: Please review the changes I have made to the docstrings...
Args:
graph_to_split (Graph | FrozenGraph): The input graph.
pop_col (str): The name of the column in the graph nodes that contains the population data.
pop_target (int | float): The target population for each subgraph.
epsilon (float): The allowed deviation from the target population as a percentage of
pop_target.
node_repeats (int, optional): Additional roots to try on each spanning tree before
drawing a new tree. Defaults to 0.
repeat_until_valid (bool, optional): Whether to repeat the bipartitioning process until a
valid bipartition is found. Defaults to True.
spanning_tree (Graph | None, optional): The spanning tree to use for bipartitioning. If
None, a random spanning tree will be generated. Defaults to None.
spanning_tree_fn (Callable, optional): The function to generate a spanning tree. Defaults
to random_spanning_tree.
find_balanced_edge_cuts_fn (Callable, optional): The function to find balanced edge cuts.
Defaults to find_balanced_edge_cuts_memoization.
rootnode_choice_fn (Callable | None): The function to choose a random element from a
list. Defaults to the supplied RNG's ``choice``.
max_attempts (int, optional): The maximum number of attempts to find a valid
bipartition. Defaults to 100000.
warn_attempts (int, optional): Number of attempts after which to issue a warning. Defaults
to 1000.
allow_pair_reselection (bool, optional): Whether to request a different district pair when
no valid cut is found. Defaults to ``False``.
Returns:
tuple[list[_Cut], _PopulatedGraph]: The possible cuts and the populated graph they were
obtained from.
Raises:
RuntimeError: If a valid bipartition cannot be found after the specified number of
attempts.
"""
# dict of node_id: population for the nodes in the subgraph
populations = {
node_id: graph_to_split.node_data(node_id)[pop_col]
for node_id in graph_to_split.node_indices
}
possible_cuts = []
if spanning_tree is None:
spanning_tree = spanning_tree_fn(graph_to_split)
h = _PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
# Give parameters and variables descritptive names to make logic clearer
#
# node_repeats specifies the number of extra tries, so we need to add
# one to this value to get the total number of times to try a given
# spanning tree
num_times_to_use_given_spanning_tree = node_repeats + 1
num_times_current_spanning_tree_has_been_tried = 0
attempts = 0
while attempts < max_attempts:
# Try to find balanced edge cuts - note that the find_balanced_edge_cuts_fn()
# typically permutes the spanning tree by selecting a root node for the tree
# at random, which means that each successive call using the same graph
# will effectively use a different spanning tree. Stated differently, the
# spanning_tree variable is in fact a graph, from which a spanning tree can
# be created by picking a root node. This is just an explanation for why it
# makes sense to repeatedly call the same function expecting different results
# each time...
#
# Only forward rootnode_choice_fn when the caller set one: passing None explicitly would
# override a custom finder's own default root chooser.
kwargs = {}
if rootnode_choice_fn is not None:
kwargs["rootnode_choice_fn"] = rootnode_choice_fn
possible_cuts = find_balanced_edge_cuts_fn(h, **kwargs)
# In the case of reversible_recom, we are OK returning None.
if not repeat_until_valid:
return (possible_cuts, h)
# If we have found at least one possible cut, return
num_cuts = len(possible_cuts)
if num_cuts != 0:
return (possible_cuts, h)
# Don't forget to change the documentation if you change this number
if attempts == warn_attempts and not allow_pair_reselection:
warnings.warn(
f"\nFailed to find a balanced cut after {warn_attempts} attempts.\n"
"This can happen if there are relatively few ways to split a graph\n"
"into 2 pieces with appropriate population balance (+/- epsilon)\n"
"for example, if there is a choke-point in the middle of the \n"
"merged districts.\n"
"\n"
"If possible, consider enabling pair reselection within your\n"
"MarkovChain proposal method to allow the algorithm to select\n"
"a different pair of districts for recombination.",
BipartitionWarning,
)
# Record that we have tried to find a solution with the current
# spanning tree, and then see if we should continue with the
# current spanning tree or create a new spanning tree.
#
num_times_current_spanning_tree_has_been_tried += 1
if num_times_current_spanning_tree_has_been_tried == num_times_to_use_given_spanning_tree:
spanning_tree = spanning_tree_fn(graph_to_split)
num_times_current_spanning_tree_has_been_tried = 0
# Regenerate the _PopulatedGraph because the find_balanced_edge_fn() is
# destructive - it can change the internals of a _PopulatedGraph.
#
h = _PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
attempts += 1
# Sanity check to catch possible future bug (a bug we were bitten by in the past)
if num_times_current_spanning_tree_has_been_tried > num_times_to_use_given_spanning_tree:
raise Exception("_get_possible_edge_cuts_and_populated_graph: Should never happen...")
if allow_pair_reselection:
raise ReselectException(
f"Failed to find a balanced cut after {max_attempts} attempts.\n"
f"Selecting a new district pair."
)
raise RuntimeError(f"Could not find a possible cut after {max_attempts} attempts.")
[docs]
@deprecated_parameters(
renamed={
"graph": "subgraph_to_split",
"balance_edge_fn": "find_balanced_edge_cuts_fn",
"one_sided_cut": "single_district_cut",
"choice": "rootnode_choice_fn",
"cut_choice": "cut_choice_fn",
}
)
def bipartition_tree_random_with_num_cuts(
subgraph_to_split: GraphLike,
pop_col: str,
pop_target: int | float,
epsilon: float,
node_repeats: int = 0,
repeat_until_valid: bool = True,
spanning_tree: GraphLike | None = None,
spanning_tree_fn: SpanningTreeFn = random_spanning_tree,
find_balanced_edge_cuts_fn: FindBalancedEdgeCutsFn = find_balanced_edge_cuts_memoization,
single_district_cut: bool = False,
rootnode_choice_fn: Callable[[Sequence[Hashable]], Hashable] | None = None,
max_attempts: int = 100000,
cut_choice_fn: CutChoiceFn | RegionAwareCutChoiceFn | None = None,
*,
rng: random.Random | int | None = None,
) -> tuple[int, AbstractSet[Hashable]]:
"""This is like `bipartition_tree` except it always chooses a random balanced cut.
This function finds a balanced 2 partition of a graph by drawing a spanning tree and finding an
edge to cut that leaves at most an epsilon imbalance between the populations of the parts.
Each spanning tree is searched once plus ``node_repeats`` additional times before a new tree
is drawn.
Builds up a connected subgraph with a connected complement whose population is ``epsilon *
pop_target`` away from ``pop_target``.
Args:
subgraph_to_split (Graph): The graph to partition.
pop_col (str): The node attribute holding the population of each node.
pop_target (int | float): The target population for the returned subset of nodes.
epsilon (float): The allowable deviation from ``pop_target`` (as a percentage of
``pop_target``) for the subgraph's population.
node_repeats (int): Additional roots to try on each spanning tree before drawing a new
tree. Defaults to 0.
repeat_until_valid (bool, optional): Determines whether to keep drawing spanning trees
until a tree with a balanced cut is found. If `True`, a set of nodes will always be
returned; if `False`, an empty set is returned when no valid cut is found on the first
try. Defaults to True.
spanning_tree (Graph | None, optional): The spanning tree for the algorithm to use (used
when the algorithm chooses a new root and for testing). Defaults to None.
spanning_tree_fn (Callable, optional): The random spanning tree algorithm to use if a
spanning tree is not provided. Defaults to `random_spanning_tree`.
find_balanced_edge_cuts_fn (FindBalancedEdgeCutsFn, optional): The algorithm used to find
balanced cut edges. Defaults to `find_balanced_edge_cuts_memoization`. A custom finder
must accept ``single_district_cut``, ``rootnode_choice_fn``, and ``rng`` as keyword
arguments; see :class:`FindBalancedEdgeCutsFn`.
single_district_cut (bool, optional): Passed to the ``find_balanced_edge_cuts_fn``.
Determines whether we are cutting off a single district when partitioning the tree. When
set to False, we check if the node we are cutting and the remaining graph are both
within epsilon of the ideal population. When set to True, we only check if the node we
are cutting is within epsilon of the ideal population. Defaults to False.
rootnode_choice_fn (Callable | None): The random choice function. Can be substituted for
testing. Defaults to the supplied RNG's ``choice``.
max_attempts (int): The max number of attempts that should be made to
bipartition. Defaults to 100,000.
cut_choice_fn (CutChoiceFn | RegionAwareCutChoiceFn | None): Function used to
select a cut. Defaults to the supplied RNG's ``choice``.
rng (random.Random | int | None, optional): Source of randomness. Pass a shared
``Random`` for repeated standalone calls; an integer restarts the stream each call.
Returns:
tuple[int, AbstractSet[Hashable]]: The number of possible cuts and the selected connected
subset of nodes.
"""
# Note: This routine is only used in one place in the GerryChain code.
#
# It is called in tree_proposals.py: reversible_recom().
rng = make_rng(rng)
if cut_choice_fn is None:
cut_choice_fn = _random_choice
num_cuts, node_ids = _internal_bipartition_tree(
subgraph_to_split=subgraph_to_split,
pop_col=pop_col,
pop_target=pop_target,
epsilon=epsilon,
node_repeats=node_repeats,
spanning_tree=spanning_tree,
spanning_tree_fn=spanning_tree_fn,
# region_surcharge: dict[str, float] | None = None,
find_balanced_edge_cuts_fn=find_balanced_edge_cuts_fn,
single_district_cut=single_district_cut,
rootnode_choice_fn=rootnode_choice_fn,
max_attempts=max_attempts,
# warn_attempts: int = 1000,
# allow_pair_reselection: bool = False,
repeat_until_valid=repeat_until_valid,
cut_choice_fn=cut_choice_fn,
rng=rng,
)
return num_cuts, node_ids
[docs]
class BalanceError(Exception):
"""Raised when a balanced cut cannot be found."""
[docs]
class PopulationBalanceError(Exception):
"""Raised when the population of a district is outside the acceptable epsilon range."""