Merging GTFS with shapefile geometries.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import csv\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import partridge as ptg\n",
    "from scipy.spatial.distance import directed_hausdorff\n",
    "import shapefile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "shp_path = '/Users/drw/Downloads/Carreiras_Lineas_WGS84/Carreiras_Lineas_WGS84.shp'\n",
    "gtfs_path = '/Users/drw/Downloads/carris-clean.zip'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = shapefile.Reader(shp_path)\n",
    "feed = ptg.feed(gtfs_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cleaned up 1 wrapped points\n",
      "cleaned up 1 wrapped points\n",
      "cleaned up 1 wrapped points\n",
      "cleaned up 1 wrapped points\n",
      "cleaned up 1 wrapped points\n",
      "cleaned up 1 wrapped points\n"
     ]
    }
   ],
   "source": [
    "S = []\n",
    "for sr in reader.iterShapeRecords():\n",
    "    points = sr.shape.points\n",
    "    \n",
    "    bad_points = 0\n",
    "    while points[0] == points[-1]:\n",
    "        bad_points += 1\n",
    "        points = points[:-1]\n",
    "\n",
    "    if bad_points:\n",
    "        print('cleaned up {} wrapped points'.format(bad_points))\n",
    "\n",
    "    S.append(np.array(points))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "pattern_by_trip_id = {}\n",
    "patterns = set()\n",
    "for trip_id, stop_times in feed.stop_times.sort_values('stop_sequence').groupby('trip_id'):\n",
    "    pattern = tuple(stop_times.stop_id)\n",
    "    patterns.add(pattern)\n",
    "    pattern_by_trip_id[trip_id] = pattern\n",
    "\n",
    "patterns = list(patterns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "stops = {}\n",
    "for _, stop in feed.stops.iterrows():\n",
    "    stops[stop.stop_id] = (stop.stop_lon, stop.stop_lat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "P = []\n",
    "for pattern in patterns:\n",
    "    P.append([stops[stop_id] for stop_id in pattern])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "C = []\n",
    "for p in P:\n",
    "    best_idx = None\n",
    "    best_dist = float('inf')\n",
    "    for j, s in enumerate(S):\n",
    "        dist = directed_hausdorff(p, s)[0]\n",
    "        if dist < best_dist:\n",
    "            best_dist = dist\n",
    "            best_idx = j\n",
    "\n",
    "    assert best_idx, 'something went wrong'\n",
    "\n",
    "    C.append(best_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes = {\n",
    "    'shape_id': [],\n",
    "    'shape_pt_sequence': [],\n",
    "    'shape_pt_lat': [],\n",
    "    'shape_pt_lon': [],\n",
    "}\n",
    "\n",
    "shape_id_by_pattern = {}\n",
    "for shape_id, pattern in enumerate(patterns):\n",
    "    shape_id_by_pattern[pattern] = shape_id\n",
    "    shape_points = S[C[shape_id]]\n",
    "    for shape_pt_sequence, (shape_pt_lon, shape_pt_lat) in enumerate(shape_points, start=1):\n",
    "        shapes['shape_id'].append(shape_id)\n",
    "        shapes['shape_pt_sequence'].append(shape_pt_sequence)\n",
    "        shapes['shape_pt_lon'].append(shape_pt_lon)\n",
    "        shapes['shape_pt_lat'].append(shape_pt_lat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_feed = ptg.raw_feed(gtfs_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdf = pd.DataFrame(shapes)\n",
    "\n",
    "output_feed.shapes['shape_id'] = sdf.shape_id\n",
    "output_feed.shapes['shape_pt_sequence'] = sdf.shape_pt_sequence\n",
    "output_feed.shapes['shape_pt_lon'] = sdf.shape_pt_lon\n",
    "output_feed.shapes['shape_pt_lat'] = sdf.shape_pt_lat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_shape_id_by_trip_id(trip_id):\n",
    "    pattern = pattern_by_trip_id[trip_id]\n",
    "    return shape_id_by_pattern[pattern]\n",
    "\n",
    "output_feed.trips['shape_id'] = output_feed.trips.trip_id.apply(get_shape_id_by_trip_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/Users/drw/Code/learning/carris-with-shapes.zip'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ptg.writers.write_feed_dangerously(output_feed, 'carris-with-shapes.zip')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}