notebooks/31_run_rosetta_ddg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Calculate features using [ProtBert](https://github.com/agemagician/ProtTrans).\n",
    "\n",
    "```bash\n",
    "export NOTEBOOK_PATH=\"$(realpath 31_run_rosetta_ddg.ipynb)\"\n",
    "\n",
    "export DATASET_NAME=\"cagi6-sherloc\"\n",
    "export DATASET_PATH=\"30_cagi6_sherloc/input-data-gby-protein.parquet\"\n",
    "export ORIGINAL_ARRAY_TASK_COUNT=4182\n",
    "\n",
    "sbatch --export DATASET_NAME,DATASET_PATH,NOTEBOOK_PATH,ORIGINAL_ARRAY_TASK_COUNT --array=1-1 --time 3:00:00 --mem=0 --exclusive ../scripts/run_notebook_cpu.sh\n",
    "\n",
    "\n",
    "# Niagara\n",
    "sbatch --export DATASET_NAME,DATASET_PATH,NOTEBOOK_PATH,ORIGINAL_ARRAY_TASK_COUNT --array=1-4182 --time 3:00:00 --mem=0 --exclusive ../scripts/run_notebook_cpu.sh\n",
    "```\n",
    "\n",
    "\n",
    "```bash\n",
    "export NOTEBOOK_PATH=\"$(realpath 31_run_rosetta_ddg.ipynb)\"\n",
    "export DATASET_NAME=\"humsavar\"\n",
    "export DATASET_PATH=\"30_humsavar/humsavar-gby-protein.parquet\"\n",
    "export ORIGINAL_ARRAY_TASK_COUNT=612\n",
    "\n",
    "sbatch --export DATASET_NAME,DATASET_PATH,NOTEBOOK_PATH,ORIGINAL_ARRAY_TASK_COUNT --array=1-612 --time 24:00:00 --ntasks-per-node=10 --mem=20G ../scripts/run_notebook_cpu.sh\n",
    "```\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import itertools\n",
    "import os\n",
    "import re\n",
    "import socket\n",
    "import subprocess\n",
    "import sys\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "\n",
    "import kmbio\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "from elaspic2.plugins.rosetta_ddg import RosettaDDG\n",
    "from kmbio import PDB\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NOTEBOOK_DIR = Path(\"31_run_rosetta_ddg\").resolve()\n",
    "NOTEBOOK_DIR.mkdir(exist_ok=True)\n",
    "\n",
    "NOTEBOOK_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (slurm_tmpdir := os.getenv(\"SLURM_TMPDIR\")) is not None:\n",
    "    os.environ[\"TMPDIR\"] = slurm_tmpdir\n",
    "    \n",
    "print(tempfile.gettempdir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"scinet\" in socket.gethostname():\n",
    "    CPU_COUNT = 40\n",
    "else:\n",
    "    CPU_COUNT = max(1, len(os.sched_getaffinity(0)))\n",
    "\n",
    "# CPU_COUNT = max(1, CPU_COUNT // 2)\n",
    "\n",
    "CPU_COUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_NAME = os.getenv(\"DATASET_NAME\")\n",
    "DATASET_PATH = os.getenv(\"DATASET_PATH\")\n",
    "TASK_ID = os.getenv(\"SLURM_ARRAY_TASK_ID\")\n",
    "TASK_ID_OFFSET = int(os.getenv(\"ARRAY_TASK_ID_OFFSET\", \"0\"))\n",
    "TASK_COUNT = os.getenv(\"ORIGINAL_ARRAY_TASK_COUNT\") or os.getenv(\"SLURM_ARRAY_TASK_COUNT\")\n",
    "\n",
    "TASK_ID = (int(TASK_ID) + TASK_ID_OFFSET) if TASK_ID is not None else None\n",
    "TASK_COUNT = int(TASK_COUNT) if TASK_COUNT is not None else None\n",
    "\n",
    "DATASET_NAME, DATASET_PATH, TASK_ID, TASK_ID_OFFSET, TASK_COUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEBUG = TASK_ID is None\n",
    "\n",
    "if DEBUG:\n",
    "    DATASET_NAME = \"humsavar\"\n",
    "    DATASET_PATH = str(\n",
    "        NOTEBOOK_DIR.parent.joinpath(\"30_humsavar\", \"humsavar-gby-protein.parquet\")\n",
    "        # NOTEBOOK_DIR.parent.joinpath(\"30_cagi6_sherloc\", \"input-data-gby-protein.parquet\")\n",
    "    )\n",
    "    TASK_ID = 612\n",
    "    TASK_COUNT = 612 # 4182\n",
    "else:\n",
    "    assert DATASET_NAME is not None\n",
    "    assert DATASET_PATH is not None\n",
    "    assert TASK_ID is not None\n",
    "    assert TASK_COUNT is not None\n",
    "    DATASET_PATH = Path(DATASET_PATH).expanduser().resolve()\n",
    "\n",
    "DATASET_NAME, DATASET_PATH, TASK_ID, TASK_COUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "output_file = NOTEBOOK_DIR.joinpath(DATASET_NAME, f\"shard-{TASK_ID}-of-{TASK_COUNT}.parquet\")\n",
    "output_file.parent.mkdir(exist_ok=True)\n",
    "\n",
    "output_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if output_file.is_file():\n",
    "    raise Exception(\"Already finished!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workspace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pfile = pq.ParquetFile(DATASET_PATH)\n",
    "\n",
    "pfile.num_row_groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert TASK_COUNT == pfile.num_row_groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_df = pfile.read_row_group(TASK_ID - 1).to_pandas(integer_object_nulls=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(input_df.head(2))\n",
    "print(len(input_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "protein_id_column = None\n",
    "\n",
    "for col in [\"protein_id\", \"uniprot_id\"]:\n",
    "    if col in input_df:\n",
    "        protein_id_column = col\n",
    "        \n",
    "assert protein_id_column is not None\n",
    "protein_id_column"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tup = next(input_df.itertuples(index=False))\n",
    "\n",
    "iterable_fields = []\n",
    "for field in tup._fields:\n",
    "    if field in [protein_id_column]:\n",
    "        continue\n",
    "    try:\n",
    "        if len(getattr(tup, field)) == len(tup.mutation):\n",
    "            iterable_fields.append(field)\n",
    "    except TypeError:\n",
    "        pass\n",
    "\n",
    "iterable_fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    for field in iterable_fields:\n",
    "        input_df[field] = input_df[field].str[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Score mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_mutation(mutation):\n",
    "    aa = \"GVALICMFWPDESTYQNKRH\"\n",
    "    if re.search(f\"^[{aa}][1-9]+[0-9]*[{aa}]$\", mutation) is None:\n",
    "        print(f\"Skipping mutation {mutation} because it appears to be malformed.\")\n",
    "        return False\n",
    "\n",
    "    if mutation[0] == mutation[-1]:\n",
    "        print(\n",
    "            f\"Skipping mutation {mutation} because the wildtype and mutant residues are the same.\"\n",
    "        )\n",
    "        return False\n",
    "\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def worker(idx, mutation_idx, mutation, data):\n",
    "    if not validate_mutation(mutation):\n",
    "        return None\n",
    "\n",
    "    try:\n",
    "        results = RosettaDDG.analyze_mutation(f\"A_{mutation}\", data)\n",
    "    except Exception as error:\n",
    "        print(f\"{error!r}\")\n",
    "        return None\n",
    "    else:\n",
    "        results = {f\"rosetta_{key}\": value for key, value in results.items()}\n",
    "        return idx, mutation_idx, mutation, results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "futures = []\n",
    "results = []\n",
    "with concurrent.futures.ProcessPoolExecutor(CPU_COUNT) as pool:\n",
    "    num_mutations = 0\n",
    "    for tup in input_df.itertuples():\n",
    "        assert all(\n",
    "            [\n",
    "                (len(getattr(tup, field)) == len(tup.mutation))\n",
    "                for field in iterable_fields\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        with tempfile.NamedTemporaryFile(suffix=\".pdb\") as tmp_file:\n",
    "            with open(tmp_file.name, \"wt\") as fout:\n",
    "                fout.write(tup.structure)\n",
    "            data = RosettaDDG.build(\n",
    "                tmp_file.name,\n",
    "                protocol=\"cartesian_ddg\",\n",
    "                energy_function=\"beta_nov16_cart\",\n",
    "                interface=0,\n",
    "            )\n",
    "\n",
    "        for mutation_idx, mutation in enumerate(tup.mutation):\n",
    "            future = pool.submit(worker, tup.Index, mutation_idx, mutation, data)\n",
    "            futures.append(future)\n",
    "            num_mutations += 1\n",
    "\n",
    "    for future in tqdm(concurrent.futures.as_completed(futures), total=num_mutations):\n",
    "        result = future.result()\n",
    "        if result is None:\n",
    "            continue\n",
    "\n",
    "        idx, mutation_idx, mutation, rosetta_result = result\n",
    "        row = input_df.loc[idx]\n",
    "\n",
    "        results.append(\n",
    "            {\n",
    "                \"protein_id\": row[protein_id_column],\n",
    "                \"mutation\": mutation,\n",
    "            }\n",
    "            | {field: row[field][mutation_idx] for field in iterable_fields}\n",
    "            | rosetta_result\n",
    "        )\n",
    "\n",
    "results_df = pd.DataFrame(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(results_df.head(2))\n",
    "print(len(results_df))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not DEBUG:\n",
    "    pq.write_table(pa.Table.from_pandas(results_df, preserve_index=False), output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}