{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The PLINDEX\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import mols2grid\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying and filtering the index\n",
    "\n",
    "Your main entry point to the dataset is the annotations table or the **index**, which is a Parquet file containing all the annotations for each system in the dataset.\n",
    "\n",
    ":::{note} The entire file has 745 columns, >1.3M rows, and takes 24G of RAM to load into memory :::\n",
    "\n",
    "We provide a `query_index` function to access the index and filter it based on the columns and criteria you need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plinder.core.scores import query_index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The core of the PLINDER dataset is a collection of Protein-Ligand Interaction (PLI) systems extracted from the Protein Data Bank (PDB). The curation process in a nutshell is as follows:\n",
    "\n",
    "1. For every PDB entry, we generate all available biological assemblies (biounits).\n",
    "2. For each biounit, we identify all ligands and all protein chains within 6 Å of any ligand.\n",
    "3. Ligands within 4 Å of each other are merged into a single PLI system.\n",
    "4. For each system, we provide a range of detailed annotations and files to facilitate in-depth analysis and to enable a variety of use-cases.\n",
    "\n",
    "Thus, a PLI **_system_** is defined as a collection of **_protein_** and **_ligand_** chains that are in close proximity to each other within a given **_biounit_** of a given PDB **_entry_**. The **_pocket_** of a system is defined as the set of protein residues within 6 Å of the ligands in the system.\n",
    "\n",
    "Let's look at all the columns that define a `system`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\n",
    "        \"system_id\",\n",
    "        \"entry_pdb_id\",\n",
    "        \"system_biounit_id\",\n",
    "        \"system_protein_chains_asym_id\",\n",
    "        \"system_ligand_chains_asym_id\",\n",
    "        \"ligand_instance_chain\",\n",
    "    ]\n",
    ")\n",
    "plindex.head(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, a system is uniquely qualified by its `system_id` which is a combination of\n",
    "- `entry_pdb_id` - the PDB identifier\n",
    "- `system_biounit_id` - the biological assembly identifier\n",
    "- `system_protein_chains_asym_id` - The list of protein chains in the system defined by their `<instance>.<label_asym_id>`\n",
    "- `system_ligand_chains_asym_id` - The list of ligand chains in the system defined by their `<instance>.<label_asym_id>`\n",
    "\n",
    "Systems containing multiple ligands (e.g `4grc__1__1.A__1.C_1.D`) span multiple rows, where each row represents a different ligand in the system (as seen in the `ligand_instance_chain` column)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Number of ligand chains: {plindex.shape[0]}\")\n",
    "print(f\"Number of systems: {plindex.system_id.nunique()}\")\n",
    "print(\n",
    "    f\"Number of biounits: {plindex[['entry_pdb_id', 'system_biounit_id']].drop_duplicates().shape[0]}\"\n",
    ")\n",
    "print(f\"Number of PDB IDs: {plindex.entry_pdb_id.nunique()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While `query_index` by default loads all the systems in the `train` and `val` splits, not every returned system may be useful for training your model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single-ligand single-protein predictors\n",
    "\n",
    "As an example, we consider the case of training single-ligand single-protein models. One way to filter training data is as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex_single = query_index(\n",
    "    filters=[\n",
    "        (\"system_num_ligand_chains\", \"==\", 1),\n",
    "        (\"system_num_protein_chains\", \"==\", 1),\n",
    "    ],\n",
    "    splits=[\"train\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex_single.system_id.nunique()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, as PLINDER also considers ions and crystallization artifacts as ligands if they are within 4 Å of a non-ion non-artifact ligand, there are also systems in PLINDER which only have one \"proper\" ligand. So, another strategy would be to only train on the proper ligands and ignore the ions and artifacts in the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex_single_proper = query_index(\n",
    "    filters=[\n",
    "        (\"system_proper_num_ligand_chains\", \"==\", 1),\n",
    "        (\"system_num_protein_chains\", \"==\", 1),\n",
    "        (\"ligand_is_proper\", \"==\", True),  # filters out all other ligands in the system\n",
    "    ],\n",
    "    splits=[\"train\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex_single_proper.system_id.nunique()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can provide up to 20% more data for training, however, the caveat is that some of the interactions made by artifacts or ions may influence the binding pose of the \"proper\" ligand.\n",
    "\n",
    "One could also choose to include multi-ligand systems but only train with one ligand at a time, and the same for multi-protein. These choices are up to the user and we provide the annotations to enable such choices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Annotations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are 745 columns in the full index. Of course, not every one will be useful for your use-case so we'll go through some common categories of them and some  use-cases requiring different annotations. You can get the full list of columns with descriptions in the [index docs](https://plinder-org.github.io/plinder/dataset.html#annotation-tables-index)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Structure quality\n",
    "\n",
    "A core principle of PLINDER is to be able to annotate which systems are of high enough experimental structure quality to be reliably used as the _ground truth_ for measuring model performance. As the quality of experimentally resolved structures [can vary significantly](https://doi.org/10.1107/S2059798322011901) and many crystal structures with ligands contain missing atoms or missing residues in the binding site, comparing prediction results to lower quality structures can incorrectly skew the perception of their performance.\n",
    "\n",
    "![quality](https://www.plinder.sh/blog/figure2.png)\n",
    "\n",
    "\n",
    "You can access all the crystal structure validation information extracted from PDB validation reports as well as crystal contact annotations by looking for columns starting with \"entry_validation\", \"system_pocket_validation\", \"system_ligand_validation\" etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simplicity, we have \"system_pass_validation_criteria\" as a column that can be used to filter systems which pass our quality definitions:\n",
    "\n",
    "```python\n",
    "class QualityCriteria:\n",
    "    max_entry_resolution: float = 3.5\n",
    "    max_entry_r: float = 0.4\n",
    "    max_entry_rfree: float = 0.45\n",
    "    max_entry_r_minus_rfree: float = 0.05\n",
    "    ligand_max_num_unresolved_heavy_atoms: int = 0  # except for covalent\n",
    "    ligand_max_alt_count: int = 1  # misnomer: this counts number of total configurations\n",
    "    ligand_min_average_occupancy: float = 0.8\n",
    "    ligand_min_average_rscc: float = 0.8\n",
    "    ligand_max_average_rsr: float = 0.3\n",
    "    ligand_max_percent_outliers_clashes: float = 0\n",
    "    ligand_max_fraction_atoms_with_crystal_contacts: float = 0\n",
    "    pocket_max_num_unresolved_heavy_atoms: int = 0\n",
    "    pocket_max_alt_count: int = 1  # same as above\n",
    "    pocket_min_average_occupancy: float = 0.8\n",
    "    pocket_min_average_rscc: float = 0.8\n",
    "    pocket_max_average_rsr: float = 0.3\n",
    "    pocket_max_percent_outliers_clashes: int = 100\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\"system_id\", \"system_pass_validation_criteria\"],\n",
    "    splits=[\"train\", \"val\", \"test\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = (\n",
    "    plindex.drop_duplicates(\"system_id\")\n",
    "    .sort_values(by=\"system_pass_validation_criteria\")\n",
    "    .groupby([\"split\", \"system_pass_validation_criteria\"])\n",
    "    .system_id.count()\n",
    "    .unstack()\n",
    ")\n",
    "data_percentage = data.div(data.sum(axis=1), axis=0) * 100\n",
    "ax = data_percentage.plot(\n",
    "    kind=\"bar\", stacked=True, figsize=(5, 3), color=[\"firebrick\", \"darkseagreen\"]\n",
    ")\n",
    "ax.set_xticklabels(\n",
    "    [label.get_text().upper() for label in ax.get_xticklabels()], rotation=0\n",
    ")\n",
    "ax.set_xlabel(\"\")\n",
    "ax.set_ylabel(\"% passing criteria\")\n",
    "ax.get_legend().remove()\n",
    "for container, count_data in zip(ax.containers, data.values.T):\n",
    "    ax.bar_label(container, labels=count_data, label_type=\"edge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Structure completeness\n",
    "\n",
    "Related to structure quality, one aspect that is important to consider when using these structures in your training is **completeness**. While the inputs for prediction may be a protein sequence and a SMILES string, the protein-ligand complex structure that you get from the PDB may not have all the residues in the protein or all the atoms in the ligand resolved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\n",
    "        \"system_id\",\n",
    "        \"ligand_num_unresolved_heavy_atoms\",\n",
    "        \"ligand_num_heavy_atoms\",\n",
    "        \"system_protein_chains_num_unresolved_residues\",\n",
    "        \"system_pocket_validation_num_unresolved_heavy_atoms\",\n",
    "        \"system_protein_chains_total_length\",\n",
    "    ],\n",
    "    splits=[\"train\", \"val\", \"test\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we calculate the fraction of unresolved residues in the protein chains and the fraction of unresolved heavy atoms in the ligand, to see their distributions across the splits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[\n",
    "    \"system_protein_chains_total_num_unresolved_residues\"\n",
    "] = plindex.system_protein_chains_num_unresolved_residues.map(sum)\n",
    "plindex[\"system_protein_chains_fraction_unresolved_residues\"] = (\n",
    "    plindex.system_protein_chains_total_num_unresolved_residues\n",
    "    / plindex.system_protein_chains_total_length\n",
    ")\n",
    "plindex[\"ligand_fraction_unresolved_heavy_atoms\"] = (\n",
    "    plindex.ligand_num_unresolved_heavy_atoms / plindex.ligand_num_heavy_atoms\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(10, 3), sharex=True, sharey=True)\n",
    "grouped_data = plindex.drop_duplicates(\"system_id\").groupby(\"split\")\n",
    "split_colors = {\n",
    "    \"train\": \"#ff9999\",\n",
    "    \"test\": \"#66b3ff\",\n",
    "    \"val\": \"#7cc27c\",\n",
    "}\n",
    "for i, (split, data) in enumerate(grouped_data):\n",
    "    data.system_protein_chains_fraction_unresolved_residues.hist(\n",
    "        ax=axes[i], density=True, color=split_colors[split]\n",
    "    )\n",
    "    axes[i].set_title(split.upper())\n",
    "    if i == 0:\n",
    "        axes[i].set_ylabel(\"Density\")\n",
    "fig.suptitle(\n",
    "    \"Fraction of Unresolved Protein Residues (SEQRES length - no. resolved residues)\"\n",
    ")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Percentage of ligands with no unresolved heavy atoms:\")\n",
    "for split in [\"train\", \"val\", \"test\"]:\n",
    "    print(\n",
    "        f'{split.capitalize()}: {100 * (plindex[plindex[\"split\"] == split].ligand_fraction_unresolved_heavy_atoms == 0).sum() / plindex[plindex[\"split\"] == split].shape[0]:.2f}%'\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pocket domains\n",
    "\n",
    "We annotated domains from different databases onto the protein chains of each system and then picked the one spanning the pocket residues of the system as the domain of the system pocket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\"system_id\", \"system_pocket_ECOD_t_name\", \"system_pocket_kinase_name\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex.drop_duplicates(\"system_id\").system_pocket_ECOD_t_name.value_counts().head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex.drop_duplicates(\"system_id\").system_pocket_kinase_name.value_counts().head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ligand properties\n",
    "\n",
    "Molecular properties and annotations are calculated from the ligand SMILES strings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "properties = [\n",
    "    \"ligand_molecular_weight\",\n",
    "    \"ligand_num_rot_bonds\",\n",
    "    \"ligand_num_rings\",\n",
    "    \"ligand_num_hbd\",\n",
    "    \"ligand_num_hba\",\n",
    "    \"ligand_num_heavy_atoms\",\n",
    "    \"ligand_crippen_clogp\",\n",
    "    \"ligand_qed\",\n",
    "    \"ligand_tpsa\",\n",
    "    \"ligand_is_kinase_inhibitor\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\"system_id\", \"ligand_instance_chain\", \"ligand_unique_ccd_code\"]\n",
    "    + properties,\n",
    "    splits=[\"train\", \"val\", \"test\"],\n",
    "    filters=[\n",
    "        (\"ligand_is_proper\", \"==\", True)  # focusing on non-ion, non-artifact ligands\n",
    "    ],\n",
    ")\n",
    "plindex.head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binwidths = {\n",
    "    \"ligand_molecular_weight\": (\"Molecular weight\", 50),\n",
    "    \"ligand_num_heavy_atoms\": (\"Number of heavy atoms\", 5),\n",
    "    \"ligand_num_rot_bonds\": (\"Number of rotatable bonds\", 2),\n",
    "    \"ligand_num_rings\": (\"Number of rings\", 1),\n",
    "    \"ligand_tpsa\": (\"Topological polar surface area\", 10),\n",
    "    \"ligand_crippen_clogp\": (\"Crippen logP\", 1),\n",
    "}\n",
    "fig, axes = plt.subplots(3, 2, figsize=(10, 7))\n",
    "axes = axes.flatten()\n",
    "for i, prop in enumerate(binwidths):\n",
    "    ax = axes[i]\n",
    "    plindex.groupby(\"split\")[prop].plot(\n",
    "        kind=\"hist\",\n",
    "        density=True,\n",
    "        alpha=0.5,\n",
    "        histtype=\"stepfilled\",\n",
    "        legend=True,\n",
    "        ax=ax,\n",
    "        color=split_colors,\n",
    "        bins=np.arange(\n",
    "            min(plindex[prop]),\n",
    "            max(plindex[prop]) + binwidths[prop][1],\n",
    "            binwidths[prop][1],\n",
    "        ),\n",
    "    )\n",
    "    ax.set_xlabel(binwidths[prop][0])\n",
    "    ax.set_ylabel(\"Density\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These properties were further used to categorize the ligands into different types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ligand_types = [\n",
    "    f\"ligand_is_{x}\"\n",
    "    for x in [\n",
    "        \"lipinski\",\n",
    "        \"covalent\",\n",
    "        \"cofactor\",\n",
    "        \"oligo\",\n",
    "        \"ion\",\n",
    "        \"fragment\",\n",
    "        \"artifact\",\n",
    "    ]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These categories were also aggregated to the system level for easier filtering, with `system_has_<category>` as a column indicating whether any ligand in the system is of that category."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\"system_id\", \"ligand_unique_ccd_code\"] + ligand_types,\n",
    "    splits=[\"train\", \"val\", \"test\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = [c.replace(\"ligand_is_\", \"\").capitalize() for c in ligand_types]\n",
    "bar_colors = plt.cm.Pastel2.colors\n",
    "split_names = [\"train\", \"val\", \"test\"]\n",
    "fig, axes = plt.subplots(1, 3, figsize=(10, 3))\n",
    "for i, split in enumerate(split_names):\n",
    "    ax = axes[i]\n",
    "    split_data = plindex[plindex[\"split\"] == split]\n",
    "    bars = ax.bar(\n",
    "        labels,\n",
    "        split_data[ligand_types].mean().mul(100).to_list(),\n",
    "        width=1,\n",
    "        color=bar_colors,\n",
    "        edgecolor=\"black\",\n",
    "        label=split,\n",
    "        linewidth=1,\n",
    "    )\n",
    "    ax.set_xticks(np.arange(len(labels)), labels, rotation=70)\n",
    "    ax.set_ylim(0, 100)\n",
    "    ax.set_title(split.upper())\n",
    "    counts = split_data[ligand_types].sum().to_list()\n",
    "    for bar, count in zip(bars, counts):\n",
    "        ax.text(\n",
    "            bar.get_x() + bar.get_width() / 2.0,\n",
    "            bar.get_height() + 2,\n",
    "            f\"{count}\",\n",
    "            ha=\"center\",\n",
    "            va=\"bottom\",\n",
    "            rotation=70,\n",
    "            fontsize=10,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PLI-specific properties\n",
    "\n",
    "There are also some ligand properties that are specific to the interaction with the particular protein pocket present in the system. These include things like:\n",
    "- The kinds of interactions the ligand may have with the protein, calculated using PLIP\n",
    "- Experimental binding affinity, pulled from BindingDB when available\n",
    "- Potential ligand-protein crystal contacts, defined as ligand-protein contacts below 5 Å which are not in the same asymmetric unit (symmetry mates) and not in the system biounit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pli_specific = [\n",
    "    \"system_fraction_atoms_with_crystal_contacts\",\n",
    "    \"system_num_crystal_contacted_residues\",\n",
    "    \"ligand_binding_affinity\",\n",
    "    \"system_has_binding_affinity\",\n",
    "    \"ligand_interactions\",\n",
    "    \"system_num_interactions\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\"system_id\"] + pli_specific,\n",
    "    splits=[\"train\", \"val\", \"test\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could filter out systems with crystal contacts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    plindex.drop_duplicates(\"system_id\").system_fraction_atoms_with_crystal_contacts > 0\n",
    ").sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your model has an additional component for predicting binding affinity, you could see how much data is available for training and evaluation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex.drop_duplicates(\"system_id\").groupby(\n",
    "    \"split\"\n",
    ").system_has_binding_affinity.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[\"ligand_binding_affinity\"].hist()\n",
    "plt.xlabel(r\"$pK_i$ or $pK_D$\")\n",
    "plt.ylabel(\"No. of system ligands\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could even look deeper into the kinds of protein-ligand interactions present in your training data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex.groupby(\"split\").system_num_interactions.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[\"ligand_interactions\"].values[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Beyond the split\n",
    "\n",
    "Here we mainly focused on systems in the train/val/test splits as the train and val splits are the ones that can be used for training models to either participate in the PLINDER leaderboard or compare your models' performance to methods in the leaderboard.\n",
    "\n",
    "However, PLINDER itself contains all PLI systems in the PDB (except those containing only crystallization artifacts), and you may be interested in this dataset as a whole. This can be queried with the `splits=[\"*\"]` option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex = query_index(\n",
    "    columns=[\n",
    "        \"system_id\",\n",
    "        \"entry_pdb_id\",\n",
    "        \"system_num_protein_chains\",\n",
    "        \"system_num_ligand_chains\",\n",
    "    ]\n",
    "    + ligand_types,\n",
    "    splits=[\"*\"],\n",
    ")\n",
    "print(\"No. ligands:\", len(plindex))\n",
    "print(\"No. systems:\", plindex.system_id.nunique())\n",
    "print(\"No. PDB entries:\", plindex.entry_pdb_id.nunique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(5, 3))\n",
    "bars = ax.bar(\n",
    "    labels,\n",
    "    plindex[ligand_types].mean().mul(100).to_list(),\n",
    "    width=1,\n",
    "    color=bar_colors,\n",
    "    edgecolor=\"black\",\n",
    "    label=split,\n",
    "    linewidth=1,\n",
    ")\n",
    "ax.set_xticks(np.arange(len(labels)), labels, rotation=70)\n",
    "ax.set_ylim(0, 100)\n",
    "ax.set_title(f\"No. PLINDER ligands: {len(plindex)}\")\n",
    "counts = plindex[ligand_types].sum().to_list()\n",
    "for bar, count in zip(bars, counts):\n",
    "    ax.text(\n",
    "        bar.get_x() + bar.get_width() / 2.0,\n",
    "        bar.get_height() + 2,\n",
    "        f\"{count}\",\n",
    "        ha=\"center\",\n",
    "        va=\"bottom\",\n",
    "        rotation=70,\n",
    "        fontsize=10,\n",
    "    )\n",
    "\n",
    "print(plindex.system_num_protein_chains.value_counts())\n",
    "print(plindex.system_num_ligand_chains.value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When looking at distributions across the entire dataset, you'll notice that there are quite a lot of ion systems, as well as systems with more than 5 protein or ligand chains.\n",
    "\n",
    ":::{note} While we provide annotations for all systems, those which are not in the train and val splits may not be used for training models to participate in the MLSB PLINDER leaderboard. :::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Similarity clusters\n",
    "\n",
    "Similarity between two protein-ligand complexes can occur at various levels, including protein sequence, structural features, binding pocket characteristics, or ligand and interaction properties. We calculated a comprehensive set of similarity metrics to cover every combination of these dimensions — from identical systems, where the protein, pocket, interactions and ligand are the same, to systems that differ across all levels.\n",
    "\n",
    "The entire set of similarity metrics is described [here](https://plinder-org.github.io/plinder/dataset.html#clusters-clusters), and they were calculated across all pairs of systems having a Foldseek or MMseqs alignment. We then used graph clustering to group systems into clusters based on their similarity with a specific **_metric_** and **_threshold_**. There are three types of clusters available: strongly connected graph components (**_strong component_**), weakly connected graph components (**_weak component_**), and communities detected using asynchronous label propagation (**_community_**).\n",
    "\n",
    "We can load the similarity clusters assigned to each system for different metrics, thresholds, and clustering types."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These cluster labels can be used in many ways with the annotations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: ATP-binding pockets\n",
    "\n",
    "For example, maybe you are interested in seeing how many different kinds of ATP-binding pockets we have. First, let's find all the analogs of ATP. We can do this by finding the 95% ECFP4 tanimoto similarity component of ATP and then getting the CCD codes of all the ligands in the same cluster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ligand_cluster_column = \"tanimoto_similarity_max__95__strong__component\"\n",
    "atp_ligand_cluster = query_index(\n",
    "    columns=[\n",
    "        ligand_cluster_column,\n",
    "    ],\n",
    "    filters=[\n",
    "        (\"ligand_unique_ccd_code\", \"==\", \"ATP\"),\n",
    "        (\"system_num_ligand_chains\", \"==\", 1),\n",
    "    ],\n",
    ")[ligand_cluster_column].values[0]\n",
    "\n",
    "atp_analogs = query_index(\n",
    "    columns=[\n",
    "        \"system_id\",\n",
    "        \"entry_pdb_id\",\n",
    "        \"ligand_unique_ccd_code\",\n",
    "        \"ligand_rdkit_canonical_smiles\",\n",
    "    ],\n",
    "    filters=[\n",
    "        (ligand_cluster_column, \"==\", atp_ligand_cluster),\n",
    "        (\"system_num_ligand_chains\", \"==\", 1),\n",
    "    ],\n",
    ").drop_duplicates(\"ligand_unique_ccd_code\")\n",
    "atp_analogs_code_set = set(atp_analogs.ligand_unique_ccd_code.unique())\n",
    "atp_analogs_code_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how they look"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = mols2grid.MolGrid(atp_analogs, smiles_col=\"ligand_rdkit_canonical_smiles\")\n",
    "grid.display(subset=[\"ligand_unique_ccd_code\", \"img\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can get the pocket clusters of all systems containing ATP analogs. Here we're using `pocket_qcov__50__weak__component`, meaning a system within a cluster has some other system in that cluster with which it shares at least 50% of pocket residues when aligned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pocket_cluster_column = \"pocket_qcov__50__weak__component\"\n",
    "plindex = query_index(\n",
    "    columns=[\n",
    "        \"system_id\",\n",
    "        \"entry_pdb_id\",\n",
    "        \"entry_release_date\",\n",
    "        pocket_cluster_column,\n",
    "        \"system_pocket_ECOD_t_name\",\n",
    "    ],\n",
    "    filters=[\n",
    "        (\n",
    "            \"ligand_unique_ccd_code\",\n",
    "            \"in\",\n",
    "            atp_analogs_code_set,\n",
    "        ),  # only ATP-binding systems\n",
    "        (\"ligand_num_interactions\", \">\", 2),  # with >2 interactions with ATP\n",
    "        (\"system_num_ligand_chains\", \"==\", 1),  # and ATP is the only ligand\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    \"No. systems:\",\n",
    "    plindex.system_id.nunique(),\n",
    "    \"No. PDB entries:\",\n",
    "    plindex.entry_pdb_id.nunique(),\n",
    "    \"No. clusters:\",\n",
    "    plindex[pocket_cluster_column].nunique(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[pocket_cluster_column].value_counts().head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we look at the ECOD topology names of the cluster \"c11\" for example, we see that these systems are GroEL equatorial domain-like domains. Interestingly, 7 systems from 2 different PDB entries don't have ECOD annotations yet as they were recently released, but are indeed [the same domain](https://doi.org/10.1038/s41467-024-45242-x)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[plindex[pocket_cluster_column] == \"c11\"].system_pocket_ECOD_t_name.value_counts(\n",
    "    dropna=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[\n",
    "    (plindex[pocket_cluster_column] == \"c11\")\n",
    "    & (plindex[\"system_pocket_ECOD_t_name\"].isna())\n",
    "].drop_duplicates(\"entry_pdb_id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plindex[plindex[\"system_pocket_ECOD_t_name\"].isna()][\"entry_pdb_id\"].nunique()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: Diversity sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `pli_unique_qcov__50__community` column clusters systems such that the protein and ligand make similar interactions with the pocket. This clustering combines protein sequence and structural similarity (needed to obtain the pocket alignment), as well as the ligand-pocket interactions, making it a good proxy for unique kinds of binding present in the dataset.\n",
    "\n",
    "Here's an example of how one might use `torch.utils.data.WeightedRandomSampler` to sample training systems based on their PLI community cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import WeightedRandomSampler\n",
    "\n",
    "cluster_column = \"pli_unique_qcov__50__community\"\n",
    "\n",
    "# Get train systems and their cluster labels\n",
    "plindex = query_index(\n",
    "    columns=[\"system_id\", \"entry_pdb_id\", cluster_column], splits=[\"train\"]\n",
    ").drop_duplicates(\"system_id\")\n",
    "\n",
    "# Add the number of systems in each cluster to the dataframe\n",
    "plindex = plindex.merge(\n",
    "    plindex[cluster_column].value_counts().rename(\"cluster_num_systems\"),\n",
    "    left_on=cluster_column,\n",
    "    right_index=True,\n",
    ").reset_index(drop=True)\n",
    "\n",
    "# Add the number of PDB entries in each cluster to the dataframe\n",
    "plindex = plindex.merge(\n",
    "    plindex.drop_duplicates(\"entry_pdb_id\")[cluster_column]\n",
    "    .value_counts()\n",
    "    .rename(\"cluster_num_entries\"),\n",
    "    left_on=cluster_column,\n",
    "    right_index=True,\n",
    ").reset_index(drop=True)\n",
    "\n",
    "# Ignore clusters with only one PDB entry\n",
    "sample_from = plindex[plindex[\"cluster_num_entries\"] > 1].reset_index(drop=True)\n",
    "\n",
    "# Calculate the weight for each cluster\n",
    "cluster_weights = 1.0 / sample_from.cluster_num_systems.values\n",
    "\n",
    "# Create a WeightedRandomSampler and sample systems from the train set\n",
    "sampler = WeightedRandomSampler(\n",
    "    weights=cluster_weights, num_samples=len(cluster_weights)\n",
    ")\n",
    "sampled_indices = list(sampler)\n",
    "sampled_plindex = (\n",
    "    sample_from.iloc[sampled_indices]\n",
    "    .reset_index(drop=True)\n",
    "    .drop(columns=[\"cluster_num_systems\", \"cluster_num_entries\"])\n",
    ")\n",
    "\n",
    "# Add the number of sampled systems in each cluster to the sampled dataframe\n",
    "sampled_plindex = sampled_plindex.merge(\n",
    "    sampled_plindex[cluster_column].value_counts().rename(\"cluster_num_systems\"),\n",
    "    left_on=cluster_column,\n",
    "    right_index=True,\n",
    ").reset_index(drop=True)\n",
    "\n",
    "# Add the number of sampled PDB entries in each cluster to the sampled dataframe\n",
    "sampled_plindex = sampled_plindex.merge(\n",
    "    sampled_plindex.drop_duplicates(\"entry_pdb_id\")[cluster_column]\n",
    "    .value_counts()\n",
    "    .rename(\"cluster_num_entries\"),\n",
    "    left_on=cluster_column,\n",
    "    right_index=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"No. of original systems: \", sample_from.system_id.nunique())\n",
    "print(\"No. of nonredundant sampled systems: \", sampled_plindex.system_id.nunique())\n",
    "print(\"No. of original clusters: \", sample_from[cluster_column].nunique())\n",
    "print(\"No. of sampled clusters: \", sampled_plindex[cluster_column].nunique())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how the sampling process has affected the distribution of cluster sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 3))\n",
    "\n",
    "cluster_sizes = (\n",
    "    sample_from.drop_duplicates(cluster_column)\n",
    "    .cluster_num_systems.value_counts()\n",
    "    .sort_index()\n",
    ")\n",
    "ax[0].hist(cluster_sizes.index, weights=cluster_sizes.values, bins=20, log=True)\n",
    "ax[0].set_xlabel(\"Cluster Size\")\n",
    "ax[0].set_ylabel(\"# Clusters (log scale)\")\n",
    "ax[0].grid(True, which=\"both\", ls=\"-\", alpha=0.2)\n",
    "\n",
    "cluster_sizes = (\n",
    "    sampled_plindex.drop_duplicates(cluster_column)\n",
    "    .cluster_num_systems.value_counts()\n",
    "    .sort_index()\n",
    ")\n",
    "ax[1].hist(cluster_sizes.index, weights=cluster_sizes.values, bins=20, log=False)\n",
    "ax[1].set_xlabel(\"Cluster Size\")\n",
    "ax[1].set_ylabel(\"# Clusters\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}