Build, visualize, and submit a tyk2 perturbation network using openfe and alchemiscale

tyk2_viz_alchemiscale.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "101b0492",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import copy\n",
    "import json\n",
    "import itertools\n",
    "from typing import Iterable, List, Tuple\n",
    "from ipywidgets import interact, widgets\n",
    "from rdkit import Chem\n",
    "from alchemiscale import AlchemiscaleClient, Scope, ScopedKey\n",
    "from openff.units import unit\n",
    "from gufe import AlchemicalNetwork, tokenization\n",
    "import openfe\n",
    "from openfe import (\n",
    "    LigandNetwork, ChemicalSystem, SmallMoleculeComponent, SolventComponent, ProteinComponent,\n",
    ")\n",
    "from openfe.setup.atom_mapping.lomap_scorers import default_lomap_score\n",
    "from openfe.setup.atom_mapping.lomap_mapper import LomapAtomMapper\n",
    "from openfe.setup.atom_mapping import LigandAtomMapper, LigandAtomMapping\n",
    "from openfe.utils.visualization_3D import view_mapping_3d\n",
    "from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol as RHTP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a22094",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_relative_network_from_names(ligands: Iterable[SmallMoleculeComponent],\n",
    "                                         connections: List[Tuple[str, str]],\n",
    "                                         mappers: Iterable[LigandAtomMapper],\n",
    "                                         scorer=None):\n",
    "    \"\"\"\n",
    "    Generate a network from an input list of tuples each containing\n",
    "    a pair of ligand names to connect with an edge.\n",
    "    Parameters\n",
    "    ----------\n",
    "    ligands : Iterable[SmallMoleculeComponent]\n",
    "        The ligands to create the network from.\n",
    "    connections : List[Tuple[str, str]]\n",
    "        The list of edges to create with each node identified by its\n",
    "        SmallMoleculeComponent name.\n",
    "    mapper : Iterable[LigandAtomMapper]\n",
    "        Mappers to use. At least 1 is required.\n",
    "    scorer : scoring function, optional\n",
    "        A callable which returns a float for any LigandAtomMapping. Used to\n",
    "        assign score to potential mappings, higher scores indicate worse\n",
    "        mappings.\n",
    "    Raises\n",
    "    ------\n",
    "    ValueError\n",
    "        If no mapping can be found for a supplied edge.\n",
    "    Returns\n",
    "    -------\n",
    "    network : LigandNetwork\n",
    "        Network of SmallMoleculeComponent transformations.\n",
    "    \"\"\"\n",
    "    edges = []\n",
    "\n",
    "    for entry in connections:\n",
    "        nodes = [idx for idx, lig in enumerate(ligands) if lig.name in entry]\n",
    "        best_score = 999.\n",
    "        for mapping in itertools.chain.from_iterable(\n",
    "            mapper.suggest_mappings(ligands[nodes[0]], ligands[nodes[1]])\n",
    "            for mapper in mappers\n",
    "        ):\n",
    "            if not scorer:\n",
    "                best_mapping = mapping\n",
    "                break\n",
    "\n",
    "            score = scorer(mapping)\n",
    "            mapping = mapping.with_annotations({\"score\": score})\n",
    "\n",
    "            if score < best_score:\n",
    "                best_mapping = mapping\n",
    "                best_score = score\n",
    "\n",
    "        if best_mapping is None:\n",
    "            raise ValueError(f\"No mapping for pair {entry}\")\n",
    "        edges.append(best_mapping)\n",
    "\n",
    "    return LigandNetwork(edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef1c82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a list of small molecule components\n",
    "supp = Chem.SDMolSupplier('tyk2_ligands.sdf', removeHs=False)\n",
    "mols = [SmallMoleculeComponent(m) for m in supp]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb43600e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A user defined list of ligand connections by SDF name field\n",
    "connections = [(\"lig_ejm_31\", \"lig_ejm_50\"),\n",
    "               (\"lig_ejm_46\", \"lig_jmc_23\"),\n",
    "               (\"lig_ejm_31\", \"lig_ejm_55\"),\n",
    "               (\"lig_ejm_31\", \"lig_ejm_48\"),\n",
    "               (\"lig_ejm_31\", \"lig_ejm_54\"),\n",
    "               (\"lig_ejm_31\", \"lig_ejm_47\"),\n",
    "               (\"lig_ejm_31\", \"lig_ejm_46\"),\n",
    "               (\"lig_ejm_46\", \"lig_jmc_27\"),\n",
    "               (\"lig_ejm_46\", \"lig_jmc_28\"),\n",
    "               (\"lig_ejm_42\", \"lig_ejm_43\"),\n",
    "               (\"lig_ejm_31\", \"lig_ejm_42\"),\n",
    "               (\"lig_ejm_45\", \"lig_ejm_55\"),]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e73f179",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a LigandNetwork from the connections using a Lomap mapper & scorer\n",
    "mappers = [LomapAtomMapper(time=20, threed=True, element_change=False, max3d=1),]\n",
    "scorer = default_lomap_score\n",
    "\n",
    "ligand_network = generate_relative_network_from_names(mols, connections, mappers=mappers, scorer=scorer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b8b9c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create method to display edges\n",
    "edges = list(ligand_network.edges)\n",
    "\n",
    "def display_edge(index):\n",
    "    view = view_mapping_3d(edges[index], spheres=True, show_atomIDs=True)\n",
    "    view.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77ebfab3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# traverse through all views\n",
    "interact(display_edge, index=widgets.IntSlider(min=0, max=len(edges)-1, step=1));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc22c93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a protein and solvent component\n",
    "protein_comp = ProteinComponent.from_pdb_file('tyk2_protein.pdb')\n",
    "solvent_comp = SolventComponent()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eafcc348",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get RHTP complex settings & protocol\n",
    "settings = RHTP.default_settings()\n",
    "settings.forcefield_settings.hydrogen_mass = 3.0\n",
    "settings.simulation_settings.equilibration_length = 1000 * unit.picosecond\n",
    "settings.simulation_settings.production_length = 5000 * unit.picosecond\n",
    "settings.integrator_settings.timestep = 4.0 * unit.femtosecond\n",
    "settings.integrator_settings.n_steps = 250 * unit.timestep\n",
    "settings.alchemical_sampler_settings.n_repeats = 1\n",
    "settings.alchemical_sampler_settings.online_analysis_interval = 200\n",
    "settings.alchemical_sampler_settings.online_analysis_target_error = 0.2 * unit.boltzmann_constant * unit.kelvin\n",
    "settings.simulation_settings.output_indices = \"not water\"\n",
    "settings.engine_settings.compute_platform = \"CUDA\"\n",
    "settings.system_settings.nonbonded_cutoff = 0.9 * unit.nanometer\n",
    "solvent_protocol = RHTP(settings=settings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "665818c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create network\n",
    "solvent_transformations = []\n",
    "complex_transformations = []\n",
    "alchemical_nodes = []\n",
    "\n",
    "for edge in edges:\n",
    "    stateA_solv = ChemicalSystem({'ligand': edge.componentA,\n",
    "                                  'solvent': solvent_comp})\n",
    "    stateA_prot = ChemicalSystem({'ligand': edge.componentA,\n",
    "                                  'protein': protein_comp,\n",
    "                                  'solvent': solvent_comp})\n",
    "    stateB_solv = ChemicalSystem({'ligand': edge.componentB,\n",
    "                                  'solvent': solvent_comp})\n",
    "    stateB_prot = ChemicalSystem({'ligand': edge.componentB,\n",
    "                                  'protein': protein_comp,\n",
    "                                  'solvent': solvent_comp})\n",
    "    alchemical_nodes.extend([stateA_prot, stateA_solv, stateB_prot, stateB_solv])\n",
    "\n",
    "    solvent_transformations.append(\n",
    "        openfe.Transformation(\n",
    "            stateA=stateA_solv, stateB=stateB_solv,\n",
    "            mapping={'ligand': edge},\n",
    "            protocol=solvent_protocol,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    complex_transformations.append(\n",
    "        openfe.Transformation(\n",
    "            stateA=stateA_prot, stateB=stateB_prot,\n",
    "            mapping={'ligand': edge},\n",
    "            protocol=solvent_protocol,\n",
    "        )\n",
    "    )\n",
    "\n",
    "network = AlchemicalNetwork(\n",
    "    edges=(solvent_transformations + complex_transformations),\n",
    "    nodes=alchemical_nodes,\n",
    "    name=\"tyk2_benchmark\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d868b77e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dump network to file\n",
    "print(\"generated network: \", network)\n",
    "with open('network.json', 'w') as f:\n",
    "    json.dump(network.to_dict(), f, cls=tokenization.JSON_HANDLER.encoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b12e42f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get alchemiscale client from env var defined ID and KEY\n",
    "user_id = os.environ['ALCHEMISCALE_ID']\n",
    "user_key = os.environ['ALCHEMISCALE_KEY']\n",
    "asc =  AlchemiscaleClient('https://api.alchemiscale.org', user_id, user_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d14aa09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get scope and create network\n",
    "scope = Scope('openfe', 'v0_7_4', 'tyk2')\n",
    "network_sk = asc.create_network(network, scope)\n",
    "print(network_sk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cb0c2bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# store the scoped key\n",
    "with open('scoped-key.dat', 'w') as f:\n",
    "    f.write(str(network_sk))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fed09c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# action out tasks\n",
    "for transform in network.edges:\n",
    "    transform_sk = asc.get_scoped_key(transform, scope)\n",
    "    tasks = asc.create_tasks(transform_sk, count=3)\n",
    "    asc.action_tasks(tasks, network_sk)"
   ]
  }
 ],
 "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}