{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# 5. Entity Aware\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import site\nsite.addsitedir(r\"E:\\AA\\AI4Water\")\nsite.addsitedir(r\"E:\\AA\\easy_mpl\")\n\nfrom ai4water.functional import Model\nfrom ai4water.preprocessing import Transformations\nfrom utils import prepare_data, eval_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "lookback = 12\nbatch_size = 64\nunits = 100\nlr = 0.0001"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tr_x, tr_y, val_x, val_y, test_x, test_y = prepare_data(\n    target=\"TP\", lookback=lookback, batch_size=batch_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "num_dyn_inputs = tr_x[0].shape[-1]\nnum_static_inputs = tr_x[1].shape[-1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "transformer = Transformations('box-cox')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "x_mean = tr_x[0].mean()\nx_std = tr_x[0].std()\nx_mean_st = tr_x[1].mean()\nx_std_st = tr_x[1].std()\n\ny_mean = tr_y.mean()\ny_std = tr_y.std()\n\ntr_x[0] = (tr_x[0] - x_mean) / x_std\ntr_x[1] = (tr_x[1] - x_mean_st) / x_std_st\n#tr_y = (tr_y - y_mean) / y_std\ntr_y = transformer.fit_transform(tr_y)\n\nval_x[0] = (val_x[0] - x_mean) / x_std\nval_x[1] = (val_x[1] - x_mean_st) / x_std_st\n#val_y = (val_y - y_mean) / y_std\nval_y = transformer.transform(val_y)\n\ntest_x[0] = (test_x[0] - x_mean) / x_std\ntest_x[1] = (test_x[1] - x_mean_st) / x_std_st\n#test_y = (test_y - y_mean) / y_std\ntest_y = transformer.transform(test_y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "layers = {\n    \"Input_dyn\": {\"batch_shape\": (batch_size, lookback, num_dyn_inputs)},\n    \"Input_cat\": {\"batch_shape\": (batch_size, num_static_inputs)} ,\n    \"EALSTM\": {\"config\": {\"units\": units, \"num_static_inputs\": num_static_inputs},\n               \"inputs\": [\"Input_dyn\", \"Input_cat\"]},\n    \"Dense\": 1\n}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model = Model(model = {\"layers\": layers},\n              batch_size=batch_size,\n              epochs=100,\n              lr=lr\n              )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "h = model.fit(x=tr_x, y=tr_y, validation_data=(val_x, val_y))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "eval_model(model, tr_x, tr_y, batch_size=batch_size, prefix=\"Training\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "eval_model(model, val_x, val_y, batch_size=batch_size, prefix=\"Validation\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "eval_model(model, test_x, test_y, batch_size=batch_size, prefix=\"Test\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model.predict(test_x, test_y, plots=['residual', 'edf', 'regression', 'prediction'\n                                     ])"
      ]
    }
  ],
  "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.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}