Esmfold2食用体验-蛋白质结构预测

# If you are working in colab, uncomment these lines to install dependencies
!pip install esm@git+https://github.com/Biohub/esm.git@c94ed8d
!pip install py3dmol

from getpass import getpass
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import py3Dmol
from esm.sdk import esmfold2_client
from esm.sdk.api import FoldingConfig
from esm.utils.structure import input_builder

from IPython.display import HTML, display


def plddt_hex(v):
    """Convert pLDDT score to hex color."""
    if v >= 90:
        return"#0053D6"
    if v >= 70:
        return"#65CBF3"
    if v >= 50:
        return"#FFDB13"
    return"#FF7D45"


PLDDT_LEGEND = (
    '<span style="color:#FF7D45;">&#9632;</span> &lt;50 &nbsp;'
    '<span style="color:#FFDB13;">&#9632;</span> 50–70 &nbsp;'
    '<span style="color:#65CBF3;">&#9632;</span> 70–90 &nbsp;'
    '<span style="color:#0053D6;">&#9632;</span> &gt;90'
)


def normalize_plddt(plddt):
    plddt = np.array(plddt)
    return plddt * 100if plddt.max() <= 1.0else plddt


def apply_plddt_style(view, plddt, add_hetflag=False):
    view.setStyle({}, {})
    for i, s in enumerate(plddt):
        view.setStyle({"resi": i + 1}, {"cartoon": {"color": plddt_hex(s)}})
    if add_hetflag:
        view.addStyle({"hetflag": True}, {"stick": {}})
    view.zoomTo()


def show_dual_viewer(v1, title1, legend1, v2, title2, legend2):
    l1 = (
        f'<p style="text-align:center; font-size:12px; margin-top:4px;">{legend1}</p>'
        if legend1
        else""
    )
    l2 = (
        f'<p style="text-align:center; font-size:12px; margin-top:4px;">{legend2}</p>'
        if legend2
        else""
    )
    display(
        HTML(
            '<div style="display:flex; gap:20px;">'
            f'<div><p style="text-align:center; font-weight:bold;">{title1}</p>'
            + v1._make_html()
            + l1
            + "</div>"
            f'<div><p style="text-align:center; font-weight:bold;">{title2}</p>'
            + v2._make_html()
            + l2
            + "</div>"
            "</div>"
        )
    )


def plot_pae(pae, chains=None, title=None):
    """Plot a PAE matrix with optional per-chain colored bars and legend.

    Args:
        pae: PAE as a numpy array (n x n).
        chains: list of (start, end, label, color) tuples. When provided, colored
            bars are drawn along the top and left axes and a legend is shown.
        title: optional plot title.
    """
    total_len = pae.shape[0]

    fig, ax = plt.subplots(figsize=(8, 8))
    im = ax.imshow(
        pae, cmap="Greens_r", vmin=0, vmax=30, origin="upper", aspect="equal"
    )

    if chains:
        bar_thickness = total_len * 0.015
        for start, end, label, color in chains:
            ax.add_patch(
                mpatches.Rectangle(
                    (start - 0.5, -bar_thickness - 1),
                    end - start,
                    bar_thickness,
                    facecolor=color,
                    clip_on=False,
                    edgecolor="black",
                    linewidth=0.5,
                )
            )
            ax.add_patch(
                mpatches.Rectangle(
                    (-bar_thickness - 1, start - 0.5),
                    bar_thickness,
                    end - start,
                    facecolor=color,
                    clip_on=False,
                    edgecolor="black",
                    linewidth=0.5,
                )
            )
        for start, *_ in chains[1:]:
            ax.axhline(start - 0.5, color="black", linewidth=1.5)
            ax.axvline(start - 0.5, color="black", linewidth=1.5)
        legend_handles = [
            mpatches.Patch(facecolor=c, edgecolor="black", label=l)
            for _, _, l, c in chains
        ]
        ax.legend(
            handles=legend_handles,
            loc="upper center",
            bbox_to_anchor=(0.5, -0.08),
            ncol=len(chains),
            frameon=True,
            fontsize=10,
            title="Chain",
        )
        ax.set_xlim(-bar_thickness - 2, total_len)
        ax.set_ylim(total_len, -bar_thickness - 2)

    ax.set_xlabel("Scored Residue", fontsize=12)
    ax.set_ylabel("Aligned Residue", fontsize=12)
    if title:
        ax.set_title(title, fontsize=14, fontweight="bold", pad=15)

    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label("Expected Position Error (Ångströms)", fontsize=12)
    cbar.set_ticks([0, 5, 10, 15, 20, 25, 30])
    cbar.ax.invert_yaxis()

    plt.tight_layout()
    plt.show()

token = getpass("Biohub token: ")

client = esmfold2_client(
    model="esmfold2-fast-2026-05", url="https://biohub.ai", token=token
)

config = FoldingConfig(num_loops=3, num_sampling_steps=32, include_pae=True)

# Define all components from PDB 4H8K (LC11-RNase H1 + RNA/DNA hybrid)

# LC11-RNase H1 protein sequence (140 amino acids, chains A and B share this sequence)
rnaseh_sequence = (
    "MNKIIIYTDGGARGNPGPAGIGVVITDEKGNTLHESSAYIGETTNNVAEYEALIRALEDLQ"
    "MFGDKLVDMEVEVRMNSELIVRQMQGVYKVKEPTLKEKFAKIAHIKMERVPNLVFVHIPRE"
    "KNARADELVNEAIDKALS"
)

# Substrate RNA sequence (14 nucleotides)
rna_sequence = "CGACACCUGAUUCC"

# Substrate DNA sequence (14 nucleotides, complementary to the RNA)
dna_sequence = "GGAATCAGGTGTCG"

# Express the homodimer as a single ProteinInput with two chain IDs
rnaseh_protein = input_builder.ProteinInput(id=["A", "B"], sequence=rnaseh_sequence)
substrate_rna = input_builder.RNAInput(id="C", sequence=rna_sequence)
substrate_dna = input_builder.DNAInput(id="D", sequence=dna_sequence)

# Create the complex input
rnaseh_complex_input = input_builder.StructurePredictionInput(
    sequences=[rnaseh_protein, substrate_rna, substrate_dna]
)

rnaseh_structure = client.fold_all_atom(rnaseh_complex_input, config=config)

# Save the structure
with open("rnaseh_complex_4h8k.cif", "w") as f:
    f.write(rnaseh_structure.complex.to_mmcif())

rnaseh_cif = open("rnaseh_complex_4h8k.cif").read()
rnaseh_plddt = normalize_plddt(rnaseh_structure.plddt.cpu().numpy())

v1 = py3Dmol.view(width=500, height=400)
v1.addModel(rnaseh_cif, "mmcif")
v1.setStyle({"chain": "A"}, {"cartoon": {"color": "#4A90E2"}})
v1.setStyle({"chain": "B"}, {"cartoon": {"color": "#7FB3E8"}})
v1.setStyle({"chain": "C"}, {"cartoon": {"color": "#F5A623"}})
v1.setStyle({"chain": "D"}, {"cartoon": {"color": "#50E3C2"}})
v1.zoomTo()

v2 = py3Dmol.view(width=500, height=400)
v2.addModel(rnaseh_cif, "mmcif")
apply_plddt_style(v2, rnaseh_plddt)

rnaseh_legend = (
    '<span style="color:#4A90E2;">&#9632;</span> RNase H1 (A) &nbsp;'
    '<span style="color:#7FB3E8;">&#9632;</span> RNase H1 (B) &nbsp;'
    '<span style="color:#F5A623;">&#9632;</span> RNA &nbsp;'
    '<span style="color:#50E3C2;">&#9632;</span> DNA'
)
show_dual_viewer(
    v1,
    "RNase H1–RNA/DNA Hybrid Complex",
    rnaseh_legend,
    v2,
    "Confidence (pLDDT)",
    PLDDT_LEGEND,
)

# Human GLP1R sequence (UniProt P43220)
# This construct includes N-terminal FLAG tag, TEV cleavage site, and C-terminal His-tag
receptor_sequence = (
    "MKTIIALSYIFCLVFADYKDDDDLEVLFQGPARPQGATVSLWETVQKWREYRRQCQRSLTEDPPPATDLFCNRTFDEYAC"
    "WPDGEPGSFVNVSCPWYLPWASSVPQGHVYRFCTAEGLWLQKDNSSLPWRDLSECEESKRGERSSPEEQLLFLYIIYTVG"
    "YALSFSALVIASAILLGFRHLHCTRNYIHLNLFASFILRALSVFIKDAALKWMYSTAAQQHQWDGLLSYQDSLSCRLVFL"
    "LMQYCVAANYYWLLVEGVYLYTLLAFSVFSEQWIFRLYVSIGWGVPLLFVVPWGIVKYLYEDEGCWTRNSNMNYWLIIRL"
    "PILFAIGVNFLIFVRVICIVVSKLKANLMCKTDIKCRLAKSTLTLIPLLGTHEVIFAFVMDEHARGTLRFIKLFTELSFT"
    "SFQGLMVAILYCFVNNEVQLEFRKSWERWRLEHLHIQRDSSMKPLKCPTSSLSSGATAGSSMYTATCQASCSPAGLEVLF"
    "QGPHHHHHHH"
)

receptor = input_builder.ProteinInput(id="A", sequence=receptor_sequence)

# Semaglutide peptide sequence
peptide_sequence = "HAEGTFTSDVSSYLEGQAAKEFIAWLVRGRG"

# AIB substitution (0-indexed)
aib_position = 1

peptide = input_builder.ProteinInput(
    id="B",
    sequence=peptide_sequence,
    modifications=[input_builder.Modification(position=aib_position, ccd="AIB")],
)

lys_idx = peptide_sequence.index("K")  # residue index in peptide (0-based)
print("Lys at residue index:", lys_idx)

def get_protractor_smiles() -> str:
    # "Protractor" refers to Novo Nordisk's fatty acid linker technology used in semaglutide.
    # This SMILES is a representative structure inspired by semaglutide's linker design,
    # not the exact compound. For the precise semaglutide structure, see PDB 7KI0 or the FDA label.
    # First atom (index 0) is the carbonyl carbon used for attachment to the peptide.
    return "C(=O)(CCOCCOCC(=O)NCCOCCOCCNCC(=O)N[C@@H](CCC(=O)NCCCCCCCCCCCCCCCCCC(=O)O)C(=O)O)"


ligand = input_builder.LigandInput(id="C", ccd=None, smiles=get_protractor_smiles())

bond = input_builder.CovalentBond(
    chain_id1="B", res_idx1=lys_idx, atom_idx1=8, chain_id2="C", res_idx2=0, atom_idx2=0
)

complex_inputs = input_builder.StructurePredictionInput(
    sequences=[receptor, peptide, ligand], covalent_bonds=[bond]
)

complex_structure = client.fold_all_atom(complex_inputs, config=config)

with open("glp1r_semaglutide.cif", "w") as f:
    f.write(complex_structure.complex.to_mmcif())

# download the MSA
!wget -q "https://drive.google.com/uc?export=download&id=1la9UUK_FnFQFR9VB35DIB797d87Zwm0b" -O ubiquitin.a3m

from esm.utils.msa import MSA
from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput

ubiquitin_sequence = (
    "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"
)

msa = MSA.from_a3m(path="ubiquitin.a3m", remove_insertions=True, max_sequences=1000)
print(f"MSA depth: {msa.depth} sequences, length: {msa.seqlen} residues")

# Sanity check — first MSA seq should match the query
print("First MSA seq matches query:", msa.sequences[0] == ubiquitin_sequence)
print("First MSA seq:", repr(msa.sequences[0][:80]))
print("Query:        ", repr(ubiquitin_sequence))

fold_result = client.fold_all_atom(
    StructurePredictionInput(
        sequences=[ProteinInput(id="A", sequence=ubiquitin_sequence, msa=msa)]
    )
)

print(f"\npTM: {fold_result.ptm:.3f}")
print(f"Average pLDDT: {fold_result.plddt.mean().item():.1f}")

with open("ubiquitin_with_msa.cif", "w") as f:
    f.write(fold_result.complex.to_mmcif())