On this post we are going to have a look at how we are able to visualize proteins on Hugging Face Spaces.
Update May 2024
While the tactic described below still works, you will probably want to avoid wasting a while and use the Molecule3D Gradio Custom Component. This component will allow users to change the protein visualization on the fly and you may more easily set the default visualization. Simply install it using:
pip install gradio_molecule3d
from gradio_molecule3d import Molecule3D
reps = [
{
"model": 0,
"chain": "",
"resname": "",
"style": "stick",
"color": "whiteCarbon",
"residue_range": "",
"around": 0,
"byres": False,
}
]
with gr.Blocks() as demo:
Molecule3D(reps=reps)
Motivation 🤗
Proteins have a big impact on our life – from medicines to washing powder. Machine learning on proteins is a rapidly growing field to assist us design latest and interesting proteins. Proteins are complex 3D objects generally composed of a series of constructing blocks called amino acids which can be arranged in 3D space to provide the protein its function. For machine learning purposes a protein can for instance be represented as coordinates, as graph or as 1D sequence of letters to be used in a protein language model.
A famous ML model for proteins is AlphaFold2 which predicts the structure of a protein sequence using a multiple sequence alignment of comparable proteins and a structure module.
Since AlphaFold2 made its debut many more such models have come out resembling OmegaFold, OpenFold etc. (see this list or this list for more).
Seeing is believing
The structure of a protein is an integral part to our understanding of what a protein does. Nowadays, there are just a few tools available to visualise proteins directly within the browser resembling mol* or 3dmol.js. On this post, you’ll learn the way to integrate structure visualization into your Hugging Face Space using 3Dmol.js and the HTML block.
Prerequisites
Make certain you might have the gradio Python package already installed and basic knowledge of Javascript/JQuery.
Taking a Take a look at the Code
Let’s take a have a look at the way to create the minimal working demo of our interface before we dive into the way to setup 3Dmol.js.
We are going to construct an easy demo app that may accept either a 4-digit PDB code or a PDB file. Our app will then retrieve the pdb file from the RCSB Protein Databank and display it or use the uploaded file for display.
import gradio as gr
def update(inp, file):
pdb_path = get_pdb(inp, file)
return molecule(pdb_path)
demo = gr.Blocks()
with demo:
gr.Markdown("# PDB viewer using 3Dmol.js")
with gr.Row():
with gr.Box():
inp = gr.Textbox(
placeholder="PDB Code or upload file below", label="Input structure"
)
file = gr.File(file_count="single")
btn = gr.Button("View structure")
mol = gr.HTML()
btn.click(fn=update, inputs=[inp, file], outputs=mol)
demo.launch()
update: That is the function that does the processing of our proteins and returns an iframe with the viewer
Our get_pdb function can be easy:
import os
def get_pdb(pdb_code="", filepath=""):
if pdb_code is None or len(pdb_code) != 4:
try:
return filepath.name
except AttributeError as e:
return None
else:
os.system(f"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb")
return f"{pdb_code}.pdb"
Now, the way to visualize the protein since Gradio doesn’t have 3Dmol directly available as a block?
We use an iframe for this.
Our molecule function which returns the iframe conceptually looks like this:
def molecule(input_pdb):
mol = read_mol(input_pdb)
x = (""" [..] """)
return f"""">
It is a bit clunky to setup but is essential due to the security rules in modern browsers.
3Dmol.js setup is pretty easy and the documentation provides a few examples.
The head of our returned document must load 3Dmol.js (which in turn also loads JQuery).
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8" />
<style>
.mol-container {
width: 100%;
height: 700px;
position: relative;
}
.mol-container select{
background-image:None;
}
style>
<script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer">script>
<script src="https://3Dmol.csb.pitt.edu/construct/3Dmol-min.js">script>
head>
The styles for .mol-container might be used to change the dimensions of the molecule viewer.
The body looks as follows:
<body>
<div id="container" class="mol-container">div>
<script>
let pdb = mol
$(document).ready(function () {
let element = $("#container");
let config = { backgroundColor: "white" };
let viewer = $3Dmol.createViewer(element, config);
viewer.addModel(pdb, "pdb");
viewer.getModel(0).setStyle({}, { cartoon: { colorscheme:"whiteCarbon" } });
viewer.zoomTo();
viewer.render();
viewer.zoom(0.8, 2000);
})
script>
body>
We use a template literal (denoted by backticks) to store our pdb file within the html document directly after which output it using 3dmol.js.
And that is it, now you may couple your favorite protein ML model to a fun and straightforward to make use of gradio app and directly visualize predicted or redesigned structures. When you are predicting properities of a structure (e.g how likely each amino acid is to bind a ligand), 3Dmol.js also allows to make use of a custom colorfunc based on a property of every atom.
You possibly can check the source code of the 3Dmol.js space for the total code.
For a production example, you may check the ProteinMPNN space where a user can upload a backbone, the inverse folding model ProteinMPNN predicts latest optimal sequences after which one can run AlphaFold2 on all predicted sequences to confirm whether or not they adopt the initial input backbone. Successful redesigns that qualitiatively adopt the identical structure as predicted by AlphaFold2 with high pLDDT rating must be tested within the lab.
Issues
When you encounter any issues with the mixing of 3Dmol.js in Gradio/HF spaces, please open a discussion in hf.space/simonduerr/3dmol.js.
If you might have problems with 3Dmol.js configuration – you’ll want to ask the developers, please, open a 3Dmol.js Issue as a substitute and describe your problem.
