cloud-quaternion-test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\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\n",
    "\n",
    "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.distributed import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(fname, type_index=None):\n",
    "    t  = pd.read_csv(fname)\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('scalar_coupling_contributions.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_index'] = pd.factorize(pd.concat([pd.Series(type_index),t['type']]))[0][len(type_index):]\n",
    "    else:\n",
    "        t['type_index'] = pd.factorize(t['type'])[0]\n",
    "    s = pd.concat([s,pd.get_dummies(s['atom'])], axis=1)\n",
    "\n",
    "    max_items = 785836 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",
    "    a_hot = ['C','F','H','N','O']\n",
    "    \n",
    "    x_xyz   = np.zeros((max_items,len(xyz),  max_atoms), dtype=np.float32)\n",
    "    x_a_hot = np.zeros((max_items,len(a_hot),max_atoms), dtype=np.float32)\n",
    "    x_type  = np.zeros((max_items,1,         max_atoms), dtype=np.float32)\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",
    "\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_a_hot[i] = ref_a[a_hot].values.T\n",
    "            x_type[i] = -1\n",
    "\n",
    "            x_xyz[i,:,:n_atoms] = (ss[xyz].values-ref_a[xyz].values).T  # xyz \n",
    "            x_a_hot[i,:,:n_atoms] = ss[a_hot].T                  # a_hot\n",
    "            x_type[i,0,a_group['atom_index_1']] = a_group['type_index']  # type \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_a_hot,x_type, m , y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential\n",
    "    else:\n",
    "        return x_xyz,x_a_hot,x_type, m, xt_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_fname = Path('train.npz')\n",
    "try:\n",
    "    npzfile = np.load(train_fname)\n",
    "    x_xyz = npzfile['x_xyz']\n",
    "    x_a_hot = npzfile['x_a_hot']\n",
    "    x_type = npzfile['x_type']\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_a_hot,x_type, m , y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential = \\\n",
    "        preprocess(train_fname.with_suffix('.csv'))\n",
    "    np.savez(train_fname, \n",
    "             x_xyz=x_xyz,\n",
    "             x_a_hot=x_a_hot,\n",
    "             x_type=x_type,\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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fname = Path('test.npz')\n",
    "try:\n",
    "    npzfile = np.load(test_fname)\n",
    "    xt_xyz = npzfile['x_xyz']\n",
    "    xt_a_hot = npzfile['x_a_hot']\n",
    "    xt_type = npzfile['x_type']\n",
    "    mt = npzfile['m']\n",
    "    #xt_ids = npzfile['x_ids']\n",
    "except:\n",
    "    train_csv = pd.read_csv(train_fname.with_suffix('.csv'))\n",
    "    xt_xyz,xt_a_hot,xt_type,mt,xt_ids = \\\n",
    "        preprocess(test_fname.with_suffix('.csv'), type_index = pd.factorize(train_csv['type'])[1])\n",
    "    np.savez(test_fname, \n",
    "             x_xyz=xt_xyz,\n",
    "             x_a_hot=xt_a_hot,\n",
    "             x_type=xt_type,\n",
    "             m=mt,\n",
    "             x_ids=xt_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    xt_ids = np.load('xt_ids.npy')\n",
    "except:\n",
    "    print(\"ASK PAVEL\")\n",
    "    assert False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(785836, 3, 29),\n",
       " (785836, 5, 29),\n",
       " (785836, 1, 29),\n",
       " (785836, 4, 29),\n",
       " (785836, 9, 29),\n",
       " (785836, 1, 29),\n",
       " (785836, 3),\n",
       " (785836, 1),\n",
       " (785836,)]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[v.shape for v in [x_xyz,x_a_hot,x_type, y_scalar, y_magnetic, y_mulliken, y_dipole, y_potential, m]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is commented b/c I want to rotate from the source atom @ (0,0) w/o norm\n",
    "# MAYBE we should at least norm magnitude (std)\n",
    "\n",
    "# TODO: Do norm in FASTAI instead of here\n",
    "# TOCHECK: Filter only valid atoms (otherwise repeated atoms may skew stats)\n",
    "#xyz_mean, xyz_std = x_xyz.mean(axis=(0,2), keepdims=True),  x_xyz.std(axis=(0,2), keepdims=True)\n",
    "#x_xyz  = (x_xyz  - xyz_mean) / xyz_std\n",
    "#xt_xyz = (xt_xyz - xyz_mean) / xyz_std\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(422550, 3, 29), (422550, 5, 29), (422550, 1, 29), (422550,)]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[v.shape for v in [xt_xyz,xt_a_hot,xt_type, mt]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quaternion playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "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",
    "    \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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x is a random 3-element vector with elements between 0..1 (in code below comes from sigmoid)\n",
    "x = torch.rand((2,3)) # batch size 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.0476, -0.0096, -0.1651, -0.9851],\n",
       "         [-0.9585, -0.2693, -0.0544, -0.0766]]), torch.Size([2, 4]))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get the quaternion q based on x\n",
    "# http://planning.cs.uiuc.edu/node198.html\n",
    "sq1_v1,sqv1,v2_2pi,v3_2pi = torch.sqrt(1-x[:,:1]),torch.sqrt(x[:,:1]),2*math.pi*x[:,1:2],2*math.pi*x[:,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=1)\n",
    "q, q.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 29])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# these are the 2 batches of points we want to rotate\n",
    "xs = Tensor(x_xyz[[2,100]])\n",
    "xs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 29])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rot_xs = qrot(q.unsqueeze(1).expand(-1,xs.shape[2],-1).contiguous(),xs.transpose(1,2).contiguous()).transpose(2,1)\n",
    "rot_xs.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(tensor([1.0919, 1.7831, 1.7832, 0.0000, 1.7831, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000]),\n",
       "  tensor([1.0919, 1.7831, 1.7832, 0.0000, 1.7831, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000])),\n",
       " (tensor([0.9626, 1.9393, 3.1752, 4.3067, 0.0000, 2.3209, 2.3206, 5.3312, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000]),\n",
       "  tensor([0.9626, 1.9393, 3.1752, 4.3067, 0.0000, 2.3209, 2.3206, 5.3312, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
       "          0.0000, 0.0000]))]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check norm of each atom in rotated baches is OK.\n",
    "[(torch.norm(rot_xs[i],dim=0), torch.norm(xs[i], dim=0)) for i in range(xs.shape[0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "dropout_rate = 0.1\n",
    "# from https://githuba.com/fxia22/pointnet.pytorch/blob/master/pointnet/model.py\n",
    "\n",
    "class Quaternion(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Quaternion, self).__init__()\n",
    "        self.conv1 = torch.nn.Conv1d(3, 64, 1)\n",
    "        self.conv2 = torch.nn.Conv1d(64, 128, 1)\n",
    "        self.conv3 = torch.nn.Conv1d(128, 1024, 1)\n",
    "        self.fc1 = nn.Linear(1024, 512)\n",
    "        self.fc2 = nn.Linear(512, 256)\n",
    "        self.fc3 = nn.Linear(256, 3)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "        self.bn1 = nn.BatchNorm1d(64)\n",
    "        self.bn2 = nn.BatchNorm1d(128)\n",
    "        self.bn3 = nn.BatchNorm1d(1024)\n",
    "        self.bn4 = nn.BatchNorm1d(512)\n",
    "        self.bn5 = nn.BatchNorm1d(256)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batchsize = x.size()[0]\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = F.relu(self.bn2(self.conv2(x)))\n",
    "        x = F.relu(self.bn3(self.conv3(x)))\n",
    "        x = torch.max(x, 2, keepdim=True)[0]\n",
    "        x = x.view(-1, 1024)\n",
    "\n",
    "        x = F.relu(self.bn4(self.fc1(x)))\n",
    "        x = F.relu(self.bn5(self.fc2(x)))\n",
    "        x = torch.sigmoid(self.fc3(x))\n",
    "        \n",
    "        #http://planning.cs.uiuc.edu/node198.html\n",
    "        sq1_v1,sqv1,v2_2pi,v3_2pi = torch.sqrt(1-x[:,:1]),torch.sqrt(x[:,:1]),2*math.pi*x[:,1:2],2*math.pi*x[:,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=1)\n",
    "        return q\n",
    "\n",
    "class STNkd(nn.Module):\n",
    "    def __init__(self, k=64):\n",
    "        super(STNkd, self).__init__()\n",
    "        self.conv1 = torch.nn.Conv1d(k, 64, 1)\n",
    "        self.conv2 = torch.nn.Conv1d(64, 128, 1)\n",
    "        self.conv3 = torch.nn.Conv1d(128, 1024, 1)\n",
    "        self.fc1 = nn.Linear(1024, 512)\n",
    "        self.fc2 = nn.Linear(512, 256)\n",
    "        self.fc3 = nn.Linear(256, k*k)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "        self.bn1 = nn.BatchNorm1d(64)\n",
    "        self.bn2 = nn.BatchNorm1d(128)\n",
    "        self.bn3 = nn.BatchNorm1d(1024)\n",
    "        self.bn4 = nn.BatchNorm1d(512)\n",
    "        self.bn5 = nn.BatchNorm1d(256)\n",
    "\n",
    "        self.k = k\n",
    "\n",
    "    def forward(self, x):\n",
    "        batchsize = x.size()[0]\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = F.relu(self.bn2(self.conv2(x)))\n",
    "        x = F.relu(self.bn3(self.conv3(x)))\n",
    "        x = torch.max(x, 2, keepdim=True)[0]\n",
    "        x = x.view(-1, 1024)\n",
    "\n",
    "        x = F.relu(self.bn4(self.fc1(x)))\n",
    "        x = F.relu(self.bn5(self.fc2(x)))\n",
    "        x = self.fc3(x)\n",
    "        x = x.view(-1, self.k, self.k)\n",
    "        return x\n",
    "\n",
    "class PointNetfeat(nn.Module):\n",
    "    def __init__(self, global_feat = True, input_transform = True, feature_transform = False):\n",
    "        super(PointNetfeat, self).__init__()\n",
    "        self.stn = Quaternion() if input_transform else None\n",
    "        self.conv1 = torch.nn.Conv1d(9, 64, 1)\n",
    "        self.conv2 = torch.nn.Conv1d(64, 128, 1)\n",
    "        self.conv3 = torch.nn.Conv1d(128, 1024, 1)\n",
    "        self.bn1 = nn.BatchNorm1d(64)\n",
    "        self.bn2 = nn.BatchNorm1d(128)\n",
    "        self.bn3 = nn.BatchNorm1d(1024)\n",
    "        self.global_feat = global_feat\n",
    "        self.feature_transform = feature_transform\n",
    "        if self.feature_transform:\n",
    "            self.fstn = STNkd(k=64)\n",
    "        self.zero = torch.zeros((1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        n_pts = x.size()[2]\n",
    "        if self.stn is not None:\n",
    "            x_xyz = x[:,:3,...]\n",
    "            q = self.stn(x_xyz).unsqueeze(1).expand(-1,n_pts,-1).contiguous()\n",
    "            x_xyz = x_xyz.transpose(1,2).contiguous()\n",
    "            x[:,:3,:] = qrot(q,x_xyz).transpose(2,1)\n",
    "        else:\n",
    "            trans = torch.eye(\n",
    "                self.k, dtype=x.dtype, device=x.device).view(1,self.k*self.k).expand((batchsize,-1))\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "\n",
    "        if self.feature_transform:\n",
    "            trans_feat = self.fstn(x)\n",
    "            x = x.transpose(2,1)\n",
    "            x = torch.bmm(x, trans_feat)\n",
    "            x = x.transpose(2,1)\n",
    "        else:\n",
    "            trans_feat = None\n",
    "\n",
    "        pointfeat = x\n",
    "        x = F.relu(self.bn2(self.conv2(x)))\n",
    "        x = self.bn3(self.conv3(x))\n",
    "        x = torch.max(x, 2, keepdim=True)[0]\n",
    "        x = x.view(-1, 1024)\n",
    "        if self.global_feat:\n",
    "            return x, trans, trans_feat\n",
    "        else:\n",
    "            x = x.view(-1, 1024, 1).expand(-1, -1, n_pts)\n",
    "            return torch.cat([x, pointfeat], 1)#, self.zero, trans_feat # FIX THIS\n",
    "    \n",
    "class PointNetDenseReg(nn.Module):\n",
    "    def __init__(self, k = 2, feature_transform=False):\n",
    "        super(PointNetDenseReg, self).__init__()\n",
    "        self.k = k\n",
    "        self.feature_transform=feature_transform\n",
    "        self.feat = PointNetfeat(global_feat=False, feature_transform=feature_transform)\n",
    "        self.conv1 = torch.nn.Conv1d(1088, 512, 1)\n",
    "        self.conv2 = torch.nn.Conv1d(512, 256, 1)\n",
    "        self.conv3 = torch.nn.Conv1d(256, 128, 1)\n",
    "        self.conv4 = torch.nn.Conv1d(128, self.k, 1)\n",
    "        self.bn1 = nn.BatchNorm1d(512)\n",
    "        self.bn2 = nn.BatchNorm1d(256)\n",
    "        self.bn3 = nn.BatchNorm1d(128)\n",
    "        self.dropout = nn.Dropout(p=dropout_rate)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        batchsize = x.size()[0]\n",
    "        n_pts = x.size()[2]\n",
    "        x = self.feat(x)\n",
    "        x = self.dropout(x)\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = F.relu(self.bn2(self.conv2(x)))\n",
    "        x = F.relu(self.bn3(self.conv3(x)))\n",
    "        x = self.conv4(x)\n",
    "        #x = x.transpose(2,1).contiguous()\n",
    "        #x = x.view(batchsize, n_pts, self.k)\n",
    "        return x#, trans, trans_feat if self.feature_transform else trans\n",
    "\n",
    "def feature_transform_regularizer(trans):\n",
    "    d = trans.size()[1]\n",
    "    batchsize = trans.size()[0]\n",
    "    I = torch.eye(d, dtype=trans.dtype, device=trans.device).unsqueeze(0)\n",
    "    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MoleculeItem(ItemBase):\n",
    "    def __init__(self,i,xyz,a_hot,type): \n",
    "        self.i, self.xyz,self.a_ahot,self.type = i,xyz,a_hot,type\n",
    "        self.data = torch.cat([Tensor(xyz), Tensor(self.a_ahot),Tensor(self.type)], dim=0)\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",
    "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(np.sum(scalar, axis=0))\n",
    "    def __str__(self):\n",
    "        res, spacer, n_couplings = '', '', 0\n",
    "        print(self.data)\n",
    "        for s in self.data:\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": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def LMAEMaskedLoss(input_outputs,target):\n",
    "    input, output = input_outputs\n",
    "    loss = 0.\n",
    "    n = 0\n",
    "    for type in range(n_types):\n",
    "        mask = (input[:,8,:] == type)\n",
    "        if mask.sum() > 0:\n",
    "            _output,_target = output[:,0,:], target\n",
    "            loss += torch.log((_output[mask] - _target[mask]).abs().mean()+1e-9)\n",
    "            n+=1\n",
    "    return loss/n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScalarCouplingList(ItemList):\n",
    "    def __init__(self, items:Iterator, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.c = 1\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": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = ItemList(items=(MoleculeItem(i,*v) for i,v in enumerate(zip(x_xyz,x_a_hot,x_type))),\n",
    "                label_cls=ScalarCouplingItem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "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": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ItemLists;\n",
       "\n",
       "Train: ItemList (707619 items)\n",
       "0 5 atoms 4 couplings,1 5 atoms 3 couplings,2 5 atoms 2 couplings,3 5 atoms 1 couplings,4 4 atoms 3 couplings\n",
       "Path: .;\n",
       "\n",
       "Valid: ItemList (78217 items)\n",
       "9 8 atoms 7 couplings,10 8 atoms 6 couplings,11 8 atoms 5 couplings,12 8 atoms 4 couplings,13 8 atoms 3 couplings\n",
       "Path: .;\n",
       "\n",
       "Test: None"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = data.split_by_idx(idx_valid_split)\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "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": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "data=data.databunch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChemCallback(Callback):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "    def on_epoch_begin(self,**kwargs):\n",
    "        pass # PAVEL -> at some point do learner.opt.clear() or equivalent reset opt internals \n",
    "    def on_batch_begin(self,**kwargs):\n",
    "        \"Save the last_input (i.e. current input) for on_loss_begin_callback\"\n",
    "        self.last_input = kwargs['last_input']\n",
    "    def on_loss_begin(self, last_output:Tensor, **kwargs):\n",
    "        \"Add last_input to last_output, i.e. input to loss function b/c we need it\"\n",
    "        return {'last_output': (self.last_input, last_output)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "net, learner = None,None\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "net = PointNetDenseReg(k = 1)\n",
    "learner = Learner(data,net, callbacks=[ChemCallback()], \n",
    "                  loss_func=partial(LMAEMaskedLoss),\n",
    "                  metrics=[partial(LMAEMaskedLoss)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "#learner = learner.to_parallel().to_fp16() # \n",
    "data.batch_size = 4096*2#*2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LR Finder is complete, type {learner_name}.recorder.plot() to see the graph.\n",
      "Min numerical gradient: 2.75E-02\n",
      "Min loss divided by 10: 3.31E-02\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYwAAAEKCAYAAAAB0GKPAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvNQv5yAAAIABJREFUeJzt3Xd8lfXd//HXJzshIQRIIiMQCFtlSAC3OOreUke1tUqLuGdr7/q7297a3lqto25ta52odbXiqKMKqIAQZCN7D0kCJITs8f39cS5vA4ZwIDm5cs55Px+P6+E51/W9zvX5mpD3udb3MuccIiIi+xLjdwEiIhIeFBgiIhIUBYaIiARFgSEiIkFRYIiISFAUGCIiEhQFhoiIBEWBISIiQVFgiIhIUOL8LqA1de3a1eXm5vpdhohI2JgzZ06xcy4zmLYRFRi5ubkUFBT4XYaISNgws3XBttUhKRERCYoCQ0REgqLAEBGRoCgwREQkKAoMEREJSsgCw8xyzOxTM/vazBab2Y1NtDEze9jMVprZAjM7rNGyy81shTddHqo6RUQkOKG8rLYOuNU595WZpQFzzOwj59ySRm1OA/p70xjgCWCMmXUGfgvkA85b923n3I4Q1isiIs0IWWA457YAW7zXZWb2NdADaBwY5wDPu8BzYmeaWScz6waMBT5yzm0HMLOPgFOBl0NV7/7asL2CD5dspaq2HucczkFCXAwjemUwPKcTCXE62icikaVNbtwzs1xgBPDlHot6ABsavd/ozdvb/JAYe9+nxMYYnTskMKJXBqccnM3Qnp2Ij42htr6B5VvLWFm4i9VF5XxTWsWa4nJmrd2+189LSYhlVG5n8ntnkJwQS2piHEf160pO5xSqauvZVl5DB29+XKyCRUTCQ8gDw8xSgTeAm5xzO/dc3MQqrpn5TX3+BGACQK9evfa7PuccR/fvyo7yWgrLqvj7F2t4etpq4mKMnM4pbC6ppLquwdsWdE1N5KCOSdzygwGcf1gPMtMSMQwzKK+u48s125m+spjPVxYzdXnRbtvq0iGBbeU1//c+KT6GUbmdGZHTiZTEOBLjYkiKjyUpPoakuFiS4mNJjPfmxcWSnBBLSkIsDc6xq6qOsuo6yqrq2LSjkrXbyklNjGNoz3S6piZS3+DokBhLRkoCsTFGTX0DxWU1fLOzks0lVRSWVZMUH0NGSgIZKfF0TIqnpLKW7eU1pCTEkp4cT3pyPJ1SEuiWnkSHxIgaFEBEDoAFjgaF6MPN4oF3gA+ccw80sfwpYIpz7mXv/TICh6PGAmOdc1c11W5v8vPzXUuHBtlZVcvUZUV8vWUna4rL6d4pmeE5nRh0UBq9uqSQGBcb9GdV1tRT29BA4c5qpiwrZPnWMnIyUsjqmEh5dT3rt1fwxcpiVhTualHNEDgcVlvfQKh+nOnJ8XTvlEyPTsn07pLCQR2T2FVdR219A/2zUxncrSN5manEa49JJKyY2RznXH5QbUMVGGZmwHPAdufcTXtpcwZwHXA6gZPeDzvnRnsnvecA31419RUw8ttzGnvTGoHhh/oGR3VdPVW1DVTV1ntTA1V1gdfV386vq6eiph7DSEuKIy0pjtTEOLI7JtGjUzLlNXUs2rSTXdV1GFBeU8eO8hoaHMTHWmDvKD2JbunJdE1NoLbesaOihh0VNeysrCM9OZ7OHRKorK2ntLKW0spaSipq2FxSxeaSSjaXVLJhRwXrt1dQVRvY64qNMeobAr9DCbEx9Onagez0JLLSEsnumEhWWhLZHRPJTPt2XpLO74i0I/sTGKE8znAU8GNgoZnN8+b9GugF4Jx7EniPQFisBCqAK7xl283sLmC2t96d+wqLcBYbY6QkxJGS0LLPSUuK54i8LkG3j4uF5IRkundK3q/tNDQ4yqrrSE2Mo8E51hSX8/WWnSzZspNVheUUlVWxYmsZRWXV1DXs/oXEDLLTkuiZkUxO5xQGHpTG8QOzGJCdSuA7hoi0VyE9JNXWwnUPI1I1NDi2V9RQuLOarWVVFO2sZnNpJRt3VLJhewUbd1SyqaQSgE4p8fTPSuWwXhmcfHA2I3IyiIlRgIiEWrs4JOUHBUb4+aa0iqnLC5m3oZQVW8uYv7GE2npHbpcUxh/dh3NG9KBjUrzfZYpELAWGhK3Sylo+WbqVZ6evY/6GEsxgQFYaxw7oytnDenBIj446dCXSihQYEvacc3y1voTPVxRTsG47M1dvo7beMeigNH58RG/OHd5Dl/qKtAIFhkSckooa3lv4DS/OXMeSLTtJS4pj3MieXDqmF/2y0vwuTyRsKTAkYgX2PHbw3PR1vL9oC7X1jhG9OnFRfg5nDO1Gms53iOwXBYZEhaKyav45dxOvFmxgZeEukuNjGd2nM6P7dObUQw4iLzPV7xJF2j0FhkQV5xzzNpTw1txNzFy9jeVbA3fO5/fO4EJvz0PnO0SapsCQqFa4s4q3vD2P1UXlpCTEctbQ7lw0OocROZ10lZVIIwoMEb473/Hq7A28s2ALFTX1DMhO5cL8HI7pn0n/rFTdHChRT4Ehsodd1XW8M38zr8zewLwNJUDg7vLzR/Tk0sN76XyHRC0Fhkgz1m0rZ/baHXy6rJAPFn1DXYOjR6dkxvTtzKVjejOyd4bfJYq0GQWGSJAKy6p4b8EWZq/dwWcrithZVcewnE6cP6IHpx/ajcy0RL9LFAkpBYbIASivruONrzYy6cv1LP2mjPhY49zhPbjquDz6ZemQlUQmBYZICy3fWsZLM9fxasEGqusaOHlINlcdl6errCTiKDBEWsm2XdU8O30tz01fy86qOoZ068glY3pxzvDuGkVXIoICQ6SV7aqu4625m3j5y/Us2bKT5PhYzhrWjZ8ckcshPdL9Lk/kgCkwRELEOcfCTaVM+nI9/5q3mcraekblZvDTI/tw8sHZeqa5hB0FhkgbKK2o5bU5G3h+xjrWb6+gR6dkrj+hHxeM7KngkLChwBBpQ/UNjk+WFvLopyuZv6GEHp2S+fkxfbhwVA4pCRrDSto3BYaID5wLBMdjn67kq/UlZKTE85MjcvnpkblkdEjwuzyRJikwRHxWsHY7T05dzcdfbyUjJZ67zj2EM4d297sske/Zn8DQgVaREMjP7cxfL8/n3zcdQ6/OKVw3aS4Tni9gTXG536WJHDAFhkgIDTqoI29cfSS3nzqIz1cWc/KDU7lz8hJKKmr8Lk1kvykwREIsLjaGq8fmMeUXYxk3sifPTl/DcfdN4a+fraamrsHv8kSCpsAQaSNZaUncff5Q3rvxGIb2TOf3737NDx6cyvsLtxBJ5xIlcikwRNrYoIM68sL4MTx7xSgS42K4+qWvuPCpGcz3ntMh0l4pMER8MnZgFu/dcAz/e96hrCku55zHvuCmV+ayuaTS79JEmqTAEPFRXGwMPxrTi09vG8u1x+fx/qJvOPH+qTz6yQqqauv9Lk9kNwoMkXYgLSmeX5wyiP/cehxjB2bypw+Xc9Yjn7Ngow5TSfuhwBBpR3pmpPDEZSN59opRlFXVcd7j07nn/aVU1mhvQ/ynwBBph8YOzOKDm4/lgsN68OTUVZzy0DQ+W1Hkd1kS5RQYIu1UenI8944bxqSfjyE2xvjx32Zxy6vz2FGum/7EHwoMkXbuyLyuvH/jMdxwQj8mL9jMKQ9NY9py7W1I21NgiISBpPhYbjl5IP+89ijSk+P5yTOz+N3bi3UllbQpBYZIGDm4ezqTrz+aK47K5dnpaznrkc9ZsbXM77IkSoQsMMzsGTMrNLNFe1meYWZvmdkCM5tlZoc0WrbWzBaa2Twz03jlIo0kxcfy27MO5vkrR7OjopbzHp/Ox0u2+l2WRIFQ7mE8C5zazPJfA/Occ0OBnwB/3mP58c654cGO0y4SbY4dkMnk64+iT9cO/PyFAh77dKXGpJKQCllgOOemAdubaTIE+I/XdimQa2bZoapHJBJ1S0/mtYlHcPaw7tz3wTKue3ku5dV1fpclEcrPcxjzgfMBzGw00Bvo6S1zwIdmNsfMJvhUn0hYSIqP5aGLhvPr0wfx/sItnPPYFzqvISHhZ2DcA2SY2TzgemAu8O1Xo6Occ4cBpwHXmtmxe/sQM5tgZgVmVlBUpEsNJTqZGROOzePF8WMoqajh7Ee/4J9zN/ldlkQY3wLDObfTOXeFc244gXMYmcAab9lm77+FwFvA6GY+52nnXL5zLj8zM7MNKhdpv47s15V3bziGQ3ukc9Or8/j1Wwv1kCZpNb4Fhpl1MrME7+3PgGnOuZ1m1sHM0rw2HYCTgSavtBKR78vumMSkn4/hquP6MunL9Vz+zCxKK2v9LksiQCgvq30ZmAEMNLONZjbezCaa2USvyWBgsZktJXDo6UZvfjbwuZnNB2YB7zrn/h2qOkUiUVxsDP912mAevGgYBeu2M+6J6awtLve7LAlzFkmX4eXn57uCAt22IdLY9FXFXPPSV9Q3OB66aDgnDtbFiPIdM5sT7O0LutNbJMIdmdeVydcdTe8uKYx/roAHPlxGfUPkfFGUtqPAEIkCOZ1TeH3ikfxwZE8e/mQlVz47m5IKjXor+0eBIRIlkuJjuXfcUP5w3iFMX1XMWY9+zuLNpX6XJWFEgSESRcyMS8f05h9XHUFtneP8x6fzxpyNfpclYUKBIRKFRvTKYPL1RzOiVydufW0+v3pjAd+UVvldlrRzCgyRKJWZlsiL48cw4di+/KNgA8fc+wl3vLVQzw+XvVJgiESxuNgYfn36YKbcdjwXj+rFpFnrGffkdDbuqPC7NGmHFBgiQq8uKdx17iH87fJ81m+r4MxHPuedBZv9LkvaGQWGiPyfEwZl8/b1R9O7SweumzSXm1+dxy4Nly4eBYaI7KZP1w68MfEIbj5pAP+at4mzH/mcpd/s9LssaQcUGCLyPXGxMdx4Un9e+tnhlFXXcc6jX/Dq7PV6ol+UU2CIyF4dkdeF9244hvzcDG5/YyG/fH2BhkuPYgoMEWlWZloiz185hutP6MdrczYy/rnZOq8RpRQYIrJPsTHGrScP5N4LhjJ91TYuemoGhWW60S/aKDBEJGgXjsrhrz/JZ3VROec/Pp3VRbv8LknakAJDRPbL8YOyeGXC4VTW1HPhUzNZWVjmd0nSRhQYIrLfhuV04tWrjsAMLn76S4VGlFBgiMgB6ZeVyss/PxxQaEQLBYaIHLB+Wam8MmEM8G1o6JxGJFNgiEiL9MtKaxQaMxUaEUyBISIt1i8rjZd//l1oLN+qw1ORSIEhIq2if/Z3oXHWI5/z1NRV1DdoKJFIosAQkVbTPzuN9244mmMHZHL3+0v50V9mUryr2u+ypJUoMESkVWV1TOLpH4/k/h8OY96GEs565HMWby71uyxpBQoMEWl1ZsYFI3vyxtVHAvDTv89mU0mlz1VJSykwRCRkDumRznNXjqaqpp7xz86mrKrW75KkBRQYIhJSA7LTeOzSw1hRuIvxzxZQrpFuw5YCQ0RC7tgBmTx40XAK1m3n8mdmaXj0MKXAEJE2cfaw7jx8yQjmbijhplfm0qBLbsOOAkNE2syZQ7tzx+mD+fjrQp6attrvcmQ/KTBEpE1dcVQuZwztxn0fLGXGqm1+lyP7QYEhIm3KzPjjBUPJ7dqB61+eS+FOPbkvXCgwRKTNpSbG8eRlIymvruO6SXOprW/wuyQJggJDRHwxIDuNey44lFlrt/PAR8v9LkeCoMAQEd+cM7wHF4/K4cmpq5i+qtjvcmQfQhYYZvaMmRWa2aK9LM8ws7fMbIGZzTKzQxotO9XMlpnZSjP7VahqFBH//easIfTp0oFbXp3PjvIav8uRZoRyD+NZ4NRmlv8amOecGwr8BPgzgJnFAo8BpwFDgEvMbEgI6xQRH6UkxPHwJSMo3lXNXe8s8bscaUbIAsM5Nw3Y3kyTIcB/vLZLgVwzywZGAyudc6udczXAK8A5oapTRPx3SI90rhmbx5tzNzFlWaHf5che+HkOYz5wPoCZjQZ6Az2BHsCGRu02evOaZGYTzKzAzAqKiopCWK6IhNK1J/SjX1Yqd7y1SEOHtFN+BsY9QIaZzQOuB+YCdYA10XavYwg45552zuU75/IzMzNDU6mIhFxiXCx/vGAom0srue/fS/0uR5oQVGCYWZ6ZJXqvx5rZDWbWqSUbds7tdM5d4ZwbTuAcRiawhsAeRU6jpj2BzS3ZloiEh5G9M7j8iFyen7mOgrXNHdEWPwS7h/EGUG9m/YC/AX2ASS3ZsJl1MrME7+3PgGnOuZ3AbKC/mfXxll8MvN2SbYlI+PjFKQPpnp7ML99YQFVtvd/lSCPBBkaDc64OOA94yDl3M9CtuRXM7GVgBjDQzDaa2Xgzm2hmE70mg4HFZraUwBVRNwJ427kO+AD4GviHc27x/nZMRMJTh8Q47j7/UFYXlXPzq/Oo16i27UZckO1qzewS4HLgLG9efHMrOOcu2cfyGUD/vSx7D3gvyNpEJMIcOyCT/z5zCHe9s4Rfv7mQey44FLOmTm9KWwo2MK4AJgJ/cM6tMbM+wIuhK0tEot34o/tQUlHDI5+spEdGMjec2OT3S2lDQQWGc24JcAME7tAG0pxz94SyMBGRW34wgE0llTzw0XL6dO3AWcO6+11SVAv2KqkpZtbRzDoTuH/i72b2QGhLE5FoZ2bcff6hjMrN4NbX5rNoU6nfJUW1YE96p3tXMJ0P/N05NxI4KXRliYgEJMbF8uRlI+mcksB1k77STX0+CjYw4sysG3Ah8E4I6xER+Z4uqYk8fMkI1m+v4P+9tRDndOWUH4INjDsJXOa6yjk328z6AitCV5aIyO5G9+nMTScN4J/zNvPB4q1+lxOVggoM59xrzrmhzrmrvfernXMXhLY0EZHdXTM2j4HZafz+3SW6qc8HwZ707uk9u6LQzLaa2Rtm1jPUxYmINBYXG8P/nHMwG3dU8uTUVX6XE3WCPST1dwLDc3QnMHLsZG+eiEibOrxvF84a1p0npqxibXG53+VElWADI9M593fnXJ03PUtgsEARkTZ3x+mDSYiN4VdvLqBBQ4e0mWADo9jMLjOzWG+6DNgWysJERPbmoPQk7jhjMDNXb2fSrPV+lxM1gg2MKwlcUvsNsAUYR2C4EBERX1w0Koej+nXhnveXUlRW7Xc5USHYq6TWO+fOds5lOueynHPn4j0tT0TED2bGneccQkVNHX/5bLXf5USFljxx75ZWq0JE5ADkZaZyzvAevDBjHcW7tJcRai0JDI01LCK+u+6EflTX1fOXadrLCLWWBIYuTRAR3+VlpnLWsO48P2Md27SXEVLNBoaZlZnZziamMgL3ZIiI+O76E/pRVVfPXz5b43cpEa3ZwHDOpTnnOjYxpTnngn34kohISPXLSuOsod15fsZatpfX+F1OxGrJISkRkXbjhhP7UVlbryumQkiBISIRoV9WGmcc2o3np2svI1QUGCISMW48sT8VtfX8VXsZIaHAEJGI0T87sJfxnPYyQkKBISIRRXsZoaPAEJGI0j87jZ9m15Pz/27DdewIMTHQsSNccw2s0jM0WsIi6dm4+fn5rqCgwO8yRMRP779PwwXjqKuuJqGh0VP54uMD0+uvw2mn+VdfO2Nmc5xz+cG01R6GiESOVatg3DhiKit2DwuA2lqoqIBx47SncYAUGCISOe6/PxAMzamthQcfbJt6IowCQ0Qix4svBhcYL7zQNvVEGAWGiESOXbtat53sRoEhIpEjNbV128luFBgiEjkuuyxwJVRz4uPhxz9um3oijAJDRCLHrbcGFxg339w29UQYBYaIRI68vMB9Fikp3wuO2phYXEpKYHlenk8FhjcFhohEltNOgwULYMKEwB3eMTHUp6Uxadip/O2xf+qmvRZQYIhI5MnLg0cfhdJSqK8ndudOZt92Jw+srqeoTI9xPVAhCwwze8bMCs1s0V6Wp5vZZDObb2aLzeyKRsvqzWyeN70dqhpFJHrc8oMBVNc18PiUlX6XErZCuYfxLHBqM8uvBZY454YBY4H7zSzBW1bpnBvuTWeHsEYRiRJ9M1MZd1hPXpq5nk0llX6XE5ZCFhjOuWnA9uaaAGlmZkCq17YuVPWIiNx4Un8AHv54hc+VhCc/z2E8CgwGNgMLgRudcw3esiQzKzCzmWZ2bnMfYmYTvLYFRUVFIS5ZRMJZ907JXHZ4b17/aiOrinS39/7yMzBOAeYB3YHhwKNm1tFb1ssbbvdHwENmttdr4JxzTzvn8p1z+ZmZmSEvWkTC2zXH55EYF8MDHy33u5Sw42dgXAG86QJWAmuAQQDOuc3ef1cDU4ARfhUpIpGla2oi44/uw7sLtrBoU6nf5YQVPwNjPXAigJllAwOB1WaWYWaJ3vyuwFHAEt+qFJGI87Nj+pKeHM/9Hy7zu5SwEsrLal8GZgADzWyjmY03s4lmNtFrchdwpJktBP4D3O6cKyZwXqPAzOYDnwL3OOcUGCLSatKT45l4XB6fLiti9trmrs2RxvSIVhGJShU1dRx33xTyMjvwyoQj/C7HN3pEq4jIPqQkxHHVsX2ZuXo78zeU+F1OWFBgiEjUumhUDmmJcfzls9V+lxIWFBgiErXSkuK5ZEwv3l/0DRt3VPhdTrunwBCRqPbTI3Mx4K+frfG7lHZPgSEiUa17p2QuOKwnz89Yy4xV2/wup11TYIhI1Pvvs4aQ27UD1788l8KdVX6X024pMEQk6qUmxvHEpSPZVV3LLf+YTyTdbtCaFBgiIsDAg9K444whfL6ymNfmbPS7nHZJgSEi4rl0dC9G53bm9+8sobBMh6b2pMAQEfHExBh3X3AoVXUN/M9kjUi0JwWGiEgjeZmpXDu2H+8u2ML0lcV+l9OuKDBERPZw1XF96dU5hd++vZja+oZ9rxAlFBgiIntIio/lN2cOYUXhLp6bvtbvctoNBYaISBNOGpLNCYOyeOjjFbo3w6PAEBHZi9+cOYSaugbufn+p36W0CwoMEZG9yO3agQnH9uWtuZuYtUYPWlJgiIg049rj+9E9PYnfvb2Y+obovgNcgSEi0ozkhFhuP20QS7bs5M2vovsOcAWGiMg+nD2sO8NzOnHfB8uoqKnzuxzfKDBERPbBzPjvMwdTWFbN09Oi9+l8CgwRkSCM7N2ZMw7txlNTV/NNaXReZqvAEBEJ0u2nDqK+wfGnD5f5XYovFBgiIkHq1SWFK47K5Y2vNrJoU6nf5bQ5BYaIyH645vh+ZKQk8Pt3l0Tdg5YUGCIi+yE9OZ6bT+rPzNXb+WjJVr/LaVMKDBGR/XTJ6F70y0rl7veXUlMXPaPZKjBERPZTXGwMd5w+mDXF5bz05Tq/y2kzCgwRkQMwdmAmx/TvykMfr6CkosbvctqEAkNE5ACYGXecMZiyqloe+WSl3+W0CQWGiMgBGnRQRy4alcPzM9ayprjc73JCToEhItICN/9gAElxsUx8YQ6lFbV+lxNSCgwRkRbISkviictGsrp4Fz97fjZVtfV+lxQyCgwRkRY6un9XHrxoOAXrdnDra/Mj9oY+BYaISCs4c2h3fnnKIN5dsIUnp0bmiLYhDQwze8bMCs1s0V6Wp5vZZDObb2aLzeyKRssuN7MV3nR5KOsUEWkNE4/ry5lDu3HvB0v5bEWR3+W0ulDvYTwLnNrM8muBJc65YcBY4H4zSzCzzsBvgTHAaOC3ZpYR4lpFRFrEzLh33FD6ZaZy22vzI+4keEgDwzk3DWjuyekOSDMzA1K9tnXAKcBHzrntzrkdwEc0HzwiIu1CSkIcD1w4nG27avjN200eXAlbfp/DeBQYDGwGFgI3OucagB7AhkbtNnrzRETavUN7pnP9Cf3517zN/HvRN36X02r8DoxTgHlAd2A48KiZdQSsibZNXnZgZhPMrMDMCoqKIu+YoYiEp2uOz2PQQWncOXlxxDwH3O/AuAJ40wWsBNYAgwjsUeQ0ateTwF7I9zjnnnbO5Tvn8jMzM0NesIhIMOJjY7jznEPYXFrFY59GxtAhfgfGeuBEADPLBgYCq4EPgJPNLMM72X2yN09EJGyM7tOZ80f04C/T1kTE0CGhvqz2ZWAGMNDMNprZeDObaGYTvSZ3AUea2ULgP8Dtzrli59x2b9lsb7rTmyciElZ+dfog4mKNhz5e7ncpLRYXyg93zl2yj+WbCew9NLXsGeCZUNQlItJWstKSuPzIXJ6cuoprj+/HgOw0v0s6YH4fkhIRiXgTjulLSnwsf/54hd+ltIgCQ0QkxDI6JHDl0X14d+EWFm8u9bucA6bAEBFpAz87ui+dOyTwu7cXh+3ghAoMEZE2kJ4Szy9PGcjstTv417wm7xJo9xQYIiJt5ML8HIb1TOcP731NWVX4jTOlwBARaSMxMcad5xxC8a7qsDwBrsAQEWlDw3I6cVF+Dn+fvpblW8v8Lme/KDBERNrYL08dRGpiHL/51yI2l1RSV9/gd0lBCemNeyIi8n2dOyRw2ykD+e9/LuLIez4hNsbISkukb2YHbvnBAEb27ux3iU1SYIiI+OCyMb3on5XKqqJdbCmpYnNpJTNWbWPckzP40ehejD+6D30zU/0uczcKDBERH5gZh/ftwuF9u/zfvF3Vdfzpg2W8MHMdL325nsP7dubhS0aQlZbkY6Xf0TkMEZF2IjUxjt+dfTAzfnUCt586iAUbS7n46Zls3VnV7HptdSOgAkNEpJ3J6pjE1WPzeO7K0WwtreKip2awvbxmtzYVNXW8Mms9P3tuNuc9Pr1N6lJgiIi0U6NyO/P8+NFsLq3i6hfnUFP33dVUt7w6n1+9uZCvt5QxPKdTm1xppcAQEWnHRvbuzL0XDOXLNdv57duLcM4xfVUx/178DTed1J/Pbz+e3519MHGxof9zrpPeIiLt3LkjerBsaxlPTFnFzqo6VhXuomdGMhOPy8PM2qwOBYaISBj45SkDyUiJ5+73l+IcPH7pYSTFx7ZpDQoMEZEwYGZMODaPAdlpLNhYymmHHNTmNSgwRETCyNiBWYwdmOXLtnXSW0REgqLAEBGRoCgwREQkKAoMEREJigJDRESCosAQEZGgKDBERCQoCgwREQmKtdU46m3BzIqAdUE0TQdKD7DNnvODfd94/revuwLFQdS7P/UF06ap+U3V19yy9t6HPd9Heh9kC3OuAAAHHElEQVTgwPuhPkR3H/o759KDqtA5F3UT8PSBttlzfrDvG89vNK+gPfRhb/U1t6y996G5PkViH1rSD/VBfQi2xmg9JDW5BW32nB/s+8nNtDkQrdmHPec116em+nOgQtmHPd+rD3unPux7XjT0YZ8i6pBUuDGzAudcvt91tIT60H5EQj/Uh/YtWvcw2oun/S6gFagP7Uck9EN9aMe0hyEiIkHRHoaIiARFgdFKzOwZMys0s0UHsO5IM1toZivN7GFr9MxFM7vezJaZ2WIzu7d1q/5eHa3eBzP7nZltMrN53nR661e+Wx0h+Tl4y28zM2dmXVuv4ibrCMXP4S4zW+D9DD40s+6tX/ludYSiD/eZ2VKvH2+ZWafWr3y3OkLRhx96/5YbzCz8znMcyOVfmpq8rO1Y4DBg0QGsOws4AjDgfeA0b/7xwMdAovc+Kwz78DvgtnD+OXjLcoAPCNzn0zXc+gB0bNTmBuDJMOzDyUCc9/qPwB/DsA+DgYHAFCA/lPWHYtIeRitxzk0DtjeeZ2Z5ZvZvM5tjZp+Z2aA91zOzbgT+Mc9wgd+o54FzvcVXA/c456q9bRSGYR/aVAj78CDwSyDkJ/1C0Qfn3M5GTTsQ4n6EqA8fOufqvKYzgZ5h2IevnXPLQll3KCkwQutp4Hrn3EjgNuDxJtr0ADY2er/RmwcwADjGzL40s6lmNiqk1TatpX0AuM47jPCMmWWErtS9alEfzOxsYJNzbn6oC21Gi38OZvYHM9sAXAr8JoS17k1r/C5960oC39zbWmv2Iezomd4hYmapwJHAa40OhSc21bSJed9++4sDMoDDgVHAP8ysr/etJeRaqQ9PAHd57+8C7ifwj71NtLQPZpYC3EHgcIgvWunngHPuDuAOM/sv4Drgt61c6l61Vh+8z7oDqANeas0a96U1+xCuFBihEwOUOOeGN55pZrHAHO/t2wT+oDbete4JbPZebwTe9AJilpk1EBinpiiUhTfS4j4457Y2Wu8vwDuhLLgJLe1DHtAHmO/9kegJfGVmo51z34S49m+1xu9SY5OAd2nDwKCV+mBmlwNnAie21RenRlr75xB+/D6JEkkTkEujE2TAdOCH3msDhu1lvdkE9iK+PUF2ujd/InCn93oAsAHv3pkw6kO3Rm1uBl4Jt5/DHm3WEuKT3iH6OfRv1OZ64PUw7MOpwBIgM9S1h/p3iTA96e17AZEyAS8DW4BaAnsG4wl8M/03MN/7Rf/NXtbNBxYBq4BHvw0FIAF40Vv2FXBCGPbhBWAhsIDAt69u4daHPdqEPDBC9HN4w5u/gMDYQT3CsA8rCXxpmudNob7SKxR9OM/7rGpgK/BBKPvQ2pPu9BYRkaDoKikREQmKAkNERIKiwBARkaAoMEREJCgKDBERCYoCQyKame1q4+391cyGtNJn1Xujyy4ys8n7Gp3VzDqZ2TWtsW2RpuiyWoloZrbLOZfaip8X574bAC+kGtduZs8By51zf2imfS7wjnPukLaoT6KP9jAk6phZppm9YWazvekob/5oM5tuZnO9/w705v/UzF4zs8nAh2Y21symmNnr3vMZXmr0vIMp3z7nwMx2eQP+zTezmWaW7c3P897PNrM7g9wLmsF3gyGmmtl/zOwrCzxz4RyvzT1AnrdXcp/X9hfedhaY2f+04v9GiUIKDIlGfwYedM6NAi4A/urNXwoc65wbQWA01/9ttM4RwOXOuRO89yOAm4AhQF/gqCa20wGY6ZwbBkwDft5o+3/2tr/PMYa8sYpOJHCnPEAVcJ5z7jACz0y53wusXwGrnHPDnXO/MLOTgf7AaGA4MNLMjt3X9kT2RoMPSjQ6CRjSaMTRjmaWBqQDz5lZfwKji8Y3Wucj51zjZyPMcs5tBDCzeQTGHPp8j+3U8N1gi3OAH3ivj+C7Z21MAv60lzqTG332HOAjb74B/+v98W8gsOeR3cT6J3vTXO99KoEAmbaX7Yk0S4Eh0SgGOMI5V9l4ppk9AnzqnDvPOx8wpdHi8j0+o7rR63qa/rdU6747Sbi3Ns2pdM4NN7N0AsFzLfAwgedZZAIjnXO1ZrYWSGpifQPuds49tZ/bFWmSDklJNPqQwPMgADCzb4erTgc2ea9/GsLtzyRwKAzg4n01ds6VEnis6m1mFk+gzkIvLI4HentNy4C0Rqt+AFzpPccBM+thZlmt1AeJQgoMiXQpZrax0XQLgT+++d6J4CUEhpEHuBe428y+AGJDWNNNwC1mNgvoBpTuawXn3FwCI6ReTODBQflmVkBgb2Op12Yb8IV3Ge59zrkPCRzymmFmC4HX2T1QRPaLLqsVaWPeU/wqnXPOzC4GLnHOnbOv9UT8pnMYIm1vJPCod2VTCW34yFqRltAehoiIBEXnMEREJCgKDBERCYoCQ0REgqLAEBGRoCgwREQkKAoMEREJyv8H07pBS9txCwIAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7f6170a82588>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learner.lr_find(num_it=200)\n",
    "learner.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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='0' class='' max='500', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      0.00% [0/500 00:00<00:00]\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>LMAEMaskedLoss</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\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='48' class='' max='86', style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      55.81% [48/86 00:38<00:30 1.8626]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "learner.fit_one_cycle(500, max_lr=3e-2)"
   ]
  },
  {
   "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}