Source code for gerrychain.updaters.county_splits
from __future__ import annotations
from collections.abc import Callable, Hashable, Sequence
from enum import Enum
from typing import TYPE_CHECKING, NamedTuple
from .._deprecated import deprecated_alias, deprecated_parameters
if TYPE_CHECKING:
from ..partition.partition import Partition
[docs]
class CountySplit(Enum):
"""
Enum to track county splits in a partition.
Attributes:
NOT_SPLIT (CountySplit): The county is not split.
NEW_SPLIT (CountySplit): The county became split after the initial partition.
OLD_SPLIT (CountySplit): The county has remained split since the initial partition.
"""
NOT_SPLIT = 0
NEW_SPLIT = 1
OLD_SPLIT = 2
class CountyInfo(NamedTuple):
"""Information about how a county intersects a partition."""
split: CountySplit
nodes: list[Hashable]
contains: set[Hashable]
[docs]
def county_splits(
partition_name: str, county_field_name: str
) -> Callable[[Partition], dict[Hashable, CountyInfo]]:
"""An updater for tracking the number of counties that are split in a partition.
Args:
partition_name (str): Name that the Partition instance will store.
county_field_name (str): Name of county ID field on the graph.
Returns:
Callable: The tracked data is a ``dict[Hashable, CountyInfo]`` keyed on the county ID. Each
``CountyInfo`` is a named tuple with fields `split` (a CountySplit enum), `nodes` (a
list of node IDs in the county), and `contains` (a ``set[Hashable]`` of the assignment
IDs the county intersects).
"""
def _get_county_splits(partition: Partition) -> dict[Hashable, CountyInfo]:
return compute_county_splits(partition, county_field_name, partition_name)
return _get_county_splits
def compute_county_splits(
partition: Partition, county_field: str, partition_field: str
) -> dict[Hashable, CountyInfo]:
"""Computes the number of counties that are split in a partition.
Args:
partition (Partition): The partition object to compute
county splits for.
county_field (str): Name of county ID field on the graph.
partition_field (str): Name of the attribute in the graph that stores the partition
information. The county split information will be computed with respect to this
division of the graph.
Returns:
dict[Hashable, CountyInfo]: Information on how counties changed between the parent and
child partitions. If there is no parent partition, only ``OLD_SPLIT`` and
``NOT_SPLIT`` are used. Split history is carried only through the parent, so a county
that stops being split and later becomes split again is reported as ``NEW_SPLIT``.
"""
# Create the initial county data containers.
if not partition.parent:
county_dict: dict[Hashable, CountyInfo] = {}
for node_id in partition.graph.node_indices:
# First figure get current status of the county's information
county = partition.graph.node_data(node_id)[county_field]
if county in county_dict:
split, nodes, seen = county_dict[county]
else:
split, nodes, seen = CountySplit.NOT_SPLIT, [], set()
# Now update "nodes" and "seen" with this node_id and the part (district) from
# partition's assignment.
nodes.append(node_id)
seen.update(set([partition.assignment.mapping[node_id]]))
# lastly, if we have "seen" more than one part (district), then the county is split
# across parts.
if len(seen) > 1:
split = CountySplit.OLD_SPLIT
# update the county_dict with new information
county_dict[county] = CountyInfo(split, nodes, seen)
return county_dict
new_county_dict: dict[Hashable, CountyInfo] = {}
parent = partition.parent
for county, county_info in parent[partition_field].items():
seen = set(partition.assignment.mapping[node_id] for node_id in county_info.nodes)
split = CountySplit.NOT_SPLIT
if len(seen) > 1:
if county_info.split != CountySplit.OLD_SPLIT:
split = CountySplit.NEW_SPLIT
else:
split = CountySplit.OLD_SPLIT
new_county_dict[county] = CountyInfo(split, county_info.nodes, seen)
return new_county_dict
[docs]
@deprecated_parameters(renamed={"reg_attr_lst": "region_attr_lst"})
def tally_region_splits(region_attr_lst: Sequence[str]) -> Callable[[Partition], dict[str, int]]:
"""A naive updater for tallying the number of times a region attribute is split.
Here "region" is the generic term for an administrative unit you would prefer not to split
(e.g. a county, municipality, or other locality). Each entry in ``region_attr_lst`` is the
name of a node attribute holding the region label; the ``county`` and locality updaters are
just specific instances of this region concept.
Args:
region_attr_lst (Sequence[str]): Region attribute names to tally splits for.
Returns:
Callable: A function that takes a partition and returns a dictionary which maps the region
name to the number of times that it is split in a a particular partition.
"""
def _get_splits(partition: Partition) -> dict[str, int]:
if "cut_edges" not in partition.updaters:
raise ValueError("The cut_edges updater must be attached to the partition")
return {
region_attr: total_region_splits(partition, region_attr)
for region_attr in region_attr_lst
}
return _get_splits
def total_region_splits(partition: Partition, region_attr: str) -> int:
"""Computes the total number of times that region_attr is split in the partition.
Args:
partition (Partition): The partition object to compute region splits for.
region_attr (str): The name of the region attribute to compute splits for. This should be a
node attribute on the graph.
"""
all_region_names = set(
partition.graph.node_data(node_id)[region_attr] for node_id in partition.graph.node_indices
)
split = {name: 0 for name in all_region_names}
# Require that the cut_edges updater is attached to the partition
for node1, node2 in partition["cut_edges"]:
if (
partition.assignment[node1] != partition.assignment[node2]
and partition.graph.node_data(node1)[region_attr]
== partition.graph.node_data(node2)[region_attr]
):
split[partition.graph.node_data(node1)[region_attr]] += 1
split[partition.graph.node_data(node2)[region_attr]] += 1
return sum(1 for value in split.values() if value > 0)
total_reg_splits = deprecated_alias(
"gerrychain.updaters.county_splits.total_reg_splits",
"total_region_splits",
total_region_splits,
)