phdos_eval_MPdata.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict phonon DoS for new materials and evaluate their specific heat capacities\n",
    "\n",
    "- `ComprehensiveEvaluation`: the function that evaluate phonon DoS and heat capacities with the input `*.cif` file\n",
    "- `AtomEmbeddingAndSumLastLayer`: the model function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch device: cuda\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "import torch\n",
    "import torch_geometric\n",
    "import torch_scatter\n",
    "\n",
    "import e3nn\n",
    "from e3nn import rs, o3\n",
    "from e3nn.point.data_helpers import DataPeriodicNeighbors\n",
    "from e3nn.networks import GatedConvParityNetwork\n",
    "from e3nn.kernel_mod import Kernel\n",
    "from e3nn.point.message_passing import Convolution\n",
    "\n",
    "import pymatgen\n",
    "from pymatgen.core.structure import Structure\n",
    "\n",
    "import time, os\n",
    "import datetime\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import numpy as np\n",
    "import h5py\n",
    "from class_evaluate_MPdata import ComprehensiveEvaluation, AtomEmbeddingAndSumLastLayer\n",
    "\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "print('torch device:' , device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import model and trained model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AtomEmbeddingAndSumLastLayer(\n",
       "  (linear): Linear(in_features=118, out_features=64, bias=True)\n",
       "  (model): GatedConvParityNetwork(\n",
       "    (layers): ModuleList(\n",
       "      (0): ModuleList(\n",
       "        (0): Convolution(\n",
       "          (kernel): Kernel (64x0e -> 32x0e,32x0e,32x1o)\n",
       "        )\n",
       "        (1): GatedBlockParity (32x0e + 32x0e + 32x1o -> 32x0e,32x1o)\n",
       "      )\n",
       "      (1): ModuleList(\n",
       "        (0): Convolution(\n",
       "          (kernel): Kernel (32x0e,32x1o -> 32x0e,64x0e,32x1e,32x1o)\n",
       "        )\n",
       "        (1): GatedBlockParity (32x0e + 64x0e + 32x1e,32x1o -> 32x0e,32x1e,32x1o)\n",
       "      )\n",
       "      (2): Convolution(\n",
       "        (kernel): Kernel (32x0e,32x1e,32x1o -> 51x0e)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (relu): ReLU()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_data = torch.load('models/200803-1018_len51max1000_fwin101ord3_trial_run_full_data.torch', map_location=device)\n",
    "\n",
    "model = AtomEmbeddingAndSumLastLayer(model_data['state']['linear.weight'].shape[1], \n",
    "                                     model_data['state']['linear.weight'].shape[0], \n",
    "                                     GatedConvParityNetwork(**model_data['model_kwargs']))\n",
    "\n",
    "model.load_state_dict(model_data['state'])\n",
    "model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import frequency points to evaluate heat capacities \n",
    "\n",
    "Please make sure to unzip the file `models/phdos_e3nn_len51max1000_fwin101ord3.zip`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('models/phdos_e3nn_len51max1000_fwin101ord3.pkl', 'rb') as f:\n",
    "    data_dict = pickle.load(f)\n",
    "phfre = data_dict['phfre']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load cif file names and filter out those materials with more than 13 atoms inside the unit cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('models/cif_unique_files.pkl', 'rb') as f: \n",
    "    ciflist_dict = pickle.load(f)\n",
    "\n",
    "cif_name = ciflist_dict.get('cif_name')\n",
    "cif_id = ciflist_dict.get('cif_id')\n",
    "num_sites = ciflist_dict.get('num_sites')\n",
    "\n",
    "cif_name_suc = [cif_name[i] for i in range(len(cif_name)) if num_sites[i] <= 13]\n",
    "cif_id_suc = [cif_id[i] for i in range(len(cif_id)) if num_sites[i] <= 13]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `T_lst`: temperature list that you want to evaluate $C_{V}$ at\n",
    "- `cif_path`: where you stored all downloaded `*.cif` files\n",
    "- `h5_file`: the `*.h5` file that you want to store calculated phonon DoS and heat capacities\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_lst = [273.15, 293.15]\n",
    "cif_path = 'data/mp_data/'\n",
    "h5_file = 'data/phdos_maxsites13_2020Aug27.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27   Calculating mp-328          \r"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-bc969913b070>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0mmaterial_id\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcif_id_suc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mchunk_evaluation\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mComprehensiveEvaluation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcif_name_suc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_data\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'model_kwargs'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcif_path\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcif_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mchunk_id\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmaterial_id\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m     \u001b[0mchunk_evaluation\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredict_phdos\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchunk_evaluation\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m     \u001b[0mchunk_evaluation\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcal_heatcap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchunk_evaluation\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mphdos\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mphfre\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtolist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mT_lst\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mchunk_evaluation\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstructures\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mh5_file\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Downloads/phonondos_e3nn-revision/class_evaluate_MPdata.py\u001b[0m in \u001b[0;36mpredict_phdos\u001b[0;34m(self, data, model, device)\u001b[0m\n\u001b[1;32m     41\u001b[0m             \u001b[0md\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch_geometric\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBatch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_data_list\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m             \u001b[0md\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 43\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mphdos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0medge_attr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn_norm\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m40\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtolist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     44\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     45\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mcal_heatcap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0momega\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mT_lst\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstructures\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Downloads/phonondos_e3nn-revision/class_evaluate_MPdata.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, batch, *args, **kwargs)\u001b[0m\n\u001b[1;32m     80\u001b[0m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     81\u001b[0m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 82\u001b[0;31m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     83\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mbatch\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     84\u001b[0m             \u001b[0mN\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/e3nn/networks.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input, *args, **kwargs)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    129\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mconv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mact\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayers\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 130\u001b[0;31m             \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    131\u001b[0m             \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mact\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    132\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/e3nn/point/message_passing.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, features, edge_index, edge_r, size, n_norm)\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0msize\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m             \u001b[0msize\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfeatures\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m         \u001b[0mk\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkernel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0medge_r\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m         \u001b[0mk\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiv_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn_norm\u001b[0m \u001b[0;34m**\u001b[0m \u001b[0;36m0.5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpropagate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msize\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/e3nn/kernel_mod.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, r, r_eps, **_kwargs)\u001b[0m\n\u001b[1;32m     92\u001b[0m             \u001b[0mkernel1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzeros\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mRs_out\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mRs_in\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     93\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 94\u001b[0;31m             \u001b[0mkernel1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mright\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mRY\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     95\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m         \u001b[0;31m# (2) Case r = 0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/e3nn/rs.py\u001b[0m in \u001b[0;36mright\u001b[0;34m(self, features_2)\u001b[0m\n\u001b[1;32m    500\u001b[0m         \u001b[0mfeatures_2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfeatures_2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md_in2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    501\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 502\u001b[0;31m         \u001b[0mmixing_matrix\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_sparse_buffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"mixing_matrix\"\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# [out, in1 * in2]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    503\u001b[0m         \u001b[0mmixing_matrix\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmixing_matrix\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msparse_reshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md_out\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0md_in1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md_in2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    504\u001b[0m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmixing_matrix\u001b[0m \u001b[0;34m@\u001b[0m \u001b[0mfeatures_2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mT\u001b[0m  \u001b[0;31m# [out * in1, batch]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/e3nn/lib/python3.8/site-packages/e3nn/util/sparse.py\u001b[0m in \u001b[0;36mget_sparse_buffer\u001b[0;34m(module, name)\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0mcol\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcol\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m         \u001b[0mvalue\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m         \u001b[0msparse_sizes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msiz\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtolist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m     )\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for i in range(0,len(cif_id_suc)):\n",
    "    material_id = cif_id_suc[i]\n",
    "    chunk_evaluation = ComprehensiveEvaluation([cif_name_suc[i]], model_data['model_kwargs'], cif_path=cif_path, chunk_id=material_id)\n",
    "    chunk_evaluation.predict_phdos(chunk_evaluation.data,model,device=device)\n",
    "    chunk_evaluation.cal_heatcap(chunk_evaluation.phdos,phfre.tolist(),T_lst,chunk_evaluation.structures)\n",
    "    if os.path.exists(h5_file):\n",
    "        with h5py.File(h5_file, 'a') as hf:\n",
    "            hf[\"material_id\"].resize((hf[\"material_id\"].shape[0]+np.array([material_id])[None,:].shape[0]),axis=0)\n",
    "            hf[\"material_id\"][-np.array([material_id])[None,:].shape[0]:] = np.array([material_id])[None,:]\n",
    "            hf[\"phdos_max1\"].resize((hf[\"phdos_max1\"].shape[0]+np.array(chunk_evaluation.phdos).shape[0]),axis=0)\n",
    "            hf[\"phdos_max1\"][-np.array(chunk_evaluation.phdos).shape[0]:] = np.array(chunk_evaluation.phdos)\n",
    "            hf[\"phdos_norm\"].resize((hf[\"phdos_norm\"].shape[0]+np.array(chunk_evaluation.phdos_norm).shape[0]),axis=0)\n",
    "            hf[\"phdos_norm\"][-np.array(chunk_evaluation.phdos_norm).shape[0]:] = np.array(chunk_evaluation.phdos_norm)\n",
    "            hf[\"heat_cap_mol\"].resize((hf[\"heat_cap_mol\"].shape[0]+np.array(chunk_evaluation.C_v_mol).shape[0]),axis=0)\n",
    "            hf[\"heat_cap_mol\"][-np.array(chunk_evaluation.C_v_mol).shape[0]:] = np.array(chunk_evaluation.C_v_mol)\n",
    "            hf[\"heat_cap_kg\"].resize((hf[\"heat_cap_kg\"].shape[0]+np.array(chunk_evaluation.C_v_kg).shape[0]),axis=0)\n",
    "            hf[\"heat_cap_kg\"][-np.array(chunk_evaluation.C_v_kg).shape[0]:] = np.array(chunk_evaluation.C_v_kg)\n",
    "            \n",
    "            print(\"{}   Calculating mp-{}          \".format(i, cif_id_suc[i]), end=\"\\r\", flush=True)\n",
    "    else:\n",
    "        with h5py.File(h5_file, 'w') as hf:\n",
    "            hf.create_dataset(\"material_id\", data=np.array([material_id])[None,:],\n",
    "                              compression=\"gzip\", compression_opts=9, chunks=True, maxshape=(None,None))\n",
    "            hf.create_dataset(\"phdos_max1\", data=np.array(chunk_evaluation.phdos),\n",
    "                              compression=\"gzip\", compression_opts=9, chunks=True, maxshape=(None,None))\n",
    "            hf.create_dataset(\"phdos_norm\", data=np.array(chunk_evaluation.phdos_norm),\n",
    "                              compression=\"gzip\", compression_opts=9, chunks=True, maxshape=(None,None))\n",
    "            hf.create_dataset(\"heat_cap_mol\", data=np.array(chunk_evaluation.C_v_mol),\n",
    "                              compression=\"gzip\", compression_opts=9, chunks=True, maxshape=(None,None))\n",
    "            hf.create_dataset(\"heat_cap_kg\", data=np.array(chunk_evaluation.C_v_kg),\n",
    "                              compression=\"gzip\", compression_opts=9, chunks=True, maxshape=(None,None))\n",
    "            hf.create_dataset(\"phfre\", data=phfre,\n",
    "                              compression=\"gzip\", compression_opts=9, chunks=True, maxshape=(None))\n",
    "            print(\"Created new h5py data\\n\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}