HiggsML_syst_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "import pandas as pd\n",
    "from pickle import dump\n",
    "\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "from util import train_test_split_evtNo\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_everything(seed=0):\n",
    "    #random.seed(seed)\n",
    "    #os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    #torch.manual_seed(seed)\n",
    "    #torch.cuda.manual_seed(seed)\n",
    "    #torch.cuda.manual_seed_all(seed)\n",
    "    #torch.backends.cudnn.deterministic = True\n",
    "    #tf keras\n",
    "seed_everything(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_syst = np.sort(np.concatenate([np.arange(88,97)/100,np.arange(970,1030)/1000,np.arange(103,107)/100,np.arange(1070,1130)/1000,np.arange(113,115)/100], axis=-1 ))\n",
    "systUp = 1.1 \n",
    "systDown = 0.8 \n",
    "z_nominal = 1.\n",
    "\n",
    "z_syst_train = \\\n",
    "        np.array([0.7, 0.74, 0.78, 0.8, 0.84, 0.88, 0.9 , 0.92, 0.94, 0.96, 0.98, 0.99, 1., 1.01, 1.02, 1.04, 1.06, 1.08, 1.09,\n",
    "              1.10, 1.11, 1.12, 1.13, 1.14])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loadTrainingData(trainingZvalues):\n",
    "    \n",
    "    dataList = []\n",
    "    for z_val in trainingZvalues:\n",
    "        dataList.append(loadTrainingSystData(z_val) )\n",
    "    \n",
    "    concatList = []\n",
    "    for i in range (len(dataList[0])): # loop over X, Y, etc \n",
    "        arr = np.concatenate([tup[i] for tup in dataList],axis=0) # concat the X for all datasets, then Y etc\n",
    "        concatList.append(arr)\n",
    "        \n",
    "    \n",
    "    assert len(concatList) == 8\n",
    "    X_syst_train, Y_syst_train, Z_syst_train, weights_train = \\\n",
    "    shuffle(*concatList[::2])\n",
    "    X_syst_test, Y_syst_test, Z_syst_test, weights_test = \\\n",
    "    shuffle(*concatList[1::2])\n",
    "    \n",
    "     # standard normalise scaler\n",
    "    scaler = StandardScaler()\n",
    "    X_syst_train = scaler.fit_transform(X_syst_train)\n",
    "    X_syst_test = scaler.transform(X_syst_test)\n",
    "    # save the scaler\n",
    "    dump(scaler, open('HiggsModels/scaler.pkl', 'wb'))\n",
    "    \n",
    "    \n",
    "    \n",
    "    return X_syst_train, X_syst_test, Y_syst_train, Y_syst_test, \\\n",
    "    Z_syst_train, Z_syst_test, weights_train, weights_test\n",
    "    \n",
    "\n",
    "\n",
    "# no standard normalisation done in this function\n",
    "def loadTrainingSystData(z_val):\n",
    "    from dataLoc import dataLoc\n",
    "    fileName = dataLoc +\"HiggsML_TES_{}.h5\".format(z_val)\n",
    "    print (\"Loading file: \",fileName)\n",
    "    \n",
    "    data = pd.read_hdf(fileName, \"data_syst\")\n",
    "    data = data.sample(frac=1).reset_index(drop=False) #shuffle the events\n",
    "    target = data[\"Label\"]\n",
    "    weights = data[\"Weight\"]\n",
    "    Z = data[\"Z\"]\n",
    "    evtNo = data[\"index\"]\n",
    "    del data[\"Label\"]\n",
    "    del data[\"Z\"]\n",
    "    del data[\"Weight\"]\n",
    "    del data[\"index\"]\n",
    "\n",
    "\n",
    "    data = data.values\n",
    "    target = target.values\n",
    "    weights = weights.values\n",
    "    Z = Z.values\n",
    "    evtNo = evtNo.values\n",
    "    \n",
    "    assert (Z == Z[0]).all() # files contain only 1 z\n",
    "    \n",
    "        \n",
    "    X_syst_train, X_syst_test, Y_syst_train, Y_syst_test, Z_syst_train, Z_syst_test, weights_train, weights_test  = train_test_split_evtNo(\n",
    "        data, target, Z, weights,evtNo=evtNo,n=3)\n",
    "    \n",
    "    class_weights = weights[target == 0].sum(), weights[target == 1].sum()\n",
    "\n",
    "    class_weights_test = weights_test[Y_syst_test == 0].sum(), weights_test[Y_syst_test == 1].sum(), weights_test[Y_syst_test == 2].sum()\n",
    "    scale_up = 1.\n",
    "    for i in xrange(2):\n",
    "        weights_train[Y_syst_train == i] *= scale_up*max(class_weights)/ class_weights[i]\n",
    "        weights_test[Y_syst_test == i] *= class_weights[i]/class_weights_test[i]\n",
    "        \n",
    "    del data, target, weights, Z\n",
    "    \n",
    "    return X_syst_train, X_syst_test, Y_syst_train, Y_syst_test, Z_syst_train, Z_syst_test, weights_train, weights_test\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.7.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.74.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.78.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.8.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.84.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.88.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.9.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.92.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.94.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.96.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.98.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_0.99.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.0.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.01.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.02.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.04.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.06.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.08.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.09.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.1.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.11.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.12.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.13.h5\n",
      "Loading file:  /data1/users/aishik/systcovariant_data/HiggsML_TES_1.14.h5\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    # Python 2\n",
    "    xrange\n",
    "except NameError:\n",
    "    # Python 3, xrange is now named range\n",
    "    xrange = range\n",
    "X_syst_train, X_syst_test, Y_syst_train, Y_syst_test, \\\n",
    "    Z_syst_train, Z_syst_test, weights_train, weights_test = \\\n",
    "        loadTrainingData(trainingZvalues=z_syst_train)\n",
    "\n",
    "Z_syst_train_scaled = Z_syst_train\n",
    "Z_syst_test_scaled = Z_syst_test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "systUpEvents = np.isclose(Z_syst_test, systUp)\n",
    "systDownEvents = np.isclose(Z_syst_test, systDown)\n",
    "nominalEvents = np.isclose(Z_syst_test, z_nominal)\n",
    "\n",
    "systUpEventsTrain = np.isclose(Z_syst_train, systUp)\n",
    "systDownEventsTrain = np.isclose(Z_syst_train, systDown)\n",
    "nominalEventsTrain = np.isclose(Z_syst_train, z_nominal)\n",
    "\n",
    "systUpHalfEventsTrain = Z_syst_train >= 1\n",
    "systDownHalfEventsTrain = Z_syst_train <= 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/shared_env/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:523: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/opt/anaconda3/envs/shared_env/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:524: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/opt/anaconda3/envs/shared_env/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/opt/anaconda3/envs/shared_env/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:526: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/opt/anaconda3/envs/shared_env/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:527: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/opt/anaconda3/envs/shared_env/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:532: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "onlyCPU = False\n",
    "\n",
    "\n",
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "if onlyCPU:\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n",
    "else:\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"3\"\n",
    "gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.5)\n",
    "sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "from tensorflow.keras.layers import Input, Dense, Concatenate, Dropout\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.callbacks import EarlyStopping\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.layers import LeakyReLU\n",
    "from tensorflow.python.keras.layers import Lambda"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on Nominal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_hidden_clf = 10\n",
    "n_nodes_clf = 512\n",
    "\n",
    "inputs = Input(shape=(X_syst_train.shape[1],))\n",
    "hidden = Dense(n_nodes_clf, activation='relu')(inputs)\n",
    "\n",
    "for i in range(n_hidden_clf -1):\n",
    "    hidden = Dense(n_nodes_clf, activation='relu', kernel_regularizer='l2')(hidden)\n",
    "predictions = Dense(1, activation='sigmoid')(hidden)\n",
    "\n",
    "nominal_model = Model(inputs=inputs, outputs=predictions)\n",
    "#nominal_model.compile(optimizer='adam',\n",
    "nominal_model.compile(optimizer='RMSProp',\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "nominal_model.load_weights(\"HiggsModels/clf.h5\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "nominal_model.fit(x=X_syst_train[nominalEventsTrain], y=Y_syst_train[nominalEventsTrain], \n",
    "                  sample_weight= weights_train[nominalEventsTrain], epochs=1000, batch_size=4*1024)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "nominal_model.save(\"HiggsModels/clf.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on SystDown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_hidden_clf = 10\n",
    "n_nodes_clf = 64\n",
    "\n",
    "inputs = Input(shape=(X_syst_train.shape[1],))\n",
    "hidden = Dense(n_nodes_clf, activation='relu')(inputs)\n",
    "\n",
    "for i in range(n_hidden_clf -1):\n",
    "    hidden = Dense(n_nodes_clf, activation='relu', kernel_regularizer='l2')(hidden)\n",
    "predictions = Dense(1, activation='sigmoid')(hidden)\n",
    "\n",
    "down_model = Model(inputs=inputs, outputs=predictions)\n",
    "#nominal_model.compile(optimizer='adam',\n",
    "down_model.compile(optimizer='RMSProp',\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "down_model = keras.load_model(\"HiggsModels/down_model.h5\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "down_model.fit(x=X_syst_train[systDownEventsTrain], y=Y_syst_train[systDownEventsTrain], \n",
    "                  sample_weight= weights_train[systDownEventsTrain], epochs=200, batch_size=4*1024)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "down_model.save(\"HiggsModels/down_model.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train on SystUp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_hidden_clf = 10\n",
    "n_nodes_clf = 512\n",
    "\n",
    "inputs = Input(shape=(X_syst_train.shape[1],))\n",
    "hidden = Dense(n_nodes_clf, activation='relu')(inputs)\n",
    "\n",
    "for i in range(n_hidden_clf -1):\n",
    "    hidden = Dense(n_nodes_clf, activation='relu', kernel_regularizer='l2')(hidden)\n",
    "predictions = Dense(1, activation='sigmoid')(hidden)\n",
    "\n",
    "up_model = Model(inputs=inputs, outputs=predictions)\n",
    "up_model.compile(optimizer='RMSProp',\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "up_model = keras.load_model(\"HiggsModels/up_model.h5\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "up_model.fit(x=X_syst_train[systUpEventsTrain], y=Y_syst_train[systUpEventsTrain], \n",
    "             sample_weight= weights_train[systUpEventsTrain], epochs=200, batch_size=4*1024)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "up_model.save(\"HiggsModels/up_model.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Data Augmented Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "#aug_NN = False\n",
    "aug_NN = True\n",
    "if (aug_NN):\n",
    "    from tensorflow import keras\n",
    "    from tensorflow.keras.layers import Input, Dense, Concatenate\n",
    "    from tensorflow.keras.models import Model\n",
    "\n",
    "    inputs = Input(shape=(X_syst_train.shape[1],))\n",
    "    n_hidden_aug = 10\n",
    "    #n_nodes_aug = 512\n",
    "    n_nodes_aug = 64\n",
    "\n",
    "    inputs = Input(shape=(X_syst_train.shape[1],))\n",
    "    hidden = Dense(n_nodes_aug, activation='relu')(inputs)\n",
    "\n",
    "    for i in range(n_hidden_aug -1):\n",
    "        hidden = Dense(n_nodes_aug, activation='relu', kernel_regularizer='l2')(hidden)\n",
    "    predictions = Dense(1, activation='sigmoid')(hidden)\n",
    "\n",
    "    augmented_model = Model(inputs=inputs, outputs=predictions)\n",
    "    augmented_model.compile(optimizer='adam',\n",
    "                  loss='binary_crossentropy',\n",
    "                  metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "augmented_model = keras.load_model(\"HiggsModels/aug.h5\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#train model\n",
    "\n",
    "\n",
    "augmented_model.fit(x=X_syst_train, y=Y_syst_train, \n",
    "                  sample_weight= weights_train, epochs=10, batch_size=4*1024)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "augmented_model.save(\"HiggsModels/aug.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Syst Aware Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_hidden_awe = 10\n",
    "n_nodes_awe = 64\n",
    "\n",
    "inputs_x = Input(shape=(X_syst_train.shape[1],))\n",
    "inputs_z = Input(shape=(1,))\n",
    "inputs = Concatenate(axis=-1)([inputs_x, inputs_z])\n",
    "\n",
    "netAweUp = Dense(n_nodes_awe, activation=\"relu\")(inputs)\n",
    "for i in range(n_hidden_awe -1):\n",
    "    netAweUp = Dense(n_nodes_awe, activation='relu', kernel_regularizer='l2')(netAweUp)\n",
    "predictions_netAweUp = Dense(1, activation='sigmoid')(netAweUp)\n",
    "\n",
    "netAweUp_model = Model(inputs=[inputs_x, inputs_z], outputs=predictions_netAweUp)\n",
    "netAweUp_model.compile(optimizer='RMSProp',\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "netAweDown = Dense(n_nodes_awe, activation=\"relu\")(inputs)\n",
    "for i in range(n_hidden_awe -1):\n",
    "    netAweDown = Dense(n_nodes_awe, activation='relu', kernel_regularizer='l2')(netAweDown)\n",
    "predictions_netAweDown = Dense(1, activation='sigmoid')(netAweDown)\n",
    "\n",
    "netAweDown_model = Model(inputs=[inputs_x, inputs_z], outputs=predictions_netAweDown)\n",
    "netAweDown_model.compile(optimizer='RMSProp',\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Input, Lambda, Multiply, Add\n",
    "\n",
    "def combineModels(ModelUp, ModelDown):\n",
    "    input1 = Input(shape=(X_syst_train.shape[1],))\n",
    "    input2       = Input(shape=(1,))\n",
    "\n",
    "    selectModel1 = Lambda(lambda x: K.greater_equal(x, K.constant(1.)))(input2)\n",
    "    selectModel2 = Lambda(lambda x: K.less(x, K.constant(1.)))(input2)\n",
    "\n",
    "    selectModel1 = Lambda(lambda x: K.cast(x, dtype='float32'))(selectModel1)\n",
    "    selectModel2 = Lambda(lambda x: K.cast(x, dtype='float32'))(selectModel2)\n",
    "\n",
    "    out1 = Multiply()([ModelUp([input1,input2]), selectModel1])\n",
    "    out2 = Multiply()([ModelDown([input1,input2]), selectModel2])\n",
    "    \n",
    "    out  = Add()([out1, out2])\n",
    "    model = Model(inputs=[input1,input2], outputs=out)\n",
    "    return model\n",
    "\n",
    "aware_model = combineModels(netAweUp_model, netAweDown_model)\n",
    "aware_model.compile(optimizer='RMSProp',\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "#aware_model = keras.models.load_model(\"HiggsModels/AwareModel.h5\")\n",
    "aware_model.load_weights(\"HiggsModels/AwareModel.h5\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#train model\n",
    "epochs_full = 60\n",
    "\n",
    "aware_model.fit(x=[X_syst_train, Z_syst_train_scaled], \n",
    "                y=Y_syst_train, sample_weight=weights_train, \n",
    "                epochs=epochs_full, batch_size=4*1024)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "aware_model.save(\"HiggsModels/AwareModel.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Pivot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_hidden_inv = 10; n_hidden_inv_R = 10\n",
    "n_nodes_inv = 64; n_nodes_inv_R = 64\n",
    "hp_lambda = 1\n",
    "\n",
    "from gradRevLayer import GradReverseTR as GradReverse\n",
    "\n",
    "inputs = Input(shape=(X_syst_train.shape[1],))\n",
    "\n",
    "Dx = Dense(n_nodes_inv, activation=\"relu\")(inputs)\n",
    "for i in range(n_hidden_inv -1):\n",
    "    Dx = Dense(n_nodes_inv, activation='relu', kernel_regularizer='l2')(Dx)\n",
    "Dx = Dense(1, activation=\"sigmoid\", name=\"Clf\")(Dx)\n",
    "inv_model = Model(inputs=inputs, outputs=Dx)\n",
    "\n",
    "GRx = GradReverse()(Dx)\n",
    "Rx = Dense(n_nodes_inv_R, activation=\"relu\")(GRx)\n",
    "for i in range(n_hidden_inv_R -1):\n",
    "    Rx = Dense(n_nodes_inv_R, activation='relu', kernel_regularizer='l2')(Rx)\n",
    "#Rx = Dense(1, activation=\"sigmoid\")(Rx)\n",
    "Rx = Dense(1, activation=\"linear\", name=\"Adv\")(Rx)\n",
    "GR = Model(inputs=inputs, outputs=[Dx, Rx])\n",
    "\n",
    "GR.compile(loss=[\"binary_crossentropy\", \"mean_squared_error\"], loss_weights=[1,hp_lambda], optimizer=\"RMSProp\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from tensorflow import keras\n",
    "inv_model = keras.models.load_model('invariant_model.h5')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "GR.fit(x=X_syst_train, y=[Y_syst_train, Z_syst_train], sample_weight=[weights_train, weights_train], epochs=100, batch_size=4*1024)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "inv_model.save('HiggsModels/invariant_model.h5')\n",
    "GR.save('HiggsModels/GR_model.h5')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.5"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}