statesplot_multicolor.py

statesplot_multicolor.py

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import geopandas as gpd
import fiona
import numpy as np

fiona.drvsupport.supported_drivers["KML"] = "rw"

title = "Title"
override_bucket0_color = "lightgrey"  # set to None to use colormap color
colormap_name = "viridis"

# data buckets from 0 to n-1
# data[0] is always empty because we just put all states into it
# since it's in the background and will be covered if the value is nonzero
data = [[], ["California", "Kentucky"], [], [], ["Ohio"]]

# string labels for the buckets e.g. ["0", "1", "2", "3", "4"]
legend_labels = [str(i) for i in range(len(data))]

# Load detailed US states data from the U.S. Census Bureau
# Source: https://www2.census.gov/geo/tiger/GENZ2022/shp/cb_2022_us_state_20m.zip
states = gpd.read_file(
    "./data/cb_2022_us_state_20m/cb_2022_us_state_20m.kml", driver="KML"
)

states_df = gpd.GeoDataFrame(states)

for i in range(len(states_df)):
    states_df.loc[i, "Name"] = (
        states_df.iloc[i]["Name"]
        .replace("<at><openparen>", "")
        .replace("<closeparen>", "")
    )


def empty_geodataframe():
    return gpd.GeoDataFrame(columns=["id", "name", "geometry"], crs="EPSG:4326")


frames = []
for i, names in enumerate(data):
    if i == 0:
        frames.append(states)
    elif len(names):
        frames.append(states[states["Name"].isin(names)])
    else:
        frames.append(empty_geodataframe())

cmap = plt.colormaps[colormap_name]
normalized_values = np.linspace(0, 1, len(data))
colors = [cmap(val) for val in normalized_values]
if override_bucket0_color is not None:
    colors[0] = override_bucket0_color

# Plotting
fig, ax = plt.subplots(figsize=(10, 8))
for i, frame in enumerate(frames):
    if len(frame):
        edgecolor = "white"
        if i > 0:
            edgecolor = "black"
        frame.plot(ax=ax, color=colors[i], edgecolor=edgecolor)

ax.set_title(title, fontsize=15, y=0.96)
ax.axis("off")
plt.tight_layout()

proxy_artists = []
for color in colors:
    proxy_artists.append(Rectangle((0, 0), 1, 1, fc=color))

ax.legend(
    handles=proxy_artists,
    labels=legend_labels,
)
plt.show()