maze_dataset.plotting.plot_dataset

 1"""`plot_dataset_mazes` will plot several mazes using `as_pixels`
 2
 3`print_dataset_mazes` will use `as_ascii` to print several mazes
 4"""
 5
 6import matplotlib.pyplot as plt  # type: ignore[import]
 7
 8from maze_dataset.dataset.maze_dataset import MazeDataset
 9
10
11def plot_dataset_mazes(
12	ds: MazeDataset,
13	count: int | None = None,
14	figsize_mult: tuple[float, float] = (1.0, 2.0),
15	title: bool | str = True,
16) -> tuple | None:
17	"plot `count` mazes from the dataset `d` in a single figure using `SolvedMaze.as_pixels()`"
18	count = count or len(ds)
19	if count == 0:
20		print("No mazes to plot for dataset")
21		return None
22	fig, axes = plt.subplots(
23		1,
24		count,
25		figsize=(count * figsize_mult[0], figsize_mult[1]),
26	)
27	if count == 1:
28		axes = [axes]
29	for i in range(count):
30		axes[i].imshow(ds[i].as_pixels())
31		# remove ticks
32		axes[i].set_xticks([])
33		axes[i].set_yticks([])
34
35	# set title
36	if title:
37		if isinstance(title, str):
38			fig.suptitle(title)
39		else:
40			kwargs: dict = {
41				"grid_n": ds.cfg.grid_n,
42				# "n_mazes": ds.cfg.n_mazes,
43				**ds.cfg.maze_ctor_kwargs,
44			}
45			fig.suptitle(
46				f"{ds.cfg.to_fname()}\n{ds.cfg.maze_ctor.__name__}({', '.join(f'{k}={v}' for k, v in kwargs.items())})",
47			)
48
49	# tight layout
50	fig.tight_layout()
51	# remove whitespace between title and subplots
52	fig.subplots_adjust(top=1.0)
53
54	return fig, axes
55
56
57def print_dataset_mazes(ds: MazeDataset, count: int | None = None) -> None:
58	"print ascii representation of `count` mazes from the dataset `d`"
59	count = count or len(ds)
60	if count == 0:
61		print("No mazes to print for dataset")
62		return
63	for i in range(count):
64		print(ds[i].as_ascii(), "\n\n-----\n")