notebooks/protein.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a5279e36-1b2c-4042-ab2a-65e3dc09db99",
   "metadata": {},
   "source": [
    "# Protein\n",
    "\n",
    "This notebooks covers the usage of the `protein.ProteinBase` API and describes how structures with different numbers of atoms are handled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57e6f8ba-553e-4c94-a4ac-eeb0a6e2e472",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3eb2c76-4b93-4b71-86d7-2ec86bbc4565",
   "metadata": {},
   "outputs": [],
   "source": [
    "from prtm import protein"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94d3e13-f096-43fa-973d-15466741df98",
   "metadata": {},
   "source": [
    "Let's grab a somewhat complex structure from the `PDB` and load it into a protein structure. The main structure represention is `Protein37` which means a structure with 37 possible non-hydrogen atoms per residue. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c4a09e8-5425-4014-a52d-28d5a23e27e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdb_str = protein.get_structure_from_pdb(\"6QNO\")\n",
    "structure = protein.Protein37.from_pdb_string(pdb_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "085c110d-e4c9-4ae8-a8cf-7021f3e74c49",
   "metadata": {},
   "source": [
    "## Some basics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc9cd21e-9f7b-4652-bbd6-0c5f2b267d19",
   "metadata": {},
   "source": [
    "Let's show the structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12575bec-344b-4754-b296-36a45207db30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unlike predicted structures, the b-factors in PDB files are not\n",
    "# confidence scores in the range 0-100\n",
    "structure.show(\n",
    "    bfactor_is_confidence=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a085ea0-2ac6-44b8-871d-ceee73aa1eef",
   "metadata": {},
   "source": [
    "The `Protein37` class is effectively a PDB file that's been parsed into numpy arrays with some extra sugar that makes it interactive in a notebook. The fields that define the structure can be seen from the `fields` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f2c5356-b9c9-4873-afdc-d86a8c06fd38",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure.fields"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "042c56f3-8fbe-41f9-bd99-89117970c38b",
   "metadata": {},
   "source": [
    "These fields were adopted from `AlphaFold` / `OpenFold` with some extras like the `hetatom` fields to support models like `RFDiffusion`. Briefly the fields represent:\n",
    "\n",
    "- `atom_positions`: The XYZ coordinates of each atom\n",
    "- `aatype`: The amino-acid type indices\n",
    "- `atom_mask`: Mask with zeros where an atom is missing / doesn't exist\n",
    "- `residue_index`: The index of each residue in `aatype`; usually but not always sequential\n",
    "- `chain_index`: The chain indices corresponding to each residue\n",
    "- `b_factors`: Residue position confidence\n",
    "- `hetatom_positions`: The XYZ coordinates of `hetatoms` (i.e., a bound ligand)\n",
    "- `hetatom_names`: The names of the `hetatoms`.\n",
    "- `parents`, `parents_chain_index`, `remark` some header fields used by PDB\n",
    "\n",
    "Most of these fields are arrays that aren't particularly human readable. We can get useful information about the structure from a few handy properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cece982f-1b71-4905-96b7-a8c9964114fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the shape of the structure\n",
    "print(\"Shape:\", structure.shape)\n",
    "# Show the chain names\n",
    "print(\"Chain names:\", structure.chains)\n",
    "# Show the sequence for a particular chain\n",
    "print(\"Chain A sequence\", structure.sequence(\"A\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a41661b4-bf43-416f-882a-b7f80b8bb7cf",
   "metadata": {},
   "source": [
    "For structures with multiple chains like this one we can do handy things like easily extract a certain chain as a new structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6742e13-255f-4d5d-98ea-41f996c1fa85",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure_b = structure.get_chain(\"B\")\n",
    "structure_b.show(\n",
    "    cmap=\"cool\",\n",
    "    bfactor_is_confidence=False,\n",
    "    show_sidechains=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97ecdf4a-72da-40a7-8583-d064695297b9",
   "metadata": {},
   "source": [
    "## Structure atom numbers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b01655ec-89bb-477c-a36f-e39edfcb847c",
   "metadata": {},
   "source": [
    "The `Protein37` structure is the main representation that we use for `PDB` files because it has the nice property that for every residue a particular index in the array always represent the same kind of atom. For example, if we wanted to get the positions of all `CG2` atoms we can index into the `atom_positions` with `structure.atom_positions[:, 7]`. The main disadvantage of the atom 37 representation is no residue has all possible atoms and there is a lot of empty space in the array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a333ad6d-0007-4490-9db8-90bf6aaf616c",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_atom_mask = structure.atom_mask.ravel()\n",
    "filled = all_atom_mask.sum()\n",
    "total = len(all_atom_mask)\n",
    "print(f\"Fraction of filled positions: {filled / total:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96af26fa-6bd2-4a8d-9da2-cca145d91293",
   "metadata": {},
   "source": [
    "In this case only 21% of the atoms represented in the array are actually needed! That's not usually an issue when working with individual structures but can be limiting when trying to train models on GPUs with finite memory. In practice, most models use a denser 14 atom representation of the structure. The only disadvantage here is that the identity of an atom depends on the residue in question. The mapping of atoms is specified by `prtm.constants.residue_constants.restype_atom14_to_atom37`. Let's see how much space using 14 atoms saves us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1305bf77-9b05-4de0-a960-6cab654e801a",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure14 = structure.to_protein14()\n",
    "all_atom_mask14 = structure14.atom_mask.ravel()\n",
    "filled = all_atom_mask14.sum()\n",
    "total = len(all_atom_mask14)\n",
    "print(f\"Fraction of filled positions: {filled / total:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5263a70-d4d6-4e15-a229-88947477cc1f",
   "metadata": {},
   "source": [
    "Nice! This structure has much less wasted space and is more than 2x smaller. As a quick sidenote there's an extension of the 14 atom structure that has 27 atoms to accomodate hydrogens (`Protein27`). This is used by some models but generally isn't very useful because our `pdb_parser` ignores hydrogens. The remaining representations to consider further simplify structures down to just the backbone with 5, 4, 3 or 1 atom.\n",
    "\n",
    "- `Protein5`: The backbone atoms including `O` and `CB`\n",
    "- `Protein4`: The backbone atoms including `O`\n",
    "- `Protein3`: The backbone atoms (`N`, `CA`, `C`)\n",
    "- `ProteinCATrace`: The `CA` atoms only\n",
    "\n",
    "We can easily convert to any of these when needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69238255-a727-4a17-a22d-3910cb659279",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure5 = structure.to_protein5()\n",
    "structure4 = structure.to_protein4()\n",
    "structure3 = structure.to_protein3()\n",
    "structure_ca = structure.to_ca_trace()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cd12750-a165-4934-9f77-a0eb59db34dd",
   "metadata": {},
   "source": [
    "Importantly, these backbone only representations are destructive. When created from `Protein37` or `Protein14` the sidechain atoms are lost permanently. We can still convert back to these types but any sidechains atoms will be masked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec785f0a-3110-4ccc-b636-9d37d192f450",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show that now even more of the atoms are masked out\n",
    "structure3_to_37 = structure3.to_protein37()\n",
    "all_atom_mask3_to_37 = structure3_to_37.atom_mask.ravel()\n",
    "filled = all_atom_mask3_to_37.sum()\n",
    "total = len(all_atom_mask3_to_37)\n",
    "print(f\"Fraction of filled positions: {filled / total:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a441211a-3cef-451e-820f-7a89d36a3872",
   "metadata": {},
   "source": [
    "So why is it important to have all of these representations? It allows us to easily chain together models with different input and output shapes. For example, `RFDiffusion` creates structures with 14 atoms but an inverse folding model like `ProteinMPNN` only accepts a structure with backbone atoms. We can easily chain inputs and outputs together by calling `structure.to_protein4()` in this case. Internally, every pipeline has a particular representation that it expects and we perform the conversion automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49dcc708-a804-4db9-8808-165f92b595b9",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}