phononDoS_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/ninarina12/phononDoS_tutorial/blob/main/phononDoS_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_v7BmDsQgutC"
   },
   "source": [
    "## Tutorial | Predicting phonon DoS with `e3nn`\n",
    "### Getting started\n",
    "*   Go to Runtime > Change runtime type, and select GPU.\n",
    "*   Clone the GitHub repository to access the tutorial files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xd4k295fOzl7"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/ninarina12/phononDoS_tutorial.git\n",
    "%cd phononDoS_tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cumoRx6GhLKW"
   },
   "source": [
    "*   Install some relevant packages (should take < 1 minute).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JobVnat5O5fT"
   },
   "outputs": [],
   "source": [
    "!pip install ase e3nn\n",
    "!pip install torch-scatter torch-cluster torch-sparse torch-spline-conv -f https://pytorch-geometric.com/whl/torch-$(python -c \"import torch; print(torch.__version__)\").html\n",
    "!pip install torch-geometric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iKp3KG4HhZqh"
   },
   "source": [
    "### Tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l45yb_JnQ0eD"
   },
   "outputs": [],
   "source": [
    "# model\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric as tg\n",
    "import torch_scatter\n",
    "import e3nn\n",
    "from e3nn import o3\n",
    "from typing import Dict, Union\n",
    "\n",
    "# crystal structure data\n",
    "from ase import Atom, Atoms\n",
    "from ase.neighborlist import neighbor_list\n",
    "from ase.visualize.plot import plot_atoms\n",
    "\n",
    "# data pre-processing and visualization\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "from IPython.display import HTML\n",
    "\n",
    "# utilities\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "from utils.utils_data import (load_data, train_valid_test_split, plot_example, plot_predictions, plot_partials,\n",
    "                              palette, colors, cmap)\n",
    "from utils.utils_model import Network, visualize_layers, train\n",
    "from utils.utils_plot import plotly_surface, plot_orbitals, get_middle_feats\n",
    "\n",
    "bar_format = '{l_bar}{bar:10}{r_bar}{bar:-10b}'\n",
    "default_dtype = torch.float64\n",
    "torch.set_default_dtype(default_dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I1Z-W3T75bfM"
   },
   "source": [
    "### Data provenance\n",
    "We train our model using the database of Density Functional Perturbation Theory (DFPT)-calculated phonon densities of states (DoS), containing approximately 1,500 crystalline solids [[Petretto et al. 2018]](https://doi.org/10.1038/sdata.2018.65)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zTXZExsvWWey"
   },
   "outputs": [],
   "source": [
    "# load data\n",
    "df, species = load_data('data/data.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nPWFGltW54MI"
   },
   "source": [
    "### Data structures\n",
    "Crystal structures are represented as [ASE](https://wiki.fysik.dtu.dk/ase/ase/atoms.html?highlight=atoms#the-atoms-object) (Atomic Simulation Environment) `Atoms` objects, which store the atomic species and positions of each atom in the unit cell, as well as the lattice vectors of the unit cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dMnXRZ8RhCGr"
   },
   "outputs": [],
   "source": [
    "# plot an example structure\n",
    "i = 12 # structure index in dataframe\n",
    "\n",
    "struct = df.iloc[i]['structure']\n",
    "symbols = np.unique(list(struct.symbols))\n",
    "z = dict(zip(symbols, range(len(symbols))))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6,5))\n",
    "norm = plt.Normalize(vmin=0, vmax=len(symbols)-1)\n",
    "color = [mpl.colors.to_hex(k) for k in cmap(norm([z[j] for j in list(struct.symbols)]))]\n",
    "plot_atoms(struct, ax, radii=0.25, colors=color, rotation=('0x,90y,0z'))\n",
    "\n",
    "ax.set_xlabel(r'$x_1\\ (\\AA)$')\n",
    "ax.set_ylabel(r'$x_2\\ (\\AA)$');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1GymBzWq6ThN"
   },
   "outputs": [],
   "source": [
    "# lattice parameter statistics\n",
    "def get_lattice_parameters(df):\n",
    "    a = []\n",
    "    for entry in df.itertuples():\n",
    "        a.append(entry.structure.cell.cellpar()[:3])\n",
    "    return np.stack(a)\n",
    "\n",
    "a = get_lattice_parameters(df)\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(5,4))\n",
    "b = 0.\n",
    "bins = 50\n",
    "for d, c, n in zip(['a', 'b', 'c'], colors.values(), [a[:,0], a[:,1], a[:,2]]):\n",
    "    color = [int(c.lstrip('#')[i:i+2], 16)/255. for i in (0,2,4)]\n",
    "    y, bins, _, = ax.hist(n, bins=bins, fc=color+[0.7], ec=color, bottom=b, label=d)\n",
    "    b += y\n",
    "ax.set_xlabel('lattice parameter')\n",
    "ax.set_ylabel('number of examples')\n",
    "ax.legend(frameon=False)\n",
    "\n",
    "print('average lattice parameter (a/b/c):', a[:,0].mean(), '/', a[:,1].mean(), '/', a[:,2].mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6dn1JFrY6ZWz"
   },
   "source": [
    "### Feature representation\n",
    "We construct the inputs to our neural network following the `e3nn` [Documentation](https://docs.e3nn.org/en/latest/guide/periodic_boundary_conditions.html) on handling point inputs with periodic boundary conditions. For a given crystal, each atom in the unit cell is associated with a feature vector that one-hot encodes its atomic mass in the index corresponding to its atomic number. The unit cell of the crystal is encoded as a graph in which two atoms (nodes) are joined by an edge if they are within a cutoff radius `r_max` of one another."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CEjH-a586Z8V"
   },
   "outputs": [],
   "source": [
    "# one-hot encoding atom type and mass\n",
    "type_encoding = {}\n",
    "specie_am = []\n",
    "for Z in tqdm(range(1, 119), bar_format=bar_format):\n",
    "    specie = Atom(Z)\n",
    "    type_encoding[specie.symbol] = Z - 1\n",
    "    specie_am.append(specie.mass)\n",
    "\n",
    "type_onehot = torch.eye(len(type_encoding))\n",
    "am_onehot = torch.diag(torch.tensor(specie_am))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2pEoARfd6eb2"
   },
   "outputs": [],
   "source": [
    "# build data\n",
    "def build_data(entry, type_encoding, type_onehot, r_max=5.):\n",
    "    symbols = list(entry.structure.symbols).copy()\n",
    "    positions = torch.from_numpy(entry.structure.positions.copy())\n",
    "    lattice = torch.from_numpy(entry.structure.cell.array.copy()).unsqueeze(0)\n",
    "\n",
    "    # edge_src and edge_dst are the indices of the central and neighboring atom, respectively\n",
    "    # edge_shift indicates whether the neighbors are in different images or copies of the unit cell\n",
    "    edge_src, edge_dst, edge_shift = neighbor_list(\"ijS\", a=entry.structure, cutoff=r_max, self_interaction=True)\n",
    "    \n",
    "    # compute the relative distances and unit cell shifts from periodic boundaries\n",
    "    edge_batch = positions.new_zeros(positions.shape[0], dtype=torch.long)[torch.from_numpy(edge_src)]\n",
    "    edge_vec = (positions[torch.from_numpy(edge_dst)]\n",
    "                - positions[torch.from_numpy(edge_src)]\n",
    "                + torch.einsum('ni,nij->nj', torch.tensor(edge_shift, dtype=default_dtype), lattice[edge_batch]))\n",
    "\n",
    "    # compute edge lengths (rounded only for plotting purposes)\n",
    "    edge_len = np.around(edge_vec.norm(dim=1).numpy(), decimals=2)\n",
    "    \n",
    "    data = tg.data.Data(\n",
    "        pos=positions, lattice=lattice, symbol=symbols,\n",
    "        x=am_onehot[[type_encoding[specie] for specie in symbols]],   # atomic mass (node feature)\n",
    "        z=type_onehot[[type_encoding[specie] for specie in symbols]], # atom type (node attribute)\n",
    "        edge_index=torch.stack([torch.LongTensor(edge_src), torch.LongTensor(edge_dst)], dim=0),\n",
    "        edge_shift=torch.tensor(edge_shift, dtype=default_dtype),\n",
    "        edge_vec=edge_vec, edge_len=edge_len,\n",
    "        phdos=torch.from_numpy(entry.phdos).unsqueeze(0)\n",
    "    )\n",
    "    \n",
    "    return data\n",
    "\n",
    "r_max = 4. # cutoff radius\n",
    "df['data'] = df.progress_apply(lambda x: build_data(x, type_encoding, type_onehot, r_max), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1AGI3yIx6ho9"
   },
   "outputs": [],
   "source": [
    "i = 12 # structure index in dataframe\n",
    "plot_example(df, i=i, label_edges=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hY_oakVo6j4Q"
   },
   "source": [
    "### Training, validation, and testing datasets\n",
    "Split the data into training, validation, and testing datasets with balanced representation of different elements in each set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1R6i_FUQ6mml"
   },
   "outputs": [],
   "source": [
    "# train/valid/test split\n",
    "idx_train, idx_valid, idx_test = train_valid_test_split(df, species, valid_size=.1, test_size=.1, seed=12, plot=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z20LxeLP61m6"
   },
   "source": [
    "For use with the trained model provided, the indices of the training, validation, and test sets are loaded below. These indices were generated with a specific seed using the above `train_valid_test_split` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NLEaj3-J62Nz"
   },
   "outputs": [],
   "source": [
    "# load train/valid/test indices\n",
    "with open('data/idx_train.txt', 'r') as f: idx_train = [int(i.split('\\n')[0]) for i in f.readlines()]\n",
    "with open('data/idx_valid.txt', 'r') as f: idx_valid = [int(i.split('\\n')[0]) for i in f.readlines()]\n",
    "with open('data/idx_test.txt', 'r') as f: idx_test = [int(i.split('\\n')[0]) for i in f.readlines()]\n",
    "\n",
    "# format dataloaders\n",
    "batch_size = 1\n",
    "dataloader_train = tg.loader.DataLoader(df.iloc[idx_train]['data'].values, batch_size=batch_size, shuffle=True)\n",
    "dataloader_valid = tg.loader.DataLoader(df.iloc[idx_valid]['data'].values, batch_size=batch_size)\n",
    "dataloader_test = tg.loader.DataLoader(df.iloc[idx_test]['data'].values, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blEhP5Bh64s3"
   },
   "outputs": [],
   "source": [
    "# calculate average number of neighbors\n",
    "def get_neighbors(df, idx):\n",
    "    n = []\n",
    "    for entry in df.iloc[idx].itertuples():\n",
    "        N = entry.data.pos.shape[0]\n",
    "        for i in range(N):\n",
    "            n.append(len((entry.data.edge_index[0] == i).nonzero()))\n",
    "    return np.array(n)\n",
    "\n",
    "n_train = get_neighbors(df, idx_train)\n",
    "n_valid = get_neighbors(df, idx_valid)\n",
    "n_test = get_neighbors(df, idx_test)\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(5,4))\n",
    "b = 0.\n",
    "bins = 50\n",
    "for (d, c), n in zip(colors.items(), [n_train, n_valid, n_test]):\n",
    "    color = [int(c.lstrip('#')[i:i+2], 16)/255. for i in (0,2,4)]\n",
    "    y, bins, _, = ax.hist(n, bins=bins, fc=color+[0.7], ec=color, bottom=b, label=d)\n",
    "    b += y\n",
    "ax.set_xlabel('number of neighbors')\n",
    "ax.set_ylabel('number of examples')\n",
    "ax.legend(frameon=False)\n",
    "\n",
    "print('average number of neighbors (train/valid/test):', n_train.mean(), '/', n_valid.mean(), '/', n_test.mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LZtSOwB9Cols"
   },
   "source": [
    "### Network architecture\n",
    "We build a model based on the `Network` described in the `e3nn` [Documentation](https://docs.e3nn.org/en/latest/api/nn/models/gate_points_2101.html), modified to incorporate the periodic boundaries we imposed on the crystal graphs. The network applies equivariant convolutions to each atomic node and finally takes an average over all nodes, normalizing the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X3xO0DoQ69uo"
   },
   "outputs": [],
   "source": [
    "class PeriodicNetwork(Network):\n",
    "    def __init__(self, in_dim, em_dim, **kwargs):            \n",
    "        # override the `reduce_output` keyword to instead perform an averge over atom contributions    \n",
    "        self.pool = False\n",
    "        if kwargs['reduce_output'] == True:\n",
    "            kwargs['reduce_output'] = False\n",
    "            self.pool = True\n",
    "            \n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "        # embed the mass-weighted one-hot encoding\n",
    "        self.em = nn.Linear(in_dim, em_dim)\n",
    "\n",
    "    def forward(self, data: Union[tg.data.Data, Dict[str, torch.Tensor]]) -> torch.Tensor:\n",
    "        data.x = F.relu(self.em(data.x))\n",
    "        data.z = F.relu(self.em(data.z))\n",
    "        output = super().forward(data)\n",
    "        output = torch.relu(output)\n",
    "        \n",
    "        # if pool_nodes was set to True, use scatter_mean to aggregate\n",
    "        if self.pool == True:\n",
    "            output = torch_scatter.scatter_mean(output, data.batch, dim=0)  # take mean over atoms per example\n",
    "        \n",
    "        maxima, _ = torch.max(output, dim=1)\n",
    "        output = output.div(maxima.unsqueeze(1))\n",
    "        \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d904KAn77BBl"
   },
   "outputs": [],
   "source": [
    "out_dim = len(df.iloc[0]['phfreq'])\n",
    "em_dim = 64  \n",
    "\n",
    "model = PeriodicNetwork(\n",
    "    in_dim=118,                            # dimension of one-hot encoding of atom type\n",
    "    em_dim=em_dim,                         # dimension of atom-type embedding\n",
    "    irreps_in=str(em_dim)+\"x0e\",           # em_dim scalars (L=0 and even parity) on each atom to represent atom type\n",
    "    irreps_out=str(out_dim)+\"x0e\",         # out_dim scalars (L=0 and even parity) to output\n",
    "    irreps_node_attr=str(em_dim)+\"x0e\",    # em_dim scalars (L=0 and even parity) on each atom to represent atom type\n",
    "    layers=2,                              # number of nonlinearities (number of convolutions = layers + 1)\n",
    "    mul=32,                                # multiplicity of irreducible representations\n",
    "    lmax=1,                                # maximum order of spherical harmonics\n",
    "    max_radius=r_max,                      # cutoff radius for convolution\n",
    "    num_neighbors=n_train.mean(),          # scaling factor based on the typical number of neighbors\n",
    "    reduce_output=True                     # whether or not to aggregate features of all atoms at the end\n",
    ")\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pIb6-dq87DqK"
   },
   "outputs": [],
   "source": [
    "# visualize tensor products of the model\n",
    "visualize_layers(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-5cuaOFOC_D_"
   },
   "source": [
    "### Training\n",
    "The model is trained using a mean-squared error loss function with an Adam optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ERK8uX1vC_Zy"
   },
   "outputs": [],
   "source": [
    "opt = torch.optim.AdamW(model.parameters(), lr=0.005, weight_decay=0.05)\n",
    "scheduler = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.96)\n",
    "\n",
    "loss_fn = torch.nn.MSELoss()\n",
    "loss_fn_mae = torch.nn.L1Loss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "38dUj1WpDBLh"
   },
   "outputs": [],
   "source": [
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "print('torch device:' , device)\n",
    "\n",
    "run_name = 'model_' + time.strftime(\"%y%m%d\", time.localtime())\n",
    "print(run_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mC7DaaymDLTP"
   },
   "outputs": [],
   "source": [
    "model.pool = True\n",
    "train(model, opt, dataloader_train, dataloader_valid, loss_fn, loss_fn_mae, run_name,\n",
    "      max_iter=1, scheduler=scheduler, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VRCWlA-xDQaK"
   },
   "outputs": [],
   "source": [
    "# load pre-trained model and plot its training history\n",
    "run_name = 'model'\n",
    "\n",
    "history = torch.load(run_name + '.torch', map_location=device)['history']\n",
    "steps = [d['step'] + 1 for d in history]\n",
    "loss_train = [d['train']['loss'] for d in history]\n",
    "loss_valid = [d['valid']['loss'] for d in history]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6,5))\n",
    "ax.plot(steps, loss_train, 'o-', label=\"Training\", color=colors['train'])\n",
    "ax.plot(steps, loss_valid, 'o-', label=\"Validation\", color=colors['valid'])\n",
    "ax.set_xlabel('epochs')\n",
    "ax.set_ylabel('loss')\n",
    "ax.legend(frameon=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5kpYlnpwEFW3"
   },
   "source": [
    "### Results\n",
    "We evaluate our model by visualizing the predicted and true DoS in each error quartile. We further compare the hidden features learned for each node to the partial DoS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p04FoU5dEFr2"
   },
   "outputs": [],
   "source": [
    "# predict on all data\n",
    "model.load_state_dict(torch.load(run_name + '.torch', map_location=device)['state'])\n",
    "model.pool = True\n",
    "\n",
    "dataloader = tg.loader.DataLoader(df['data'].values, batch_size=64)\n",
    "df['mse'] = 0.\n",
    "df['phdos_pred'] = np.empty((len(df), 0)).tolist()\n",
    "\n",
    "model.to(device)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    i0 = 0\n",
    "    for i, d in tqdm(enumerate(dataloader), total=len(dataloader), bar_format=bar_format):\n",
    "        d.to(device)\n",
    "        output = model(d)\n",
    "        loss = F.mse_loss(output, d.phdos, reduction='none').mean(dim=-1).cpu().numpy()\n",
    "        df.loc[i0:i0 + len(d.phdos) - 1, 'phdos_pred'] = [[k] for k in output.cpu().numpy()]\n",
    "        df.loc[i0:i0 + len(d.phdos) - 1, 'mse'] = loss\n",
    "        i0 += len(d.phdos)\n",
    "        \n",
    "df['phdos_pred'] = df['phdos_pred'].map(lambda x: x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zyQH05qpEIoD"
   },
   "outputs": [],
   "source": [
    "plot_predictions(df, idx_train, 'Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lcZPaHO3EKy4"
   },
   "outputs": [],
   "source": [
    "plot_predictions(df, idx_valid, 'Validation')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "akcuwjRrEK2y"
   },
   "outputs": [],
   "source": [
    "plot_predictions(df, idx_test, 'Testing')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l2WWbtwsEQVV"
   },
   "outputs": [],
   "source": [
    "# compare to partial DoS\n",
    "model.load_state_dict(torch.load(run_name + '.torch', map_location=device)['state'])\n",
    "model.pool = False\n",
    "\n",
    "# plot example predicted and true partial dos\n",
    "plot_partials(model, df, idx_train, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vg3iENN5ESa2"
   },
   "source": [
    "### Alloys\n",
    "The current framework extends easily to the representation of alloy structures. As an example, we will predict the phonon DoS of the Mg<sub>3</sub>(Bi,Sb)<sub>2</sub> system, incrementally varying the relative fractions of Bi and Sb. Note that both parent compounds, Mg<sub>3</sub>Sb<sub>2</sub> and Mg<sub>3</sub>Bi<sub>2</sub>, are present in our training data. We will check the validity by comparing the predicted and calculated phonon DoS of Mg<sub>3</sub>Bi<sub>1.5</sub>Sb<sub>0.5</sub>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9qgk-qwcERhz"
   },
   "outputs": [],
   "source": [
    "# load calculated alloy example\n",
    "df_alloy, _ = load_data('data/data_alloy.csv')\n",
    "df_alloy.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YU3ynm06EZLu"
   },
   "outputs": [],
   "source": [
    "# get indices of parent structures\n",
    "idx_Mg3Sb2 = df.loc[df['mp_id'] == 'mp-2646'].index.to_numpy()[0]\n",
    "idx_Mg3Bi2 = df.loc[df['mp_id'] == 'mp-569018'].index.to_numpy()[0]\n",
    "print(f'index of Mg3Sb2: {idx_Mg3Sb2}', f'\\nindex of Mg3Bi2: {idx_Mg3Bi2}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t9Jl5-HLEa-2"
   },
   "outputs": [],
   "source": [
    "# interpolate atomic positions and lattice constants\n",
    "# 2-hot encode the atomic mass, weighted by the fraction of each species\n",
    "data_alloy = []\n",
    "x_Bi = np.linspace(0.01, 0.99, 99)\n",
    "\n",
    "for i, p in tqdm(enumerate(x_Bi), total=len(x_Bi), bar_format=bar_format):\n",
    "    symbols = df['data'][idx_Mg3Bi2].symbol.copy()\n",
    "    positions = torch.lerp(df['data'][idx_Mg3Sb2].pos.clone(), df['data'][idx_Mg3Bi2].pos.clone(), p)\n",
    "    lattice = torch.lerp(df['data'][idx_Mg3Sb2].lattice.clone(), df['data'][idx_Mg3Bi2].lattice.clone(), p)\n",
    "\n",
    "    # edge_src and edge_dst are the indices of the central and neighboring atom, respectively\n",
    "    # edge_shift indicates whether the neighbors are in different images or copies of the unit cell\n",
    "    struct = df.iloc[idx_Mg3Bi2].structure.copy()\n",
    "    struct.positions = positions.numpy().copy()\n",
    "    struct.cell = lattice.numpy().squeeze().copy()\n",
    "    edge_src, edge_dst, edge_shift = neighbor_list(\"ijS\", a=struct, cutoff=r_max, self_interaction=True)\n",
    "    \n",
    "    # compute the relative distances and unit cell shifts from periodic boundaries\n",
    "    edge_batch = positions.new_zeros(positions.shape[0], dtype=torch.long)[torch.from_numpy(edge_src)]\n",
    "    edge_vec = (positions[torch.from_numpy(edge_dst)]\n",
    "                - positions[torch.from_numpy(edge_src)]\n",
    "                + torch.einsum('ni,nij->nj', torch.tensor(edge_shift, dtype=default_dtype), lattice[edge_batch]))\n",
    "\n",
    "    # compute edge lengths (rounded only for plotting purposes)\n",
    "    edge_len = np.around(edge_vec.norm(dim=1).numpy(), decimals=2)\n",
    "    \n",
    "    data_alloy.append(\n",
    "            tg.data.Data(\n",
    "            pos=positions, \n",
    "            lattice=lattice, \n",
    "            symbol=symbols,\n",
    "            x=torch.lerp(df['data'][idx_Mg3Sb2].x, df['data'][idx_Mg3Bi2].x, p),\n",
    "            z=torch.lerp(df['data'][idx_Mg3Sb2].z, df['data'][idx_Mg3Bi2].z, p),\n",
    "            edge_index=torch.stack([torch.LongTensor(edge_src), torch.LongTensor(edge_dst)], dim=0),\n",
    "            edge_shift=torch.tensor(edge_shift, dtype=default_dtype),\n",
    "            edge_vec=edge_vec, edge_len=edge_len,\n",
    "            phdos=df['data'][idx_Mg3Bi2].phdos.clone()\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JU4IhPPlEcvU"
   },
   "outputs": [],
   "source": [
    "# predict on all alloy structures\n",
    "model.load_state_dict(torch.load(run_name + '.torch', map_location=device)['state'])\n",
    "model.pool = True\n",
    "\n",
    "dataloader = tg.loader.DataLoader([df.iloc[idx_Mg3Sb2]['data']] + data_alloy + [df.iloc[idx_Mg3Bi2]['data']],\n",
    "                                  batch_size=32)\n",
    "\n",
    "output = np.zeros((len(data_alloy) + 2, len(df_alloy['phdos'][0])))\n",
    "model.to(device)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    i0 = 0\n",
    "    for i, d in tqdm(enumerate(dataloader), total=len(dataloader), bar_format=bar_format):\n",
    "        d.to(device)\n",
    "        output[i0:i0 + len(d.phdos),:] = model(d).cpu().numpy()\n",
    "        i0 += len(d.phdos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wVkxRig3Ekdb"
   },
   "outputs": [],
   "source": [
    "# plot predictions, and compare with calculated result for selected compound\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12,6), gridspec_kw={'width_ratios': [1,2]})\n",
    "color = cmap(np.linspace(0, 1, len(output)))\n",
    "f = df_alloy['phfreq'][0]\n",
    "\n",
    "# waterfall plot of alloy predictions\n",
    "s = 2./len(x_Bi)\n",
    "for i in range(len(output)):\n",
    "    ax1.plot(f, output[i]/output[i].max() + i*s, c=color[i])\n",
    "ax1.set_yticklabels([])\n",
    "ax1.set_xlabel('$Frequency\\ (cm^{-1})$')\n",
    "ax1.set_ylabel('$Intensity$')\n",
    "\n",
    "sm = mpl.cm.ScalarMappable(cmap=cmap)\n",
    "sm.set_array([])\n",
    "cax = inset_axes(ax1, width=\"40%\", height=\"4%\", loc=3, bbox_to_anchor=(0.5,0.9,1,1), bbox_transform=ax1.transAxes) \n",
    "cbar = fig.colorbar(sm, cax=cax, aspect=16, orientation='horizontal', pad=-0.1)\n",
    "cbar.ax.set_xlabel('$x_{Bi}$', fontsize=16, labelpad=-5)\n",
    "    \n",
    "# comparison to calculation\n",
    "p = x_Bi.tolist().index(0.75)\n",
    "ax2.remove()\n",
    "ax2 = fig.add_subplot(122, projection='3d')\n",
    "\n",
    "# plot calculations\n",
    "ax2.plot(f, [0.75]*len(f), df_alloy['phdos'][0], lw=1.5, c='black', label='Calculated')\n",
    "ax2.plot(f, [0]*len(f), df.iloc[idx_Mg3Sb2]['phdos'], lw=1.5, c='black')\n",
    "ax2.plot(f, [1]*len(f), df.iloc[idx_Mg3Bi2]['phdos'], lw=1.5, c='black')\n",
    "\n",
    "# plot predictions\n",
    "ax2.plot(f, [0.75]*len(f), output[p]/output[p].max(), lw=2, c=palette[1], label='Predicted (alloy)')\n",
    "ax2.plot(f, [0]*len(f), output[0]/output[0].max(), lw=2, c=palette[0], label='Predicted (pure)')\n",
    "ax2.plot(f, [1]*len(f), output[-1]/output[-1].max(), lw=2, c=palette[0])\n",
    "\n",
    "ax2.view_init(elev=20, azim=-50)\n",
    "ax2.w_xaxis.set_pane_color((1., 1., 1., 1.))\n",
    "ax2.w_yaxis.set_pane_color((1., 1., 1., 1.))\n",
    "ax2.w_zaxis.set_pane_color((0.9, 0.9, 0.9, 1.))\n",
    "ax2.grid(False)\n",
    "ax2.w_xaxis.line.set_color('dimgray'); ax2.w_yaxis.line.set_color('dimgray'); ax2.w_zaxis.line.set_color('dimgray')\n",
    "    \n",
    "ax2.set_xlabel('$Frequency\\ (cm^{-1})$', labelpad=14)\n",
    "ax2.set_ylabel('$x_{Bi}$', labelpad=10)\n",
    "ax2.set_zlabel('$Intensity$', labelpad=10)\n",
    "ax2.legend(frameon=False, bbox_to_anchor=(0.9,0.4), bbox_transform=fig.transFigure);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wN1vfDM5FyFs"
   },
   "source": [
    "### Visualization of intermediate features\n",
    "We can visualize the intermediate features on each node projected onto the basis of spherical harmonics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XHt1kAUHFvDP"
   },
   "outputs": [],
   "source": [
    "d = next(iter(dataloader_train))\n",
    "specie = d.symbol[0]\n",
    "sts, st_feats = get_middle_feats(d, model, normalize=True)\n",
    "\n",
    "for sts_idx in range(len(sts)):\n",
    "    traces, traces_species = plotly_surface(sts[sts_idx], st_feats[sts_idx].detach().cpu(), centers=d.pos.cpu(),\n",
    "                                            res=20, radius=True, species=specie)\n",
    "    fig_html = plot_orbitals(traces, traces_species, title_str=f'feature: {str(sts[sts_idx])}')\n",
    "    \n",
    "    with open(f'feature_{str(sts[sts_idx])}.html', 'w') as f:\n",
    "        f.write(fig_html)\n",
    "\n",
    "imgs = [f'feature_{str(sts[sts_idx])}.html' for sts_idx in range(len(sts))]\n",
    "\n",
    "for img in imgs:\n",
    "    display(HTML(filename=img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5_CCJ5mEkjEN"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyP9WdY8pU0J3fdzaTpnaMJ8",
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "phononDoS_colab.ipynb",
   "provenance": []
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}