mofreinforce/tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6122f2e6-08f2-440d-bc1f-7328f723b4cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce0484b9-a389-44b3-b4b7-6714c0987d2d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import shutil\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from ase.io import read\n",
    "from ase.visualize.plot import plot_atoms\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "from collections import Counter\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import RDConfig\n",
    "sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))\n",
    "import sascorer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfc6f6d8-7301-4ee7-9c8a-ad6c4bed0dc1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.rcParams['font.sans-serif'] = \"Arial\"\n",
    "plt.rcParams['font.family'] = \"sans-serif\"\n",
    "plt.rcParams[\"font.size\"] = 15\n",
    "plt.rcParams[\"xtick.major.size\"] = 0\n",
    "plt.rcParams[\"ytick.major.size\"] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e841e725-99e8-4bf3-a743-bac4cc1c51b1",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Construction of the generated MOFs by PORMAKE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb3ea000-c7a1-42de-88f1-40d0f6d9668d",
   "metadata": {},
   "source": [
    "After executing a `test` of reinforcement learning, the results will be saved in JSON format at the `log_dir` directory. The JSON file contains information on the MOFs generated from a test set of 10,000 data, including:\n",
    "```\n",
    "$ python run_reinforce.py with v0_qkh_round3 log_dir=test test_only=True load_path=model/reinforce/best_v0_qkh_round3.ckpt\n",
    "```\n",
    "\n",
    "`rewards` : the rewards of generated MOFs\n",
    "\n",
    "`pres` : the estimated target values of generated MOFs by the predictor\n",
    "\n",
    "`gen_sms` : the smiles of organic linkers of generated MOFs\n",
    "\n",
    "`gen_mcs` : the metal clusters of generated MOFs\n",
    "\n",
    "`gen_topos` : the topologies of generated MOFs"
   ]
  },
  {
   "cell_type": "raw",
   "id": "9bf1f281-7daf-40ab-805d-c2ccf3cc4c31",
   "metadata": {
    "tags": []
   },
   "source": [
    "To construct the generated MOFs, it is necessary to install PORMAKE. The installation instructions for PORMAKE can be found at [PORMAKE](https://github.com/Sangwon91/PORMAKE) Github repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0e01eee-c9ef-4ac7-9236-9824f21cc111",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "path_json = \"test/results_v0_qkh_round3_seed0_from_best_v0_qkh_round3.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ea5e96-56ee-4443-be2b-e1ff5b93f6ba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results_optimized = json.load(open(path_json))\n",
    "ret = np.array( list(zip(*results_optimized[\"preds\"])) +\n",
    "    [\n",
    "        results_optimized[\"gen_topos\"],\n",
    "        results_optimized[\"gen_mcs\"],\n",
    "        results_optimized[\"gen_sms\"],\n",
    "    ]\n",
    ").T\n",
    "len(ret)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f347692-7ace-4ba7-aaf9-2deb0cdf9c77",
   "metadata": {},
   "source": [
    "## 0. Analysis of generated MOFs for top 500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "897658e9-9329-4ca3-bf71-9f1ac2d5f2c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_n = 500\n",
    "ret = np.unique(ret, axis=0)\n",
    "sorted_ret = ret[np.argsort(ret[:, 0].astype(float))][:top_n]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c33bb729-49a7-43a4-8a7f-0efddd21fc2e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### top common organic linkers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe6354e-3b66-4a1a-b0a0-65aae0a29c59",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "counter = Counter()\n",
    "counter.update(sorted_ret[:,-1].tolist())\n",
    "top_ol_sm, c = zip(*counter.most_common())\n",
    "top_ol = [Chem.MolFromSmiles(sm) for sm in top_ol_sm]\n",
    "len(top_ol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b168c370-05f4-44a1-b9dd-06fafcbf41b3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "img = Chem.Draw.MolsToGridImage(top_ol[:16], molsPerRow=4,  subImgSize=[200, 200], useSVG=True)\n",
    "img"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17923b67-c36e-431c-a19f-7001b4caa9e7",
   "metadata": {
    "tags": []
   },
   "source": [
    "### top common metal clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cfbc706-57a3-4044-aa0e-07ccdf6524a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pormake as pm\n",
    "from ase.visualize.plot import plot_atoms\n",
    "database = pm.Database()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6730ea05-bb4f-4150-9dc8-c76bf5f42f5d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# topology\n",
    "counter = Counter()\n",
    "counter.update(sorted_ret[:, 1].tolist())\n",
    "top_topo, count = zip(*counter.most_common())\n",
    "len(top_topo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5fc4f0a-721b-4ab0-b98e-d591cd402642",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_topo = top_topo[:16]\n",
    "fig, axes = plt.subplots(len(top_topo) // 4, 4, figsize=(8, 8), constrained_layout=True)\n",
    "for i, t in enumerate(top_topo):\n",
    "    ax = axes[i//4, i%4]\n",
    "    topo_ = database.get_topo(t)\n",
    "    plot_atoms(topo_.atoms, ax=ax, radii=0.2)\n",
    "    ax.set_title(f\"{t}\\n({round(count[i]/sum(count) * 100, 1)} %)\", fontsize=20)\n",
    "    ax.set_axis_off()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25f554d7-840d-43f7-8baf-c0f71c9bc376",
   "metadata": {
    "tags": []
   },
   "source": [
    "### top common metal clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eb6b9ac-14ed-4fbb-bca5-1dfcade416f2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# metal cluster\n",
    "counter = Counter()\n",
    "counter.update(sorted_ret[:, 2].tolist())\n",
    "top_mc, count = zip(*counter.most_common())\n",
    "len(top_mc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43bca93-c8af-4645-9bf1-c5e37ca11895",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "top_mc = top_mc[:16]\n",
    "fig, axes = plt.subplots(len(top_mc) // 4, 4, figsize=(8, 8), constrained_layout=True)\n",
    "for i, mc in enumerate(top_mc):\n",
    "    ax = axes[i//4, i%4]\n",
    "    mc_ = database.get_bb(mc)\n",
    "    plot_atoms(mc_.atoms, ax=ax)\n",
    "    ax.set_title(f\"{mc}\\n({round(count[i]/sum(count) * 100, 1)} %)\", fontsize=20)\n",
    "    ax.set_axis_off()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adcbfc9f-36d5-40ea-8c17-998fd9320ca7",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1. contruction of the generated MOFs by PORMAKE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a231e4f-c8ee-44db-8722-93d750041f34",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from rdkit import RDLogger  \n",
    "RDLogger.DisableLog('rdApp.*')                                                                                                                                                           \n",
    "pm.log.disable_print()\n",
    "pm.log.disable_file_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aed87cc3-a7d6-4737-af36-126a5bc3a0b0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_dir_bb = \"results/qkh/bb_dir\"\n",
    "save_dir_gen_mofs = \"results/qkh/gen_mofs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52a61f77-e502-443b-a4a5-548e82bd5d6e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# move node bbs from origin bb_dir (PORMAKE) to new bb_dir\n",
    "if os.path.exists(save_dir_bb):\n",
    "    shutil.rmtree(save_dir_bb)\n",
    "shutil.copytree(f\"{pm.__path__[0]}/database/bbs\", save_dir_bb)\n",
    "# bb_dir\n",
    "database = pm.Database(bb_dir=Path(save_dir_bb))\n",
    "# save_dir for generated MOFs\n",
    "os.makedirs(save_dir_gen_mofs, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c37740cf-73a2-4c68-9f92-b63ef7a52ac5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def smiles_to_xyz(smiles, save_dir, bb_name=\"tmp\"):\n",
    "    # smiles to mol\n",
    "    mol = Chem.MolFromSmiles(smiles)\n",
    "    # mol to 3D mol\n",
    "    m = Chem.AddHs(mol)\n",
    "    AllChem.EmbedMolecule(m)\n",
    "    AllChem.MMFFOptimizeMolecule(m)\n",
    "    # mol to molblock\n",
    "    mol_block = Chem.MolToMolBlock(m)\n",
    "    lines = mol_block.splitlines()\n",
    "    # write xyz file\n",
    "    line = lines[3]\n",
    "    num_atoms = int(line[:3])\n",
    "    num_bonds = int(line[3:6])\n",
    "    \n",
    "    save_path = os.path.join(save_dir, f\"{bb_name}.xyz\")\n",
    "    with open(save_path, \"w\") as f:\n",
    "        f.write(f\"{num_atoms}\\n\")\n",
    "        f.write(f\"mol to xyz file\\n\")\n",
    "        # coords\n",
    "        for line in lines[4:4+num_atoms]:\n",
    "            tokens = line.split()\n",
    "            # change dummy atoms R to X\n",
    "            if tokens[3] == \"R\":\n",
    "                tokens[3] = \"X\"\n",
    "            f.write(f\"{tokens[3]:<10}    {tokens[0]:<10}    {tokens[1]:<10}    {tokens[2]:<10}\\n\")\n",
    "        # bonds\n",
    "        for line in lines[4+num_atoms:4+num_atoms+num_bonds]:\n",
    "            tokens = [int(line[:3]), int(line[3:6]), int(line[6:9])]\n",
    "            # bond type\n",
    "            if tokens[2] == 1:\n",
    "                bond_type = \"S\"\n",
    "            elif tokens[2] == 2:\n",
    "                bond_type = \"D\"\n",
    "            elif tokens[2] == 3:\n",
    "                bond_type = \"T\"\n",
    "            elif tokens[2] == 4:\n",
    "                bond_type = \"A\"\n",
    "            else:\n",
    "                raise Exception(\"bond type error\")\n",
    "            # find index of atom\n",
    "            idx_1 = int(tokens[0]) - 1\n",
    "            idx_2 = int(tokens[1]) - 1\n",
    "            f.write(f\"{idx_1:<10}{idx_2:<6}{bond_type:<6}\\n\")\n",
    "        f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60a7da05-3a81-48e7-aa22-e5751a7297c3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def construct_mofs(final_ret, save_dir_bb, save_dir_gen_mofs):\n",
    "    e0 = 0 # build error\n",
    "    e1 = 0\n",
    "    e2 = 0\n",
    "    e3 = 0\n",
    "    e4 = 0\n",
    "\n",
    "    idx = 0\n",
    "    vocab_sm = {}\n",
    "    for p, topo_, mc_, sm_ in tqdm(final_ret):\n",
    "        print(p, topo_, mc_, sm_)\n",
    "        # save smiles to xyz file\n",
    "        try:\n",
    "            if sm_ not in vocab_sm.keys():\n",
    "                smiles_to_xyz(sm_, save_dir=save_dir_bb, bb_name=f\"{len(vocab_sm)}\")\n",
    "                vocab_sm[sm_] = f\"{len(vocab_sm)}\"\n",
    "        except Exception as e:\n",
    "            e0 += 1\n",
    "            print(\"The smile of organice linker can't be converted to xyz files\")\n",
    "            continue\n",
    "\n",
    "\n",
    "        # get topo, mc, ol\n",
    "        topo = database.get_topo(topo_)\n",
    "        mc = database.get_bb(mc_)\n",
    "        ol = database.get_bb(vocab_sm[sm_])\n",
    "\n",
    "        # check connection point matching\n",
    "        topo_cn = list(topo.unique_cn)\n",
    "        if len(topo_cn) == 1:\n",
    "            topo_cn.append(2)\n",
    "        mc_cn = mc.n_connection_points\n",
    "        ol_cn = sm_.count(\"*\")\n",
    "\n",
    "        if set(topo_cn) != set([mc_cn, ol_cn]):\n",
    "            print(f\"fail : {topo_cn, mc_cn, ol_cn}\")\n",
    "            continue\n",
    "\n",
    "        # assingn node and edge\n",
    "        if mc_cn == topo_cn[0] and topo_cn[1] == 2:\n",
    "            node_bbs = {\n",
    "                0 : mc,\n",
    "            }\n",
    "            edge_bbs = {\n",
    "                tuple(topo.unique_edge_types[0]) : ol,\n",
    "            }\n",
    "        elif mc_cn == topo_cn[0] and topo_cn[1] != 2:\n",
    "            node_bbs = {\n",
    "                0 : mc,\n",
    "                1 : ol,\n",
    "            }\n",
    "            edge_bbs = {}\n",
    "        else:\n",
    "            node_bbs = {\n",
    "                0 : ol,\n",
    "                1 : mc,\n",
    "            }\n",
    "            edge_bbs = {}\n",
    "        # build MOF\n",
    "        builder = pm.Builder()\n",
    "        try:\n",
    "            gen_mof = builder.build_by_type(topology=topo, node_bbs=node_bbs, edge_bbs=edge_bbs)\n",
    "        except Exception as e:\n",
    "            e0 += 1\n",
    "            continue\n",
    "\n",
    "        # check criterion\n",
    "        # (1) SAscore < 6\n",
    "        m = Chem.MolFromSmiles(sm_)\n",
    "        score = sascorer.calculateScore(m)\n",
    "        if score > 6:\n",
    "            e1 += 1\n",
    "            continue\n",
    "        # (2) rmsd\n",
    "        if gen_mof.info[\"max_rmsd\"] > 0.3:\n",
    "            e2 += 1\n",
    "            continue\n",
    "        # (3) # of atoms <= 3000\n",
    "        if len(gen_mof.atoms) > 3000:\n",
    "            e3 += 1\n",
    "            continue\n",
    "        # (4) length of cells < 60 A\n",
    "        if gen_mof.atoms.get_cell_lengths_and_angles()[:3].max() > 60:\n",
    "            e4 += 1\n",
    "            continue\n",
    "\n",
    "        # write cif\n",
    "        try:\n",
    "            filename = f\"{str(idx).zfill(3)}_{topo_}+{mc_}+{vocab_sm[sm_]}.cif\"\n",
    "            print(f\"write_cif {filename}\")\n",
    "            gen_mof.write_cif(f\"{save_dir_gen_mofs}/{filename}\")\n",
    "            idx += 1\n",
    "        except:\n",
    "            e0 += 1\n",
    "    # write vocab for smiles of organice linker\n",
    "    json.dump(vocab_sm, open(f\"{save_dir_gen_mofs}/vocab_sm.json\", \"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba9d6354-fec9-4b54-8cf5-a8054712fd62",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# construct top 100 MOFs\n",
    "final_ret = sorted_ret[:100]\n",
    "construct_mofs(final_ret, save_dir_bb, save_dir_gen_mofs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a307d9d5-12be-4260-8232-f0ca0f184176",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 2. visualize the constructed MOFs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6926a7fc-10c1-4869-b509-2d918e749f8f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filenames = sorted(list(Path(save_dir_gen_mofs).glob(\"*.cif\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3132f833-cd14-42cd-8fdb-f7322faee9ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "p = filenames[1]\n",
    "atoms = read(p)\n",
    "plot_atoms(atoms)\n",
    "print(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3eda8f-9c4b-4fed-a774-976eb5a53e08",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3, 3, figsize=(16, 16))\n",
    "for i in range(9):\n",
    "    ax = axes[i//3, i%3]\n",
    "    atoms = read(filenames[i])\n",
    "    plot_atoms(atoms, ax=ax)\n",
    "    ax.set_title(filenames[i].name.split(\"/\")[-1].split(\".\")[0], fontsize=20)\n",
    "    ax.set_axis_off()\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}