00_qm9_final_postprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c7c0172d-ad8e-42f8-b796-b4ec7833e33e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%config Completer.use_jedi = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f387788b-c16a-45d2-8e70-e309ca3b944d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "import sys\n",
    "import pickle\n",
    "\n",
    "import numpy as np\n",
    "from pymatgen.core.structure import Molecule\n",
    "from scipy.interpolate import InterpolatedUnivariateSpline\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ecad749-6306-40bf-bb70-d9e5f7378648",
   "metadata": {},
   "source": [
    "Append the `home` path of this project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "77708ce8-462d-4d5a-9f97-6ae089a9a852",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append(str(Path.cwd().parent))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a8d5d44-b24b-45c3-b82e-a6542ee345f4",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "571da3ab-8146-44d6-b8cd-f63aa063848a",
   "metadata": {},
   "source": [
    "This notebook contains the final post-processing scripts for preparing the machine learning (ML) inputs for the QM9 datasets. If you want to generate all files from scratch, you'll need to run the `XAS-NNE/data/qm9/00_process_qm9_data.py` script from within the `XAS-NNE/data/qm9/` directory.\n",
    "\n",
    "However, we have constructed the ML inputs already and stored them in the GitHub repository. These inputs are all derived directly from the final output `pickle` file from `XAS-NNE/data/qm9/00_process_qm9_data.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "946441e3-11c6-4b93-8add-73a8faea523f",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Construct the ACSF feature vectors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ada6597-bb90-4960-89d7-e40cac6cd590",
   "metadata": {},
   "source": [
    "We use the Atom-centered Symmetry Functions (ACSF) feature vectors as inputs for the ML models. See [here](https://singroup.github.io/dscribe/latest/tutorials/descriptors/acsf.html) for the `Dscribe` library implementation docs. This is the original paper:\n",
    "\n",
    "> Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. J. Chem. Phys., 134(7):074106, 2011."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d8d2400d-d088-4052-b833-5817642918a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dscribe.descriptors import ACSF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a965bb42-f8a4-4db7-ac53-e2c127104522",
   "metadata": {},
   "source": [
    "Read in the results from the outputs of the previous scripts. This is a relatively large `pickle` file (~20 GB) so proceed with caution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "214b42de-70af-49a2-b8fa-070975298ac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "path = Path(\"data/qm9/XANES-220622-C-N-O.pkl\")\n",
    "print(path.exists())\n",
    "data = pickle.load(open(path, \"rb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39ce699a-a499-45a1-b522-9fdecbcea1ce",
   "metadata": {},
   "source": [
    "Setup some common grids for interpolating all of the spectra onto."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5c8e846d-4024-46ce-ae6c-470c77f0dc03",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 200\n",
    "grids = {\n",
    "    \"O\": np.linspace(528, 582, N),\n",
    "    \"N\": np.linspace(395, 449, N),\n",
    "    \"C\": np.linspace(275, 329, N)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eceec05-5f2f-425e-b13f-d2c3a1ca4860",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Construct a lookup table for the molecular spectra if desired"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e83bab79-e041-49ca-a2b7-4c713f6941bb",
   "metadata": {},
   "source": [
    "Note this isn't really needed in this work but we created it for convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b405072c-7b51-4ae7-8e0c-d32a2660d07c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lookup_table_data = {\n",
    "    dat[\"smiles\"]: {\n",
    "        \"qm9id\": qm9id,\n",
    "        \"canon_smiles\": None,\n",
    "        \"C-XANES\": None,\n",
    "        \"N-XANES\": None,\n",
    "        \"O-XANES\": None\n",
    "    } for qm9id, dat in data.items()\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f1f380-e778-4178-b68e-8f4740695439",
   "metadata": {},
   "outputs": [],
   "source": [
    "for qm9id, datum in tqdm(data.items()):\n",
    "    smiles = datum[\"smiles\"]\n",
    "    lookup_table_data[smiles][\"canon_smiles\"] = datum[\"canon_smiles\"]\n",
    "    molecule = Molecule.from_dict(datum[\"molecule\"])\n",
    "    atoms = [site.specie.symbol for site in molecule]\n",
    "    \n",
    "    for central_atom in [\"C\", \"N\", \"O\"]:\n",
    "        central_atom_indexes = [ii for ii, atom_type in enumerate(atoms) if atom_type == central_atom]\n",
    "        if len(central_atom_indexes) == 0:\n",
    "            continue\n",
    "\n",
    "        # Each of these comes with a spectrum\n",
    "        molecular_spectrum = []\n",
    "        molecule_valid = True\n",
    "        for ii in central_atom_indexes:\n",
    "            key = f\"{ii}_{central_atom}\"\n",
    "            \n",
    "            # Get the spectrum corresponding to it\n",
    "            s = np.array(datum[\"xanes\"][key][\"spectrum\"])\n",
    "            \n",
    "            try:\n",
    "                spline = InterpolatedUnivariateSpline(s[:, 0], s[:, 3])\n",
    "            except IndexError:\n",
    "                molecule_valid = False\n",
    "                break\n",
    "                \n",
    "            res = spline(grids[central_atom])\n",
    "            \n",
    "            # Oxygen and Carbon screening condition for unphysical/outlier results\n",
    "            if CENTRAL_ATOM == \"O\":\n",
    "                if np.any(res[:10] > 5.0):\n",
    "                    molecule_valid = False\n",
    "                    break\n",
    "            elif CENTRAL_ATOM == \"C\":\n",
    "                if np.any(res[:35] > 10.0):\n",
    "                    molecule_valid = False\n",
    "                    break\n",
    "            \n",
    "            molecular_spectrum.append(res)\n",
    "            \n",
    "        if molecule_valid:\n",
    "            molecular_spectrum = np.array(molecular_spectrum).mean(axis=0)\n",
    "            molecular_spectrum[molecular_spectrum < 0.0] = 0.0\n",
    "            lookup_table_data[smiles][f\"{central_atom}-XANES\"] = molecular_spectrum.tolist()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8637f22-c377-4fff-820c-aa5b9f457eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_lookup_table_data = {\n",
    "    \"data\": lookup_table_data,\n",
    "    \"grids\": grids\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "897a21c2-b858-44a0-8d1d-fbee6ea7665e",
   "metadata": {},
   "outputs": [],
   "source": [
    "nones = {\"C\": 0, \"N\": 0, \"O\": 0}\n",
    "for central_atom in [\"C\", \"N\", \"O\"]:\n",
    "    for smiles, datum in final_lookup_table_data[\"data\"].items():\n",
    "        if datum[f\"{central_atom}-XANES\"] is None:\n",
    "            nones[central_atom] += 1\n",
    "print(nones)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "726807bc-b89d-4e82-96cc-b7720e5c7171",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(final_lookup_table_data, open(\"qm9_molecule_xanes.pkl\", \"wb\"), protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db585587-1c60-493f-83ad-224660380371",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_lookup_table_data[\"data\"][\"CCCC\"].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d962698d-48ab-4d34-b72f-93486fed9a3a",
   "metadata": {},
   "source": [
    "## Construct for each absorbing atom type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a9a0f15-d710-45f0-b18e-6d279271c335",
   "metadata": {},
   "outputs": [],
   "source": [
    "CENTRAL_ATOM = \"C\"   # C, N or O\n",
    "grid = grids[CENTRAL_ATOM]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85d080d7-92b2-4335-806d-2acf51ce1357",
   "metadata": {},
   "outputs": [],
   "source": [
    "species = [\"H\", \"C\", \"O\", \"N\", \"F\"]\n",
    "rcut = 6.0\n",
    "g2_params = [[1.0, 0], [0.1, 0], [0.01, 0]]\n",
    "g4_params=[\n",
    "    [0.001, 1.0, -1.0],\n",
    "    [0.001, 2.0, -1.0],\n",
    "    [0.001, 4.0, -1.0],\n",
    "    [0.01, 1.0, -1.0],\n",
    "    [0.01, 2.0, -1.0],\n",
    "    [0.01, 4.0, -1.0],\n",
    "    [0.1, 1.0, -1.0],\n",
    "    [0.1, 2.0, -1.0],\n",
    "    [0.1, 3.0, -1.0]\n",
    "]  # aenet paper\n",
    "acsf = ACSF(\n",
    "    species=species,\n",
    "    rcut=rcut,\n",
    "    g2_params=g2_params,\n",
    "    g4_params=g4_params\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c8e26ca-83cc-46d8-888e-9c0654b40c9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase import Atom, Atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ea4ffa3-ac56-4bdd-8128-f191707aab86",
   "metadata": {},
   "outputs": [],
   "source": [
    "origin_smiles = []\n",
    "molecule_site_pairs = []\n",
    "acsf_array = []\n",
    "spectra = []\n",
    "\n",
    "# cc = 0\n",
    "for qm9id, datum in tqdm(data.items()):\n",
    "    molecule = Molecule.from_dict(datum[\"molecule\"])\n",
    "    atoms = []\n",
    "    central_atom_indexes = []\n",
    "    for ii, site in enumerate(molecule):\n",
    "        atom = Atom(site.specie.symbol, site.coords)\n",
    "        atoms.append(atom)\n",
    "        if site.specie.symbol == CENTRAL_ATOM:\n",
    "            central_atom_indexes.append(ii)\n",
    "    atoms = Atoms(atoms)\n",
    "    \n",
    "    if len(central_atom_indexes) == 0:\n",
    "        continue\n",
    "\n",
    "    tmp_acsf = acsf.create(atoms, positions=central_atom_indexes)\n",
    "    \n",
    "    for idx, ii in enumerate(central_atom_indexes):\n",
    "        key = f\"{ii}_{CENTRAL_ATOM}\"\n",
    "        s = np.array(datum[\"xanes\"][key][\"spectrum\"])\n",
    "\n",
    "        try:\n",
    "            spline = InterpolatedUnivariateSpline(s[:, 0], s[:, 3])\n",
    "        except IndexError:\n",
    "            continue\n",
    "            \n",
    "        res = spline(grid)\n",
    "        \n",
    "        # Oxygen and Carbon screening condition for unphysical/outlier results\n",
    "        if CENTRAL_ATOM == \"O\":\n",
    "            if np.any(res[:10] > 5.0):\n",
    "                continue\n",
    "        elif CENTRAL_ATOM == \"C\":\n",
    "            if np.any(res[:35] > 10.0):\n",
    "                continue\n",
    "        \n",
    "        origin_smiles.append(datum[\"smiles\"])\n",
    "        spectra.append(res)\n",
    "        acsf_array.append(tmp_acsf[idx])\n",
    "        molecule_site_pairs.append(f\"{qm9id}_{ii}\")\n",
    "    \n",
    "#     cc += 1\n",
    "    \n",
    "#     if cc > 20:\n",
    "#         break\n",
    "\n",
    "acsf_array = np.array(acsf_array)\n",
    "spectra = np.array(spectra)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fded89fa-e0db-4f6e-a68a-302fd68ca39c",
   "metadata": {},
   "source": [
    "And finally save to disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41edd69d-bc8e-4843-8d2b-e5584fcd2253",
   "metadata": {},
   "outputs": [],
   "source": [
    "now = datetime.now().strftime(\"%y%m%d\")\n",
    "fname = f\"../data/qm9/XANES-{now}-ACSF-{CENTRAL_ATOM}.pkl\"\n",
    "print(fname)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2166fdd7-8b79-415c-9d69-5955678e0883",
   "metadata": {},
   "source": [
    "We take the convention that `\"x\"` is the input and `\"y\"` is the output. These are the only two required keys for the ML pipeline. The rest is considered metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab3da61-f077-4a3a-83d9-311b55b52fa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(\n",
    "    {\"grid\": grid, \"y\": spectra, \"x\": acsf_array, \"names\": molecule_site_pairs, \"origin_smiles\": origin_smiles},\n",
    "    open(fname, \"wb\"),\n",
    "    protocol=pickle.HIGHEST_PROTOCOL\n",
    ")"
   ]
  }
 ],
 "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.12"
  },
  "toc-autonumbering": true,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}