notebooks/20_sudoku_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook we load a network trained to solve Sudoku puzzles and use this network to solve a single Sudoku.\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eoBMUoW2Hvhp"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TbKxMUZWHvhq"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import tempfile\n",
    "import time\n",
    "from collections import deque\n",
    "from pathlib import Path\n",
    "\n",
    "import tqdm\n",
    "from torch_geometric.data import DataLoader\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from matplotlib import cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 334
    },
    "colab_type": "code",
    "id": "rczy7pPiHvhs",
    "outputId": "8e4673d7-4cc7-4f8e-ffc4-de94302e5fab"
   },
   "outputs": [],
   "source": [
    "import proteinsolver\n",
    "import proteinsolver.datasets\n",
    "from proteinsolver.utils import gen_sudoku_graph_featured"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "try:\n",
    "    inline_rc\n",
    "except NameError:\n",
    "    inline_rc = mpl.rcParams.copy()\n",
    "    \n",
    "mpl.rcParams.update({\"font.size\": 12})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dsEY3dtLHvhy"
   },
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNIQUE_ID = \"c8de7e56\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path(tempfile.gettempdir())\n",
    "DATA_ROOT.mkdir(exist_ok=True)\n",
    "DATA_ROOT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NOTEBOOK_NAME = \"sudoku_demo\"\n",
    "NOTEBOOK_PATH = Path(NOTEBOOK_NAME)\n",
    "NOTEBOOK_PATH.mkdir(exist_ok=True)\n",
    "NOTEBOOK_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATAPKG_DATA_DIR = Path(f\"~/datapkg_data_dir\").expanduser().resolve()\n",
    "DATAPKG_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proteinsolver.settings.data_url = DATAPKG_DATA_DIR.as_posix()\n",
    "proteinsolver.settings.data_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run sudoku_train/{UNIQUE_ID}/model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_files = sorted(\n",
    "    Path(\"sudoku_train\").joinpath(UNIQUE_ID).glob(\"*.state\"),\n",
    "    key=lambda s: (int(s.stem.split(\"-\")[3].strip(\"amv\")), int(s.stem.split(\"-\")[2].strip(\"d\"))),\n",
    ")\n",
    "\n",
    "state_files[-5:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_file = state_files[-1]\n",
    "state_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net(\n",
    "    x_input_size=13, adj_input_size=3, hidden_size=162, output_size=9, batch_size=8\n",
    ").to(device)\n",
    "\n",
    "net.load_state_dict(torch.load(state_file, map_location=device))\n",
    "net = net.eval()\n",
    "net = net.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_puzzle(puzzle):\n",
    "    puzzle = puzzle - 1\n",
    "    puzzle = torch.where(puzzle >= 0, puzzle, torch.tensor(9))\n",
    "    return puzzle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode_puzzle(puzzle):\n",
    "    puzzle = (puzzle + 1) % 10\n",
    "    return puzzle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "puzzle = torch.tensor([1,1,1])\n",
    "assert torch.equal(decode_puzzle(encode_puzzle(puzzle)), puzzle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_sudoku(net, puzzle):\n",
    "    sudoku_graph = torch.from_numpy(gen_sudoku_graph_featured()).to_sparse(2)\n",
    "    edge_index = sudoku_graph.indices()\n",
    "    edge_attr = sudoku_graph.values()\n",
    "\n",
    "    output = net(\n",
    "        encode_puzzle(puzzle).view(-1).to(device), edge_index.clone().to(device), edge_attr.clone().to(device)\n",
    "    ).to(\"cpu\")\n",
    "    output = torch.softmax(output, dim=1)\n",
    "    _, predicted = output.max(dim=1)\n",
    "\n",
    "    return decode_puzzle(predicted).reshape(9, 9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_sudoku(puzzle, solved=None, pred=None, ax=None):\n",
    "    # Simple plotting statement that ingests a 9x9 array (n), and plots a sudoku-style grid around it.\n",
    "    if ax is None:\n",
    "        _, ax = plt.subplots()\n",
    "\n",
    "    for y in range(10):\n",
    "        ax.plot([-0.05, 9.05], [y, y], color=\"black\", linewidth=1)\n",
    "\n",
    "    for y in range(0, 10, 3):\n",
    "        ax.plot([-0.05, 9.05], [y, y], color=\"black\", linewidth=3)\n",
    "\n",
    "    for x in range(10):\n",
    "        ax.plot([x, x], [-0.05, 9.05], color=\"black\", linewidth=1)\n",
    "\n",
    "    for x in range(0, 10, 3):\n",
    "        ax.plot([x, x], [-0.05, 9.05], color=\"black\", linewidth=3)\n",
    "\n",
    "    ax.axis(\"image\")\n",
    "    ax.axis(\"off\")  # drop the axes, they're not important here\n",
    "\n",
    "    for x in range(9):\n",
    "        for y in range(9):\n",
    "            puzzle_element = puzzle[8 - y][\n",
    "                x\n",
    "            ]  # need to reverse the y-direction for plotting\n",
    "            if puzzle_element > 0:  # ignore the zeros\n",
    "                T = f\"{puzzle_element}\"\n",
    "                ax.text(x + 0.3, y + 0.2, T, fontsize=20)\n",
    "            elif solved is not None and pred is not None:\n",
    "                solved_element = solved[8 - y][x]\n",
    "                pred_element = pred[8 - y][x]\n",
    "                if solved_element == pred_element:\n",
    "                    T = f\"{solved_element}\"\n",
    "                    ax.text(x + 0.3, y + 0.2, T, fontsize=20, color=\"b\")\n",
    "                else:\n",
    "                    ax.text(x + 0.1, y + 0.3, f\"{pred_element}\", fontsize=13, color=\"r\")\n",
    "                    ax.text(\n",
    "                        x + 0.55, y + 0.3, f\"{solved_element}\", fontsize=13, color=\"g\"\n",
    "                    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve a simple puzzle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input sudoku puzzle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "puzzle = torch.tensor(\n",
    "    [\n",
    "        [0, 8, 0, 0, 3, 2, 0, 0, 1],\n",
    "        [7, 0, 3, 0, 8, 0, 0, 0, 2],\n",
    "        [5, 0, 0, 0, 0, 7, 0, 3, 0],\n",
    "        [0, 5, 0, 0, 0, 1, 9, 7, 0],\n",
    "        [6, 0, 0, 7, 0, 9, 0, 0, 8],\n",
    "        [0, 4, 7, 2, 0, 0, 0, 5, 0],\n",
    "        [0, 2, 0, 6, 0, 0, 0, 0, 9],\n",
    "        [8, 0, 0, 0, 9, 0, 3, 0, 5],\n",
    "        [3, 0, 0, 8, 2, 0, 0, 1, 0],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input solution (for reference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = torch.tensor(\n",
    "    [\n",
    "        [4, 8, 9, 5, 3, 2, 7, 6, 1],\n",
    "        [7, 1, 3, 4, 8, 6, 5, 9, 2],\n",
    "        [5, 6, 2, 9, 1, 7, 8, 3, 4],\n",
    "        [2, 5, 8, 3, 4, 1, 9, 7, 6],\n",
    "        [6, 3, 1, 7, 5, 9, 2, 4, 8],\n",
    "        [9, 4, 7, 2, 6, 8, 1, 5, 3],\n",
    "        [1, 2, 5, 6, 7, 3, 4, 8, 9],\n",
    "        [8, 7, 6, 1, 9, 4, 3, 2, 5],\n",
    "        [3, 9, 4, 8, 2, 5, 6, 1, 7],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert proteinsolver.utils.sudoku.sudoku_is_solved(solution.to(\"cpu\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_sudoku(puzzle, solution, solution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use trained network to solve puzzle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solution_pred = solve_sudoku(net, puzzle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proteinsolver.utils.sudoku.sudoku_is_solved(solution_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_sudoku(puzzle, solution, solution_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "1fXJz1l9HvhW",
    "u9q4iTifHvhl"
   ],
   "name": "Copy of 2019-03-30-sudoku-4xEdgeConv-09862+.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}