UFOs in the US

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from mpl_toolkits.basemap import Basemap

import matplotlib
from matplotlib.colors import rgb2hex
from matplotlib.patches import Polygon

# Read the UFOs dataset, filter by US
ufos = pd.read_csv("ufo_sightings_scrubbed.csv", low_memory = False)
ufos = ufos[["datetime", "country", "state", "latitude", "longitude "]]
ufos.columns = ["datetime", "country", "state", "latitude", "longitude"]
ufos = ufos[ufos.country == "us"]

# Map state abbreviations to names
state_names = pd.read_csv("states.csv")
state_names.abbreviation = state_names.abbreviation.str.lower()
state_names_dict = {state.abbreviation: state["name"] for index, state in state_names.iterrows()}
ufos.state.replace(state_names_dict, inplace = True)

# Get total UFO sightings by state
num_sightings_by_state = ufos.groupby("state").count().datetime

# Set up the map
fig = plt.figure(figsize = (15, 10))
m = Basemap(projection = "merc", llcrnrlon = -130, llcrnrlat = 23, urcrnrlon = -64, urcrnrlat = 50)
us_info = m.readshapefile("st99_d00", "states", drawbounds = True)
colors = {}
state_names = []
cmap = plt.cm.Greens
vmin = num_sightings_by_state.min()
vmax = num_sightings_by_state.max()

# Calculate colors
for shape_dict in m.states_info:
    state_name = shape_dict["NAME"]
    sightings = num_sightings_by_state[num_sightings_by_state.index == state_name][0]
    colors[state_name] = cmap(np.sqrt((sightings - vmin) / (vmax - vmin)))[:3]
    state_names.append(state_name)

# Show polygons with their respective colors
ax = plt.gca()
for nshape, seg in enumerate(m.states):
    color = rgb2hex(colors[state_names[nshape]])
    poly = Polygon(seg, facecolor = color, edgecolor = color)
    ax.add_patch(poly)

# Add title and colorbar
plt.title("UFO Sightings by State (Contiguous US)")
colorbar_ax = fig.add_axes([0.95, 0.15, 0.02, 0.7])
matplotlib.colorbar.ColorbarBase(colorbar_ax, cmap = cmap, norm = matplotlib.colors.Normalize(vmin, vmax))
plt.show()