Diffdock Agent Skill

# Recommended: clone and install from source
git clone https://github.com/gcorso/DiffDock.git
cd DiffDock
conda create -n diffdock python=3.9
conda activate diffdock
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

# Download pretrained model weights
python -c "from utils.download import download_pretrained; download_pretrained()"

from Bio import PDB
from Bio.PDB import PDBParser, PDBIO, Select

class NonHetSelect(Select):
    """Remove HETATM records (ligands, water) — keep only protein atoms."""
    def accept_residue(self, residue):
        return residue.id[0] == " "

def clean_pdb(input_pdb: str, output_pdb: str):
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("protein", input_pdb)
    io = PDBIO()
    io.set_structure(structure)
    io.save(output_pdb, NonHetSelect())
    print(f"Cleaned PDB saved to: {output_pdb}")

clean_pdb("raw_protein.pdb", "protein_clean.pdb")

from rdkit import Chem
from rdkit.Chem import AllChem, SDWriter

def smiles_to_sdf(smiles: str, output_sdf: str, n_confs: int = 1):
    """Convert SMILES to 3D SDF for DiffDock input."""
    mol = Chem.MolFromSmiles(smiles)
    mol = Chem.AddHs(mol)
    AllChem.EmbedMolecule(mol, AllChem.ETKDGv3())
    AllChem.MMFFOptimizeMolecule(mol)
    writer = SDWriter(output_sdf)
    writer.write(mol)
    writer.close()
    print(f"Ligand SDF written to: {output_sdf}")
    return mol

# Example: ibuprofen
smiles = "CC(C)Cc1ccc(cc1)C(C)C(=O)O"
mol    = smiles_to_sdf(smiles, "ligand.sdf")
print(f"Ligand formula: {Chem.rdMolDescriptors.CalcMolFormula(mol)}")

# Command-line inference (run from the DiffDock directory)
python inference.py \
    --protein_path protein_clean.pdb \
    --ligand       "CC(C)Cc1ccc(cc1)C(C)C(=O)O" \
    --out_dir      results/ \
    --inference_steps 20 \
    --samples_per_complex 40 \
    --batch_size 10 \
    --no_final_step_noise

import subprocess

def run_diffdock(protein_pdb: str, ligand_smiles: str, out_dir: str,
                 n_samples: int = 40, n_steps: int = 20):
    cmd = [
        "python", "inference.py",
        "--protein_path",      protein_pdb,
        "--ligand",            ligand_smiles,
        "--out_dir",           out_dir,
        "--inference_steps",   str(n_steps),
        "--samples_per_complex", str(n_samples),
        "--batch_size",        "10",
        "--no_final_step_noise",
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, cwd="DiffDock/")
    if result.returncode == 0:
        print(f"DiffDock complete. Results in: {out_dir}")
    else:
        print(f"Error: {result.stderr}")
    return result

run_diffdock("protein_clean.pdb", "CC(C)Cc1ccc(cc1)C(C)C(=O)O", "results/")

import re
from pathlib import Path
import pandas as pd

def parse_diffdock_results(out_dir: str) -> pd.DataFrame:
    """Parse DiffDock output SDF files and confidence scores."""
    out_path = Path(out_dir)
    records  = []

    # DiffDock names output files: rank{N}_confidence{score}.sdf
    for sdf_file in sorted(out_path.glob("rank*_confidence*.sdf")):
        name = sdf_file.stem
        # Extract rank and confidence from filename
        rank_match = re.search(r"rank(\d+)", name)
        conf_match = re.search(r"confidence(-?[\d.]+)", name)
        if rank_match and conf_match:
            records.append({
                "rank":       int(rank_match.group(1)),
                "confidence": float(conf_match.group(1)),
                "sdf_file":   str(sdf_file),
            })

    df = pd.DataFrame(records).sort_values("rank")
    print(f"Found {len(df)} poses")
    print(df[["rank", "confidence", "sdf_file"]].head(10))
    return df

df_results = parse_diffdock_results("results/")

from rdkit import Chem
from rdkit.Chem import AllChem
from Bio.PDB import PDBParser
import numpy as np

def get_binding_residues(protein_pdb: str, ligand_sdf: str, cutoff_angstrom: float = 4.0):
    """Find protein residues within cutoff distance of the top-ranked ligand pose."""
    parser    = PDBParser(QUIET=True)
    structure = parser.get_structure("prot", protein_pdb)
    prot_atoms = [(atom.get_coord(), residue.resname, residue.id[1])
                  for chain in structure for residue in chain
                  for atom in residue.get_atoms()]

    mol = Chem.SDMolSupplier(ligand_sdf, removeHs=False)[0]
    lig_coords = mol.GetConformer().GetPositions()

    contacts = []
    for prot_coord, resname, resnum in prot_atoms:
        dists = np.linalg.norm(lig_coords - prot_coord, axis=1)
        if dists.min() <= cutoff_angstrom:
            contacts.append((resnum, resname))

    contacts = sorted(set(contacts))
    print(f"Binding site residues within {cutoff_angstrom} A: {contacts[:10]}")
    return contacts

# Use top-ranked pose
top_sdf = df_results.loc[df_results.rank == 1, "sdf_file"].iloc[0]
contacts = get_binding_residues("protein_clean.pdb", top_sdf)

import nglview as nv
from rdkit import Chem

# Load protein + top pose in Jupyter notebook
view = nv.NGLWidget()
view.add_pdbfile("protein_clean.pdb")

top_sdf  = df_results.loc[df_results.rank == 1, "sdf_file"].iloc[0]
view.add_component(top_sdf)
view.representations = [
    {"type": "cartoon", "params": {"color": "chainindex"}},
    {"type": "ball+stick", "params": {"sele": "ligand"}},
]
print(f"Visualizing top pose: confidence={df_results.confidence.iloc[0]:.3f}")
view

import pandas as pd
import subprocess

smiles_list = [
    ("compound_1", "CC(C)Cc1ccc(cc1)C(C)C(=O)O"),
    ("compound_2", "CC(C)Cc1ccc(cc1)C(C)C(=O)N"),
    ("compound_3", "CC(C)Cc1ccc(cc1)C(C)C(=O)OC"),
]

results = []
for name, smiles in smiles_list:
    out = f"results/{name}"
    cmd = ["python", "inference.py",
           "--protein_path", "protein_clean.pdb",
           "--ligand", smiles,
           "--out_dir", out,
           "--inference_steps", "20",
           "--samples_per_complex", "20"]
    subprocess.run(cmd, cwd="DiffDock/", capture_output=True)
    # Parse top confidence score
    df_r = parse_diffdock_results(out)
    if not df_r.empty:
        top_conf = df_r.loc[df_r.rank == 1, "confidence"].iloc[0]
        results.append({"name": name, "smiles": smiles, "top_confidence": top_conf})

df_batch = pd.DataFrame(results).sort_values("top_confidence", ascending=False)
df_batch.to_csv("batch_docking_results.csv", index=False)
print(df_batch)

# Confidence > 0 generally indicates a plausible binding pose
# DiffDock confidence scores: higher = more confident; ~0 is marginal; < -1 is poor
high_conf = df_results[df_results["confidence"] > 0.0]
print(f"High-confidence poses: {len(high_conf)} / {len(df_results)}")
print(high_conf[["rank", "confidence", "sdf_file"]])

# Convert SDF to PDBQT using OpenBabel
obabel rank1_confidence0.75.sdf -O rank1_ligand.pdbqt
obabel protein_clean.pdb -O protein.pdbqt -xr

# Rescore (no docking search, just energy evaluation)
vina --receptor protein.pdbqt --ligand rank1_ligand.pdbqt \
     --score_only --out rank1_rescored.pdbqt