Add info/benchmarking mode to SimulationSet

src/ridepy/extras/io.py

@@ -164,7 +164,7 @@ def sort_params(params: dict) -> dict:
 
 def create_params_json(*, params: dict, sort=True) -> str:
     """
-    Create a dictionary containing simulation parameters to pretty JSON.
+    Convert a dictionary containing simulation parameters to pretty JSON.
     Parameter dictionaries may contain anything that is supported
     by `.ParamsJSONEncoder` and `.ParamsJSONDecoder`, e.g. `RequestGenerator`,
     `TransportSpace`s and dispatchers. For additional detail, see :ref:`Executing Simulations`.
@@ -251,3 +251,15 @@ def read_events_json(jsonl_path: Path) -> list[dict]:
     """
     with jsonl_path.open("r", encoding="utf-8") as f:
         return list(map(json.loads, f.readlines()))
+
+
+def create_info_json(info: dict) -> str:
+    """
+    Convert a dictionary containing simulation info to pretty JSON.
+
+    Parameters
+    ----------
+    info
+        dictionary containing the info to save
+    """
+    return json.dumps(info, indent=4)

src/ridepy/extras/simulation_set.py

@@ -53,6 +53,7 @@
     read_params_json,
     ParamsJSONEncoder,
     ParamsJSONDecoder,
+    create_info_json,
 )
 
 logger = logging.getLogger(__name__)
@@ -214,9 +215,12 @@ def perform_single_simulation(
     data_dir: Path,
     jsonl_chunksize: int = 1000,
     debug: bool = False,
-    param_path_suffix: str = "_params.json",
     event_path_suffix: str = ".jsonl",
+    param_path_suffix: str = "_params.json",
+    info_path_suffix: str = "_info.json",
     dry_run: bool = False,
+    info: bool = False,
+    info_contents: Optional[dict[str, Any]] = None,
 ) -> str:
     """
     Execute a single simulation run based on a parameter dictionary
@@ -252,6 +256,10 @@ def perform_single_simulation(
         Simulation events will be stored under "data_dir/<simulation_id><event_path_suffix>"
     dry_run
         If True, do not actually simulate. Just pretend to and return the corresponding ID.
+    info
+        Info/benchmark mode. If true, record the time it took to run the simulation run in a separate info file.
+    info_contents
+        Additional contents to write to the info file.
 
     Returns
     -------
@@ -264,6 +272,7 @@ def perform_single_simulation(
     sim_id = make_sim_id(params_json)
     event_path = data_dir / f"{sim_id}{event_path_suffix}"
     param_path = data_dir / f"{sim_id}{param_path_suffix}"
+    info_path = data_dir / f"{sim_id}{info_path_suffix}"
 
     if (
         param_path.exists()
@@ -341,9 +350,22 @@ def perform_single_simulation(
         else:
             raise ValueError("Must *either* specify `n_reqs` *or* `t_cutoff`")
 
+        simulation_start_time = time()
+
         while chunk := list(it.islice(simulation, jsonl_chunksize)):
             save_events_json(jsonl_path=event_path, events=chunk)
 
+        simulation_end_time = time()
+
+        if info:
+            simulation_duration = simulation_end_time - simulation_start_time
+            with info_path.open("w") as f:
+                info = {
+                    "simulation_duration": simulation_duration,
+                    "jsonl_chunksize": jsonl_chunksize,
+                } | (info_contents or {})
+                f.write(create_info_json(info))
+
         with open(str(param_path), "w") as f:
             f.write(params_json)
     tock = time()
@@ -363,6 +385,7 @@ def simulate_parameter_combinations(
     event_path_suffix: str = ".jsonl",
     param_path_suffix: str = "_params.json",
     dry_run: bool = False,
+    info: bool = False,
 ):
     """
     Run simulations for different parameter combinations using multiprocessing.
@@ -388,6 +411,8 @@ def simulate_parameter_combinations(
         Simulation events will be stored under "data_dir/<simulation_id><event_path_suffix>"
     dry_run
         If True, do not actually simulate. Just pretend to and return the corresponding IDs.
+    info
+        Info/benchmark mode. If true, record the time it took to run each simulation run.
 
     Returns
     -------
@@ -404,6 +429,11 @@ def simulate_parameter_combinations(
                     param_path_suffix=param_path_suffix,
                     event_path_suffix=event_path_suffix,
                     dry_run=dry_run,
+                    info=info,
+                    info_contents={
+                        "process_chunksize": process_chunksize,
+                        "max_workers": max_workers,
+                    },
                 ),
                 param_combinations,
                 chunksize=process_chunksize,
@@ -745,7 +775,7 @@ def param_combinations() -> Iterator[dict[str, dict[str, Any]]]:
     def __next__(self):
         return next(self._param_combinations)
 
-    def run(self, dry_run=False):
+    def run(self, dry_run: bool = False, info: bool = False):
         """
         Run the simulations configured through `base_params`, `zip_params` and `product_params` using
         multiprocessing. The parameters and resulting output events are written to disk
@@ -760,6 +790,8 @@ def run(self, dry_run=False):
         ----------
         dry_run
             If True, do not actually simulate.
+        info
+            Info/benchmark mode. If true, record the time it took to run each simulation run.
         """
 
         self._simulation_ids = simulate_parameter_combinations(
@@ -772,6 +804,7 @@ def run(self, dry_run=False):
             event_path_suffix=self._event_path_suffix,
             param_path_suffix=self._param_path_suffix,
             dry_run=dry_run,
+            info=info,
         )
 
     def __len__(self) -> int:

src/ridepy/util/spaces_cython/spaces.pyx

@@ -558,7 +558,7 @@ cdef class Graph(TransportSpace):
         if make_attribute_distance is None:
             weights = None
         else:
-            weights = [G[u][v][make_attribute_distance] for u, v in G.edges()]
+            weights = [distance for u, v, distance in G.edges.data(data=make_attribute_distance)]
 
         return cls(
             vertices=list(G.nodes()),