pipelines/mpnn_af2/02_prep_af2_inputs.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parse ProteinMPNN sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import glob\n",
    "import re\n",
    "import os\n",
    "\n",
    "mpnn_results = \"output/3HB/seqs/3HB.fa\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of sequences: 1000\n"
     ]
    }
   ],
   "source": [
    "(\n",
    "    T_values,\n",
    "    sample_values,\n",
    "    score_values,\n",
    "    global_score_values,\n",
    "    seq_recovery_values,\n",
    "    sequences,\n",
    ") = ([], [], [], [], [], [])\n",
    "\n",
    "with open(mpnn_results, \"r\") as file:\n",
    "    lines = file.readlines()\n",
    "\n",
    "    # Skip the first two rows\n",
    "    lines = lines[2:]\n",
    "\n",
    "    # Iterate over the rest\n",
    "    for i, line in enumerate(lines):\n",
    "        if line.startswith(\">\"):\n",
    "            scores = re.split(\",\\s*|,\", line[1:])\n",
    "\n",
    "            # Extract the desired fields from the header\n",
    "            T_value = float(scores[0].split(\"=\")[1])\n",
    "            sample_value = int(scores[1].split(\"=\")[1])\n",
    "            score_value = float(scores[2].split(\"=\")[1])\n",
    "            global_score_value = float(scores[3].split(\"=\")[1])\n",
    "            seq_recovery_value = float(scores[4].split(\"=\")[1])\n",
    "\n",
    "            # Save the values\n",
    "            T_values.append(T_value)\n",
    "            sample_values.append(sample_value)\n",
    "            score_values.append(score_value)\n",
    "            global_score_values.append(global_score_value)\n",
    "            seq_recovery_values.append(seq_recovery_value)\n",
    "\n",
    "            # Save the next line as the sequence\n",
    "            sequence = lines[i + 1].strip()\n",
    "            sequences.append(sequence)\n",
    "\n",
    "# Create a dataframe from the collected data\n",
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"t\": T_values,\n",
    "        \"sample\": sample_values,\n",
    "        \"score\": score_values,\n",
    "        \"global_Score\": global_score_values,\n",
    "        \"seq_Recovery\": seq_recovery_values,\n",
    "        \"sequence\": sequences,\n",
    "    }\n",
    ")\n",
    "\n",
    "df.to_csv(f\"{os.path.dirname(mpnn_results)}/seqs.csv\", index=False)\n",
    "print(f\"number of sequences: {df.shape[0]}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare inputs for AF2 predictions\n",
    "#### Group sequences for batch prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Slurm will run: 34 AF2 jobs, with each one predicting up to 30 sequences\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(f\"{os.path.dirname(mpnn_results)}/seqs.csv\")\n",
    "\n",
    "\n",
    "# Set number of sequences to be predicted per job\n",
    "grouped_sequences = 30\n",
    "num_seq = df.shape[0]\n",
    "# num_seq=1000\n",
    "\n",
    "array_jobs = num_seq // grouped_sequences\n",
    "if num_seq % grouped_sequences != 0:\n",
    "    array_jobs += 1\n",
    "\n",
    "print(\n",
    "    f\"Slurm will run: {array_jobs} AF2 jobs, with each one predicting up to {grouped_sequences} sequences\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 34 fasta files in folder: output/7urv_FMC63/tmp/af2_in\n"
     ]
    }
   ],
   "source": [
    "# Set AF2 input folder\n",
    "basename = os.path.basename(mpnn_results).split(\".\")[0]\n",
    "af2_input_folder = f\"output/{basename}/tmp/af2_in\"\n",
    "os.makedirs(af2_input_folder, exist_ok=True)\n",
    "\n",
    "# Set output folder and save sequences\n",
    "af2_df = pd.DataFrame(columns=[\"id\", \"seq\"])\n",
    "af2_output_folder = f\"output/{basename}/af2\"\n",
    "os.makedirs(af2_output_folder, exist_ok=True)\n",
    "\n",
    "# Iterate over sequences and generate AF2 inputs - fasta files\n",
    "for job_index in range(array_jobs):\n",
    "    start_index = job_index * grouped_sequences\n",
    "    end_index = start_index + grouped_sequences\n",
    "    sequences = df[\"sequence\"][start_index:end_index]\n",
    "\n",
    "    # Generate file path\n",
    "    file_path = f\"{af2_input_folder}/fasta_{job_index + 1}.fasta\"\n",
    "\n",
    "    # Write sequences to file with unique names\n",
    "    with open(file_path, \"w\") as file:\n",
    "        for i, seq in enumerate(sequences, start=start_index + 1):\n",
    "            file.write(f\">{job_index}_{i}_{basename}\\n{seq}\\n\")\n",
    "            # Append id and seq to the dataframe\n",
    "            af2_df = pd.concat(\n",
    "                [\n",
    "                    af2_df,\n",
    "                    pd.DataFrame({\"id\": [f\"{job_index}_{i}_{basename}\"], \"seq\": [seq]}),\n",
    "                ],\n",
    "                ignore_index=True,\n",
    "            )\n",
    "\n",
    "af2_df.to_csv(f\"{af2_output_folder}/af2_sequences.csv\", index=False)\n",
    "print(f\"Generated {array_jobs} fasta files in folder: {af2_input_folder}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run AF2 prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 139666\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "CompletedProcess(args='sbatch --array=1-34%18 helper_scripts/AF2_array.sh output/7urv_FMC63/tmp/af2_in ', returncode=0)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import subprocess\n",
    "\n",
    "# Setup arguments\n",
    "bash_arguments = f\"--array=1-{array_jobs}%18\"\n",
    "colabfold_arguments = f\"\"\n",
    "\n",
    "command = f\"sbatch {bash_arguments} helper_scripts/AF2_array.sh {af2_input_folder} {colabfold_arguments}\"\n",
    "\n",
    "# Run the array bash script\n",
    "subprocess.run(command, shell=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)\n",
      " 139666_[19-34%18]       gpu AF2_arra  tsatler PD       0:00      1 (JobArrayTaskLimit)\n",
      "          139666_1       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_2       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_3       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_4       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_5       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_6       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_7       gpu AF2_arra  tsatler  R       0:07      1 compute-0-12\n",
      "          139666_8       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "          139666_9       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_10       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_11       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_12       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_13       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_14       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_15       gpu AF2_arra  tsatler  R       0:07      1 compute-0-10\n",
      "         139666_16       gpu AF2_arra  tsatler  R       0:07      1 compute-0-11\n",
      "         139666_17       gpu AF2_arra  tsatler  R       0:07      1 compute-0-11\n",
      "         139666_18       gpu AF2_arra  tsatler  R       0:07      1 compute-0-11\n"
     ]
    }
   ],
   "source": [
    "!squeue --me"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyro",
   "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.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}