Creates graph network of neighboring U.S. counties

create_us_counties_network.py

import json

import networkx as nx
import requests
from shapely.geometry import shape


def fetch_geojson():
    """Fetches GeoJSON for U.S. Counties.
    Example feature: {'GEO_ID': '0400000US23', 'STATE': '23', 'NAME': 'Maine',
                      'LSAD': '', 'CENSUSAREA': 30842.923}

    :return: GeoJSON object.
    """
    url = 'https://eric.clst.org/assets/wiki/uploads/Stuff/gz_2010_us_050_00_20m.json'
    r = requests.get(url)
    return r.json()


def enumerate_features(geojson, offset=0):
    """Enumerate-like generator for GeoJSON features.

    :param geojson: GeoJSON object.
    :param offset: Offset value for enumerate generator.
    :return: Generator with index and feature pair.
    """
    for idx in range(offset, len(geojson['features'])):
        yield idx, geojson['features'][idx]


def find_intersects(geojson):
    """Finds intersections of features in GeoJSON object. Intersections
    will allow to find neighboring counties.

    :param geojson: GeoJSON object.
    :return: Generator that yields tuple of GeoJSON features.
    """
    for feature_0_idx, feature_0 in enumerate_features(geojson):
        shape_0 = shape(feature_0['geometry'])
        for feature_1_idx, feature_1 in enumerate_features(geojson, feature_0_idx + 1):
            shape_1 = shape(feature_1['geometry'])
            if shape_0.intersects(shape_1):
                yield feature_0, feature_1


def generate_edges(geojson):
    """Generates tuple of edges repesented by GEO_ID of GeoJSON feature.

    :param geojson: GeoJSON object.
    :return: Generator that yields tuple of GeoJSON feature IDs (U.S. County GEO_ID).
    """
    for feature_0, feature_1 in find_intersects(geojson):
        yield [feature_0['properties']['GEO_ID'], feature_1['properties']['GEO_ID']]


def generate_nodes(geojson):
    """Generates GeoJSON feature properties, which includes U.S. County attributes
    (GEO_ID, STATE, NAME, LSAD, CENSUSAREA).

    :param geojson: GeoJSON object.
    :return: Generator that yields GeoJSON feature properties.
    """
    for feature_idx, feature in enumerate_features(geojson):
        yield feature['properties']


def create_network_json(geojson):
    """Creates JSON object from GeoJSON object representing edges and nodes of
    U.S. Counties. Nodes represent a County, and edges represent neighboring
    counties.

    :param geojson: GeoJSON object.
    :return: Dictionary of edges and nodes of U.S. Counties.
    """
    return {
        'edges': list(generate_edges(geojson)),
        'nodes': list(generate_nodes(geojson)),
    }


if __name__ == '__main__':
    """Create network graph and JSON to represent U.S. Counties
    """
    geojson = fetch_geojson()
    network_json = create_network_json(geojson)

    network_json_fp = open('data/counties_network.json', 'w')
    json.dump(network_json, network_json_fp)

    G = nx.Graph()
    for node in network_json['nodes']:
        geo_id = node.pop('GEO_ID')
        G.add_node(geo_id, **node)
    for feature_0, feature_1 in network_json['edges']:
        G.add_edge(feature_0, feature_1)