maze_dataset.benchmark.speed

  1"benchmark the speed of maze generation"
  2
  3import functools
  4import random
  5import timeit
  6from pathlib import Path
  7from typing import Any, Sequence
  8
  9from tqdm import tqdm
 10
 11from maze_dataset import MazeDataset, MazeDatasetConfig
 12from maze_dataset.generation.default_generators import DEFAULT_GENERATORS
 13from maze_dataset.generation.generators import GENERATORS_MAP
 14
 15_BASE_CFG_KWARGS: dict = dict(
 16	grid_n=None,
 17	n_mazes=None,
 18)
 19
 20_GENERATE_KWARGS: dict = dict(
 21	gen_parallel=False,
 22	pool_kwargs=None,
 23	verbose=False,
 24	# do_generate = True,
 25	# load_local = False,
 26	# save_local = False,
 27	# zanj = None,
 28	# do_download = False,
 29	# local_base_path = "INVALID",
 30	# except_on_config_mismatch = True,
 31	# verbose = False,
 32)
 33
 34
 35def time_generation(
 36	base_configs: list[tuple[str, dict]],
 37	grid_n_vals: list[int],
 38	n_mazes_vals: list[int],
 39	trials: int = 10,
 40	verbose: bool = False,
 41) -> list[dict[str, Any]]:
 42	"time the generation of mazes for various configurations"
 43	# assemble configs
 44	configs: list[MazeDatasetConfig] = list()
 45
 46	for b_cfg in base_configs:
 47		for grid_n in grid_n_vals:
 48			for n_mazes in n_mazes_vals:
 49				configs.append(
 50					MazeDatasetConfig(
 51						name="benchmark",
 52						grid_n=grid_n,
 53						n_mazes=n_mazes,
 54						maze_ctor=GENERATORS_MAP[b_cfg[0]],
 55						maze_ctor_kwargs=b_cfg[1],
 56					),
 57				)
 58
 59	# shuffle configs (in place) (otherwise progress bar is annoying)
 60	random.shuffle(configs)
 61
 62	# time generation for each config
 63	times: list[dict[str, Any]] = list()
 64	total: int = len(configs)
 65	for idx, cfg in tqdm(
 66		enumerate(configs),
 67		desc="Timing generation",
 68		unit="config",
 69		total=total,
 70		disable=verbose,
 71	):
 72		if verbose:
 73			print(f"Timing generation for config {idx + 1}/{total}\n{cfg}")
 74
 75		t: float = (
 76			timeit.timeit(
 77				stmt=functools.partial(MazeDataset.generate, cfg, **_GENERATE_KWARGS),  # type: ignore[arg-type]
 78				number=trials,
 79			)
 80			/ trials
 81		)
 82
 83		if verbose:
 84			print(f"avg time: {t:.3f} s")
 85
 86		times.append(
 87			dict(
 88				cfg_name=cfg.name,
 89				grid_n=cfg.grid_n,
 90				n_mazes=cfg.n_mazes,
 91				maze_ctor=cfg.maze_ctor.__name__,
 92				maze_ctor_kwargs=cfg.maze_ctor_kwargs,
 93				trials=trials,
 94				time=t,
 95			),
 96		)
 97
 98	return times
 99
100
101def run_benchmark(
102	save_path: str,
103	base_configs: list[tuple[str, dict]] | None = None,
104	grid_n_vals: Sequence[int] = (2, 3, 4, 5, 8, 10, 16, 25, 32),
105	n_mazes_vals: Sequence[int] = tuple(range(1, 12, 2)),
106	trials: int = 10,
107	verbose: bool = True,
108) -> "pd.DataFrame":  # type: ignore[name-defined] # noqa: F821
109	"run the benchmark and save the results to a file"
110	import pandas as pd
111
112	if base_configs is None:
113		base_configs = DEFAULT_GENERATORS
114
115	times: list[dict] = time_generation(
116		base_configs=base_configs,
117		grid_n_vals=list(grid_n_vals),
118		n_mazes_vals=list(n_mazes_vals),
119		trials=trials,
120		verbose=verbose,
121	)
122
123	df: pd.DataFrame = pd.DataFrame(times)
124
125	# print the whole dataframe contents to console as csv
126	print(df.to_csv())
127
128	# save to file
129	Path(save_path).parent.mkdir(parents=True, exist_ok=True)
130	df.to_json(save_path, orient="records", lines=True)
131
132	return df