RL18_replication.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is intended as a simplified and straightforward to use implementation of the main neural network-based model considered in the paper (the `nn_aux_emb` model). There are minor changes in the setup compared to the original implementation for the paper (detailed below), see the folder `nn_postprocessing` for the original code)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read data (can be downloaded from https://doi.org/10.6084/m9.figshare.13516301.v1)\n",
    "\n",
    "data = pd.read_feather(path = \"data_RL18.feather\")\n",
    "data.station = pd.to_numeric(data.station, downcast = 'integer')\n",
    "\n",
    "# drop soil moisture predictions due to missing values\n",
    "# Note that this is a minor change compared to the paper, but does not have a significant effect\n",
    "data = data.drop(['sm_mean', 'sm_var'], axis=1)\n",
    "\n",
    "# split into train and test data\n",
    "eval_start = 1626724\n",
    "train_end = 1626723\n",
    "\n",
    "train_features_raw = data.iloc[:train_end,3:42].to_numpy()\n",
    "train_targets = data.iloc[:train_end,2].to_numpy()\n",
    "train_IDs = data.iloc[:train_end,1].to_numpy()\n",
    "\n",
    "test_features_raw = data.iloc[eval_start:,3:42].to_numpy()\n",
    "test_targets = data.iloc[eval_start:,2].to_numpy()\n",
    "test_IDs = data.iloc[eval_start:,1].to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalize data\n",
    "\n",
    "def normalize(data, method=None, shift=None, scale=None):\n",
    "    result = np.zeros(data.shape)\n",
    "    if method == \"MAX\":\n",
    "        scale = np.max(data, axis=0)\n",
    "        shift = np.zeros(scale.shape)\n",
    "    for index in range(len(data[0])):\n",
    "        result[:,index] = (data[:,index] - shift[index]) / scale[index]\n",
    "    return result, shift, scale\n",
    "\n",
    "train_features, train_shift, train_scale = normalize(train_features_raw[:,:], method=\"MAX\")\n",
    "\n",
    "test_features = normalize(test_features_raw[:,:], shift=train_shift, scale=train_scale)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper functions for NN models\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow.keras.layers import Input, Dense, Embedding, Flatten, Concatenate\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "\n",
    "\n",
    "def crps_cost_function(y_true, y_pred, theano=False):\n",
    "    \"\"\"Compute the CRPS cost function for a normal distribution defined by\n",
    "    the mean and standard deviation.\n",
    "\n",
    "    Code inspired by Kai Polsterer (HITS).\n",
    "\n",
    "    Args:\n",
    "        y_true: True values\n",
    "        y_pred: Tensor containing predictions: [mean, std]\n",
    "        theano: Set to true if using this with pure theano.\n",
    "\n",
    "    Returns:\n",
    "        mean_crps: Scalar with mean CRPS over batch\n",
    "    \"\"\"\n",
    "\n",
    "    # Split input\n",
    "    mu = y_pred[:, 0]\n",
    "    sigma = y_pred[:, 1]\n",
    "    # Ugly workaround for different tensor allocation in keras and theano\n",
    "    if not theano:\n",
    "        y_true = y_true[:, 0]   # Need to also get rid of axis 1 to match!\n",
    "\n",
    "    # To stop sigma from becoming negative we first have to \n",
    "    # convert it the the variance and then take the square\n",
    "    # root again. \n",
    "    var = sigma ** 2\n",
    "    # The following three variables are just for convenience\n",
    "    loc = (y_true - mu) / tf.sqrt(var)\n",
    "    phi = 1.0 / tf.sqrt(2.0 * np.pi) * tf.exp(-loc ** 2 / 2.0)\n",
    "    Phi = 0.5 * (1.0 + tf.math.erf(loc / tf.sqrt(2.0)))\n",
    "    # First we will compute the crps for each input/target pair\n",
    "    crps =  tf.sqrt(var) * (loc * (2. * Phi - 1.) + 2 * phi - 1. / tf.sqrt(np.pi))\n",
    "    # Then we take the mean. The cost is now a scalar\n",
    "    return tf.math.reduce_mean(crps)\n",
    "\n",
    "def build_emb_model(n_features, n_outputs, hidden_nodes, emb_size, max_id,\n",
    "                    compile=False, lr=0.01,\n",
    "                    loss=crps_cost_function,\n",
    "                    activation='relu', reg=None):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        n_features: Number of features\n",
    "        n_outputs: Number of outputs\n",
    "        hidden_nodes: int or list of hidden nodes\n",
    "        emb_size: Embedding size\n",
    "        max_id: Max embedding ID\n",
    "        compile: If true, compile model\n",
    "        lr: learning rate\n",
    "        loss: loss function\n",
    "        activation: Activation function for hidden layer\n",
    "\n",
    "    Returns:\n",
    "        model: Keras model\n",
    "    \"\"\"\n",
    "    if type(hidden_nodes) is not list:\n",
    "        hidden_nodes = [hidden_nodes]\n",
    "\n",
    "    features_in = Input(shape=(n_features,))\n",
    "    id_in = Input(shape=(1,))\n",
    "    emb = Embedding(max_id + 1, emb_size)(id_in)\n",
    "    emb = Flatten()(emb)\n",
    "    x = Concatenate()([features_in, emb])\n",
    "    for h in hidden_nodes:\n",
    "        x = Dense(h, activation=activation, kernel_regularizer=reg)(x)\n",
    "    x = Dense(n_outputs, activation='linear', kernel_regularizer=reg)(x)\n",
    "    model = Model(inputs=[features_in, id_in], outputs=x)\n",
    "\n",
    "    if compile:\n",
    "        opt = Adam(learning_rate=lr)\n",
    "        model.compile(optimizer=opt, loss=loss)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.backend import clear_session\n",
    "\n",
    "# training multiple models in a loop\n",
    "\n",
    "emb_size = 2\n",
    "max_id = int(tf.math.reduce_max([train_IDs.max(), test_IDs.max()]))\n",
    "n_features = train_features.shape[1]\n",
    "n_outputs = 2\n",
    "\n",
    "nreps = 10\n",
    "trn_scores = []\n",
    "test_scores = []\n",
    "preds = []\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "for i in tqdm(range(nreps)):\n",
    "    clear_session()\n",
    "    \n",
    "    features_in = Input(shape=(n_features,))\n",
    "    id_in = Input(shape=(1,))\n",
    "    emb = Embedding(max_id + 1, emb_size)(id_in)\n",
    "    emb = Flatten()(emb)\n",
    "    x = Concatenate()([features_in, emb])\n",
    "    x = Dense(512, activation='relu')(x)\n",
    "    x = Dense(n_outputs, activation='linear')(x)\n",
    "    nn_aux_emb = Model(inputs=[features_in, id_in], outputs=x)\n",
    "\n",
    "    opt = Adam(learning_rate=0.002)\n",
    "    nn_aux_emb.compile(optimizer=opt, loss=crps_cost_function)\n",
    "    \n",
    "    nn_aux_emb.fit([train_features, train_IDs], train_targets, epochs=15, batch_size=4096, verbose=0)   \n",
    "    \n",
    "    trn_scores.append(nn_aux_emb.evaluate([train_features, train_IDs], train_targets, 4096, verbose=0))\n",
    "    test_scores.append(nn_aux_emb.evaluate([test_features, test_IDs], test_targets, 4096, verbose=0))\n",
    "    preds.append(nn_aux_emb.predict([test_features, test_IDs], 4096, verbose=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "test_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate ensemble of models\n",
    "\n",
    "from scipy.stats import norm\n",
    "\n",
    "def crps_normal(mu, sigma, y):\n",
    "    \"\"\"\n",
    "    Compute CRPS for a Gaussian distribution. \n",
    "    \"\"\"\n",
    "    # Make sure sigma is positive\n",
    "    sigma = np.abs(sigma)\n",
    "    loc = (y - mu) / sigma\n",
    "    crps = sigma * (loc * (2 * norm.cdf(loc) - 1) + \n",
    "                    2 * norm.pdf(loc) - 1. / np.sqrt(np.pi))\n",
    "    return crps\n",
    "\n",
    "preds = np.array(preds)\n",
    "preds[:, :, 1] = np.abs(preds[:, :, 1]) # Make sure std is positive\n",
    "mean_preds = np.mean(preds, 0)\n",
    "ens_score = crps_normal(mean_preds[:, 0], mean_preds[:, 1], test_targets).mean()\n",
    "print(f'Ensemble test score = {ens_score}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}