Source code for gerrychain.updaters.compactness

from __future__ import annotations

import collections
from collections.abc import Hashable
from typing import TYPE_CHECKING

from .._deprecated import legacy_flips
from .cut_edges import on_edge_flow
from .flows import on_flow

if TYPE_CHECKING:
    from ..partition.partition import Partition

flips = legacy_flips




def boundary_nodes(partition: Partition, alias: str = "boundary_nodes") -> set[Hashable]:
    """Return set of nodes in the partition that are on the boundary.

    Args:
        partition (Partition): A partition of a Graph
        alias (str, optional): The name of the attribute that the boundary nodes are stored under.
            Default is 'boundary_nodes'.

    Returns:
        set: The set of nodes in the partition that are on the boundary.
    """

    # Note that the "alias" parameter is used as the attribute name
    # on the partition - using this "alias" you can retrieve the
    # the data stored by an updater that uses this routine...

    if partition.parent:
        return partition.parent[alias]
    else:
        result = {
            node_id
            for node_id in partition.graph.node_indices
            if partition.graph.node_data(node_id)["boundary_node"]
        }
        return result



def initialize_exterior_boundaries_as_a_set(partition: Partition) -> dict[Hashable, set[Hashable]]:
    """Return exterior boundary nodes for each part in the partition.

    Args:
        partition (Partition): A partition of a Graph

    Returns:
        dict[Hashable, set[Hashable]]: Parts mapped to their exterior boundary nodes.
    """
    part_boundaries: collections.defaultdict[Hashable, set[Hashable]] = collections.defaultdict(set)
    for node in partition["boundary_nodes"]:
        part_boundaries[partition.assignment.mapping[node]].add(node)

    return part_boundaries




@on_flow(initialize_exterior_boundaries_as_a_set, alias="exterior_boundaries_as_a_set")
def exterior_boundaries_as_a_set(
    partition: Partition,
    previous: set[Hashable],
    inflow: set[Hashable],
    outflow: set[Hashable],
) -> set[Hashable]:
    """Updater function that responds to the flow of nodes between different partitions.

    Args:
        partition (Partition): A partition of a Graph
        previous (set[Hashable]): Previous exterior boundary nodes for a fixed part.
        inflow (set[Hashable]): Nodes that flowed into the part.
        outflow (set[Hashable]): Nodes that flowed out of the part.

    Returns:
        set[Hashable]: The updated exterior boundary nodes for the part.
    """
    # Compute the new set of boundary nodes for the partition.
    #
    # The term, (inflow & graph_boundary), computes new nodes that are boundary nodes.
    #
    # the term, (previous | (inflow & graph_boundary)), adds those new boundary nodes to the
    # set of previous boundary nodes.
    #
    # Then all you need to do is subtract all of the nodes in the outflow to remove any of those
    # that happen to be boundary nodes...

    graph_boundary = partition["boundary_nodes"]
    return (previous | (inflow & graph_boundary)) - outflow



def initialize_exterior_boundaries(partition: Partition) -> dict[Hashable, float]:
    """Return A dictionary mapping each part of a partition to the total perimeter of the boundary.

    Args:
        partition (Partition): A partition of a Graph

    Returns:
        dict[Hashable, float]: Parts mapped to their exterior boundary perimeter.
    """
    graph_boundary = partition["boundary_nodes"]
    boundaries: collections.defaultdict[Hashable, float] = collections.defaultdict(float)
    for node in graph_boundary:
        part = partition.assignment.mapping[node]
        boundaries[part] += partition.graph.node_data(node)["boundary_perim"]
    return boundaries




@on_flow(initialize_exterior_boundaries, alias="exterior_boundaries")
def exterior_boundaries(
    partition: Partition,
    previous: float,
    inflow: set[Hashable],
    outflow: set[Hashable],
) -> float:
    """Computes the total perimeter of the boundary nodes in each part of the partition.

    Args:
        partition (Partition): A partition of a Graph
        previous (float): Previous exterior boundary perimeter for the part.
        inflow (set[Hashable]): Nodes that flowed into the part.
        outflow (set[Hashable]): Nodes that flowed out of the part.

    Returns:
        float: The updated exterior boundary perimeter for the part.
    """
    graph_boundary = partition["boundary_nodes"]
    added_perimeter = sum(
        partition.graph.node_data(node)["boundary_perim"] for node in inflow & graph_boundary
    )
    removed_perimeter = sum(
        partition.graph.node_data(node)["boundary_perim"] for node in outflow & graph_boundary
    )
    return previous + added_perimeter - removed_perimeter



def initialize_interior_boundaries(partition: Partition) -> dict[Hashable, float]:
    """Return A dictionary mapping each part of a partition to the total perimeter the given part.

    Args:
        partition (Partition): A partition of a Graph

    Returns:
        dict[Hashable, float]: Parts mapped to the perimeter shared with other parts.
    """

    # RustworkX Note:
    #
    # The old NX code did not distinguish between edges and edge_ids - they were one
    # and the same.  However, in RX an edge is a tuple and an edge_id is an integer.
    # The edges stored in partition["cut_edges_by_part"] are edges (tuples), so
    # we need to get the edge_id for each edge in order to access the data for the edge.

    # Get edge_ids for each edge (tuple)
    edge_ids_for_part = {
        part: [
            partition.graph.get_edge_id_from_edge(edge)
            for edge in partition["cut_edges_by_part"][part]
        ]
        for part in partition.parts
    }

    # Compute length of the shared perimeter of each part
    shared_perimeters_for_part = {
        part: float(
            sum(
                partition.graph.edge_data(edge_id)["shared_perim"]
                for edge_id in edge_ids_for_part[part]
            )
        )
        for part in partition.parts
    }

    return shared_perimeters_for_part




@on_edge_flow(initialize_interior_boundaries, alias="interior_boundaries")
def interior_boundaries(
    partition: Partition,
    previous: float,
    new_edges: set[tuple[int, int]],
    old_edges: set[tuple[int, int]],
) -> float:
    """Computes the total perimeter of the shared boundary between different parts of the partition.

    Args:
        partition (Partition): A partition of a Graph
        previous (float): Previous interior boundary perimeter for the part.
        new_edges (set[tuple[int, int]]): Edges that flowed into the part.
        old_edges (set[tuple[int, int]]): Edges that flowed out of the part.

    Returns:
        float: The updated interior boundary perimeter for the part.
    """

    added_perimeter = sum(
        partition.graph.edge_data(partition.graph.get_edge_id_from_edge(edge))["shared_perim"]
        for edge in new_edges
    )
    removed_perimeter = sum(
        partition.graph.edge_data(partition.graph.get_edge_id_from_edge(edge))["shared_perim"]
        for edge in old_edges
    )
    return previous + added_perimeter - removed_perimeter



def perimeter_of_part(partition: Partition, part: Hashable) -> float:
    """Totals up the perimeter of the part in the partition.

    .. Warning::

        Requires that 'boundary_perim' be a node attribute, 'shared_perim' be an edge attribute,
        'cut_edges' be an updater, and 'exterior_boundaries' be an updater.

    Args:
        partition (Partition): A partition of a Graph
        part (Hashable): The part whose perimeter to compute.

    Returns:
        float: The perimeter of the desired part.
    """

    # frm: TODO: Refactoring:   Add code to enforce the warning in the docstring above

    exterior_perimeter = partition["exterior_boundaries"][part]
    interior_perimeter = partition["interior_boundaries"][part]

    return exterior_perimeter + interior_perimeter




def perimeter(partition: Partition) -> dict[Hashable, float]:
    """Computes the perimeter of each part in the partition.

    Args:
        partition (Partition): A partition of a Graph

    Returns:
        dict[Hashable, float]: Parts mapped to their perimeter.
    """
    return {part: perimeter_of_part(partition, part) for part in partition.parts}