copolymer_informatics/meta_learner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Meta learner\n",
    "\n",
    "## Prepare data set: split, scaler, and create tensorflow data sets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.model_selection import train_test_split, StratifiedKFold\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "tf.random.set_seed(123)\n",
    "\n",
    "# Read the data set\n",
    "# df = pd.read_pickle('pkls/dataset.pkl')\n",
    "\n",
    "df = df.sample(frac=1, random_state=123)\n",
    "\n",
    "df['scaled_value'] = df['val']\n",
    "skf = StratifiedKFold(n_splits=5)\n",
    "\n",
    "props = df.prop.unique().tolist()\n",
    "\n",
    "datasets = []\n",
    "def get_dataset(df):\n",
    "    # return TF data set\n",
    "    fps_mix = np.stack(df.fps_mix.values).astype(np.float32)\n",
    "    selector = np.stack(df.dummy.values).astype(np.float32)   \n",
    "    target = df.scaled_value.astype(np.float32)[:, np.newaxis]\n",
    "    dataset = tf.data.Dataset.from_tensor_slices(({'sel': selector, 'fps_mix': fps_mix, 'prop': df.prop}, target))\n",
    "\n",
    "    dataset = dataset.cache().batch(200).prefetch(tf.data.experimental.AUTOTUNE)\n",
    "    return dataset\n",
    "\n",
    "\n",
    "training_df, test_df = train_test_split(df, test_size=0.2, stratify=df.prop, random_state=123)\n",
    "training_df, test_df = training_df.copy(), test_df.copy()\n",
    "\n",
    "for train_index, val_index in skf.split(training_df, training_df.prop):\n",
    "    # iterte over 5 splits\n",
    "    train_df = training_df.iloc[train_index].copy()\n",
    "    val_df = training_df.iloc[val_index].copy()\n",
    "    \n",
    "    # scale target values\n",
    "    property_scaler = {}\n",
    "    for prop in props:\n",
    "        property_scaler[prop] = MinMaxScaler()\n",
    "\n",
    "        # train\n",
    "        cond = train_df[train_df.prop == prop].index\n",
    "        train_df.loc[cond, ['scaled_value']] = property_scaler[prop].fit_transform(train_df.loc[cond, ['scaled_value']])\n",
    "        \n",
    "        # val\n",
    "        cond = val_df[val_df.prop == prop].index\n",
    "        val_df.loc[cond, ['scaled_value']] = property_scaler[prop].transform(val_df.loc[cond, ['scaled_value']])\n",
    "\n",
    "    datasets.append({'train': get_dataset(train_df), 'val': get_dataset(val_df), 'property_scaler': property_scaler})\n",
    "\n",
    "# Create final dataset for meta learner\n",
    "property_scaler_final = {}\n",
    "for prop in props:\n",
    "    property_scaler_final[prop] = MinMaxScaler()\n",
    "    \n",
    "   # train\n",
    "    cond = training_df[training_df.prop == prop].index\n",
    "    training_df.loc[cond, ['scaled_value']] = property_scaler_final[prop].fit_transform(training_df.loc[cond, ['scaled_value']])\n",
    "\n",
    "    # val\n",
    "    cond = test_df[test_df.prop == prop].index\n",
    "    test_df.loc[cond, ['scaled_value']] = property_scaler_final[prop].transform(test_df.loc[cond, ['scaled_value']])\n",
    "    \n",
    "datasets_final = {'train': get_dataset(training_df), 'test': get_dataset(test_df), 'property_scaler': property_scaler_final}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition of the meta learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tensorflow.keras as tfk\n",
    "import tensorflow as tf\n",
    "from datetime import datetime \n",
    "from tensorflow.python.keras.engine import data_adapter\n",
    "from kerastuner import HyperModel\n",
    "from kerastuner.tuners import Hyperband, RandomSearch\n",
    "\n",
    "\n",
    "\n",
    "class MetaModel(tfk.Model):\n",
    "    \"\"\"Meta learner class\"\"\"\n",
    "\n",
    "    def __init__(self, hp):\n",
    "        super().__init__()\n",
    "        self.base_models = []\n",
    "        for num in range(5):\n",
    "            # load all cross-validation models\n",
    "            model = tf.keras.models.load_model(f'models/fp/{num}')\n",
    "            model.trainable = False\n",
    "            self.base_models.append(model)\n",
    "        \n",
    "\n",
    "        self.my_layers = []\n",
    "        for i in range(hp.Int('num_layers', 1, 2)): \n",
    "            new_step = [               \n",
    "            tf.keras.layers.Dense(units=hp.Int('units_' + str(i),\n",
    "                                            min_value=64,\n",
    "                                            max_value=544,\n",
    "                                            step=64),),\n",
    "            \n",
    "            tf.keras.layers.PReLU(),\n",
    "            tf.keras.layers.Dropout(hp.Float(\n",
    "                'dropout_' + str(i),\n",
    "                min_value=0.0,\n",
    "                max_value=0.5,\n",
    "                default=0.25,\n",
    "                step=0.05,\n",
    "            )),]\n",
    "            self.my_layers.append(new_step)\n",
    "            \n",
    "        self.my_layers.append([tf.keras.layers.Dense(1)])\n",
    "\n",
    "    def call(self, inputs, training=None): \n",
    "        \n",
    "        # drop prop if there\n",
    "        if 'prop' in inputs:\n",
    "            del inputs['prop']\n",
    "\n",
    "        x = []\n",
    "        for base in self.base_models:\n",
    "            if training:\n",
    "                res = base.call(inputs, training)\n",
    "            else:\n",
    "                res = base.call(inputs)\n",
    "            x.append(res)\n",
    "        x = tf.concat(x, -1)\n",
    "        \n",
    "        for num, layer_step in enumerate(self.my_layers):\n",
    "            for layer in layer_step:\n",
    "                x = layer(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "    def predict_step(self, data):\n",
    "        data = data_adapter.expand_1d(data)\n",
    "        x, _, _ = data_adapter.unpack_x_y_sample_weight(data)\n",
    "\n",
    "        # drop prop here\n",
    "        prop = x['prop']\n",
    "        del x['prop']\n",
    "        return self(x, training=True), data[-1], prop\n",
    "\n",
    "        \n",
    "\n",
    "def build_model(hp):\n",
    "    # returns the compiled TF model\n",
    "    model = MetaModel(hp)\n",
    "    opt = tf.keras.optimizers.Adam(\n",
    "            hp.Choice('learning_rate',\n",
    "                      values=[1e-3]))\n",
    "    opt = tfa.optimizers.SWA(opt)\n",
    "\n",
    "    model.compile(\n",
    "        optimizer=opt,\n",
    "        loss='mse',)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and optimize the meta learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import IPython\n",
    "from sklearn.metrics import mean_squared_error, r2_score\n",
    "results, property_metric = [], []\n",
    "\n",
    "\n",
    "tuner = Hyperband(\n",
    "    build_model,\n",
    "    objective='val_loss',\n",
    "    max_epochs=300,\n",
    "    seed=10,\n",
    "    directory=f'hyperparameter_search_meta_learner',\n",
    "    project_name='fold_0',\n",
    "    )\n",
    "\n",
    "reduce_lr = tfk.callbacks.ReduceLROnPlateau(\n",
    "    factor=0.9,\n",
    "    monitor=\"val_loss\",\n",
    "    verbose=1,\n",
    ")\n",
    "\n",
    "earlystop = tf.keras.callbacks.EarlyStopping(monitor=\"val_loss\", patience=40)\n",
    "\n",
    "class ClearTrainingOutput(tf.keras.callbacks.Callback):\n",
    "    def on_train_end(*args, **kwargs):\n",
    "        IPython.display.clear_output(wait = True)\n",
    "\n",
    "# Create an instance of the model and optimize the hyperspace\n",
    "tuner.search(datasets_final['test'],\n",
    "            epochs=300,\n",
    "            validation_data=datasets_final['train'],\n",
    "            callbacks=[earlystop, reduce_lr, ClearTrainingOutput()],\n",
    "            verbose=0\n",
    "            )\n",
    "\n",
    "# Post processing\n",
    "best_values = tuner.get_best_hyperparameters()[0].values\n",
    "best_model = tuner.get_best_models(1)[0]\n",
    "\n",
    "# predict\n",
    "res = np.concatenate(best_model.predict(datasets_final['train']), -1)\n",
    "\n",
    "# save for deployment\n",
    "best_model.save(f'models/meta_model/', include_optimizer=False)\n",
    "\n",
    "\n",
    "_df = pd.DataFrame(res, columns=['pred', 'target', 'prop'])\n",
    "_df['prop'] = _df.prop.apply(lambda x: x.decode('utf-8'))\n",
    "props = _df.prop.unique()\n",
    "\n",
    "property_scaler = datasets_final['property_scaler']\n",
    "for prop in props:\n",
    "\n",
    "    cond = _df[_df.prop == prop].index\n",
    "    rmse_scaled = mean_squared_error(_df.loc[cond, ['target']], _df.loc[cond, ['pred']], squared=False)\n",
    "    r2_scaled = r2_score(_df.loc[cond, ['target']], _df.loc[cond, ['pred']])\n",
    "    \n",
    "    _df.loc[cond, ['pred']] = property_scaler[prop].inverse_transform(_df.loc[cond, ['pred']])\n",
    "    _df.loc[cond, ['target']] = property_scaler[prop].inverse_transform(_df.loc[cond, ['target']])\n",
    "    \n",
    "    rmse = mean_squared_error(_df.loc[cond, ['target']], _df.loc[cond, ['pred']], squared=False)\n",
    "    r2 = r2_score(_df.loc[cond, ['target']], _df.loc[cond, ['pred']])\n",
    "    property_metric.append({'name': f'ensemble_{what}', 'prop': prop, 'rmse': rmse, 'r2':r2, 'rmse_scaled': rmse_scaled, 'r2_scaled': r2_scaled})\n",
    "    \n",
    "# Not scaled back\n",
    "rmse = mean_squared_error(res[:,0], res[:,1], squared=False)\n",
    "r2 = r2_score(res[:,0], res[:,1])\n",
    "\n",
    "results.append({'name': f'ensemble_{what}','r2': r2, 'rmse':rmse})\n",
    "\n",
    "pd.DataFrame(results)    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.8 64-bit ('.venv': venv)",
   "language": "python",
   "name": "python37864bitvenvvenv0c9e6ba843324ba8b75a3e51791043c5"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "output_auto_scroll": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}