add traffic_light_mode to control tl generation in sumo map converting

mosstool/map/_map_util/junctions/gen_traffic_light.py

@@ -213,7 +213,7 @@ def _gen_fixed_program(
         Literal["green_red"],
         Literal["green_yellow_red"],
         Literal["green_yellow_clear_red"],
-    ] = "green_yellow_clear_red",
+    ],
 ):
     """
     Generate fixed program traffic-light
@@ -370,6 +370,11 @@ def _convert_fixed_program(
     yellow_time: float,
     id_mapping: dict,
     traffic_light_path: Optional[str],
+    traffic_light_mode: Union[
+        Literal["green_red"],
+        Literal["green_yellow_red"],
+        Literal["green_yellow_clear_red"],
+    ],
 ):
     for jid, j in juncs.items():
         j["fixed_program"] = {}
@@ -496,7 +501,7 @@ def in_get_vector(line):
                 "phases": out_phases,
             }
     # Complete the default fixed program traffic-light for all junctions
-    _gen_fixed_program(lanes, roads, juncs, green_time, yellow_time)
+    _gen_fixed_program(lanes, roads, juncs, green_time, yellow_time, traffic_light_mode)
 
 
 def _gen_available_phases(lanes: dict, juncs: dict, min_direction_group: int):
@@ -737,6 +742,11 @@ def convert_traffic_light(
     green_time: float,
     yellow_time: float,
     min_direction_group: int,
+    traffic_light_mode: Union[
+        Literal["green_red"],
+        Literal["green_yellow_red"],
+        Literal["green_yellow_clear_red"],
+    ],
     traffic_light_path: Optional[str] = None,
 ):
     # Generating available phases for MP
@@ -745,5 +755,12 @@ def convert_traffic_light(
     # Generating fixed program traffic-lights or convert SUMO traffic-lights
     logging.info("Generating fixed program")
     _convert_fixed_program(
-        lanes, roads, juncs, green_time, yellow_time, id_mapping, traffic_light_path
+        lanes=lanes,
+        roads=roads,
+        juncs=juncs,
+        green_time=green_time,
+        yellow_time=yellow_time,
+        id_mapping=id_mapping,
+        traffic_light_path=traffic_light_path,
+        traffic_light_mode=traffic_light_mode,
     )

mosstool/map/sumo/map.py

@@ -2,7 +2,7 @@
 import time
 from collections import defaultdict
 from multiprocessing import cpu_count
-from typing import Optional
+from typing import Literal, Optional, Union
 from xml.dom.minidom import parse
 
 import numpy as np
@@ -35,6 +35,11 @@ def __init__(
         merge_aoi: bool = False,
         enable_tqdm: bool = False,
         multiprocessing_chunk_size: int = 500,
+        traffic_light_mode: Union[
+            Literal["green_red"],
+            Literal["green_yellow_red"],
+            Literal["green_yellow_clear_red"],
+        ] = "green_yellow_clear_red",
         workers: int = cpu_count(),
     ):
         """
@@ -50,6 +55,7 @@ def __init__(
         - merge_aoi (bool): merge nearby aois
         - enable_tqdm (bool): when enabled, use tqdm to show the progress bars
         - multiprocessing_chunk_size (int): the maximum size of each multiprocessing chunk
+        - traffic_light_mode (str): fixed time traffic-light generation mode. `green_red` means only green and red light will be generated, `green_yellow_red` means there will be yellow light between green and red light, `green_yellow_clear_red` add extra pedestrian clear red light.
         - workers (int): number of workers
         """
         self._lane_width = default_lane_width
@@ -62,6 +68,7 @@ def __init__(
         self._merge_aoi = merge_aoi
         self.multiprocessing_chunk_size = multiprocessing_chunk_size
         self._enable_tqdm = enable_tqdm
+        self._traffic_light_mode = traffic_light_mode
         self._workers = workers
         logging.info(f"Reading net file from {net_path}")
         dom_tree = parse(net_path)
@@ -470,14 +477,15 @@ def _add_traffic_light(self):
         roads = {rid: d for rid, d in self._output_roads.items()}
         juncs = {jid: d for jid, d in self._output_junctions.items()}
         convert_traffic_light(
-            lanes,
-            roads,
-            juncs,
-            self._id2uid,
-            self._green_time,
-            self._yellow_time,
-            self._traffic_light_min_direction_group,
-            self._traffic_light_path,
+            lanes=lanes,
+            roads=roads,
+            juncs=juncs,
+            id_mapping=self._id2uid,
+            green_time=self._green_time,
+            yellow_time=self._yellow_time,
+            min_direction_group=self._traffic_light_min_direction_group,
+            traffic_light_mode=self._traffic_light_mode,  # type:ignore
+            traffic_light_path=self._traffic_light_path,
         )
 
     def _get_output_map(self):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mosstool"
-version = "1.3.2"
+version = "1.3.3"
 description = "MObility Simulation System toolbox "
 authors = ["Jun Zhang <[email protected]>","Junbo Yan <[email protected]>"]
 license = "MIT"