Coverage for tests/unit/plotting/test_maze_plot.py: 100%

1"""mostly taken from `demo_latticemaze.ipynb`""" 

2 

3import os 

4 

5import matplotlib.pyplot as plt 

6import numpy as np 

7 

8from maze_dataset.generation import LatticeMazeGenerators 

9from maze_dataset.maze import SolvedMaze, TargetedLatticeMaze 

10from maze_dataset.plotting import MazePlot 

11 

12FIG_SAVE: str = "tests/_temp/figures/" 

13 

14 

15def test_maze_plot(): 

16 N: int = 10 

17 

18 os.makedirs(FIG_SAVE, exist_ok=True) # noqa: PTH103 

19 

20 maze = LatticeMazeGenerators.gen_dfs(np.array([N, N])) 

21 tgt_maze: TargetedLatticeMaze = TargetedLatticeMaze.from_lattice_maze( 

22 maze, 

23 (0, 0), 

24 (N - 1, N - 1), 

25 ) 

26 solved_maze: SolvedMaze = SolvedMaze.from_targeted_lattice_maze(tgt_maze) 

27 

28 fig, ax = plt.subplots(1, 3, figsize=(15, 5)) 

29 

30 for ax_i, temp_maze in zip(ax, [maze, tgt_maze, solved_maze], strict=False): 

31 ax_i.set_title(temp_maze.as_ascii(), fontfamily="monospace") 

32 ax_i.imshow(temp_maze.as_pixels()) 

33 

34 assert temp_maze == temp_maze.__class__.from_pixels(temp_maze.as_pixels()) 

35 assert temp_maze == temp_maze.__class__.from_ascii(temp_maze.as_ascii()) 

36 

37 plt.savefig(FIG_SAVE + "pixels_and_ascii.png") 

38 

39 MazePlot(maze).plot() 

40 plt.savefig(FIG_SAVE + "mazeplot-pathless.png") 

41 

42 true_path = maze.find_shortest_path(c_start=(0, 0), c_end=(3, 3)) 

43 

44 MazePlot(solved_maze).plot() 

45 plt.savefig(FIG_SAVE + "mazeplot-solvedmaze.png") 

46 

47 pred_path1: list[tuple[int, int]] = [ 

48 (0, 0), 

49 (1, 0), 

50 (2, 0), 

51 (3, 0), 

52 (3, 1), 

53 (3, 2), 

54 (3, 3), 

55 ] 

56 pred_path2: list[tuple[int, int]] = [ 

57 (0, 0), 

58 (0, 1), 

59 (0, 2), 

60 (0, 3), 

61 (1, 3), 

62 (2, 3), 

63 (2, 2), 

64 (3, 2), 

65 (3, 3), 

66 ] 

67 ( 

68 MazePlot(maze) 

69 .add_true_path(true_path) 

70 .add_predicted_path(pred_path1) 

71 .add_predicted_path(pred_path2) 

72 .plot() 

73 ) 

74 plt.savefig(FIG_SAVE + "mazeplot-fakepaths.png") 

75 

76 node_values = np.random.uniform(size=maze.grid_shape) 

77 

78 MazePlot(maze).add_node_values(node_values, color_map="Blues").plot() 

79 plt.savefig(FIG_SAVE + "mazeplot-nodevalues.png") 

80 

81 MazePlot(maze).add_node_values( 

82 node_values, 

83 color_map="Blues", 

84 target_token_coord=np.array([2, 0]), 

85 preceeding_tokens_coords=np.array([[0, 0], [3, 1]]), 

86 ).plot() 

87 plt.savefig(FIG_SAVE + "mazeplot-nodevalues_target.png") 

88 

89 pred_paths: list[list[tuple[int, int]]] = [pred_path1, pred_path2] 

90 MazePlot(maze).add_multiple_paths(pred_paths).plot() 

91 plt.savefig(FIG_SAVE + "mazeplot-multipath.png") 

92 

93 ascii_maze = MazePlot(maze).to_ascii() 

94 print(ascii_maze)