transformer-experiments.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:15.951117Z",
     "start_time": "2019-08-23T16:38:15.300381Z"
    }
   },
   "outputs": [],
   "source": [
    "#%load_ext autoreload\n",
    "#%autoreload 2\n",
    "import os\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.utils.data\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:16.485168Z",
     "start_time": "2019-08-23T16:38:16.433372Z"
    }
   },
   "outputs": [],
   "source": [
    "from torch.nn import Parameter\n",
    "\n",
    "def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,\n",
    "                          missing_keys, unexpected_keys, error_msgs):\n",
    "    r\"\"\"Copies parameters and buffers from :attr:`state_dict` into only\n",
    "    this module, but not its descendants. This is called on every submodule\n",
    "    in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this\n",
    "    module in input :attr:`state_dict` is provided as :attr:`local_metadata`.\n",
    "    For state dicts without metadata, :attr:`local_metadata` is empty.\n",
    "    Subclasses can achieve class-specific backward compatible loading using\n",
    "    the version number at `local_metadata.get(\"version\", None)`.\n",
    "\n",
    "    .. note::\n",
    "        :attr:`state_dict` is not the same object as the input\n",
    "        :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So\n",
    "        it can be modified.\n",
    "\n",
    "    Arguments:\n",
    "        state_dict (dict): a dict containing parameters and\n",
    "            persistent buffers.\n",
    "        prefix (str): the prefix for parameters and buffers used in this\n",
    "            module\n",
    "        local_metadata (dict): a dict containing the metadata for this module.\n",
    "            See\n",
    "        strict (bool): whether to strictly enforce that the keys in\n",
    "            :attr:`state_dict` with :attr:`prefix` match the names of\n",
    "            parameters and buffers in this module\n",
    "        missing_keys (list of str): if ``strict=True``, add missing keys to\n",
    "            this list\n",
    "        unexpected_keys (list of str): if ``strict=True``, add unexpected\n",
    "            keys to this list\n",
    "        error_msgs (list of str): error messages should be added to this\n",
    "            list, and will be reported together in\n",
    "            :meth:`~torch.nn.Module.load_state_dict`\n",
    "    \"\"\"\n",
    "    for hook in self._load_state_dict_pre_hooks.values():\n",
    "        hook(state_dict, prefix, local_metadata, strict, missing_keys,\n",
    "             unexpected_keys, error_msgs)\n",
    "\n",
    "    local_name_params = itertools.chain(self._parameters.items(),\n",
    "                                        self._buffers.items())\n",
    "    local_state = {k: v.data for k, v in local_name_params if v is not None}\n",
    "\n",
    "    for name, param in local_state.items():\n",
    "        key = prefix + name\n",
    "        if key in state_dict:\n",
    "            input_param = state_dict[key]\n",
    "\n",
    "            # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+\n",
    "            if len(param.shape) == 0 and len(input_param.shape) == 1:\n",
    "                input_param = input_param[0]\n",
    "\n",
    "            if input_param.shape != param.shape:\n",
    "                # local shape should match the one in checkpoint\n",
    "                error_msgs.append(\n",
    "                    'size mismatch for {}: copying a param with shape {} from checkpoint, '\n",
    "                    'the shape in current model is {}.'.format(\n",
    "                        key, input_param.shape, param.shape))\n",
    "                #if not strict:\n",
    "                #    continue\n",
    "\n",
    "            if isinstance(input_param, Parameter):\n",
    "                # backwards compatibility for serialized parameters\n",
    "                input_param = input_param.data\n",
    "\n",
    "            try:\n",
    "                if False:\n",
    "                    input_param = input_param[torch.randperm(input_param.size()[0])]\n",
    "                param.copy_(input_param)\n",
    "            except Exception:\n",
    "                error_msgs.append(\n",
    "                    'While copying the parameter named \"{}\", '\n",
    "                    'whose dimensions in the model are {} and '\n",
    "                    'whose dimensions in the checkpoint are {}.'.format(\n",
    "                        key, param.size(), input_param.size()))\n",
    "                # PG load partially\n",
    "\n",
    "                if len(input_param.size()) == 3:\n",
    "                    error_msgs.append(\n",
    "                        'Partially copying the parameter named \"{}\", '\n",
    "                        'whose dimensions in the model are {} and '\n",
    "                        'whose dimensions in the checkpoint are {}. - trying {}'\n",
    "                        .format(\n",
    "                            key, param.size(), input_param.size(),\n",
    "                            param[:input_param.size()[0], :input_param.size(\n",
    "                            )[1], :input_param.size()[2]].shape))\n",
    "                else:\n",
    "                    error_msgs.append(\n",
    "                        'Partially copying the parameter named \"{}\", '\n",
    "                        'whose dimensions in the model are {} and '\n",
    "                        'whose dimensions in the checkpoint are {}. - trying {}'\n",
    "                        .format(key, param.size(), input_param.size(),\n",
    "                                param[:input_param.size()[0]].shape))\n",
    "\n",
    "                try:\n",
    "                    new_input_param = torch.empty_like(param)\n",
    "                    new_input_param = torch.nn.init.normal_(new_input_param,\n",
    "                                                            mean=input_param.mean(),\n",
    "                                                            std=input_param.std())\n",
    "\n",
    "                    if len(input_param.size()) == 3:\n",
    "                        new_input_param[:input_param.size()[0], :input_param.\n",
    "                                        size()[1], :input_param.size(\n",
    "                                        )[2]] = input_param\n",
    "                    else:\n",
    "                        new_input_param[:input_param.size()[0]] = input_param\n",
    "                    param.copy_(new_input_param)\n",
    "                except Exception as e:\n",
    "                    assert e\n",
    "                    error_msgs.append(\n",
    "                        'Failed to load weights partially {}'.format(e))\n",
    "        elif strict:\n",
    "            missing_keys.append(key)\n",
    "\n",
    "    if strict:\n",
    "        for key in state_dict.keys():\n",
    "            if key.startswith(prefix):\n",
    "                input_name = key[len(prefix):]\n",
    "                input_name = input_name.split(\n",
    "                    '.', 1)[0]  # get the name of param/buffer/child\n",
    "                if input_name not in self._modules and input_name not in local_state:\n",
    "                    unexpected_keys.append(key)\n",
    "\n",
    "def load_state_dict(self, state_dict, strict=True):\n",
    "    r\"\"\"Copies parameters and buffers from :attr:`state_dict` into\n",
    "    this module and its descendants. If :attr:`strict` is ``True``, then\n",
    "    the keys of :attr:`state_dict` must exactly match the keys returned\n",
    "    by this module's :meth:`~torch.nn.Module.state_dict` function.\n",
    "\n",
    "    Arguments:\n",
    "        state_dict (dict): a dict containing parameters and\n",
    "            persistent buffers.\n",
    "        strict (bool, optional): whether to strictly enforce that the keys\n",
    "            in :attr:`state_dict` match the keys returned by this module's\n",
    "            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n",
    "\n",
    "    Returns:\n",
    "        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n",
    "            * **missing_keys** is a list of str containing the missing keys\n",
    "            * **unexpected_keys** is a list of str containing the unexpected keys\n",
    "    \"\"\"\n",
    "    missing_keys = []\n",
    "    unexpected_keys = []\n",
    "    error_msgs = []\n",
    "\n",
    "    # copy state_dict so _load_from_state_dict can modify it\n",
    "    metadata = getattr(state_dict, '_metadata', None)\n",
    "    state_dict = state_dict.copy()\n",
    "    if metadata is not None:\n",
    "        state_dict._metadata = metadata\n",
    "\n",
    "    def load(module, prefix=''):\n",
    "        local_metadata = {} if metadata is None else metadata.get(\n",
    "            prefix[:-1], {})\n",
    "        module._load_from_state_dict(state_dict, prefix, local_metadata, True,\n",
    "                                     missing_keys, unexpected_keys, error_msgs)\n",
    "        for name, child in module._modules.items():\n",
    "            if child is not None:\n",
    "                load(child, prefix + name + '.')\n",
    "                \n",
    "    load(self)\n",
    "\n",
    "    if strict:\n",
    "        if len(unexpected_keys) > 0:\n",
    "            error_msgs.insert(\n",
    "                0, 'Unexpected key(s) in state_dict: {}. '.format(', '.join(\n",
    "                    '\"{}\"'.format(k) for k in unexpected_keys)))\n",
    "        if len(missing_keys) > 0:\n",
    "            error_msgs.insert(\n",
    "                0, 'Missing key(s) in state_dict: {}. '.format(', '.join(\n",
    "                    '\"{}\"'.format(k) for k in missing_keys)))\n",
    "\n",
    "    if strict and len(error_msgs) > 0:\n",
    "        raise RuntimeError(\n",
    "            'Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n",
    "                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:16.923971Z",
     "start_time": "2019-08-23T16:38:16.919052Z"
    }
   },
   "outputs": [],
   "source": [
    "nn.Module._load_from_state_dict = _load_from_state_dict\n",
    "nn.Module.load_state_dict = load_state_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:18.528699Z",
     "start_time": "2019-08-23T16:38:17.514897Z"
    }
   },
   "outputs": [],
   "source": [
    "from fastai import *\n",
    "from fastai.basic_train import *\n",
    "from fastai.basic_data import *\n",
    "from fastai.data_block import *\n",
    "from fastai.torch_core import *\n",
    "from fastai.train import *\n",
    "from fastai.callback import *\n",
    "from fastai.callbacks import *\n",
    "from fastai.distributed import *\n",
    "from fastai.layers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:18.564608Z",
     "start_time": "2019-08-23T16:38:18.531218Z"
    }
   },
   "outputs": [],
   "source": [
    "fname_ext = lambda fname,ext: f'{str(fname)[:-4]}{ext}{str(fname)[-4:]}'\n",
    "\n",
    "def preprocess(fname, type_index=None, ext=''):\n",
    "    t  = pd.read_csv(fname_ext(fname,ext))\n",
    "    s  = pd.read_csv('structures.csv')\n",
    "    \n",
    "    has_y = 'scalar_coupling_constant' in t.columns\n",
    "\n",
    "    if has_y:\n",
    "        # atom-atom level\n",
    "        # molecule_name,atom_index_0,atom_index_1,type,fc,sd,pso,dso\n",
    "        scalar_couplings = pd.read_csv(f'scalar_coupling_contributions{ext}.csv') # fc,sd,pso,dso\n",
    "\n",
    "        # atom level\n",
    "        # molecule_name,atom_index,XX,YX,ZX,XY,YY,ZY,XZ,YZ,ZZ\n",
    "        magnetic_shielding = pd.read_csv('magnetic_shielding_tensors.csv')\n",
    "        # molecule_name,atom_index,mulliken_charge\n",
    "        mulliken_charges = pd.read_csv('mulliken_charges.csv')\n",
    "\n",
    "        # molecule level\n",
    "        # molecule_name,X,Y,Z\n",
    "        dipole_moments = pd.read_csv('dipole_moments.csv')\n",
    "        # molecule_name,potential_energy\n",
    "        potential_energy = pd.read_csv('potential_energy.csv')\n",
    "\n",
    "    t['molecule_index'] = pd.factorize(t['molecule_name'])[0] + t['id'].min()\n",
    "    # make sure we use the same indexes in train/test (test needs to provide type_index)\n",
    "    if type_index is not None:\n",
    "        t['type_idx'] = t['type'].apply(lambda x: type_index.index(x) ) # pd.factorize(pd.concat([pd.Series(type_index),t['type']]))[0][len(type_index):]\n",
    "    else:\n",
    "        t['type_idx'] = pd.factorize(t['type'])[0]\n",
    "\n",
    "    s['atom_idx'] = s['atom'].apply(lambda x: atoms.index(x) )\n",
    "\n",
    "    max_items = len(t.groupby(['molecule_name', 'atom_index_0']))# if has_y else 422550\n",
    "    max_atoms = int(s.atom_index.max() + 1)\n",
    "\n",
    "    if has_y:\n",
    "        contributions = ['fc','sd','pso','dso']\n",
    "        magnetic_tensors = ['XX','YX','ZX','XY','YY','ZY','XZ','YZ','ZZ']\n",
    "        XYZ = ['X','Y','Z']\n",
    "    xyz = ['x', 'y', 'z']\n",
    "    \n",
    "    x_xyz   = np.zeros((max_items,len(xyz),  max_atoms), dtype=np.float32)\n",
    "    x_type  = np.zeros((max_items,1,         max_atoms), dtype=np.int8)\n",
    "    x_ext   = np.zeros((max_items,1,         max_atoms), dtype=np.bool_)\n",
    "    x_atom  = np.empty((max_items,1,         max_atoms), dtype=np.int8)\n",
    "    x_atom[:] = -1\n",
    "\n",
    "    if has_y:\n",
    "        y_scalar   = np.zeros((max_items,len(contributions)   ,max_atoms), dtype=np.float32)\n",
    "        y_magnetic = np.zeros((max_items,len(magnetic_tensors),max_atoms), dtype=np.float32)\n",
    "        y_mulliken = np.zeros((max_items,1                    ,max_atoms), dtype=np.float32)\n",
    "\n",
    "        y_dipole   = np.zeros((max_items,len(XYZ)), dtype=np.float32)\n",
    "        y_potential= np.zeros((max_items,1       ), dtype=np.float32)\n",
    "\n",
    "        y_magnetic[...] = np.nan\n",
    "        y_mulliken[...] = np.nan\n",
    "    else:\n",
    "        xt_ids = np.zeros((max_items, max_atoms), dtype=np.int32)\n",
    "\n",
    "    m = np.zeros((max_items,), dtype=np.int32)\n",
    "    i = j = 0\n",
    "    \n",
    "    for (m_name, m_index) ,m_group in tqdm(t.groupby(['molecule_name', 'molecule_index'])):\n",
    "        ss = s[s.molecule_name==m_name]\n",
    "        n_atoms = len(ss)\n",
    "        if has_y:\n",
    "            magnetic = magnetic_shielding[\n",
    "                    (magnetic_shielding['molecule_name']==m_name)][magnetic_tensors].values.T\n",
    "\n",
    "            mulliken = mulliken_charges[\n",
    "                    (mulliken_charges['molecule_name']==m_name)]['mulliken_charge'].values.T\n",
    "\n",
    "            scs = scalar_couplings[scalar_couplings['molecule_name']==m_name]\n",
    "            \n",
    "            y_dipole[j,:]= dipole_moments[dipole_moments['molecule_name']==m_name][XYZ].values\n",
    "            y_potential[j,:]=potential_energy[\n",
    "                potential_energy['molecule_name']==m_name]['potential_energy'].values\n",
    "        \n",
    "        for a_name,a_group in m_group.groupby('atom_index_0'):\n",
    "            \n",
    "            ref_a = ss[ss['atom_index']==a_name]\n",
    "            \n",
    "            x_xyz[i] = 0.\n",
    "            x_type[i] = -1\n",
    "            x_ext[i] =  True\n",
    "            \n",
    "            x_xyz[i,:,:n_atoms] = (ss[xyz].values-ref_a[xyz].values).T  # xyz \n",
    "            x_type[i,0,a_group['atom_index_1']] = a_group['type_idx']  # type \n",
    "            x_ext[i,0,a_group['atom_index_1']] = a_group['ext']  # ext \n",
    "            x_atom[i,:,:n_atoms] = ss['atom_idx'].T                \n",
    "\n",
    "            if has_y:\n",
    "                y_scalar[i,:,a_group['atom_index_1']] = scs[scs['atom_index_0']==a_name][contributions]\n",
    "                y_magnetic[i,:,:n_atoms] = magnetic\n",
    "                y_mulliken[i,:,:n_atoms] = mulliken\n",
    "            else:\n",
    "                xt_ids[i,a_group['atom_index_1']] = a_group['id']  \n",
    "\n",
    "            m[i] = m_index\n",
    "            i+=1\n",
    "        j += 1\n",
    "    assert i == max_items\n",
    "    print(i,max_items)\n",
    "    if has_y:\n",
    "        return x_xyz,x_type,x_ext,x_atom, m, y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential\n",
    "    else:\n",
    "        return x_xyz,x_type,x_ext,x_atom, m, xt_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define where you want to use original training set '' or extended ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:18.960049Z",
     "start_time": "2019-08-23T16:38:18.956943Z"
    }
   },
   "outputs": [],
   "source": [
    "ext = '_ext' # or ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load preprocessed or preprocess and save for later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:24.742471Z",
     "start_time": "2019-08-23T16:38:20.072205Z"
    }
   },
   "outputs": [],
   "source": [
    "train_fname = Path('train.npz')\n",
    "types = ['1JHC', '2JHH', '1JHN', '2JHN', '2JHC', '3JHH', '3JHC', '3JHN'] \n",
    "atoms = 'CFHNO'\n",
    "\n",
    "try:\n",
    "    npzfile = np.load(fname_ext(train_fname, ext))\n",
    "    x_xyz   = npzfile['x_xyz']\n",
    "    x_type  = npzfile['x_type']\n",
    "    x_ext   = npzfile['x_ext']\n",
    "    x_atom  = npzfile['x_atom']\n",
    "\n",
    "    y_scalar    = npzfile['y_scalar']\n",
    "    y_magnetic  = npzfile['y_magnetic']\n",
    "    y_mulliken  = npzfile['y_mulliken']\n",
    "    y_dipole    = npzfile['y_dipole']\n",
    "    y_potential = npzfile['y_potential']\n",
    "    m = npzfile['m']\n",
    "    max_items, max_atoms = x_xyz.shape[0], x_xyz.shape[-1]\n",
    "except:\n",
    "    x_xyz,x_type,x_ext,x_atom, m, y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential = \\\n",
    "        preprocess(train_fname.with_suffix('.csv'), type_index=types, ext=ext)\n",
    "    np.savez(fname_ext(train_fname, ext), \n",
    "             x_xyz=x_xyz,\n",
    "             x_type=x_type,\n",
    "             x_ext=x_ext,\n",
    "             x_atom=x_atom,\n",
    "             y_scalar=y_scalar,\n",
    "             y_magnetic=y_magnetic,\n",
    "             y_mulliken=y_mulliken,\n",
    "             y_dipole=y_dipole,\n",
    "             y_potential=y_potential,\n",
    "             m=m)\n",
    "n_types = int(x_type[~np.isnan(x_type)].max() + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:24.751433Z",
     "start_time": "2019-08-23T16:38:24.744879Z"
    }
   },
   "outputs": [],
   "source": [
    "use_memmap = True\n",
    "try:\n",
    "    load_fn = np.load if not use_memmap else partial(np.lib.format.open_memmap, mode='r')\n",
    "    x_coulombmat = load_fn(f'x_coulombmat32{ext}.npy')\n",
    "except:\n",
    "    x_coulombmat = np.load(f'x_coulombmat{ext}.npy', allow_pickle=True)\n",
    "    x_coulombmat = np.array(x_coulombmat.tolist()).astype(np.float32)\n",
    "    np.save(f'x_coulombmat32{ext}.npy', x_coulombmat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:24.757374Z",
     "start_time": "2019-08-23T16:38:24.754107Z"
    }
   },
   "outputs": [],
   "source": [
    "x_qm9_mulliken = load_fn(f'x_qm9_mulliken{ext}.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:24.767162Z",
     "start_time": "2019-08-23T16:38:24.759924Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1405126, 3, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 4, 29),\n",
       " (1405126, 9, 29),\n",
       " (1405126, 1, 29),\n",
       " (1405126, 3),\n",
       " (1405126, 1),\n",
       " (1405126,)]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[v.shape for v in [x_xyz,x_type,x_ext,x_atom,x_qm9_mulliken, \n",
    "                   y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential, m]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:25.801808Z",
     "start_time": "2019-08-23T16:38:24.769243Z"
    }
   },
   "outputs": [],
   "source": [
    "x_xyz_mean, x_xyz_std = Tensor(x_xyz.mean(axis=(0,2),keepdims=True)), Tensor(x_xyz.std(axis=(0,2),keepdims=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:25.953583Z",
     "start_time": "2019-08-23T16:38:25.803811Z"
    }
   },
   "outputs": [],
   "source": [
    "x_qm9_mulliken_mean = Tensor(x_qm9_mulliken.mean(axis=(0,2),keepdims=True))\n",
    "x_qm9_mulliken_std  = Tensor(x_qm9_mulliken.std( axis=(0,2),keepdims=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:25.959654Z",
     "start_time": "2019-08-23T16:38:25.955551Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 1])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_xyz_std.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fastai classes (this should should be done into its own `application` but who has time?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:25.978083Z",
     "start_time": "2019-08-23T16:38:25.962189Z"
    }
   },
   "outputs": [],
   "source": [
    "class MoleculeItem(ItemBase):\n",
    "    def __init__(self,i,xyz,type,ext,atom,qm9_mulliken,coulomb): \n",
    "        self.i,self.xyz,self.type,self.ext, self.atom,self.qm9_mulliken,self.coulomb = \\\n",
    "            i,xyz,type,ext,atom,qm9_mulliken,coulomb\n",
    "        self.data = [Tensor(xyz), LongTensor((type)), \n",
    "                     Tensor(ext), LongTensor((atom)),Tensor(qm9_mulliken), Tensor(coulomb)]\n",
    "    def __str__(self):\n",
    "        # TODO: count n_atoms correctly. \n",
    "        n_atoms = np.count_nonzero(np.sum(np.absolute(self.xyz), axis=0))+1\n",
    "        n_couplings = np.sum((self.type!=-1))\n",
    "        return f'{self.i} {n_atoms} atoms {n_couplings} couplings'\n",
    "    \n",
    "    def apply_tfms(self, tfms:Collection, **kwargs):\n",
    "        x = self.clone()\n",
    "        for t in tfms:\n",
    "            if t: x.data = t(x.data)\n",
    "        return x\n",
    "    \n",
    "    def clone(self):\n",
    "        return self.__class__(self.i,self.xyz,self.type,self.ext,self.atom,self.qm9_mulliken,self.coulomb)\n",
    "    \n",
    "class ScalarCouplingItem(ItemBase):\n",
    "    def __init__(self,scalar,magnetic,mulliken,dipole,potential,**kwargs): \n",
    "        self.scalar,self.magnetic,self.mulliken,self.dipole,self.potential = \\\n",
    "            scalar,magnetic,mulliken,dipole,potential\n",
    "        self.data = (Tensor(scalar), Tensor(magnetic), Tensor(dipole), Tensor(potential))\n",
    "    def __str__(self):\n",
    "        res, spacer, n_couplings = '', '', 0\n",
    "        for s in self.data[0].sum(dim=0):\n",
    "            if s==0.: spacer = ' * '\n",
    "            else: \n",
    "                res += f'{spacer}{s:.4f}'\n",
    "                spacer = ' '\n",
    "                n_couplings +=1\n",
    "        return f'{n_couplings}: {res}'\n",
    "    def __hash__(self): return hash(str(self))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:25.996906Z",
     "start_time": "2019-08-23T16:38:25.980347Z"
    }
   },
   "outputs": [],
   "source": [
    "class LMAEMaskedLoss(Module):\n",
    "    def __init__(self,\n",
    "                 contrib_w=0., magnetic_w=0., dipole_w=0., potential_w=0., \n",
    "                 types_w = [1]*n_types, return_all=False, proxy_log=torch.log, exclude_ext=False):\n",
    "        self.contrib_w,self.magnetic_w,self.dipole_w,self.potential_w = contrib_w,magnetic_w,dipole_w,potential_w\n",
    "        self.types_w = types_w\n",
    "        self.return_all = return_all\n",
    "        self.proxy_log = proxy_log\n",
    "        self.exclude_ext = exclude_ext\n",
    "    \n",
    "    def forward(self, input_outputs, t_scalar, t_magnetic, t_dipole, t_potential):    \n",
    "        type, ext, p_scalar, p_magnetic, p_dipole, p_potential = input_outputs\n",
    "        loss = 0.\n",
    "        n = 0\n",
    "        j_loss = [0] * n_types\n",
    "        for t in range(n_types):\n",
    "            mask = (type == t).squeeze(1) if not self.exclude_ext else ((type == t) & (ext == 0)).squeeze(1)\n",
    "            if mask.sum() > 0:\n",
    "                _output,_target = p_scalar.transpose(1,2)[mask], t_scalar.transpose(1,2)[mask] # scalars at the end\n",
    "                # LMAE scalar\n",
    "                s_loss = self.proxy_log((_output.sum(dim=-1) - _target.sum(dim=-1)).abs().mean()+1e-9)\n",
    "                loss += self.types_w[t] * s_loss\n",
    "                j_loss[t] += s_loss\n",
    "                # LMAE scalar contributions\n",
    "                for i_contrib in range(_output.shape[-1]):\n",
    "                    loss += self.contrib_w * \\\n",
    "                        self.proxy_log((_output[...,i_contrib] - _target[...,i_contrib]).abs().mean()+1e-9)\n",
    "                n+=1\n",
    "        loss /= n\n",
    "        \n",
    "        if self.magnetic_w > 0:\n",
    "            mask = ~torch.isnan(t_magnetic)\n",
    "            loss += self.magnetic_w * MSELossFlat()(p_magnetic[mask], t_magnetic[mask])\n",
    "            \n",
    "        if self.dipole_w    > 0: loss += self.dipole_w    * MSELossFlat()(p_dipole,    t_dipole)\n",
    "        if self.potential_w > 0: loss += self.potential_w * MSELossFlat()(p_potential, t_potential)\n",
    "\n",
    "        return loss if not self.return_all else (loss, *j_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:26.004471Z",
     "start_time": "2019-08-23T16:38:25.998949Z"
    }
   },
   "outputs": [],
   "source": [
    "class ScalarCouplingList(ItemList):\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.loss_func = LMAEMaskedLoss\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        return ScalarCouplingItem(*o)\n",
    "\n",
    "    def reconstruct(self,t): return 0; # TODO for viz !!!! ScalarCouplingItem(t.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quaterions allow us to rotate 3d points randoming with a nice uniform distribution of 3 numbers hece we use them, however it's still to be seen if are useful here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:38:26.041887Z",
     "start_time": "2019-08-23T16:38:26.006694Z"
    }
   },
   "outputs": [],
   "source": [
    "#https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py\n",
    "def qrot(q, v):\n",
    "    \"\"\"\n",
    "    Rotate vector(s) v about the rotation described by quaternion(s) q.\n",
    "    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,\n",
    "    where * denotes any number of dimensions.\n",
    "    Returns a tensor of shape (*, 3).\n",
    "    \"\"\"\n",
    "    assert q.shape[-1] == 4\n",
    "    assert v.shape[-1] == 3\n",
    "    assert q.shape[:-1] == v.shape[:-1]\n",
    "    \n",
    "    original_shape = list(v.shape)\n",
    "    q = q.view(-1, 4)\n",
    "    v = v.view(-1, 3)\n",
    "    \n",
    "    qvec = q[:, 1:]\n",
    "    uv = torch.cross(qvec, v, dim=1)\n",
    "    uuv = torch.cross(qvec, uv, dim=1)\n",
    "    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)\n",
    "\n",
    "def random_rotation(data):\n",
    "    x_xyz = data[0].transpose(0,1)\n",
    "    r = torch.rand(3)\n",
    "    sq1_v1,sqv1,v2_2pi,v3_2pi = torch.sqrt(1-r[:1]),torch.sqrt(r[:1]),2*math.pi*r[1:2],2*math.pi*r[2:3]\n",
    "    q = torch.cat([sq1_v1*torch.sin(v2_2pi), sq1_v1*torch.cos(v2_2pi), \n",
    "                   sqv1  *torch.sin(v3_2pi), sqv1  *torch.cos(v3_2pi)], dim=0).unsqueeze(0)\n",
    "    x_xyz = qrot(q.expand(x_xyz.shape[0],-1), x_xyz).squeeze(0).transpose(0,1)\n",
    "    return (x_xyz, *data[1:])\n",
    "\n",
    "def normalize(data):\n",
    "    sq = False\n",
    "    if data[0].ndim < 3:\n",
    "        data[0].unsqueeze_(0)\n",
    "        data[4].unsqueeze_(0)\n",
    "        sq = True\n",
    "    x_xyz      = (data[0] - x_xyz_mean)          / x_xyz_std\n",
    "    x_mulliken = (data[4] - x_qm9_mulliken_mean) / x_qm9_mulliken_std\n",
    "    if sq:\n",
    "        x_xyz.squeeze_(0)\n",
    "        x_mulliken.squeeze_(0)\n",
    "    return (x_xyz, data[1],data[2],data[3],x_mulliken,data[5])\n",
    "\n",
    "def canonize(data):\n",
    "    xyz,type,ext,atom,mulliken,coulomb = data\n",
    "    mask = (atom == -1).squeeze(0)\n",
    "    mask_atoms = ~mask.unsqueeze(0)\n",
    "    n_atoms = mask_atoms.sum()\n",
    "    i = torch.nonzero(type.squeeze(0) == -1)[0] # pick first one w/o j-coupling\n",
    "    xyz[:,mask], type[:,mask],ext[:,mask],atom[:,mask],mulliken[:,mask] = \\\n",
    "        0,-1,1,-1,0\n",
    "#        xyz[:,i],type[:,i],   ext[:,i],   atom[:,i],   mulliken[:,i]\n",
    "    return (xyz,type,ext,atom,mulliken,coulomb, mask_atoms, n_atoms)\n",
    "\n",
    "def canonize2(data):\n",
    "    xyz,type,ext,atom,mulliken = data\n",
    "    mask = (atom == -1).squeeze(0)\n",
    "    i_max_atom = torch.nonzero(atom != -1).max() + 1\n",
    "    #mask_atoms = torch.ones((max_atoms,   max_atoms), dtype=torch.uint8)\n",
    "    #zeros      = torch.zeros((i_max_atom, i_max_atom), dtype=torch.uint8)\n",
    "    #mask_atoms[:zeros.shape[0],:zeros.shape[1]] = zeros\n",
    "\n",
    "    mask_atoms = torch.ones ((max_atoms, ), dtype=torch.uint8)\n",
    "    zeros      = torch.zeros((i_max_atom,), dtype=torch.uint8)\n",
    "    mask_atoms[:zeros.shape[0],] = zeros\n",
    "    #mask_atoms[(type==-1).squeeze(0),:] = 1\n",
    "    n_atoms = i_max_atom\n",
    "    xyz[:,mask], type[:,mask],ext[:,mask],atom[:,mask],mulliken[:,mask] = 0,-1,1,-1,0\n",
    "    return (xyz,type,ext,atom,mulliken, mask_atoms, n_atoms)\n",
    "\n",
    "def canonize_and_shuffle(data):\n",
    "    xyz,type,ext,atom,mulliken,coulomb = data\n",
    "    mask = (atom == -1).squeeze(0)\n",
    "    mask_atoms = ~mask.unsqueeze(0)\n",
    "    n_atoms = mask_atoms.sum()\n",
    "    i = torch.nonzero(type.squeeze(0) == -1)[0] # pick first one w/o j-coupling\n",
    "    xyz[:,mask], type[:,mask],ext[:,mask],atom[:,mask],mulliken[:,mask] = \\\n",
    "        0,-1,1,-1,0\n",
    "    \n",
    "    perm = torch.randperm(n_atoms)\n",
    "    xyz[:,~mask] = xyz[:,~mask][:,perm]\n",
    "    type[:,~mask] = type[:,~mask][:,perm]\n",
    "    ext[:,~mask] = ext[:,~mask][:,perm]\n",
    "    atom[:,~mask] = atom[:,~mask][:,perm]\n",
    "    mulliken[:,~mask] = mulliken[:,~mask][:,perm]\n",
    "    \n",
    "    return (xyz,type,ext,atom,mulliken,coulomb, mask_atoms, n_atoms)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build `data` bunch etc. for fastai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:39:57.221762Z",
     "start_time": "2019-08-23T16:38:26.353364Z"
    }
   },
   "outputs": [],
   "source": [
    "data = ItemList(items=(MoleculeItem(i,*v) for i,v in \n",
    "                       enumerate(zip(x_xyz,x_type,x_ext,x_atom,x_qm9_mulliken,x_coulombmat))),\n",
    "                label_cls=ScalarCouplingItem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:39:57.596889Z",
     "start_time": "2019-08-23T16:39:57.223751Z"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "_, idx_valid_split = train_test_split(range(m.max()+1), test_size=0.1, random_state=13)\n",
    "idx_valid_split = np.argwhere(np.isin(m, idx_valid_split)).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:39:57.602893Z",
     "start_time": "2019-08-23T16:39:57.599833Z"
    },
    "pixiedust": {
     "displayParams": {}
    }
   },
   "outputs": [],
   "source": [
    "#data = data.split_by_idx(idx_valid_split)\n",
    "data = data.split_none()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:01.484048Z",
     "start_time": "2019-08-23T16:39:57.605156Z"
    }
   },
   "outputs": [],
   "source": [
    "data = data.label_from_func(\n",
    "    func=lambda o: (y_scalar[o.i], y_magnetic[o.i], y_mulliken[o.i], y_dipole[o.i], y_potential[o.i]),\n",
    "    label_cls=ScalarCouplingList)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:01.489979Z",
     "start_time": "2019-08-23T16:40:01.485802Z"
    }
   },
   "outputs": [],
   "source": [
    "if False:\n",
    "    try:\n",
    "        xt_coulombmat = load_fn(f'xt_coulombmat32{ext}.npy')\n",
    "    except:\n",
    "        xt_coulombmat = np.load(f'xt_coulombmat{ext}.npy', allow_pickle=True)\n",
    "        xt_coulombmat = np.array(xt_coulombmat.tolist()).astype(np.float32)\n",
    "        np.save(f'xt_coulombmat32{ext}.npy', xt_coulombmat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:01.497536Z",
     "start_time": "2019-08-23T16:40:01.491898Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<function normalize at 0x7fda172c3ae8>, <function canonize at 0x7fda172c3d90>] [<function normalize at 0x7fda172c3ae8>, <function canonize at 0x7fda172c3d90>]\n"
     ]
    }
   ],
   "source": [
    "tfms = [normalize, canonize]\n",
    "tta_tfms = list(tfms)\n",
    "#tta_tfms.insert(0,random_rotation)\n",
    "#tta_tfms[-1] = canonize_and_shuffle\n",
    "print(tta_tfms, tfms)\n",
    "data = data.transform((tta_tfms, tfms))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:01.612592Z",
     "start_time": "2019-08-23T16:40:01.499281Z"
    }
   },
   "outputs": [],
   "source": [
    "data=data.databunch(num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:17.581796Z",
     "start_time": "2019-08-23T16:40:01.615559Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LabelList (785836 items)\n",
       "x: ItemList\n",
       "1 5 atoms 4 couplings,2 5 atoms 4 couplings,3 5 atoms 4 couplings,4 5 atoms 4 couplings,6 4 atoms 3 couplings\n",
       "y: ScalarCouplingList\n",
       "4: 84.8076 * -11.2569 -11.2549 -11.2543,4: 84.8074 -11.2569 * -11.2542 -11.2548,4: 84.8093 -11.2549 -11.2542 * -11.2543,4: 84.8095 -11.2543 -11.2548 -11.2543,3: 32.6888 * -11.1867 -11.1757\n",
       "Path: ."
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.train_ds.filter_by_func(lambda item, _: len(item.data[2].cpu().numpy()[item.data[2].cpu().numpy()==0]) == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-22T07:20:35.131065Z",
     "start_time": "2019-08-22T07:20:33.484583Z"
    }
   },
   "outputs": [],
   "source": [
    "data.valid_ds.filter_by_func(lambda item, _: len(item.data[2].cpu().numpy()[item.data[2].cpu().numpy()==0]) == 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model\n",
    "Whole model here, self-contained (needs some cleanup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:17.599714Z",
     "start_time": "2019-08-23T16:40:17.583854Z"
    }
   },
   "outputs": [],
   "source": [
    "class LMAEMetric(LearnerCallback):\n",
    "    _order=-20 # Needs to run before the recorder\n",
    "    def __init__(self, learn, val_only=True):\n",
    "        super().__init__(learn)\n",
    "        self.val_only=val_only\n",
    "        self.metric = LMAEMaskedLoss(return_all=True, exclude_ext=True)\n",
    "\n",
    "    def on_train_begin(self, **kwargs):\n",
    "        if not self.val_only: self.learn.recorder.add_metric_names(['tLMAE'])\n",
    "        self.learn.recorder.add_metric_names(['👉🏻LMAE👈🏻'] + [f'lmae{i}' for i in range(n_types)])\n",
    "            \n",
    "    def on_batch_end(self, train, last_output, last_target, **kwargs):\n",
    "        if self.val_only and train: return \n",
    "        preds,targs = self.preds[int(train)], self.targs[int(train)] # 0 val 1 train\n",
    "        if preds is None:\n",
    "            targs, preds = listify(last_target), listify(last_output)\n",
    "            targs,preds = [t.detach() for t in targs],[t.detach() for t in preds]\n",
    "        else:\n",
    "            for i,(o,t) in enumerate(zip(last_output, last_target)):\n",
    "                preds[i] = torch.cat([preds[i], o.detach()], dim=0)\n",
    "                targs[i] = torch.cat([targs[i], t.detach()], dim=0)\n",
    "        self.preds[int(train)], self.targs[int(train)] = preds,targs\n",
    "        \n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.targs, self.preds = [None, None], [None, None]\n",
    "\n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        mets = []\n",
    "        if self.preds[1]: mets.append(self.metric.forward(self.preds[1], *self.targs[1])[0]) # just tLMAE\n",
    "        if self.preds[0]: mets.extend(self.metric.forward(self.preds[0], *self.targs[0]))\n",
    "        return add_metrics(last_metrics, mets) if mets else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:17.658423Z",
     "start_time": "2019-08-23T16:40:17.601574Z"
    }
   },
   "outputs": [],
   "source": [
    "Activation = Enum('Activation', 'ReLU Swish GeLU LearnedSwish')\n",
    "\n",
    "class PositionalEncoding(Module):\n",
    "    \"Implement the PE function.\"\n",
    "    def __init__(self, d_model, dropout, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        # Compute the positional encodings once in log space.\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        \n",
    "        position = torch.arange(0., max_len).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model))\n",
    "\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)[:, :-1]\n",
    "        pe = pe.unsqueeze(0)\n",
    "        self.register_buffer('pe', pe)\n",
    "        print(pe.size())\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.transpose(2, 1)\n",
    "        x = x + Variable(self.pe[:, :x.size(1)], \n",
    "                         requires_grad=False)\n",
    "        return self.dropout(x).transpose(2, 1)\n",
    "    \n",
    "class GeLU(Module):\n",
    "    def forward(self, x): return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))\n",
    "\n",
    "class Swish(Module):\n",
    "    def forward(self, x): return x * torch.sigmoid(x)\n",
    "    \n",
    "class LearnedSwish(Module):\n",
    "    def __init__(self, slope = 0.5):\n",
    "        super().__init__()\n",
    "        self.slope = slope * torch.nn.Parameter(torch.ones(1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.slope.cuda() * x * torch.sigmoid(x)\n",
    "\n",
    "_activ_func = {Activation.ReLU:nn.ReLU(inplace=True), Activation.GeLU:GeLU(), Activation.Swish: Swish(), Activation.LearnedSwish: LearnedSwish()}\n",
    "\n",
    "def feed_forward(d_model:int, d_ff:int, ff_p:float=0., act:Activation=Activation.ReLU, double_drop:bool=True):\n",
    "    layers = [nn.Linear(d_model, d_ff), _activ_func[act]]\n",
    "    if double_drop: layers.append(nn.Dropout(ff_p))\n",
    "    return SequentialEx(*layers, nn.Linear(d_ff, d_model), nn.Dropout(ff_p), MergeLayer(), nn.LayerNorm(d_model))\n",
    "\n",
    "class MultiHeadAttention(Module):\n",
    "    \"MutiHeadAttention.\"\n",
    "    def __init__(self, n_heads:int, d_model:int, d_head:int=None, resid_p:float=0., attn_p:float=0., bias:bool=False,\n",
    "                 scale:bool=True):\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.n_heads,self.d_head,self.scale = n_heads,d_head,scale\n",
    "\n",
    "        self.attention1 = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "        self.attention2 = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "        self.attention3 = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "        self.attention4 = nn.Linear(d_model, d_model, bias=bias)\n",
    "        self.out = nn.Linear(n_heads * d_head, d_model, bias=bias)\n",
    "        self.drop_att,self.drop_res = nn.Dropout(attn_p),nn.Dropout(resid_p)\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ln(x + self.drop_res(self.out(self._apply_attention(x, mask=mask, **kwargs))))\n",
    "\n",
    "    def _apply_attention(self, x:Tensor, mask:Tensor=None):\n",
    "        if False:\n",
    "            if mask is not None:\n",
    "                x= x.masked_fill(mask[:,:,0].squeeze(1).unsqueeze(-1).repeat(1, 1, x.size(2)),  0.)\n",
    "        bs,x_len = x.size(0),x.size(1)\n",
    "        wq,wk,wv = self.attention1(x), self.attention2(x), self.attention3(x)\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        wq,wk,wv = wq.permute(0, 2, 1, 3),wk.permute(0, 2, 3, 1),wv.permute(0, 2, 1, 3)\n",
    "        attn_score = torch.matmul(wq, wk)\n",
    "        if self.scale: attn_score.div_(self.d_head ** 0.5)\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask,  -1.0E9).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=-1))\n",
    "        attn_vec = torch.matmul(attn_prob, wv)\n",
    "        return self.attention4(attn_vec.permute(0, 2, 1, 3).contiguous().contiguous().view(bs, x_len, -1))\n",
    "\n",
    "    def _attention_einsum(self, x, mask=None):\n",
    "        # Permute and matmul is a little bit faster but this implementation is more readable\n",
    "        bs,x_len = x.size(0),x.size(1)\n",
    "        wq,wk,wv = torch.chunk(self.attention(x), 3, dim=-1)\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        attn_score = torch.einsum('bind,bjnd->bijn', (wq, wk))\n",
    "        if self.scale: attn_score.mul_(1/(self.d_head ** 0.5))\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask, -float('-inf')).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=2))\n",
    "        attn_vec = torch.einsum('bijn,bjnd->bind', (attn_prob, wv))\n",
    "        return attn_vec.contiguous().view(bs, x_len, -1)\n",
    "\n",
    "\n",
    "class DecoderLayer(Module):\n",
    "    \"Basic block of a Transformer model.\"\n",
    "    #Can't use Sequential directly cause more than one input...\n",
    "    def __init__(self, n_heads:int, d_model:int, d_head:int, d_inner:int, resid_p:float=0., attn_p:float=0., ff_p:float=0.,\n",
    "                 bias:bool=True, scale:bool=True, act:Activation=Activation.ReLU, double_drop:bool=True,\n",
    "                 attn_cls:Callable=MultiHeadAttention):\n",
    "        self.mhra = attn_cls(n_heads, d_model, d_head, resid_p=resid_p, attn_p=attn_p, bias=bias, scale=scale)\n",
    "        self.ff   = feed_forward(d_model, d_inner, ff_p=ff_p, act=act, double_drop=double_drop)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs): return self.ff(self.mhra(x, mask=mask, **kwargs))\n",
    "\n",
    "class Transformer(Module):\n",
    "    \"Transformer model: https://arxiv.org/abs/1706.03762.\"\n",
    "    def __init__(self, n_layers:int, n_heads:int, d_model:int, d_head:int, d_inner:int,\n",
    "                 resid_p:float=0., attn_p:float=0., ff_p:float=0., embed_p:float=0., bias:bool=True, scale:bool=True,\n",
    "                 act:Activation=Activation.ReLU, double_drop:bool=True, attn_cls:Callable=MultiHeadAttention,\n",
    "                 learned_pos_enc:bool=True, mask:bool=True):\n",
    "        self.mask = mask\n",
    "        #self.encoder = nn.Embedding(vocab_sz, d_model)\n",
    "        #self.pos_enc = nn.Embedding(ctx_len, d_model) if learned_pos_enc else PositionalEncoding(d_model)\n",
    "        #self.pos_enc = PositionalEncoding(29, 0.1)\n",
    "        self.drop_emb = nn.Dropout(embed_p)\n",
    "        self.layers = nn.ModuleList([DecoderLayer(n_heads, d_model, d_head, d_inner, resid_p=resid_p, attn_p=attn_p,\n",
    "                      ff_p=ff_p, bias=bias, scale=scale, act=act, double_drop=double_drop,\n",
    "                      attn_cls=attn_cls) for k in range(n_layers)])\n",
    "\n",
    "    def reset(self): pass\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        #bs, x_len = x.size()\n",
    "        #pos = torch.arange(0, x_len, device=x.device, dtype=x.dtype)\n",
    "        inp = self.drop_emb(x)# + self.pos_enc(pos)[None]) #.mul_(self.d_model ** 0.5)\n",
    "        #if False:\n",
    "        #    inp += self.pos_enc(x)[None]\n",
    "        #mask = None #torch.triu(x.new_ones(x_len, x_len), diagonal=1).byte()[None,None] if self.mask else None\n",
    "        #[None,:,:None] for einsum implementation of attention\n",
    "        for layer in self.layers: inp = layer(inp, mask=mask)\n",
    "        return inp #For the LinearDecoder\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:17.689812Z",
     "start_time": "2019-08-23T16:40:17.660300Z"
    }
   },
   "outputs": [],
   "source": [
    "class MultiHeadAttentionOld(Module):\n",
    "    \"MutiHeadAttention.\"\n",
    "    def __init__(self, n_heads:int, d_model:int, d_head:int=None, resid_p:float=0., attn_p:float=0., bias:bool=True,\n",
    "                 scale:bool=True):\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.n_heads,self.d_head,self.scale = n_heads,d_head,scale\n",
    "        self.attention = nn.Linear(d_model, 3 * n_heads * d_head, bias=bias)\n",
    "        self.out = nn.Linear(n_heads * d_head, d_model, bias=bias)\n",
    "        self.drop_att,self.drop_res = nn.Dropout(attn_p),nn.Dropout(resid_p)\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ln(x + self.drop_res(self.out(self._apply_attention(x, mask=mask, **kwargs))))\n",
    "\n",
    "    def _apply_attention(self, x:Tensor, mask:Tensor=None):\n",
    "        bs,x_len = x.size(0),x.size(1)\n",
    "        wq,wk,wv = torch.chunk(self.attention(x), 3, dim=-1)\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        wq,wk,wv = wq.permute(0, 2, 1, 3),wk.permute(0, 2, 3, 1),wv.permute(0, 2, 1, 3)\n",
    "        attn_score = torch.matmul(wq, wk)\n",
    "        if self.scale: attn_score.div_(self.d_head ** 0.5)\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask,  -1.0E9).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=-1))\n",
    "        attn_vec = torch.matmul(attn_prob, wv)\n",
    "        return attn_vec.permute(0, 2, 1, 3).contiguous().contiguous().view(bs, x_len, -1)\n",
    "    \n",
    "class AtomTransformerOld(Module):\n",
    "    def __init__(self,  n_heads,d_model, d_head=None, **kwargs):\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.transformer = Transformer(n_heads=n_heads,d_model=d_model, d_head=d_head, **kwargs)\n",
    "        \n",
    "        #self.scalar    = nn.Conv1d(d_model+ n_types + 1, 4, 1)\n",
    "        self.scalar    = nn.Conv1d(d_model, 4, 1)\n",
    "        self.magnetic  = nn.Conv1d(d_model, 9, 1)\n",
    "        self.dipole    = nn.Linear(d_model*max_atoms, 3)\n",
    "        self.potential = nn.Linear(d_model*max_atoms, 1)\n",
    "        \n",
    "        self.n_atom_embedding = d_model//2\n",
    "        self.n_type_embedding = d_model - self.n_atom_embedding - 6\n",
    "        self.type_embedding = nn.Embedding(len(types)+1,self.n_type_embedding)\n",
    "        self.atom_embedding = nn.Embedding(len(atoms)+1,self.n_atom_embedding)\n",
    "    \n",
    "        #n_pos_encoder = d_model - n_type_embedding - n_atom_embedding\n",
    "        #self.pos_encoder = nn.Sequential(\n",
    "        #    nn.Conv1d(3+1+1,n_pos_encoder, 1), nn.ReLU(), nn.BatchNorm1d(n_pos_encoder),\n",
    "        #)\n",
    "        \n",
    "    def forward(self,xyz,type,ext,atom,mulliken,coulomb, mask_atoms, n_atoms):\n",
    "        bs, _, n_pts = xyz.shape        \n",
    "        t = self.type_embedding((type+1).squeeze(1)) #* math.sqrt(self.n_atom_embedding) #.transpose(1,2)\n",
    "        a = self.atom_embedding((atom+1).squeeze(1)) #* math.sqrt(self.n_type_embedding) #.transpose(1,2)\n",
    "        \n",
    "        x = torch.cat([xyz, mulliken, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)        \n",
    "        #x = self.pos_encoder(x).transpose(1,2)\n",
    "\n",
    "        x = torch.cat([x.transpose(1,2), t, a], dim=-1) \n",
    "\n",
    "        mask = (coulomb == 0).unsqueeze(1)\n",
    "        #mask = torch.triu(x.new_ones(max_atoms, max_atoms), diagonal=1).byte()[None,:,:,None]#[None,None] \n",
    "        #print(mask.shape, mask)\n",
    "        x = self.transformer(x, mask).transpose(1,2).contiguous()\n",
    "        \n",
    "        #t_one_hot = torch.zeros(bs,n_types+1,n_pts,device=type.device,dtype=x.dtype).scatter_(1,type+1, 1.)\n",
    "        #scalar    = self.scalar(torch.cat([x, t_one_hot], dim=1))\n",
    "        scalar    = self.scalar(x)\n",
    "        \n",
    "        magnetic  = self.magnetic(x) \n",
    "        dipole    = self.dipole(x.view(bs,-1))\n",
    "        potential = self.potential(x.view(bs,-1))\n",
    "                \n",
    "        return type,ext,scalar,magnetic,dipole,potential\n",
    "    \n",
    "    def reset(self): pass\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:17.760459Z",
     "start_time": "2019-08-23T16:40:17.691804Z"
    }
   },
   "outputs": [],
   "source": [
    "class AtomTransformer(Module):\n",
    "    def __init__(self,  n_heads,d_model, d_head=None, **kwargs):\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.transformer = Transformer(n_heads=n_heads,d_model=d_model, d_head=d_head, **kwargs)\n",
    "        \n",
    "        self.scalar    = nn.Conv1d(d_model+ n_types + 1, 4, 1)\n",
    "        self.magnetic  = nn.Conv1d(d_model, 9, 1)\n",
    "        self.dipole    = nn.Linear(d_model*max_atoms, 3)\n",
    "        self.potential = nn.Linear(d_model*max_atoms, 1)\n",
    "        \n",
    "        self.n_atom_embedding = d_model//2\n",
    "        self.n_type_embedding = d_model - self.n_atom_embedding - 3 # -3 best model\n",
    "        self.type_embedding = nn.Embedding(len(types)+1,self.n_type_embedding)\n",
    "        self.atom_embedding = nn.Embedding(len(atoms)+1,self.n_atom_embedding)\n",
    "    \n",
    "        #n_pos_encoder = d_model - n_type_embedding - n_atom_embedding\n",
    "        #self.pos_encoder = nn.Sequential(\n",
    "        #    nn.Conv1d(3+1+1,n_pos_encoder, 1), nn.ReLU(), nn.BatchNorm1d(n_pos_encoder),\n",
    "        #)\n",
    "        \n",
    "    def forward(self,xyz,type,ext,atom,mulliken,coulomb, mask_atoms, n_atoms):\n",
    "        bs, _, n_pts = xyz.shape        \n",
    "        t = self.type_embedding((type+1).squeeze(1)) #* math.sqrt(self.n_atom_embedding) #.transpose(1,2)\n",
    "        a = self.atom_embedding((atom+1).squeeze(1)) #* math.sqrt(self.n_type_embedding) #.transpose(1,2)\n",
    "        \n",
    "        #x = torch.cat([xyz, mulliken, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)               \n",
    "        #x = torch.cat([xyz, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)               \n",
    "        x = xyz \n",
    "        #x = self.pos_encoder(x).transpose(1,2)\n",
    "\n",
    "        x = torch.cat([x.transpose(1,2), t, a], dim=-1) \n",
    "\n",
    "        mask = (coulomb == 0).unsqueeze(1)\n",
    "        #mask = torch.triu(x.new_ones(max_atoms, max_atoms), diagonal=1).byte()[None,:,:,None]#[None,None] \n",
    "        #print(mask.shape, mask)\n",
    "        x = self.transformer(x, mask).transpose(1,2).contiguous()\n",
    "        t_one_hot = torch.zeros(bs,n_types+1,n_pts,device=type.device,dtype=x.dtype).scatter_(1,type+1, 1.)\n",
    "        \n",
    "        scalar    = self.scalar(torch.cat([x, t_one_hot], dim=1))\n",
    "        magnetic  = self.magnetic(x) \n",
    "        dipole    = self.dipole(x.view(bs,-1))\n",
    "        potential = self.potential(x.view(bs,-1))\n",
    "                \n",
    "        return type,ext,scalar,magnetic,dipole,potential\n",
    "    \n",
    "    def reset(self): pass\n",
    "    \n",
    "    \n",
    "class AtomTransformer2(Module):\n",
    "    def __init__(self,  n_heads,d_model, d_head=None, **kwargs):\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.transformer = Transformer(n_heads=n_heads,d_model=d_model, d_head=d_head, **kwargs)\n",
    "        \n",
    "        self.scalar    = nn.Conv1d(d_model+ n_types + 1, 4, 1)\n",
    "        self.magnetic  = nn.Conv1d(d_model, 9, 1)\n",
    "        self.dipole    = nn.Linear(d_model*max_atoms, 3)\n",
    "        self.potential = nn.Linear(d_model*max_atoms, 1)\n",
    "        \n",
    "        self.n_atom_embedding = d_model//2\n",
    "        self.n_type_embedding = d_model - self.n_atom_embedding - 6 # -3 best model\n",
    "        self.type_embedding = nn.Embedding(len(types)+1,self.n_type_embedding)\n",
    "        self.atom_embedding = nn.Embedding(len(atoms)+1,self.n_atom_embedding)\n",
    "    \n",
    "        #n_pos_encoder = d_model - n_type_embedding - n_atom_embedding\n",
    "        #self.pos_encoder = nn.Sequential(\n",
    "        #    nn.Conv1d(3+1+1,n_pos_encoder, 1), nn.ReLU(), nn.BatchNorm1d(n_pos_encoder),\n",
    "        #)\n",
    "        \n",
    "    def forward(self,xyz,type,ext,atom,mulliken,coulomb, mask_atoms, n_atoms):\n",
    "        bs, _, n_pts = xyz.shape        \n",
    "        t = self.type_embedding((type+1).squeeze(1)) #* math.sqrt(self.n_atom_embedding) #.transpose(1,2)\n",
    "        a = self.atom_embedding((atom+1).squeeze(1)) #* math.sqrt(self.n_type_embedding) #.transpose(1,2)\n",
    "        \n",
    "        x = torch.cat([xyz, mulliken, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)               \n",
    "        #x = torch.cat([xyz, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)               \n",
    "        #x = xyz \n",
    "        #x = self.pos_encoder(x).transpose(1,2)\n",
    "\n",
    "        x = torch.cat([x.transpose(1,2), t, a], dim=-1) \n",
    "\n",
    "        mask = (coulomb == 0).unsqueeze(1)\n",
    "        #mask = torch.triu(x.new_ones(max_atoms, max_atoms), diagonal=1).byte()[None,:,:,None]#[None,None] \n",
    "        #print(mask.shape, mask)\n",
    "        x = self.transformer(x, mask).transpose(1,2).contiguous()\n",
    "        t_one_hot = torch.zeros(bs,n_types+1,n_pts,device=type.device,dtype=x.dtype).scatter_(1,type+1, 1.)\n",
    "        \n",
    "        scalar    = self.scalar(torch.cat([x, t_one_hot], dim=1))\n",
    "        magnetic  = self.magnetic(x) \n",
    "        dipole    = self.dipole(x.view(bs,-1))\n",
    "        potential = self.potential(x.view(bs,-1))\n",
    "                \n",
    "        return type,ext,scalar,magnetic,dipole,potential\n",
    "    \n",
    "    def reset(self): pass\n",
    "    \n",
    "\n",
    "class MultiHeadAttention3(Module):\n",
    "    \"MutiHeadAttention.\"\n",
    "    def __init__(self, n_heads:int, d_model:int, d_head:int=None, resid_p:float=0., attn_p:float=0., bias:bool=False,\n",
    "                 scale:bool=False):\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.n_heads,self.d_head,self.scale = n_heads,d_head,scale\n",
    "\n",
    "        self.attention1 = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "        self.attention2 = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "        self.attention3 = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "        self.attention4 = nn.Linear(d_model, d_model, bias=bias)\n",
    "        self.out = nn.Linear(n_heads * d_head, d_model, bias=bias)\n",
    "        self.drop_att,self.drop_res = nn.Dropout(attn_p),nn.Dropout(resid_p)\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        \n",
    "        self.scale=False\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ln(x + self.drop_res(self.out(self._apply_attention(x, mask=mask, **kwargs))))\n",
    "\n",
    "    def _apply_attention(self, x:Tensor, mask:Tensor=None):\n",
    "        if False:\n",
    "            if mask is not None:\n",
    "                x= x.masked_fill(mask[:,:,0].squeeze(1).unsqueeze(-1).repeat(1, 1, x.size(2)),  0.)\n",
    "        bs,x_len = x.size(0),x.size(1)\n",
    "        wq,wk,wv = self.attention1(x), self.attention2(x), self.attention3(x)\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        wq,wk,wv = wq.permute(0, 2, 1, 3),wk.permute(0, 2, 3, 1),wv.permute(0, 2, 1, 3)\n",
    "        attn_score = torch.matmul(wq, wk)\n",
    "        if self.scale: attn_score.div_(self.d_head ** 0.5)\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask,  -1.0E9).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=-1))\n",
    "        attn_vec = torch.matmul(attn_prob, wv)\n",
    "        return self.attention4(attn_vec.permute(0, 2, 1, 3).contiguous().contiguous().view(bs, x_len, -1))\n",
    "\n",
    "    def _attention_einsum(self, x, mask=None):\n",
    "        # Permute and matmul is a little bit faster but this implementation is more readable\n",
    "        bs,x_len = x.size(0),x.size(1)\n",
    "        wq,wk,wv = torch.chunk(self.attention(x), 3, dim=-1)\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        attn_score = torch.einsum('bind,bjnd->bijn', (wq, wk))\n",
    "        if self.scale: attn_score.mul_(1/(self.d_head ** 0.5))\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask, -float('-inf')).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=2))\n",
    "        attn_vec = torch.einsum('bijn,bjnd->bind', (attn_prob, wv))\n",
    "        return attn_vec.contiguous().view(bs, x_len, -1)\n",
    "\n",
    "    \n",
    "class AtomTransformer3(Module):\n",
    "    def __init__(self,  n_heads,d_model, d_head=None, **kwargs):\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        d_head = ifnone(d_head, d_model//n_heads)\n",
    "        self.transformer = Transformer(n_heads=n_heads,d_model=d_model, d_head=d_head, **kwargs)\n",
    "        \n",
    "        self.scalar    = nn.Conv1d(d_model+ n_types + 1, 4, 1)\n",
    "        self.magnetic  = nn.Conv1d(d_model, 9, 1)\n",
    "        self.dipole    = nn.Linear(d_model*max_atoms, 3)\n",
    "        self.potential = nn.Linear(d_model*max_atoms, 1)\n",
    "        \n",
    "        self.n_atom_embedding = d_model//2\n",
    "        self.n_type_embedding = d_model - self.n_atom_embedding - 3 # -3 best model\n",
    "        self.type_embedding = nn.Embedding(len(types)+1,self.n_type_embedding)\n",
    "        self.atom_embedding = nn.Embedding(len(atoms)+1,self.n_atom_embedding)\n",
    "        \n",
    "        self.dropout = nn.Dropout(0.05)\n",
    "    \n",
    "        #n_pos_encoder = d_model - n_type_embedding - n_atom_embedding\n",
    "        #self.pos_encoder = nn.Sequential(\n",
    "        #    nn.Conv1d(3+1+1,n_pos_encoder, 1), nn.ReLU(), nn.BatchNorm1d(n_pos_encoder),\n",
    "        #)\n",
    "        \n",
    "    def forward(self,xyz,type,ext,atom,mulliken,coulomb, mask_atoms, n_atoms):\n",
    "        bs, _, n_pts = xyz.shape        \n",
    "        t = self.type_embedding((type+1).squeeze(1)) #* math.sqrt(self.n_atom_embedding) #.transpose(1,2)\n",
    "        a = self.atom_embedding((atom+1).squeeze(1)) #* math.sqrt(self.n_type_embedding) #.transpose(1,2)\n",
    "        \n",
    "        #x = torch.cat([xyz, mulliken, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)               \n",
    "        #x = torch.cat([xyz, ext, mask_atoms.float()], dim=1) #* math.sqrt(self.d_model)               \n",
    "        x = xyz \n",
    "        #x = self.pos_encoder(x).transpose(1,2)\n",
    "\n",
    "        x = torch.cat([x.transpose(1,2), t, a], dim=-1) \n",
    "\n",
    "        mask = (coulomb == 0).unsqueeze(1)\n",
    "        #mask = torch.triu(x.new_ones(max_atoms, max_atoms), diagonal=1).byte()[None,:,:,None]#[None,None] \n",
    "        #print(mask.shape, mask)\n",
    "        x = self.transformer(x, mask).transpose(1,2).contiguous()\n",
    "        x = self.dropout(x)\n",
    "        t_one_hot = torch.zeros(bs,n_types+1,n_pts,device=type.device,dtype=x.dtype).scatter_(1,type+1, 1.)\n",
    "        \n",
    "        scalar    = self.scalar(torch.cat([x, t_one_hot], dim=1))\n",
    "        magnetic  = self.magnetic(x) \n",
    "        dipole    = self.dipole(x.view(bs,-1))\n",
    "        potential = self.potential(x.view(bs,-1))\n",
    "                \n",
    "        return type,ext,scalar,magnetic,dipole,potential\n",
    "    \n",
    "    def reset(self): pass\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This callback allows to insert multiple stateful (not averaged) metrics in one pass. Addditionally we could add metrics for train if we want to."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Model instantiation: where's all your TPUs/GPUs when you need a decent hyperparam sweep?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:17.774724Z",
     "start_time": "2019-08-23T16:40:17.762255Z"
    }
   },
   "outputs": [],
   "source": [
    "from fastai.callbacks import *\n",
    "#from fastai.callbacks import SaveModelCallback\n",
    "class SaveModelCustomCallback(TrackerCallback):\n",
    "    \"A `TrackerCallback` that saves the model when monitored quantity is best.\"\n",
    "    def __init__(self, learn:Learner, monitor:str='valid_loss', mode:str='auto', every:str='improvement', name:str='bestmodel'):\n",
    "        super().__init__(learn, monitor=monitor, mode=mode)\n",
    "        self.every,self.name = every,name\n",
    "        if self.every not in ['improvement', 'epoch']:\n",
    "            warn(f'SaveModel every {self.every} is invalid, falling back to \"improvement\".')\n",
    "            self.every = 'improvement'\n",
    "                 \n",
    "    def jump_to_epoch(self, epoch:int)->None:\n",
    "        try: \n",
    "            self.learn.load(f'{self.name}_{epoch-1}', purge=False)\n",
    "            print(f\"Loaded {self.name}_{epoch-1}\")\n",
    "        except: print(f'Model {self.name}_{epoch-1} not found.')\n",
    "\n",
    "    def on_epoch_end(self, epoch:int, **kwargs:Any)->None:\n",
    "        \"Compare the value monitored to its best score and maybe save the model.\"\n",
    "        if self.every==\"epoch\": self.learn.save(f'{self.name}_{epoch}')\n",
    "        else: #every=\"improvement\"\n",
    "            current = self.get_monitor_value()\n",
    "            if current is not None and self.operator(current, self.best):\n",
    "                print(f'Better model found at epoch {epoch} with {self.monitor} value: {current}.')\n",
    "                self.best = current\n",
    "                self.learn.save(f'{self.name}_{epoch}_{self.best}')\n",
    "\n",
    "    def on_train_end(self, **kwargs):\n",
    "        \"Load the best model.\"\n",
    "        if False and self.every==\"improvement\" and (self.learn.path/f'{self.learn.model_dir}/{self.name}.pth').is_file():\n",
    "            self.learn.load(f'{self.name}', purge=False)\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:25.426831Z",
     "start_time": "2019-08-23T16:40:17.776541Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py:25: UserWarning: nn.init.xavier_uniform is now deprecated in favor of nn.init.xavier_uniform_.\n"
     ]
    }
   ],
   "source": [
    "net, learner = None,None\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "net = AtomTransformer(n_layers=6, n_heads=32,d_model=1024,d_inner=4096, act=Activation.LearnedSwish, attn_cls=MultiHeadAttention3)\n",
    "#net = AtomTransformer3(n_layers=6, n_heads=16,d_model=1024,d_inner=4096, attn_cls=MultiHeadAttention3)\n",
    "#net = AtomTransformer(n_layers=6, n_heads=16,d_model=2048,d_inner=2048)\n",
    "#net = AtomTransformer(n_layers=6, n_heads=8, d_model=512,d_inner=2048)\n",
    "\n",
    "#net = AtomTransformerOld(n_layers=6, n_heads=8,d_model=512,d_inner=2048, attn_cls=MultiHeadAttentionOld)\n",
    "\n",
    "#net = AtomTransformer2(n_layers=6, n_heads=12,d_model=768,d_inner=3072)\n",
    "\n",
    "\n",
    "#net = AtomTransformer(n_layers=6, n_heads=12,d_model=768,d_inner=3072)\n",
    "\n",
    "#net = AtomTransformer(n_layers=6, n_heads=16,d_model=1024,d_inner=4096, attn_cls=MultiHeadAttentionOld)\n",
    "\n",
    "learner = Learner(data,net, loss_func=LMAEMaskedLoss(),)\n",
    "\n",
    "learner.callbacks.append(LMAEMetric(learner))\n",
    "\n",
    "for p in learner.model.parameters():\n",
    "    if p.dim() > 1:\n",
    "        nn.init.xavier_uniform(p)\n",
    "        \n",
    "\n",
    "learner.callbacks.append(SaveModelCustomCallback(learner, monitor='👉🏻LMAE👈🏻', mode='min'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:30:01.490730Z",
     "start_time": "2019-08-23T16:29:59.907957Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:25.431404Z",
     "start_time": "2019-08-23T16:40:25.428794Z"
    }
   },
   "outputs": [],
   "source": [
    "#sub_fname = \"loss-5.9943val-2.7766\" # uncomment or set None to skip loading trained net\n",
    "#sub_fname = \"loss-5.7552val-2.9950\"\n",
    "#sub_fname = \"loss-5.9943val-2.7766\"\n",
    "#sub_fname = \"loss-4.9516val-3.0042\"\n",
    "#sub_fname = \"loss-4.0414val-2.7287\"\n",
    "#sub_fname = \"loss-4.9044val-2.5880\"\n",
    "\n",
    "#sub_fname = \"loss-5.0862val-3.0047\"\n",
    "\n",
    "#sub_fname = 'loss-4.9516val-2.8229'\n",
    "#sub_fname = 'bestmodel_62_-2.2235493659973145'\n",
    "#sub_fname = 'loss-5.6540val-3.0131'\n",
    "#sub_fname = 'bestmodel_6_-0.6834144592285156'\n",
    "\n",
    "#sub_fname = 'loss-5.5435val-3.0776'\n",
    "sub_fname = 'loss-5.6017'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:35.077420Z",
     "start_time": "2019-08-23T16:40:25.433358Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/torch/serialization.py:256: UserWarning: Couldn't retrieve source code for container of type LearnedSwish. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
      "/opt/conda/lib/python3.7/site-packages/torch/serialization.py:256: UserWarning: Couldn't retrieve source code for container of type LMAEMaskedLoss. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
      "/opt/conda/lib/python3.7/site-packages/torch/serialization.py:256: UserWarning: Couldn't retrieve source code for container of type AtomTransformer. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
      "/opt/conda/lib/python3.7/site-packages/torch/serialization.py:256: UserWarning: Couldn't retrieve source code for container of type Transformer. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
      "/opt/conda/lib/python3.7/site-packages/torch/serialization.py:256: UserWarning: Couldn't retrieve source code for container of type DecoderLayer. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
      "/opt/conda/lib/python3.7/site-packages/torch/serialization.py:256: UserWarning: Couldn't retrieve source code for container of type MultiHeadAttention3. It won't be checked for correctness upon loading.\n",
      "  \"type \" + obj.__name__ + \". It won't be checked \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attempting to load: loss-5.6017... Loaded\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(f\"Attempting to load: {sub_fname}... \", end=\"\")\n",
    "    learner.load(sub_fname, strict=False,with_opt=False)\n",
    "    print(\"Loaded\")\n",
    "except Exception as e:\n",
    "    print(\"NOT loaded! \", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-19T06:40:14.614754Z",
     "start_time": "2019-08-19T06:40:14.609905Z"
    }
   },
   "outputs": [],
   "source": [
    "if False:\n",
    "    for name, param in learner.model.named_parameters():\n",
    "\n",
    "        if 'transformer.layers.5.mhra.attention4' not in name:\n",
    "            #param.requires_grad = False\n",
    "            pass\n",
    "        else:\n",
    "            print(name)\n",
    "            try:\n",
    "                nn.init.xavier_uniform_(param)\n",
    "            except:\n",
    "                pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:35.083371Z",
     "start_time": "2019-08-23T16:40:35.079634Z"
    }
   },
   "outputs": [],
   "source": [
    "learner = learner.to_parallel() #b/c it NaNs loss (probably would need to change fp16 settings)\n",
    "data.batch_size = int(4096//2)\n",
    "data.batch_size = data.batch_size // 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Real loss func. Need to test different auxiliary tasks weights: `magnetic_w`, `dipole_w`, `potential`, weights of indivial `lmae`s: `types_w` and maybe `input_transform_w` and `feature_transform_w`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-17T06:02:27.635352Z",
     "start_time": "2019-08-17T06:02:27.626543Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:35.089578Z",
     "start_time": "2019-08-23T16:40:35.085523Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.loss_func = LMAEMaskedLoss(magnetic_w=0, dipole_w=0, potential_w=0, \n",
    "                                  types_w = [1] + [1] * (n_types-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-19T13:55:31.203434Z",
     "start_time": "2019-08-19T13:53:36.687673Z"
    }
   },
   "outputs": [],
   "source": [
    "#learner.data.batch_size = 512\n",
    "#learner.opt_func = RAdam\n",
    "learner.lr_find()\n",
    "learner.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T16:40:35.104301Z",
     "start_time": "2019-08-23T16:40:35.091786Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[functools.partial(<class 'fastai.basic_train.Recorder'>, add_time=True, silent=False),\n",
       " __main__.ShowGraph]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class ShowGraph(LearnerCallback):\n",
    "    \"Update a graph of learner stats and metrics after each epoch.\"\n",
    "    def on_epoch_end(self, n_epochs:int, last_metrics:MetricsList, **kwargs)->bool:\n",
    "        \"If we have `last_metrics` plot them in our pbar graph\"\n",
    "        if last_metrics is not None and last_metrics[0] is not None:\n",
    "            rec = self.learn.recorder\n",
    "            iters = range_of(rec.losses)\n",
    "            val_iter = np.array(rec.nb_batches).cumsum()\n",
    "            x_bounds = (0, (n_epochs - len(rec.nb_batches)) * rec.nb_batches[-1] + len(rec.losses))\n",
    "            y_bounds = (min((min(Tensor(rec.losses)), min(Tensor(rec.val_losses)))), \n",
    "                        max((max(Tensor(rec.losses)), max(Tensor(rec.val_losses)))))\n",
    "            rec.pbar.update_graph([(iters, rec.losses), (val_iter, rec.val_losses)], x_bounds, y_bounds)\n",
    "        return {}\n",
    "learner.callback_fns.append(ShowGraph)\n",
    "learner.callback_fns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-23T16:38:57.885Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='53' class='' max='100', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      53.00% [53/100 12:05:10<10:43:04]\n",
       "    </div>\n",
       "    \n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>👉🏻LMAE👈🏻</th>\n",
       "      <th>lmae0</th>\n",
       "      <th>lmae1</th>\n",
       "      <th>lmae2</th>\n",
       "      <th>lmae3</th>\n",
       "      <th>lmae4</th>\n",
       "      <th>lmae5</th>\n",
       "      <th>lmae6</th>\n",
       "      <th>lmae7</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>-2.231656</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:49</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>-2.550897</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:38</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>-2.752684</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>-2.915701</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>-3.047175</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>-3.165217</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>-3.258894</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>-3.339325</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>-3.399309</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>-3.447084</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>-3.469729</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>11</td>\n",
       "      <td>-3.500457</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>-3.517563</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>13</td>\n",
       "      <td>-3.500508</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>14</td>\n",
       "      <td>-3.522196</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:43</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>-3.493659</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>16</td>\n",
       "      <td>-3.493721</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17</td>\n",
       "      <td>-3.480302</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>18</td>\n",
       "      <td>-3.464704</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>19</td>\n",
       "      <td>-3.449000</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>-3.433270</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>21</td>\n",
       "      <td>-3.378308</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>22</td>\n",
       "      <td>-3.369914</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>23</td>\n",
       "      <td>-3.351868</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:46</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24</td>\n",
       "      <td>-3.346385</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>25</td>\n",
       "      <td>-3.319162</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>26</td>\n",
       "      <td>-3.328486</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>27</td>\n",
       "      <td>-3.308582</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>28</td>\n",
       "      <td>-3.314197</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:43</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>29</td>\n",
       "      <td>-3.299612</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>-3.326965</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>31</td>\n",
       "      <td>-3.319751</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>32</td>\n",
       "      <td>-3.219625</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>33</td>\n",
       "      <td>-3.344137</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:43</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>34</td>\n",
       "      <td>-3.357655</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>35</td>\n",
       "      <td>-3.352365</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>36</td>\n",
       "      <td>-3.359392</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>37</td>\n",
       "      <td>-3.388980</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:41</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>38</td>\n",
       "      <td>-3.362045</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>39</td>\n",
       "      <td>-3.405196</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>-3.425964</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>41</td>\n",
       "      <td>-3.396528</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>42</td>\n",
       "      <td>-3.448561</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>43</td>\n",
       "      <td>-3.432447</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>44</td>\n",
       "      <td>-3.479277</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>45</td>\n",
       "      <td>-3.488188</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>46</td>\n",
       "      <td>-3.495797</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>47</td>\n",
       "      <td>-3.534294</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>48</td>\n",
       "      <td>-3.554801</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:44</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>49</td>\n",
       "      <td>-3.556508</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>-3.565556</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:39</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>51</td>\n",
       "      <td>-3.607949</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>52</td>\n",
       "      <td>-3.562529</td>\n",
       "      <td>#na#</td>\n",
       "      <td>13:40</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>\n",
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='696' class='' max='877', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      79.36% [696/877 10:50<02:49 -3.6698]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "#learner.data.batch_size = 512*2\n",
    "#learner.data.batch_size = 512\n",
    "#learner.opt_func = RAdam\n",
    "#learner.opt_func = partial(torch.optim.Adam, betas=(0.9,0.99), eps=1e-2)\n",
    "learner.data.batch_size =1024-128\n",
    "learner.fit_one_cycle(100, 1.0E-04)#, moms=(0.75,0.70))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T06:03:37.370Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder.plot_lr(show_moms=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:15:20.112334Z",
     "start_time": "2019-08-24T16:15:19.141497Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYAAAAEGCAYAAABsLkJ6AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8QZhcZAAAgAElEQVR4nO3deXxU1d3H8c8vCwQCYQ0KRAzgwr6EiLsCIoKoFKVWqk8Va2lp+3TR2qLWpVUrVatWa2vdQPsouFC1FSmiQnGXRUSUVQgSQAhRCBBCtvP8MTfDhCSQQGbuJPf7fr3mlZlzt9+9XOY359xzzzXnHCIiEjwJfgcgIiL+UAIQEQkoJQARkYBSAhARCSglABGRgEryO4C6aN++vcvMzPQ7DBGRBmXx4sXbnXPpB5Y3qASQmZnJokWL/A5DRKRBMbMN1ZWrCUhEJKCUAEREAkoJQEQkoBrUNQARaRxKSkrIzc2lqKjI71AalZSUFDIyMkhOTq7V/L4mADN7ErgA2Oac6+NnLCISO7m5ubRs2ZLMzEzMzO9wGgXnHPn5+eTm5tK1a9daLeN3E9A0YKTPMYhIjBUVFdGuXTt9+dcjM6Ndu3Z1qlX5mgCccwuAr/2MQUT8oS//+lfXY+p3DeCQzGyimS0ys0V5eXmHtY43V2zlr/PX1nNkIiINW9wnAOfco865bOdcdnp6lRvZamX+qjwef3t9PUcmIg1Vfn4+AwYMYMCAARx99NF07tw5/Lm4uLhW65gwYQKrVq2KcqTRFYheQAkGZeV68I2IhLRr146lS5cCcNttt9GiRQt+9atfVZrHOYdzjoSE6n8nT506NepxRlvc1wDqQ0KCUa4nn4nIIaxdu5ZevXpx+eWX07t3b7Zs2cLEiRPJzs6md+/e/P73vw/Pe8YZZ7B06VJKS0tp3bo1kydPpn///px66qls27bNx72oPb+7gU4HhgDtzSwXuNU590R9byfBDH3/i8Sn3/37Mz7fXFCv6+zVKY1bL+x9WMuuXLmSp59+muzsbACmTJlC27ZtKS0tZejQoYwbN45evXpVWmbnzp2cffbZTJkyhWuvvZYnn3ySyZMnH/F+RJvfvYDGO+c6OueSnXMZ0fjyBzUBiUjtde/ePfzlDzB9+nSysrLIyspixYoVfP7551WWadasGaNGjQJg0KBB5OTkxCrcIxKMawBqAhKJW4f7Sz1aUlNTw+/XrFnDn//8Zz766CNat27NFVdcUW0/+yZNmoTfJyYmUlpaGpNYj1QwrgGYEoCI1F1BQQEtW7YkLS2NLVu2MGfOHL9DqleBqAEkmqEWIBGpq6ysLHr16kWPHj049thjOf300/0OqV6Za0C/jLOzs93hPBDmvtdX8dC8tay/a3QUohKRulqxYgU9e/b0O4xGqbpja2aLnXPZB84biCYg83oBNaRkJyISbYFIAIkJofEx1AwkIrJfIBKA9/2vrqAiIhGCkQDCNQAlABGRCsFIAKYEICJyoEAkgETTNQARkQMFIgFUPCNBNQARARg6dGiVm7oeeOABJk2aVOMyLVq0AGDz5s2MGzeu2nmGDBnCobqqP/DAAxQWFoY/n3/++ezYsaO2oderQCSAcBOQqgAiAowfP54ZM2ZUKpsxYwbjx48/5LKdOnXixRdfPOxtH5gAXnvtNVq3bn3Y6zsSgUgA6gYqIpHGjRvHrFmzwg9/ycnJYfPmzQwcOJBzzjmHrKws+vbtyyuvvFJl2ZycHPr06QPA3r17ueyyy+jZsydjx45l79694fkmTZoUHkb61ltvBeDBBx9k8+bNDB06lKFDhwKQmZnJ9u3bAbjvvvvo06cPffr04YEHHghvr2fPnvzgBz+gd+/ejBgxotJ2jkQghoJQN1CRODZ7Mnz1af2u8+i+MGpKjZPbtm3L4MGDmT17NmPGjGHGjBlceumlNGvWjJdeeom0tDS2b9/OKaecwkUXXVTjs3b/9re/0bx5c1asWMGyZcvIysoKT7vzzjtp27YtZWVlnHPOOSxbtoyf/exn3HfffcybN4/27dtXWtfixYuZOnUqH374Ic45Tj75ZM4++2zatGnDmjVrmD59Oo899hiXXnopM2fO5IorrjjiwxSIGkBFN1DdCSwiFSKbgSqaf5xz3HjjjfTr14/hw4ezadMmtm7dWuM6FixYEP4i7tevH/369QtPe/7558nKymLgwIF89tln1Q4jHemdd95h7NixpKam0qJFCy6++GLefvttALp27cqAAQOA+h1uOiA1gFACKFMCEIk/B/mlHk1jxozhl7/8JUuWLKGwsJBBgwYxbdo08vLyWLx4McnJyWRmZlY7/POhrF+/nnvvvZeFCxfSpk0brrrqqsNaT4WmTZuG3ycmJtZbE1AwagDhXkD+xiEi8aNFixYMHTqUq6++Onzxd+fOnXTo0IHk5GTmzZvHhg0bDrqOs846i2effRaA5cuXs2zZMiA0jHRqaiqtWrVi69atzJ49O7xMy5Yt2bVrV5V1nXnmmbz88ssUFhayZ88eXnrpJc4888z62t1qBaoGoF5AIhJp/PjxjB07NtwUdPnll3PhhRfSt29fsrOz6dGjx0GXnzRpEhMmTKBnz5707NmTQYMGAdC/f38GDhxIjx49OOaYYyoNIz1x4kRGjhxJp06dmDdvXrg8KyuLq666isGDBwNwzTXXMHDgwKg+XSwQw0HPXJzLdS98wn+vH8Kx7VIPvYCIRJWGg44eDQd9AHUDFRGpKhAJwNQNVESkCl8TgJmNNLNVZrbWzCZHazuJGg1UJO40pObnhqKux9S3BGBmicDDwCigFzDezHpFY1sVg8GpBiASH1JSUsjPz1cSqEfOOfLz80lJSan1Mn72AhoMrHXOrQMwsxnAGODgd0schooagBKASHzIyMggNzeXvLw8v0NpVFJSUsjIyKj1/H4mgM7AxojPucDJB85kZhOBiQBdunQ5rA0pAYjEl+TkZLp27ep3GIEX9xeBnXOPOueynXPZ6enph7WOigRQqgQgIhLmZwLYBBwT8TnDK6t3SQmh3dRFYBGR/fxMAAuB482sq5k1AS4D/hWNDXnf/5SWKQGIiFTw7RqAc67UzH4KzAESgSedc59FY1uqAYiIVOXrWEDOudeA16K9ncSKGoCuAYiIhMX9ReD6kFhRA1ACEBEJC0QCSFIvIBGRKgKRAMIPhCkv9zkSEZH4EYgEkJRYkQB8DkREJI4EIgFU1ABKVQMQEQkLRAJI0migIiJVBCIBhIeC0I1gIiJhgUoAqgGIiOwXiASgbqAiIlUFIgEkaDhoEZEqApEAkpQARESqCEQCUA1ARKSqQCQA1QBERKoKRALYfyOYEoCISIVAJADVAEREqgpEAtBD4UVEqgpEAjAzEkwJQEQkUiASAIQeC1mmO4FFRMICkwASElQDEBGJFJgEkJSQoAQgIhIhMAkgMcGUAEREIviSAMzs22b2mZmVm1l2LLaZmGB6IIyISAS/agDLgYuBBbHaYKgGEKutiYjEvyQ/NuqcWwGh7pmxkmimh8KLiESI+2sAZjbRzBaZ2aK8vLzDXo9qACIilUWtBmBmbwBHVzPpJufcK7Vdj3PuUeBRgOzs7MO+ihtKAMoAIiIVopYAnHPDo7Xuw5GUYBoMTkQkQtw3AdWXxATTM4FFRCL41Q10rJnlAqcCs8xsTrS3mZhglJYpAYiIVPCrF9BLwEux3KZqACIilQWqCUjXAERE9gtUAtBQECIi+wUmASQpAYiIVBKYBKCLwCIilQUmATRJSqRYtwKLiIQFJgE0TUpgX6kSgIhIhcAkgCZJCewrLfM7DBGRuBGYBNA0KYF9JaoBiIhUCFACSFQTkIhIhAAlgASK1QQkIhIWnASQrIvAIiKRgpMAEkMJwGk8IBERIEgJIDkRQPcCiIh4gpMAkkK7WqSeQCIiQIASQPMmoZGv9xbrQrCICAQoAaQ1CyWAgqISnyMREYkPwUkAKckA7FICEBEBgpQAmoUSQMHeUp8jERGJD8FJAClqAhIRiRScBBCuASgBiIhAgBJAS68GsFMJQEQE8CkBmNk9ZrbSzJaZ2Utm1jra22yalEiLpknk7ymO9qZERBoEv2oAc4E+zrl+wGrghlhs9Ki0pmwr2BeLTYmIxD1fEoBz7nXnXEV3nA+AjFhs96i0FL4qKIrFpkRE4l48XAO4Gphd00Qzm2hmi8xsUV5e3hFt6Ki0FL7aqQQgIgJRTABm9oaZLa/mNSZinpuAUuCZmtbjnHvUOZftnMtOT08/opg6t27GVwVFFGtYaBERkmozk5l1B3Kdc/vMbAjQD3jaObejpmWcc8MPsc6rgAuAc1yMxmju2j6VsnLHxm8K6Z7eIhabFBGJW7WtAcwEyszsOOBR4Bjg2cPdqJmNBH4NXOScKzzc9dRVt/RUANbl7YnVJkVE4lZtE0C5d9F2LPCQc+56oOMRbPcvQEtgrpktNbNHjmBdtdbN+9W/dtvuWGxORCSu1aoJCCgxs/HAlcCFXlny4W7UOXfc4S57JFp5dwP/8T8rmTSkux8hiIjEjdrWACYApwJ3OufWm1lX4B/RC0tERKKtVgnAOfe5c+5nzrnpZtYGaOmc+2OUY4uKYT06AKg7qIgEXq0SgJnNN7M0M2sLLAEeM7P7ohtadJzX+ygAFqw5snsKREQauto2AbVyzhUAFxPq/nkycNBunvHqov6dAXjj860+RyIi4q/aJoAkM+sIXAq8GsV4oq5Zk0RSmySy5MtviNHtByIicam2CeD3wBzgC+fcQjPrBqyJXljRdfu3+rB9dzGLNnzjdygiIr6p7UXgF5xz/Zxzk7zP65xzl0Q3tOg5r/fRNG+SyMzFuX6HIiLim9peBM7wxu3f5r1mmllMRvCMhtSmSYzq05FZy7ZQVFLmdzgiIr6obRPQVOBfQCfv9W+vrMG6JKszu/aVMlcXg0UkoGqbANKdc1Odc6XeaxpwZENz+uyUbu3o1CqFmUvUDCQiwVTbBJBvZleYWaL3ugLIj2Zg0ZaQYIzN6syC1Xls26WbwkQkeGqbAK4m1AX0K2ALMA64KkoxxczFWRmUO/jX0s1+hyIiEnO17QW0wTl3kXMu3TnXwTn3LaDB9gKq0D29BZ1bN+OOWSt0T4CIBM6RPBHs2nqLwkfn9goNDfHdxz70ORIRkdg6kgRg9RaFj26+oBcA76/LZ/12PShGRILjSBJAo2gzSUww7v12fwCG3jvf32BERGLooAnAzHaZWUE1r12E7gdoFMYN2n9P287CEh8jERGJnYMmAOdcS+dcWjWvls652j5NrEF48UenAnD/G6t9jkREJDaOpAmoURnYpQ0A097LYd7KbT5HIyISfUoAnsQE466L+wIwYdpCNuTrgrCING5KABHGD+4Sfn/2PfN1b4CINGq+JAAzu93MlpnZUjN73czi5oJyzpTR4fe53+z1MRIRkejyqwZwj/d8gQGEnjB2i09xVOulH58GwJl3z6O8XLUAEWmcfEkA3vOFK6QSZ/cU9MtoHX7f7cbXfIxERCR6fLsGYGZ3mtlG4HIOUgMws4lmtsjMFuXl5cUktsQE4/0bhoU/Z06eFZPtiojEUtQSgJm9YWbLq3mNAXDO3eScOwZ4BvhpTetxzj3qnMt2zmWnp8fuEQQdWzVj6S3nhj8v+VLPDxaRxiVqCcA5N9w516ea1ysHzPoMcTqyaOvmTXjrurMBuPiv7+l6gIg0Kn71Ajo+4uMYYKUfcdRGt/QW+9/f+BrXPrfUx2hEROqPX9cApnjNQcuAEcDPfYqjVtbfdX74/T8/3sSmHeoe+mV+Ia99uoXbX/2cj9Z/XeflS8vKdZ+FiM/86gV0idcc1M85d6FzbpMfcdSWmZEzZTTXn3ciAKdPecvniGJjX2kZn2/e32HrrZVb2VFYjHOOs+6Zx4+fWcIT76zn0r+/T4+bZ1dadmHO1zy38Esg9GU/7d317Cst45Wlm9haUMRxN83miXfWx3R/RKQya0i/wrKzs92iRYt8235RSRk9bv5PpbLIG8cam1teWc7T729gysV9Oa5DC8Y98j4nZbbhofFZnHLXm1Xm790pjYv6d2JYjw6ce/8CAD6++VwG3j4XgKPTUviqoPLzl3OmjKa0rJztu4s5ulVK9HdKJIDMbLFzLrtKuRJA3XyRt5tz/vTf8OebL+jF98/o6mNE0TPygQWs/GpXva0vLSWJgqLS8Of2LZqy6LfDGfan+azL28O0CSdx5vHpdL/xNb6TfQzNmiRy3YgTaJmSXGk9zjmcg4QEwzmHWaN4NpFI1CgB1LMLH3qHTzftBOCotKa8dd0QUpvG/wjZBUUlLFidxwX9QqNv7Csto2lSIhB6FkJJeTkb8gv5cH0+989dTUmZ/+fHE1dmc1LXtvxh1gomj+rBgN+HahR//59BTH13Pcs3FfDOb4byk2eXkNWlDdeeewKfbS6gT+dWPkcuEh+UAOqZc46uN1S+S3jhTcNJb9nUp4iqV1xaztaCIo5p2xyAH/5jEXM+28pb153NHbNW8FYjHPq6ZUoSu4pKueq0TKa9lwPA278eGj4GACVl5SzL3cGgY9v6FKVI7CgBREFRSRn3z13N3xesq1T+yBWDGNnnaJ+iqqyiHf+TW0fwwqKN3DFrRZ3XkWDQuU0zNn4d6v10z7h+zFi4kdsu7M2CNXlceVomn23aSdf2qbRv0ZRNO/ayYE0eF/XvxI+fWcJdF/clf3cxYx5+t753r9Z+Mfx4vn9GV/aVljP13fU89vZ6ikvLmXhWN7bv3sdPhh7HYwvWccuFvUhKSOC9L7Yz5MQOvsUrUp+UAKKorNzR/YAxg6ZedRJDe8T2C2RvcRlf5O1m7bbd/OK5pTx19WCufPKjw1rXzEmn8uaKbYzu15HenaLXlLJ5x17ydxeze18pb6/J46/zvwDgNyN7MLzn/ovJsZJ9bBsWbQjd9f2DM7ty0+heOOcoKimnWZPEmMYiUl+UAKKsqKSMIffMr9LLBeDN686me8QNZfXtL2+tIW/XPp56f0Odl73twl50adecISd0ICEhfi+m5u/ex39X57F5x156d2rFhGkLeXD8QLYVFJHWLJlfv7gsKtt94spsVm3dxd3/WcX8Xw0hs31qVLYjEk1KADF0sMHjrj/vRMYNyiApwWib2qRKD5ayckdiNV/EWwuKeG7hRu6bG3pmcWa75nz35C784bW63UT9j+8P5ozj2jfanjNrt+3mvS+289zCjQA8dfVgHn97PY/894sjXndyorHmzv03BX6zp5j8Pfs4rkPLI163SDQpAfiktiOJPvzdLH7y7JLw596d0vhsc8FBlji44T07cN2IE+ncphlpB3SjDKKCohL63fY6D40fyI7CYj5Y9zVd2jXnb/PrnhiuPfcEhvXowAUPvQPAu5OH0bl1M6DmBC7iJyUAH81buY0eHVty6l3Ru4N4xsRTGHRsG5IT9ZTPutiycy/tWzQlOTGB0+56k807i5hweibPL9zInuKyOq2r4ka3iWd148bzewKhXljrtu+mx9FpFJWUkZKs6wgSe0oAccQ5x52zVvD4YQyFMHPSaQw6tg1FJWUkJybo12aU5e/ex+YdRXy6aSc3vvRprZf71YgTuPf11eHPF2d15p9LNvHU1YPp1TEt7roLS+OmBNAAOecod+hLPo6UlTueW7iRbump/PK5pWzZWfWif208OH4gP5v+Me/fMIx2qU35eo+GwpDoUQIQiaKycsfMJbk8/vY6Vm/dXevlLs3O4PlFuQAsuflccr8prPRIUpH6oAQgEiM7Cou59/VVbCvYx5gBnbnuhaUUlZTXevmVt48kwYyH3lrDpdnHsHNviYa1kCOiBCASB3bvK6XPrXPqvNwLPzqVtqlNono/iTReSgAiccQ5xxsrtvHBuvw6PRfh24My6JbegklDukcxOmlslABE4lhRSRnvr8vn+heWsX33vlotc9lJx/C7Mb3ZV1pOolmDGI1W/KEEINLALFidx/dqOZZTx1YpzL9+CHuLy2jdvEmUI5OGRglApAFzzjH9o421uhdh6Inp7Coq5cVJp8UgMmkIlABEGpFtBUUM/kPVx3IeKGfKaB6et5azT0hXT6IAUwIQaaReWbqJn89Yesj5JpyeySVZGRyVlkLr5skaNiRA4jIBmNl1wL1AunNu+6HmVwIQObiXPs6ldfMm/Pal5WzasbdWy6y5c5SSQSNXUwLw7V/dzI4BRgBf+hWDSGMzdmAGQ0/swLuTh7HqjpHcMKrHIZd56M01fL2nmIbUGiD1w7cagJm9CNwOvAJkqwYgEl3rt+9h6L3zDzrPHy/pS/f0FnRu04wOLVM0DlUjEVdNQGY2BhjmnPu5meVwkARgZhOBiQBdunQZtGFD3Z96JSKV5X5TyKT/W8Knm3bWOM8Pz+7GDaN6xjAqiZaYJwAzewOo7snoNwE3AiOcczsPlQAiqQYgUv8Ki0v58TNLmL8qr9rprZsnc8+4/pzb66gYRyb1JW5qAGbWF3gTKPSKMoDNwGDn3FcHW1YJQCR6du8rZdL/LSZ/dzGfb6n+aXT3f6c/YwdmxDgyOVJxkwCqBKAagEjcKS93zFh48BvPIp98JvEt7noBiUj8SkgwvntyF3KmjObFH51a7TyPLljHD/+xiIKikhhHJ/XF9xpAXagGIOKf4tJyfvn8UmYt21LjPB/ddA4dWurJZvEmbpuA6kIJQCQ+FBaX0uuWmp9r8N7kYXRq3SyGEcnBKAGISFRs/LqQM++eV+204T2P4vErq3zvSIwpAYhIVDnnGP/YB3yw7usq0/p2bsXvxvQmq0sbHyITJQARiYkN+XuYMnsls5dX36t75e0jSUlOjHFUwaYEICIxt3jD11zyt/ernbbuD+eToKEmYkIJQER89b/TP+bfn2yuUp6WksTSW0YoGUSREoCIxIVdRSX0ve31KuV/+nZ/Plr/NXeM7aPhqetZTQlAT5EWkZhqmZLMuj+cz1cFRfxixlI+ygldNL7uhU8AyGyfyqQh3f0MMTBUAxARXx3sOsEb157FcR1axjiixkdNQCIS9/reOodd+0qrlM+cdCqDjm3rQ0SNgxKAiDQIzjle+/QrfvLskirT1IX08GgwOBFpEMyM0f06svL2kVWm9bj5P2ROnkVJWbkPkTU+SgAiEpdSkhPJmTKanCmjuf68EytNO/6m2Tz1Xg43v7yc8vKG04oRb9QEJCINRnXPNe6f0YpXfnqGPwE1EGoCEpEGr2v7VJ79wcmVyj7J3Unm5FlkTp7FihqeZCbVUw1ARBqsa55axBsrtlYqe+x72RSVlHF+344k6u5iQL2ARKQR21tcRs9b/lOprHPrZrw7eZhPEcUXNQGJSKPVrEkiH910TqWyTTv2hpuGdhQW+xRZfFMCEJFGoUPLlHCvoWE9OlSaNuD3cylTb6Eq1AQkIo1SaVk5v5n5KTOX5FYqnzbhJIac2KGGpRqnuGoCMrPbzGyTmS31Xuf7EYeINF5JiQn86dL+PPa9yt97V01dSObkWWz8utCnyOKHLzUAM7sN2O2cu7cuy6kGICJHInPyrCplOVNG+xBJbMVVDUBExA/r76ra2FBxofj5hRt9iMhffiaAn5rZMjN70sz0pGgRiTozC18ovuKULpWm/XrmMrYWFLG3uMyn6GIvak1AZvYGcHQ1k24CPgC2Aw64HejonLu6hvVMBCYCdOnSZdCGDRuiEq+IBE9xaTkLVudxzdOVm5avOi2T2y7q7VNU9S9ubwQzs0zgVedcn0PNq2sAIhIN3+wpZuDtc6uUP/CdAYwZ0Amzhn1HcVxdAzCzjhEfxwLL/YhDRASgTWqTcNPQxVmdw+W/eG4pXW94zcfIosuvZwLfbWYDCDUB5QA/9CkOEZFK7h3Xn6PTUvjr/C/CZRW9h64+vSu3XNjLr9Dqne9NQHWhJiARiaVXlm7i5zOWVipbdcdImiQmNKhmoZqagPyqAYiIxL0xAzqT1iyZbu1TueDBd9i1r5QTfxsadO76807kJ0OP8znCI6P7AEREDmLoiR04tl0qb/7q7Erl98xZxbtrtzNzcW6DfSqZmoBEROroh/9YxJzPKj+HIJ7vKI6rXkAiIg3Z3/+nyncpmZNnMeSeeSzftNOHiA6PEoCIyGGo6DZ6ywX7ewXl5BdywUPvNJgmITUBiYgcoZVfFfDUezlM/6jqeELx0DSkJiARkSjpcXQad13cj8W/HV5l2qMLvmDPvlIfojo01QBEROrZzr0l3P7q57y4uPLDaNbeOYqkxNj/7lYNQEQkRlo1S+aecf2qlB9302yy75jLf5Z/5UNUVakGICISZdt37yP7jjcqlS25+VzapjaJyfZVAxAR8Un7Fk25/zv9K5Vl3T43/DCaXUUlvsSlGoCISAyVlTu631jzCKOr7xhFk6T6/W0et88DqAslABFpTPaVloXHFqpOv4xWPPo/2by4eCM/HXb8YW9Hg8GJiMSZpkmJ4fsEcr8p5Iw/zqs0fVnuTk65600AEhMSmDSke71uXwlARCQOZLRpTs6U0RQWl7Io5xt+M3MZW3YWhadfc2bXet+mEoCISBxp3iSJs05I5/0bzon6ttQLSEQkoJQAREQCSglARCSglABERAJKCUBEJKCUAEREAkoJQEQkoJQAREQCqkGNBWRmecCGw1y8PbC9HsNpLHRcqtIxqUrHpHoN5bgc65xLP7CwQSWAI2Fmi6obDCnodFyq0jGpSsekeg39uKgJSEQkoJQAREQCKkgJ4FG/A4hTOi5V6ZhUpWNSvQZ9XAJzDUBERCoLUg1AREQiKAGIiARUIBKAmY00s1VmttbMJvsdT30zs2PMbJ6ZfW5mn5nZz73ytmY218zWeH/beOVmZg96x2OZmWVFrOtKb/41ZnZlRPkgM/vUW+ZBM7PY72ndmVmimX1sZq96n7ua2YfefjxnZk288qbe57Xe9MyIddzgla8ys/MiyhvceWVmrc3sRTNbaWYrzOxUnSdgZr/0/u8sN7PpZpYSiHPFOdeoX0Ai8AXQDWgCfAL08juuet7HjkCW974lsBroBdwNTPbKJwN/9N6fD8wGDDgF+NArbwus8/628d638aZ95M1r3rKj/N7vWh6ba4FngVe9z88Dl3nvH62aLd8AAAdXSURBVAEmee9/DDzivb8MeM5738s7Z5oCXb1zKbGhnlfAU8A13vsmQOugnydAZ2A90CziHLkqCOdKEGoAg4G1zrl1zrliYAYwxueY6pVzbotzbon3fhewgtBJPYbQf3i8v9/y3o8BnnYhHwCtzawjcB4w1zn3tXPuG2AuMNKbluac+8CFzvSnI9YVt8wsAxgNPO59NmAY8KI3y4HHpOJYvQic480/BpjhnNvnnFsPrCV0TjW488rMWgFnAU8AOOeKnXM7CPh54kkCmplZEtAc2EIAzpUgJIDOwMaIz7leWaPkVUcHAh8CRznntniTvgKO8t7XdEwOVp5bTXm8ewD4NVDufW4H7HDOlXqfI/cjvO/e9J3e/HU9VvGsK5AHTPWaxR43s1QCfp445zYB9wJfEvri3wksJgDnShASQGCYWQtgJvAL51xB5DTvF1lg+vya2QXANufcYr9jiSNJQBbwN+fcQGAPoSafsKCdJwDeNY8xhBJkJyAVGOlrUDEShASwCTgm4nOGV9aomFkyoS//Z5xz//SKt3rVcry/27zymo7JwcozqimPZ6cDF5lZDqEq9zDgz4SaMZK8eSL3I7zv3vRWQD51P1bxLBfIdc596H1+kVBCCPJ5AjAcWO+cy3POlQD/JHT+NPpzJQgJYCFwvHdFvwmhizb/8jmmeuW1Pz4BrHDO3Rcx6V9ARQ+NK4FXIsq/5/XyOAXY6TUBzAFGmFkb71fRCGCON63AzE7xtvW9iHXFJefcDc65DOdcJqF/87ecc5cD84Bx3mwHHpOKYzXOm9955Zd5PT+6AscTutDZ4M4r59xXwEYzO9ErOgf4nACfJ54vgVPMrLkXd8Vxafznit9XoWPxItSbYTWhK/E3+R1PFPbvDELV9mXAUu91PqF2yTeBNcAbQFtvfgMe9o7Hp0B2xLquJnTxai0wIaI8G1juLfMXvLvIG8ILGML+XkDdCP2nXAu8ADT1ylO8z2u96d0ilr/J2+9VRPRqaYjnFTAAWOSdKy8T6sUT+PME+B2w0ov9H4R68jT6c0VDQYiIBFQQmoBERKQaSgAiIgGlBCAiElBKACIiAaUEICISUEoAElfMrMzMlprZJ2a2xMxOO8T8rc3sx7VY73wza7AP744GM5tmZuMOPac0VkoAEm/2OucGOOf6AzcAdx1i/taERmeMSxF3korEHSUAiWdpwDcQGufIzN70agWfmlnFaIpTgO5ereEeb97fePN8YmZTItb3bTP7yMxWm9mZ3ryJZnaPmS200Jj3P/TKO5rZAm+9yyvmj2RmOWZ2t7etj8zsOK98mpk9YmYfAndbaLz9l731f2Bm/SL2aaq3/DIzu8QrH2Fm73v7+oI3xhNmNsVCz3xYZmb3emXf9uL7xMwWHGKfzMz+YqFx6d8AOtTnP5Y0PPp1IvGmmZktJXS3ZUdCY/gAFAFjnXMFZtYe+MDM/kVoMLM+zrkBAGY2itDAXic75wrNrG3EupOcc4PN7HzgVkJjwHyf0BAHJ5lZU+BdM3sduJjQ8AZ3mlkioSGCq7PTOdfXzL5HaPTRC7zyDOA051yZmT0EfOyc+5aZDSM0TPIA4OaK5b3Y23j79ltguHNuj5n9BrjWzB4GxgI9nHPOzFp727kFOM85tymirKZ9GgicSGjc+qMIDXfwZK3+VaRRUgKQeLM34sv8VOBpM+tDaFiCP5jZWYSGd+7M/mGLIw0HpjrnCgGcc19HTKsYJG8xkOm9HwH0i2gLb0VoDJeFwJMWGmTvZefc0hrinR7x9/6I8hecc2Xe+zOAS7x43jKzdmaW5sV6WcUCzrlvLDSKaS9CX9oQeoDI+4SGHC4CnrDQ081e9RZ7F5hmZs9H7F9N+3QWMN2La7OZvVXDPklAKAFI3HLOve/9Ik4nNJZKOjDIOVdioVE+U+q4yn3e3zL2n/sG/K9zbs6BM3vJZjShL9j7nHNPVxdmDe/31DG28GYJPWxlfDXxDCY0UNk44KfAMOfcj8zsZC/OxWY2qKZ98mo+ImG6BiBxy8x6EHqcXj6hX7HbvC//ocCx3my7CD0Gs8JcYIKZNffWEdkEVJ05wCTvlz5mdoKZpZrZscBW59xjhJ4ollXD8t+J+Pt+DfO8DVzurX8IsN2FntcwF/hJxP62AT4ATo+4npDqxdQCaOWcew34JdDfm97dOfehc+4WQg97OaamfQIWAN/xrhF0BIYe4thII6cagMSbimsAEPole6XXjv4M8G8z+5TQaJYrAZxz+Wb2rpktB2Y75643swHAIjMrBl4DbjzI9h4n1By0xEJtLnmEHv03BLjezEqA3YSGNq5OGzNbRqh2UeVXu+c2Qs1Jy4BC9g8lfAfwsBd7GfA759w/zewqYLrXfg+hawK7gFfMLMU7Ltd60+4xs+O9sjcJPW92WQ379BKhayqfExoCuaaEJQGh0UBFDpPXDJXtnNvudywih0NNQCIiAaUagIhIQKkGICISUEoAIiIBpQQgIhJQSgAiIgGlBCAiElD/D9TkQ5LaCaC2AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "learner.recorder.plot_losses()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-22T18:21:22.276722Z",
     "start_time": "2019-08-22T18:02:21.676Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder.plot_metrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit_one_cycle(500, 2e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-15T11:15:02.866857Z",
     "start_time": "2019-08-15T11:15:02.810887Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional Fine tune regular fit "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-21T06:37:06.107163Z",
     "start_time": "2019-08-21T06:37:06.103534Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.callbacks.append(\n",
    "    ReduceLROnPlateauCallback(learner, monitor='train_loss', mode='min', factor=0.2, patience=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-21T07:03:49.373544Z",
     "start_time": "2019-08-21T06:37:06.109477Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#learner.opt_func = RAdam\n",
    "learner.fit(100, 1.0E-06)# , moms=(0.75,0.70))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit_one_cycle(300, 5e-4)#, moms=(0.75,0.70))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-10T11:24:45.530739Z",
     "start_time": "2019-08-10T11:24:44.737186Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder.plot_lr(show_moms=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-16T16:57:30.084191Z",
     "start_time": "2019-08-16T16:57:29.686756Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder.plot_losses()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-16T16:57:43.396559Z",
     "start_time": "2019-08-16T16:57:41.386234Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.recorder.plot_metrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.fit(10,1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-22T04:04:24.821143Z",
     "start_time": "2019-08-22T04:02:37.936699Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.to_fp32()\n",
    "\n",
    "val = learner.validate()\n",
    "print(val)\n",
    "val = val[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-22T04:04:25.377768Z",
     "start_time": "2019-08-22T04:04:24.824619Z"
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    sub_fname = f'loss{learner.recorder.losses[-1]:.04f}val{val:.04f}'\n",
    "except:\n",
    "    sub_fname = f'val{val:.04f}'\n",
    "learner.save(sub_fname)\n",
    "print(sub_fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:35.969Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_fname = f'loss{learner.recorder.losses[-1]:.04f}'\n",
    "learner.save(sub_fname)\n",
    "print(sub_fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:38.157Z"
    }
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference\n",
    "\n",
    "Make sure `tranforms` are activated to test set otherwise TTA > 1 will be as TTA =1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:40.536Z"
    }
   },
   "outputs": [],
   "source": [
    "test_fname = Path('test.npz')\n",
    "try:\n",
    "    npzfile  = np.load(fname_ext(test_fname, ext))\n",
    "    xt_xyz   = npzfile['x_xyz']\n",
    "    xt_type  = npzfile['x_type']\n",
    "    xt_ext   = npzfile['x_ext']\n",
    "    xt_atom  = npzfile['x_atom']\n",
    "    mt = npzfile['m']\n",
    "    xt_ids = npzfile['x_ids']\n",
    "except:\n",
    "    xt_xyz,xt_type,xt_ext,xt_atom,mt,xt_ids = \\\n",
    "        preprocess(test_fname.with_suffix('.csv'), type_index=types,ext=ext)\n",
    "    np.savez(fname_ext('_'+test_fname, ext), \n",
    "             x_xyz  = xt_xyz,\n",
    "             x_type = xt_type,\n",
    "             x_ext  = xt_ext,\n",
    "             x_atom = xt_atom,\n",
    "             m=mt,\n",
    "             x_ids=xt_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:42.544Z"
    }
   },
   "outputs": [],
   "source": [
    "if True:\n",
    "    try:\n",
    "        xt_coulombmat = load_fn(f'xt_coulombmat32{ext}.npy')\n",
    "    except:\n",
    "        xt_coulombmat = np.load(f'xt_coulombmat{ext}.npy', allow_pickle=True)\n",
    "        xt_coulombmat = np.array(xt_coulombmat.tolist()).astype(np.float32)\n",
    "        np.save(f'xt_coulombmat32{ext}.npy', xt_coulombmat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:43.229Z"
    }
   },
   "outputs": [],
   "source": [
    "xt_qm9_mulliken = load_fn(f'xt_qm9_mulliken{ext}.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:44.297Z"
    }
   },
   "outputs": [],
   "source": [
    "f'xt_qm9_mulliken{ext}.npy'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:44.980Z"
    }
   },
   "outputs": [],
   "source": [
    "[v.shape for v in [xt_xyz,xt_type,xt_ext,xt_atom, xt_qm9_mulliken,xt_ids, mt]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:46.006Z"
    }
   },
   "outputs": [],
   "source": [
    "TTA_N = 1\n",
    "\n",
    "learner.data.add_test(ItemList(items=(MoleculeItem(i,*v) for i,v in \n",
    "                              enumerate(zip(xt_xyz,xt_type,xt_ext,xt_atom,xt_qm9_mulliken,xt_coulombmat)))))\n",
    "learner.data.test_ds.tfms = tta_tfms if TTA_N > 1 else tfms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-24T10:38:46.692Z"
    }
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "del xt_xyz,xt_type,xt_ext,xt_atom, xt_qm9_mulliken, mt\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T15:50:04.292728Z",
     "start_time": "2019-08-23T15:50:04.289894Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.model = nn.DataParallel(learner.model)\n",
    "#data.batch_size = int(4096*2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:19.656865Z",
     "start_time": "2019-08-24T16:02:56.399455Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>scalar_coupling_constant</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4658147</td>\n",
       "      <td>23.877972</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>4658148</td>\n",
       "      <td>191.218170</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4658149</td>\n",
       "      <td>5.330479</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4658150</td>\n",
       "      <td>191.442947</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4658151</td>\n",
       "      <td>24.374382</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        id  scalar_coupling_constant\n",
       "0  4658147                 23.877972\n",
       "1  4658148                191.218170\n",
       "2  4658149                  5.330479\n",
       "3  4658150                191.442947\n",
       "4  4658151                 24.374382"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub = defaultdict(int)\n",
    "xt_ids_not_extended = (xt_ids!=0) & (xt_ids<=7163688) # TODO\n",
    "ids = xt_ids[xt_ids_not_extended]\n",
    "\n",
    "mb = master_bar(range(TTA_N))\n",
    "for tta in mb:\n",
    "    test_preds = np.zeros((0, 29), dtype=np.float32)\n",
    "\n",
    "    for batch_idx, batch in progress_bar(\n",
    "        enumerate(learner.dl(DatasetType.Test)), total=len(learner.dl(DatasetType.Test)), parent=mb):\n",
    "        _, _, preds_,_,_,_ = learner.pred_batch(ds_type=DatasetType.Test, batch=batch)\n",
    "        preds_ = preds_.sum(dim=1)\n",
    "        test_preds = np.concatenate([test_preds, preds_.data.cpu().numpy()], axis = 0)\n",
    "\n",
    "    preds = test_preds[xt_ids_not_extended]\n",
    "    for k in range(len(ids)):\n",
    "        sub[int(ids[k])] += preds[k]\n",
    "    \n",
    "for k in range(len(ids)):\n",
    "    sub[int(ids[k])] = sub[int(ids[k])]/TTA_N\n",
    "\n",
    "sub_df = pd.DataFrame(sub.items(), columns=['id', 'scalar_coupling_constant'])\n",
    "sub_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit to Kaggle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:19.663161Z",
     "start_time": "2019-08-24T16:13:19.659123Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'loss-5.6396'"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#sub_fname = 'loss-4.9044val-2.5880'\n",
    "#sub_fname='loss-4.9516val-2.8229'\n",
    "sub_fname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:32.180602Z",
     "start_time": "2019-08-24T16:13:19.665337Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_df.to_csv(sub_fname, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:32.185869Z",
     "start_time": "2019-08-24T16:13:32.182918Z"
    }
   },
   "outputs": [],
   "source": [
    "comp = 'champs-scalar-coupling'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:32.198097Z",
     "start_time": "2019-08-24T16:13:32.188040Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>scalar_coupling_constant</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2505537</th>\n",
       "      <td>7163684</td>\n",
       "      <td>0.737902</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2505538</th>\n",
       "      <td>7163685</td>\n",
       "      <td>4.346299</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2505539</th>\n",
       "      <td>7163686</td>\n",
       "      <td>1.841039</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2505540</th>\n",
       "      <td>7163687</td>\n",
       "      <td>3.900772</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2505541</th>\n",
       "      <td>7163688</td>\n",
       "      <td>120.185211</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              id  scalar_coupling_constant\n",
       "2505537  7163684                  0.737902\n",
       "2505538  7163685                  4.346299\n",
       "2505539  7163686                  1.841039\n",
       "2505540  7163687                  3.900772\n",
       "2505541  7163688                120.185211"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:32.202615Z",
     "start_time": "2019-08-24T16:13:32.200216Z"
    }
   },
   "outputs": [],
   "source": [
    "TTA_N=1\n",
    "ext='_ext'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:13:37.480872Z",
     "start_time": "2019-08-24T16:13:32.204943Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████| 64.0M/64.0M [00:02<00:00, 27.6MB/s]\n",
      "Successfully submitted to Predicting Molecular Properties"
     ]
    }
   ],
   "source": [
    "!kaggle competitions submit -c {comp} -f {sub_fname} -m 'no QM9 no ext no max_atoms tta {TTA_N} {ext} transformer + attn no scale + 1024'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:14:33.062380Z",
     "start_time": "2019-08-24T16:13:37.485597Z"
    }
   },
   "outputs": [
    {
     "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-56-f9d1b88c5181>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m60\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      3\u001b[0m \u001b[0mget_ipython\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msystem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'kaggle competitions submissions -c {comp} -v > submissions-{comp}.csv'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import time\n",
    "time.sleep(60)\n",
    "!kaggle competitions submissions -c {comp} -v > submissions-{comp}.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-24T16:14:33.064841Z",
     "start_time": "2019-08-24T10:39:24.896Z"
    }
   },
   "outputs": [],
   "source": [
    "submissions = pd.read_csv(f'submissions-{comp}.csv')\n",
    "submissions.iloc[0].publicScore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# L2 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T15:58:30.962770Z",
     "start_time": "2019-08-23T15:58:28.830093Z"
    }
   },
   "outputs": [],
   "source": [
    "learner.data.valid_dl = data.valid_dl.new(shuffle=False)\n",
    "TTA_N =1                                          \n",
    "sub = defaultdict(int)\n",
    "targets = defaultdict(int)\n",
    "targets_types = defaultdict(int)\n",
    "\n",
    "mb = master_bar(range(TTA_N))\n",
    "for tta in mb:\n",
    "    test_preds = np.zeros((0, 29), dtype=np.float32)\n",
    "    test_exts = np.zeros((0, 29), dtype=np.float32)\n",
    "    test_types = np.zeros((0, 29), dtype=np.float32)\n",
    "    test_targets = np.zeros((0, 29), dtype=np.float32)\n",
    "\n",
    "    for batch_idx, batch in progress_bar(\n",
    "        enumerate(learner.dl(DatasetType.Valid)), total=len(learner.dl(DatasetType.Valid)), parent=mb):\n",
    "\n",
    "        types_, ext_, preds_,_,_,_ = learner.pred_batch(ds_type=DatasetType.Valid, batch=batch)\n",
    "        preds_ = preds_.sum(dim=1)\n",
    "        test_preds = np.concatenate([test_preds, preds_.data.cpu().numpy()], axis = 0)\n",
    "\n",
    "        targets_ = batch[1][0].sum(dim=1)\n",
    "        test_targets = np.concatenate([test_targets, targets_.data.cpu().numpy()], axis = 0)\n",
    "        \n",
    "        test_types = np.concatenate([test_types, types_.data.cpu().numpy().squeeze(1)], axis = 0)\n",
    "        test_exts = np.concatenate([test_exts, ext_.data.cpu().numpy().squeeze(1)], axis = 0)\n",
    "    \n",
    "    test_preds = test_preds.flatten()\n",
    "    test_types = test_types.flatten()\n",
    "    test_targets = test_targets.flatten()\n",
    "    test_exts = test_exts.flatten()\n",
    "    \n",
    "    mask = (test_types != -1) & (test_exts == 0) #.squeeze(1)\n",
    "    preds = test_preds[mask]\n",
    "    test_targets = test_targets[mask]\n",
    "    test_types = test_types[mask]\n",
    "    for k in range(len(preds)):\n",
    "        sub[k] += preds[k]\n",
    "        targets[k] = test_targets[k]\n",
    "        targets_types[k] = test_types[k]\n",
    "    \n",
    "for k in range(len(sub.keys())):\n",
    "    sub[k] = sub[k]/TTA_N\n",
    "\n",
    "sub_df = pd.DataFrame(sub.items(), columns=['id', 'scalar_coupling_constant'])\n",
    "sub_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T15:58:30.965159Z",
     "start_time": "2019-08-23T15:50:37.862Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_targets = pd.DataFrame(targets.items(), columns=['id', 'scalar_coupling_constant'])\n",
    "sub_targets.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-22T18:21:22.311470Z",
     "start_time": "2019-08-22T18:15:45.443Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_targets = pd.DataFrame(targets.items(), columns=['id', 'scalar_coupling_constant'])\n",
    "sub_targets.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-21T03:24:10.685807Z",
     "start_time": "2019-08-21T03:24:10.363096Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_targets_types = pd.DataFrame(targets_types.items(), columns=['id', 'type'])\n",
    "sub_targets_types.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-05T21:42:00.805651Z",
     "start_time": "2019-08-05T21:41:59.063328Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_targets_types.to_csv('validation_types', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-22T04:18:28.140915Z",
     "start_time": "2019-08-22T04:18:28.133365Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-21T03:24:10.704696Z",
     "start_time": "2019-08-21T03:24:10.697572Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_targets.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-03T20:07:20.048256Z",
     "start_time": "2019-08-03T20:07:20.041023Z"
    }
   },
   "outputs": [],
   "source": [
    "mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-03T20:09:32.533038Z",
     "start_time": "2019-08-03T20:09:30.294063Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_targets.to_csv('validation_targets', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T15:58:30.966364Z",
     "start_time": "2019-08-23T15:50:47.589Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_df.to_csv(sub_fname + '_val', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-23T15:58:30.967796Z",
     "start_time": "2019-08-23T15:50:48.285Z"
    }
   },
   "outputs": [],
   "source": [
    "sub_fname + '_val'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-05T21:00:08.361139Z",
     "start_time": "2019-08-05T21:00:07.258433Z"
    }
   },
   "outputs": [],
   "source": [
    "!head loss-4.9516val-3.0042_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-03T20:11:05.245077Z",
     "start_time": "2019-08-03T20:11:04.340447Z"
    }
   },
   "outputs": [],
   "source": [
    "!head loss-5.7552val-2.9950_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-03T20:11:18.206277Z",
     "start_time": "2019-08-03T20:11:17.307404Z"
    }
   },
   "outputs": [],
   "source": [
    "!head validation_targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}