docs/example/azoswitch/3_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3: Analyse Results and Visualize Predictions\n",
    "\n",
    "* Time to run the cells: ~ 5 mins\n",
    "\n",
    "First thing's first; set the absolute path of the ``rho_learn`` directory on\n",
    "your local machine, for instance:\n",
    "\n",
    "``RHOLEARN_DIR = \"/Users/joe.abbott/Documents/phd/code/qml/rho_learn/\"``\n",
    "\n",
    "In this notebook we will plot a few analysis plots using data outputted from the\n",
    "model training procedure. \n",
    "\n",
    "Then we will load our validation structure, make a prediction on it, and\n",
    "visualize the predicted, target, and delta electron densities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import ase.io\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import equistore\n",
    "import qstack\n",
    "\n",
    "import rholearn.io\n",
    "from rholearn import analysis, io, plots, utils\n",
    "from settings import RASCAL_HYPERS, DATA_SETTINGS, ML_SETTINGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create directory to save plots and visualizations\n",
    "plot_dir = os.path.join(ML_SETTINGS[\"run_dir\"], \"plots\")\n",
    "rholearn.io.check_or_create_dir(plot_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Analysis\n",
    "\n",
    "Based on the learning exercise run in the previous notebook, we can produce some\n",
    "analysis plots.\n",
    "\n",
    "First, plot a log-log figure of the train and test losses against epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the range of exercises and subsets to compile data for\n",
    "exercises = [0]\n",
    "subsets = [0, 1, 2]\n",
    "\n",
    "# Compile data\n",
    "train, test = analysis.compile_loss_data(ML_SETTINGS[\"run_dir\"], exercises, subsets)\n",
    "mean_train = analysis.average_losses(train)\n",
    "mean_test = analysis.average_losses(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot losses\n",
    "fig, ax = plt.subplots(1, 1, sharey=True)\n",
    "for subset in mean_train.keys():\n",
    "    loss_train = mean_train[subset]\n",
    "    loss_test = mean_test[subset]\n",
    "    ax.loglog(loss_train, label=f\"linear, train, subset{subset}\")\n",
    "    ax.loglog(loss_test, label=f\"linear, test, subset{subset}\")\n",
    "ax.set_ylabel(\"Coulomb Loss / Ha\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, plot the learning curve for the learning exercises."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Log-log learning curve plot of loss vs training set size\n",
    "point = \"final\"  # take the final epoch loss, as opposed to \"best\" (i.e. lowest)\n",
    "fig, ax = plots.learning_curve(\n",
    "    [mean_train, mean_test],\n",
    "    np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"subset_sizes_train.npy\")),\n",
    "    point=point,\n",
    ")\n",
    "\n",
    "# Format\n",
    "fig.tight_layout()\n",
    "ax.set_ylabel(point + r\" loss\")\n",
    "ax.legend(labels=[\"train\", \"test\"])\n",
    "\n",
    "# Save\n",
    "plots.save_fig_mpltex(fig, os.path.join(plot_dir, \"learning_curve\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation Structure\n",
    "\n",
    "Now we can load the validation structure we created earlier, make a prediction\n",
    "on it using our model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the lambda-SOAP TensorMap for the validation structure and its QM (i.e. the\n",
    "target) electron density. Then, load the best model and make a prediction on the\n",
    "validation structure. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the lambda-SOAP TensorMaps for the input and output validation structure\n",
    "in_val = io.load_tensormap_to_torch(\n",
    "    os.path.join(DATA_SETTINGS[\"data_dir\"], \"in_val.npz\"), **ML_SETTINGS[\"torch\"]\n",
    ")\n",
    "out_val = equistore.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"out_val.npz\"))\n",
    "\n",
    "# Load the model from the largest training subset and latest epoch. In principle\n",
    "# we could average the weights from multiple models from different exercises,\n",
    "# but here we'll just load one.\n",
    "model = io.load_torch_object(\n",
    "    os.path.join(ML_SETTINGS[\"run_dir\"], \"exercise_0\", \"subset_2\", \"epoch_10\", \"model.pt\"),\n",
    "    device=ML_SETTINGS[\"torch\"][\"device\"],\n",
    "    torch_obj_str=\"model\",\n",
    ")\n",
    "\n",
    "# Make a prediction on the validation structure (invariants standardized)\n",
    "out_val_pred = model(in_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a parity plot of the target electron density coefficients against\n",
    "predicted, coloured by $\\lambda$ value $\\in [0, ..., 5]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can also color by \"species_center\"\n",
    "fig, ax = plots.parity_plot(\n",
    "    target=utils.standardize_invariants(\n",
    "        utils.tensor_to_numpy(out_val),\n",
    "        equistore.load(\n",
    "            os.path.join(\n",
    "                DATA_SETTINGS[\"data_dir\"], \"exercise_0\", \"subset_2\", \"inv_means.npz\"\n",
    "            )\n",
    "        ),\n",
    "    ),\n",
    "    predicted=utils.tensor_to_numpy(out_val_pred),\n",
    "    color_by=\"spherical_harmonics_l\",\n",
    ")\n",
    "lim = [-0.2, 0.2]\n",
    "ax.set_xlim(lim)\n",
    "ax.set_ylim(lim)\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_xlabel(\"target density coefficient value\")\n",
    "ax.set_ylabel(\"predicted density coefficient value\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also colour the plot according to the elemental species, where index 1\n",
    "indicates Hydrogen, 6 Carbon, 7 Nitrogen, 8 Oxygen, 16 Sulfur"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization of the Electron Density\n",
    "\n",
    "Having made a prediction on the validation structure using the linear model, we\n",
    "can convert the TensorMap into a cube file format using Q-Stack, and visualize\n",
    "the electron density.\n",
    "\n",
    "First, build a molecule object using the Q-Stack ``compound`` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the validation structure index from file\n",
    "val_structure_idx = np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"structure_idxs_val.npy\"))[0][0]\n",
    "\n",
    "# Load the xyz file of the validation structure\n",
    "with open(os.path.join(DATA_SETTINGS[\"data_dir\"], \"molecule_list.dat\"), \"r\") as molecule_list:\n",
    "    val_xyz = molecule_list.read().splitlines()[val_structure_idx]\n",
    "\n",
    "# Build a molecule object using Q-Stack. The basis used in density fitting for\n",
    "# this particular dataset was 'ccpvqz jkfit'\n",
    "molecule = qstack.compound.xyz_to_mol(\n",
    "    os.path.join(DATA_SETTINGS[\"data_dir\"], \"xyz\", val_xyz), basis=\"ccpvqz jkfit\"\n",
    ")\n",
    "\n",
    "# Unstandardize the predicted density coefficients\n",
    "out_val_pred = utils.standardize_invariants(\n",
    "    utils.tensor_to_numpy(out_val_pred),\n",
    "    equistore.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"exercise_0\", \"subset_2\", \"inv_means.npz\")),\n",
    "    reverse=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do some renaming of TensorMaps to make naming conventions match and build a\n",
    "delta electron density, given by subtraction of coefficients between the\n",
    "predicted and target (i.e. QM) electron densities. This will help to visualize\n",
    "areas of the molecule that the ML model isn't performing so well at performing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rename tensors to fit Q-Stack naming convention\n",
    "new_key_names = [\"spherical_harmonics_l\", \"element\"]\n",
    "out_val = utils.rename_tensor(\n",
    "    utils.drop_metadata_name(out_val, axis=\"samples\", name=\"structure\"), keys_names=new_key_names\n",
    ")\n",
    "out_val_pred = utils.rename_tensor(\n",
    "    utils.drop_metadata_name(out_val_pred, axis=\"samples\", name=\"structure\"),\n",
    "    keys_names=new_key_names,\n",
    ")\n",
    "\n",
    "# Generate a delta electron density: QM - ML. This helps to visualize where the\n",
    "# ML model isn't working so well\n",
    "out_val_delta = equistore.abs(equistore.subtract(utils.tensor_to_numpy(out_val_pred), utils.tensor_to_numpy(out_val)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vectorize the density coefficients and convert them to cube file format using Q-Stack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qstack import equio\n",
    "from qstack.fields import density2file\n",
    "\n",
    "# Vectorize the coefficients from each of the TensorMaps\n",
    "vect_coeffs_target = equio.tensormap_to_vector(molecule, out_val)\n",
    "vect_coeffs_input = equio.tensormap_to_vector(molecule, out_val_pred)\n",
    "vect_coeffs_delta = equio.tensormap_to_vector(molecule, out_val_delta)\n",
    "\n",
    "# Define a number of grid points to represent the electon density on\n",
    "n = 80\n",
    "\n",
    "# Convert the basis function coefficients to a cube file\n",
    "for (coeffs, filename) in [\n",
    "    (vect_coeffs_target, \"out_val.cube\"),\n",
    "    (vect_coeffs_input, \"out_val_pred.cube\"),\n",
    "    (vect_coeffs_delta, \"out_val_delta.cube\"),\n",
    "]:\n",
    "    density2file.coeffs_to_cube(\n",
    "        molecule,\n",
    "        coeffs,\n",
    "        os.path.join(plot_dir, filename),\n",
    "        nx=n,\n",
    "        ny=n,\n",
    "        nz=n,\n",
    "        resolution=None,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use Py3Dmol to visualize the cube files for the target, predicted, and delta\n",
    "electron densities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import py3Dmol\n",
    "\n",
    "# Visualize the target density\n",
    "for filename in [\"out_val.cube\", \"out_val_pred.cube\", \"out_val_delta.cube\"]:\n",
    "    print(filename)\n",
    "    v = py3Dmol.view()\n",
    "    v.addModelsAsFrames(open(os.path.join(plot_dir, filename), \"r\").read(), \"cube\")\n",
    "    v.setStyle({\"stick\": {}})\n",
    "    v.addVolumetricData(\n",
    "        open(os.path.join(plot_dir, filename), \"r\").read(),\n",
    "        \"cube\",\n",
    "        {\"isoval\": 0.01, \"color\": \"blue\", \"opacity\": 0.8},\n",
    "    )\n",
    "    v.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.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "576c71a426691bc103e620abf31b98f592c88b3903fdf6bf41ae71c4b8043fe6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}