{"cells":[{"metadata":{},"cell_type":"markdown","source":"This kernel only shows training on 1JHN for the first 20 epochs.\nIf you want a description of our model, you can view this discussion:\nhttps://www.kaggle.com/c/champs-scalar-coupling/discussion/106271#latest-610727"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nfrom timeit import default_timer as timer","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!conda install pytorch torchvision -c pytorch --yes\n!pip install torch-scatter\n!pip install torch_geometric\n!pip install torch_sparse\n!pip install torch_cluster","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -U ase==3.17.0\n!pip install schnetpack","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import dataclasses\nfrom multiprocessing import cpu_count\nfrom pprint import pformat\nfrom time import time\nfrom typing import Dict, Callable, List\nimport schnetpack.atomistic as atm\nimport schnetpack.nn.blocks\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch_geometric.nn as gnn\nfrom schnetpack.data import Structure, AtomsLoader\nfrom schnetpack.datasets import *\nfrom schnetpack.nn import shifted_softplus, Dense\nfrom schnetpack.representation import SchNetInteraction\nfrom sklearn.model_selection import KFold\nfrom tensorboardX import SummaryWriter\nfrom torch.optim import Adam\nfrom torch_scatter import scatter_mean\nfrom sklearn.preprocessing import LabelEncoder, OneHotEncoder\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def time_to_str(t, mode='min'):\n    if mode=='min':\n        t  = int(t)/60\n        hr = t//60\n        min = t%60\n        return '%2d hr %02d min'%(hr,min)\n\n    elif mode=='sec':\n        t   = int(t)\n        min = t//60\n        sec = t%60\n        return '%2d min %02d sec'%(min,sec)\n\n\n    else:\n        raise NotImplementedError\n        \n\ndef ensure_dir(dir_name):\n    if not os.path.exists(dir_name):\n        os.makedirs(dir_name)\n\n        \nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n\n    def __init__(self):\n        self.count = 0\n        self.sum = 0\n        self.avg = 0\n        self.val = 0\n\n    def reset(self):\n        self.count = 0\n        self.sum = 0\n        self.avg = 0\n        self.val = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __format__(self, fmt):\n        return \"{self.val:{format}} ({self.avg:{format}})\".format(self=self, format=fmt)\n\n\nclass AverageMeterSet:\n    def __init__(self):\n        self.meters = {}\n\n    def __getitem__(self, key) -> AverageMeter:\n        return self.meters[key]\n\n    def update(self, name, value, n=1):\n        if not name in self.meters:\n            self.meters[name] = AverageMeter()\n        self.meters[name].update(value, n)\n\n    def reset(self):\n        for meter in self.meters.values():\n            meter.reset()\n\n    def values(self, postfix=''):\n        return {name + postfix: meter.val for name, meter in self.meters.items()}\n\n    def averages(self, postfix='/avg'):\n        return {name + postfix: meter.avg for name, meter in self.meters.items()}\n\n    def sums(self, postfix='/sum'):\n        return {name + postfix: meter.sum for name, meter in self.meters.items()}\n\n    def counts(self, postfix='/count'):\n        return {name + postfix: meter.count for name, meter in self.meters.items()}\n    \nimport math\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\n\nclass CyclicLR(_LRScheduler):\n    \"\"\"Sets the learning rate of each parameter group according to\n    cyclical learning rate policy (CLR). The policy cycles the learning\n    rate between two boundaries with a constant frequency, as detailed in\n    the paper `Cyclical Learning Rates for Training Neural Networks`_.\n    The distance between the two boundaries can be scaled on a per-iteration\n    or per-cycle basis.\n    Cyclical learning rate policy changes the learning rate after every batch.\n    `step` should be called after a batch has been used for training.\n    This class has three built-in policies, as put forth in the paper:\n    \"triangular\":\n        A basic triangular cycle w/ no amplitude scaling.\n    \"triangular2\":\n        A basic triangular cycle that scales initial amplitude by half each cycle.\n    \"exp_range\":\n        A cycle that scales initial amplitude by gamma**(cycle iterations) at each\n        cycle iteration.\n    This implementation was adapted from the github repo: `bckenstler/CLR`_\n    Args:\n        optimizer (Optimizer): Wrapped optimizer.\n        base_lr (float or list): Initial learning rate which is the\n            lower boundary in the cycle for each parameter group.\n        max_lr (float or list): Upper learning rate boundaries in the cycle\n            for each parameter group. Functionally,\n            it defines the cycle amplitude (max_lr - base_lr).\n            The lr at any cycle is the sum of base_lr\n            and some scaling of the amplitude; therefore\n            max_lr may not actually be reached depending on\n            scaling function.\n        step_size_up (int): Number of training iterations in the\n            increasing half of a cycle. Default: 2000\n        step_size_down (int): Number of training iterations in the\n            decreasing half of a cycle. If step_size_down is None,\n            it is set to step_size_up. Default: None\n        mode (str): One of {triangular, triangular2, exp_range}.\n            Values correspond to policies detailed above.\n            If scale_fn is not None, this argument is ignored.\n            Default: 'triangular'\n        gamma (float): Constant in 'exp_range' scaling function:\n            gamma**(cycle iterations)\n            Default: 1.0\n        scale_fn (function): Custom scaling policy defined by a single\n            argument lambda function, where\n            0 <= scale_fn(x) <= 1 for all x >= 0.\n            If specified, then 'mode' is ignored.\n            Default: None\n        scale_mode (str): {'cycle', 'iterations'}.\n            Defines whether scale_fn is evaluated on\n            cycle number or cycle iterations (training\n            iterations since start of cycle).\n            Default: 'cycle'\n        cycle_momentum (bool): If ``True``, momentum is cycled inversely\n            to learning rate between 'base_momentum' and 'max_momentum'.\n            Default: True\n        base_momentum (float or list): Lower momentum boundaries in the cycle\n            for each parameter group. Note that momentum is cycled inversely\n            to learning rate; at the peak of a cycle, momentum is\n            'base_momentum' and learning rate is 'max_lr'.\n            Default: 0.8\n        max_momentum (float or list): Upper momentum boundaries in the cycle\n            for each parameter group. Functionally,\n            it defines the cycle amplitude (max_momentum - base_momentum).\n            The momentum at any cycle is the difference of max_momentum\n            and some scaling of the amplitude; therefore\n            base_momentum may not actually be reached depending on\n            scaling function. Note that momentum is cycled inversely\n            to learning rate; at the start of a cycle, momentum is 'max_momentum'\n            and learning rate is 'base_lr'\n            Default: 0.9\n        last_epoch (int): The index of the last batch. This parameter is used when\n            resuming a training job. Since `step()` should be invoked after each\n            batch instead of after each epoch, this number represents the total\n            number of *batches* computed, not the total number of epochs computed.\n            When last_epoch=-1, the schedule is started from the beginning.\n            Default: -1\n    Example:\n        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)\n        >>> scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.01, max_lr=0.1)\n        >>> data_loader = torch.utils.data.DataLoader(...)\n        >>> for epoch in range(10):\n        >>>     for batch in data_loader:\n        >>>         train_batch(...)\n        >>>         scheduler.step()\n    .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186\n    .. _bckenstler/CLR: https://github.com/bckenstler/CLR\n    \"\"\"\n\n    def __init__(self,\n                 optimizer,\n                 base_lr,\n                 max_lr,\n                 step_size_up=2000,\n                 step_size_down=None,\n                 mode='triangular',\n                 gamma=1.,\n                 scale_fn=None,\n                 scale_mode='cycle',\n                 cycle_momentum=True,\n                 base_momentum=0.8,\n                 max_momentum=0.9,\n                 last_epoch=-1):\n\n        if not isinstance(optimizer, Optimizer):\n            raise TypeError('{} is not an Optimizer'.format(\n                type(optimizer).__name__))\n        self.optimizer = optimizer\n\n        base_lrs = self._format_param('base_lr', optimizer, base_lr)\n        if last_epoch == -1:\n            for lr, group in zip(base_lrs, optimizer.param_groups):\n                group['lr'] = lr\n\n        self.max_lrs = self._format_param('max_lr', optimizer, max_lr)\n\n        step_size_up = float(step_size_up)\n        step_size_down = float(step_size_down) if step_size_down is not None else step_size_up\n        self.total_size = step_size_up + step_size_down\n        self.step_ratio = step_size_up / self.total_size\n\n        if mode not in ['triangular', 'triangular2', 'exp_range'] \\\n                and scale_fn is None:\n            raise ValueError('mode is invalid and scale_fn is None')\n\n        self.mode = mode\n        self.gamma = gamma\n\n        if scale_fn is None:\n            if self.mode == 'triangular':\n                self.scale_fn = self._triangular_scale_fn\n                self.scale_mode = 'cycle'\n            elif self.mode == 'triangular2':\n                self.scale_fn = self._triangular2_scale_fn\n                self.scale_mode = 'cycle'\n            elif self.mode == 'exp_range':\n                self.scale_fn = self._exp_range_scale_fn\n                self.scale_mode = 'iterations'\n        else:\n            self.scale_fn = scale_fn\n            self.scale_mode = scale_mode\n\n        self.cycle_momentum = cycle_momentum\n        if cycle_momentum:\n            if 'momentum' not in optimizer.defaults:\n                raise ValueError('optimizer must support momentum with `cycle_momentum` option enabled')\n\n            base_momentums = self._format_param('base_momentum', optimizer, base_momentum)\n            if last_epoch == -1:\n                for momentum, group in zip(base_momentums, optimizer.param_groups):\n                    group['momentum'] = momentum\n            self.base_momentums = list(map(lambda group: group['momentum'], optimizer.param_groups))\n            self.max_momentums = self._format_param('max_momentum', optimizer, max_momentum)\n\n        super(CyclicLR, self).__init__(optimizer, last_epoch)\n\n    def _format_param(self, name, optimizer, param):\n        \"\"\"Return correctly formatted lr/momentum for each param group.\"\"\"\n        if isinstance(param, (list, tuple)):\n            if len(param) != len(optimizer.param_groups):\n                raise ValueError(\"expected {} values for {}, got {}\".format(\n                    len(optimizer.param_groups), name, len(param)))\n            return param\n        else:\n            return [param] * len(optimizer.param_groups)\n\n    def _triangular_scale_fn(self, x):\n        return 1.\n\n    def _triangular2_scale_fn(self, x):\n        return 1 / (2. ** (x - 1))\n\n    def _exp_range_scale_fn(self, x):\n        return self.gamma ** (x)\n\n    def get_lr(self):\n        \"\"\"Calculates the learning rate at batch index. This function treats\n        `self.last_epoch` as the last batch index.\n        If `self.cycle_momentum` is ``True``, this function has a side effect of\n        updating the optimizer's momentum.\n        \"\"\"\n        cycle = math.floor(1 + self.last_epoch / self.total_size)\n        x = 1. + self.last_epoch / self.total_size - cycle\n        if x <= self.step_ratio:\n            scale_factor = x / self.step_ratio\n        else:\n            scale_factor = (x - 1) / (self.step_ratio - 1)\n\n        lrs = []\n        for base_lr, max_lr in zip(self.base_lrs, self.max_lrs):\n            base_height = (max_lr - base_lr) * scale_factor\n            if self.scale_mode == 'cycle':\n                lr = base_lr + base_height * self.scale_fn(cycle)\n            else:\n                lr = base_lr + base_height * self.scale_fn(self.last_epoch)\n            lrs.append(lr)\n\n        if self.cycle_momentum:\n            momentums = []\n            for base_momentum, max_momentum in zip(self.base_momentums, self.max_momentums):\n                base_height = (max_momentum - base_momentum) * scale_factor\n                if self.scale_mode == 'cycle':\n                    momentum = max_momentum - base_height * self.scale_fn(cycle)\n                else:\n                    momentum = max_momentum - base_height * self.scale_fn(self.last_epoch)\n                momentums.append(momentum)\n            for param_group, momentum in zip(self.optimizer.param_groups, momentums):\n                param_group['momentum'] = momentum\n\n        return lrs\nARTIFACTS_DIR = 'data/artifacts'\nTYPES = ['1JHN', '2JHN', '3JHN', '1JHC', '2JHC', '3JHC', '2JHH', '3JHH']\n\nTYPE_WEIGHT = [\n    6.566171,\n    107.422157,\n    4.083679,\n    12.321967,\n    39.061047,\n    3.084091,\n    7.886997,\n    27.991149,\n]\nfrom abc import ABC, abstractmethod\n\nimport numpy as np\nimport optuna\nfrom optuna import Trial, Study\nfrom optuna.trial import FixedTrial\n\n\nclass Estimator(ABC):\n    @abstractmethod\n    def predict(self, X: np.ndarray):\n        return NotImplemented\n\n\nclass CVOptimizer(ABC):\n    def __init__(self, study: Study, is_one_cv=False, prune=False):\n        # self.X = X\n        # self.y = y\n        self.study = study\n        self.is_one_cv = is_one_cv\n        self.prune = prune\n\n    @abstractmethod\n    def fit(self, trial: Trial, X_train, y_train, X_val, y_val, step) -> Estimator:\n        return NotImplemented\n\n    @abstractmethod\n    def loss_fn(self, y_true, y_pred, trial: Trial):\n        return NotImplemented\n\n    @abstractmethod\n    def split(self, X, y):\n        return NotImplemented\n\n    @abstractmethod\n    def make_xy(self, trial: Trial):\n        return NotImplemented\n\n    def on_after_trial(self, trial: Trial, cv_models, cv_preds, loss_val):\n        pass\n\n    def best_models(self, return_preds=False):\n        fixed_trial = FixedTrial(self.study.best_params)\n        loss, models, cv_preds = self.objective(fixed_trial, return_model=True)\n\n        if return_preds:\n            return models, cv_preds\n        return models\n\n    def objective(self, trial: Trial, return_model=False):\n        X, y = self.make_xy(trial)\n        # X, y = self.X, self.y\n\n        loss_cv_train = []\n\n        cv_preds = np.zeros_like(y)\n        cv_preds.fill(np.nan)\n\n        models = []\n\n        for step, (train_idx, val_idx) in enumerate(self.split(X, y)):\n            X_train = X[train_idx]\n            X_val = X[val_idx]\n            y_train = y[train_idx]\n            y_val = y[val_idx]\n\n            model = self.fit(trial, X_train, y_train, X_val, y_val, step)\n\n            y_train_preds = model.predict(X_train)\n            y_val_preds = model.predict(X_val)\n            # from IPython.core.debugger import Pdb;\n            # Pdb().set_trace()\n            cv_preds[val_idx] = y_val_preds\n\n            mask_done = ~np.isnan(cv_preds)\n            intermediate_loss_train = self.loss_fn(y_train, y_train_preds, trial)\n            intermediate_loss_val = self.loss_fn(y[mask_done], cv_preds[mask_done], trial)\n\n            loss_cv_train.append(intermediate_loss_train)\n\n            trial.report(intermediate_loss_val, step)\n            if self.prune and trial.should_prune(step):\n                raise optuna.structs.TrialPruned()\n\n            models.append(model)\n\n            if self.is_one_cv:\n                break\n\n        mask_done = ~np.isnan(cv_preds)\n        loss_train = float(np.mean(loss_cv_train))\n        loss_val = float(self.loss_fn(y[mask_done], cv_preds[mask_done], trial))\n\n        trial.set_user_attr('train_loss', loss_train)\n        trial.set_user_attr('val_loss', loss_val)\n        trial.set_user_attr('is_one_cv', int(self.is_one_cv))\n\n        self.on_after_trial(trial, models, cv_preds, loss_val)\n\n        if return_model:\n            return loss_val, models, cv_preds\n\n        return loss_val\n\n    def optimize(self, n_trials=100, n_jobs=1):\n        self.study.optimize(self.objective, n_trials=n_trials, n_jobs=n_jobs)\n\nimport os\n\nimport feather\nimport numpy as np\nimport pandas as pd\nfrom sklearn.preprocessing import StandardScaler, MinMaxScaler\n\ndef get_main_train_df(coupling_type=None):\n    df = pd.read_csv('data/input/train.csv')\n    if coupling_type is None:\n        return df\n    if isinstance(coupling_type, str):\n        coupling_type = [coupling_type]\n\n    return df[df.type.isin(coupling_type)].reset_index(drop=True)\n\n\ndef get_main_test_df(coupling_type=None):\n    df = pd.read_csv('data/input/test.csv')\n    if coupling_type is None:\n        return df\n    if isinstance(coupling_type, str):\n        coupling_type = [coupling_type]\n\n    return df[df.type.isin(coupling_type)].reset_index(drop=True)\n\n\nclass CouplingProvider:\n    DF_CACHE_PATH = 'tmp/df_CouplingProvider.fth'\n\n    def __init__(self, scaler='std'):\n        self.cols = [\n            'scalar_coupling_constant',\n            'fc',\n            'sd',\n            'pso',\n            'dso',\n        ]\n        self.scaler = scaler\n        df = self.init_df()\n        self.grouped = df.groupby('molecule_name')\n        self.type_encoder = get_type_onehot_encoder()\n\n    def get_coupling_values(self, mol_name, types):\n        g = self.grouped.get_group(mol_name)\n        g = g[g.type.isin(types)]\n\n        # Case for H2O\n        if len(g.type.values) == 0:\n            return np.array([]).reshape(-1, 2), np.array([]).reshape(-1, 13)\n\n        encoded_types = self.type_encoder.transform(g.type.values.reshape(-1, 1))\n        feats = np.concatenate((\n            g[self.cols].values,\n            encoded_types,\n        ), axis=1)\n        edge_index = g[['atom_index_0', 'atom_index_1']].values\n\n        return edge_index, feats\n\n    def init_df(self):\n        if os.path.exists(self.DF_CACHE_PATH):\n            print('load df from cache')\n            return feather.read_dataframe(self.DF_CACHE_PATH)\n\n        df = get_main_train_df()\n        df = df.merge(pd.read_csv('data/input/scalar_coupling_contributions.csv').drop(columns=['type']),\n                      on=['molecule_name', 'atom_index_0', 'atom_index_1'])\n        types = sorted(df.type.unique())\n\n        # Scale in for each types\n        for t in types:\n            df_tmp = df[df.type == t]\n            if self.scaler == 'std':\n                scaled = StandardScaler().fit_transform(df_tmp[self.cols])\n            elif self.scaler == 'minmax':\n                scaled = MinMaxScaler().fit_transform(df_tmp[self.cols])\n            else:\n                raise Exception('not supported scaler')\n            df.loc[df.type == t, self.cols] = scaled\n\n        df[self.cols] = df[self.cols].astype(np.float32)\n        df.to_feather(self.DF_CACHE_PATH)\n        return df\nimport torch\nimport torch.nn as nn\nfrom schnetpack.datasets import *\nfrom schnetpack.nn.activations import shifted_softplus\n\n\ndef tile(a, dim, n_tile):\n    init_dim = a.size(dim)\n    repeat_idx = [1] * a.dim()\n    repeat_idx[dim] = n_tile\n    a = a.repeat(*repeat_idx)\n    order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)]))\n    return torch.index_select(a, dim, order_index)\n\n\nclass ShiftedSoftplus(nn.Module):\n    def __init__(self, inplace=False):\n        super(ShiftedSoftplus, self).__init__()\n\n    def forward(self, input):\n        return shifted_softplus(input)\n\n\ndef l1_loss(y_pred, y_true):\n    # TODO should use weight for each types?\n    loss = torch.abs(y_true - y_pred)\n    loss = loss.mean(dim=0)\n    loss = torch.log(loss)\n    loss = torch.sum(loss)\n\n    return loss\n\n\ndef smooth_l1_loss(delta=1.):\n    def loss_fn(y_pred, y_true):\n        loss = torch.abs(y_true - y_pred)\n        loss = torch.where(loss < delta, 0.5 * (loss ** 2), delta * loss - 0.5 * (delta ** 2))\n        loss = loss.mean(dim=0)\n        loss = torch.log(loss)\n        loss = torch.sum(loss)\n        return loss\n\n    return loss_fn\n\n\ndef log_cosh_loss(y_pred, y_true):\n    loss = torch.cosh(y_pred - y_true)\n    loss = torch.log(loss)\n    loss = loss.mean(dim=0)\n    loss = torch.log(loss)\n    loss = torch.sum(loss)\n    return loss\nimport numpy as np\n\n\nclass EarlyStopping(object):\n    def __init__(self, mode='min', min_delta=0, patience=10):\n        self.mode = mode\n        self.min_delta = min_delta\n        self.patience = patience\n        self.best = None\n        self.num_bad_epochs = 0\n        self.is_better = None\n        self._init_is_better(mode, min_delta)\n\n        if patience == 0:\n            self.is_better = lambda a, b: True\n\n    def step(self, metrics):\n        if self.best is None:\n            self.best = metrics\n            return False\n\n        if np.isnan(metrics):\n            return True\n\n        if self.is_better(metrics, self.best):\n            self.num_bad_epochs = 0\n            self.best = metrics\n        else:\n            self.num_bad_epochs += 1\n\n        if self.num_bad_epochs >= self.patience:\n            return True\n\n        return False\n\n    def _init_is_better(self, mode, min_delta):\n        if mode not in {'min', 'max'}:\n            raise ValueError('mode ' + mode + ' is unknown!')\n        if mode == 'min':\n            self.is_better = lambda a, best: a < best - min_delta\n        if mode == 'max':\n            self.is_better = lambda a, best: a > best + min_delta\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PROPERTIES = [\n    'scalar_coupling_constants',\n    'scalar_coupling_types',\n    'mulliken_charges',\n    'gasteiger_charges',\n    'dihedral_edge_index',\n    'dihedral_edge_attr',\n    'bond_edge_index',\n    'bond_edge_attr',\n    'sigma_iso',\n    'log_omega',\n    'kappa',\n    # 'potential_energy',\n    # 'dipole_moments',\n    # 'homo',\n    # 'lumo',\n    # 'U0',\n    # 'paths',\n]\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nnp.random.seed(0)\n\ntorch.manual_seed(0)\ntorch.backends.cudnn.deterministic = True\ntorch.backends.cudnn.benchmark = False\n\npd.options.display.max_rows = 999\npd.options.display.max_columns = 999\npd.options.display.width = 999","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def l1_loss(y_pred, y_true):\n    # TODO should use weight for each types?\n    loss = torch.abs(y_true - y_pred)\n    loss = loss.mean(dim=0)\n    loss = torch.log(loss)\n    loss = torch.sum(loss)\n\n    return loss\n\n\ndef smooth_l1_loss(delta=1.):\n    def loss_fn(y_pred, y_true):\n        loss = torch.abs(y_true - y_pred)\n        loss = torch.where(loss < delta, 0.5 * (loss ** 2), delta * loss - 0.5 * (delta ** 2))\n        loss = loss.mean(dim=0)\n        loss = torch.log(loss)\n        loss = torch.sum(loss)\n        return loss\n\n    return loss_fn\n\n\ndef _logcosh(x):\n    return x + F.softplus(-2. * x) - np.log(2.)\n\n\ndef log_cosh_loss(y_pred, y_true):\n    loss = _logcosh(y_pred - y_true)\n    loss = loss.mean(dim=0)\n    loss = torch.log(loss)\n    loss = torch.sum(loss)\n    # loss = loss.mean()\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SchNet(nn.Module):\n    def __init__(self, n_atom_basis=128, n_filters=128, n_interactions=1, cutoff=5.0, n_gaussians=25,\n                 normalize_filter=False, coupled_interactions=False,\n                 return_intermediate=False, max_z=100, interaction_block=SchNetInteraction, trainable_gaussians=False,\n                 distance_expansion=None):\n        super(SchNet, self).__init__()\n\n        n_extra_embeddings = 1\n        # n_extra_embeddings = 0\n\n        # atom type embeddings\n        self.embedding = nn.Embedding(max_z, n_atom_basis - n_extra_embeddings, padding_idx=0)\n\n        # spatial features\n        self.distances = schnetpack.nn.neighbors.AtomDistances()\n        if distance_expansion is None:\n            self.distance_expansion = schnetpack.nn.acsf.GaussianSmearing(0.0, cutoff, n_gaussians,\n                                                                          trainable=trainable_gaussians)\n        else:\n            self.distance_expansion = distance_expansion\n\n        self.return_intermediate = return_intermediate\n\n        # interaction network\n        if coupled_interactions:\n            self.interactions = nn.ModuleList([\n                                                  SchNetInteraction(n_atom_basis=n_atom_basis,\n                                                                    n_spatial_basis=n_gaussians,\n                                                                    n_filters=n_filters,\n                                                                    normalize_filter=normalize_filter)\n                                              ] * n_interactions)\n        else:\n            self.interactions = nn.ModuleList([\n                SchNetInteraction(n_atom_basis=n_atom_basis, n_spatial_basis=n_gaussians,\n                                  n_filters=n_filters, normalize_filter=normalize_filter)\n                for _ in range(n_interactions)\n            ])\n\n        n_dihedral_attrs = 2\n        n_dihedral_feats = 16\n        self.dihedral_conv = gnn.NNConv(n_atom_basis, n_dihedral_feats, nn.Sequential(\n            Dense(n_dihedral_attrs, n_atom_basis, activation=F.relu),\n            Dense(n_atom_basis, n_atom_basis * n_dihedral_feats),\n        ))\n\n        n_bond_attrs = 5\n        n_bond_feats = 32\n        self.bond_conv = gnn.NNConv(n_atom_basis, n_bond_feats, nn.Sequential(\n            Dense(n_bond_attrs, n_atom_basis, activation=F.relu),\n            Dense(n_atom_basis, n_atom_basis * n_bond_feats),\n        ))\n\n        self.init_fc = Dense(n_atom_basis + n_dihedral_feats + n_bond_feats, n_atom_basis, activation=shifted_softplus)\n\n    def forward(self, inputs):\n        atomic_numbers = inputs[Structure.Z]\n        positions = inputs[Structure.R]\n        cell = inputs[Structure.cell]\n        cell_offset = inputs[Structure.cell_offset]\n        neighbors = inputs[Structure.neighbors]\n        neighbor_mask = inputs[Structure.neighbor_mask]\n        \n        # atom embedding\n        x = self.embedding(atomic_numbers)\n        x = torch.cat((\n            x,\n            inputs['mulliken_charges'],\n            # inputs['gasteiger_charges'],\n        ), dim=2)\n\n        n_batch, n_atoms, n_embeddings = x.shape\n        # print(n_batch, n_atoms, n_embeddings)\n\n        # -------\n        batch_range = torch.arange(0, n_batch).to(device)\n\n        d_edge_index = inputs['dihedral_edge_index'].long()\n        d_edge_index = d_edge_index + batch_range.view(-1, 1, 1) * n_atoms\n        d_edge_index = d_edge_index.reshape(-1, 2)\n        d_edge_mask = d_edge_index[:, 0] != d_edge_index[:, 1]\n        d_edge_index = d_edge_index[d_edge_mask]\n\n        assert inputs['dihedral_edge_attr'].shape[2] == 3, 'n dihedral feats is not 3.'\n        d_edge_attr = inputs['dihedral_edge_attr'][:, :, :2]\n        d_edge_attr = d_edge_attr.reshape(-1, 2)[d_edge_mask]\n\n        # noinspection PyCallingNonCallable\n        dihedral_feat = self.dihedral_conv(x.reshape(-1, n_embeddings), torch.t(d_edge_index), d_edge_attr)\n        dihedral_feat = F.relu(dihedral_feat.view(n_batch, n_atoms, -1))\n        # print(dihedral_feat.shape)\n        # -------\n\n        # -------\n        a_edge_index = inputs['bond_edge_index'].long()\n        a_edge_index = a_edge_index + batch_range.view(-1, 1, 1) * n_atoms\n        a_edge_index = a_edge_index.reshape(-1, 2)\n        a_edge_mask = a_edge_index[:, 0] != a_edge_index[:, 1]\n        a_edge_index = a_edge_index[a_edge_mask]\n        \n        assert inputs['bond_edge_attr'].shape[2] == 5 # , 'n angle feats is not 2.'\n        a_edge_attr = inputs['bond_edge_attr'][:, :, :5]\n        a_edge_attr = a_edge_attr.reshape(-1, 5)[a_edge_mask]\n        \n        # noinspection PyCallingNonCallable\n        bond_feat = self.bond_conv(x.reshape(-1, n_embeddings), torch.t(a_edge_index), a_edge_attr)\n        bond_feat = F.relu(bond_feat.view(n_batch, n_atoms, -1))\n        # print(angle_feat.shape)\n        # -------\n\n        # spatial features\n        r_ij = self.distances(positions, neighbors, cell, cell_offset)\n        f_ij = self.distance_expansion(r_ij)\n\n        x = self.init_fc(torch.cat((x, dihedral_feat, bond_feat), dim=2))\n\n        for interaction in self.interactions:\n            v = interaction(x, r_ij, neighbors, neighbor_mask, f_ij=f_ij)\n            x = x + v\n\n\n\n        return x\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AtomPairwise(atm.OutputModule):\n    def __init__(self, n_in=128, n_out=1, n_layers=2, n_neurons=None,\n                 activation=schnetpack.nn.activations.shifted_softplus, return_contributions=False,\n                 requires_dr=False, create_graph=False, mean=None, stddev=None, atomref=None, max_z=100,\n                 train_embeddings=False):\n        super(AtomPairwise, self).__init__(n_in, n_out, requires_dr)\n\n        self.n_layers = n_layers\n\n        self.out_net = nn.Sequential(\n            schnetpack.nn.base.GetItem('atom_pair_rep'),\n            schnetpack.nn.blocks.MLP(n_in * 2, n_out, n_neurons, n_layers, activation),\n            # schnetpack.nn.blocks.MLP(n_in, n_out, n_neurons, n_layers, activation),\n        )\n\n        self.contributions_to_coupling_constant = Dense(4, 1)\n\n    def forward(self, inputs):\n        # atomic_numbers = inputs[Structure.Z]\n        # atom_mask = inputs[Structure.atom_mask]\n        atom_rep = inputs['representation']\n\n        n_batch = atom_rep.shape[0]\n\n        atom_index_0 = inputs['scalar_coupling_constants'][:, :, 0].long()\n        atom_index_1 = inputs['scalar_coupling_constants'][:, :, 1].long()\n        # atom_2 = inputs['paths'][:, :, 1].long()\n\n        coupling_batch_idx = tile(torch.arange(0, n_batch), 0, atom_index_0.shape[1])\n\n        atom_rep_0 = atom_rep[coupling_batch_idx, atom_index_0.view(-1)]\n        atom_rep_1 = atom_rep[coupling_batch_idx, atom_index_1.view(-1)]\n        #add pooling layer like set2set\n        \n        # atom_rep_2 = atom_rep[coupling_batch_idx, atom_2.view(-1)]\n\n        atom_pair_rep = torch.cat((atom_rep_0, atom_rep_1), dim=1)\n        # atom_pair_rep = torch.cat((atom_rep_0, atom_rep_1, atom_rep_2), dim=1)\n        # atom_pair_rep = atom_pair_rep.view(n_batch, -1, atom_pair_rep.shape[-1])\n        # print(atom_pair_rep.shape)\n\n        scc = inputs['scalar_coupling_constants'].view(-1, 7)  # idx0, idx1, cc, fc, sd, pso, dso\n        types = inputs['scalar_coupling_types'].long().view(-1)\n        mask = scc[:, 0] != scc[:, 1]\n\n        atom_pair_rep = atom_pair_rep[mask]\n\n        inputs['atom_pair_rep'] = atom_pair_rep\n\n        #predict just scc\n        contributions = self.out_net(inputs)\n\n        y = self.contributions_to_coupling_constant(contributions)\n        y = torch.cat((y, contributions), dim=1)\n\n        result = {\n            'y_pred': y,\n            'y_true': scc[mask, 2:],\n            'types': types[mask]\n        }\n\n        return result\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_atomwise_result(batch, result):\n    atom_mask = batch[Structure.atom_mask]\n    mask = atom_mask.byte()\n\n    y_pred = result['yi'][mask]  # sigma_iso, log_omega, kappa\n    y_true = torch.cat((\n        batch['sigma_iso'][mask],\n        batch['log_omega'][mask],\n        batch['kappa'][mask],\n    ), dim=1)\n\n    return y_pred, y_true\n\ndef get_type_encoder():\n    enc = LabelEncoder()\n    enc.fit(TYPES)\n\n\ndef calc_loss_atomwise_detail(y_pred, y_true):\n    with torch.no_grad():\n        loss = torch.abs(y_pred - y_true)\n        loss = loss.mean(dim=0)\n        # loss = torch.log(loss)\n        loss = loss.detach().cpu().numpy()\n    detail = {\n        'sigma_iso': loss[0],\n        'log_omega': loss[1],\n        'kappa': loss[2],\n    }\n    return detail\n\n\ndef calc_loss_contribution_detail(y_pred, y_true):\n    with torch.no_grad():\n        loss = torch.abs(y_pred - y_true)\n        loss = loss.mean(dim=0)\n        # loss = torch.log(loss)\n        loss = loss.detach().cpu().numpy()\n    detail = {\n        'cc': loss[0],\n        'fc': loss[1],\n        'sd': loss[2],\n        'pso': loss[3],\n        'dso': loss[4],\n    }\n    return detail\n\n\ndef calc_loss_energy_detail(y_pred, y_true):\n    with torch.no_grad():\n        loss = torch.abs(y_pred - y_true)\n        loss = loss.mean(dim=0)\n        loss = torch.log(loss)\n        loss = loss.detach().cpu().numpy()\n    detail = {\n        'alpha': loss[0],\n        'r2': loss[1]\n    }\n    return detail\n\n\ndef calc_loss_type_detail(y_pred, y_true, types):\n    abs_errs = torch.abs(y_pred - y_true)\n    abs_errs_cc = abs_errs.detach().cpu().numpy()[:, 0]\n    types = types.cpu().numpy().astype(int)\n\n    maes = pd.DataFrame({\n        'errs': abs_errs_cc,\n        'types': types,\n    }).groupby('types').agg({\n        'errs': [np.mean, 'size']\n    })\n    maes = maes.reset_index()\n    maes.columns = ['type', 'mae', 'n_data']\n    # maes['log_mae'] = np.log(maes['log_mae'])\n\n    return maes\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ARTIFACTS_DIR = 'data/artifacts'\nTYPES = ['1JHN', '2JHN', '3JHN', '1JHC', '2JHC', '3JHC', '2JHH', '3JHH']\n\nTYPE_WEIGHT = [\n    6.566171,\n    107.422157,\n    4.083679,\n    12.321967,\n    39.061047,\n    3.084091,\n    7.886997,\n    27.991149,\n]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_energy_result(batch, result):\n    y_pred = result['y'].to(device)  # homo, lumo, U0\n    \n    features=[]\n    if(VALIDATION):\n        molNam=molNameVal\n    else:\n        molNam=molNameTrain\n    for i in batch[\"_idx\"]:\n        features.append(qm9.get_group(molNam[i])[cols])\n    features=np.concatenate(features,axis=0)\n    features=torch.from_numpy(features).to(device)\n    features=features.type(torch.FloatTensor)\n\n    y_true = torch.cat((\n        features[:,0].reshape(-1,1),\n        features[:,1].reshape(-1,1)\n    ), dim=1)\n    y_true=y_true.to(device)\n    y_pred=y_pred.to(device)\n    return y_pred, y_true\n\ndef l1_loss(y_pred, y_true):\n    # TODO should use weight for each types?\n    loss = torch.abs(y_true - y_pred)\n    loss = loss.mean(dim=0)\n    loss = torch.log(loss)\n    loss = torch.sum(loss)\n\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_type_encoder():\n        enc = LabelEncoder()\n        enc.fit(TYPES)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@dataclasses.dataclass()\nclass Conf:\n        \n    lr: float = 1e-4\n    weight_decay: float = 1e-4\n\n    clr_max_lr: float = 3e-3\n    clr_base_lr: float = 3e-6\n    clr_gamma: float = 0.99994\n\n    train_batch: int = 32\n    val_batch: int = 256\n\n    n_atom_basis: int = 128\n    n_interactions: int = 1\n    coupled_interactions: bool = True\n    n_filters: int = 128\n    cutoff: float = 5.0\n    n_gaussians: int = 25\n\n    pairwise_layers: int = 2\n    atomwise_layers: int = 2\n    energy_layers: int = 2\n\n    atomwise_weight: float = 1.\n\n    pre_trained_path: str = None\n\n    optim: str = 'adam'\n    loss_fn: Callable = l1_loss\n\n    epochs: int = 25\n    is_save_epoch_fn: Callable = None\n    resume_from: Dict[str, int] = None\n\n    types: List[str] = None\n    db_path: str = None\n\n    seed: int = 0\n\n    is_one_cv: bool = True\n\n    device: str = device\n\n    exp_name: str = 'schnet'\n    exp_time: float = time()\n\n    logger_step = None\n    logger_epoch = None\n    \n    @staticmethod\n    def create_logger(name, filename):\n        logger = logging.getLogger(name)\n        logger.setLevel(logging.DEBUG)\n        if not logger.hasHandlers():\n            logger.addHandler(logging.FileHandler(filename))\n        return logger\n    \n    def get_type_encoder():\n        enc = LabelEncoder()\n        enc.fit(TYPES)\n\n    def __post_init__(self):\n        self.out_dir=os.getcwd()\n        if self.resume_from is not None:\n            assert os.path.exists(self.out_dir), '{} does not exist.'.format(self.out_dir)\n\n        ensure_dir(self.out_dir)\n\n        self.logger_step = self.create_logger('step_logger_{}'.format(self.exp_time),\n                                              '{}/step.log'.format(self.out_dir))\n        self.logger_epoch = self.create_logger('epoch_logger_{}'.format(self.exp_time),\n                                               '{}/epoch.log'.format(self.out_dir))\n        self.type_encoder = get_type_encoder()\n\n        with open('{}/conf.txt'.format(self.out_dir), 'w') as f:\n            f.write(str(self))\n\n        global device\n        device = self.device\n\n    #def out_dir(self):\n    #    return 'data/experiments/{}/{}'.format(self.exp_name, self.exp_time)\n\n    def __str__(self):\n        return pformat(dataclasses.asdict(self))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"conf=Conf(\n    is_one_cv=True,\n\n    device='cuda',\n\n    # train_batch=32,\n    train_batch=64,\n    # train_batch=128,\n    val_batch=256,\n\n    lr=.001,\n    clr_max_lr=3e-3,\n    clr_base_lr=3e-6,\n    # lr=3e-5,\n    # clr_max_lr=1e-3,\n    # clr_base_lr=1e-6,\n    clr_gamma=0.999991,\n    weight_decay=1e-4,\n    # weight_decay=3e-3,\n\n    n_atom_basis=256 + 64,\n    n_interactions=6,\n    coupled_interactions=False,\n    n_filters=256,#256*2,\n    pairwise_layers=3,#4,\n    atomwise_layers=3,#4,\n    n_gaussians=25,#40,\n\n    # atomwise_weight=0.3,\n\n    # loss_fn=log_cosh_loss,\n    # loss_fn=smooth_l1_loss(),\n    loss_fn=l1_loss,\n    optim='adam',\n\n    epochs= 20,    #CHANGE THIS TO 400-500\n    # is_save_epoch_fn=lambda x: x % 20 == 19,\n\n    # pre_trained_path='data/experiments/schnet-qm9/1563077515.9908571/0-380.ckpt',  # n_filter=128\n    # pre_trained_path='data/experiments/schnet-qm9/1563103994.1368518/0-420.ckpt',  # n_filter=256\n    # pre_trained_path='data/experiments/schnet-qm9/1563622942.702352/0-300.ckpt',  # n_filter=256\n    # pre_trained_path='data/experiments/schnet-qm9/1563624328.5389972/0-300.ckpt',  # n_filter=256 basis=256\n    types=['1JHN'],\n    db_path=\"../input/schnet-data/CHAMPS_with_bond_train_1JHN.db\",\n    exp_time=time()\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def log_hist(df_hist, logger: logging.Logger, types):\n    df_hist[\"time\"]=time_to_str((timer() - start),'min')\n    last = df_hist.tail(1)\n    best = df_hist.sort_values('val_loss_cc', ascending=True).head(1)\n    summary = pd.concat((last, best)).reset_index(drop=True)\n    summary['name'] = ['Last', 'Best']\n    logger.debug(summary[[\n                             'name',\n                             'epoch',\n                             'train_loss_coupling',\n                             'train_loss_atomwise',\n                             # 'train_loss_energy',\n                             'val_loss_coupling',\n                             'val_loss_atomwise',\n                             # 'train_loss_total',\n                             # 'val_loss_total',\n                         ] + [\n                             'train_loss_{}'.format(t) for t in types\n                         ] + [\n                             'val_loss_{}'.format(t) for t in types\n                         ]])\n    \n    print(summary[[          \n                             'name',\n                             'epoch',\n                             'train_loss_coupling',\n                             # 'train_loss_energy',\n                             'val_loss_coupling',\n                             'val_loss_atomwise',\n                             \"time\",\n                             # 'train_loss_total',\n                             # 'val_loss_total',\n                         ] + [\n                             'train_loss_{}'.format(t) for t in types\n                         ] + [\n                             'val_loss_{}'.format(t) for t in types\n                         ]])\n\n\ndef write_on_board(df_hist, writer, conf: Conf):\n    row = df_hist.tail(1).iloc[0]\n\n    writer.add_scalars('{}/lr'.format(conf.exp_name), {\n        '{}'.format(conf.exp_time): row.lr,\n    }, row.epoch)\n\n    for tag in ['cc', 'fc', 'sd', 'pso', 'dso']:\n        writer.add_scalars('{}/loss/coupling/{}'.format(conf.exp_name, tag), {\n            '{}_train'.format(conf.exp_time): row['train_loss_{}'.format(tag)],\n            '{}_val'.format(conf.exp_time): row['val_loss_{}'.format(tag)],\n        }, row.epoch)\n    writer.add_scalars('{}/loss/coupling/total'.format(conf.exp_name), {\n        '{}_train'.format(conf.exp_time): row.train_loss_coupling,\n        '{}_val'.format(conf.exp_time): row.val_loss_coupling,\n    }, row.epoch)\n\n    for tag in ['sigma_iso', 'log_omega', 'kappa']:\n        writer.add_scalars('{}/loss/atomwise/{}'.format(conf.exp_name, tag), {\n            '{}_train'.format(conf.exp_time): row['train_loss_{}'.format(tag)],\n            '{}_val'.format(conf.exp_time): row['val_loss_{}'.format(tag)],\n        }, row.epoch)\n    writer.add_scalars('{}/loss/atomwise/total'.format(conf.exp_name), {\n        '{}_train'.format(conf.exp_time): row.train_loss_atomwise,\n        '{}_val'.format(conf.exp_time): row.val_loss_atomwise,\n    }, row.epoch)\n\n    # for tag in ['homo', 'lumo', 'U0']:\n    #     writer.add_scalars('{}/loss/energy/{}'.format(conf.exp_name, tag), {\n    #         '{}_train'.format(conf.exp_time): row['train_loss_{}'.format(tag)],\n    #         '{}_val'.format(conf.exp_time): row['val_loss_{}'.format(tag)],\n    #     }, row.epoch)\n    # writer.add_scalars('{}/loss/energy/total'.format(conf.exp_name), {\n    #     '{}_train'.format(conf.exp_time): row.train_loss_energy,\n    #     '{}_val'.format(conf.exp_time): row.val_loss_energy,\n    # }, row.epoch)\n\n    # writer.add_scalars('{}/loss_dipole'.format(exp_name), {\n    #     '{}_train'.format(exp_id): row.train_loss_dipole,\n    #     '{}_val'.format(exp_id): row.val_loss_dipole,\n    # }, row.epoch)\n\n    for tag in conf.types:\n        writer.add_scalars('{}/loss/type/{}'.format(conf.exp_name, tag), {\n            '{}_train'.format(conf.exp_time): row['train_loss_{}'.format(tag)],\n            '{}_val'.format(conf.exp_time): row['val_loss_{}'.format(tag)],\n        }, row.epoch)\n    writer.add_scalars('{}/loss/type/total'.format(conf.exp_name), {\n        '{}_train'.format(conf.exp_time): row['train_loss_total'],\n        '{}_val'.format(conf.exp_time): row['val_loss_total'],\n    }, row.epoch)\n\n\ndef load_pre_trained(conf: Conf):\n    ckpt = torch.load(conf.pre_trained_path, map_location=device)\n\n    reps = SchNet(\n        n_interactions=conf.n_interactions,\n        coupled_interactions=False,\n        n_filters=conf.n_filters,\n        n_atom_basis=conf.n_atom_basis\n    )\n    model = atm.AtomisticModel(reps, [\n        atm.Atomwise(\n            return_contributions=True,\n            n_in=conf.n_atom_basis,\n            n_out=1,\n            # n_layers=3,\n        ),\n        atm.Atomwise(\n            return_contributions=False,\n            n_in=conf.n_atom_basis,\n            n_out=3,\n            # n_layers=3,\n        ),\n    ])\n    model.load_state_dict(ckpt['model'], strict=False)\n\n    return model\n\nfrom sklearn.preprocessing import LabelEncoder, OneHotEncoder\ndef ensure_dir(dir_name):\n    if not os.path.exists(dir_name):\n        os.makedirs(dir_name)\ndef get_type_encoder() -> LabelEncoder:\n    enc = LabelEncoder()\n    enc.fit(TYPES)\n    return enc\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"VALIDATION=False\ndef train(loader, model: atm.AtomisticModel, optimizer, scheduler, conf: Conf):\n    global VALIDATION\n    VALIDATION=False\n    meters = AverageMeterSet()\n    model.train()\n\n    for i, batch in enumerate(loader):\n        scheduler.step()\n        meters.update('lr', optimizer.param_groups[0]['lr'])\n\n        batch = {\n            k: v.to(device)\n            for k, v in batch.items()\n        }\n        result = model(batch)\n\n        loss_fn = conf.loss_fn\n        \n        # ----- Coupling Loss -----\n        # y_pred_coupling, y_true_coupling, types = make_coupling_result(batch, result[model.output_modules[0]])\n        coupling_res = result[model.output_modules[0]]\n        n_pairs = len(coupling_res['types'])\n        coupling_loss = loss_fn(coupling_res['y_pred'], coupling_res['y_true'])\n        meters.update('loss_coupling', coupling_loss.item(), n_pairs)\n \n        \n        # ----- Atomwise Loss -----\n        y_pred_atomwise, y_true_atomwise = make_atomwise_result(batch, result[model.output_modules[1]])\n        n_atoms = len(y_pred_atomwise)\n        atomwise_loss = loss_fn(y_pred_atomwise, y_true_atomwise)\n        meters.update('loss_atomwise', atomwise_loss.item(), n_atoms)\n\n    \n\n        # ----- Dipole Moment Loss -----\n        # dipole_loss = loss_fn(torch.abs(batch['dipole_moments'] - result[model.output_modules[2]]['y']))\n        # n_mols = len(batch['dipole_moments'])\n\n        # ----- Metric of coupling contributions -----\n        contribution_detail = calc_loss_contribution_detail(coupling_res['y_pred'], coupling_res['y_true'])\n        meters.update('loss_cc', contribution_detail['cc'], n_pairs)\n        meters.update('loss_fc', contribution_detail['fc'], n_pairs)\n        meters.update('loss_sd', contribution_detail['sd'], n_pairs)\n        meters.update('loss_pso', contribution_detail['pso'], n_pairs)\n        meters.update('loss_dso', contribution_detail['dso'], n_pairs)\n\n        # ----- Metric of atomwise -----\n        atomwise_detail = calc_loss_atomwise_detail(y_pred_atomwise, y_true_atomwise)\n        meters.update('loss_sigma_iso', atomwise_detail['sigma_iso'], n_atoms)\n        meters.update('loss_log_omega', atomwise_detail['log_omega'], n_atoms)\n        meters.update('loss_kappa', atomwise_detail['kappa'], n_atoms)\n\n        # ----- Metric of energy -----\n        # meters.update('loss_homo', energy_detail['homo'])\n        # meters.update('loss_lumo', energy_detail['lumo'])\n        # meters.update('loss_U0', energy_detail['U0'])\n\n        # ----- Metric of dipole -----\n        # meters.update('loss_dipole', dipole_loss.item(), n_mols)\n\n        # ----- Metric for each types -----\n        type_detail = calc_loss_type_detail(coupling_res['y_pred'], coupling_res['y_true'], coupling_res['types'])\n        conf.type_encoder = get_type_encoder()\n        type_detail['type_name'] = conf.type_encoder.inverse_transform(type_detail.type.values)\n        for _, row in type_detail.iterrows():\n            meters.update('loss_{}'.format(row.type_name), row.mae, row.n_data)\n\n        # ----- Total Loss -----\n        # loss = coupling_loss\n        loss = coupling_loss + atomwise_loss * conf.atomwise_weight\n        # loss = coupling_loss + atomwise_loss + energy_loss\n        # loss = coupling_loss + atomwise_loss + dipole_loss\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n    return {\n        'lr': meters['lr'].avg,\n        'train_loss_cc': meters['loss_cc'].avg,\n        'train_loss_fc': meters['loss_fc'].avg,\n        'train_loss_sd': meters['loss_sd'].avg,\n        'train_loss_pso': meters['loss_pso'].avg,\n        'train_loss_dso': meters['loss_dso'].avg,\n        'train_loss_coupling': meters['loss_coupling'].avg,\n\n        'train_loss_sigma_iso': meters['loss_sigma_iso'].avg,\n        'train_loss_log_omega': meters['loss_log_omega'].avg,\n        'train_loss_kappa': meters['loss_kappa'].avg,\n        'train_loss_atomwise': meters['loss_atomwise'].avg,\n\n        # 'train_loss_homo': meters['loss_homo'].avg,\n        # 'train_loss_lumo': meters['loss_lumo'].avg,\n        # 'train_loss_U0': meters['loss_U0'].avg,\n\n        **{\n            'train_loss_{}'.format(t): np.log(meters['loss_{}'.format(t)].avg)\n            for t in conf.types\n        },\n        'train_loss_total': np.mean([\n            np.log(meters['loss_{}'.format(t)].avg)\n            for t in conf.types\n        ]),\n\n        # 'train_loss_dipole': meters['loss_dipole'].avg,\n    }\n\ndef validate(loader, model, conf: Conf):\n    global VALIDATION\n    validation=True\n    \n    meters = AverageMeterSet()\n    model.eval()\n\n    for i, batch in enumerate(loader):\n        batch = {\n            k: v.to(device)\n            for k, v in batch.items()\n        }\n\n        loss_fn = conf.loss_fn\n\n        with torch.no_grad():\n            result = model(batch)\n            # ----- Coupling Loss -----\n            # y_pred_coupling, y_true_coupling, types = make_coupling_result(batch, result[model.output_modules[0]])\n            coupling_res = result[model.output_modules[0]]\n            n_pairs = len(coupling_res['types'])\n\n            coupling_loss = loss_fn(coupling_res['y_pred'], coupling_res['y_true'])\n            meters.update('loss_coupling', coupling_loss.item(), n_pairs)\n            \n            \n            # ----- Atomwise Loss -----\n            y_pred_atomwise, y_true_atomwise = make_atomwise_result(batch, result[model.output_modules[1]])\n            n_atoms = len(y_pred_atomwise)\n            atomwise_loss = loss_fn(y_pred_atomwise, y_true_atomwise)\n            meters.update('loss_atomwise', atomwise_loss.item(), n_atoms)\n\n\n            # ----- Dipole Moment Loss -----\n            # dipole_loss = loss_fn(torch.abs(batch['dipole_moments'] - result[model.output_modules[2]]['y']))\n            # n_mols = len(batch['dipole_moments'])\n\n            # ----- Metric of coupling contributions -----\n            contribution_detail = calc_loss_contribution_detail(coupling_res['y_pred'], coupling_res['y_true'])\n            meters.update('loss_cc', contribution_detail['cc'], n_pairs)\n            meters.update('loss_fc', contribution_detail['fc'], n_pairs)\n            meters.update('loss_sd', contribution_detail['sd'], n_pairs)\n            meters.update('loss_pso', contribution_detail['pso'], n_pairs)\n            meters.update('loss_dso', contribution_detail['dso'], n_pairs)\n\n            # ----- Metric of atomwise -----\n            atomwise_detail = calc_loss_atomwise_detail(y_pred_atomwise, y_true_atomwise)\n            meters.update('loss_sigma_iso', atomwise_detail['sigma_iso'], n_atoms)\n            meters.update('loss_log_omega', atomwise_detail['log_omega'], n_atoms)\n            meters.update('loss_kappa', atomwise_detail['kappa'], n_atoms)\n\n            # ----- Metric of energy -----\n            # energy_detail = calc_loss_energy_detail(y_pred_energy, y_true_energy)\n            # meters.update('loss_homo', energy_detail['homo'])\n            # meters.update('loss_lumo', energy_detail['lumo'])\n            # meters.update('loss_U0', energy_detail['U0'])\n\n            # ----- Metric of dipole -----\n            # meters.update('loss_dipole', dipole_loss.item(), n_mols)\n\n            # ----- Metric for each types -----\n            type_detail = calc_loss_type_detail(coupling_res['y_pred'], coupling_res['y_true'], coupling_res['types'])\n            type_detail['type_name'] = conf.type_encoder.inverse_transform(type_detail.type.values)\n            for _, row in type_detail.iterrows():\n                meters.update('loss_{}'.format(row.type_name), row.mae, row.n_data)\n\n    return {\n        'val_loss_cc': meters['loss_cc'].avg,\n        'val_loss_fc': meters['loss_fc'].avg,\n        'val_loss_sd': meters['loss_sd'].avg,\n        'val_loss_pso': meters['loss_pso'].avg,\n        'val_loss_dso': meters['loss_dso'].avg,\n        'val_loss_coupling': meters['loss_coupling'].avg,\n\n        'val_loss_sigma_iso': meters['loss_sigma_iso'].avg,\n        'val_loss_log_omega': meters['loss_log_omega'].avg,\n        'val_loss_kappa': meters['loss_kappa'].avg,\n        'val_loss_atomwise': meters['loss_atomwise'].avg,\n\n        #'val_loss_homo': meters['loss_homo'].avg,\n        # 'val_loss_lumo': meters['loss_lumo'].avg,\n        # 'val_loss_U0': meters['loss_U0'].avg,\n\n        **{\n            'val_loss_{}'.format(t): np.log(meters['loss_{}'.format(t)].avg)\n            for t in conf.types\n        },\n        'val_loss_total': np.mean([\n            np.log(meters['loss_{}'.format(t)].avg)\n            for t in conf.types\n        ]),\n\n        # 'val_loss_dipole': meters['loss_dipole'].avg,\n    }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"start = timer()\nconf.db_path = '../input/schnet-data/CHAMPS_with_bond_train_1JHN.db'\ndb = schnetpack.data.AtomsData(conf.db_path, properties=PROPERTIES)\nfolds = KFold(n_splits=4, random_state=1, shuffle=True)\nfor cv, (train_idx, val_idx) in enumerate(folds.split(range(len(db)))):\n    # Use partial data for small experiments\n    # train_idx = np.random.RandomState(conf.seed).permutation(train_idx)[:21420]\n    # val_idx = np.random.RandomState(conf.seed).permutation(val_idx)[:7140]\n\n    # if conf.types == ['1JHC']:\n    #     train_idx = np.setdiff1d(train_idx, [14])\n    #     val_idx = np.setdiff1d(val_idx, [14])\n    # if conf.types == ['2JHC']:\n    #     train_idx = np.setdiff1d(train_idx, [12])\n    #     val_idx = np.setdiff1d(val_idx, [12])\n    # if conf.types == ['3JHC']:\n    #     train_idx = np.setdiff1d(train_idx, [6])\n    #     val_idx = np.setdiff1d(val_idx, [6])\n\n    train_data = db.create_subset(train_idx)\n    val_data = db.create_subset(val_idx)\n    print(cv, len(train_data), len(val_data))\n    train_loader = AtomsLoader(train_data, batch_size=conf.train_batch, shuffle=True, num_workers=4)\n    val_loader = AtomsLoader(val_data, batch_size=conf.val_batch, num_workers=4)\n\n    reps = SchNet(\n        n_atom_basis=conf.n_atom_basis,\n        n_interactions=conf.n_interactions,\n        coupled_interactions=conf.coupled_interactions,\n        n_filters=conf.n_filters,\n        cutoff=conf.cutoff,\n        n_gaussians=conf.n_gaussians,\n    )\n    if conf.pre_trained_path is not None:\n        pre_trained = load_pre_trained(conf)\n        reps.load_state_dict(pre_trained.representation.state_dict(), strict=False)\n        # n_in = conf.n_atom_basis + 17\n        n_in = conf.n_atom_basis + 33\n    else:\n        # n_in = conf.n_atom_basis + 16\n        n_in = conf.n_atom_basis\n        # n_in = conf.n_atom_basis + 32\n\n    model = atm.AtomisticModel(reps, [\n        AtomPairwise(\n            n_in=n_in,\n            n_out=4,\n            n_layers=conf.pairwise_layers,\n        ),\n        atm.Atomwise(\n            return_contributions=True,\n            n_in=n_in,\n            n_out=3,\n            n_layers=conf.atomwise_layers,\n        ),\n            # atm.Energy(n_in=128),\n            # atm.DipoleMoment(n_in=128),\n            # atm.ElementalDipoleMoment(n_in=128),\n        ])\n    model = model.to(device)\n\n    if conf.optim == 'adam':\n        opt = Adam(model.parameters(), lr=conf.lr, weight_decay=conf.weight_decay)\n    elif conf.optim == 'adamw':\n        opt = optim.AdamW(model.parameters(), lr=conf.lr, weight_decay=conf.weight_decay)\n    else:\n        raise Exception(\"Not supported optim {}\".format(conf.optim))\n    scheduler = CyclicLR(\n        opt,\n        base_lr=conf.clr_base_lr,\n        max_lr=conf.clr_max_lr,\n        step_size_up=len(train_loader) * 10,\n        mode=\"exp_range\",\n        gamma=conf.clr_gamma,\n        cycle_momentum=False,\n    )\n    early_stopping = EarlyStopping(patience=100)\n\n    if conf.resume_from is not None:\n        cv_resume = conf.resume_from['cv']\n        start_epoch = conf.resume_from['epoch']\n        if cv_resume != cv:\n            continue\n        ckpt = torch.load('{}/{}-{:03d}.ckpt'.format(conf.out_dir, cv, start_epoch))\n\n        model.load_state_dict(ckpt['model'])\n        opt.load_state_dict(ckpt['optimizer'])\n        scheduler.load_state_dict(ckpt['scheduler'])\n\n        writer = SummaryWriter(logdir=ckpt['writer_logdir'], purge_step=start_epoch)\n        hist = pd.read_csv('{}/{}.csv'.format(conf.out_dir, cv)).to_dict('records')\n\n        print('Loaded checkpoint cv {}, epoch {} from {}'.format(cv, start_epoch, conf.out_dir))\n    else:\n        hist = []\n#         writer = SummaryWriter()\n        start_epoch = 0\n\n    for epoch in range(start_epoch, conf.epochs):\n        train_result = train(train_loader, model, opt, scheduler, conf)\n        val_result = validate(val_loader, model, conf)\n        hist.append({\n            'epoch': epoch,\n            **train_result,\n            **val_result,\n        })\n        df_hist = pd.DataFrame(hist)\n#         print(df_hist)\n        log_hist(df_hist, conf.logger_epoch, conf.types)\n#         write_on_board(df_hist, writer, conf)\n\n        if conf.is_save_epoch_fn is not None and conf.is_save_epoch_fn(epoch):\n            torch.save({\n                'model': model.state_dict(),\n                'optimizer': opt.state_dict(),\n                'scheduler': scheduler.state_dict(),\n                'writer_logdir': writer.logdir,\n            }, '{}/{}-{:03d}.ckpt'.format(conf.out_dir, cv, epoch + 1))\n            df_hist.to_csv('{}/{}.csv'.format(conf.out_dir, cv))\n            print('Saved checkpoint {}/{}-{:03d}.ckpt'.format(conf.out_dir, cv, epoch + 1))\n\n        should_stop = early_stopping.step(val_result['val_loss_total'])\n        if should_stop:\n            print('Early stopping at {}'.format(epoch))\n            break\n\n    df_hist = pd.DataFrame(hist)\n    best = df_hist.sort_values('val_loss_cc', ascending=True).head(1).iloc[0]\n    print(best)\n\n#     writer.close()\n    if conf.is_one_cv:\n        break\n","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}