From f90cd0d6edab4537dce01ff18ead0d00efb76e93 Mon Sep 17 00:00:00 2001 From: Amit Rai <32944233+raiamitgit@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:54:13 -0400 Subject: [PATCH] Add AlphaFold 3 quickstart notebook (#4572) * Add AlphaFold 3 quickstart notebook * Update CODEOWNERS --------- Co-authored-by: Amit Rai --- notebooks/community/CODEOWNERS | 1 + notebooks/community/alphafold3/README.md | 35 ++ .../cloudai_alphafold3_vai_quickstart.ipynb | 572 ++++++++++++++++++ 3 files changed, 608 insertions(+) create mode 100644 notebooks/community/alphafold3/README.md create mode 100644 notebooks/community/alphafold3/cloudai_alphafold3_vai_quickstart.ipynb diff --git a/notebooks/community/CODEOWNERS b/notebooks/community/CODEOWNERS index eabdf619f..957372edd 100644 --- a/notebooks/community/CODEOWNERS +++ b/notebooks/community/CODEOWNERS @@ -28,6 +28,7 @@ /vertex_endpoints/optimized_tensorflow_runtime @vlasenkoalexey /notebooks/community/alphagenome/cloudai_alphagenome_vai_quickstart.ipynb @dpanigra /notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb @dpanigra +/notebooks/community/alphafold3/cloudai_alphafold3_vai_quickstart.ipynb @raiamitgit /notebooks/community/weathernext/weathernext_2_early_access_program.ipynb @dpanigra /notebooks/community/weathernext/weathernext_2_ic_early_access_program.ipynb @dpanigra /notebooks/community/weathernext/weathernext_2_dws.ipynb @dpanigra diff --git a/notebooks/community/alphafold3/README.md b/notebooks/community/alphafold3/README.md new file mode 100644 index 000000000..fb63d0f21 --- /dev/null +++ b/notebooks/community/alphafold3/README.md @@ -0,0 +1,35 @@ +# AlphaFold 3 +[**Overview**](#overview) | [**Use cases**](#use-cases) | [**Documentation**](#documentation) | [**Prerequisites**](#prerequisites) | [**Quick start**](#quick-start) + +## Overview +AlphaFold 3 is a revolutionary model developed by Google DeepMind and Isomorphic Labs that predicts the 3D structures and interactions of proteins, DNA, RNA, ligands, and chemical modifications. + +By modeling these molecules and their interactions together in a unified diffusion-based architecture, AlphaFold 3 provides a comprehensive view of cellular machinery, enabling researchers to understand biological processes at atomic resolution. + +AlphaFold 3 is available for commercial use on [Gemini Enterprise Agent Platform](https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/open-models/alphafold-3). + +## Use cases +* **Protein-Ligand Interaction Prediction**: Model the binding of small molecule ligands to proteins, enabling drug discovery and development. +* **Nucleic Acid Interaction Prediction**: Predict the complex structures of proteins interacting with DNA and RNA sequences. +* **Chemical Modifications**: Predict structures containing modified residues, ions, and covalent linkages. +* **Antibody-Antigen Modeling**: Map the 3D structures of antibody-antigen complexes to support therapeutic antibody design. + +## Documentation +The examples provided here demonstrate how to deploy and use AlphaFold 3 on Gemini Enterprise Agent Platform. + +### Links +* Read the [Nature journal paper](https://doi.org/10.1038/s41586-024-07487-w) +* Read the [Google DeepMind blog post](https://blog.google/technology/ai/google-deepmind-isomorphic-alphafold-3-ai-model/) +* Explore the [AlphaFold Server](https://alphafoldserver.com/welcome) +* View the open-source code and non-commercial weights on [GitHub](https://github.com/google-deepmind/alphafold3) + +## Prerequisites +To deploy and use AlphaFold 3 on Vertex AI: +1. **Request Access**: Submit the [AlphaFold 3 Request Form](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/alphafold3-request) and work with your Google Cloud account team for commercial subscription allowlisting. +2. **Hardware Quota**: Deployments require an `a3-highgpu-1g` machine type (1x NVIDIA H100 80GB GPU) with 750 GB Local SSD provisioned for database caching. +3. **Endpoint Configuration**: Deploy the model to a Dedicated Endpoint and configure the inference timeout to 3,600 seconds. + +## Quick start +| Notebook | Description | Links | +| :--- | :--- | :--- | +| [AlphaFold 3 Quickstart](cloudai_alphafold3_vai_quickstart.ipynb) | End-to-end protein-ligand docking prediction (KRAS G12C covalent complex with Sotorasib), output handling, and 3D visualization. | Open in Colab | diff --git a/notebooks/community/alphafold3/cloudai_alphafold3_vai_quickstart.ipynb b/notebooks/community/alphafold3/cloudai_alphafold3_vai_quickstart.ipynb new file mode 100644 index 000000000..4057edb49 --- /dev/null +++ b/notebooks/community/alphafold3/cloudai_alphafold3_vai_quickstart.ipynb @@ -0,0 +1,572 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "id": "ab7ebb3e" + }, + "outputs": [], + "source": [ + "# Copyright 2026 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", + "# https://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": { + "id": "abd825f8" + }, + "source": [ + "# AlphaFold 3: Protein-Ligand Docking Prediction Example\n", + "\n", + " \n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Google
Open in Colab\n", + " \n", + " 		\n", + " \n", + " \"Google
Open in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"Vertex
Open in Agent Platform Workbench\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6340d859" + }, + "source": [ + "## Overview\n", + "\n", + "AlphaFold 3 predicts the 3D structure of proteins, DNA, RNA, chemical modifications, and small-molecule ligands. It enables researchers to simulate physical interactions directly from sequence inputs and chemical descriptions.\n", + "\n", + "This tutorial covers an end-to-end prediction workflow for AlphaFold 3 deployed on Model Garden on the Google Cloud Gemini Enterprise Agent Platform.\n", + "\n", + "This tutorial covers:\n", + "1. Authenticating and connecting to the Google Cloud environment.\n", + "2. Submitting a prediction request. The example shows a covalent complex of the **KRAS G12C mutant protein** bound to the inhibitor **Sotorasib (AMG-510)**.\n", + "3. Understanding Output Structure & Artifacts.\n", + "4. Comparing & Filtering Candidate Samples.\n", + "5. Extracting & Rendering 3D Coordinates." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "26d7d56d" + }, + "source": [ + "## Getting Started\n", + "\n", + "### Install Dependencies\n", + "First, install the libraries required to download files from GCS, manipulate matrices, and visualize 3D molecular coordinates." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "id": "9a6b1c46" + }, + "outputs": [], + "source": [ + "%pip install -q py3Dmol biopython matplotlib numpy google-cloud-storage google-cloud-aiplatform" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d60b56b3" + }, + "source": [ + "### Authenticate the Notebook Environment (Colab Only)\n", + "If running this notebook in Google Colab, authenticate the account to grant access to Google Cloud resources." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "id": "b0f67731" + }, + "outputs": [], + "source": [ + "# Import required libraries\n", + "import json\n", + "import sys\n", + "\n", + "from google.cloud import aiplatform, storage\n", + "\n", + "# Authenticate environment for Google Colab only\n", + "if \"google.colab\" in sys.modules:\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e8a0bbcd" + }, + "source": [ + "## Prerequisite: Deploy AlphaFold 3 on Model Garden\n", + "\n", + "AlphaFold 3 is available as a self-deployed model on Model Garden on Gemini Enterprise Agent Platform. Users can deploy the model into their own secure Google Cloud environments. Refer to the official documentation for more details on deploying and scaling the model.\n", + "\n", + "**Note:** The next steps in this notebook assume an AlphaFold 3 model is deployed and actively running on a dedicated endpoint." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b8f1341b" + }, + "source": [ + "### Configure Project Variables\n", + "Set the Google Cloud Project ID, the region where the AlphaFold 3 endpoint is deployed, and the GCS bucket where inference outputs will be saved." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "id": "0d2f75b0" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "LOCATION = \"\" # @param {type:\"string\"}\n", + "ENDPOINT_ID = \"\" # @param {type:\"string\"}\n", + "BUCKET_NAME = \"\" # @param {type:\"string\"}\n", + "JOB_NAME = \"\" # @param {type:\"string\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "14a24e4b" + }, + "source": [ + "## Step 1: Submit the Prediction Request\n", + "\n", + "Configure the AlphaFold 3 request JSON payload. The example showcases how to predict interactions for two molecules:\n", + "1. **KRAS G12C Protein Chain (Entity A)**: A 189 amino acid mutant protein.\n", + "2. **Sotorasib Ligand (Entity B)**: Defined by its CCD code `MOV`.\n", + "\n", + "To model the covalent interaction, explicitly define the covalent bond in the `bondedAtomPairs` list. Specify that the Sulfur-Gamma (`SG`) atom of residue 12 (Cysteine) on the protein chain (`A`) forms a covalent bond with the Carbon 25 (`C25`) atom of the Sotorasib ligand (`B`).\n", + "\n", + "**Note on Output Handling:** While the response can be returned directly through the API, it is highly recommended to provide a GCS bucket (`output_dir`) to persist the output. This also enables capturing a more detailed output dataset." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "c3017b66" + }, + "outputs": [], + "source": [ + "# --- Define AlphaFold 3 Inputs ---\n", + "SEEDS = [42, 43]\n", + "\n", + "PROTEIN_ID = \"A\"\n", + "PROTEIN_SEQ = \"MTEYKLVVVGACGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGETCLLDILDTAGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHHYREQIKRVKDSEDVPMVLVGNKCDLPSRTVDTKQAQDLARSYGIPFIETSAKTRQRVEDAFYTLVREIRQYRLKKISKEEKTPGCVKIKKCIIM\"\n", + "PROTEIN_DESC = \"Human KRAS4B G12C mutant (UniProt P01116 with G12C, 189 aa)\"\n", + "\n", + "LIGAND_ID = \"B\"\n", + "LIGAND_CCD = \"MOV\"\n", + "LIGAND_DESC = \"Sotorasib / AMG-510 (CCD: MOV), covalent KRAS G12C inhibitor\"\n", + "\n", + "# Define covalent bond pair: [atom_1, atom_2] where atom = [entity_id, residue_index, atom_name]\n", + "BONDED_ATOM_PAIRS = [[[\"A\", 12, \"SG\"], [\"B\", 1, \"C25\"]]]\n", + "\n", + "OUTPUT_DIR = f\"gs://{BUCKET_NAME}/my_dir\"\n", + "\n", + "# --- Construct AlphaFold 3 Prediction Request Payload ---\n", + "# --- Overwrites the output directory, if one exists already ---\n", + "payload = {\n", + " \"instances\": [\n", + " {\n", + " \"name\": JOB_NAME,\n", + " \"dialect\": \"alphafold3\",\n", + " \"version\": 4,\n", + " \"modelSeeds\": SEEDS,\n", + " \"sequences\": [\n", + " {\n", + " \"protein\": {\n", + " \"id\": PROTEIN_ID,\n", + " \"sequence\": PROTEIN_SEQ,\n", + " \"description\": PROTEIN_DESC,\n", + " }\n", + " },\n", + " {\n", + " \"ligand\": {\n", + " \"id\": LIGAND_ID,\n", + " \"ccdCodes\": [LIGAND_CCD],\n", + " \"description\": LIGAND_DESC,\n", + " }\n", + " },\n", + " ],\n", + " \"bondedAtomPairs\": BONDED_ATOM_PAIRS,\n", + " }\n", + " ],\n", + " \"parameters\": {\n", + " \"output_dir\": OUTPUT_DIR,\n", + " \"force_output_dir\": True,\n", + " \"run_data_pipeline\": True,\n", + " },\n", + "}\n", + "\n", + "# --- Submit Prediction Request ---\n", + "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n", + "endpoint = aiplatform.Endpoint(ENDPOINT_ID)\n", + "\n", + "response = endpoint.predict(\n", + " instances=payload[\"instances\"], parameters=payload[\"parameters\"]\n", + ")\n", + "\n", + "print(\"Prediction Request Submitted Successfully:\")\n", + "print(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PVcgUiuPqHvV" + }, + "source": [ + "## Step 2: Inline Response vs. Saved Artifacts\n", + "\n", + "AlphaFold 3 provides two delivery mechanisms for retrieving structural prediction outputs. By default, the API response returns prediction results inline over HTTP. For production workloads, specifying a Google Cloud Storage (GCS) bucket exports the complete multi-sample artifact dataset." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "id": "7zYSywvOqHvV" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Downloaded top-ranked model.cif, confidence.json, and summary.json successfully.\n" + ] + } + ], + "source": [ + "def download_af3_outputs(\n", + " bucket_name, job_name, seed=None, sample_idx=None, local_dir=\".\"\n", + "):\n", + " \"\"\"Downloads CIF coordinates, confidences, and summary files from GCS.\n", + "\n", + " If seed and sample_idx are None, downloads the overall #1 top-ranked candidate\n", + " model directly from the top-level job directory.\n", + " \"\"\"\n", + " if bucket_name == \"your-bucket-name\":\n", + " print(\"Using local sample data files for demonstration.\")\n", + " print(\n", + " \"Note: Update BUCKET_NAME with your Google Cloud Storage bucket to download live prediction outputs.\"\n", + " )\n", + " return\n", + "\n", + " client = storage.Client(project=PROJECT_ID)\n", + " bucket = client.bucket(bucket_name)\n", + "\n", + " if seed is not None and sample_idx is not None:\n", + " gcs_prefix = f\"my_dir/{job_name}/seed-{seed}_sample-{sample_idx}\"\n", + " file_prefix = f\"{job_name}_seed-{seed}_sample-{sample_idx}\"\n", + " else:\n", + " gcs_prefix = f\"my_dir/{job_name}\"\n", + " file_prefix = job_name\n", + "\n", + " cif_blob = bucket.blob(f\"{gcs_prefix}/{file_prefix}_model.cif\")\n", + " cif_blob.download_to_filename(\"model.cif\")\n", + "\n", + " conf_blob = bucket.blob(f\"{gcs_prefix}/{file_prefix}_confidences.json\")\n", + " conf_blob.download_to_filename(\"confidence.json\")\n", + "\n", + " summary_blob = bucket.blob(f\"{gcs_prefix}/{file_prefix}_summary_confidences.json\")\n", + " summary_blob.download_to_filename(\"summary.json\")\n", + "\n", + " # Download ranking scores CSV if present\n", + " try:\n", + " ranking_blob = bucket.blob(f\"my_dir/{job_name}/{job_name}_ranking_scores.csv\")\n", + " ranking_blob.download_to_filename(\"ranking_scores.csv\")\n", + " except Exception:\n", + " pass\n", + "\n", + " print(\n", + " \"Downloaded top-ranked model.cif, confidence.json, and summary.json successfully.\"\n", + " )\n", + "\n", + "\n", + "# Run the download function for the prediction outputs (fetches overall top-ranked structure)\n", + "download_af3_outputs(BUCKET_NAME, JOB_NAME)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1udHnz9EqHvV" + }, + "source": [ + "## Step 3: Comparing Candidate Samples\n", + "\n", + "Because AlphaFold 3 utilizes a generative diffusion module, each run produces an ensemble of candidate structures across random seeds. Inspect `ranking_scores.csv` and `summary.json` **first** to identify the top-ranked candidate and verify that the model is free of steric clashes (`has_clash == 0`)." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "id": "ZNhOJL15qHvV" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "=== Candidate Samples Ledger (ranking_scores.csv) ===\n", + "seed sample ranking_score\n", + " 42 0 0.9765861758441248\n", + " 43 3 0.9737541527867797\n", + " 43 4 0.9733747951235855\n", + " 42 2 0.9732408715975229\n", + " 42 3 0.9723226792181007\n", + " 42 1 0.970259597191794\n", + " 43 0 0.968876439745061\n", + " 43 2 0.9686427701415588\n", + " 43 1 0.9618608651497567\n", + " 42 4 0.956112724247255\n", + "\n", + "Promoted Top Candidate: Seed 42, Sample 0 (Score: 0.9766)\n", + "\n", + "=== Top Candidate Quality Validation (summary.json) ===\n", + "Top Candidate Metrics:\n", + " • Ranking Score : 0.98\n", + " • Global pTM : 0.85\n", + " • Interface ipTM: 0.95\n", + " • Steric Clash : None Detected (PASS)\n" + ] + } + ], + "source": [ + "import csv\n", + "\n", + "print(\"=== Candidate Samples Ledger (ranking_scores.csv) ===\")\n", + "with open(\"ranking_scores.csv\", newline=\"\", encoding=\"utf-8\") as f:\n", + " rows = sorted(\n", + " csv.DictReader(f), key=lambda x: float(x[\"ranking_score\"]), reverse=True\n", + " )\n", + "\n", + "# Calculate column widths cleanly and readably\n", + "headers = list(rows[0].keys())\n", + "widths = {}\n", + "for col in headers:\n", + " lengths = [len(col)] + [len(r[col]) for r in rows]\n", + " widths[col] = max(lengths)\n", + "\n", + "print(\" \".join(col.rjust(widths[col]) for col in headers))\n", + "for r in rows:\n", + " print(\" \".join(r[col].rjust(widths[col]) for col in headers))\n", + "\n", + "top = rows[0]\n", + "print(\n", + " f\"\\nPromoted Top Candidate: Seed {int(top['seed'])}, Sample\"\n", + " f\" {int(top['sample'])} (Score: {float(top['ranking_score']):.4f})\"\n", + ")\n", + "\n", + "print(\"\\n=== Top Candidate Quality Validation (summary.json) ===\")\n", + "with open(\"summary.json\", encoding=\"utf-8\") as f:\n", + " summary = json.load(f)\n", + "\n", + "clash_str = \"DETECTED (FAIL)\" if summary.get(\"has_clash\") else \"None Detected (PASS)\"\n", + "print(\"Top Candidate Metrics:\")\n", + "print(f\" • Ranking Score : {summary.get('ranking_score', 'N/A')}\")\n", + "print(f\" • Global pTM : {summary.get('ptm', 'N/A')}\")\n", + "print(f\" • Interface ipTM: {summary.get('iptm', 'N/A')}\")\n", + "print(f\" • Steric Clash : {clash_str}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ThhRv85wqHvV" + }, + "source": [ + "## Step 4: Inspecting 3D Structures\n", + "\n", + "When reviewing the downloaded atomic structure `model.cif` for the top-ranked candidate, note key structural conventions: heavy-atom representation ($x,y,z$) and per-atom local confidence (pLDDT) stored in the B-factor column (`_atom_site.B_iso_or_equiv`)." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "id": "Hi1ViaklqHvV" + }, + "outputs": [ + { + "data": { + "application/3dmoljs_load.v0": "
\n

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n
\n", + "text/html": [ + "
\n", + "

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

\n", + "
\n", + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "None" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import py3Dmol\n", + "from Bio.PDB import MMCIFParser\n", + "from IPython.display import display\n", + "\n", + "with open(\"model.cif\") as f:\n", + " cif_data = f.read()\n", + "\n", + "# Map AlphaFold pLDDT confidence color bands\n", + "color_map = {\n", + " \"#ff7d45\": [], # Very Low (<50)\n", + " \"#ffdb13\": [], # Low (50-70)\n", + " \"#65cbf3\": [], # Confident (70-90)\n", + " \"#0053d6\": [], # Very High (>=90)\n", + "}\n", + "\n", + "# Parse mmCIF file using BioPython\n", + "parser = MMCIFParser(QUIET=True)\n", + "structure = parser.get_structure(\"af3\", \"model.cif\")\n", + "\n", + "# Extract atom serial numbers based on B-factor (pLDDT)\n", + "for atom in structure.get_atoms():\n", + " serial = atom.get_serial_number()\n", + " b_factor = atom.bfactor\n", + "\n", + " if b_factor < 50:\n", + " color_map[\"#ff7d45\"].append(serial)\n", + " elif b_factor < 70:\n", + " color_map[\"#ffdb13\"].append(serial)\n", + " elif b_factor < 90:\n", + " color_map[\"#65cbf3\"].append(serial)\n", + " else:\n", + " color_map[\"#0053d6\"].append(serial)\n", + "\n", + "# Initialize py3Dmol viewer\n", + "view = py3Dmol.view(width=800, height=500)\n", + "view.addModel(cif_data, \"cif\")\n", + "\n", + "# Apply pLDDT cartoon styling to polymer backbones\n", + "view.setStyle({\"model\": -1}, {})\n", + "for color, atoms in color_map.items():\n", + " if atoms:\n", + " view.addStyle({\"serial\": atoms}, {\"cartoon\": {\"color\": color}})\n", + "\n", + "# Render covalent ligand Sotorasib (CCD: MOV) explicitly as green-carbon sticks\n", + "view.addStyle({\"hetflag\": True}, {\"stick\": {\"colorscheme\": \"greenCarbon\"}})\n", + "\n", + "view.zoomTo()\n", + "view.spin(True)\n", + "display(view.show())" + ] + } + ], + "metadata": { + "colab": { + "name": "cloudai_alphafold3_vai_quickstart.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +}