CoolFace
Apppublic

Tsimech2000/virulent_molecular_structures

sourceHugging Facemitupdated 2mo agoView on Hugging Face
0likes
docking.py232 linesDownload Raw Back to root
1"""Validated and reproducible AutoDock Vina workflow."""2 3from __future__ import annotations4 5import hashlib6import importlib.metadata7import io8import json9import platform10import tempfile11import uuid12import zipfile13from dataclasses import asdict, dataclass14from datetime import datetime, timezone15from pathlib import Path16 17import pandas as pd18from openbabel import openbabel, pybel19from vina import Vina20 21 22class DockingError(RuntimeError):23    """Raised when preparation or docking cannot be completed safely."""24 25 26@dataclass(frozen=True)27class DockingConfig:28    center: tuple[float, float, float]29    size: tuple[float, float, float]30    exhaustiveness: int = 1631    num_modes: int = 932    seed: int = 2026080133    energy_range: float = 3.034    retain_receptor_heteroatoms: bool = False35 36    def __post_init__(self) -> None:37        if any(not 5 <= float(value) <= 60 for value in self.size):38            raise DockingError("Every search-box dimension must be between 5 and 60 Å.")39        if not 1 <= self.num_modes <= 20:40            raise DockingError("The number of modes must be between 1 and 20.")41        if self.exhaustiveness < 8:42            raise DockingError("Exhaustiveness must be at least 8.")43 44 45@dataclass46class DockingResult:47    run_id: str48    scores: pd.DataFrame49    poses: list[str]50    receptor_pdb: str51    bundle: bytes52 53 54def _sha256(data: bytes) -> str:55    return hashlib.sha256(data).hexdigest()56 57 58def _first_molecule(path: Path, file_format: str):59    try:60        return next(pybel.readfile(file_format, str(path)))61    except (StopIteration, OSError, RuntimeError) as exc:62        raise DockingError(f"Open Babel could not read {path.name}: {exc}") from exc63 64 65def _prepare_receptor(input_path: Path, cleaned_path: Path, output_path: Path, retain_heteroatoms: bool) -> None:66    lines = input_path.read_text(errors="replace").splitlines()67    cleaned_lines = []68    in_first_model = True69    for line in lines:70        record = line[:6].strip().upper()71        if record == "MODEL" and cleaned_lines:72            in_first_model = False73        if record == "ENDMDL":74            in_first_model = False75        if not in_first_model:76            continue77        if record == "ATOM" or (retain_heteroatoms and record == "HETATM") or record == "TER":78            cleaned_lines.append(line)79    cleaned_lines.append("END")80    cleaned_path.write_text("\n".join(cleaned_lines) + "\n")81    mol = _first_molecule(cleaned_path, "pdb")82    if not mol.atoms:83        raise DockingError("The receptor contains no atoms.")84    mol.addh()85    mol.write("pdbqt", str(output_path), overwrite=True)86    # Open Babel writes torsion-tree directives for all PDBQT molecules. A Vina87    # receptor must be rigid and may contain only atom records and remarks.88    rigid_lines = [89        line90        for line in output_path.read_text(errors="replace").splitlines()91        if line.startswith(("ATOM", "HETATM", "REMARK"))92    ]93    output_path.write_text("\n".join(rigid_lines) + "\n")94 95 96def _prepare_ligand(input_path: Path, output_path: Path, file_format: str) -> None:97    mol = _first_molecule(input_path, file_format)98    if not mol.atoms:99        raise DockingError("The ligand contains no atoms.")100    mol.OBMol.AddHydrogens()101    charge_model = openbabel.OBChargeModel.FindType("gasteiger")102    if charge_model is None or not charge_model.ComputeCharges(mol.OBMol):103        raise DockingError("Gasteiger charges could not be assigned to the ligand.")104    mol.write("pdbqt", str(output_path), overwrite=True)105 106 107def _split_and_convert_poses(pdbqt_path: Path) -> list[str]:108    text = pdbqt_path.read_text(errors="replace")109    blocks = text.split("MODEL")110    if len(blocks) == 1:111        blocks = ["", " 1\n" + text]112    poses: list[str] = []113    for block in blocks[1:]:114        pdbqt = "MODEL" + block115        with tempfile.NamedTemporaryFile(suffix=".pdbqt") as handle:116            handle.write(pdbqt.encode())117            handle.flush()118            mol = _first_molecule(Path(handle.name), "pdbqt")119            poses.append(mol.write("pdb"))120    return poses121 122 123def _make_bundle(files: dict[str, bytes | str]) -> bytes:124    buffer = io.BytesIO()125    with zipfile.ZipFile(buffer, "w", compression=zipfile.ZIP_DEFLATED) as archive:126        for name, value in files.items():127            archive.writestr(name, value)128    return buffer.getvalue()129 130 131def run_docking_workflow(132    receptor_bytes: bytes,133    ligand_bytes: bytes,134    receptor_name: str,135    ligand_name: str,136    config: DockingConfig,137) -> DockingResult:138    run_id = uuid.uuid4().hex[:12]139    ligand_format = Path(ligand_name).suffix.lower().lstrip(".")140    if ligand_format not in {"sdf", "mol"}:141        raise DockingError("Ligand format must be SDF or MOL.")142 143    with tempfile.TemporaryDirectory(prefix=f"dock_{run_id}_") as directory:144        work = Path(directory)145        receptor_path = work / "receptor_original.pdb"146        ligand_path = work / f"ligand_original.{ligand_format}"147        receptor_pdbqt = work / "receptor_prepared.pdbqt"148        receptor_cleaned = work / "receptor_cleaned.pdb"149        ligand_pdbqt = work / "ligand_prepared.pdbqt"150        output_pdbqt = work / "docked_poses.pdbqt"151        receptor_path.write_bytes(receptor_bytes)152        ligand_path.write_bytes(ligand_bytes)153 154        try:155            _prepare_receptor(156                receptor_path,157                receptor_cleaned,158                receptor_pdbqt,159                config.retain_receptor_heteroatoms,160            )161            _prepare_ligand(ligand_path, ligand_pdbqt, ligand_format)162            vina = Vina(sf_name="vina", seed=config.seed, verbosity=1)163            vina.set_receptor(str(receptor_pdbqt))164            vina.set_ligand_from_file(str(ligand_pdbqt))165            vina.compute_vina_maps(center=list(config.center), box_size=list(config.size))166            vina.dock(exhaustiveness=config.exhaustiveness, n_poses=config.num_modes)167            vina.write_poses(str(output_pdbqt), n_poses=config.num_modes, energy_range=config.energy_range, overwrite=True)168            energies = vina.energies(n_poses=config.num_modes, energy_range=config.energy_range)169        except Exception as exc:170            raise DockingError(f"Docking failed: {exc}") from exc171 172        if not output_pdbqt.exists() or output_pdbqt.stat().st_size == 0:173            raise DockingError("Vina completed without producing poses.")174 175        scores = pd.DataFrame(176            [177                {178                    "pose": index,179                    "vina_score_kcal_mol": round(float(row[0]), 3),180                    "intermolecular_kcal_mol": round(float(row[1]), 3),181                    "intramolecular_kcal_mol": round(float(row[2]), 3),182                    "torsional_kcal_mol": round(float(row[3]), 3),183                    "unbound_kcal_mol": round(float(row[4]), 3),184                }185                for index, row in enumerate(energies, start=1)186            ]187        )188        poses = _split_and_convert_poses(output_pdbqt)189        if not poses:190            raise DockingError("Docked poses could not be converted for visualization.")191 192        metadata = {193            "run_id": run_id,194            "created_utc": datetime.now(timezone.utc).isoformat(),195            "workflow": "AutoDock Vina protein-ligand docking",196            "config": asdict(config),197            "inputs": {198                "receptor_filename": Path(receptor_name).name,199                "receptor_sha256": _sha256(receptor_bytes),200                "ligand_filename": Path(ligand_name).name,201                "ligand_sha256": _sha256(ligand_bytes),202            },203            "environment": {204                "python": platform.python_version(),205                "vina": importlib.metadata.version("vina"),206                "openbabel": openbabel.OBReleaseVersion(),207                "biopython": importlib.metadata.version("biopython"),208                "pandas": importlib.metadata.version("pandas"),209            },210            "interpretation_warning": "Scores rank modeled poses and are not experimental binding free energies.",211        }212        readme = (213            "REPRODUCIBLE DOCKING RUN\n\n"214            "Inspect all preparation choices and the search box before interpreting results. "215            "Computational docking generates hypotheses and requires controls and experimental validation.\n"216        )217        bundle = _make_bundle(218            {219                "README.txt": readme,220                "metadata.json": json.dumps(metadata, indent=2),221                "scores.csv": scores.to_csv(index=False),222                "inputs/receptor_original.pdb": receptor_bytes,223                f"inputs/ligand_original.{ligand_format}": ligand_bytes,224                "prepared/receptor.pdbqt": receptor_pdbqt.read_bytes(),225                "prepared/receptor_cleaned.pdb": receptor_cleaned.read_bytes(),226                "prepared/ligand.pdbqt": ligand_pdbqt.read_bytes(),227                "results/docked_poses.pdbqt": output_pdbqt.read_bytes(),228                **{f"results/pose_{i:02d}.pdb": pose for i, pose in enumerate(poses, 1)},229            }230        )231        return DockingResult(run_id, scores, poses, receptor_bytes.decode(errors="replace"), bundle)232