"""
This module provides tools for working with graphs in the context of geographic data.
It extends the functionality of the NetworkX library, adding support for spatial data structures,
geographic projections, and serialization to and from JSON format.
This module is designed to be used in conjunction with geopandas, shapely, and pandas libraries,
facilitating the integration of graph-based algorithms with geographic information systems (GIS).
Note:
This module relies on NetworkX, pandas, and geopandas, which should be installed and
imported as required.
"""
import functools
import json
import warnings
from typing import Any, Iterable, List, Optional, Set, Tuple, Union
import networkx
import pandas as pd
from networkx.classes.function import frozen
from networkx.readwrite import json_graph
from .adjacency import neighbors
from .geo import GeometryError, invalid_geometries, reprojected
from .schema import validate_graph
def json_serialize(input_object: Any) -> Optional[int]:
"""
This function is used to handle one of the common issues that
appears when trying to convert a pandas dataframe into a JSON
serializable object. Specifically, it handles the issue of converting
the pandas int64 to a python int so that JSON can serialize it.
This is specifically used so that we can write graphs out to JSON
files.
:param input_object: The object to be converted
:type input_object: Any (expected to be a pd.Int64Dtype)
:returns: The converted pandas object or None if input is not of type
pd.Int64Dtype
:rtype: Optional[int]
"""
if pd.api.types.is_integer_dtype(input_object): # handle int64
return int(input_object)
return None
[docs]
class Graph(networkx.Graph):
"""
Represents a graph to be partitioned, extending the :class:`networkx.Graph`.
This class includes additional class methods for constructing graphs from shapefiles,
and for saving and loading graphs in JSON format.
"""
def __repr__(self):
return "<Graph [{} nodes, {} edges]>".format(len(self.nodes), len(self.edges))
[docs]
@classmethod
def from_networkx(cls, graph: networkx.Graph) -> "Graph":
"""
Create a Graph instance from a networkx.Graph object.
:param graph: The networkx graph to be converted.
:type graph: networkx.Graph
:returns: The converted graph as an instance of this class.
:rtype: Graph
"""
g = cls(graph)
return g
[docs]
@classmethod
def create(cls, source, **kwargs) -> "Graph":
"""
Smart constructor that dispatches based on the type of ``source``.
- ``str`` (filename ending in .shp, .geojson, .gpkg, etc.) -> ``from_file``
- ``geopandas.GeoDataFrame`` -> ``from_geodataframe``
- ``networkx.Graph`` -> ``from_networkx``
- ``dict`` with key ``edges`` -> ``from_data``
- existing ``Graph`` -> returned as-is
:param source: Input data of various types.
:param kwargs: Forwarded to the underlying constructor.
:returns: A Graph object.
:rtype: Graph
Example::
graph = Graph.create("districts.shp", demand_col="POP")
graph = Graph.create(my_geodataframe, candidate_col="is_candidate")
graph = Graph.create({"edges": [(1,2),(2,3)], "demand": {1:10, 2:20, 3:30},
"candidates": [1, 3]})
"""
if isinstance(source, cls):
return source
if isinstance(source, str):
return cls.from_file(source, **kwargs)
try:
import geopandas as gp
if isinstance(source, gp.GeoDataFrame):
return cls.from_geodataframe(source, **kwargs)
except ImportError:
pass
if isinstance(source, networkx.Graph):
g = cls.from_networkx(source)
if kwargs.get("validate", True):
validate_graph(g, strict=True)
return g
if isinstance(source, dict) and "edges" in source:
return cls.from_data(**source, **kwargs)
raise TypeError(
f"Cannot build a Graph from {type(source).__name__}. "
"Use a filename, GeoDataFrame, networkx.Graph, or dict with 'edges' key."
)
[docs]
@classmethod
def from_data(
cls,
edges,
demand,
candidates=None,
super_candidates=None,
coordinates=None,
area=None,
extra_attributes=None,
validate: bool = True,
) -> "Graph":
"""
Build a Graph from raw data dictionaries. The recommended entry point
for users who don't have a GeoDataFrame.
:param edges: Iterable of (u, v) tuples defining the adjacency.
:param demand: Dict mapping node -> demand value.
:param candidates: Iterable of node IDs that are level-1 facility
candidates (F^1 ⊂ V^1), or dict node -> bool. Defaults to no
candidates.
:param super_candidates: Iterable of node IDs that are level-2
(super-) facility candidates (F^2 ⊂ V^1), or dict node -> bool.
Independent of ``candidates`` — a node can be in either, both, or
neither. Defaults to no super-candidates.
:param coordinates: Dict mapping node -> (x, y). Defaults to (0, 0).
:param area: Dict mapping node -> area. Defaults to 1.0.
:param extra_attributes: Dict of {attribute_name: {node: value}} for
additional node attributes (e.g., custom features).
:returns: A Graph with node attributes ``demand``, ``area``,
``candidate``, ``super_candidate``, ``C_X``, ``C_Y``, plus any
extras.
:rtype: Graph
Example::
graph = Graph.from_data(
edges=[(1, 2), (2, 3), (3, 1)],
demand={1: 100, 2: 150, 3: 80},
candidates=[1, 3], # level-1 candidates
super_candidates=[3], # level-2 candidate (must be V^1)
coordinates={1: (0, 0), 2: (1, 0), 3: (0.5, 1)},
)
"""
g = cls()
g.add_edges_from(edges)
# Ensure all nodes referenced in demand also exist
for node in demand:
if node not in g:
g.add_node(node)
# Normalize candidate sets
def _to_set(value):
if value is None:
return set()
if isinstance(value, dict):
return {n for n, v in value.items() if v}
return set(value)
candidate_set = _to_set(candidates)
super_candidate_set = _to_set(super_candidates)
for node in g.nodes:
g.nodes[node]["demand"] = float(demand.get(node, 0))
g.nodes[node]["area"] = float((area or {}).get(node, 1.0))
g.nodes[node]["candidate"] = 1 if node in candidate_set else 0
g.nodes[node]["super_candidate"] = 1 if node in super_candidate_set else 0
x, y = (coordinates or {}).get(node, (0.0, 0.0))
g.nodes[node]["C_X"] = float(x)
g.nodes[node]["C_Y"] = float(y)
if extra_attributes:
for attr_name, attr_dict in extra_attributes.items():
for node, value in attr_dict.items():
if node in g:
g.nodes[node][attr_name] = value
if validate:
validate_graph(g, strict=True)
return g
[docs]
def set_partition(self, assignment, teams_per_district=None, capacity_level=None):
"""
Write partition state into the graph as node and graph attributes.
After calling this, the graph IS the partition — you can serialize
it as a single artifact.
:param assignment: Dict mapping node -> district ID.
:param teams_per_district: Optional dict district -> team count.
:param capacity_level: Optional max teams per district.
"""
for node, dist_id in assignment.items():
if node in self:
self.nodes[node]["district"] = dist_id
if teams_per_district is not None:
self.graph["teams_per_district"] = dict(teams_per_district)
if capacity_level is not None:
self.graph["capacity_level"] = int(capacity_level)
[docs]
def get_partition_data(self):
"""
Read partition state from graph attributes. Returns None if no
``district`` attribute is present.
:returns: Dict with ``assignment``, ``teams_per_district``,
``capacity_level``, or None if no partition data exists.
"""
sample_node = next(iter(self.nodes), None)
if sample_node is None or "district" not in self.nodes[sample_node]:
return None
assignment = {n: self.nodes[n]["district"] for n in self.nodes
if "district" in self.nodes[n]}
return {
"assignment": assignment,
"teams_per_district": self.graph.get("teams_per_district", {}),
"capacity_level": self.graph.get("capacity_level"),
}
[docs]
@classmethod
def from_json(cls, json_file: str) -> "Graph":
"""
Load a graph from a JSON file in the NetworkX json_graph format.
:param json_file: Path to JSON file.
:type json_file: str
:returns: The loaded graph as an instance of this class.
:rtype: Graph
"""
with open(json_file) as f:
data = json.load(f)
g = json_graph.adjacency_graph(data)
graph = cls.from_networkx(g)
graph.issue_warnings()
return graph
[docs]
def to_json(
self, json_file: str, *, include_geometries_as_geojson: bool = False
) -> None:
"""
Save a graph to a JSON file in the NetworkX json_graph format.
:param json_file: Path to target JSON file.
:type json_file: str
:param bool include_geometry_as_geojson: Whether to include
any :mod:`shapely` geometry objects encountered in the graph's node
attributes as GeoJSON. The default (``False``) behavior is to remove
all geometry objects because they are not serializable. Including the
GeoJSON will result in a much larger JSON file.
:type include_geometries_as_geojson: bool, optional
:returns: None
"""
data = json_graph.adjacency_data(self)
if include_geometries_as_geojson:
convert_geometries_to_geojson(data)
else:
remove_geometries(data)
with open(json_file, "w") as f:
json.dump(data, f, default=json_serialize)
[docs]
@classmethod
def from_file(
cls,
filename: str,
adjacency: str = "rook",
demand_col: str = "demand",
candidate_col: str = "candidate",
super_candidate_col: Optional[str] = "super_candidate",
area_col: Optional[str] = None,
cols_to_add: Optional[List[str]] = None,
reproject: bool = False,
ignore_errors: bool = False,
validate: bool = True,
) -> "Graph":
"""
Create a :class:`Graph` from a shapefile (or GeoPackage, or GeoJSON, or
any other library that :mod:`geopandas` can read. See :meth:`from_geodataframe`
for more details.
:param filename: Path to the shapefile / GeoPackage / GeoJSON / etc.
:type filename: str
:param adjacency: The adjacency type to use ("rook" or "queen"). Default is "rook"
:type adjacency: str, optional
:param cols_to_add: The names of the columns that you want to
add to the graph as node attributes. Default is None.
:type cols_to_add: Optional[List[str]], optional
:param reproject: Whether to reproject to a UTM projection before
creating the graph. Default is False.
:type reproject: bool, optional
:param ignore_errors: Whether to ignore all invalid geometries and try to continue
creating the graph. Default is False.
:type ignore_errors: bool, optional
:returns: The Graph object of the geometries from `filename`.
:rtype: Graph
.. Warning::
This method requires the optional ``geopandas`` dependency.
This method requires ``geopandas``. Install it with:
.. code-block:: console
pip install geopandas
"""
import geopandas as gp
df = gp.read_file(filename)
graph = cls.from_geodataframe(
df,
adjacency=adjacency,
demand_col=demand_col,
candidate_col=candidate_col,
super_candidate_col=super_candidate_col,
area_col=area_col,
cols_to_add=cols_to_add,
reproject=reproject,
ignore_errors=ignore_errors,
validate=validate,
)
graph.graph["crs"] = df.crs.to_json()
return graph
[docs]
@classmethod
def from_geodataframe(
cls,
dataframe: pd.DataFrame,
adjacency: str = "rook",
demand_col: str = "demand",
candidate_col: str = "candidate",
super_candidate_col: Optional[str] = "super_candidate",
area_col: Optional[str] = None,
cols_to_add: Optional[List[str]] = None,
reproject: bool = False,
ignore_errors: bool = False,
crs_override: Optional[Union[str, int]] = None,
validate: bool = True,
) -> "Graph":
"""
Creates the adjacency :class:`Graph` of geometries described by `dataframe`.
The areas of the polygons are included as node attributes (with key `area`).
The shared perimeter of neighboring polygons are included as edge attributes
(with key `shared_perim`).
Nodes corresponding to polygons on the boundary of the union of all the geometries
(e.g., the state, if your dataframe describes VTDs) have a `boundary_node` attribute
(set to `True`) and a `boundary_perim` attribute with the length of this "exterior"
boundary.
By default, areas and lengths are computed in a UTM projection suitable for the
geometries. This prevents the bizarro area and perimeter values that show up when
you accidentally do computations in Longitude-Latitude coordinates. If the user
specifies `reproject=False`, then the areas and lengths will be computed in the
GeoDataFrame's current coordinate reference system. This option is for users who
have a preferred CRS they would like to use.
:param dataframe: The GeoDataFrame to convert.
:type dataframe: :class:`geopandas.GeoDataFrame`
:param adjacency: The adjacency type to use ("rook" or "queen"). Default "rook".
:param demand_col: Name of the column holding demand values. The column
is renamed to ``demand`` on the resulting graph. Default "demand".
:param candidate_col: Name of the column holding level-1 facility
candidate flags (0/1 or True/False). Renamed to ``candidate``.
Default "candidate".
:param super_candidate_col: Name of the column holding level-2
(super-) facility candidate flags. Renamed to ``super_candidate``.
Optional — if the column is absent from the dataframe, every node
gets ``super_candidate=0``. Default "super_candidate".
:param area_col: Name of the column holding pre-computed areas. If None,
areas are computed from geometry. Default None.
:param cols_to_add: Additional columns to copy as node attributes.
:param reproject: Reproject to a UTM projection before creating the graph.
:param ignore_errors: Ignore invalid geometries.
:param crs_override: Override the CRS of the GeoDataFrame.
:param validate: Verify required FalcomChain attributes after construction.
:returns: A Graph with required attributes ``demand`` and ``candidate``,
plus computed ``area``, ``boundary_node``, ``boundary_perim``, and
edge ``shared_perim``.
:rtype: Graph
:raises SchemaValidationError: If ``validate=True`` and required
attributes are missing or invalid.
"""
# Validate geometries before reprojection
if not ignore_errors:
invalid = invalid_geometries(dataframe)
if len(invalid) > 0:
raise GeometryError(
"Invalid geometries at rows {} before "
"reprojection. Consider repairing the affected geometries with "
"`.buffer(0)`, or pass `ignore_errors=True` to attempt to create "
"the graph anyways.".format(invalid)
)
# Project the dataframe to an appropriate UTM projection unless
# explicitly told not to.
if reproject:
df = reprojected(dataframe)
if ignore_errors:
invalid_reproj = invalid_geometries(df)
warnings.warn(f"Invalid geometries after reprojection: {invalid_reproj}")
if len(invalid_reproj) > 0:
raise GeometryError(
"Invalid geometries at rows {} after "
"reprojection. Consider reloading the GeoDataFrame with "
"`reproject=False` or repairing the affected geometries "
"with `.buffer(0)`.".format(invalid_reproj)
)
else:
df = dataframe
# Generate dict of dicts of dicts with shared perimeters according
# to the requested adjacency rule
adjacencies = neighbors(df, adjacency)
graph = cls(adjacencies)
graph.geometry = df.geometry
graph.issue_warnings()
# Add "exterior" perimeters to the boundary nodes
add_boundary_perimeters(graph, df.geometry)
# Add area data to the nodes (computed from geometry, unless overridden)
if area_col is not None and area_col in df.columns:
networkx.set_node_attributes(graph, name="area", values=df[area_col].to_dict())
else:
areas = df.geometry.area.to_dict()
networkx.set_node_attributes(graph, name="area", values=areas)
# Always include the schema-required columns plus user extras
required_cols = []
rename_map = {}
if demand_col in df.columns:
required_cols.append(demand_col)
if demand_col != "demand":
rename_map[demand_col] = "demand"
if candidate_col in df.columns:
required_cols.append(candidate_col)
if candidate_col != "candidate":
rename_map[candidate_col] = "candidate"
# Super-candidate column is optional. If the dataframe doesn't have it,
# every node gets super_candidate=0 below.
if super_candidate_col is not None and super_candidate_col in df.columns:
required_cols.append(super_candidate_col)
if super_candidate_col != "super_candidate":
rename_map[super_candidate_col] = "super_candidate"
all_cols = list(set(required_cols + (cols_to_add or [])))
if all_cols:
graph.add_data(df, columns=all_cols)
# Rename columns to schema-canonical names if needed
for src, dst in rename_map.items():
for n in graph.nodes:
if src in graph.nodes[n]:
graph.nodes[n][dst] = graph.nodes[n].pop(src)
# Default missing super_candidate to 0
for n in graph.nodes:
graph.nodes[n].setdefault("super_candidate", 0)
if crs_override is not None:
df.set_crs(crs_override, inplace=True)
if df.crs is None:
warnings.warn(
"GeoDataFrame has no CRS. Did you forget to set it? "
"If you're sure this is correct, you can ignore this warning. "
"Otherwise, please set the CRS using the `crs_override` parameter. "
"Attempting to proceed without a CRS."
)
graph.graph["crs"] = None
else:
graph.graph["crs"] = df.crs.to_json()
if validate:
validate_graph(graph, strict=True)
return graph
[docs]
def lookup(self, node: Any, field: Any) -> Any:
"""
Lookup a node/field attribute.
:param node: Node to look up.
:type node: Any
:param field: Field to look up.
:type field: Any
:returns: The value of the attribute `field` at `node`.
:rtype: Any
"""
return self.nodes[node][field]
@property
def node_indices(self):
return set(self.nodes)
@property
def edge_indices(self):
return set(self.edges)
[docs]
def add_data(
self, df: pd.DataFrame, columns: Optional[Iterable[str]] = None
) -> None:
"""
Add columns of a DataFrame to a graph as node attributes
by matching the DataFrame's index to node ids.
:param df: Dataframe containing given columns.
:type df: :class:`pandas.DataFrame`
:param columns: List of dataframe column names to add. Default is None.
:type columns: Optional[Iterable[str]], optional
:returns: None
"""
if columns is None:
columns = list(df.columns)
check_dataframe(df[columns])
column_dictionaries = df.to_dict("index")
networkx.set_node_attributes(self, column_dictionaries)
if hasattr(self, "data"):
self.data[columns] = df[columns] # type: ignore
else:
self.data = df[columns]
[docs]
def join(
self,
dataframe: pd.DataFrame,
columns: Optional[List[str]] = None,
left_index: Optional[str] = None,
right_index: Optional[str] = None,
) -> None:
"""
Add data from a dataframe to the graph, matching nodes to rows when
the node's `left_index` attribute equals the row's `right_index` value.
:param dataframe: DataFrame.
:type dataframe: :class:`pandas.DataFrame`
:columns: The columns whose data you wish to add to the graph.
If not provided, all columns are added. Default is None.
:type columns: Optional[List[str]], optional
:left_index: The node attribute used to match nodes to rows.
If not provided, node IDs are used. Default is None.
:type left_index: Optional[str], optional
:right_index: The DataFrame column name to use to match rows
to nodes. If not provided, the DataFrame's index is used. Default is None.
:type right_index: Optional[str], optional
:returns: None
"""
if right_index is not None:
df = dataframe.set_index(right_index)
else:
df = dataframe
if columns is not None:
df = df[columns]
check_dataframe(df)
column_dictionaries = df.to_dict()
if left_index is not None:
ids_to_index = networkx.get_node_attributes(self, left_index)
else:
# When the left_index is node ID, the matching is just
# a redundant {node: node} dictionary
ids_to_index = dict(zip(self.nodes, self.nodes))
node_attributes = {
node_id: {
column: values[index] for column, values in column_dictionaries.items()
}
for node_id, index in ids_to_index.items()
}
networkx.set_node_attributes(self, node_attributes)
@property
def islands(self) -> Set:
"""
:returns: The set of degree-0 nodes.
:rtype: Set
"""
return set(node for node in self if self.degree[node] == 0)
[docs]
def warn_for_islands(self) -> None:
"""
:returns: None
:raises: UserWarning if the graph has any islands (degree-0 nodes).
"""
islands = self.islands
if len(self.islands) > 0:
warnings.warn(
"Found islands (degree-0 nodes). Indices of islands: {}".format(islands)
)
[docs]
def issue_warnings(self) -> None:
"""
:returns: None
:raises: UserWarning if the graph has any red flags (right now, only islands).
"""
self.warn_for_islands()
def add_boundary_perimeters(graph: Graph, geometries: pd.Series) -> None:
"""
Add shared perimeter between nodes and the total geometry boundary.
:param graph: NetworkX graph
:type graph: :class:`Graph`
:param geometries: :class:`geopandas.GeoSeries` containing geometry information.
:type geometries: :class:`pandas.Series`
:returns: The updated graph.
:rtype: Graph
"""
from shapely.ops import unary_union
from shapely.prepared import prep
prepared_boundary = prep(unary_union(geometries).boundary)
boundary_nodes = geometries.boundary.apply(prepared_boundary.intersects)
for node in graph:
graph.nodes[node]["boundary_node"] = bool(boundary_nodes[node])
if boundary_nodes[node]:
total_perimeter = geometries[node].boundary.length
shared_perimeter = sum(
neighbor_data["shared_perim"] for neighbor_data in graph[node].values()
)
boundary_perimeter = total_perimeter - shared_perimeter
graph.nodes[node]["boundary_perim"] = boundary_perimeter
def check_dataframe(df: pd.DataFrame) -> None:
"""
:returns: None
:raises: UserWarning if the dataframe has any NA values.
"""
for column in df.columns:
if sum(df[column].isna()) > 0:
warnings.warn("NA values found in column {}!".format(column))
def remove_geometries(data: networkx.Graph) -> None:
"""
Remove geometry attributes from NetworkX adjacency data object,
because they are not serializable. Mutates the ``data`` object.
Does nothing if no geometry attributes are found.
:param data: an adjacency data object (returned by
:func:`networkx.readwrite.json_graph.adjacency_data`)
:type data: networkx.Graph
:returns: None
"""
for node in data["nodes"]:
bad_keys = []
for key in node:
# having a ``__geo_interface__``` property identifies the object
# as being a ``shapely`` geometry object
if hasattr(node[key], "__geo_interface__"):
bad_keys.append(key)
for key in bad_keys:
del node[key]
def convert_geometries_to_geojson(data: networkx.Graph) -> None:
"""
Convert geometry attributes in a NetworkX adjacency data object
to GeoJSON, so that they can be serialized. Mutates the ``data`` object.
Does nothing if no geometry attributes are found.
:param data: an adjacency data object (returned by
:func:`networkx.readwrite.json_graph.adjacency_data`)
:type data: networkx.Graph
:returns: None
"""
for node in data["nodes"]:
for key in node:
# having a ``__geo_interface__``` property identifies the object
# as being a ``shapely`` geometry object
if hasattr(node[key], "__geo_interface__"):
# The ``__geo_interface__`` property is essentially GeoJSON.
# This is what :func:`geopandas.GeoSeries.to_json` uses under
# the hood.
node[key] = node[key].__geo_interface__
[docs]
class FrozenGraph:
"""
Represents an immutable graph to be partitioned. It is based off :class:`Graph`.
This speeds up chain runs and prevents having to deal with cache invalidation issues.
This class behaves slightly differently than :class:`Graph` or :class:`networkx.Graph`.
Not intended to be a part of the public API.
:ivar graph: The underlying graph.
:type graph: Graph
:ivar size: The number of nodes in the graph.
:type size: int
Note
----
The class uses `__slots__` for improved memory efficiency.
"""
__slots__ = ["graph", "size"]
def __init__(self, graph: Graph) -> None:
"""
Initialize a FrozenGraph from a Graph.
:param graph: The mutable Graph to be converted into an immutable graph
:type graph: Graph
:returns: None
"""
self.graph = networkx.classes.function.freeze(graph)
self.graph.join = frozen
self.graph.add_data = frozen
self.size = len(self.graph)
def __len__(self) -> int:
return self.size
def __getattribute__(self, __name: str) -> Any:
try:
return object.__getattribute__(self, __name)
except AttributeError:
return object.__getattribute__(self.graph, __name)
def __getitem__(self, __name: str) -> Any:
return self.graph[__name]
def __iter__(self) -> Iterable[Any]:
yield from self.node_indices
@functools.lru_cache(16384)
def neighbors(self, n: Any) -> Tuple[Any, ...]:
return tuple(self.graph.neighbors(n))
@functools.cached_property
def node_indices(self) -> Iterable[Any]:
return self.graph.node_indices
@functools.cached_property
def edge_indices(self) -> Iterable[Any]:
return self.graph.edge_indices
@functools.lru_cache(16384)
def degree(self, n: Any) -> int:
return self.graph.degree(n)
@functools.lru_cache(65536)
def lookup(self, node: Any, field: str) -> Any:
return self.graph.nodes[node][field]
def subgraph(self, nodes: Iterable[Any]) -> "FrozenGraph":
return FrozenGraph(self.graph.subgraph(nodes))