Convert matplotlib contours into valid (compressed) topojson.

contours.py

import logging
import matplotlib.pyplot as plt
import numpy as np
import os
import scipy.stats as stats
import sys


def read_data(filename):
    """Reads a data file assumed to have at least 2 columns: 1) lat, 2) lng."""
    logging.info("reading %s. . . ", filename)
    # lat, lng, (etc) => lng, lat (via usecols)
    return np.genfromtxt(filename, delimiter=',', skip_header=1, usecols=(1, 0))


def make_contour(data, levels=9, res=64):
    """Data should be a numpy array with 2 cols, lat/lng.
       levels is the desired number of contour levels.
       res is the output resolution, which has a significant perf impact.
            doubling it roughly doubles runtime

       Returns: the output of plt.contour.
    """
    logging.info("Computing contour (%d rows) ... %s", len(data), data.shape)
    kde = stats.kde.gaussian_kde(data.T, bw_method='scott')

    resolution = complex(0, res)
    # Random numbers here are for US Bounding box, including AK and HI
    x_flat = np.r_[-170:-60:resolution]
    y_flat = np.r_[  15: 75:resolution]
    x, y = np.meshgrid(x_flat, y_flat)
    grid_coords = np.append(x.reshape(-1, 1), y.reshape(-1, 1), axis=1)
    z = kde(grid_coords.T)
    z = z.reshape(res, res)

    return plt.contour(x, y, z, levels, linewidths=0.5, colors='k')


def get_bounds(segments):
    """Compute and return bounding box for segments."""
    x0, x1 = sys.maxint, -sys.maxint
    y0, y1 = x0, x1
    for boundary in segments:
        for line in boundary:
            for (x, y) in line:
                if x < x0:
                    x0 = x
                elif x1 < x:
                    x1 = x

                if y < y0:
                    y0 = y
                elif y1 < y:
                    y1 = y

    return x0, x1, y0, y1


def compute_transform(segments, q=1e4):
    """Compute transform needed for compression.

        Quantization factor (q=1e4) means that there are a
            maximum of 10k differentiable values along each dimension.
            Use a higher q-value for higher resolution output.
    """
    x0, x1, y0, y1 = get_bounds(segments)
    return {
        "scale": [
            1 / ((q - 1) / (x1 - x0)),
            1 / ((q - 1) / (y1 - y0))
        ],
        "translate": [x0, y0],
    }

def compress_boundary(transform, boundary):
    """Quantize and delta-encode boundary."""
    kx, ky = transform['scale']
    x0, y0 = transform['translate']

    arcs = []
    for line in boundary:
        arc = []
        dx = dy = 0  # delta-encoding.
        for x, y in line:
            x = int(round((x - x0) / kx - dx))
            y = int(round((y - y0) / ky - dy))
            dx += x
            dy += y
            arc.append([x, y])
        arcs.append(arc)

    return arcs

def convert_boundary(boundary):
    """Convert boundary from numpy arrays to lists for JSON dump.

        Uses the same API as compress_boundary so it's easy
        to turn compression off.
    """
    return [[[x, y] for (x, y) in line] for line in boundary]


def convert_contour(contour, use_transform=True):
    """Save the generated contour to topojson.

       use_transform shrinks output data size by ~80%.
    """
    logging.info("Contour. With transform: %s", use_transform)

    if use_transform:
        transform = compute_transform(contour.allsegs)
        process_boundary = lambda b: compress_boundary(transform, b)
    else:
        transform = None
        process_boundary = convert_boundary

    # Convert contour boundaries to topojson objects.
    num_arcs = 0
    objects, arcs = {}, []
    for level, boundary in enumerate(contour.allsegs):
        new_arcs = process_boundary(boundary)
        arcs.extend(new_arcs)

        objects['level_%d' % level] = {
            "type": "LineString",
            "id": level,
            "arcs": range(num_arcs, len(arcs)),
            "properties": {"LEVEL": level}
        }
        num_arcs += len(new_arcs)

    topo_json = {
        "type": "Topology",
        "transform": transform,
        "objects": objects,
        "arcs": arcs
    }
    if not use_transform:
        del topo_json['transform']

    return topo_json
    

def main():
    """Try computing some contours."""
    data = read_data('data.csv')
    topojson = convert_contour(make_contour(data))
    with open('contours.json', 'w') as f:
        json.dump(topojson, f)
    

if __name__ == '__main__':
    main()