RUNME_unconditional_generation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eebefaa7-cd7c-4d1f-8e8d-da2ddffbef55",
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d \n",
    "from shepherd_score_utils.generate_point_cloud import (\n",
    "    get_atom_coords, \n",
    "    get_atomic_vdw_radii, \n",
    "    get_molecular_surface,\n",
    "    get_electrostatics,\n",
    "    get_electrostatics_given_point_charges,\n",
    ")\n",
    "from shepherd_score_utils.pharm_utils.pharmacophore import get_pharmacophores\n",
    "from shepherd_score_utils.conformer_generation import update_mol_coordinates\n",
    "\n",
    "print('importing rdkit')\n",
    "import rdkit\n",
    "from rdkit.Chem import rdDetermineBonds\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print('importing torch')\n",
    "import torch\n",
    "import torch_geometric\n",
    "from torch_geometric.nn import radius_graph\n",
    "import torch_scatter\n",
    "\n",
    "import pickle\n",
    "from copy import deepcopy\n",
    "import os\n",
    "import multiprocessing\n",
    "from tqdm import tqdm\n",
    "\n",
    "import sys\n",
    "sys.path.insert(-1, \"model/\")\n",
    "sys.path.insert(-1, \"model/equiformer_v2\")\n",
    "\n",
    "print('importing lightning')\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import CSVLogger\n",
    "\n",
    "from lightning_module import LightningModule\n",
    "from datasets import HeteroDataset\n",
    "\n",
    "import importlib\n",
    "\n",
    "from inference import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18e33ab4-ba3d-42aa-b7ea-37f5a547b7d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "47d84cc9-4a50-4d38-b62b-90d1220afd70",
   "metadata": {},
   "source": [
    "## P(x1,x2), P(x1,x3), P(x1,x4) models trained on ShEPhERD-GDB-17"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b94d3287-0d1a-4ae3-8119-6e753ab29c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pick one\n",
    "#chkpt = 'shepherd_chkpts/x1x2_diffusion_gdb17_20240824_submission.ckpt'\n",
    "#chkpt = 'shepherd_chkpts/x1x3_diffusion_gdb17_20240824_submission.ckpt'\n",
    "chkpt = 'shepherd_chkpts/x1x4_diffusion_gdb17_20240824_submission.ckpt'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac107bc0-f454-468a-9287-d46ed385f902",
   "metadata": {},
   "source": [
    "## P(x1,x3,x4) model trained on ShEPhERD-MOSES-aq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "120a2d86-2a6c-4bcb-b5b6-3bde2d51d15d",
   "metadata": {},
   "outputs": [],
   "source": [
    "chkpt = 'shepherd_chkpts/x1x3x4_diffusion_mosesaq_20240824_submission.ckpt' # checkpoint used for evaluations in preprint\n",
    "#chkpt = 'shepherd_chkpts/x1x3x4_diffusion_mosesaq_20240824_30epochs_latest.ckpt' # latest checkpoint that was trained for 2-3X longer than the original version in the preprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37fd3aca-9bec-42fd-b737-5af786021921",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "209e1bfb-3f5c-46cf-9be4-b6256c83cba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "model_pl = LightningModule.load_from_checkpoint(chkpt)\n",
    "params = model_pl.params\n",
    "model_pl.to(device)\n",
    "model_pl.model.device = device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b65e6d-cae7-4de7-bc5d-d0ccb00c0324",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 10\n",
    "n_atoms = 60\n",
    "num_pharmacophores = 10 # set to 5 (just a dummy value) if using a model that does not model x4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50631996-8cbf-4df7-97bd-a614f56649b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use to break symmetry during unconditional generation\n",
    "T = params['noise_schedules']['x1']['T'] # T == 400\n",
    "inject_noise_at_ts = list(np.arange(130, 80, -1)) # [150]\n",
    "inject_noise_scales = [1.0] * len(inject_noise_at_ts)\n",
    "harmonize = True\n",
    "harmonize_ts = [80]\n",
    "harmonize_jumps = [20]\n",
    "\n",
    "# to NOT break symmetry (expect spherical molecules with low diversity), use this instead:\n",
    "\"\"\"\n",
    "inject_noise_at_ts = []\n",
    "inject_noise_scales = []\n",
    "harmonize = False\n",
    "harmonize_ts = []\n",
    "harmonize_jumps = []\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "generated_samples = inference_sample(\n",
    "    model_pl,\n",
    "    batch_size = batch_size,\n",
    "    \n",
    "    N_x1 = n_atoms,\n",
    "    N_x4 = num_pharmacophores, \n",
    "    \n",
    "    unconditional = True,\n",
    "    \n",
    "    prior_noise_scale = 1.0,\n",
    "    denoising_noise_scale = 1.0,\n",
    "    \n",
    "    # use to break symmetry during unconditional generation\n",
    "    inject_noise_at_ts = inject_noise_at_ts,\n",
    "    inject_noise_scales = inject_noise_scales, \n",
    "    harmonize = harmonize,\n",
    "    harmonize_ts = harmonize_ts,\n",
    "    harmonize_jumps = harmonize_jumps,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81d1e269-bbab-47ec-a1e0-cfd301ae2677",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(generated_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db8b465e-2d03-4894-af27-03942832410b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x1']['atoms']) # atomic numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7ab362-2b92-4638-a9dd-3a986fc50723",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x1']['positions']) # atomic coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "158c927b-a010-40fb-9d37-82e9768b7af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x2']['positions']) # shape surface point coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230086e4-ef0b-4fc8-950a-3a06d7fde112",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x3']['positions']) # ESP surface point coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c329952-a56d-4cd6-811b-80e485b7c9ea",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x3']['charges']) # ESP values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e8c66d3-2178-410b-89d6-1c54692527ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x4']['types']) # pharmacophore types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5672f111-498f-4d72-80bf-e5adbe38912b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x4']['positions']) # pharmacophore positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52aadc47-6d8b-4096-ae9d-4e18cad095ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_samples[0]['x4']['directions']) # pharmacophore directions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af20e867-08d6-49e0-8626-b959ac518738",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0fe3005-0e75-4d72-81e0-688134383885",
   "metadata": {},
   "outputs": [],
   "source": [
    "# quick visualization of generated samples\n",
    "# full analyses, including extensive validity checks, can be performed by following https://github.com/coleygroup/shepherd-score\n",
    "\n",
    "for b,sample_dict in enumerate(generated_samples):\n",
    "    \n",
    "    xyz = '' \n",
    "    \n",
    "    x_ = sample_dict['x1']['atoms']\n",
    "    pos_ = sample_dict['x1']['positions']\n",
    "    \n",
    "    xyz += f'{len(x_)}\\n{b+1}\\n'\n",
    "    for a in range(len(x_)):\n",
    "        atomic_number_ = int(x_[a])\n",
    "        position_ = pos_[a]\n",
    "        \n",
    "        xyz+= f'{rdkit.Chem.Atom(atomic_number_).GetSymbol()} {str(position_[0].round(3))} {str(position_[1].round(3))} {str(position_[2].round(3))}\\n'\n",
    "    xyz+= '\\n'\n",
    "    \n",
    "    try:\n",
    "        mol_ = rdkit.Chem.MolFromXYZBlock(xyz)\n",
    "    except Exception as e:\n",
    "        mol_ = None\n",
    "        print(f'invalid molecule: {e}')\n",
    "        continue\n",
    "    \n",
    "    try:\n",
    "        for c in [0, 1, -1, 2, -2]:\n",
    "            mol__ = deepcopy(mol_)\n",
    "            try:\n",
    "                rdkit.Chem.rdDetermineBonds.DetermineBonds(mol__, charge = c, embedChiral = True)\n",
    "            except:\n",
    "                continue\n",
    "            if mol__ is not None:\n",
    "                print(c)\n",
    "                break \n",
    "    except Exception as e:\n",
    "        mol_ = None\n",
    "        print(f'invalid molecule: {e}')\n",
    "        continue\n",
    "    \n",
    "    mol_ = mol__\n",
    "    try:\n",
    "        assert sum([a.GetNumRadicalElectrons() for a in mol_.GetAtoms()]) == 0, 'has radical electrons'\n",
    "        mol_.UpdatePropertyCache()\n",
    "        rdkit.Chem.GetSymmSSSR(mol_)\n",
    "        \n",
    "    except Exception as e:\n",
    "        mol_ = None\n",
    "        print(f'invalid molecule: {e}')\n",
    "        continue\n",
    "\n",
    "    display(rdkit.Chem.MolFromSmiles(rdkit.Chem.MolToSmiles(mol_)))\n",
    "    \n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7327069-27a7-4c23-92cb-0b6ebfe30ce0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-shepherd]",
   "language": "python",
   "name": "conda-env-.conda-shepherd-py"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}