notebooks/02_generate_difficult_sudokus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import tqdm\n",
    "\n",
    "import gcsfs\n",
    "import pandas as pd\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/data/workspace/proteinsolver/notebooks/generate_difficult_sudokus')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NOTEBOOK_PATH = Path(\"generate_difficult_sudokus\").resolve()\n",
    "NOTEBOOK_PATH.mkdir(exist_ok=True)\n",
    "NOTEBOOK_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_EXISTING_VALIDATION_DATASET = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_LOCAL_DATA_DIR = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('deep-protein-gen', PosixPath('/data/datapkg_data_dir/deep-protein-gen'))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "project = 'deep-protein-gen'\n",
    "bucket = 'deep-protein-gen'\n",
    "if USE_LOCAL_DATA_DIR:\n",
    "    bucket = Path(f\"~/datapkg_data_dir/{bucket}\").expanduser().resolve()\n",
    "    bucket.joinpath(\"sudoku_difficult\").mkdir(exist_ok=True)\n",
    "\n",
    "project, bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "if USE_LOCAL_DATA_DIR:\n",
    "    fs = None\n",
    "else:\n",
    "    fs = gcsfs.GCSFileSystem(project=project)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "beluga\tniagara-1  old_951\t\tstrokach-cloudtop-2\n",
      "new\tniagara-2  strokach-cloudtop-1\n"
     ]
    }
   ],
   "source": [
    "!ls generate_difficult_sudokus/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read all generated sudokus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "271\n",
      "[PosixPath('/data/workspace/proteinsolver/notebooks/generate_difficult_sudokus/beluga/sodoku_20070_beluga.csv'), PosixPath('/data/workspace/proteinsolver/notebooks/generate_difficult_sudokus/beluga/sodoku_20121_beluga.csv'), PosixPath('/data/workspace/proteinsolver/notebooks/generate_difficult_sudokus/beluga/sodoku_20090_beluga.csv')]\n"
     ]
    }
   ],
   "source": [
    "sudoku_files = (\n",
    "    list(NOTEBOOK_PATH.joinpath(\"beluga\").glob(\"*.csv\")) + \n",
    "    list(NOTEBOOK_PATH.joinpath(\"niagara-1\").glob(\"*.csv\")) + \n",
    "    list(NOTEBOOK_PATH.joinpath(\"niagara-2\").glob(\"*.csv\")) + \n",
    "    list(NOTEBOOK_PATH.joinpath(\"strokach-cloudtop-1\").glob(\"*.csv\")) +\n",
    "    list(NOTEBOOK_PATH.joinpath(\"strokach-cloudtop-2\").glob(\"*.csv\"))\n",
    ")\n",
    "print(len(sudoku_files))\n",
    "print(sudoku_files[:3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3a39d94277c4784849c505140e5a656",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=271), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "dfs = []\n",
    "for file in tqdm.tqdm_notebook(sudoku_files, total=len(sudoku_files)):\n",
    "    df = pd.read_csv(file)\n",
    "    df = df.drop_duplicates()\n",
    "    df[\"is_951\"] = (df[\"difficulty\"] == 951).all()\n",
    "    dfs.append(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "26776593\n"
     ]
    }
   ],
   "source": [
    "df = pd.concat(dfs, ignore_index=True, sort=False)\n",
    "print(len(df))  # 21482631"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23860522\n"
     ]
    }
   ],
   "source": [
    "df = df.sort_values(\"is_951\").drop_duplicates(subset=[\"puzzle\"], keep=\"first\")\n",
    "print(len(df))  # 18566560"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>puzzle</th>\n",
       "      <th>solution</th>\n",
       "      <th>difficulty</th>\n",
       "      <th>is_951</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>13388296</th>\n",
       "      <td>0300700500000100260045001000907000002003010080...</td>\n",
       "      <td>1326798545894137267645821938937654122573419686...</td>\n",
       "      <td>655</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17827660</th>\n",
       "      <td>0900564004023000000000009000500280040306070508...</td>\n",
       "      <td>1982564734623795813758149627591286342316478598...</td>\n",
       "      <td>453</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                     puzzle  \\\n",
       "13388296  0300700500000100260045001000907000002003010080...   \n",
       "17827660  0900564004023000000000009000500280040306070508...   \n",
       "\n",
       "                                                   solution  difficulty  \\\n",
       "13388296  1326798545894137267645821938937654122573419686...         655   \n",
       "17827660  1982564734623795813758149627591286342316478598...         453   \n",
       "\n",
       "          is_951  \n",
       "13388296   False  \n",
       "17827660   False  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23860522\n"
     ]
    }
   ],
   "source": [
    "display(df.head(2))\n",
    "print(len(df))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a validation dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "if not USE_EXISTING_VALIDATION_DATASET:\n",
    "    validation_df = df.sample(n=1_000, random_state=42)\n",
    "else:\n",
    "    validation_df = pq.read_table(f\"{bucket}/sudoku_difficult/valid_0.parquet\").to_pandas()\n",
    "validation_solutions = set(validation_df[\"solution\"])\n",
    "print(len(validation_solutions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Puzzles that are not in validation are in training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_df = df[~df[\"solution\"].isin(validation_solutions)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23859522"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(training_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(training_df) <= len(df) - len(validation_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write final datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_table(df, outfile, filesystem, chunk_size=100_000):\n",
    "    writer = None\n",
    "    for i in range(0, len(df), chunk_size):\n",
    "        chunk = df[i : i + chunk_size]\n",
    "        table = pa.Table.from_pandas(chunk, preserve_index=False)\n",
    "        if writer is None:\n",
    "            writer = pq.ParquetWriter(\n",
    "                outfile, schema=table.schema, filesystem=filesystem\n",
    "            )\n",
    "        writer.write_table(table)\n",
    "    if writer is not None:\n",
    "        writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Chunk size: 1192977.\n",
      "Chunk 0 size: 1192977.\n",
      "Chunk 1 size: 1192977.\n",
      "Chunk 2 size: 1192977.\n",
      "Chunk 3 size: 1192977.\n",
      "Chunk 4 size: 1192977.\n",
      "Chunk 5 size: 1192977.\n",
      "Chunk 6 size: 1192977.\n",
      "Chunk 7 size: 1192977.\n",
      "Chunk 8 size: 1192977.\n",
      "Chunk 9 size: 1192977.\n",
      "Chunk 10 size: 1192977.\n",
      "Chunk 11 size: 1192977.\n",
      "Chunk 12 size: 1192977.\n",
      "Chunk 13 size: 1192977.\n",
      "Chunk 14 size: 1192977.\n",
      "Chunk 15 size: 1192977.\n",
      "Chunk 16 size: 1192977.\n",
      "Chunk 17 size: 1192977.\n",
      "Chunk 18 size: 1192977.\n",
      "Chunk 19 size: 1192959.\n"
     ]
    }
   ],
   "source": [
    "num_chunks = 20\n",
    "\n",
    "tmp_df = training_df.sample(frac=1, random_state=42)\n",
    "chunk_size = len(tmp_df) // num_chunks + 1\n",
    "print(f\"Chunk size: {chunk_size}.\")\n",
    "\n",
    "for i in range(num_chunks):\n",
    "    chunk_df = tmp_df.iloc[chunk_size * i:chunk_size * (i + 1)]\n",
    "    print(f\"Chunk {i} size: {len(chunk_df)}.\")\n",
    "    write_table(chunk_df, f\"{bucket}/sudoku_difficult/train_{i}.parquet\", filesystem=fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not USE_EXISTING_VALIDATION_DATASET:\n",
    "    write_table(validation_df, f\"{bucket}/sudoku_difficult/valid_0.parquet\", filesystem=fs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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
}