Interpolate Bangkok Pollution with IDW (and Voronoi)

bkk_aqi_data.csv

bkk_base_map.png

script.py

#!/usr/bin/env python3

import csv
from PIL import Image, ImageDraw, ImageFont


BORDER = (0, 0, 0, 192)


def read_data(filename):
    reader = csv.reader(open(filename))
    next(reader)
    return {tuple(map(int, rs[:2])): (int(rs[2]), rs[3]) for rs in reader}


def partition(value, start, end, left, right):
    if value <= start:
        return left
    if value > end:
        return right
    return left + (right-left) * (value-start) // (end-start)


def get_color(value):
    alpha = 160
    if value <= 95:
        red = partition(value, 0, 63, 95, 255)
    else:
        red = partition(value, 127, 255, 255, 95)
    green = partition(value, 63, 127, 255, 0)
    blue = partition(value, 127, 191, 0, 63)
    return (red, green, blue, alpha)


def draw_text(data, image):
    font = ImageFont.truetype('Pillow/Tests/fonts/DejaVuSans.ttf', 18)
    black = (0, 0, 0, 255)
    drawer = ImageDraw.Draw(image)
    for (x, y), (aqi, note) in data.items():
        w, h = drawer.textsize(note, font=font)
        pos = (x-w/2, y+h/2)
        drawer.text(pos, note, font=font, fill=black)
        w, h = drawer.textsize(str(aqi), font=font)
        pos = (x-w/2, y-h/2)
        drawer.text(pos, str(aqi), font=font, fill=black)


def dist(p, q):
    return ( (p[0]-q[0])**2 + (p[1]-q[1])**2 )**0.5


def voronoi_cell(x, y, data):
    cells = sorted((dist((x,y), pos), pos) for pos, _ in data.items())
    if abs(cells[0][0] - cells[1][0]) <= 1:
        return None
    return cells[0][1]


def draw_voronoi(data, image):
    pixels = image.load()
    for x in range(image.size[0]):
        for y in range(image.size[1]):
            cell = voronoi_cell(x, y, data)
            if cell is None:
                pixels[x,y] = BORDER
            else:
                pixels[x,y] = get_color(data[cell][0])


def render_voronoi(data, background):
    overlay = Image.new('RGBA', background.size)
    draw_voronoi(data, overlay)
    draw_text(data, overlay)
    Image.alpha_composite(background, overlay).save('voronoi.png')


def get_weights(x, y, data):
    return {pos: 1/dist((x, y), pos)**2 for pos in data.keys()}


def idw(x, y, data):
    if (x, y) in data:
        return data[x,y][0]
    weights = get_weights(x, y, data)
    return sum(weights[pos]*aqi for pos, (aqi, _) in data.items()) / sum(weights.values())


def draw_idw(data, image):
    pixels = image.load()
    for x in range(image.size[0]):
        for y in range(image.size[1]):
            value = int(idw(x, y, data))
            if value % 10 == 0:
                pixels[x,y] = BORDER
            else:
                pixels[x,y] = get_color(value)


def render_idw(data, background):
    overlay = Image.new('RGBA', background.size)
    draw_idw(data, overlay)
    draw_text(data, overlay)
    Image.alpha_composite(background, overlay).save('idw.png')


def main():
    data = read_data('bkk_aqi_data.csv')
    background = Image.open('bkk_base_map.png').convert('RGBA')
    render_voronoi(data, background)
    render_idw(data, background)


if __name__ == '__main__':
    main()