Illustrative_unified_dataset_traffic_assignment.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "geeMM17tc_EL",
        "Dnj71Z6qdDMs",
        "upKfqqAfdHNM",
        "bIKkBQxJdOEj",
        "GiTmXiYJdU1W",
        "EmeYGBL9dao8",
        "C3r8cjq1dhzT",
        "dGKTe0tAdl7o",
        "4zfWLPEEdrHH"
      ],
      "authorship_tag": "ABX9TyNx2/Nez/sKbskbFrk9/+eq",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/kelvinfkr/Unified_UE_US_cities/blob/main/Illustrative_unified_dataset_traffic_assignment.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Traffic Assignment Model Implementation\n",
        "\n",
        "This Colab file provides an implementation of the traffic assignment model with [AequilibraE](https://github.com/AequilibraE/aequilibrae) described in the paper [\"A unified dataset for the city-scale traffic assignment model in 20 U.S. cities\"](https://doi.org/10.1038/s41597-024-03149-8) by Xu, X., Zheng, Z., Hu, Z. et al. (Sci Data 11, 325 (2024)). The purpose of this Colab file is to facilitate the understanding and learning of traffic assignment concepts for anyone interested in this topic.\n",
        "\n",
        "## Dataset\n",
        "The dataset used in this implementation can be found at:\n",
        "- [Dataset Address](https://figshare.com/ndownloader/files/45086305)\n",
        "\n",
        "## Contact Information\n",
        "For inquiries related to the paper, please contact:\n",
        "- Zhenjie Zheng: zzj17.zheng@polyu.edu.hk\n",
        "- Wei Ma: wei.w.ma@polyu.edu.hk\n",
        "\n",
        "The majority of the code in this Colab file is sourced from:\n",
        "- [GitHub Repository](https://github.com/xuxiaotong/A_unified_and_validated_traffic_dataset_for_20_U.S._cities)\n",
        "\n",
        "For questions regarding the GitHub code, please contact:\n",
        "- Xiaotong Xu: kid-a.xu@connect.polyu.hk\n",
        "\n",
        "If you have any questions or are interested in this Colab file, feel free to reach out to me:\n",
        "- Kairui Feng: kelvinfkr2015@gmail.com\n",
        "\n",
        "## Usage\n",
        "If you are new to Colab or Python, it is recommended to compile this notebook sequentially, line by line, to better understand the code and its functionality.\n",
        "\n",
        "## License\n",
        "This code is provided under a non-commercial license. You are free to use, modify, and distribute the code for non-commercial purposes. However, if you use this code in your research or any other work, please make sure to cite the following paper:\n",
        "\n",
        "Xu, X., Zheng, Z., Hu, Z. et al. (2024). A unified dataset for the city-scale traffic assignment model in 20 U.S. cities. Sci Data 11, 325. https://doi.org/10.1038/s41597-024-03149-8"
      ],
      "metadata": {
        "id": "-yhINJvz9DNs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# download dataset and unzip"
      ],
      "metadata": {
        "id": "geeMM17tc_EL"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "shaOgrmI85XQ"
      },
      "outputs": [],
      "source": [
        "!wget https://figshare.com/ndownloader/files/45086305\n",
        "!mv 45086305 downloaded_file.zip\n",
        "!unzip downloaded_file.zip\n",
        "#this will creat a series of folders under the root"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "# Install dependency"
      ],
      "metadata": {
        "id": "Dnj71Z6qdDMs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install osmnx\n",
        "!pip install geopandas\n",
        "# if you want to use system optimal for aequilibrea, do not compile the following line\n",
        "!pip install aequilibrae\n",
        "# if you want to get system optimal driven assignment, you need to download the bpr2,pyx\n",
        "# and substitue that one in the aequilibrae package, paths folder. https://github.com/kelvinfkr/Unified_UE_US_cities/blob/main/bpr2.pyx\n",
        "# the package aequlibrae is here: https://github.com/AequilibraE/aequilibrae\n",
        "# the following code shows how to do it.\n",
        "# !git clone https://github.com/AequilibraE/aequilibrae.git\n",
        "# then substitue the file\n",
        "# !wget -O /content/aequilibrae/aequilibrae/paths/bpr2.pyx https://raw.githubusercontent.com/kelvinfkr/Unified_UE_US_cities/main/bpr2.pyx\n",
        "# Then install the package\n",
        "# !pip install -e /content/aequilibrae/"
      ],
      "metadata": {
        "id": "GJS5lJQwINA9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Read the city data"
      ],
      "metadata": {
        "id": "upKfqqAfdHNM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import shutil\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import osmnx as ox\n",
        "import geopandas as gpd\n",
        "import numpy as np\n",
        "\n",
        "folders = [\n",
        "    \"01_San_Francisco\",\n",
        "    \"02_Seattle\",\n",
        "    \"03_Portland\",\n",
        "    \"04_Las_Vegas\",\n",
        "    \"05_Chicago\",\n",
        "    \"06_New_Orleans\",\n",
        "    \"07_Austin\",\n",
        "    \"08_Minneapolis\",\n",
        "    \"09_Dallas\",\n",
        "    \"10_Milwaukee\",\n",
        "    \"11_New_York\",\n",
        "    \"12_Washington\",\n",
        "    \"13_Boston\",\n",
        "    \"14_Philadelphia\",\n",
        "    \"15_Pittsburgh\",\n",
        "    \"16_Miami\",\n",
        "    \"17_Atlanta\",\n",
        "    \"18_Phoenix\",\n",
        "    \"19_Denver\",\n",
        "    \"20_Honolulu\"\n",
        "]\n",
        "\n",
        "# if you have other folders with the similar name structure then use the following line to find all the cities.\n",
        "#import re\n",
        "#folders = sorted([folder for folder in os.listdir() if re.match(r'^\\d+_', folder)], key=lambda x: int(x.split('_')[0]))"
      ],
      "metadata": {
        "id": "ayGxbucVJq--"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creat your own folder to store the results"
      ],
      "metadata": {
        "id": "bIKkBQxJdOEj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# List of folder names\n",
        "\n",
        "\n",
        "# Iterate over each folder\n",
        "for folder in folders:\n",
        "    # Change directory to the current folder\n",
        "    os.chdir(folder)\n",
        "\n",
        "    # Specify the source and destination folder names\n",
        "    source_folder = \"03_AequilibraE_results\"\n",
        "    destination_folder = \"04_my_AequilibraE_results\"\n",
        "\n",
        "    # Check if the source folder exists\n",
        "    if os.path.exists(source_folder):\n",
        "        # Copy the folder and its contents to the destination folder\n",
        "        shutil.copytree(source_folder, destination_folder)\n",
        "        print(f\"Folder '{source_folder}' copied and renamed to '{destination_folder}' successfully in folder '{folder}'.\")\n",
        "    else:\n",
        "        print(f\"Folder '{source_folder}' does not exist in folder '{folder}'.\")\n",
        "\n",
        "    # Change directory back to the parent folder\n",
        "    os.chdir(\"..\")"
      ],
      "metadata": {
        "id": "zzvX3tp-KoXz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Set Parameters for each city"
      ],
      "metadata": {
        "id": "GiTmXiYJdU1W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# predefine the traffic parameters for each city\n",
        "\n",
        "road_1_capacity_precondition = [2200,2200,2200,2200,2000,2200,2200,2200,2200,2200,2200,1800,2200,2000,2200,1800,2200,2200,2000,2200]\n",
        "alpha_precondition = [0.5,0.6,0.5,0.5,0.5,0.6,0.5,0.15,0.6,0.5,0.25,0.5,0.25,0.5,0.5,0.5,0.2,0.15,0.5,0.5]\n",
        "beta_precondition  = [1.8,3.0,1.2,1.3,1.2,1.8,1.5,1.80,1.3,1.5,1.50,1.5,2.00,1.2,2.0,1.5,1.5,1.20,1.5,1.5]\n",
        "road_1_capacity_precondition = [2200,2200,2200,2200,2000,2200,2200,2200,2200,2200,2200,1800,2200,2000,2200,1800,2200,2200,2000,2200]\n",
        "road_1_speed_precondition = [90,90,90,90,90,90,90,90,90,90,90,60,60,90,90,65,70,90,90,90]       # free flow speed (km/h)\n",
        "road_2_capacity_precondition = [2000,2000,2000,2000,1900,2000,2000,2000,2000,2000,2000,1500,2000,1800,2000,1500,2000,2000,1800,2000]\n",
        "road_2_speed_precondition = [60,65,65,60,60,60,65,65,65,65,60,40,45,60,60,50,50,65,60,60]\n",
        "road_3_capacity_precondition = [1400,1400,1400,1400,1400,1400,1400,1300,1400,1400,1400,600,1300,1200,1200,900,1400,1500,1300,1400]\n",
        "road_3_speed_precondition = [40,45,45,40,40,40,40,40,45,45,40,25,30,30,30,35,35,45,35,40]\n",
        "od_multiplier_precondition = [0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6,0.6]\n",
        "\n",
        "### if you have other cities, just add a column later, remember to number your city 21,22,23,24,25....\n",
        "\n",
        "city_params= {\n",
        "    folder: {\n",
        "        'alpha': alpha,\n",
        "        'beta': beta,\n",
        "        'road_1_capacity': road_1_capacity,\n",
        "        'road_1_speed': road_1_speed,\n",
        "        'road_2_capacity': road_2_capacity,\n",
        "        'road_2_speed': road_2_speed,\n",
        "        'road_3_capacity': road_3_capacity,\n",
        "        'road_3_speed': road_3_speed,\n",
        "        'od_multiplier': od_multiplier\n",
        "    }\n",
        "    for folder, alpha, beta, road_1_capacity, road_1_speed, road_2_capacity, road_2_speed, road_3_capacity, road_3_speed, od_multiplier in zip(\n",
        "        folders,\n",
        "        alpha_precondition,\n",
        "        beta_precondition,\n",
        "        road_1_capacity_precondition,\n",
        "        road_1_speed_precondition,\n",
        "        road_2_capacity_precondition,\n",
        "        road_2_speed_precondition,\n",
        "        road_3_capacity_precondition,\n",
        "        road_3_speed_precondition,\n",
        "        od_multiplier_precondition\n",
        "    )\n",
        "}\n",
        "\n",
        "# you could use something like city_params['01_San_Francisco']['alpha'] to access the detaled parameter for each city"
      ],
      "metadata": {
        "id": "23BPWyQsKr2s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Define a traffic assignment modular with AequilibraE"
      ],
      "metadata": {
        "id": "EmeYGBL9dao8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from aequilibrae.matrix import AequilibraeMatrix\n",
        "from aequilibrae.paths import Graph\n",
        "from aequilibrae.paths import TrafficAssignment\n",
        "from aequilibrae.paths.traffic_class import TrafficClass\n",
        "\n",
        "# the traffic assignment wrap-up, the input city_path should be something like \"01_San_Francisco\", a number _ a city name\n",
        "\n",
        "def traffic_assignment(city_path, vdf_type=\"BPR\"):\n",
        "    AE_path = city_path + '/' + '04_my_AequilibraE_results'\n",
        "\n",
        "    ######## step 2: define the od matrix file and generate od index #############\n",
        "    aemfile = AE_path + '/' + 'od_demand.aem'\n",
        "    od_data = pd.read_csv(city_path + '/' + '01_Input_data' + '/' + city_path[city_path.index(\"_\") + 1:] + '_od.csv')\n",
        "    zones = int(max(od_data.O_ID.max() - 9999999, od_data.D_ID.max() - 9999999))\n",
        "    index = np.arange(zones) + 1\n",
        "\n",
        "    ####### step 3: define the traffic assignment parameters #####\n",
        "    connector_speed = 5\n",
        "    road_1_capacity = city_params[city_path]['road_1_capacity'] # 3 lanes in total, capacity for each lane\n",
        "    road_1_speed = city_params[city_path]['road_1_speed'] # free flow speed (km/h)\n",
        "    road_2_capacity = city_params[city_path]['road_2_capacity'] # 2 lanes in total, capacity for each lane\n",
        "    road_2_speed = city_params[city_path]['road_2_speed']\n",
        "    road_3_capacity = city_params[city_path]['road_3_capacity'] # capacity for one lane\n",
        "    road_3_speed = city_params[city_path]['road_3_speed']\n",
        "    alpha = city_params[city_path]['alpha'] # alpha in BPR function\n",
        "    beta = city_params[city_path]['beta'] # beta in BPR function\n",
        "    od_multiplier = city_params[city_path]['od_multiplier'] # OD multiplier\n",
        "\n",
        "    ######## step 4: select the initial network shapefile path #############\n",
        "    net = gpd.read_file(AE_path + '/'+ city_path[city_path.index(\"_\") + 1:] + '.shp')\n",
        "    net.drop('geometry', axis=1, inplace=True)\n",
        "    net.columns = [\"ID\", \"DIR\", \"length\", \"LINK_ID\", \"a_node\", \"b_node\", \"capacity\", \"LENGTH1\", \"speed\", \"Lanes\", \"link_type\"]\n",
        "\n",
        "    ######## step 5: modify the parameters for traffic assignment #############\n",
        "    ########### connector ##############\n",
        "    net.loc[net['link_type'] > 10, 'speed'] = connector_speed\n",
        "    ########### road class 1 ##############\n",
        "    net.loc[net['link_type'].isin([1, 2]), 'capacity'] = road_1_capacity\n",
        "    net.loc[net['link_type'].isin([1, 2]), 'speed'] = road_1_speed\n",
        "    ########### road class 2 ##############\n",
        "    net.loc[net['link_type'] == 3, 'capacity'] = road_2_capacity\n",
        "    net.loc[net['link_type'] == 3, 'speed'] = road_2_speed\n",
        "    ########### road class 3 ##############\n",
        "    net.loc[net['link_type'].isin([4, 5]), 'capacity'] = road_3_capacity\n",
        "    net.loc[net['link_type'].isin([4, 5]), 'speed'] = road_3_speed\n",
        "    ########### BPR function ##############\n",
        "    net['free_flow_time'] = net['length'] / net['speed'] * 60\n",
        "    net['a'] = alpha\n",
        "    net['b'] = beta\n",
        "\n",
        "    ########### organize the network data fields ##############\n",
        "    network = net[['a_node', 'b_node', \"capacity\", 'free_flow_time', \"a\", \"b\"]]\n",
        "    network = network.assign(direction=1)\n",
        "    network[\"link_id\"] = network.index + 1\n",
        "    network = network.astype({\"a_node\": \"int64\", \"b_node\": \"int64\"})\n",
        "\n",
        "    ########### step 6: conduct the traffic assignment in AE ##############\n",
        "    g = Graph()\n",
        "    g.cost = network['free_flow_time'].values\n",
        "    g.capacity = network['capacity'].values\n",
        "    g.free_flow_time = network['free_flow_time'].values\n",
        "    g.network = network\n",
        "    g.network_ok = True\n",
        "    g.status = 'OK'\n",
        "    g.prepare_graph(index)\n",
        "    g.set_graph(\"free_flow_time\")\n",
        "    g.cost = np.array(g.cost, copy=True)\n",
        "    g.set_skimming([\"free_flow_time\"])\n",
        "    g.set_blocked_centroid_flows(False)\n",
        "    g.network[\"id\"] = g.network.link_id\n",
        "\n",
        "    aem = AequilibraeMatrix()\n",
        "    aem.load(aemfile)\n",
        "    aem.computational_view([\"matrix\"])\n",
        "\n",
        "    assigclass = TrafficClass(\"car\", g, aem)\n",
        "    assig = TrafficAssignment()\n",
        "    assig.set_classes([assigclass])\n",
        "    assig.set_vdf(vdf_type)\n",
        "    assig.set_vdf_parameters({\"alpha\": \"a\", \"beta\": \"b\"})\n",
        "    assig.set_capacity_field(\"capacity\")\n",
        "    assig.set_time_field(\"free_flow_time\")\n",
        "    assig.set_algorithm(\"bfw\")\n",
        "    assig.max_iter = 500\n",
        "    assig.rgap_target = 0.001\n",
        "    assig.execute()\n",
        "\n",
        "    assignment_results = assig.results()\n",
        "    network_data = network.copy()\n",
        "\n",
        "    ########### step 7: save the output ##############\n",
        "    network.to_csv(AE_path + '\\\\' + 'network.csv')\n",
        "    assig.results().to_csv(AE_path+ '\\\\' + 'assignment_result.csv')\n",
        "    return assignment_results, network_data,aem"
      ],
      "metadata": {
        "id": "JByovyjwSKRP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Run the traffic assignment"
      ],
      "metadata": {
        "id": "C3r8cjq1dhzT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "city_ind = 0 # calculate the UE for the city you like\n",
        "# if you want to add other cities, name the city as 21_cityname for example and follow the same folder structures\n",
        "city_path =  folders[city_ind]\n",
        "assignment_results,network_data,aem = traffic_assignment(city_path)# if you are just aiming for UE, then use this\n",
        "# if you installed SO and want to use SO\n",
        "#traffic_assignment(city_path,\"BPR2\")"
      ],
      "metadata": {
        "id": "Q94GepazMD5k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# A simple application: average traffic time"
      ],
      "metadata": {
        "id": "dGKTe0tAdl7o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# a simple use of the assignment results, which is to calculate the commuting cost for each vehicle in the system\n",
        "def calculate_average_traffic_time(city_path, assignment_results, network_data, aem):\n",
        "    # Step 1: Reset the index of assignment_results and rename the index column to \"link_id\"\n",
        "    assignment_results = assignment_results.copy()  # Create a copy to avoid modifying the original DataFrame\n",
        "    assignment_results.reset_index(inplace=True)\n",
        "    assignment_results.rename(columns={\"index\": \"link_id\"}, inplace=True)\n",
        "\n",
        "    # Step 2: Merge the assignment results with the network data\n",
        "    network_data = network_data.merge(assignment_results[[\"link_id\", \"VOC_AB\", \"PCE_tot\"]], on=\"link_id\")\n",
        "\n",
        "    # Step 3: Define the BPR function parameters\n",
        "    alpha = city_params[city_path]['alpha']  # alpha in BPR function\n",
        "    beta = city_params[city_path]['beta']  # beta in BPR function\n",
        "\n",
        "    # Step 4: Calculate the congested travel time using the BPR function, considering connector links\n",
        "    def calculate_congested_time(row):\n",
        "        if row[\"a_node\"] > 100000 or row[\"b_node\"] > 100000:  # Connector link\n",
        "            return row[\"free_flow_time\"]\n",
        "        else:  # Regular link\n",
        "            return row[\"free_flow_time\"] * (1 + alpha * (row[\"VOC_AB\"]) ** beta)\n",
        "\n",
        "    network_data[\"congested_time\"] = network_data.apply(calculate_congested_time, axis=1)\n",
        "\n",
        "    # Step 5: Load the OD matrix\n",
        "    od_matrix = aem.matrix[\"matrix\"]\n",
        "\n",
        "    # Step 6: Calculate the total OD flow\n",
        "    total_od_flow = od_matrix.sum().sum()\n",
        "\n",
        "    # Step 7: Calculate the total travel time\n",
        "    total_travel_time = (network_data[\"congested_time\"] * network_data[\"PCE_tot\"]).sum()\n",
        "\n",
        "    # Step 8: Calculate the total traffic volume\n",
        "    total_traffic_volume = network_data[\"PCE_tot\"].sum()\n",
        "\n",
        "    # Step 9: Calculate the average traffic time\n",
        "    average_traffic_time = total_travel_time / total_od_flow\n",
        "\n",
        "    return average_traffic_time"
      ],
      "metadata": {
        "id": "0b7Iow1cTew6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "average_time = calculate_average_traffic_time(city_path, assignment_results, network_data, aem)\n",
        "print(\"Average Traffic Time (User Equilibrium):\", average_time, \"minutes\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-HslGBodVYcm",
        "outputId": "8b8c5904-61df-494f-bb53-31bd73bedd12"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Average Traffic Time (User Equilibrium): 24.361353481084972 minutes\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Run over all the cities"
      ],
      "metadata": {
        "id": "4zfWLPEEdrHH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# a simple code to go through all the cities and store the results in a dictionary.\n",
        "results_store = {}\n",
        "\n",
        "for city_path in folders:\n",
        "    assignment_results, network_data, aem = traffic_assignment(city_path)\n",
        "    average_time = calculate_average_traffic_time(city_path, assignment_results, network_data, aem)\n",
        "    print(f\"Average Traffic Time (User Equilibrium) for {city_path}: {average_time} minutes\")\n",
        "\n",
        "    results_store[city_path] = {\n",
        "        \"assignment_results\": assignment_results,\n",
        "        \"network_data\": network_data,\n",
        "        \"aem\": aem,\n",
        "        \"average_time\": average_time\n",
        "    }"
      ],
      "metadata": {
        "id": "CQ9B9XLkWvgu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#if you can not wait, you could use a bit of parallel computing, if you are colab pro user and have 8 kernels.\n",
        "import multiprocessing\n",
        "\n",
        "def process_city(city_path):\n",
        "    assignment_results, network_data, aem = traffic_assignment(city_path)\n",
        "    average_time = calculate_average_traffic_time(city_path, assignment_results, network_data, aem)\n",
        "    print(f\"Average Traffic Time (User Equilibrium) for {city_path}: {average_time} minutes\")\n",
        "\n",
        "    return {\n",
        "        \"assignment_results\": assignment_results,\n",
        "        \"network_data\": network_data,\n",
        "        \"aem\": aem,\n",
        "        \"average_time\": average_time\n",
        "    }\n",
        "\n",
        "pool = multiprocessing.Pool()\n",
        "results = pool.map(process_city, folders)\n",
        "pool.close()\n",
        "pool.join()\n",
        "\n",
        "results_store = dict(zip(folders, results))"
      ],
      "metadata": {
        "id": "pFzAXRZaW9Po"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#if you want an so also stored, make sure you changed the bpr2.pyx\n",
        "results_store = {}\n",
        "\n",
        "for city_path in folders:\n",
        "    assignment_results, network_data, aem = traffic_assignment(city_path)\n",
        "    average_time = calculate_average_traffic_time(city_path, assignment_results, network_data, aem)\n",
        "    assignment_results_so, network_data_so, aem = traffic_assignment(city_path, \"BPR2\")\n",
        "    average_time_so = calculate_average_traffic_time(city_path, assignment_results, network_data, aem)\n",
        "    print(f\"Average Traffic Time (User Equilibrium) for {city_path}: {average_time} minutes\")\n",
        "    print(f\"Average Traffic Time (SO) for {city_path}: {average_time_so} minutes\")\n",
        "\n",
        "    results_store[city_path] = {\n",
        "        \"assignment_results\": assignment_results,\n",
        "        \"network_data\": network_data,\n",
        "        \"aem\": aem,\n",
        "        \"average_time\": average_time,\n",
        "        \"assignment_results_so\": assignment_results_so,\n",
        "        \"network_data_so\": network_data_so,\n",
        "        \"average_time_so\": average_time_so\n",
        "    }"
      ],
      "metadata": {
        "id": "J_-KF-ovXOwN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# a parallel version\n",
        "import multiprocessing\n",
        "\n",
        "def process_city(city_path):\n",
        "    assignment_results, network_data, aem = traffic_assignment(city_path)\n",
        "    average_time = calculate_average_traffic_time(city_path, assignment_results, network_data, aem)\n",
        "\n",
        "    assignment_results_so, network_data_so, _ = traffic_assignment(city_path, \"BPR2\")\n",
        "    average_time_so = calculate_average_traffic_time(city_path, assignment_results_so, network_data_so, aem)\n",
        "\n",
        "    print(f\"Average Traffic Time (User Equilibrium) for {city_path}: {average_time} minutes\")\n",
        "    print(f\"Average Traffic Time (SO) for {city_path}: {average_time_so} minutes\")\n",
        "\n",
        "    return {\n",
        "        \"assignment_results\": assignment_results,\n",
        "        \"network_data\": network_data,\n",
        "        \"aem\": aem,\n",
        "        \"average_time\": average_time,\n",
        "        \"assignment_results_so\": assignment_results_so,\n",
        "        \"network_data_so\": network_data_so,\n",
        "        \"average_time_so\": average_time_so\n",
        "    }\n",
        "\n",
        "pool = multiprocessing.Pool()\n",
        "results = pool.map(process_city, folders)\n",
        "pool.close()\n",
        "pool.join()\n",
        "\n",
        "results_store = dict(zip(folders, results))"
      ],
      "metadata": {
        "id": "DVEI_ZsvXmL-"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}