compute-regions.py

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

class Blockmap:
	'''
	Models a map of blocks. Maps each location to a square area.
	'''
	block_factor = 0.005
	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.
		'''
		block_coordinate = lambda x: int(x / self.block_factor)
		return block_coordinate(point.x), block_coordinate(point.y)

def minmax(data):
	'''
	Computes the minimum and maximum values in one-pass using only
	1.5*len(data) comparisons
	'''
	import itertools
	it = iter(data)
	try:
		lo = hi = next(it)
	except StopIteration:
		raise ValueError('minmax() arg is an empty sequence')
	for x, y in itertools.zip_longest(it, it, fillvalue = lo):
		if x > y:
			x, y = y, x
		lo = min(x, lo)
		hi = max(y, hi)
	return lo, hi

def blocks_in_shape(blockmap, shape):
	'''
	Finds all blocks inside the bounding box of a shape.
	'''
	from itertools import product
	min_x, max_x = minmax(point.x for point in shape.points)
	min_y, max_y = minmax(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))
	range_x = range(min_blockpoint[0], max_blockpoint[0] + 1)
	range_y = range(min_blockpoint[1], max_blockpoint[1] + 1)
	yield from (blockmap[x, y] for x, y in product(range_x, range_y))

def location_from_row(row):
	return Location(float(row['stop_lat']), float(row['stop_lon']))

def read_bus_stops(archive_path):
	bus_stops = dict()
	with ZipFile(archive_path) as archive:
		with archive.open('stops.txt') as file:
			return {
				row['stop_id']: location_from_row(row)
				for row in read_csv(map(bytes.decode, file))
			}

def create_blockmap(regions):
	'''
	Creates a blockmap of regions
	'''
	blockmap = Blockmap()
	for region in regions.values():
		for block in blocks_in_shape(blockmap, region['shape']):
			set.add(block, region['name'])
	return blockmap

def get_args():
	from argparse import ArgumentParser
	parser = ArgumentParser()
	parser.add_argument('gtfs_zip')
	parser.add_argument('mapfile')
	return parser.parse_args()

def compute_bus_stop_regions(regions, bus_stops):
	# Find the region every node is in
	blockmap = create_blockmap(regions)
	bus_stop_regions = dict()
	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
	return bus_stop_regions

def percentage(a):
	return '%.1f%%' % (a * 100)

def main():
	from regions import parse_regions
	args = get_args()
	regions = parse_regions(args.mapfile)
	bus_stops = read_bus_stops(args.gtfs_zip)
	bus_stop_regions = compute_bus_stop_regions(regions, bus_stops)
	covered = sum(1 if value else 0 for value in bus_stop_regions.values())
	total = len(bus_stops)
	print('%s bus stops covered.' % percentage(covered / total),
		file = sys.stderr)
	json.dump(bus_stop_regions, sys.stdout, indent = 2)

if __name__ == '__main__':
	main()