foldcomp-py-examples.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/steineggerlab/foldcomp/blob/master/foldcomp-py-examples.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6IdoP-YCJ3Ps"
      },
      "source": [
        "# Foldcomp Python Example Notebook\n",
        "\n",
        "Foldcomp compresses protein structures with torsion angles effectively. It compresses the backbone atoms to 8 bytes and the side chain to additionally 4-5 byes per residue, thus an averaged-sized protein of 350 residues requires ~4.2kb.\n",
        "\n",
        "In this notebook we will show you the basic usage of the Foldcomp Python API."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "37KdADYrKHyI"
      },
      "outputs": [],
      "source": [
        "# Installing foldcomp\n",
        "%pip install -q foldcomp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ew55RBITKjhj"
      },
      "outputs": [],
      "source": [
        "# Download example PDB file\n",
        "!wget -q https://raw.githubusercontent.com/steineggerlab/foldcomp/master/test/test.pdb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mJ3y04ghRvxq"
      },
      "source": [
        "## Roundtrip a PDB file through Foldcomp\n",
        "\n",
        "In this example we first compress and then decompress a PDB file, afterwards we load the atoms of both files into BioPython to compute a RMSD."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9CcutAfUKymB"
      },
      "outputs": [],
      "source": [
        "# load foldcomp module\n",
        "import foldcomp\n",
        "\n",
        "# read example PDB file\n",
        "with open(\"test.pdb\", \"r\") as f:\n",
        "  original = f.read()\n",
        "\n",
        "# Compress input with reset points every 25 residues\n",
        "# Should give a RMSD ~0.07A. A reset point every 200 residues will give a RMSD ~0.2A\n",
        "fcz = foldcomp.compress(\"test.pdb\", original, anchor_residue_threshold=25)\n",
        "\n",
        "# Decompress again\n",
        "(name, pdb) = foldcomp.decompress(fcz)\n",
        "\n",
        "# Save as a PDB file again\n",
        "with open(name + \"_new.pdb\", \"w\") as f:\n",
        "    f.write(pdb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aTTmbJXTMf3a"
      },
      "outputs": [],
      "source": [
        "# install biopython\n",
        "%pip install -q biopython"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5V6ZpDw4M9y5"
      },
      "outputs": [],
      "source": [
        "# compute the RMSD between the two versions of the structure\n",
        "import Bio.PDB\n",
        "\n",
        "parser = Bio.PDB.PDBParser(QUIET = True)\n",
        "reference = parser.get_structure(\"original\", \"test.pdb\")\n",
        "roundtrip = parser.get_structure(\"foldcomp\", \"test.pdb_new.pdb\")\n",
        "\n",
        "ref_atoms = [residue['CA'] for chain in reference[0] for residue in chain]\n",
        "fcz_atoms = [residue['CA'] for chain in roundtrip[0] for residue in chain]\n",
        "\n",
        "superposition = Bio.PDB.Superimposer()\n",
        "superposition.set_atoms(ref_atoms, fcz_atoms)\n",
        "superposition.apply([])\n",
        "\n",
        "print(superposition.rms)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wRfH6sFCOofi"
      },
      "source": [
        "# Dealing with Foldcomp Databases\n",
        "\n",
        "Dealing with Databases of enormous size such as the AlphaFold Database can be very painful due to the large number of files contained in it. We offer the AlphaFold Database in our own database format for easier download and handling.\n",
        "\n",
        "Due to the limited amout of disk space in Google Colab, we will only download the Swiss-Prot subset of the database."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rhLfy7sPJv7"
      },
      "outputs": [],
      "source": [
        "# use the built-in downloader\n",
        "# You can just call \n",
        "#   foldcomp.setup(\"afdb_swissprot_v4\")\n",
        "# in a local python environment. Google Colab requires a async call to work correctly\n",
        "await foldcomp.setup_async(\"afdb_swissprot_v4\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dPXwuTNrPqDP"
      },
      "outputs": [],
      "source": [
        "with foldcomp.open(\"afdb_swissprot_v4\") as db:\n",
        "  # Iterate through database\n",
        "  i = 0\n",
        "  for (name, pdb) in db:\n",
        "    # save entries as seperate pdb files\n",
        "    with open(name + \".pdb\", \"w\") as f:\n",
        "      f.write(pdb)\n",
        "    i += 1\n",
        "    if i % 10 == 0:\n",
        "      break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JSAUA_DQc1o5"
      },
      "source": [
        "The database setup command downloaded a database file containing all compressed structures in a single file and a few accompanying files. The \"afdb_swissprot_v4.lookup\" contains all the AlphaFold database accession mapping to the Foldcomp database. By giving a list of accessions to the `uniprot_ids` parameter in `foldcomp.open`, you can iterate over a user defined subset of accessions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OVKkQuFlQ2fF"
      },
      "outputs": [],
      "source": [
        "!head afdb_swissprot_v4.lookup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hPvc2qvPQuED"
      },
      "outputs": [],
      "source": [
        "ids = [\"AF-Q53M11-F1-model_v4\", \"AF-Q8IYB0-F1-model_v4\"]\n",
        "with foldcomp.open(\"afdb_swissprot_v4\", ids = ids) as db:\n",
        "  for (name, pdb) in db:\n",
        "    with open(name + \".pdb\", \"w\") as f:\n",
        "      f.write(pdb)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Using `foldcomp.get(data)`, you can get a dictionary containing the precomputed\n",
        "features of the structure. Here is an example of drawing Ramachandran plots for\n",
        "a subset of the structures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install matplotlib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.tri as tri\n",
        "import numpy as np\n",
        "\n",
        "plt.style.use(\n",
        "    'https://github.com/dhaitz/matplotlib-stylesheets/raw/master/pitayasmoothie-light.mplstyle'\n",
        ")\n",
        "\n",
        "# load the database\n",
        "db_swissprot = foldcomp.open(\"afdb_swissprot_v4\")\n",
        "N_PROTEINS = 4\n",
        "db = [db_swissprot[i] for i in range(N_PROTEINS)]\n",
        "list_of_data_dicts = []\n",
        "for (name, pdb) in db:\n",
        "    list_of_data_dicts.append((name, foldcomp.get_data(pdb)))\n",
        "\n",
        "def set_axis_for_ramachandran(ax):\n",
        "    ax.set_xlim(-180, 180)\n",
        "    ax.set_ylim(-180, 180)\n",
        "    ax.set_aspect(\"equal\")\n",
        "    ax.set_xticks([-180, -90, 0, 90, 180])\n",
        "    ax.set_yticks([-180, -90, 0, 90, 180])\n",
        "    ax.set_xlabel(r\"$\\phi$\")\n",
        "    ax.set_ylabel(r\"$\\psi$\")\n",
        "\n",
        "# plot the ramachandran plots\n",
        "fig, ax = plt.subplots(figsize=(5, 5))\n",
        "set_axis_for_ramachandran(ax)\n",
        "for i, (name, data) in enumerate(list_of_data_dicts[:N_PROTEINS]):\n",
        "    ax.scatter(data[\"phi\"], data[\"psi\"], s=1, label=name)\n",
        "ax.legend()\n",
        "plt.tight_layout()\n",
        "plt.savefig(\"ramachandran.png\", dpi=300)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}