diff --git a/main.ipynb b/main.ipynb
index c760046..51c4e53 100644
--- a/main.ipynb
+++ b/main.ipynb
@@ -309,7 +309,8 @@
     "## Visualisation tools\n",
     "\n",
     "In this section a visualisation help was implemented for debugging of the game and a proper display of the results.\n",
-    "For this visualisation ChatGPT was used as a prompted code generator that was later reviewed and refactored by hand to integrate seamlessly into the project as a whole."
+    "For this visualisation ChatGPT was used as a prompted code generator that was later reviewed and refactored by hand to integrate seamlessly into the project as a whole.\n",
+    "White stones represent the player, black stones the enemy. A single plot can be used as a subplot when the `ax` argument is used."
    ],
    "metadata": {
     "collapsed": false
@@ -340,6 +341,7 @@
     "        board: The bord that should be plotted. Only a single games is allowed. A numpy array of the form 8x8 is expected.\n",
     "        ax: If needed a matplotlib axis object can be defined that is used to place the board as a sublot into a bigger context.\n",
     "    \"\"\"\n",
+    "    assert board.shape == (8, 8)\n",
     "    plot_all = False\n",
     "    if ax is None:\n",
     "        fig_size = 3\n",
@@ -347,18 +349,17 @@
     "        fig, ax = plt.subplots(figsize=(fig_size, fig_size))\n",
     "\n",
     "    ax.set_facecolor(\"#66FF00\")\n",
-    "    for i in range(BOARD_SIZE):\n",
-    "        for j in range(BOARD_SIZE):\n",
-    "            if board[i, j] == -1:\n",
-    "                color = \"white\"\n",
-    "            elif board[i, j] == 1:\n",
-    "                color = \"black\"\n",
-    "            else:\n",
-    "                continue\n",
-    "            ax.scatter(j, i, s=300 if plot_all else 150, c=color)\n",
-    "    for i in range(-1, 8):\n",
-    "        ax.axhline(i + 0.5, color=\"black\", lw=2)\n",
-    "        ax.axvline(i + 0.5, color=\"black\", lw=2)\n",
+    "    for x_pos, y_pos in itertools.product(range(BOARD_SIZE), range(BOARD_SIZE)):\n",
+    "        if board[x_pos, y_pos] == -1:\n",
+    "            color = \"white\"\n",
+    "        elif board[x_pos, y_pos] == 1:\n",
+    "            color = \"black\"\n",
+    "        else:\n",
+    "            continue\n",
+    "        ax.scatter(y_pos, x_pos, s=300 if plot_all else 150, c=color)\n",
+    "    for x_pos in range(-1, 8):\n",
+    "        ax.axhline(x_pos + 0.5, color=\"black\", lw=2)\n",
+    "        ax.axvline(x_pos + 0.5, color=\"black\", lw=2)\n",
     "    ax.set_xlim(-0.5, 7.5)\n",
     "    ax.set_ylim(7.5, -0.5)\n",
     "    ax.set_xticks(np.arange(8))\n",
@@ -380,7 +381,17 @@
    "outputs": [],
    "source": [
     "def plot_othello_boards(boards: np.ndarray) -> None:\n",
+    "    \"\"\"Plots multiple boards into subplots.\n",
+    "\n",
+    "    The plots are shown directly.\n",
+    "\n",
+    "    Args:\n",
+    "        boards: Plots the boards given into subplots. The maximum number of boards accepted is 70.\n",
+    "    \"\"\"\n",
+    "    assert len(boards.shape) == 3\n",
+    "    assert boards.shape[1:] == (BOARD_SIZE, BOARD_SIZE)\n",
     "    assert boards.shape[0] < 70\n",
+    "\n",
     "    plots_per_row = 4\n",
     "    rows = int(np.ceil(boards.shape[0] / plots_per_row))\n",
     "    fig, axs = plt.subplots(rows, plots_per_row, figsize=(12, 3 * rows))\n",
@@ -485,7 +496,7 @@
     "\n",
     "\n",
     "def get_possible_turns(boards: np.ndarray) -> np.ndarray:\n",
-    "    \"\"\"Check where turns are possible on a board.\n",
+    "    \"\"\"Analyses a stack of boards.\n",
     "\n",
     "    Args:\n",
     "        boards: A stack of boards to check.\n",
@@ -493,6 +504,9 @@
     "    Returns:\n",
     "        A stack of game boards containing boolean values showing where turns are possible for the player.\n",
     "    \"\"\"\n",
+    "    assert len(boards.shape) == 3\n",
+    "    assert boards.shape[:2] == (BOARD_SIZE, BOARD_SIZE)\n",
+    "\n",
     "    _poss_turns = boards == 0  # checks where fields are empty.\n",
     "    _poss_turns &= binary_dilation(\n",
     "        boards == -1, SURROUNDING\n",