alphafold_interaction_scores_for alphafold-multimer data_v0.31.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# jupyter-notebook to handle data derived from original AlphaFold-Multimer. \n",
    "# pkl will be used to get iptm, ptm, plddt (thanks to Katerina Atallah-Yunes).\n",
    "# pae information will be loaded from pkl instead of json (updated on 2024/03/28). \n",
    "# the lines for loadning pae-related json were disabled (thanks to Yaqing Cheng).\n",
    "\n",
    "import os\n",
    "import json\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from Bio import PDB\n",
    "\n",
    "def reverse_and_scale_matrix(matrix: np.ndarray, pae_cutoff: float = 12.0) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Scale the values in the matrix such that:\n",
    "    0 becomes 1, pae_cutoff becomes 0, and values greater than pae_cutoff are also 0.\n",
    "    \n",
    "    Args:\n",
    "    - matrix (np.ndarray): Input numpy matrix.\n",
    "    - pae_cutoff (float): Threshold above which values become 0.\n",
    "    \n",
    "    Returns:\n",
    "    - np.ndarray: Transformed matrix.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Scale the values to [0, 1] for values between 0 and cutoff\n",
    "    scaled_matrix = (pae_cutoff - matrix) / pae_cutoff\n",
    "    scaled_matrix = np.clip(scaled_matrix, 0, 1)  # Ensures values are between 0 and 1\n",
    "    \n",
    "    return scaled_matrix\n",
    "\n",
    "\n",
    "def process_alphafold_output(base_directory: str, model_numbers: int = 5, recycling_numbers: int = 5, Protein_1 = \"A\", Protein_2 = \"B\", pae_cutoff: float = 12.0, lis_threshold: float = 0.203, lia_threshold: float = 3432.0) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Process AlphaFold output files and return a DataFrame containing relevant information.\n",
    "    \n",
    "    Args:\n",
    "    - base_directory (str): Base directory where the AlphaFold output files are stored.\n",
    "    - model_numbers (int): Number of models to process.\n",
    "    - recycling_numbers (int): Number of recycling indices to process.\n",
    "    - pae_cutoff (float): Threshold for PAE matrix above which values become 0.\n",
    "    - lis_threshold (float): Threshold for LIS (Local Interaction Score) to filter results.\n",
    "    - lia_threshold (float): Threshold for LIA (Local Interaction Area) to filter results.\n",
    "    \n",
    "    Returns:\n",
    "    - pd.DataFrame: DataFrame containing processed data.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Initialize a list to hold the Pandas Series objects\n",
    "    series_list = []\n",
    "\n",
    "    # Loop over each model number and recycling number\n",
    "    for model_num in range(1, model_numbers+1):\n",
    "        for recycling_num in range(0, recycling_numbers):\n",
    "            # print(f\"Model: {model_num}, Recycle: {recycling_num}\")\n",
    "            \n",
    "            # pae_name = f\"pae_model_{model_num}_multimer_v3_pred_{recycling_num}.json\"\n",
    "            # pae_file_path = os.path.join(base_directory, pae_name)\n",
    "            # with open(pae_file_path, 'r') as file:\n",
    "            #     json_data = json.load(file)\n",
    "\n",
    "\n",
    "            pdb_name = f\"unrelaxed_model_{model_num}_multimer_v3_pred_{recycling_num}.pdb\"\n",
    "            pdb_file_path = os.path.join(base_directory, pdb_name)\n",
    "\n",
    "\n",
    "            parser = PDB.PDBParser(QUIET=True)\n",
    "            structure = parser.get_structure(\"example\", pdb_file_path)\n",
    "            \n",
    "            chain_lengths = {}\n",
    "            \n",
    "            for model in structure:\n",
    "                for chain in model:\n",
    "                    chain_id = chain.get_id()\n",
    "                    chain_length = sum(1 for _ in chain.get_residues())  # Calculate chain length\n",
    "                    chain_lengths[chain_id] = chain_length\n",
    "                    # Accessing the length of chain 'A' from the dictionary\n",
    "                    protein_a_len = chain_lengths.get('B', 0)  # Default to 0 if 'A' chain is not found\n",
    "            # print(chain_lengths)\n",
    "            # print(protein_a_len)\n",
    "\n",
    "            pkl_name = f\"result_model_{model_num}_multimer_v3_pred_{recycling_num}.pkl\"\n",
    "            pkl_file_path = os.path.join(base_directory, pkl_name)\n",
    "            d = pickle.load(open(pkl_file_path,'rb'))\n",
    "            iptm = d.get('iptm')\n",
    "            ptm = d.get('ptm')\n",
    "            # added on 2024/03/28 #\n",
    "            pae = d.get('predicted_aligned_error')\n",
    "            # added on 2024/03/28 #\n",
    "            plddt = np.mean(d.get('plddt'))\n",
    "            confidence = d.get('ranking_confidence')\n",
    "\n",
    "            # print(pae_file_path)\n",
    "            # print(pdb_file_path)\n",
    "            # print(pkl_file_path)\n",
    "            # print(iptm)\n",
    "            # print(ptm)\n",
    "            # print(plddt)\n",
    "\n",
    "            # pae = np.array(json_data[0]['predicted_aligned_error'])\n",
    "            # pae_data[f\"model_{model_num}_recycle_{recycling_num}\"] = pae_matrix\n",
    "\n",
    "            pae_cutoff = 12\n",
    "\n",
    "            thresholded_pae = np.where(pae < pae_cutoff, 1, 0)\n",
    "\n",
    "            # Calculate the interaction amino acid numbers\n",
    "            local_interaction_protein_a = np.count_nonzero(thresholded_pae[:protein_a_len, :protein_a_len])\n",
    "            local_interaction_protein_b = np.count_nonzero(thresholded_pae[protein_a_len:, protein_a_len:])\n",
    "            local_interaction_interface_1 = np.count_nonzero(thresholded_pae[:protein_a_len, protein_a_len:])\n",
    "            local_interaction_interface_2 = np.count_nonzero(thresholded_pae[protein_a_len:, :protein_a_len])\n",
    "            local_interaction_interface_avg = (\n",
    "                local_interaction_interface_1 + local_interaction_interface_2\n",
    "            )\n",
    "\n",
    "            \n",
    "            # Calculate average thresholded_pae for each region\n",
    "            average_thresholded_protein_a = thresholded_pae[:protein_a_len,:protein_a_len].mean() * 100\n",
    "            average_thresholded_protein_b = thresholded_pae[protein_a_len:,protein_a_len:].mean() * 100\n",
    "            average_thresholded_interaction1 = thresholded_pae[:protein_a_len,protein_a_len:].mean() * 100\n",
    "            average_thresholded_interaction2 = thresholded_pae[protein_a_len:,:protein_a_len].mean() * 100\n",
    "            average_thresholded_interaction_total = (average_thresholded_interaction1 + average_thresholded_interaction2) / 2\n",
    "            \n",
    "\n",
    "            pae_protein_a = np.mean( pae[:protein_a_len,:protein_a_len] )\n",
    "            pae_protein_b = np.mean( pae[protein_a_len:,protein_a_len:] )\n",
    "            pae_interaction1 = np.mean(pae[:protein_a_len,protein_a_len:])\n",
    "            pae_interaction2 = np.mean(pae[protein_a_len:,:protein_a_len])\n",
    "            pae_interaction_total = ( pae_interaction1 + pae_interaction2 ) / 2\n",
    "\n",
    "            # For pae_A\n",
    "            selected_values_protein_a = pae[:protein_a_len, :protein_a_len][thresholded_pae[:protein_a_len, :protein_a_len] == 1]\n",
    "            average_selected_protein_a = np.mean(selected_values_protein_a)\n",
    "\n",
    "            # For pae_B\n",
    "            selected_values_protein_b = pae[protein_a_len:, protein_a_len:][thresholded_pae[protein_a_len:, protein_a_len:] == 1]\n",
    "            average_selected_protein_b = np.mean(selected_values_protein_b)\n",
    "\n",
    "            # For pae_interaction1\n",
    "            selected_values_interaction1 = pae[:protein_a_len, protein_a_len:][thresholded_pae[:protein_a_len, protein_a_len:] == 1]\n",
    "            average_selected_interaction1 = np.mean(selected_values_interaction1) if selected_values_interaction1.size > 0 else pae_cutoff\n",
    "\n",
    "            # For pae_interaction2\n",
    "            selected_values_interaction2 = pae[protein_a_len:, :protein_a_len][thresholded_pae[protein_a_len:, :protein_a_len] == 1]\n",
    "            average_selected_interaction2 = np.mean(selected_values_interaction2) if selected_values_interaction2.size > 0 else pae_cutoff\n",
    "\n",
    "            # For pae_interaction_total\n",
    "            average_selected_interaction_total = (average_selected_interaction1 + average_selected_interaction2) / 2\n",
    "\n",
    "        # At this point, plddt_data and pae_data dictionaries will have the extracted data\n",
    "            print_results = False\n",
    "            if print_results:\n",
    "                # Print the total results\n",
    "                print(\"Total pae_A : {:.2f}\".format(pae_protein_a))\n",
    "                print(\"Total pae_B : {:.2f}\".format(pae_protein_b))\n",
    "                print(\"Total pae_i_1 : {:.2f}\".format(pae_interaction1))\n",
    "                print(\"Total pae_i_2 : {:.2f}\".format(pae_interaction2))\n",
    "                print(\"Total pae_i_avg : {:.2f}\".format(pae_interaction_total))\n",
    "\n",
    "                # Print the local results\n",
    "                print(\"Local pae_A : {:.2f}\".format(average_selected_protein_a))\n",
    "                print(\"Local pae_B : {:.2f}\".format(average_selected_protein_b))\n",
    "                print(\"Local pae_i_1 : {:.2f}\".format(average_selected_interaction1))\n",
    "                print(\"Local pae_i_2 : {:.2f}\".format(average_selected_interaction2))\n",
    "                print(\"Local pae_i_avg : {:.2f}\".format(average_selected_interaction_total))\n",
    "\n",
    "                # Print the >PAE-cutoff area\n",
    "                print(\"Local interaction area (Protein A):\", local_interaction_protein_a)\n",
    "                print(\"Local interaction area (Protein B):\", local_interaction_protein_b)\n",
    "                print(\"Local interaction area (Interaction 1):\", local_interaction_interface_1)\n",
    "                print(\"Local interaction area (Interaction 2):\", local_interaction_interface_2)\n",
    "                print(\"Total Interaction area (Interface):\", local_interaction_interface_avg)\n",
    "\n",
    "\n",
    "            # Transform the pae matrix\n",
    "            scaled_pae = reverse_and_scale_matrix(pae, pae_cutoff)\n",
    "\n",
    "            # For local interaction score for protein_a\n",
    "            selected_values_protein_a = scaled_pae[:protein_a_len, :protein_a_len][thresholded_pae[:protein_a_len, :protein_a_len] == 1]\n",
    "            average_selected_protein_a_score = np.mean(selected_values_protein_a)\n",
    "\n",
    "            # For local interaction score for protein_b\n",
    "            selected_values_protein_b = scaled_pae[protein_a_len:, protein_a_len:][thresholded_pae[protein_a_len:, protein_a_len:] == 1]\n",
    "            average_selected_protein_b_score = np.mean(selected_values_protein_b)\n",
    "\n",
    "            # For local interaction score1\n",
    "            selected_values_interaction1_score = scaled_pae[:protein_a_len, protein_a_len:][thresholded_pae[:protein_a_len, protein_a_len:] == 1]\n",
    "            average_selected_interaction1_score = np.mean(selected_values_interaction1_score) if selected_values_interaction1_score.size > 0 else 0\n",
    "\n",
    "            # For local interaction score2\n",
    "            selected_values_interaction2_score = scaled_pae[protein_a_len:, :protein_a_len][thresholded_pae[protein_a_len:, :protein_a_len] == 1]\n",
    "            average_selected_interaction2_score = np.mean(selected_values_interaction2_score) if selected_values_interaction2_score.size > 0 else 0\n",
    "\n",
    "            # For average local interaction score\n",
    "            average_selected_interaction_total_score = (average_selected_interaction1_score + average_selected_interaction2_score) / 2\n",
    "\n",
    "            # Append the data to the series list\n",
    "            series_list.append(pd.Series({\n",
    "                'Protein_1': Protein_1,\n",
    "                'Protein_2': Protein_2,\n",
    "                'LIS': round(average_selected_interaction_total_score, 3), # Local Interaction Score (LIS)\n",
    "                'LIA': local_interaction_interface_avg, # Local Interaction Area (LIA)\n",
    "                'ipTM': round(float(iptm), 3),\n",
    "                'Confidence': round(float(iptm*0.8 + ptm*0.2),3),\n",
    "                'pTM': round(float(ptm), 3),\n",
    "                'pLDDT': round(plddt, 2),\n",
    "                'Model': model_num,\n",
    "                'Recycle': recycling_num,\n",
    "                'saved folder': os.path.dirname(pdb_file_path),\n",
    "                'pdb': os.path.basename(pdb_file_path),\n",
    "                'pkl': os.path.basename(pkl_file_path),\n",
    "            }))\n",
    "\n",
    "    # Concatenate all Pandas Series objects into a single DataFrame\n",
    "    result_df = pd.concat(series_list, axis=1).T\n",
    "    \n",
    "    # Filter rows based on the specified LIS and LIA thresholds\n",
    "    result_df_filtered = result_df[(result_df['LIS'] >= lis_threshold) & (result_df['LIA'] >= lia_threshold)]\n",
    "    \n",
    "    return result_df, result_df_filtered\n",
    "\n",
    "# Create a DataFrame for the additional data\n",
    "metrics_data = {\n",
    "    'Metric': ['Average LIS', 'Best LIS', 'Average LIA', 'Best LIA', 'Average ipTM', 'Best ipTM', 'Average Confidence', 'Best Confidence', 'Average pDockQ', 'Best pDockQ', 'Average pDockQ2', 'Best pDockQ2'],\n",
    "    'Optimal Threshold': [0.0734, 0.203, 1610.4, 3432, 0.322, 0.38, 0.3672, 0.432, 0.133427109, 0.148516258, 0.015093248, 0.02106924],\n",
    "    'Specificity': [0.926011561, 0.919075145, 0.876300578, 0.855491329, 0.937572254, 0.823121387, 0.951445087, 0.85433526, 0.804624277, 0.865895954, 0.917919075, 0.895953757],\n",
    "    'Sensitivity': [0.786516854, 0.730337079, 0.767790262, 0.775280899, 0.711610487, 0.734082397, 0.674157303, 0.685393258, 0.68164794, 0.666666667, 0.692883895, 0.670411985],\n",
    "    'AUC': [0.910601632, 0.890710312, 0.888699097, 0.86613626, 0.891301336, 0.862501353, 0.859056959, 0.84053387, 0.79601221, 0.818267628, 0.841622827, 0.832053863],\n",
    "    \"Youden's Index\": [0.712528415, 0.649412223, 0.64409084, 0.630772228, 0.649182741, 0.557203784, 0.62560239, 0.539728519, 0.486272218, 0.53256262, 0.61080297, 0.566365742]\n",
    "}\n",
    "metrics_data_df = pd.DataFrame(metrics_data)\n",
    "metrics_data_df = metrics_data_df.round(3)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage:\n",
    "base_directory = 'YOUR ALPHAFOLD OUTPUT FOLDER'\n",
    "model_numbers = 5\n",
    "recycling_numbers = 5\n",
    "Protein_1 = \"A\"\n",
    "Protein_2 = \"B\"\n",
    "pae_cutoff = 12 \n",
    "lis_threshold = 0.203 # might be too strict for weak or transient interaction\n",
    "lia_threshold = 3432.0 \n",
    "total_prediction, positive_prediction = process_alphafold_output(base_directory, model_numbers, recycling_numbers, Protein_1, Protein_2, pae_cutoff, lis_threshold, lia_threshold,)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write DataFrames to an Excel file with three sheets\n",
    "with pd.ExcelWriter('alphafold_predictions_test.xlsx') as writer:\n",
    "    positive_prediction.to_excel(writer, sheet_name='Positive_PPI', index=False)\n",
    "    total_prediction.to_excel(writer, sheet_name='Total_Prediction', index=False)\n",
    "    \n",
    "    # Write metrics data with wrapped header\n",
    "    metrics_data_df.to_excel(writer, sheet_name='Optimal Thresholds', index=False)\n",
    "    \n",
    "    # Get the workbook and the worksheet\n",
    "    workbook = writer.book\n",
    "    worksheet = writer.sheets['Optimal Thresholds']\n",
    "    \n",
    "    # Get the header range\n",
    "    header_range = f'A1:{chr(64 + len(metrics_data_df.columns))}1'\n",
    "    \n",
    "    # Wrap the header text\n",
    "    header_format = workbook.add_format({'text_wrap': True, 'valign': 'vcenter', 'align': 'center', 'bold': True})\n",
    "    worksheet.set_row(0, None, header_format)\n",
    "    \n",
    "    # Auto-adjust column width\n",
    "    for i, col in enumerate(metrics_data_df.columns):\n",
    "        column_len = metrics_data_df[col].astype(str).str.len().max()\n",
    "        column_width = max(column_len, len(col))\n",
    "        worksheet.set_column(i, i, column_width + 2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}