{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "evaluation-00",
   "metadata": {},
   "source": [
    "# Evaluate ligand poses and protein interfaces\n",
    "\n",
    "Compare predictions with PLINDER references using OpenStructure, and check ligand plausibility with PoseBusters.\n",
    "\n",
    "Evaluate two Boltz predictions for the 8c3u protein–ligand complex from [Runs N' Poses](https://github.com/plinder-org/runs-n-poses) and a GeoDock prediction for the 2e31 protein interface from [PINDER's evaluation tutorial](https://github.com/pinder-org/pinder/blob/main/examples/pinder-eval.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-01",
   "metadata": {},
   "source": [
    "## Set up\n",
    "\n",
    "Install OpenStructure 2.12.0 or newer and the evaluation extra in your notebook environment:\n",
    "\n",
    "```bash\n",
    "conda install -c conda-forge -c bioconda 'openstructure>=2.12.0'\n",
    "pip install 'plinder[eval]'\n",
    "```\n",
    "\n",
    "The `ost` command must be on your `PATH`. Set `PLINDER_DATA_DIR` to use a local release; otherwise PLINDER uses the configured release. The [evaluation reference](../evaluation.md) describes the options and output files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "evaluation-02",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "import yaml\n",
    "import pandas as pd\n",
    "\n",
    "from plinder.core import PlinderRelease\n",
    "from plinder.eval import evaluate_predictions\n",
    "\n",
    "data_dir = os.environ.get(\"PLINDER_DATA_DIR\")\n",
    "release = PlinderRelease(data_dir=Path(data_dir)) if data_dir else PlinderRelease()\n",
    "work_dir = Path(\n",
    "    os.environ.get(\"PLINDER_EVAL_WORK_DIR\")\n",
    "    or tempfile.mkdtemp(prefix=\"plinder-evaluation-\")\n",
    ").resolve()\n",
    "predictions = work_dir / \"predictions\"\n",
    "work_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-03",
   "metadata": {},
   "source": [
    "## Choose prediction files\n",
    "\n",
    "Search and evaluation share an input table with `input_id`, `structure_path`, and optional `ligand_path`. Evaluation adds a `reference_id`: a ligand-system ID, an interface ID, or a PDB ID. The {ref}`evaluation reference <arrange-your-predictions>` describes the columns and folder alternative.\n",
    "\n",
    "```text\n",
    "predictions/\n",
    "├── <ligand-system ID>/\n",
    "│   ├── boltz_seed_1372115236.cif\n",
    "│   └── boltz_seed_555476360.cif\n",
    "└── <interface ID>/\n",
    "    └── geodock_model_1.pdb\n",
    "```\n",
    "\n",
    "For your own models, use the same folder layout and skip this example-data setup. Interface evaluation accepts PDB or mmCIF files. For ligand evaluation, pair a receptor `model.pdb` or `model.cif` with `model.sdf`, or a `model.ligands/` folder containing one SDF per ligand. Each prediction then carries its own ligand chemistry and coordinates.\n",
    "\n",
    "Complete-mmCIF predictions are also supported, as shown with the two Boltz models below. Their input YAML provides the ligand SMILES."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "evaluation-04",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_id = \"8c3u__1__1.A__1.C\"\n",
    "revision = \"78bfe19e15b086e39de4e03f30142e0dcb38beb0\"\n",
    "example_url = f\"https://raw.githubusercontent.com/plinder-org/runs-n-poses/{revision}/examples\"\n",
    "ligand_folder = predictions / system_id\n",
    "ligand_folder.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "for seed in (1372115236, 555476360):\n",
    "    model = ligand_folder / f\"boltz_seed_{seed}.cif\"\n",
    "    if not model.is_file():\n",
    "        urlretrieve(\n",
    "            f\"{example_url}/outputs/boltz/{system_id}/{seed}/\"\n",
    "            \"boltz_results_input/predictions/input/input_model_0.cif\",\n",
    "            model,\n",
    "        )\n",
    "\n",
    "input_yaml = work_dir / \"boltz_input.yaml\"\n",
    "if not input_yaml.is_file():\n",
    "    urlretrieve(f\"{example_url}/inputs/boltz/{system_id}/input.yaml\", input_yaml)\n",
    "boltz_input = yaml.safe_load(input_yaml.read_text())\n",
    "ligand_smiles = next(\n",
    "    item[\"ligand\"][\"smiles\"] for item in boltz_input[\"sequences\"] if \"ligand\" in item\n",
    ")\n",
    "\n",
    "interface_id = \"2e31__1__1.A--1.B\"\n",
    "interface_folder = predictions / interface_id\n",
    "interface_folder.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "pinder_revision = \"20487ac403516f453d8be1b0bdc07a9f7a77d9df\"\n",
    "pinder_id = \"2e31__A1_Q80UW2--2e31__B1_P63208\"\n",
    "geodock_pdb = interface_folder / \"geodock_model_1.pdb\"\n",
    "if not geodock_pdb.is_file():\n",
    "    urlretrieve(\n",
    "        f\"https://raw.githubusercontent.com/pinder-org/pinder/{pinder_revision}/\"\n",
    "        f\"tests/test_data/method_eval/geodock/{pinder_id}/holo_decoys/model_1.pdb\",\n",
    "        geodock_pdb,\n",
    "    )\n",
    "\n",
    "input_table = pd.DataFrame([\n",
    "    {\n",
    "        \"input_id\": path.stem,\n",
    "        \"structure_path\": str(path.resolve()),\n",
    "        \"reference_id\": path.parent.name,\n",
    "    }\n",
    "    for path in sorted(predictions.glob(\"*/*\")) if path.is_file()\n",
    "])\n",
    "input_table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-05",
   "metadata": {},
   "source": [
    "## Evaluate the predictions\n",
    "\n",
    "One call evaluates both kinds of prediction. Choose `mode=\"ligands\"` or `mode=\"interfaces\"` for just one kind. `num_workers` limits the number of predictions processed at once.\n",
    "\n",
    "By default, ligand results cover proper **reference** ligands. Set `include_all_ligands=True` to include reference ions and artifacts too.\n",
    "\n",
    "The two Boltz predictions share the same ligand, so pass its SMILES as `ligand_smiles={\"LIG\": ligand_smiles}`. A SMILES must follow the CIF heavy-atom order. With complete-mmCIF inputs, `ligand_ccd_codes` supplies CCD templates and `ligand_chains=[\"B\"]` selects a peptide ligand by its label asym ID. For receptor + SDF inputs, the SDF files supply the chemistry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "evaluation-06",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = evaluate_predictions(\n",
    "    input_table,\n",
    "    output_dir=work_dir / \"results\",\n",
    "    release=release,\n",
    "    mode=\"both\",\n",
    "    num_workers=2,\n",
    "    ligand_smiles={\"LIG\": ligand_smiles},\n",
    ")\n",
    "\n",
    "display(results[\"failures\"])\n",
    "assert results[\"failures\"].empty, \"Inspect the errors and tool logs before interpreting results.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-07",
   "metadata": {},
   "source": [
    "## Read ligand results\n",
    "\n",
    "Each row represents one reference ligand and one prediction. `rmsd` is the binding-site-superposed, symmetry-corrected ligand RMSD; `bb_rmsd` describes the backbone fit, `lddt_lp` the local binding-site structure, and `lddt_pli` the protein–ligand contacts. Lower RMSD is better; lDDT values range from zero to one.\n",
    "\n",
    "`status=\"success\"` means the comparison completed. Both Boltz examples here have a large ligand RMSD. A missing or unassigned reference ligand has a row with unavailable metrics and an unassigned reason."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "evaluation-08",
   "metadata": {},
   "outputs": [],
   "source": [
    "ligands = results[\"ligands\"]\n",
    "ligands[\n",
    "    [\n",
    "        \"prediction\", \"ligand_id\", \"status\",\n",
    "        \"rmsd\", \"bb_rmsd\", \"lddt_lp\", \"lddt_pli\",\n",
    "        \"rmsd_unassigned\", \"lddt_pli_unassigned\",\n",
    "    ]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-09",
   "metadata": {},
   "source": [
    "## Attach PoseBusters checks\n",
    "\n",
    "PoseBusters checks all candidate predicted ligands. To attach those checks to a reference-ligand result, use the predicted ligand selected by the metric you are analysing. RMSD and lDDT-PLI can choose different assignments.\n",
    "\n",
    "Below, the join follows `rmsd_model_ligand`. For lDDT-PLI, use `lddt_pli_model_ligand` instead. PoseBusters uses each prediction's own receptor coordinates. Failed checks appear as `False`; PoseBusters fields for missing ligands are empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "evaluation-10",
   "metadata": {},
   "outputs": [],
   "source": [
    "checks = results[\"posebusters\"][\n",
    "    [\n",
    "        \"prediction\", \"model_ligand\", \"sanitization\",\n",
    "        \"internal_steric_clash\", \"minimum_distance_to_protein\",\n",
    "    ]\n",
    "]\n",
    "ligands_with_checks = ligands.merge(\n",
    "    checks,\n",
    "    left_on=[\"prediction\", \"rmsd_model_ligand\"],\n",
    "    right_on=[\"prediction\", \"model_ligand\"],\n",
    "    how=\"left\",\n",
    "    validate=\"many_to_one\",\n",
    ")\n",
    "ligands_with_checks[\n",
    "    [\n",
    "        \"prediction\", \"ligand_id\", \"rmsd_model_ligand\", \"rmsd\",\n",
    "        \"sanitization\", \"internal_steric_clash\", \"minimum_distance_to_protein\",\n",
    "    ]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-11",
   "metadata": {},
   "source": [
    "## Read interface results\n",
    "\n",
    "OST maps the GeoDock model's chains `R` and `L` to reference chains `1.A` and `1.B`.\n",
    "\n",
    "`ilddt` describes local agreement at the interface; QS-score describes assembly agreement. Use `dockq_ave_full` or `dockq_wave_full` when comparing predictions that may have missing chains: these include the penalty for unmapped reference interfaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "evaluation-12",
   "metadata": {},
   "outputs": [],
   "source": [
    "results[\"interfaces\"][\n",
    "    [\n",
    "        \"prediction\", \"system_id\", \"status\", \"lddt\", \"ilddt\",\n",
    "        \"qs_global\", \"qs_best\", \"dockq_ave_full\", \"dockq_wave_full\",\n",
    "    ]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-13",
   "metadata": {},
   "source": [
    "Access reference-ligand assignments in `results[\"ligands\"]` using `rmsd_model_ligand` and `lddt_pli_model_ligand`, with reasons for unmatched ligands in `rmsd_unassigned` and `lddt_pli_unassigned`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "evaluation-15",
   "metadata": {},
   "source": [
    "Reload saved tables with pandas, for example `pd.read_parquet(work_dir / \"results\" / \"ligands.parquet\")`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}