2-metadata-exploration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2022 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Metadata Exploration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install and import required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install packages\n",
    "! pip install -U google-cloud-firestore google-cloud-storage ipywidgets\n",
    "! pip install \"colabfold[alphafold] @ git+https://github.com/sokrypton/ColabFold\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reload the kernel before proceeding\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import packages\n",
    "import py3Dmol\n",
    "\n",
    "from google.cloud import firestore\n",
    "from google.cloud import storage\n",
    "\n",
    "from IPython import display\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "from colabfold.colabfold import plot_plddt_legend\n",
    "from colabfold.colabfold import pymol_color_list, alphabet_list\n",
    "\n",
    "# Instantiate firestore client\n",
    "fs_client = firestore.Client()\n",
    "gcs_client = storage.Client()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Succeeded pipelines\n",
    "\n",
    "List 10 experiments that completed successfully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can change to 'RUNNING', 'FAILED' or 'SUCCEEDED'\n",
    "status = 'SUCCEEDED'\n",
    "\n",
    "# Run Query\n",
    "docs = fs_client.collection('colabfold-experiments').where(\n",
    "    'status', '==', status).limit(10).get()\n",
    "\n",
    "if len(docs) == 0:\n",
    "    raise ValueError('No experiment with state SUCCEEDED')\n",
    "\n",
    "# Print results\n",
    "print(f'Found {len(docs)} experiment(s).')\n",
    "for doc in docs:\n",
    "    doc_dict = doc.to_dict()\n",
    "    print(f'=> Experiment: {doc_dict[\"job_id\"]}')\n",
    "    print(f'   - Firestore Reference: {doc_dict[\"firestore_ref\"]}')\n",
    "    job_id = doc_dict['firestore_ref']"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell will:\n",
    " - Query one specific job and bring its information from Firestore.\n",
    " - Get the list of all the files generated by this job.\n",
    " - Get the list of PDB files generated by this job\n",
    " - List the models by its ranking condifence\n",
    " - Display a sample image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (optional) Set the job_id with the Firestore reference you want.\n",
    "# job_id = '<firestore reference>'\n",
    "\n",
    "# Run Query\n",
    "doc = fs_client.collection('colabfold-experiments').document(job_id).get()\n",
    "\n",
    "if not doc.exists:\n",
    "    raise ValueError('Document does not exist.')\n",
    "\n",
    "# Print results\n",
    "doc_dict = doc.to_dict()\n",
    "\n",
    "blobs_listing = gcs_client.list_blobs(\n",
    "    doc_dict['job_gcs_path'],\n",
    "    prefix=doc_dict['job_id'] + '/',\n",
    "    delimiter='/')\n",
    "filenames = [blob.name for blob in blobs_listing]\n",
    "pdb_filenames = []\n",
    "\n",
    "print('Model Ranking')\n",
    "i = 1\n",
    "for f in filenames:\n",
    "    if 'pdb' in f:\n",
    "        print(f'Rank {i}: {Path(f).name}')\n",
    "        pdb_filenames.append(f)\n",
    "        i+=1\n",
    "\n",
    "bucket = gcs_client.bucket(doc_dict['job_gcs_path'])\n",
    "image_blobs = [i for i in filenames if '.png' in i]\n",
    "\n",
    "# Download all the images\n",
    "for i in image_blobs:\n",
    "    blob = bucket.blob(i)\n",
    "    blob.download_to_filename(Path(i).name)\n",
    "    if 'plddt.png' in Path(i).name:\n",
    "        image_to_display = Path(i).name\n",
    "\n",
    "display.Image(image_to_display, width=800)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Search for jobs by its labels\n",
    "\n",
    "List all the experiments with experiment name equals to the variable \"experiment_name\" and print the result of the execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a different label to search\n",
    "experiment_name = 'experiment-name'\n",
    "\n",
    "# Run Query\n",
    "docs = fs_client.collection('colabfold-experiments').where(\n",
    "    'labels.experiment_name', '==', experiment_name).limit(10).get()\n",
    "\n",
    "if len(docs) == 0:\n",
    "    raise ValueError('No experiment with this specific label.')\n",
    "\n",
    "print(f'Experiments with experience_name equals to: {experiment_name}')\n",
    "# Print results\n",
    "for idx, doc in enumerate(docs):\n",
    "    doc_dict = doc.to_dict()\n",
    "    print(f'=> Experiment {idx}: {doc_dict[\"job_id\"]}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize 3D protein"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next you will visualise the prediction & confidence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy PDB file locally\n",
    "blob = bucket.blob(pdb_filenames[0])\n",
    "blob.download_to_filename(Path(pdb_filenames[0]).name)\n",
    "prediction_filename = Path(pdb_filenames[0]).name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display protein and confidence\n",
    "rank_num = 1\n",
    "color = \"lDDT\"\n",
    "show_sidechains = False\n",
    "show_mainchains = False\n",
    "\n",
    "def show_pdb(rank_num=1, show_sidechains=False, show_mainchains=False, color=\"lDDT\"):\n",
    "  model_name = f\"rank_{rank_num}\"\n",
    "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
    "  view.addModel(open(prediction_filename,'r').read(),'pdb')\n",
    "\n",
    "  if color == \"lDDT\":\n",
    "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})\n",
    "\n",
    "  if show_sidechains:\n",
    "    BB = ['C','O','N']\n",
    "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
    "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
    "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
    "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
    "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
    "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})  \n",
    "  if show_mainchains:\n",
    "    BB = ['C','O','N','CA']\n",
    "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
    "\n",
    "  view.zoomTo()\n",
    "  return view\n",
    "\n",
    "show_pdb(rank_num, show_sidechains, show_mainchains, color).show()\n",
    "if color == \"lDDT\":\n",
    "  plot_plddt_legend().show() "
   ]
  }
 ],
 "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.7.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}