testing_loading_data_batches.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:15:24.938802Z",
     "iopub.status.busy": "2024-04-11T10:15:24.938419Z",
     "iopub.status.idle": "2024-04-11T10:16:30.665186Z",
     "shell.execute_reply": "2024-04-11T10:16:30.661491Z",
     "shell.execute_reply.started": "2024-04-11T10:15:24.938770Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting supervenn\n",
      "  Downloading supervenn-0.5.0-py3-none-any.whl.metadata (1.2 kB)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from supervenn) (1.26.4)\n",
      "Requirement already satisfied: matplotlib>=2.2.5 in /opt/conda/lib/python3.10/site-packages (from supervenn) (3.7.5)\n",
      "Requirement already satisfied: pandas in /opt/conda/lib/python3.10/site-packages (from supervenn) (2.2.1)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (1.2.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (4.47.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (1.4.5)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (21.3)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (9.5.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (3.1.1)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.2.5->supervenn) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/conda/lib/python3.10/site-packages (from pandas->supervenn) (2023.3.post1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /opt/conda/lib/python3.10/site-packages (from pandas->supervenn) (2023.4)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib>=2.2.5->supervenn) (1.16.0)\n",
      "Downloading supervenn-0.5.0-py3-none-any.whl (17 kB)\n",
      "Installing collected packages: supervenn\n",
      "Successfully installed supervenn-0.5.0\n",
      "Collecting upsetplot\n",
      "  Downloading UpSetPlot-0.9.0.tar.gz (23 kB)\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: pandas>=0.23 in /opt/conda/lib/python3.10/site-packages (from upsetplot) (2.2.1)\n",
      "Requirement already satisfied: matplotlib>=2.0 in /opt/conda/lib/python3.10/site-packages (from upsetplot) (3.7.5)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (1.2.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (4.47.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (1.4.5)\n",
      "Requirement already satisfied: numpy<2,>=1.20 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (1.26.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (21.3)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (9.5.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (3.1.1)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.10/site-packages (from matplotlib>=2.0->upsetplot) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/conda/lib/python3.10/site-packages (from pandas>=0.23->upsetplot) (2023.3.post1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /opt/conda/lib/python3.10/site-packages (from pandas>=0.23->upsetplot) (2023.4)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib>=2.0->upsetplot) (1.16.0)\n",
      "Building wheels for collected packages: upsetplot\n",
      "  Building wheel for upsetplot (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for upsetplot: filename=UpSetPlot-0.9.0-py3-none-any.whl size=24817 sha256=38438c63fbe8cdb82838e6f2b166e6c71122c66477f9085774a295c50bfde16d\n",
      "  Stored in directory: /root/.cache/pip/wheels/73/42/9f/1c9718ea27f30466d2787e0f7d88a7cb11942e3460c17e0ef6\n",
      "Successfully built upsetplot\n",
      "Installing collected packages: upsetplot\n",
      "Successfully installed upsetplot-0.9.0\n"
     ]
    }
   ],
   "source": [
    "# Install librariess\n",
    "!pip install supervenn\n",
    "!pip install upsetplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:16:30.674689Z",
     "iopub.status.busy": "2024-04-11T10:16:30.673041Z",
     "iopub.status.idle": "2024-04-11T10:16:30.897252Z",
     "shell.execute_reply": "2024-04-11T10:16:30.892834Z",
     "shell.execute_reply.started": "2024-04-11T10:16:30.674634Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "../input/leash-BELKA/sample_submission.csv\n",
      "../input/leash-BELKA/train.parquet\n",
      "../input/leash-BELKA/test.parquet\n"
     ]
    }
   ],
   "source": [
    "# This Python 3 environment comes with many helpful analytics libraries installed\n",
    "# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n",
    "# For example, here's several helpful packages to load\n",
    "\n",
    "import numpy as np # linear algebra\n",
    "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "\n",
    "# Input data files are available in the read-only \"../input/\" directory\n",
    "# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "from tqdm.auto import tqdm\n",
    "from pathlib import Path\n",
    "from supervenn import supervenn\n",
    "import os\n",
    "for dirname, _, filenames in os.walk('../input'):\n",
    "    for filename in filenames[:3]:\n",
    "        print(os.path.join(dirname, filename))\n",
    "\n",
    "# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n",
    "# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:16:30.906071Z",
     "iopub.status.busy": "2024-04-11T10:16:30.905017Z",
     "iopub.status.idle": "2024-04-11T10:16:30.920093Z",
     "shell.execute_reply": "2024-04-11T10:16:30.917130Z",
     "shell.execute_reply.started": "2024-04-11T10:16:30.905965Z"
    }
   },
   "outputs": [],
   "source": [
    "SAVEDIR = Path(\"train\")\n",
    "SAVEDIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:16:30.923833Z",
     "iopub.status.busy": "2024-04-11T10:16:30.923433Z",
     "iopub.status.idle": "2024-04-11T10:16:30.939929Z",
     "shell.execute_reply": "2024-04-11T10:16:30.936468Z",
     "shell.execute_reply.started": "2024-04-11T10:16:30.923799Z"
    }
   },
   "outputs": [],
   "source": [
    "protein_names = [\"BRD4\", \"HSA\", \"sEH\"]\n",
    "smiles_col_names = [\n",
    "    \"buildingblock1_smiles\", \n",
    "    \"buildingblock2_smiles\", \n",
    "    \"buildingblock3_smiles\", \n",
    "    \"molecule_smiles\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:16:30.946608Z",
     "iopub.status.busy": "2024-04-11T10:16:30.945614Z",
     "iopub.status.idle": "2024-04-11T10:16:30.971186Z",
     "shell.execute_reply": "2024-04-11T10:16:30.968342Z",
     "shell.execute_reply.started": "2024-04-11T10:16:30.946485Z"
    }
   },
   "outputs": [],
   "source": [
    "schema_single_target = pa.schema([\n",
    "    (col_name, pa.string())\n",
    "    for col_name in smiles_col_names\n",
    "])\n",
    "\n",
    "schema_mixed_target = pa.schema([\n",
    "    (col_name, pa.string())\n",
    "    for col_name in smiles_col_names\n",
    "] + [(prot_name, pa.int8()) for prot_name in protein_names])\n",
    "\n",
    "schema_unprocessed = pa.schema([\n",
    "    (col_name, pa.string())\n",
    "    for col_name in smiles_col_names\n",
    "] + [\n",
    "    ('protein_name', pa.string()),\n",
    "    ('binds', pa.int8())\n",
    "])\n",
    "\n",
    "KEYS_INFO = [\n",
    "    (\"all0\", schema_single_target), \n",
    "    (\"all1\", schema_single_target), \n",
    "    (\"all_mixed\", schema_mixed_target),\n",
    "    ('unprocessed', schema_unprocessed)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:16:30.975886Z",
     "iopub.status.busy": "2024-04-11T10:16:30.975340Z",
     "iopub.status.idle": "2024-04-11T10:16:31.009896Z",
     "shell.execute_reply": "2024-04-11T10:16:31.006223Z",
     "shell.execute_reply.started": "2024-04-11T10:16:30.975854Z"
    }
   },
   "outputs": [],
   "source": [
    "class PqDataSaverWrapper:\n",
    "    def __init__(self, savedir=\".\", keys_info=KEYS_INFO, prefix=\"train\"):\n",
    "        self.file_handlers = None\n",
    "        if isinstance(savedir, str):\n",
    "            savedir = Path(savedir)\n",
    "        self.savedir = savedir\n",
    "        self.keys_info = keys_info\n",
    "        self.prefix = prefix\n",
    "\n",
    "    def add_data(self, df_dict):\n",
    "        if self.file_handlers is None:\n",
    "            raise Exception(\"Attempt to call 'add_data' outside of with-block\")\n",
    "        for key in df_dict:\n",
    "            if not key in self.file_handlers:\n",
    "                raise Exception(f\"Key should be equal to one of the provided in __init__ method, got '{key}' instead\")\n",
    "            if df_dict[key].shape[0] > 0:\n",
    "                pq_batch = pa.RecordBatch.from_pandas(df_dict[key])\n",
    "                self.file_handlers[key].write(pq_batch)\n",
    "\n",
    "    def __enter__(self):\n",
    "        self.file_handlers = {}\n",
    "        for key, schema in self.keys_info:\n",
    "            filename = self.savedir / f\"{self.prefix}_{key}.parquet\"\n",
    "            if filename.exists():\n",
    "                raise Exception(f\"File '{filename}' for key '{key}' already exists, remove it if you want to regenerate everything\")\n",
    "            self.file_handlers[key] = pq.ParquetWriter(filename.as_posix(), schema)\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        for pq_writer in self.file_handlers.values():\n",
    "            pq_writer.close()\n",
    "        self.file_handlers = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-11T10:16:31.017807Z",
     "iopub.status.busy": "2024-04-11T10:16:31.015548Z",
     "iopub.status.idle": "2024-04-11T10:35:28.868290Z",
     "shell.execute_reply": "2024-04-11T10:35:28.866402Z",
     "shell.execute_reply.started": "2024-04-11T10:16:31.017690Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ec33911221a42df9391bd9b3b908140",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/282 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_pq = pq.ParquetFile(\"../input/leash-BELKA/train.parquet\")\n",
    "\n",
    "partial_df = None\n",
    "additional_records = []\n",
    "with PqDataSaverWrapper(savedir=SAVEDIR, keys_info=KEYS_INFO) as pq_wrapper:\n",
    "    for i in tqdm(np.arange(train_pq.num_row_groups), total=train_pq.num_row_groups):\n",
    "        batch = train_pq.read_row_group(i)\n",
    "        batch_df = batch.to_pandas()\n",
    "        if partial_df is not None and partial_df.shape[0] > 0:\n",
    "            batch_df = pd.concat([partial_df, batch_df], ignore_index=True)\n",
    "            partial_df = None\n",
    "        batch_df = batch_df.pivot_table(\n",
    "            columns=\"protein_name\", \n",
    "            values=\"binds\", \n",
    "            index=smiles_col_names).reset_index(drop=False)\n",
    "        batch_df.columns.name = None\n",
    "        missing_protein_data_ids = batch_df[protein_names].isnull().any(axis=1)\n",
    "\n",
    "        filename = SAVEDIR / \"all_present\"/ f\"part_{i}.csv\"\n",
    "        new_data = {}\n",
    "        if (~missing_protein_data_ids).any():\n",
    "            all0_ids = batch_df[protein_names].max(1) == 0\n",
    "            all1_ids = batch_df[protein_names].min(1) == 1\n",
    "            new_data['all0'] = batch_df[(~missing_protein_data_ids) & all0_ids][smiles_col_names].reset_index(drop=True)\n",
    "            new_data['all1'] = batch_df[(~missing_protein_data_ids) & all1_ids][smiles_col_names].reset_index(drop=True)\n",
    "            all_mixed_df = batch_df[(~missing_protein_data_ids) & (~all1_ids) & (~all0_ids)].reset_index(drop=True)\n",
    "            all_mixed_df[protein_names] = all_mixed_df[protein_names].astype(np.int8)\n",
    "            new_data['all_mixed'] = all_mixed_df\n",
    "\n",
    "        if missing_protein_data_ids.any():\n",
    "            partial_df = batch_df[missing_protein_data_ids].melt(\n",
    "                id_vars=smiles_col_names,\n",
    "                value_vars=protein_names,\n",
    "                value_name=\"binds\",\n",
    "                var_name=\"protein_name\"\n",
    "            )\n",
    "            partial_df = partial_df[~partial_df.binds.isnull()].reset_index(drop=True)\n",
    "            partial_df['binds'] = partial_df['binds'].astype(np.int8)\n",
    "        pq_wrapper.add_data(new_data)\n",
    "\n",
    "    if partial_df is not None and partial_df.shape[0] > 0:\n",
    "        pq_wrapper.add_data({\n",
    "            'unprocessed': partial_df\n",
    "        })\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "none",
   "dataSources": [
    {
     "databundleVersionId": 8006601,
     "sourceId": 67356,
     "sourceType": "competition"
    }
   ],
   "dockerImageVersionId": 30684,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}