Generate trips from frequencies.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import copy\n",
    "import math\n",
    "\n",
    "from networkx.utils import pairwise\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import partridge as ptg\n",
    "\n",
    "\n",
    "def seconds_to_gtfs_time(total_seconds):\n",
    "    if math.isnan(total_seconds):\n",
    "        return total_seconds  # TODO: What to do here?\n",
    "    minutes, seconds = divmod(total_seconds, 60)\n",
    "    hours, minutes = divmod(minutes, 60)\n",
    "    time = list(map(lambda x: str(x).rjust(2, '0'), [int(hours), int(minutes), int(seconds)]))\n",
    "    return f'{time[0]}:{time[1]}:{time[2]}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "inpath = \"/Users/drw/Downloads/daviz.cg_gtfs (1).zip\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "feed = ptg.load_feed(inpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "trips_by_id = {}\n",
    "for _, trip in feed.trips.iterrows():\n",
    "    trips_by_id[trip.trip_id] = dict(trip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "trip_patterns = {}\n",
    "for trip_id, stop_times in feed.stop_times.sort_values(\"stop_sequence\").groupby(\"trip_id\"):\n",
    "    stops = tuple(stop_times.stop_id)\n",
    "    mintime = stop_times.arrival_time.min()\n",
    "    times = tuple(t - mintime for t in stop_times.arrival_time)\n",
    "    trip_patterns[trip_id] = (stops, times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "freq_trips = []\n",
    "for _, freq in feed.frequencies.iterrows():\n",
    "    window_start = int(freq.start_time)\n",
    "    window_end = int(freq.end_time)\n",
    "    for start in range(window_start, window_end, freq.headway_secs):\n",
    "        freq_trips.append({\n",
    "            \"trip_id\": freq.trip_id,\n",
    "            \"start\": start,\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_trips = []\n",
    "new_stop_times = []\n",
    "for i, ftrip in enumerate(freq_trips, start=1):\n",
    "    new_trips.append(copy(trips_by_id[ftrip[\"trip_id\"]]))\n",
    "    new_trips[-1][\"trip_id\"] = i # override trip_id\n",
    "\n",
    "    stops, times = trip_patterns[ftrip[\"trip_id\"]]\n",
    "    for j in range(len(stops)):\n",
    "        t = seconds_to_gtfs_time(times[j] + ftrip[\"start\"])\n",
    "        new_stop_times.append({\n",
    "            \"trip_id\": i,\n",
    "            \"stop_id\": stops[j],\n",
    "            \"arrival_time\": t,\n",
    "            \"departure_time\": t,\n",
    "            \"stop_sequence\": j + 1,\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/Users/drw/Downloads/SPTrans-generated-trips.zip'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trips_df = pd.DataFrame(new_trips)\n",
    "stop_times_df = pd.DataFrame(new_stop_times)\n",
    "empty_frequencies_df = ptg.utilities.empty_df()\n",
    "\n",
    "new_feed = ptg.load_raw_feed(inpath)\n",
    "new_feed.set(\"trips.txt\", trips_df)\n",
    "new_feed.set(\"stop_times.txt\", stop_times_df)\n",
    "new_feed.set(\"frequencies.txt\", empty_frequencies_df) # we don't want frequencies.txt\n",
    "\n",
    "ptg.writers.write_feed_dangerously(new_feed, \"/Users/drw/Downloads/SPTrans-generated-trips.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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}