notebooks/02_prepare_dataset_manual.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import tqdm\n",
    "\n",
    "import numpy as np\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": [],
   "source": [
    "SUBSET = \"training\"\n",
    "RANDOM_SEED = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab_sysmaster/datapkg_output_dir')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATAPKG_OUTPUT_DIR = Path(\"~/datapkg_output_dir\").expanduser().resolve()\n",
    "DATAPKG_OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab_sysmaster/datapkg_output_dir/deep-protein-gen')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DEEP_PROTEIN_GEN_DIR = DATAPKG_OUTPUT_DIR / \"deep-protein-gen\"\n",
    "DEEP_PROTEIN_GEN_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab_sysmaster/datapkg_output_dir/deep-protein-gen/training_data')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATA_DIR = DEEP_PROTEIN_GEN_DIR / f\"{SUBSET}_data\"\n",
    "DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/kimlab_sysmaster/datapkg_output_dir/deep-protein-gen/training_data_rs0')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "OUTPUT_DATA_DIR = DEEP_PROTEIN_GEN_DIR / f\"{SUBSET}_data_rs{RANDOM_SEED}\"\n",
    "OUTPUT_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['database_id=G3DSA%3A1.10.10.180',\n",
       " 'database_id=G3DSA%3A1.10.10.190',\n",
       " 'database_id=G3DSA%3A1.10.10.520']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "folder_list = os.listdir(DATA_DIR)\n",
    "folder_list[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['uniparc_id',\n",
       " 'structure_id',\n",
       " 'sequence',\n",
       " 'residue_idx_1_corrected',\n",
       " 'residue_idx_2_corrected',\n",
       " 'distances']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "columns = [\n",
    "    \"uniparc_id\",\n",
    "    \"structure_id\",\n",
    "    \"sequence\",\n",
    "    \"residue_idx_1_corrected\",\n",
    "    \"residue_idx_2_corrected\",\n",
    "    \"distances\",\n",
    "]\n",
    "columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_table(table, output_file, max_tries=5):\n",
    "    n_tries = 0\n",
    "    while True:\n",
    "        n_tries += 1\n",
    "        try:\n",
    "            pq.write_table(table, output_file)\n",
    "            return\n",
    "        except OSError:\n",
    "            if n_tries > max_tries:\n",
    "                raise\n",
    "            else:\n",
    "                time.sleep(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ae314e95b5648da815e131f8ffc62e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=998), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'AssertionError'> \n",
      "1.10.10.190 11 11 11\n",
      "<class 'AssertionError'> \n",
      "1.10.287.120 9 9 9\n",
      "<class 'AssertionError'> \n",
      "1.10.3200.10 29 29 29\n",
      "<class 'AssertionError'> \n",
      "1.10.3400.10 19 19 19\n",
      "<class 'AssertionError'> \n",
      "1.10.3960.10 37 37 37\n",
      "<class 'AssertionError'> \n",
      "1.20.10.10 16 16 16\n",
      "<class 'AssertionError'> \n",
      "1.20.120.590 32 32 32\n",
      "<class 'AssertionError'> \n",
      "1.20.120.630 93 93 93\n",
      "<class 'AssertionError'> \n",
      "1.20.1280.10 57 57 57\n",
      "<class 'AssertionError'> \n",
      "1.20.1400.10 8 8 8\n",
      "<class 'AssertionError'> \n",
      "1.20.1480.10 28 28 28\n",
      "<class 'AssertionError'> \n",
      "1.20.200.20 77 77 77\n",
      "1.20.900.10 42645 3023 3023\n",
      "1.20.91.20 1437 531 531\n",
      "1.20.910.10 41989 2839 2839\n",
      "1.20.920.10 37910 2721 2721\n",
      "1.20.930.10 7256 1226 1226\n",
      "1.20.930.20 850 396 396\n",
      "1.20.960.10 7669 1197 1197\n",
      "1.25.10.10 574964 10803 10803\n",
      "1.25.10.20 14284 1645 1645\n",
      "1.25.20.10 32814 2483 2483\n",
      "1.25.40.170 608 340 340\n",
      "1.25.40.190 2201 656 656\n",
      "1.25.40.250 1485 535 535\n",
      "1.25.40.30 2006 647 647\n",
      "1.25.40.40 1297 504 504\n",
      "1.25.40.90 28804 2353 2353\n",
      "1.25.50.10 10778 1490 1490\n",
      "1.40.10.10 303 240 240\n",
      "1.50.10.100 22821 2129 2129\n",
      "1.50.10.130 8045 1212 1212\n",
      "1.50.40.10 88382 4195 4195\n",
      "2.10.10.10 5388 1035 1035\n",
      "2.10.150.10 9831 1351 1351\n",
      "2.10.240.10 8307 1329 1329\n",
      "2.10.90.10 28411 2321 2321\n",
      "2.102.20.10 2390 693 693\n",
      "2.110.10.10 15813 1760 1760\n",
      "2.115.10.20 111982 4627 4627\n",
      "2.120.10.30 380422 8775 8775\n",
      "2.120.10.80 120415 5112 5112\n",
      "2.130.10.10 1143087 15252 15252\n",
      "2.130.10.80 27026 2340 2340\n",
      "2.150.10.10 218532 6543 6543\n",
      "2.160.20.20 45940 3014 3014\n",
      "2.160.20.50 272488 7457 7457\n",
      "2.160.20.60 26919 2328 2328\n",
      "2.170.11.10 2976 748 748\n",
      "2.170.120.12 8001 1250 1250\n",
      "2.170.130.10 239212 6801 6801\n",
      "2.170.130.20 4161 933 933\n",
      "2.170.140.10 34506 2604 2604\n",
      "2.170.180.11 2019 629 629\n",
      "2.170.210.10 1656 582 582\n",
      "2.170.240.10 3516 846 846\n",
      "2.170.30.10 4007 914 914\n",
      "2.20.100.10 69083 3678 3678\n",
      "2.20.160.10 248 221 221\n",
      "2.20.180.10 721 366 366\n",
      "2.20.20.10 1090 466 466\n",
      "2.20.25.100 2483 686 686\n",
      "2.20.25.30 4489 990 990\n",
      "2.20.28.20 25047 2181 2181\n",
      "2.20.28.40 1810 573 573\n",
      "2.20.50.20 1737 596 596\n",
      "2.20.60.10 386 271 271\n",
      "2.30.110.10 204278 6493 6493\n",
      "2.30.110.20 13677 1704 1704\n",
      "2.30.140.10 6980 1170 1170\n",
      "2.30.140.20 478 319 319\n",
      "2.30.180.10 36492 2751 2751\n",
      "2.30.230.10 3942 863 863\n",
      "2.30.24.10 6331 1074 1074\n",
      "2.30.26.10 314 236 236\n",
      "2.30.270.10 3806 912 912\n",
      "2.30.280.10 4199 937 937\n",
      "2.30.29.50 4530 1008 1008\n",
      "2.30.290.10 3908 860 860\n",
      "2.30.30.130 6126 1101 1101\n",
      "2.30.30.150 288 249 249\n",
      "2.30.30.210 2204 678 678\n",
      "2.30.30.230 7064 1214 1214\n",
      "2.30.30.290 4167 955 955\n",
      "2.30.30.320 1635 539 539\n",
      "2.30.30.360 3383 781 781\n",
      "2.30.30.70 3999 887 887\n",
      "2.30.31.10 8033 1346 1346\n",
      "2.30.32.30 858 408 408\n",
      "2.30.320.10 8618 1356 1356\n",
      "2.30.33.40 24690 2122 2122\n",
      "2.30.330.10 19375 1982 1982\n",
      "2.30.40.10 266217 7450 7450\n",
      "2.30.60.10 2158 650 650\n",
      "2.40.10.250 264 223 223\n",
      "2.40.100.10 94822 4343 4343\n",
      "2.40.128.110 29529 2387 2387\n",
      "2.40.128.120 507 299 299\n",
      "2.40.128.140 3003 817 817\n",
      "2.40.128.20 61637 3599 3599\n",
      "2.40.128.30 2209 675 675\n",
      "2.40.128.50 1255 471 471\n",
      "2.40.128.80 866 394 394\n",
      "2.40.15.10 614 360 360\n",
      "2.40.150.20 19590 2033 2033\n",
      "2.40.160.20 127308 5139 5139\n",
      "2.40.170.10 8794 1293 1293\n",
      "2.40.170.20 352649 8312 8312\n",
      "2.40.250.10 740 381 381\n",
      "2.40.260.10 29603 2401 2401\n",
      "2.40.280.10 18368 1819 1819\n",
      "2.40.30.100 4622 995 995\n",
      "2.40.30.110 40830 2911 2911\n",
      "2.40.30.20 65236 3555 3555\n",
      "2.40.30.40 10260 1378 1378\n",
      "2.40.30.60 17341 1908 1908\n",
      "2.40.30.70 10236 1433 1433\n",
      "2.40.30.80 402 290 290\n",
      "2.40.310.10 369 280 280\n",
      "2.40.33.40 2256 680 680\n",
      "2.40.330.10 21410 2047 2047\n",
      "2.40.340.10 21076 2024 2024\n",
      "2.40.350.10 2147 659 659\n",
      "2.40.37.10 81871 4111 4111\n",
      "2.40.370.10 10211 1457 1457\n",
      "2.40.390.10 3082 762 762\n",
      "2.40.40.10 45626 2981 2981\n",
      "2.40.480.10 6272 1154 1154\n",
      "2.40.50.150 12657 1674 1674\n",
      "2.40.50.220 3718 884 884\n",
      "2.40.50.230 14014 1669 1669\n",
      "2.40.50.240 893 440 440\n",
      "2.40.50.60 75 75 75\n",
      "2.40.70.10 333798 8251 8251\n",
      "2.50.10.10 106 106 106\n",
      "2.50.20.10 48495 3102 3102\n",
      "2.60.11.10 1810 576 576\n",
      "2.60.110.10 7972 1241 1241\n",
      "2.60.120.230 6261 1168 1168\n",
      "2.60.120.240 348 250 250\n",
      "2.60.120.260 437173 9186 9186\n",
      "2.60.120.290 75418 3845 3845\n",
      "2.60.120.310 3996 895 895\n",
      "2.60.120.320 7873 1267 1267\n",
      "2.60.120.330 70319 3732 3732\n",
      "2.60.120.390 339 260 260\n",
      "2.60.120.40 24977 2277 2277\n",
      "2.60.120.400 743 390 390\n",
      "2.60.120.420 141 141 141\n",
      "2.60.120.460 16011 1721 1721\n",
      "2.60.120.470 3314 785 785\n",
      "2.60.120.480 7117 1163 1163\n",
      "2.60.120.530 697 409 409\n",
      "2.60.120.570 385 277 277\n",
      "2.60.120.590 26866 2287 2287\n",
      "2.60.120.600 13650 1608 1608\n",
      "2.60.120.610 2838 696 696\n",
      "2.60.130.10 27300 2274 2274\n",
      "2.60.15.10 21245 1990 1990\n",
      "2.60.169.10 2197 662 662\n",
      "2.60.175.20 4018 886 886\n",
      "2.60.200.30 22732 2118 2118\n",
      "2.60.210.10 22967 2248 2248\n",
      "2.60.220.10 16263 1765 1765\n",
      "2.60.240.10 237 219 219\n",
      "2.60.30.10 534 337 337\n",
      "2.60.300.12 34593 2615 2615\n",
      "2.60.310.10 369 268 268\n",
      "2.60.34.10 46704 2931 2931\n",
      "2.60.370.10 6861 1159 1159\n",
      "2.60.390.10 1444 530 530\n",
      "2.60.40.1250 13978 1687 1687\n",
      "2.60.40.1260 21821 1996 1996\n",
      "2.60.40.1280 8084 1300 1300\n",
      "2.60.40.1390 2794 730 730\n",
      "2.60.40.1400 6384 1134 1134\n",
      "2.60.40.1450 1493 569 569\n",
      "2.60.40.1480 2041 672 672\n",
      "2.60.40.1490 1827 582 582\n",
      "2.60.40.150 138143 5274 5274\n",
      "2.60.40.1600 865 405 405\n",
      "2.60.40.1680 182 182 182\n",
      "2.60.40.1700 1710 616 616\n",
      "2.60.40.180 11509 1501 1501\n",
      "2.60.40.20 452 314 314\n",
      "2.60.40.200 18137 1931 1931\n",
      "2.60.40.230 4181 878 878\n",
      "2.60.40.340 3840 883 883\n",
      "2.60.40.350 2582 773 773\n",
      "2.60.40.380 19451 1945 1945\n",
      "2.60.40.420 250875 7104 7104\n",
      "2.60.40.50 785 391 391\n",
      "2.60.40.630 2180 652 652\n",
      "2.60.40.690 7066 1187 1187\n",
      "2.60.40.700 95 95 95\n",
      "2.60.40.710 3391 793 793\n",
      "2.60.40.780 1356 518 518\n",
      "2.60.40.790 73467 3865 3865\n",
      "2.60.40.820 7708 1199 1199\n",
      "2.60.40.830 2481 711 711\n",
      "2.60.40.840 2932 768 768\n",
      "2.60.410.10 15442 1749 1749\n",
      "2.60.430.10 99 99 99\n",
      "2.60.440.10 2135 650 650\n",
      "2.60.490.10 4165 909 909\n",
      "2.60.510.10 65 65 65\n",
      "2.60.60.20 19471 1959 1959\n",
      "2.60.90.10 676 364 364\n",
      "2.60.98.20 10346 1433 1433\n",
      "2.70.100.10 7367 1219 1219\n",
      "2.70.130.10 13988 1759 1759\n",
      "2.70.170.10 37668 2698 2698\n",
      "2.70.180.10 1181 478 478\n",
      "2.70.20.10 38323 2783 2783\n",
      "2.70.210.12 22515 2095 2095\n",
      "2.70.220.10 5190 1029 1029\n",
      "2.70.240.10 2801 731 731\n",
      "2.70.250.10 24 24 24\n",
      "2.70.40.20 91 91 91\n",
      "2.70.50.40 2087 657 657\n",
      "2.70.9.10 2323 702 702\n",
      "2.70.98.10 114704 4837 4837\n",
      "2.70.98.40 20588 2081 2081\n",
      "3.10.100.10 124956 4978 4978\n",
      "3.10.110.10 61869 3585 3585\n",
      "3.10.120.10 27218 2295 2295\n",
      "3.10.130.10 3143 789 789\n",
      "3.10.180.10 389549 9003 9003\n",
      "3.10.20.250 91 91 91\n",
      "3.10.20.270 232 216 216\n",
      "3.10.20.280 5732 1086 1086\n",
      "3.10.20.30 263657 7256 7256\n",
      "3.10.20.70 33580 2591 2591\n",
      "3.10.20.80 21888 2052 2052\n",
      "3.10.200.10 21560 2073 2073\n",
      "3.10.25.10 31025 2470 2470\n",
      "3.10.250.10 15126 1743 1743\n",
      "3.10.260.10 3235 859 859\n",
      "3.10.28.10 29961 2461 2461\n",
      "3.10.280.10 3060 794 794\n",
      "3.10.290.10 142687 5419 5419\n",
      "3.10.290.30 7720 1252 1252\n",
      "3.10.300.10 8555 1305 1305\n",
      "3.10.320.10 31701 2502 2502\n",
      "3.10.330.10 8549 1293 1293\n",
      "3.10.350.10 134637 5179 5179\n",
      "3.10.380.10 995 441 441\n",
      "3.10.390.10 9362 1323 1323\n",
      "3.10.420.10 9090 1394 1394\n",
      "3.10.440.10 3227 790 790\n",
      "3.10.450.140 1054 496 496\n",
      "3.10.450.20 11 11 11\n",
      "3.10.450.70 6280 1105 1105\n",
      "3.10.460.10 119 119 119\n",
      "3.10.470.10 482 319 319\n",
      "3.10.50.10 4764 961 961\n",
      "3.10.500.10 474 308 308\n",
      "3.10.510.10 509 341 341\n",
      "3.10.530.10 1663 597 597\n",
      "3.10.550.10 311 253 253\n",
      "3.20.100.10 1140 477 477\n",
      "3.20.120.10 604 344 344\n",
      "3.20.130.10 13521 1668 1668\n",
      "3.20.14.10 4136 883 883\n",
      "3.20.180.10 6749 1140 1140\n",
      "3.20.19.10 63496 3514 3514\n",
      "3.20.190.10 31904 2498 2498\n",
      "3.20.20.100 240459 7041 7041\n",
      "3.20.20.105 34155 2649 2649\n",
      "3.20.20.110 67768 3660 3660\n",
      "3.20.20.120 130753 5177 5177\n",
      "3.20.20.20 66053 3552 3552\n",
      "3.20.20.240 30745 2518 2518\n",
      "3.20.20.30 186824 6122 6122\n",
      "3.20.20.300 82272 3976 3976\n",
      "3.20.20.330 40931 2828 2828\n",
      "3.20.20.350 1995 633 633\n",
      "3.20.20.410 19577 2051 2051\n",
      "3.20.20.440 1622 585 585\n",
      "3.20.20.450 174799 5936 5936\n",
      "3.20.20.460 667 362 362\n",
      "3.30.1160.10 2667 739 739\n",
      "3.30.1200.10 8741 1364 1364\n",
      "3.30.1220.10 25047 2195 2195\n",
      "3.30.1250.10 2895 781 781\n",
      "3.30.1280.10 11996 1525 1525\n",
      "3.30.1300.20 468 298 298\n",
      "3.30.1310.10 22796 2103 2103\n",
      "3.30.1320.10 18457 1959 1959\n",
      "3.30.1330.10 102904 4440 4440\n",
      "3.30.1330.110 3858 841 841\n",
      "3.30.1330.20 37662 2763 2763\n",
      "3.30.1330.30 69212 3733 3733\n",
      "3.30.1330.40 90566 4228 4228\n",
      "3.30.1330.60 100547 4491 4491\n",
      "3.30.1330.70 10901 1433 1433\n",
      "3.30.1330.90 28051 2346 2346\n",
      "3.30.1340.10 24963 2271 2271\n",
      "3.30.1340.20 2693 766 766\n",
      "3.30.1350.10 358 271 271\n",
      "3.30.1360.10 37100 2758 2758\n",
      "3.30.1360.100 3119 765 765\n",
      "3.30.1360.110 2692 775 775\n",
      "3.30.1360.20 15630 1867 1867\n",
      "3.30.1360.40 49584 3125 3125\n",
      "3.30.1360.60 27544 2355 2355\n",
      "3.30.1360.70 26185 2333 2333\n",
      "3.30.1360.80 6699 1147 1147\n",
      "3.30.1370.10 78104 4012 4012\n",
      "3.30.1370.20 1942 642 642\n",
      "3.30.1370.70 8036 1256 1256\n",
      "3.30.1370.80 277 229 229\n",
      "3.30.1380.10 48386 3107 3107\n",
      "3.30.1390.20 17406 1885 1885\n",
      "3.30.1400.10 3681 840 840\n",
      "3.30.1440.10 25370 2236 2236\n",
      "3.30.1470.10 870 403 403\n",
      "3.30.1480.10 18235 1911 1911\n",
      "3.30.1490.100 38864 2726 2726\n",
      "3.30.1490.120 3437 795 795\n",
      "3.30.1490.140 688 360 360\n",
      "3.30.1490.40 5420 1083 1083\n",
      "3.30.1490.50 1898 634 634\n",
      "3.30.1490.80 2115 628 628\n",
      "3.30.1500.10 723 398 398\n",
      "3.30.1520.10 33408 2631 2631\n",
      "3.30.1530.10 1238 485 485\n",
      "3.30.1550.10 14627 1718 1718\n",
      "3.30.1570.10 116 116 116\n",
      "3.30.1580.10 1879 606 606\n",
      "3.30.160.100 20232 2030 2030\n",
      "3.30.160.120 96 96 96\n",
      "3.30.160.140 3623 879 879\n",
      "3.30.160.160 7548 1229 1229\n",
      "3.30.160.170 5468 1075 1075\n",
      "3.30.160.220 429 280 280\n",
      "3.30.160.40 18674 1943 1943\n",
      "3.30.160.70 27885 2287 2287\n",
      "3.30.160.80 17 17 17\n",
      "3.30.1600.10 3214 812 812\n",
      "3.30.1610.10 2064 663 663\n",
      "3.30.1660.10 10658 1426 1426\n",
      "3.30.170.10 1870 581 581\n",
      "3.30.1700.10 12155 1565 1565\n",
      "3.30.1750.10 224 211 211\n",
      "3.30.1760.10 348 257 257\n",
      "3.30.1780.10 20024 2051 2051\n",
      "3.30.1790.10 20 20 20\n",
      "3.30.1880.10 53 53 53\n",
      "3.30.190.10 5127 1030 1030\n",
      "3.30.1920.10 245 219 219\n",
      "3.30.1940.10 1760 615 615\n",
      "3.30.1960.10 1840 612 612\n",
      "3.30.1970.10 21 21 21\n",
      "3.30.1980.10 1619 560 560\n",
      "3.30.2000.10 2724 727 727\n",
      "3.30.2020.10 3724 819 819\n",
      "3.30.210.10 6362 1184 1184\n",
      "3.30.2120.10 811 404 404\n",
      "3.30.2190.10 410 277 277\n",
      "3.30.2230.10 5624 1101 1101\n",
      "3.30.2280.10 2380 703 703\n",
      "3.30.2290.10 26315 2365 2365\n",
      "3.30.230.10 323323 8031 8031\n",
      "3.30.230.20 12121 1590 1590\n",
      "3.30.230.30 16731 1748 1748\n",
      "3.30.230.60 1819 581 581\n",
      "3.30.230.70 75859 3851 3851\n",
      "3.30.2380.10 2404 708 708\n",
      "3.30.2390.10 20395 1969 1969\n",
      "3.30.240.20 17766 1871 1871\n",
      "3.30.250.10 16176 1742 1742\n",
      "3.30.260.10 4188 907 907\n",
      "3.30.280.10 8081 1243 1243\n",
      "3.30.30.50 2434 691 691\n",
      "3.30.30.60 681 361 361\n",
      "3.30.300.100 345 264 264\n",
      "3.30.300.130 47348 2950 2950\n",
      "3.30.300.20 167193 5682 5682\n",
      "3.30.300.50 7724 1233 1233\n",
      "3.30.300.70 14755 1713 1713\n",
      "3.30.310.100 1769 564 564\n",
      "3.30.310.110 2016 600 600\n",
      "3.30.310.50 58558 3355 3355\n",
      "3.30.310.60 130 130 130\n",
      "3.30.310.70 8847 1286 1286\n",
      "3.30.350.10 2489 719 719\n",
      "3.30.360.30 1885 632 632\n",
      "3.30.365.10 188284 6218 6218\n",
      "3.30.370.10 7130 1224 1224\n",
      "3.30.379.10 30292 2470 2470\n",
      "3.30.380.10 206 201 201\n",
      "3.30.387.10 1230 489 489\n",
      "3.30.390.60 1561 584 584\n",
      "3.30.40.10 561602 10491 10491\n",
      "3.30.40.20 135 135 135\n",
      "3.30.40.30 673 348 348\n",
      "3.30.420.110 22445 2047 2047\n",
      "3.30.420.140 28392 2336 2336\n",
      "3.30.420.180 9479 1383 1383\n",
      "3.30.428.20 953 434 434\n",
      "3.30.450.40 364173 8575 8575\n",
      "3.30.457.10 29029 2429 2429\n",
      "3.30.465.10 200610 6286 6286\n",
      "3.30.479.20 43109 2915 2915\n",
      "3.30.50.10 45153 3029 3029\n",
      "3.30.56.30 2523 681 681\n",
      "3.30.56.60 511 326 326\n",
      "3.30.565.10 978915 14298 14298\n",
      "3.30.565.30 793 399 399\n",
      "3.30.572.10 28399 2438 2438\n",
      "3.30.62.10 9614 1346 1346\n",
      "3.30.66.10 3668 871 871\n",
      "3.30.67.10 1137 477 477\n",
      "3.30.70.1000 39 39 39\n",
      "3.30.70.1010 2506 684 684\n",
      "3.30.70.1040 550 329 329\n",
      "3.30.70.1050 18792 1936 1936\n",
      "3.30.70.1070 33369 2590 2590\n",
      "3.30.70.1130 2234 654 654\n",
      "3.30.70.1220 1104 473 473\n",
      "3.30.70.1230 77320 4044 4044\n",
      "3.30.70.1240 2690 738 738\n",
      "3.30.70.1250 7594 1260 1260\n",
      "3.30.70.1290 46322 3072 3072\n",
      "3.30.70.1340 3269 845 845\n",
      "3.30.70.1350 39020 2811 2811\n",
      "3.30.70.1360 1169 486 486\n",
      "3.30.70.140 2930 773 773\n",
      "3.30.70.1450 72617 3852 3852\n",
      "3.30.70.1480 158 158 158\n",
      "3.30.70.1490 4225 900 900\n",
      "3.30.70.150 31057 2537 2537\n",
      "3.30.70.340 6329 1134 1134\n",
      "3.30.70.380 22302 2144 2144\n",
      "3.30.70.400 1182 478 478\n",
      "3.30.70.420 5793 1080 1080\n",
      "3.30.70.440 10 10 10\n",
      "3.30.70.470 1612 588 588\n",
      "3.30.70.520 2019 614 614\n",
      "3.30.70.560 28703 2455 2455\n",
      "3.30.70.60 30025 2493 2493\n",
      "3.30.70.600 14469 1722 1722\n",
      "3.30.70.640 20089 2000 2000\n",
      "3.30.70.650 840 406 406\n",
      "3.30.70.660 29046 2334 2334\n",
      "3.30.70.670 2750 725 725\n",
      "3.30.70.680 349 264 264\n",
      "3.30.70.790 13873 1679 1679\n",
      "3.30.70.80 23244 2215 2215\n",
      "3.30.70.860 8669 1259 1259\n",
      "3.30.70.960 11728 1557 1557\n",
      "3.30.70.970 3974 883 883\n",
      "3.30.70.980 38640 2808 2808\n",
      "3.30.700.20 5286 1002 1002\n",
      "3.30.730.10 20732 2006 2006\n",
      "3.30.740.10 5593 1081 1081\n",
      "3.30.750.24 94663 4351 4351\n",
      "3.30.750.60 635 348 348\n",
      "3.30.760.10 7943 1286 1286\n",
      "3.30.780.10 10066 1393 1393\n",
      "3.30.800.10 3115 768 768\n",
      "3.30.810.10 8655 1301 1301\n",
      "3.30.860.10 18339 1948 1948\n",
      "3.30.870.20 2793 756 756\n",
      "3.30.900.10 5415 1008 1008\n",
      "3.30.910.10 4150 927 927\n",
      "3.30.910.20 18354 1910 1910\n",
      "3.30.920.10 5309 976 976\n",
      "3.30.920.20 2806 720 720\n",
      "3.30.950.10 82497 4068 4068\n",
      "3.30.960.10 1905 580 580\n",
      "3.30.990.10 3297 813 813\n",
      "3.40.10.10 16849 1839 1839\n",
      "3.40.1000.20 53 53 53\n",
      "3.40.1010.10 95485 4396 4396\n",
      "3.40.1050.10 35579 2577 2577\n",
      "3.40.1060.10 54763 3363 3363\n",
      "3.40.109.10 118196 4858 4858\n",
      "3.40.1120.10 3571 816 816\n",
      "3.40.1160.10 133131 5249 5249\n",
      "3.40.1170.10 24121 2199 2199\n",
      "3.40.1180.10 38604 2783 2783\n",
      "3.40.1230.10 6812 1210 1210\n",
      "3.40.1260.10 32442 2656 2656\n",
      "3.40.1370.10 40361 2776 2776\n",
      "3.40.1380.20 26378 2261 2261\n",
      "3.40.1390.20 14915 1781 1781\n",
      "3.40.140.20 46158 3123 3123\n",
      "3.40.140.30 553 308 308\n",
      "3.40.1400.10 23115 2157 2157\n",
      "3.40.1410.10 64775 3544 3544\n",
      "3.40.1440.10 62156 3496 3496\n",
      "3.40.1480.10 7063 1210 1210\n",
      "3.40.1490.10 9386 1400 1400\n",
      "3.40.1510.10 897 430 430\n",
      "3.40.1520.10 1615 599 599\n",
      "3.40.1530.10 84 84 84\n",
      "3.40.1540.10 447 296 296\n",
      "3.40.1550.10 21590 2130 2130\n",
      "3.40.1560.10 62 62 62\n",
      "3.40.1610.10 3741 845 845\n",
      "3.40.1650.10 10404 1488 1488\n",
      "3.40.1690.10 15504 1745 1745\n",
      "3.40.1700.10 8573 1306 1306\n",
      "3.40.1760.10 2261 668 668\n",
      "3.40.1790.10 8142 1305 1305\n",
      "3.40.1810.10 14777 1767 1767\n",
      "3.40.1830.10 8736 1281 1281\n",
      "3.40.1840.10 2881 754 754\n",
      "3.40.198.10 430 288 288\n",
      "3.40.225.10 60963 3558 3558\n",
      "3.40.250.10 189747 6220 6220\n",
      "3.40.35.10 10191 1387 1387\n",
      "3.40.350.10 74599 3884 3884\n",
      "3.40.390.10 164664 5815 5815\n",
      "3.40.390.30 23439 2056 2056\n",
      "3.40.430.10 100984 4567 4567\n",
      "3.40.462.10 11576 1526 1526\n",
      "3.40.47.10 694315 11759 11759\n",
      "3.40.5.10 14001 1666 1666\n",
      "3.40.5.20 183 183 183\n",
      "3.40.50.1010 174186 5928 5928\n",
      "3.40.50.10140 45649 3009 3009\n",
      "3.40.50.10150 2112 595 595\n",
      "3.40.50.10190 71901 3603 3603\n",
      "3.40.50.10220 2235 634 634\n",
      "3.40.50.10230 10460 1425 1425\n",
      "3.40.50.10240 9622 1360 1360\n",
      "3.40.50.10300 28516 2388 2388\n",
      "3.40.50.10310 13626 1691 1691\n",
      "3.40.50.10320 39494 2846 2846\n",
      "3.40.50.10360 3318 853 853\n",
      "3.40.50.10380 31046 2535 2535\n",
      "3.40.50.10400 5187 987 987\n",
      "3.40.50.10540 100412 4575 4575\n",
      "3.40.50.10610 9230 1348 1348\n",
      "3.40.50.10740 17738 1830 1830\n",
      "3.40.50.10780 4103 911 911\n",
      "3.40.50.10790 5520 1031 1031\n",
      "3.40.50.10840 12028 1539 1539\n",
      "3.40.50.1170 24027 2182 2182\n",
      "3.40.50.1220 184471 6076 6076\n",
      "3.40.50.1240 169358 5939 5939\n",
      "3.40.50.1260 56672 3359 3359\n",
      "3.40.50.1370 101373 4561 4561\n",
      "3.40.50.1380 53440 3228 3228\n",
      "3.40.50.1450 15163 1720 1720\n",
      "3.40.50.170 86417 4208 4208\n",
      "3.40.50.1760 2544 691 691\n",
      "3.40.50.180 24732 2162 2162\n",
      "3.40.50.1950 45014 2951 2951\n",
      "3.40.50.200 126800 4885 4885\n",
      "3.40.50.261 79440 4030 4030\n",
      "3.40.50.360 206851 6469 6469\n",
      "3.40.50.40 18714 1933 1933\n",
      "3.40.50.410 193328 6387 6387\n",
      "3.40.50.510 21587 2180 2180\n",
      "3.40.50.670 48726 3185 3185\n",
      "3.40.50.700 25960 2308 2308\n",
      "3.40.50.7700 30298 2530 2530\n",
      "3.40.50.790 16937 1869 1869\n",
      "3.40.50.850 90300 4255 4255\n",
      "3.40.50.880 471743 9733 9733\n",
      "3.40.50.9100 20843 2091 2091\n",
      "3.40.50.960 22731 2141 2141\n",
      "3.40.525.10 33369 2579 2579\n",
      "3.40.532.10 6664 1155 1155\n",
      "3.40.580.10 491 307 307\n",
      "3.40.605.10 365753 8472 8472\n",
      "3.40.630.20 11394 1553 1553\n",
      "3.40.80.10 47465 3091 3091\n",
      "3.40.800.20 47539 2944 2944\n",
      "3.40.910.10 8208 1284 1284\n",
      "3.40.920.10 29738 2378 2378\n",
      "3.40.930.10 58272 3404 3404\n",
      "3.40.970.10 5946 1087 1087\n",
      "3.40.980.10 74669 3840 3840\n",
      "3.50.14.10 2505 697 697\n",
      "3.50.20.10 1849 597 597\n",
      "3.50.30.50 25239 2185 2185\n",
      "3.50.30.60 16696 1843 1843\n",
      "3.50.30.70 1155 501 501\n",
      "3.50.40.10 36717 2717 2717\n",
      "3.50.80.10 21481 2009 2009\n",
      "3.55.10.10 3414 837 837\n",
      "3.55.20.10 18 18 18\n",
      "3.55.30.10 13512 1702 1702\n",
      "3.60.100.10 292 240 240\n",
      "3.60.110.10 121762 4955 4955\n",
      "3.60.120.10 81876 4009 4009\n",
      "3.60.140.10 26594 2336 2336\n",
      "3.60.150.10 32297 2535 2535\n",
      "3.60.20.10 228250 6762 6762\n",
      "3.60.21.10 364843 8567 8567\n",
      "3.60.70.12 32685 2605 2605\n",
      "3.60.9.10 6878 1183 1183\n",
      "3.60.90.10 10609 1506 1506\n",
      "3.65.10.10 109365 4651 4651\n",
      "3.65.10.20 7405 1250 1250\n",
      "3.75.10.20 6352 1122 1122\n",
      "3.80.10.10 1112392 15045 15045\n",
      "3.80.20.20 25622 2225 2225\n",
      "3.90.100.10 12489 1563 1563\n",
      "3.90.1030.10 18240 1890 1890\n",
      "3.90.1110.10 13154 1644 1644\n",
      "3.90.1140.10 34674 2628 2628\n",
      "3.90.1150.10 994570 14206 14206\n",
      "3.90.1150.20 63 63 63\n",
      "3.90.1170.20 23867 2177 2177\n",
      "3.90.1170.40 17298 1902 1902\n",
      "3.90.1170.50 48284 3155 3155\n",
      "3.90.1220.10 2707 717 717\n",
      "3.90.1290.10 16931 1828 1828\n",
      "3.90.1300.10 12495 1553 1553\n",
      "3.90.1330.10 2408 728 728\n",
      "3.90.1340.10 12591 1667 1667\n",
      "3.90.1360.10 167 167 167\n",
      "3.90.1370.10 101 101 101\n",
      "3.90.1380.10 13758 1633 1633\n",
      "3.90.1400.10 806 394 394\n",
      "3.90.1420.10 4524 940 940\n",
      "3.90.1460.10 1798 610 610\n",
      "3.90.148.10 221 212 212\n",
      "3.90.15.10 7436 1202 1202\n",
      "3.90.1560.10 5070 1055 1055\n",
      "3.90.1600.10 24451 2247 2247\n",
      "3.90.1650.10 457 285 285\n",
      "3.90.1670.10 5150 1028 1028\n",
      "3.90.1680.10 28259 2430 2430\n",
      "3.90.1780.10 1047 452 452\n",
      "3.90.190.10 123809 5059 5059\n",
      "3.90.190.20 133101 5082 5082\n",
      "3.90.199.10 44908 2945 2945\n",
      "3.90.209.20 37279 2729 2729\n",
      "3.90.215.10 24042 2215 2215\n",
      "3.90.245.10 34324 2660 2660\n",
      "3.90.280.10 25503 2258 2258\n",
      "3.90.290.10 7945 1294 1294\n",
      "3.90.310.10 2368 709 709\n",
      "3.90.320.10 63713 3649 3649\n",
      "3.90.320.20 364 271 271\n",
      "3.90.330.10 4920 969 969\n",
      "3.90.350.10 54106 3291 3291\n",
      "3.90.360.10 1667 553 553\n",
      "3.90.390.10 316 248 248\n",
      "3.90.420.10 36719 2764 2764\n",
      "3.90.430.10 702 421 421\n",
      "3.90.45.10 40786 2804 2804\n",
      "3.90.460.10 1669 587 587\n",
      "3.90.470.10 24911 2212 2212\n",
      "3.90.470.20 62768 3560 3560\n",
      "3.90.520.10 2830 805 805\n",
      "3.90.530.10 1951 601 601\n",
      "3.90.550.10 844526 13180 13180\n",
      "3.90.56.10 1155 461 461\n",
      "3.90.570.10 1200 500 500\n",
      "3.90.600.10 22950 2222 2222\n",
      "3.90.640.20 8100 1306 1306\n",
      "3.90.650.10 111182 4729 4729\n",
      "3.90.730.10 11085 1534 1534\n",
      "3.90.740.10 63335 3616 3616\n",
      "3.90.760.10 1665 624 624\n",
      "3.90.770.10 15322 1758 1758\n",
      "3.90.78.10 22027 2088 2088\n",
      "3.90.80.10 19420 1962 1962\n",
      "3.90.810.10 5776 1053 1053\n",
      "3.90.840.10 1767 579 579\n",
      "3.90.850.10 92940 4312 4312\n",
      "3.90.870.10 76323 3767 3767\n",
      "3.90.930.1 618 341 341\n",
      "3.90.930.12 44576 2945 2945\n",
      "3.90.940.10 13975 1695 1695\n",
      "3.90.940.20 2515 708 708\n",
      "3.90.950.10 72941 3886 3886\n",
      "3.90.960.10 43981 2960 2960\n",
      "3.90.970.10 528 327 327\n",
      "4.10.10.10 569 341 341\n",
      "4.10.1000.10 55431 3365 3365\n",
      "4.10.1010.10 258 236 236\n",
      "4.10.1040.10 1690 599 599\n",
      "4.10.1050.10 14227 1695 1695\n",
      "4.10.1080.10 34373 2654 2654\n",
      "4.10.110.20 1316 484 484\n",
      "4.10.1110.10 8885 1364 1364\n",
      "4.10.1120.10 348 247 247\n",
      "4.10.1150.10 1072 504 504\n",
      "4.10.1170.10 802 405 405\n",
      "4.10.1210.10 1050 490 490\n",
      "4.10.1230.10 1153 460 460\n",
      "4.10.1240.10 8404 1280 1280\n",
      "4.10.140.10 336 260 260\n",
      "4.10.160.10 183 183 183\n",
      "4.10.20.10 8615 1298 1298\n",
      "4.10.220.20 723 373 373\n",
      "4.10.230.10 1146 476 476\n",
      "4.10.240.10 76032 3942 3942\n",
      "4.10.280.10 50951 3096 3096\n",
      "4.10.320.10 39814 2875 2875\n",
      "4.10.372.10 1937 636 636\n",
      "4.10.410.10 29588 2384 2384\n",
      "4.10.480.10 3438 862 862\n",
      "4.10.49.10 1102 484 484\n",
      "4.10.490.10 1828 615 615\n",
      "4.10.540.10 591 338 338\n",
      "4.10.60.10 123429 5018 5018\n",
      "4.10.640.10 15312 1733 1733\n",
      "4.10.670.10 164 164 164\n",
      "4.10.680.10 360 258 258\n",
      "4.10.720.10 254 223 223\n",
      "4.10.740.10 3203 805 805\n",
      "4.10.75.10 9041 1329 1329\n",
      "4.10.780.10 3746 850 850\n",
      "4.10.80.10 2117 678 678\n",
      "4.10.81.10 575 329 329\n",
      "4.10.820.10 86 86 86\n",
      "4.10.860.20 823 406 406\n",
      "4.10.870.10 979 468 468\n",
      "4.10.91.10 1551 567 567\n",
      "4.10.93.10 946 401 401\n",
      "4.10.95.10 1920 610 610\n",
      "4.10.970.10 670 371 371\n",
      "\n"
     ]
    }
   ],
   "source": [
    "random_state = np.random.RandomState(RANDOM_SEED)\n",
    "\n",
    "num_rows_per_group = 200\n",
    "\n",
    "for folder in tqdm.tqdm_notebook(folder_list, total=len(folder_list)):\n",
    "    file_list = list((DATA_DIR / folder).glob(\"*.parquet\"))\n",
    "    assert len(file_list) == 1\n",
    "\n",
    "    database_id = folder.replace(\"database_id=G3DSA%3A\", \"\")\n",
    "\n",
    "    output_file = OUTPUT_DATA_DIR / f\"{database_id}.parquet\"\n",
    "    if output_file.is_file():\n",
    "        try:\n",
    "            df = pq.read_table(output_file, use_threads=False).to_pandas(\n",
    "                integer_object_nulls=True\n",
    "            )\n",
    "            assert len(df) > 0\n",
    "            continue\n",
    "        except Exception as e:\n",
    "            print(type(e), e)\n",
    "\n",
    "    pqfile = pq.ParquetFile(file_list[0])\n",
    "    num_rows = pqfile.metadata.num_rows\n",
    "    mask = random_state.rand(num_rows) < (\n",
    "        num_rows_per_group / num_rows * max(1, num_rows / num_rows_per_group) ** 0.5\n",
    "    )\n",
    "\n",
    "    dfs = []\n",
    "    start = 0\n",
    "    for row_group in range(pqfile.num_row_groups):\n",
    "        df = pqfile.read_row_group(\n",
    "            row_group, use_threads=False, columns=columns\n",
    "        ).to_pandas(integer_object_nulls=True)\n",
    "        end = start + len(df)\n",
    "        df = df.loc[mask[start:end], :]\n",
    "        assert len(df.index) == len(df.index.drop_duplicates())\n",
    "        dfs.append(df)\n",
    "        start = end\n",
    "    assert start == num_rows\n",
    "    df = pd.concat(dfs, ignore_index=True, sort=False)\n",
    "    df[\"database_id\"] = database_id\n",
    "\n",
    "    print(database_id, num_rows, mask.sum(), len(df))\n",
    "    table = pa.Table.from_pandas(df, preserve_index=False)\n",
    "    write_table(table, output_file)"
   ]
  },
  {
   "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
}