Modeling spatial structure in binary data with an H3 hexagonal coordinate system

import geopandas as gpd
import h3
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import pymc3 as pm
import shapely

from sklearn.neighbors import BallTree

%matplotlib inline
%config InlineBackend.figure_format='retina'

%load_ext watermark
%watermark -iv

def xy_from_gdf(gdf):
    '''
    Returns Nx2 matrix of X,Y coordinates from a GeoDataFrame
    '''
    return np.stack([gdf.centroid.x, gdf.centroid.y], axis=1)

def lat_lng_to_h3(point, h3_level):
    '''
    Applies H3's geocoding to determine the hexagonal cell
    containing a given point. The h3 level determines the
    size of the hexagonal lattice used.
    '''
    return h3.geo_to_h3(
      point.geometry.centroid.y, point.geometry.centroid.x, h3_level)

def add_geometry(row):
    '''
    Creates a vector feature from the H3 hexagonal coordinates.
    '''
    points = h3.h3_to_geo_boundary(
      row['h3'], True)
    return shapely.geometry.Polygon(points)

def nearest_neighbor_centroid(gdf1, gdf2, k=4):
    '''
    Vectorized operation for identifying the nearest points in gdf2 relative to gdf1.
    '''
    X_proposed, X_base = xy_from_gdf(gdf1), xy_from_gdf(gdf2)        
    nearest = BallTree(X_proposed, leaf_size=2).query(X_base, k=k, return_distance=False)
    return nearest

def sigmoid(x):
    return np.exp(x) / (1 + np.exp(x))

def adjacency_via_buffer(gdf, very_small_distance=0.0003):
    '''
    Uses a spatial buffering and intersection operator to determine
    which features share a boundary in a GeoDataFrame.
    '''
    N = len(gdf)
    W = np.zeros([N,N], dtype=int)

    buffered = gdf.copy()
    buffered.geometry = buffered.buffer(very_small_distance)

    # Find neighbors by buffering and locating non-null overlap
    nearby  = buffered.geometry.apply(lambda x: np.where(buffered.intersection(x).area > 0)[0])
    nearest = nearest_neighbor_centroid(gdf, gdf, k=2)[:, 1:]

    for i, neighbors in enumerate(nearby):
        if len(neighbors) > 1:
            W[i, neighbors] += 1
            W[i,i]          -= 1 # self-neighboring is not allowed
        else:
            W[i, nearest[i]] = 1

    W = W+W.T > 0.
    W = W.astype(int)
    return W

def connect_components(W, geom_series):
    '''
    Iteratively add edges between nodes to the network
    until only a single edge-connected component covers the entire graph. This
    is critical for usage of the CAR model, which can fail if there
    are "islands" disconnected from each other in the network / adjacency matrix.
    '''
    connected = False
    while not connected:

        # Find the largest component and drop it from
        # the list of islands
        G = nx.convert_matrix.from_numpy_matrix(W)
        
        if nx.is_connected(G):
            break
            
        components = list(nx.connected_components(G))
        sizes      = [len(x) for x in components]
        largest    = np.argmax(sizes)
        components.pop(largest)

        # For each island, find the nearest node not on
        # the island and hook it up
        for island in components:
            element_on_island = list(island)[0]
            geom_on_island    = geom_series.iloc[[element_on_island]]
            repeated  = geom_on_island.iloc[np.zeros(geom_series.shape[0])]
            distances = geom_series.geometry.apply(lambda x: geom_on_island.distance(x))
            ordered_by_dist = np.argsort(distances.values[:,0])

            connected_for_island = False
            ctr = 0
            while not connected_for_island:
                proposed = ordered_by_dist[ctr]
                if proposed not in island:
                    W[element_on_island, proposed] = 1
                    W[proposed, element_on_island] = 1
                    print('Match for element {0} is {1}'.format(element_on_island, proposed))
                    connected_for_island = True
                ctr += 1

        G = nx.convert_matrix.from_numpy_matrix(W)
        connected = nx.is_connected(G)
    return W

'''
Census tract shapefile taken from https://opendata.arcgis.com/datasets/f33d847161174e81ad59c9ea9c1f5a00_36.zip
'''
census_tract_path = "./data/Preliminary_2020_Census_Tract/Preliminary_2020_Census_Tract.shp"
tract_gdf = gpd.read_file(census_tract_path)

tract_gdf.head()

base_pregnancy_rate = 11.4 / 1000 # births per thousand people
years = np.arange(2010, 2019)

birth_coords = []
birth_points = []

for year in years:
    
    for i, tract in tract_gdf.iterrows():
        tract_boundary = tract.geometry
        left, lower, right, upper = tract_boundary.bounds
        n_pregnancies = int(base_pregnancy_rate * tract['POP10'])
        
        for j in range(n_pregnancies):
            is_in_bounds = False
            
            while not is_in_bounds:
                coords = np.random.uniform(low=[left,lower],high=[right,upper])
                sample = shapely.geometry.Point(coords)
                is_in_bounds = tract.geometry.contains((sample))
                
            birth_points.append(sample)
            birth_coords.append(list(coords)+[year + np.random.uniform()])
                     
birth_gdf = gpd.GeoDataFrame(data=birth_coords, columns=['lat','lon','year'], geometry=birth_points)
birth_gdf["year_int"] = birth_gdf['year'].astype(int)

birth_gdf.head()

birth_df   = pd.DataFrame(birth_gdf).drop(['geometry', 'year_int'], axis=1)
scales     = birth_df.std()
means      =  birth_df.mean()
zscore_gdf = (birth_df -means)/scales

# coefs are for lat, lon, and year respectively.
true_coefficients = [0.2, 0.2, 0.0]

# this value was chosen by hand to roughly line up with ~12% preterm births, on average
true_intercept    = -2
logits = zscore_gdf.dot(true_coefficients)+true_intercept

birth_gdf['preterm_prob'] = sigmoid(logits)

birth_gdf['preterm'] = np.random.binomial(1, birth_gdf['preterm_prob'])
birth_gdf['preterm'].mean()

fig, axes = plt.subplots(2,2, figsize=(20,22), sharex=True, sharey=True)
axes = axes.ravel()

for i, year in enumerate(years[0:4]):
    birth_gdf.query(f"year_int=={year} & preterm==1").plot(markersize=5,
                                            color='red', ax=axes[i],zorder=1)
    
    birth_gdf.query(f"year_int=={year} & preterm==0").plot(markersize=2,
                                            color='blue', ax=axes[i],zorder=1, alpha=0.5)
    
    tract_gdf.plot(ax=axes[i], facecolor='none', edgecolor='k',zorder=2)
    preterm_frac = birth_gdf.query(f"year_int=={year}")['preterm'].mean()
    axes[i].set_title(f'Simulated births for {year}\n(Preterm fraction: {int(preterm_frac*100)}%)', fontsize=24)
    axes[i].axis('off')
    
plt.tight_layout()

h3_level = 9
gdf = birth_gdf
gdf['h3']      = gdf.apply(lat_lng_to_h3, h3_level=h3_level, axis=1)
gdf['hexagon'] = gdf.apply(add_geometry, axis=1)

hex_only = gdf[['hexagon','h3']].drop_duplicates(subset='h3')
hex_only = gpd.GeoDataFrame(geometry=hex_only['hexagon'], data=hex_only['h3'])
hex_only.sort_values(by='h3', inplace=True)

h3_to_int = {code: integer for integer, code in enumerate(np.sort(gdf['h3'].unique()))}
gdf['h3_int'] = gdf['h3'].apply(lambda x: h3_to_int[x])
print(hex_only.shape)

fig, axes = plt.subplots(1,2, figsize=(10,5), sharey=True)
hex_only.plot(ax=axes[0],edgecolor='k'), axes[0].set_title('H3 spatial cells')
tract_gdf.plot(ax=axes[1],edgecolor='k'), axes[1].set_title('DC census tracts')
plt.tight_layout()

# In geographic coordinate system
very_small_distance = 0.0003
W = adjacency_via_buffer(hex_only, very_small_distance=very_small_distance)
W = connect_components(W, hex_only)

neighbors_per_cell = W.sum(axis=0)

assert neighbors_per_cell.min() >= 1 & neighbors_per_cell.max() <= 6

preterm_counts = gdf.groupby('h3_int')['preterm'].sum()
total_counts   = gdf.groupby('h3_int')['preterm'].count()

n = len(preterm_counts)

with pm.Model() as model:
    # Hyperparameters on spatial correlation, random effect size, and model intercept
    alpha       = pm.Uniform('alpha',lower=-0.95, upper=0.95)
    scale       = pm.HalfNormal('scale', sd=2)
    intercept   = pm.Normal('intercept', sd=3)
    
    # Spatially-smoothing prior on logit of preterm birth probability
    spatial_effect = pm.CAR('spatial_effect', mu=0., W=W, alpha=alpha, tau=1., sparse=True, shape=n)
    likelihood  = pm.Binomial('likelihood', p=pm.math.sigmoid(spatial_effect*scale + intercept),
                             observed=preterm_counts.values, n=total_counts.values)
    
    # Applies Markov chain Monte Carlo to draw from the posterior distribution
    trace = pm.sample(target_accept=0.95, tune=2000)

with model:
    print(pm.summary(trace, var_names=['scale', 'alpha', 'intercept']))

with model:
    pm.plot_trace(trace, var_names=['scale', 'alpha', 'intercept'])

sigma_cutoff = 2

estimated_intercept = trace['intercept'].mean()

hex_only['preterm_fraction'] = (preterm_counts / total_counts).values
hex_only['estimate']   = trace['spatial_effect'].mean(axis=0)
hex_only['stdevs']     = np.abs(trace['spatial_effect'].mean(axis=0) / trace['spatial_effect'].std(axis=0))
hex_only['is_sig']     = hex_only['stdevs'] > sigma_cutoff
hex_only['delta_prob'] = sigmoid(hex_only['estimate']+estimated_intercept) - sigmoid(estimated_intercept)

fig, ax = plt.subplots(1, 1, figsize=(25,25))
hex_only.plot('preterm_fraction', ax=ax)
handle = tract_gdf.plot(facecolor='none', edgecolor='k',zorder=2,ax=ax, alpha=0.2, legend=True);

row_ctr = 0

for i, row in hex_only.iterrows():
    cent = row['geometry'].centroid 
    ax.text(cent.x, cent.y,f'{preterm_counts.iloc[row_ctr]} / {total_counts.iloc[row_ctr]}',
                 ha='center', va='center',fontsize=2, fontweight='bold', color='w')
    row_ctr += 1
    

ax = hex_only.plot('estimate',figsize=(25,25))
tract_gdf.plot(facecolor='none', edgecolor='k',zorder=2,ax=ax, alpha=0.2, legend=True)

for i, row in hex_only.iterrows():
    cent = row['geometry'].centroid 
    if row['is_sig']:
        sig_str = '*'
    else:
        sig_str = ''
    se = row['delta_prob']
    ax.text(cent.x, cent.y, sig_str, ha='center', va='center',fontsize=16, fontweight='bold')

axes[0].set_title('Simulated preterm birth ratio', fontsize=24)
axes[1].set_title('Inferred change in probability of preterm birth', fontsize=24);