structure_seer_code_examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31af1da0-8f61-45b2-bf58-dc36bd594a85",
   "metadata": {},
   "source": [
    "#### Demonstration of the Structure Seer model capabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4957dc86-0a5a-49fd-8245-e7ab1554aa6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch\n",
    "!pip install rdkit\n",
    "!pip install pubchempy\n",
    "!pip install networkx[default]\n",
    "!pip install tqdm\n",
    "!pip install matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2efb87ba-7699-4c86-a369-189d8f602c60",
   "metadata": {},
   "source": [
    " - Initiate the model class and load QM-9 pre-trained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d90f60b-604f-45f3-9380-93b49f41af7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from models import StructureSeer\n",
    "from utils import MolGraph\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "qm9_model = StructureSeer()\n",
    "qm9_model.encoder.load_state_dict(\n",
    "    torch.load(\n",
    "        \"weights/structure_seer/qm9/qm9_structure_seer_encoder.weights\",\n",
    "        map_location=\"cpu\",\n",
    "    )\n",
    ")\n",
    "qm9_model.decoder.load_state_dict(\n",
    "    torch.load(\n",
    "        \"weights/structure_seer/qm9/qm9_structure_seer_decoder.weights\",\n",
    "        map_location=\"cpu\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b35cc62-2229-40b2-9730-33df7c50aef4",
   "metadata": {},
   "source": [
    "- Element and shielding constants vectors declaration (input for the model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be952acb-55a1-4ac6-b865-0a291949558f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dsgdb9nsd_121697.xyz\n",
    "elements_vector = torch.tensor([6, 6, 6, 6, 6, 7, 7, 8, 8])\n",
    "shielding_constants_vector = torch.tensor([ 79.994, 93.406, 179.437, 225.061, 227.187, 55.066, 198.318, -204.424, 344.265])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5cc720d-cd45-49c3-8e10-fa4e4630e207",
   "metadata": {},
   "source": [
    "Pre-processing of input\n",
    "- Padding\n",
    "- Normalisation of shielding constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7ff88f6-dda7-4364-9bde-4dfe73228f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_input(elements_vector, shielding_constants_vector):\n",
    "    size = elements_vector.size(dim=0)\n",
    "    pad_size = 54 - size\n",
    "    # Normalise shielding to (0, 1) range for shielding in (-1000, 1000)\n",
    "    normalised_shielding = torch.tensor([(x + 1000) / 2000 for x in shielding_constants_vector])\n",
    "    input_vectors = [ F.pad(elements_vector, (0, pad_size), value=0), F.pad(normalised_shielding, (0, pad_size), value=0)]\n",
    "    return input_vectors, size\n",
    "\n",
    "input_vectors, size = preprocess_input(elements_vector, shielding_constants_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "763cca17-fc66-479f-a576-d4f4621c2b7f",
   "metadata": {},
   "source": [
    "Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5ef0a3f-2230-4b01-b39a-b38bf61bea46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import show_adjacency_matrix\n",
    "%matplotlib inline\n",
    "\n",
    "prediction = qm9_model.predict(input_vectors)\n",
    "show_adjacency_matrix(matrix=prediction, title=\"Prediction\", size=size)\n",
    "nodes = torch.zeros((size,2))\n",
    "for i in range(size):\n",
    "    nodes[i][0] = elements_vector[i]\n",
    "    nodes[i][1] = shielding_constants_vector[i]\n",
    "    \n",
    "prediction_graph = MolGraph.from_adjacency_matrix(nodes=nodes, adjacency_matrix=prediction)\n",
    "prediction_graph.show(size=4, title=\"Predicted molecular graph\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4063119d-c608-41e0-87ed-884129304c56",
   "metadata": {},
   "source": [
    "- Find and show target adjacency matrix and molecular graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99582328-e75c-4017-bc20-e76c543d4c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import read_sdf_compounds\n",
    "compounds = read_sdf_compounds(\"example_datasets/demo_compounds_qm9.sdf\")\n",
    "for compound in compounds:\n",
    "    if compound.GetProp(\"_Name\") == \"dsgdb9nsd_121697.xyz\":\n",
    "        target_graph = MolGraph.from_mol(compound, shielding=True).sort(shielding=True)[0]\n",
    "        target_matrix = target_graph.adjacency_matrix()\n",
    "        show_adjacency_matrix(matrix=target_matrix, title=\"Target\", size=size)\n",
    "        target_graph.show(size=4, title=\"Target molecular graph\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33d53f7b-9f0f-4b4c-be80-4422aac5c1b1",
   "metadata": {},
   "source": [
    "The same workflow for the PubChem dataset sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee57a0b2-a01a-44ce-aa89-00062671c489",
   "metadata": {},
   "outputs": [],
   "source": [
    "pubchem_model = StructureSeer()\n",
    "pubchem_model.encoder.load_state_dict(\n",
    "    torch.load(\n",
    "        \"weights/structure_seer/pubchem/pubchem_structure_seer_encoder.weights\",\n",
    "        map_location=\"cpu\",\n",
    "    )\n",
    ")\n",
    "pubchem_model.decoder.load_state_dict(\n",
    "    torch.load(\n",
    "        \"weights/structure_seer/pubchem/pubchem_structure_seer_decoder.weights\",\n",
    "        map_location=\"cpu\",\n",
    "    )\n",
    ")\n",
    "\n",
    "# pubchem_cid_86033414\n",
    "elements_vector = torch.tensor([ 6,  6,  6,  6,  8,  9,  9, 16])\n",
    "shielding_constants_vector = torch.tensor([127.9520, 180.2580, 221.0570, 223.8610, 420.3100, 437.6430, 439.1160,\n",
    "        635.8940])\n",
    "input_vectors, size = preprocess_input(elements_vector, shielding_constants_vector)\n",
    "\n",
    "prediction = pubchem_model.predict(input_vectors)\n",
    "show_adjacency_matrix(matrix=prediction, title=\"Prediction\", size=size)\n",
    "nodes = torch.zeros((size,2))\n",
    "for i in range(size):\n",
    "    nodes[i][0] = elements_vector[i]\n",
    "    nodes[i][1] = shielding_constants_vector[i]\n",
    "    \n",
    "prediction_graph = MolGraph.from_adjacency_matrix(nodes=nodes, adjacency_matrix=prediction)\n",
    "prediction_graph.show(size=4, title=\"Predicted molecular graph\")\n",
    "\n",
    "compounds = read_sdf_compounds(\"example_datasets/demo_compounds_pubchem.sdf\")\n",
    "for compound in compounds:\n",
    "    if compound.GetProp(\"_Name\") == \"pubchem_cid_86033414\":\n",
    "        target_graph = MolGraph.from_mol(compound, shielding=True).sort(shielding=True)[0]\n",
    "        target_matrix = target_graph.adjacency_matrix()\n",
    "        show_adjacency_matrix(matrix=target_matrix, title=\"Target\", size=size)\n",
    "        target_graph.show(size=4, title=\"Target molecular graph\")\n",
    "        break\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b23fca10-2233-4a1c-a819-a6d669f67aa3",
   "metadata": {},
   "source": [
    "The model can be used for peak attribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4768db-7c49-4f50-9ff4-da8e9b922331",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import generate_shielding_permutations, generate_adjacency_matrix_permutations\n",
    "# dsgdb9nsd_004997.xyz\n",
    "compounds = read_sdf_compounds(\"example_datasets/demo_compounds_qm9.sdf\")\n",
    "for compound in compounds:\n",
    "    if compound.GetProp(\"_Name\") == \"dsgdb9nsd_004997.xyz\":\n",
    "        compound_for_attribution = compound\n",
    "        break\n",
    "        \n",
    "graph_for_attribution = MolGraph.from_mol(compound_for_attribution, shielding=False).sort()[0]\n",
    "matrix = graph_for_attribution.adjacency_matrix()\n",
    "\n",
    "# Spectra is specified\n",
    "elements_vector = torch.tensor([6, 6, 6, 6, 6, 7, 7, 8])\n",
    "shielding_constants_vector = torch.tensor([79.9410,  91.0000,  92.8660, 115.5880, 217.7200,   1.4700,  15.2030, 206.3840])\n",
    "\n",
    "input_vectors, size = preprocess_input(elements_vector, shielding_constants_vector)\n",
    "nodes = torch.zeros((size,2))\n",
    "for i in range(size):\n",
    "    nodes[i][0] = elements_vector[i]\n",
    "    nodes[i][1] = shielding_constants_vector[i]\n",
    "\n",
    "permutations = generate_shielding_permutations(nodes)\n",
    "perm_count = len(permutations)\n",
    "\n",
    "input_permutations = generate_adjacency_matrix_permutations(elements_vector, permutations, matrix)\n",
    "t1 = torch.argmax(input_permutations, dim=-1)\n",
    "\n",
    "output = qm9_model.predict(input_vectors)\n",
    "prediction_tensor = torch.stack([output] * perm_count, dim=0)\n",
    "t2 = torch.argmax(prediction_tensor, dim=-1)\n",
    "\n",
    "diff = torch.abs(t2 - t1)\n",
    "s_diff = torch.reshape(diff, (diff.size(dim=0), diff.size(dim=1) * diff.size(dim=2)))\n",
    "summs = torch.sum(s_diff, dim=-1)\n",
    "top_idx = torch.argmin(summs)\n",
    "top_permutation = input_permutations[top_idx]\n",
    "\n",
    "predicted_attribution = MolGraph.from_adjacency_matrix(nodes=nodes, adjacency_matrix=top_permutation)\n",
    "target_attribution = MolGraph.from_mol(compound_for_attribution, shielding=True)\n",
    "\n",
    "graph_for_attribution.show(size=4, title=\"Input molecular graph\")\n",
    "predicted_attribution.show(size=4, title=\"Predicted attribution\")\n",
    "target_attribution.show(size=4, title=\"Target attribution\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ea52d5-b0e8-4e82-bcc2-e6068cde7460",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}