compute-regions.py

#!/usr/bin/env python3
import sys, json
from misc import *
from geometry import *
from zipfile import ZipFile
from configparser import ConfigParser
from regions import parse_regions

representatives = {}

regions = parse_regions(sys.argv[2])
bus_stops = {}
block_factor = 0.005

class Blockmap:
	''' Models a map of blocks. Maps each location to a square area. '''
	def __init__(self, default = None):
		from collections import defaultdict
		self.blocks = default or defaultdict(set)
	def __getitem__(self, blockid):
		''' Returns a block for block coordinates. The block is a set that can contain anything. '''
		return self.blocks[blockid]
	def blockpoint(self, point):
		''' Returns a block point for a location. The block point is a coordinate in the blockmap. '''
		return int(point.x / block_factor), int(point.y / block_factor)

def blocks_in_shape(blockmap, shape):
	''' Finds all blocks inside the bounding box of a shape. '''
	min_x = min(point.x for point in shape.points)
	min_y = min(point.y for point in shape.points)
	max_x = max(point.x for point in shape.points)
	max_y = max(point.y for point in shape.points)
	min_blockpoint = blockmap.blockpoint(Location(min_x, min_y))
	max_blockpoint = blockmap.blockpoint(Location(max_x, max_y))
	for x in range(min_blockpoint[0], max_blockpoint[0] + 1):
		for y in range(min_blockpoint[1], max_blockpoint[1] + 1):
			yield blockmap[x, y]

def main():
	with ZipFile(sys.argv[1]) as archive:
		with archive.open('stops.txt') as file:
			for row in read_csv(map(bytes.decode, file)):
				location = Location(float(row['stop_lat']), float(row['stop_lon']))
				reference = row['stop_id']
				bus_stops[reference] = location
	region_shapes = list()
	districts = dict()
	bus_stop_regions = dict()
	blockmap = Blockmap()
	# Find all regions for every block.
	for region in regions.values():
		for block in blocks_in_shape(blockmap, region['shape']):
			set.add(block, region['name'])
	# Find the region every node is in
	for stop_id, stop_position in bus_stops.items():
		for region_name in blockmap[blockmap.blockpoint(stop_position)]:
			region = regions[region_name]
			if region['shape'].contains_point(stop_position):
				bus_stop_regions[stop_id] = region['name']
				break
		else:
			bus_stop_regions[stop_id] = None
	covered = sum(1 if value else 0 for value in bus_stop_regions.values())
	total = len(bus_stops)
	print('%.1f%% bus stops covered.' % (covered * 100 / total), file = sys.stderr)
	json.dump(bus_stop_regions, sys.stdout, indent = 2)

if __name__ == '__main__':
	main()