Go from lat lon pairs to a geographical heatmap

latlon_pairs_to_heatmap.py

# Python tools
from matplotlib import pyplot as plt
import pandas as pd
import math
import numpy as np
from collections import defaultdict
from functools import partial
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib as mpl

# Geometry, shapefiles and projections
import fiona
from shapely.geometry import shape
from shapely.geometry import Point
import pyproj
from descartes import PolygonPatch
from shapely.ops import transform

# Shape files for TTWAs, GB and NI
gb_filename = 'Countries_December_2016_Ultra_Generalised_Clipped_Boundaries_in_Great_Britain.shp'
ni_filename = 'OSNI_Open_Data_Largescale_Boundaries__Country_2016.shp'

# Generate a function to create a UK East/North point from Lon/Lat
wgs84 = pyproj.Proj(init = 'epsg:4326')
ukgrid = pyproj.Proj(init = 'epsg:27700')
EnPoint = lambda lon, lat : Point(*pyproj.transform(wgs84, ukgrid, lon, lat))

project = partial(
    pyproj.transform,
    pyproj.Proj(init='epsg:29902'), # source coordinate system
    pyproj.Proj(init='epsg:27700')) # destination coordinate system

# Calculate great circle distance between lat/lon points
def haversine(lon1, lat1, lon2, lat2, in_radians=False):
    # convert decimal degrees to radians 
    if not in_radians:
        lon1, lat1, lon2, lat2 = map(math.radians, [lon1, lat1, lon2, lat2])
    # haversine formula 
    dlon = lon2 - lon1 
    dlat = lat2 - lat1 
    a = math.sin(dlat/2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon/2)**2
    c = 2 * math.asin(math.sqrt(a)) 
    r = 6371 # Radius of earth in kilometers. Use 3956 for miles
    return c * r    

# Calculate great circle distance between lat/lon points
def latlonline_to_points(lon1, lat1, lon2, lat2, max_distance = 1e6, min_distance=0):
    
    lon1, lat1, lon2, lat2 = map(math.radians, [lon1, lat1, lon2, lat2])
    
    r = 6371
    total_distance = haversine(lon1, lat1, lon2, lat2, in_radians=True)
    delta = total_distance / r
    
    cc1 = np.cos(lon1)*np.cos(lat1)
    cs1 = np.cos(lon1)*np.sin(lat1)
    cc2 = np.cos(lon2)*np.cos(lat2)
    cs2 = np.cos(lon2)*np.sin(lat2)
    s1 = np.sin(lon1)
    s2 = np.sin(lon2)

    points = []
    if total_distance > max_distance:
        return points
    if total_distance < min_distance:
        return points
    
    for f in np.arange(0,1.0001,0.0001):
    #total = 0
    #while total < total_distance:
    #    f = total/total_distance
    #    total += 0.01
        
        a = np.sin((1-f)*delta) / np.sin(delta)
        b = np.sin(f*delta) / np.sin(delta)
        x = a*cc1 + b*cc2
        y = a*cs1 + b*cs2
        z = a*s1 + b*s2    
        lon = np.arctan2(z, np.sqrt(x**2 + y**2))
        lat = np.arctan2(y, x)
        lon, lat = map(math.degrees, [lon, lat])
        lon, lat = pyproj.transform(wgs84, ukgrid, lon, lat)
        points.append((lon,lat))
    return points  

def gen_plot(points,ax,title,cmap):
    #fig,ax = plt.subplots(figsize=(8,16))

    # Lists of x,y bounds in order to set fig-ax lims
    xs = []
    ys = []
    # Plot GB
    with fiona.open(gb_filename) as gb:
        for country in gb:
            s = shape(country['geometry'])
            p = PolygonPatch(s, fc="white", ec="grey")
            ax.add_patch(p)
            xs += [s.bounds[0],s.bounds[2]]
            ys += [s.bounds[1],s.bounds[3]]

    # Plot NI, correcting for different east/northing zone
    with fiona.open(ni_filename) as ni:
         for country in ni:
            pyproj.Proj("+proj=utm +zone=23K, +south +ellps=WGS84 +datum=WGS84 +units=m +no_defs")
            s = shape(country['geometry'])
            s = transform(project, s)
            p = PolygonPatch(s, fc="white", ec="grey")
            ax.add_patch(p)
            xs += [s.bounds[0],s.bounds[2]]
            ys += [s.bounds[1],s.bounds[3]]

    ax.set_xlim(min(xs),max(xs))
    ax.set_ylim(min(ys),max(ys))

    _points = [(x,y) for x,y in points 
               if not (pd.isnull(y) or np.isinf(x))]

    xpoints = [x for x,y in _points]
    ypoints = [y for x,y in _points]
    weights, bins_x, bins_y = np.histogram2d(x=xpoints,y=ypoints, bins=60)
    #weights = np.log10(weights)
    new_weights = []
    for row in weights.T:
        new_row = []
        for x in row:
            if x != 0:            
                x = np.log10(x)
            else:
                x = 0
            new_row.append(x)
        new_weights.append(new_row)
    qcs = ax.contourf(bins_x[0:-1],bins_y[0:-1],new_weights,zorder=10,alpha=0.75,cmap=cmap,
                      levels=[x for x in np.arange(2.75,10,0.25)],vmax=6.5)
    
    #ax.set_facecolor('black')
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
    
    # Shorten title
    _title = []
    n = 0
    for c in title:
        n += 1
        if n > 14 and c == " ":
            _title.append("\n")
            n = 0
        else:
            _title.append(c)
    title = "".join(_title)
    ax.text(x=min(xs)+20000,y=0.95*max(ys),s=title,
            fontdict=dict(color='black',size=14),va="center")
        
    return ax

def gen_white(points,cmap,sample=None):
    fig = plt.figure(figsize=(20,17))
    gs = gridspec.GridSpec(3, 4, hspace=0.02,
                           wspace=0.02, height_ratios=[1, 1, 0.18]) 
    i = 0
    for (k,v),g in zip(points,gs):
        # Take a random sample 
        if sample is not None:
            if sample < len(v):
                v = random.sample(v,sample)
        print(k,len(v))
        ax = plt.subplot(g)
        gen_plot(v, ax, k, cmap)
        i += 1

    # Bar along the bottom
    ax = plt.subplot(gs[i:i+4])
    ax_cbar = inset_axes(ax, width="48%", height="25%", loc=8)
    ax_cbar.set_axis_off()
    norm = mpl.colors.Normalize(vmin=min(ax_cbar.get_xlim()), vmax=max(ax_cbar.get_xlim()))
    cb = mpl.colorbar.ColorbarBase(ax_cbar, cmap=cmap, norm=norm, orientation='horizontal',)
    # Add label and arrow to colourbar
    ax.text(s="Increasing directional collaboration density",
            x=0.5,y=0.75,va="center",ha="center",fontdict={"size":20})
    ax.arrow(0.3,0.55,0.4,0,color="black",head_width=0.1,head_length=0.03)
    ax.set_axis_off()
    plt.savefig("something.png")

# Lists of x,y bounds in order to set fig-ax lims
xs = []
ys = []
# Plot GB
with fiona.open(gb_filename) as gb:
    for country in gb:
        s = shape(country['geometry'])
        xs += [s.bounds[0],s.bounds[2]]
        ys += [s.bounds[1],s.bounds[3]]

# Plot NI, correcting for different east/northing zone
with fiona.open(ni_filename) as ni:
     for country in ni:
        pyproj.Proj("+proj=utm +zone=23K, +south +ellps=WGS84 +datum=WGS84 +units=m +no_defs")
        s = shape(country['geometry'])
        s = transform(project, s)
        xs += [s.bounds[0],s.bounds[2]]
        ys += [s.bounds[1],s.bounds[3]]

# Generate points data
# Assumes dataframe has columns `lon`,`lat` `plon`,`plat` (which are the 'from' and 'to' coordinates)
# and has the column ce_label
points = {}
for _,row in df.iterrows():
    try:
        x,y = pyproj.transform(wgs84, ukgrid, row.lon, row.lat)
    except:
        continue
    if not((x > min(xs) and x < max(xs)) 
           and (y > min(ys) and y < max(ys))):
        continue
        
    x,y = pyproj.transform(wgs84, ukgrid, row.plon, row.plat)
    if not((x > min(xs) and x < max(xs)) 
           and (y > min(ys) and y < max(ys))):
        continue
    
    if row.ce_label not in points:
        points[row.ce_label] = []
    points[row.ce_label] += latlonline_to_points(row.lon,row.lat,row.plon,row.plat)

points = [(k,v) for k,v in points.items()]
gen_white(points,'gist_heat_r')