From f951b296f74015903b9ae0b83e30842bdb3761cc Mon Sep 17 00:00:00 2001
From: Dustin Carlino <dabreegster@gmail.com>
Date: Thu, 1 Oct 2020 15:13:18 -0700
Subject: [PATCH] Cap all roads in the Python experiment. Right now, this'll
 unintentionally merge adjacent zones with the same cap into zones.

---
 headless/examples/cap_experiment.py | 47 +++++++++++++++++++++++------
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/headless/examples/cap_experiment.py b/headless/examples/cap_experiment.py
index a76aa45233..ceca4c2e21 100755
--- a/headless/examples/cap_experiment.py
+++ b/headless/examples/cap_experiment.py
@@ -24,6 +24,7 @@ def main():
     parser.add_argument('--hours', type=int, default=24)
     parser.add_argument('--cap_pct', type=int, default=80)
     parser.add_argument('--rounds', type=int, default=10)
+    parser.add_argument('--cap_all_roads', type=bool, default=True)
     args = parser.parse_args()
     print('Simulating {} hours of data/system/scenarios/{}/weekday.bin'.format(args.hours, args.map_name))
 
@@ -32,20 +33,27 @@ def main():
 
     for _ in range(args.rounds):
         print('')
-        print('{} roads have a {}% cap'.format(
-            len(edits['commands']), args.cap_pct))
+        print('{} roads have a cap'.format(len(edits['commands'])))
         experiment = abst_helpers.run_sim(args, edits=edits)
         baseline.compare(experiment)
         print('{:,} trips changed due to the caps'.format(
             len(experiment.capped_trips)))
 
-        # Cap the busiest road
-        busiest_road, thruput = find_busiest_road(args)
-        cmd = get(args, '/map/get-edit-road-command',
-                  params={'id': busiest_road}).json()
-        cmd['ChangeRoad']['new']['access_restrictions']['cap_vehicles_per_hour'] = int(
-            (args.cap_pct / 100.0) * thruput)
-        edits['commands'].append(cmd)
+        if args.cap_all_roads:
+            edits['commands'] = cap_all_roads(args)
+        else:
+            # Cap the busiest road
+            busiest_road, thruput = find_busiest_road(args)
+            cmd = get(args, '/map/get-edit-road-command',
+                      params={'id': busiest_road}).json()
+            cmd['ChangeRoad']['new']['access_restrictions']['cap_vehicles_per_hour'] = int(
+                (args.cap_pct / 100.0) * thruput)
+            edits['commands'].append(cmd)
+
+    # Write the final edits
+    f = open('cap_edits.json', 'w')
+    f.write(get(args, '/map/get-edits').text)
+    f.close()
 
 
 # Find the road with the most car traffic in any one hour period
@@ -64,5 +72,26 @@ def find_busiest_road(args):
     return (max_key[0], max_value)
 
 
+# Cap all roads to some percent of their max throughput over any one hour period
+def cap_all_roads(args):
+    thruput = get(
+        args, '/data/get-road-thruput').json()['counts']
+    max_per_road = {}
+    for r, agent, hr, count in thruput:
+        if agent == 'Car':
+            if r not in max_per_road or count > max_per_road[r]:
+                max_per_road[r] = count
+    commands = []
+    for r, count in max_per_road.items():
+        cap = int((args.cap_pct / 100.0) * count)
+        # Don't go too low
+        if cap > 10:
+            cmd = get(args, '/map/get-edit-road-command',
+                      params={'id': r}).json()
+            cmd['ChangeRoad']['new']['access_restrictions']['cap_vehicles_per_hour'] = cap
+            commands.append(cmd)
+    return commands
+
+
 if __name__ == '__main__':
     main()