docs/example/azoswitch/1_data.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1: Molecules, Structural Representations, and Training Sets\n",
    "\n",
    "* Time to run the cells: ~ 1 minute\n",
    "\n",
    "First thing's first, in `settings.py`, set the `RHOLEARN_DIR` to be the absolute\n",
    "path of `.../rholearn/` on your local machine. Inspect the other options set in\n",
    "this file. The settings relevant to this notebook are the `RASCAL_HYPERS` and\n",
    "`DATA_SETTINGS`.\n",
    "\n",
    "`RASCAL_HYPERS` sets the hyperparameters used to generate the $\\lambda$-SOAP\n",
    "structural representation. `DATA_SETTINGS` contains settings for performing a\n",
    "train-test(-validation) split, partitioning the data ready for learning\n",
    "exercises, and where this data should be written.\n",
    "\n",
    "Provided you have set the correct `RHOLEARN_DIR`, these settings should allow\n",
    "these notebook tutorials to run out of the box.\n",
    "\n",
    "First, import all the necessary packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Useful standard and scientific ML libraries\n",
    "import os\n",
    "import ase.io\n",
    "import numpy as np\n",
    "\n",
    "# M-Stack packages\n",
    "import chemiscope\n",
    "import equistore\n",
    "from equistore import Labels\n",
    "\n",
    "from rholearn import io, features, pretraining, utils\n",
    "from settings import RASCAL_HYPERS, DATA_SETTINGS"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference Data\n",
    "\n",
    "### Electron Densities"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data used here is a 10-molecule subset of a largest dataset of azoswitch\n",
    "molecules used in the electron density learning of excited state properties. You\n",
    "can read the paper at __\"Learning the Exciton Properties of Azo-dyes\"__, J.\n",
    "Phys. Chem. Lett. 2021, 12, 25, 5957–5962. DOI:\n",
    "[10.1021/acs.jpclett.1c01425](https://doi.org/10.1021/acs.jpclett.1c01425). \n",
    "\n",
    "For the purposes of this workflow we are focussing on predicting only the\n",
    "ground-state electron density, but can easily be extended to first- and\n",
    "second-excited state hole and particle densities, for which there is reference\n",
    "QM data at the above source.\n",
    "\n",
    "All the data needed to run this proof-of-concept workflow is shipped in the\n",
    "GitHub repo, stored in the ``rho_learn/docs/example/azoswitch/data/`` directory.\n",
    "Inspect this directory. There is a file called ``molecule_list.dat`` containing the\n",
    "filenames of 10 structures, a subfolder ``xyz/`` containing these ``.xyz``\n",
    "files, a folder containing some QM-calculated Coulomb repulsion matrices, and\n",
    "the QM-calculated (i.e. reference) ground state electron density coefficients of\n",
    "the moelcules included in the training set.\n",
    "\n",
    "Both the Coulomb matrices and electron density are stored in equistore TensorMap\n",
    "format. Let's load and inspect the structure of the electron density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = equistore.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"e_densities.npz\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorMaps are main object users interact with when using equistore, storing in\n",
    "principle any kind of data useful in atomistic simulations and their associated\n",
    "metadata. \n",
    "\n",
    "A TensorMap is a collection of TensorBlocks, each of which is indexed by a key\n",
    "and contains atomistic data on a subset of a system of interest. In our case,\n",
    "the electron density TensorMap has blocks for each combination of spherical\n",
    "harmonic channel, $l$, and chemical species. \n",
    "\n",
    "Run the cell below. Notice how the $l$ values run from 0 -> 5 (inclusive) and\n",
    "the chemical species (or 'species_center') span values 1, 6, 7, 8, 16, for\n",
    "elements H, C, N, O, S respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output.keys"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at a specific block. TensorBlock contain three axis: the first is\n",
    "a single dimension, the samples. The last is also a single axis, the properties.\n",
    "And all other intermediate dimensions are the components. In general, samples\n",
    "are used to describe what we are representing, i.e. atomic environments in a\n",
    "given structure, and properties are used to describe how we are representing it.\n",
    "\n",
    "In this example, a set of coefficients for the expansion of the electron density\n",
    "on a set of basis functions are given as the learning targets and therefore the\n",
    "data that appears in the TensorMaps. For a given structure, $A$\n",
    "\n",
    "$ \\rho_A (x) = \\sum_{inlm} c^i_{nlm} \\phi_{nlm}(x - r_i)$\n",
    "\n",
    "where $c^i_{nlm}$ are the expansion coefficients, $\\phi$ the basis functions.\n",
    "$i$ is an atomic index for the atoms in a molecule, $n$ the radial index, and\n",
    "$l$ and $m$ the spherical harmonics indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output.block(0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The samples contain 'structure' (i.e. $A$ in the equation above) and 'center' ($i$) indices. The\n",
    "components contains 'spherical_harmonics_m' ($m$) indices, and the properties\n",
    "contains 'n' (i.e. radial channel $n$) indices. Remember from above that the\n",
    "keys of the TensorMap store the sparse indices for 'spherical_harmonics_l' (i.e.\n",
    "$l$) as well as 'species_center' - the latter because often different basis\n",
    "functions are used for different chemical species."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coulomb Metrics"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each structure in the training set, a Coulomb repulsion metric can be\n",
    "calculated betweeen pairs of basis functions indexed by ${n_1l_1m_1}$ and\n",
    "${n_2l_2m_2}$. The provided Coulomb matrices contains these repulsions, measured\n",
    "in Hartree units of energy.\n",
    "\n",
    "These metrics will be used to define a physically-inspired loss\n",
    "function used in model training (in the second example notebook).\n",
    "\n",
    "Because these matrices are quite large, they had to be split up in order to be\n",
    "stored on GitHub. Run the cell below to recombine them, and observe the keys.\n",
    "Notice how each block is indexed by a pair of $l$ values and chemical species now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azoswitch_utils import recombine_coulomb_metrics\n",
    "\n",
    "coulomb_metrics = recombine_coulomb_metrics(DATA_SETTINGS[\"data_dir\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just view the first 10 keys (as there are > 600 of them)\n",
    "coulomb_metrics.keys[:10]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now inspect a single block. Samples monitor a single structure index, and the 2\n",
    "atomic center indices the basis functions belong to. Note only a single\n",
    "structure index is present here because it doesn't make sense to calculate\n",
    "repulsion between atoms in different structures. The components index the $m$\n",
    "value for the 2 basis functions, and properties indexes the $n$ values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coulomb_metrics.block(0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Structural Descriptors\n",
    "\n",
    "Now we can build a $\\lambda$-SOAP structural representation of the input data,\n",
    "using only the ``.xyz`` files. First, we load the filenames from\n",
    "``molecule_list.dat``. The order of the filenames as listed dictates their structure\n",
    "index, of which all will run from 0 -> 9."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the filenames from molecule_list.dat \n",
    "with open(os.path.join(DATA_SETTINGS[\"data_dir\"], \"molecule_list.dat\"), \"r\") as molecule_list:\n",
    "    xyz_files = molecule_list.read().splitlines()\n",
    "xyz_files"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of these ``.xyz`` files can be read into an ASE object, or 'frame', and\n",
    "these frames can be visualized with chemiscope. Use the slider to have a look at\n",
    "each molecule in turn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read xyz structures into ASE frames\n",
    "frames = [ase.io.read(os.path.join(DATA_SETTINGS[\"data_dir\"], \"xyz\", f)) for f in xyz_files]\n",
    "\n",
    "# Display molecules with chemiscope\n",
    "chemiscope.show(\n",
    "    frames,\n",
    "    properties={\n",
    "        \"Number of atoms\": [f.get_global_number_of_atoms() for f in frames],\n",
    "        \"Molecular mass / u\": [np.sum(f.get_masses()) for f in frames],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the unique chemical species present in the dataset\n",
    "unique_species = list(set([specie for f in frames for specie in f.get_atomic_numbers()]))\n",
    "unique_species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute lambda-SOAP: uses rascaline to compute a SphericalExpansion (~ 15 secs)\n",
    "input = features.lambda_soap_vector(\n",
    "    frames, RASCAL_HYPERS, neighbor_species=unique_species, even_parity_only=True\n",
    ")\n",
    "# Drop the block for l=5, Hydrogen as this isn't included in the output electron density\n",
    "input = equistore.drop_blocks(input, keys=Labels(input.keys.names, np.array([[5, 1]])))\n",
    "\n",
    "# Load the output data (i.e. electron density)\n",
    "output = equistore.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"e_densities.npz\"))\n",
    "\n",
    "# Check that the metadata of input and output match along the samples and components axes\n",
    "assert equistore.equal_metadata(input, output, check=[\"samples\", \"components\"])\n",
    "\n",
    "# Save lambda-SOAP descriptor to file\n",
    "equistore.save(os.path.join(DATA_SETTINGS[\"data_dir\"], \"lambda_soap.npz\"), input)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform data partitioning"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input and output data has been defined, cleaned, and checked for metadata\n",
    "consistency. Now we need to perform a train-test-validation split and, in order\n",
    "to perform a learning exercise, create some subsets of the training data.\n",
    "\n",
    "In the `DATA_SETTINGS` dict of `settings.py` are the options used to perform\n",
    "this data partitioning, that we will provide to the function ``partition_data``.\n",
    "\n",
    "* `axis` controls the TensorMap axis the train-test split should be performed\n",
    "  along. As we want to split our data by structure, we specify `axis=\"samples\"`.\n",
    "* `names` dictates the names of the samples we want to split according to.\n",
    "  Again, we want to split by structure here, so set `names=\"structures\"`.\n",
    "* `n_groups` is how many groups to split the data into. We want to perform a\n",
    "  train-test-validation split, so specify `n_groups=3`.\n",
    "* `group_sizes` controls the number of our named splitting index (in this case\n",
    "  the structures along the samples axis) in each group. We have a dataset of\n",
    "  size 10, and want 7, 2, and 1 structure(s) to be in teh train, test, and\n",
    "  validation sets respectively, so set `group_sizes=[7, 2, 1]`. We could also\n",
    "  pass relative sizes as floats, for instance `group_sizes=[0.7, 0.2, 0.1]`.\n",
    "* `seed` defines the numpy random seed used for shuffling the structure indices\n",
    "  before splitting. Passing this as none gives no shuffling.\n",
    "* `n_exercises` specifies how many learning exercises should be performed, and\n",
    "  thus how many top-level directories with partitioned data should be created.\n",
    "  For each exercise, the data is shuffled differently, leading to different\n",
    "  train, test, and validation data.\n",
    "* `n_subsets` controls how many subsets of the training data should be creted\n",
    "  for each exercise. The size of the subsets, relative to the total number of\n",
    "  training structures, are equally spaced along a log (base 10) scale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Runtime approx 15 seconds\n",
    "pretraining.partition_data(\n",
    "    input_path=os.path.join(DATA_SETTINGS[\"data_dir\"], \"lambda_soap.npz\"),\n",
    "    output_path=os.path.join(DATA_SETTINGS[\"data_dir\"], \"e_densities.npz\"),\n",
    "    data_settings=DATA_SETTINGS,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect how the data was partitioned. In the \"partitions\" folder, a numpy\n",
    "array called \"subset_sizes.npy\" was saved. This stores the sizes (i.e. number of\n",
    "training structures) of each of the training subsets.\n",
    "\n",
    "You can see that, of the 7 structures that we designated as the the total\n",
    "training set, 2, 4, and 6 structures were assigned to each of the training\n",
    "subsets to be used in a learning exercise (provided the random `seed=10` in\n",
    "`DATA_SETTINGS`). While these seem evenly spaced in linear space, in practice\n",
    "the ``partition_data`` function ensures that the sizes of training subsets are\n",
    "evenly spaced along a *log* (base ``e``) scale, to the nearest integer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"subset_sizes_train.npy\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each of the 2 learning exercises, the training structures indices were\n",
    "shuffled before subsets were created. Let's check this by printing the ordered\n",
    "structure indices from which the training set was partitioned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"train structure idxs:\")\n",
    "print(\"exercise 0: \", np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"exercise_0\", \"structure_idxs_train.npy\")))\n",
    "print(\"exercise 1: \", np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"exercise_1\", \"structure_idxs_train.npy\")))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the 7 structure indices are equivalent across both of the lists, the order is\n",
    "different. That means, for instance, when the first subset of size 2 is created,\n",
    "structures 2 and 5 will be present in the training set for exercise 0,\n",
    "and structures 1 and 9 for exercise 1.\n",
    "\n",
    "Just as a sanity check, let's print the test and validation structure indices.\n",
    "We see indices 6, 0, and 8 returned, none of which are present in the training\n",
    "indices above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"test structure idxs: \", np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"structure_idxs_test.npy\")))\n",
    "print(\"val structure idxs: \", np.load(os.path.join(DATA_SETTINGS[\"data_dir\"], \"structure_idxs_val.npy\")))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the data has been partitioned, we are ready to move on to building and\n",
    "training models."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rho",
   "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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "576c71a426691bc103e620abf31b98f592c88b3903fdf6bf41ae71c4b8043fe6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}