deep_bucket_lab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Bucket Lab: \n",
    "## A Playground for Understanding Hydrologic Process Representations with Deep Learning\n",
    "\n",
    "Welcome to the Deep Bucket Lab, where the complexities of hydrologic systems are unraveled through the lens of deep learning. This notebook is an educational journey designed to help you grasp the fundamentals of simulating dynamic hydrological systems using advanced machine learning techniques.\n",
    "\n",
    "Our tutorial models a 'leaking bucket' – a simplified yet effective representation of a hydrological system. This approachable model allows us to evaluate and understand the dynamics of water flow in a controlled, yet insightful manner.\n",
    "\n",
    "Here, we generate all data synthetically, using numerical simulations to mimic various hydrological scenarios. You'll witness how different bucket characteristics and environmental variables, like precipitation, influence the system's behavior.\n",
    "\n",
    "We employ a single-layer Long Short-Term Memory (LSTM) network, a type of deep learning model, to learn and predict the dynamics of water levels and fluxes in our simulated buckets, based on the 'ground truth' data we generate.\n",
    "\n",
    "This notebook is your playground to experiment and explore. Modify the bucket model attributes, tweak the modeling setup, and delve into the fascinating world of hydrologic process representations and their predictability. Through interactive graphs and hands-on experiments, you'll gain a deeper understanding of hydrologic principles and the power of deep learning in simulating them.\n",
    "\n",
    "Whether you're a beginner in hydrology or looking to expand your knowledge, this lab is designed to offer valuable insights in an engaging and accessible manner. By the end of this tutorial, you'll have a firmer grasp of hydrologic simulations, deep learning applications in hydrology, and the confidence to apply these concepts in real-world scenarios. Let's dive in!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# 1. Setup\n",
    "\n",
    "The first thing to do is set up the notebook environment with all the necessary libraries, declare model global parameters, settings, variables, and functions that define the bucket system we want to represent. We also define the hyperparameters and structure for the deep learning model.\n",
    "\n",
    "Note: In a typical full-scale modeling framework, these would be declared in a configuration file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Importing Libraries\n",
    "\n",
    "We import standard libraries for data management, calculations, and plotting. Each library plays a specific role in our modeling process: `numpy` for efficient numerical computations, `pandas` for data handling, `matplotlib` for visualizing our results, and `IPython.display` for interactive outputs in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing Machine Learning Libraries\n",
    "\n",
    "In this section, we import machine learning libraries that will be used throughout this notebook. These libraries provide tools for building our model architecture, as well as for training, validating, and testing:\n",
    "\n",
    "- `torch` and its submodules (`nn`, `optim`, etc.) from PyTorch, a leading deep learning library, for defining and training neural network models.\n",
    "- `torch.utils.data` for efficient data handling and batching.\n",
    "- `torch.autograd.Variable` for automatic differentiation: compute the derivatives of functions whose exact forms are known. Applies the **chain rule** to break down expressions into simpler ones.\n",
    "- `sklearn.preprocessing.StandardScaler` from Scikit-learn, for scaling our data to something consistent for the learning model to use.\n",
    "- `tqdm.notebook` for displaying progress bars during training and validation.\n",
    "\n",
    "These libraries will be used in the development and evaluation of our deep learning model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data_utils\n",
    "from torch.autograd import Variable \n",
    "import sklearn\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from tqdm.notebook import trange, tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Defining the Physical Bucket Model System\n",
    "\n",
    "In this section, we establish the elements of our physical bucket model system.\n",
    "\n",
    "**Global Variables:**\n",
    "\n",
    "- `g`: Gravitational acceleration, denoted in meters per second squared \\([m/s^2]\\).\n",
    "- `time step`: The time increment for each step in our simulation, measured in seconds \\([s]\\)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = 9.807\n",
    "time_step = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Parameters for the Forcing Process (Precipitation)**\n",
    "\n",
    "In our synthetic system, precipitation is converted to a hydrologic response. To simulate this, we use a stochastic (random) process that generates precipitation data based on probabilistic models. This process not only determines the occurrence of rain but also quantifies the total amount of rainfall during each event.\n",
    "\n",
    "We categorize precipitation for likelihood (`rain_probability_range`) and intensity (`rain_depth_range`):\n",
    "- **None**: The probability range for this type is set to ensure that dry periods are realistically simulated.\n",
    "- **Light**: More frequent events with lower intensity.\n",
    "- **Heavy**: Less frequent with higher intensity. Driving \"overtopping\" events.\n",
    "\n",
    "These parameters allow us to create a diverse set of precipitation scenarios, each affecting the hydrological system in different ways.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rain_probability_range = {\"None\": [0.6, 0.7], \n",
    "                          \"Light\": [0.5, 0.8], \n",
    "                          \"Heavy\": [0.2, 0.3]}\n",
    "\n",
    "rain_depth_range = {\"Light\": [0, 2], \"Heavy\": [2, 8]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Bucket Physical Attributes**\n",
    "\n",
    "To simulate the dynamics of our leaky bucket system, we define a set of physical attributes that characterize each bucket. These attributes determine how the bucket interacts with the hydrological inputs (like precipitation) and influences the overall water flow dynamics. The attributes include:\n",
    "\n",
    "- `A_bucket`: The cross-sectional area of the bucket, measured in square meters \\([m^2]\\). This determines the bucket's capacity to hold water.\n",
    "- `A_spigot`: The area of the spigot, also in square meters \\([m^2]\\). The spigot area affects flow out of the bucket.\n",
    "- `H_bucket`: The height of the bucket, measured in meters \\([m]\\). This defines the maximum water level the bucket can sustain.\n",
    "- `H_spigot`: The height of the spigot from the bucket's base, in meters \\([m]\\). It influences how water level changes affect the outflow.\n",
    "- `K_infiltration`: The infiltration rate, indicating how quickly water seeps through the bucket, measured in millimeters per hour \\([mm/hr]\\).\n",
    "- `ET_parameter`: The evapotranspiration parameter, representing the rate of water loss due to evaporation and transpiration, measured in millimeters per day \\([mm/day]\\).\n",
    "\n",
    "To introduce variability and simulate different hydrological scenarios, we will generate a range of leaky buckets by randomly selecting values for these physical attributes from predefined ranges. This approach allows us to explore a wide spectrum of bucket behaviors and their responses to various hydrological conditions.\n",
    "\n",
    "The following code block defines the possible ranges for each attribute, ensuring a diverse set of buckets for our simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_attributes_range = {\"A_bucket\": [1.0, 2.0],\n",
    "                           \"A_spigot\": [0.1, 0.2],\n",
    "                           \"H_bucket\": [5.0, 6.0],\n",
    "                           \"H_spigot\": [1.0, 3.0],\n",
    "                           \"K_infiltration\": [1e-7, 1e-9],\n",
    "                           \"ET_parameter\": [7, 9]\n",
    "                          }\n",
    "\n",
    "bucket_attributes_list = list(bucket_attributes_range.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Model Input and Output Variables**\n",
    "\n",
    "The following variables represent fluxes and the state of the system:\n",
    "\n",
    "**Inputs:**\n",
    "- `precip`: Precipitation into the bucket. This variable represents the external hydrological input driving the system.\n",
    "- `et` (Evapotranspiration): The actual loss of water from the bucket due to evaporation and transpiration.\n",
    "- `pet` (Potential Evapotranspiration): The potential loss to evaporation, providing an upper limit to the `et` variable.\n",
    "\n",
    "**Outputs:**\n",
    "- `q_overflow`: The water flow over the bucket's rim. This occurs when the water level exceeds the bucket's capacity.\n",
    "- `q_spigot`: The water flow out of the bucket's spigot. It's a controlled outflow dependent on the water level and spigot characteristics.\n",
    "- `h_bucket` (Water Head): The **state** of the water level in the bucket.\n",
    "\n",
    "These variables represent the processes and states we aim to simulate and understand. The following code block initializes these lists and calculates the number of input and output variables, setting the stage for the subsequent modeling steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_vars = ['precip', 'et']\n",
    "input_vars.extend(bucket_attributes_list)\n",
    "output_vars = ['q_overflow', 'q_spigot']\n",
    "n_input = len(input_vars)\n",
    "n_output = len(output_vars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Data Noise**\n",
    "\n",
    "In real-world hydrological systems, data often contains noise due to various environmental and measurement uncertainties. To make our synthetic data more representative of actual scenarios, we introduce noise into the system. We implement this by muadjusting the synthetic fluxes and state values by a random factor drawn from a normal distribution. This method simulates **some** of the random fluctuations and uncertainties inherent in real hydrological data, but not all types of biases and errors in the data are accounted for using this method.\n",
    "\n",
    "The following code block sets up the noise parameters for our model. The `noise` dictionary specifies the standard deviation of the noise for each variable (`pet`, `et`, `q`, and `head`). By adjusting these parameters, we can simulate different levels of data quality and observe how the model responds to these changes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise = {\"pet\": 0.1, \"et\": 0.1, \"q\": 0.1, \"head\": 0.1} "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Illustration of the Leaking Bucket Hydrological System**\n",
    "\n",
    "Below is a schematic representation of the leaking bucket system used in our simulations. This diagram provides a visual understanding of the various components and dynamics of the system, including the bucket's structure, the spigot, and the flow of water.\n",
    "\n",
    "<img src=\"figs/bucket_schematic.png\" alt=\"Leaking Bucket System Schematic\" width=\"250\"/>\n",
    "\n",
    "\n",
    "*Figure: Schematic of the Leaking Bucket Hydrological System. The diagram illustrates the key elements of the system, such as the bucket area, spigot, water inflow (precipitation), and outflows (spigot flow and overflow). Understanding this setup is crucial for comprehending how the system responds to different hydrological inputs and conditions.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Defining the Modeling Setup\n",
    "\n",
    "**Deep Learning Model (LSTM) Hyperparameters**\n",
    "\n",
    "In this section, we define the hyperparameters for our LSTM (Long Short-Term Memory) deep learning model. These hyperparameters influence the model's learning process, its ability to capture complex patterns in the data, and its overall performance. Here's a brief overview of each hyperparameter:\n",
    "\n",
    "- `device`: Specifies the computing device for training and inference, such as CPU or CUDA (GPU). Utilizing CUDA can significantly speed up the training process if a compatible GPU is available.\n",
    "- `hidden_state_size`: The number of units in the hidden state of the LSTM. This determines the capacity of the model to learn and remember information over time.\n",
    "- `num_layers`: The number of LSTM layers in the model. More layers can increase the model's complexity and ability to learn more intricate patterns.\n",
    "- `num_epochs`: The number of complete passes through the training dataset. This affects how long the model will learn from the data.\n",
    "- `batch_size`: The number of samples processed before the model's internal parameters are updated. It's a balance between training speed and memory usage.\n",
    "- `seq_length`: The length of the input sequences. This defines how much past information the model considers for making predictions.\n",
    "- `learning_rate`: The step size at each iteration while moving toward a minimum of the loss function. It's crucial for convergence of the model during training.\n",
    "- `num_classes`: The number of output classes or predictions the model makes.\n",
    "- `input_size`: The size of the input layer, which should match the number of features in the dataset.\n",
    "\n",
    "Adjusting these hyperparameters **can** optimize the LSTM model's accuracy and generalization capabilities for our specific task and dataset. The following code block sets these parameters, preparing our model for training and evaluation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\")\n",
    "    print(\"Using CUDA device: \", torch.cuda.get_device_name(0))\n",
    "else:\n",
    "    device = 'cpu'\n",
    "hidden_state_size = 16\n",
    "num_layers = 8\n",
    "num_epochs = 5  \n",
    "batch_size = 256 \n",
    "seq_length = 24\n",
    "learning_rate = np.linspace(start=0.1, stop=0.01, num=num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Size of Data Records and Splits**\n",
    "\n",
    "In deep learning, particularly for time-series data like our hydrological simulations, it's crucial to split the data into distinct sets for different stages of model development and evaluation. These splits help in effectively training the model, tuning its hyperparameters, and assessing its performance. We divide our data into three sets:\n",
    "\n",
    "- **Training Set (`train`)**: This is the primary dataset used for training the model. Backpropagation, the process of adjusting the model's weights based on the error in its predictions, is performed on this set.\n",
    "- **Validation Set (`val`)**: This dataset is used for hyperparameter tuning. It helps in optimizing the model without using the test set, thus preventing information leakage and overfitting.\n",
    "- **Testing Set (`test`)**: The final evaluation of the model's accuracy is done on this dataset. Note that this data is **never** used in the training process to ensure an unbiased assessment of the model's performance.\n",
    "\n",
    "We will define the number of randomly generated bucket configurations for each of these sets (`n_buckets_split`) as well as the length of the simulations (number of time steps) for each set (`time_splits`).\n",
    "\n",
    "For instance, if we have 10 different buckets in the training set, each with 1000 time steps, we will train the model on these 10 bucket configurations, each simulated over 1000 time steps. This approach allows us to capture a diverse range of system behaviors and responses, ensuring that our model is robust and generalizable.\n",
    "\n",
    "The following code block sets the number of buckets and time steps for each data split, laying the groundwork for our model training and evaluation process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_buckets_split = {\"train\": 10, \"val\": 5,\"test\": 1}\n",
    "time_splits = {\"train\": 1000, \"val\": 500,\"test\": 500}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to manage our data for training, validation, and testing, we calculate the total length of the data record we need to generate, as well as the total number of bucket configurations we will use across all splits. These calculations are based on the settings we defined earlier for the time steps and bucket splits.\n",
    "\n",
    "- `num_records`: This variable represents the total length of the data record. It is calculated by summing the number of time steps allocated for the training, validation, and testing sets, along with additional steps to account for the sequence length in each phase. This ensures we have a comprehensive dataset that covers all aspects of our model's training and evaluation process.\n",
    "\n",
    "- `n_buckets`: This variable indicates the total number of unique bucket configurations we will use. It is the sum of the number of buckets allocated for training, validation, and testing. This diversity in bucket configurations is key to ensuring that our model is trained and tested on a wide range of scenarios, enhancing its robustness and generalizability.\n",
    "\n",
    "<img src=\"./figs/bucket_splits.png\" alt=\"Sample Image\" width=\"200\"/>\n",
    "\n",
    "The following code block performs these calculations, setting the stage for the subsequent data generation and modeling steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_records = time_splits[\"train\"] + time_splits[\"val\"] + time_splits[\"test\"] + seq_length * 3\n",
    "n_buckets = n_buckets_split[\"train\"] + n_buckets_split[\"val\"] + n_buckets_split[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To effectively train, validate, and test our deep learning model, we use separate datasets for each phase. This process involves calculating specific parameters for bucket configurations and time steps that correspond to our training, validation, and testing sets. These parameters will ensure that each dataset is distinct and suitable for its respective purpose in the model development process.\n",
    "\n",
    "- **Bucket Splits**: We calculate the indices of buckets to be used for each set (training, validation, testing) based on the predefined number of buckets for each phase. This ensures a diverse range of bucket configurations for the model to learn from and be evaluated against.\n",
    "\n",
    "- **Time Splits**: We determine the time range for each dataset, ensuring that the model is trained, validated, and tested on different segments of the data. This segmentation is crucial for assessing the model's ability to generalize and perform accurately on unseen data.\n",
    "\n",
    "The following code block defines a function `split_parameters` that performs these calculations. It creates lists of bucket indices for each set and determines the time range for each set based on the given time splits. The function then organizes these split parameters into separate lists for the training, validation, and testing sets, facilitating the subsequent data preparation and modeling steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_parameters():\n",
    "    # create lists of bucket indices for each set based on the given bucket splits\n",
    "    buckets_for_training = list(range(0, n_buckets_split['train'] + 1))\n",
    "    buckets_for_val = list(range(n_buckets_split['train'] + 1, \n",
    "                                 n_buckets_split['train'] + n_buckets_split['val'] + 1))\n",
    "    buckets_for_test = list(range(n_buckets - n_buckets_split['test'], n_buckets))\n",
    "\n",
    "    # determine the time range for each set based on the given time splits\n",
    "    train_start = seq_length\n",
    "    train_end   = time_splits[\"train\"]\n",
    "    val_start   = train_end + seq_length\n",
    "    val_end     = val_start + time_splits[\"val\"]\n",
    "    test_start  = val_end + seq_length\n",
    "    test_end    = test_start + time_splits[\"test\"]\n",
    "    \n",
    "    # organize the split parameters into separate lists for each set\n",
    "    train_split_parameters = [buckets_for_training, train_start, train_end]\n",
    "    val_split_parameters = [buckets_for_val, val_start, val_end]\n",
    "    test_split_parameters = [buckets_for_test, test_start, test_end]\n",
    "    \n",
    "    return [train_split_parameters, val_split_parameters, test_split_parameters]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have defined our function to calculate the split parameters, it's time to run it. Executing the `split_parameters` function will allocate the specific bucket indices and time ranges for our training, validation, and testing sets. This step organizes our data into structured segments, each tailored for a specific phase of the model's development and evaluation process.\n",
    "\n",
    "By running this function, we'll obtain the following:\n",
    "\n",
    "- A list of bucket indices for the training set (`buckets_for_training`) and the corresponding start and end time steps (`train_start`, `train_end`).\n",
    "- A similar list and time steps for the validation set (`buckets_for_val`, `val_start`, `val_end`).\n",
    "- And the same for the testing set (`buckets_for_test`, `test_start`, `test_end`).\n",
    "\n",
    "This organization ensures that each part of our dataset is used effectively in the model's training and evaluation, contributing to a more robust and accurate learning process.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[[buckets_for_training, train_start, train_end],\n",
    "[buckets_for_val, val_start, val_end],\n",
    "[buckets_for_test, test_start, test_end]]= split_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Creating a Sample of Diverse Buckets\n",
    "\n",
    "With the parameters for our bucket and time splits established, our next step is to generate a diverse sample of buckets for our simulations. This diversity is key to creating a robust model that can generalize well across various hydrological scenarios.\n",
    "\n",
    "We achieve this by randomly sampling from the possible ranges for each bucket attribute that we defined earlier. This process involves creating buckets with varying physical characteristics such as area, height, spigot size, and infiltration rates. By doing so, we ensure that our model encounters and learns from a wide array of system behaviors and responses.\n",
    "\n",
    "The following code block will generate this sample of buckets, each with its unique set of boundary and initial conditions. This step builds a comprehensive dataset that reflects the variability and complexity of real-world hydrological systems, thereby enhancing the learning and predictive capabilities of our LSTM model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_buckets(n_buckets):\n",
    "    # Boundary conditions\n",
    "    buckets = {bucket_attribute:[] for bucket_attribute in bucket_attributes_list}\n",
    "    for i in range(n_buckets):\n",
    "        for attribute in bucket_attributes_list:\n",
    "            buckets[attribute].append(np.random.uniform(bucket_attributes_range[attribute][0], \n",
    "                                                        bucket_attributes_range[attribute][1]))\n",
    "\n",
    "    # Initial conditions\n",
    "    h_water_level = [np.random.random() for i in range(n_buckets)]\n",
    "    mass_overflow = [np.random.random() for i in range(n_buckets)]\n",
    "    return buckets, h_water_level, mass_overflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we proceed to run the `setup_buckets` function. This function is responsible for initializing our bucket configurations based on the number of buckets we've determined to use in our simulations. It sets up the initial conditions and physical attributes for each bucket in our model.\n",
    "\n",
    "By executing `setup_buckets`, we will achieve the following:\n",
    "\n",
    "- `buckets`: This variable will hold the array of bucket configurations, each with its unique set of physical attributes as defined earlier.\n",
    "- `h_water_level`: This represents the initial water level in each bucket.\n",
    "- `mass_overflow`: This variable tracks the overflow mass from each bucket.\n",
    "\n",
    "The setup of these buckets establishes the diverse conditions under which our LSTM model will learn and make predictions. This diversity is key to ensuring that our model is robust and can generalize well across different hydrological systems.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "buckets, h_water_level, mass_overflow = setup_buckets(n_buckets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.5 Creating the Synthetic \"Precipitation\"\n",
    "\n",
    "Synthetic precipitation data serves as the main input (forcing data) for each bucket model. This step involves creating time series data that mimic real-world rainfall patterns, influencing how our model simulates water flow and dynamics within each bucket.\n",
    "\n",
    "The process of generating this synthetic precipitation data includes two main steps:\n",
    "\n",
    "1. **Assigning Rainfall Parameters**: For each bucket, we randomly assign rainfall parameters based on the probability ranges specified for each precipitation type (None, Light, and Heavy). These parameters determine the likelihood and intensity of rainfall events, thereby influencing the hydrological behavior of each bucket.\n",
    "\n",
    "2. **Generating Forcing Data**: Using the assigned rainfall parameters, we then generate the synthetic input time series for each bucket. This involves a stochastic process that simulates realistic precipitation patterns over time, based on the probability and depth ranges of each precipitation type.\n",
    "\n",
    "The function `pick_rain_params` defined in the following code block is responsible for randomly selecting these rainfall parameters for each bucket. This function ensures that each bucket experiences a unique set of precipitation conditions, adding to the diversity and realism of our simulations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pick_rain_params():\n",
    "    buck_rain_params = [rain_depth_range,\n",
    "                        np.random.uniform(rain_probability_range[\"None\"][0],\n",
    "                                            rain_probability_range[\"None\"][1]),\n",
    "                        np.random.uniform(rain_probability_range[\"Heavy\"][0],\n",
    "                                            rain_probability_range[\"Heavy\"][1]),\n",
    "                        np.random.uniform(rain_probability_range[\"Light\"][0],\n",
    "                                            rain_probability_range[\"Light\"][1])\n",
    "                 ]\n",
    "    return buck_rain_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Random Rain\n",
    "\n",
    "The next step in our simulation is to generate random rain events. This process is key to creating realistic and varied hydrological conditions for each bucket. The random rain function simulates precipitation events based on probabilities and patterns that mimic real-world weather conditions.\n",
    "\n",
    "The `random_rain` function defined in the following code block takes into account:\n",
    "\n",
    "- The probability of having no rain, light rain, or heavy rain, based on the parameters assigned to each bucket.\n",
    "- The intensity of the rain, which is determined by a random process within specified depth ranges for light and heavy rain.\n",
    "- The influence of preceding rainfall conditions. For instance, the likelihood of a certain type of rain (light or heavy) can depend on the rainfall of the previous day.\n",
    "\n",
    "This stochastic approach to generating rain ensures that each bucket experiences a dynamic and varied set of precipitation events. Such variability is essential for testing and training our model to handle a wide range of hydrological scenarios, enhancing its robustness and predictive accuracy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_rain(preceding_rain, bucket_rain_params):\n",
    "    depth_range, no_rain_probability, light_rain_probability, heavy_rain_probability = bucket_rain_params\n",
    "    # some percent of time we have no rain at all\n",
    "    if np.random.uniform(0.01, 0.99) < no_rain_probability:\n",
    "        rain = 0\n",
    "\n",
    "    # When we do have rain, the probability of heavy or light rain depends on the previous day's rainfall\n",
    "    else:\n",
    "        # If yesterday was a light rainy day, or no rain, then we are likely to have light rain today\n",
    "        if preceding_rain < depth_range[\"Light\"][1]:\n",
    "            if np.random.uniform(0, 1) < light_rain_probability:\n",
    "                rain = np.random.uniform(0, 1)\n",
    "            else:\n",
    "                # But if we do have heavy rain, then it could be very heavy\n",
    "                rain = np.random.uniform(depth_range[\"Heavy\"][0], depth_range[\"Heavy\"][1])\n",
    "\n",
    "        # If it was heavy rain yesterday, then we might have heavy rain again today\n",
    "        else:\n",
    "            if np.random.uniform(0, 1) < heavy_rain_probability:\n",
    "                rain = np.random.uniform(0, 1)\n",
    "            else:\n",
    "                rain = np.random.uniform(depth_range[\"Light\"][0], depth_range[\"Light\"][1])\n",
    "    return rain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have established the mechanism for generating random rain events, our next step is to create a rainfall input time series for each bucket in our simulation. This time series will drive the hydrological processes within each bucket.\n",
    "\n",
    "We accomplish this by iterating over each bucket and generating a sequence of rainfall events using the `random_rain` function. The parameters for each bucket's rainfall events are determined by the `pick_rain_params` function, ensuring that each bucket experiences a unique pattern of precipitation.\n",
    "\n",
    "The generated rainfall data for each bucket is stored in a dictionary, `in_list`, where each key corresponds to a bucket, and the value is the list of rainfall amounts over time. This structure allows us to efficiently manage and access the rainfall data for each bucket, facilitating the subsequent steps in our simulation and modeling process.\n",
    "\n",
    "The following code block iterates through all the buckets, generating and storing the random rainfall time series for each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_list = {}\n",
    "for ibuc in range(n_buckets):\n",
    "    bucket_rain_params = pick_rain_params()\n",
    "    in_list[ibuc] = [0]\n",
    "    for i in range(1, num_records):\n",
    "        in_list[ibuc].append(random_rain(in_list[ibuc][i-1], bucket_rain_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.6 Running Numerical Simulations of the Bucket Model to Generate \"Ground Truth\" Data\n",
    "\n",
    "To develop a reliable and accurate deep learning model, we need \"ground truth\" data. In our case, the \"ground truth\" data is generated by running numerical simulations of the bucket model.\n",
    "\n",
    "For each bucket configuration, the simulation iterates over each time step, updating the water level based on factors such as precipitation, evapotranspiration, infiltration, overflow, and spigot outflow. The results of these simulations, including key variables like water level, overflow, and spigot flow, are recorded and stored in a data frame. This provides a concise yet comprehensive representation of the bucket model's behavior over time.\n",
    "\n",
    "The function `run_bucket_simulation` defined in the following code block is responsible for performing these simulations. It takes into account all the physical attributes and dynamic inputs of each bucket, ensuring that the generated data accurately reflects the complex interactions and processes within the hydrological system. This ground truth data provides the real-world scenarios and responses that the model needs to learn and replicate.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_bucket_simulation(ibuc):\n",
    "    columns = ['precip', 'et', 'h_bucket', 'q_overflow', 'q_spigot']\n",
    "    columns.extend(bucket_attributes_list)\n",
    "    # Memory to store model results\n",
    "    df = pd.DataFrame(index=list(range(len(in_list[ibuc]))), columns=columns)\n",
    "    \n",
    "    # Main loop through time\n",
    "    for t, precip_in in enumerate(in_list[ibuc]):\n",
    "        \n",
    "        # Add the input mass to the bucket\n",
    "        h_water_level[ibuc] = h_water_level[ibuc] + precip_in\n",
    "\n",
    "        # Lose mass out of the bucket. Some periodic type loss, evaporation, and some infiltration...\n",
    "        et = np.max([0, (buckets[\"A_bucket\"][ibuc] / buckets[\"ET_parameter\"][ibuc]) * np.sin(t) * np.random.normal(1, noise['pet'])])\n",
    "        infiltration = h_water_level[ibuc] * buckets[\"K_infiltration\"][ibuc]\n",
    "        h_water_level[ibuc] = np.max([0 , (h_water_level[ibuc] - et)])\n",
    "        h_water_level[ibuc] = np.max([0 , (h_water_level[ibuc] - infiltration)])\n",
    "        h_water_level[ibuc] = h_water_level[ibuc] * np.random.normal(1, noise['et'])\n",
    "\n",
    "        # Overflow if the bucket is too full\n",
    "        if h_water_level[ibuc] > buckets[\"H_bucket\"][ibuc]:\n",
    "            mass_overflow[ibuc] = h_water_level[ibuc] - buckets[\"H_bucket\"][ibuc]\n",
    "            h_water_level[ibuc] = buckets[\"H_bucket\"][ibuc] \n",
    "            h_water_level[ibuc] = h_water_level[ibuc] - np.random.uniform(0, noise['q'])\n",
    "\n",
    "        # Calculate head on the spigot\n",
    "        h_head_over_spigot = (h_water_level[ibuc] - buckets[\"H_spigot\"][ibuc] ) \n",
    "        h_head_over_spigot = h_head_over_spigot * np.random.normal(1, noise['head'])\n",
    "\n",
    "        # Calculate water leaving bucket through spigot\n",
    "        if h_head_over_spigot > 0:\n",
    "            velocity_out = np.sqrt(2 * g * h_head_over_spigot)\n",
    "            spigot_out = velocity_out *  buckets[\"A_spigot\"][ibuc] * time_step\n",
    "            h_water_level[ibuc] = h_water_level[ibuc] - spigot_out\n",
    "        else:\n",
    "            spigot_out = 0\n",
    "\n",
    "        # Save the data in time series\n",
    "        df.loc[t,'precip'] = precip_in\n",
    "        df.loc[t,'et'] = et\n",
    "        df.loc[t,'h_bucket'] = h_water_level[ibuc]\n",
    "        df.loc[t,'q_overflow'] = mass_overflow[ibuc]\n",
    "        df.loc[t,'q_spigot'] = spigot_out\n",
    "        for attribute in bucket_attributes_list:\n",
    "            df.loc[t, attribute] = buckets[attribute][ibuc]\n",
    "\n",
    "        mass_overflow[ibuc] = 0\n",
    "        \n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having established the function to run our bucket model simulations, the next step is to execute these simulations for each bucket configuration and store the results.\n",
    "\n",
    "We use a dictionary, `bucket_dictionary`, to store the simulation results. Each key in this dictionary corresponds to a specific bucket index, and the value is the data frame containing the simulation results for that bucket. This structure allows for efficient organization and retrieval of the simulation data, which is essential for the subsequent analysis and model training.\n",
    "\n",
    "The following code block iterates through all the buckets, running the `run_bucket_simulation` function for each. This function generates the time series data for each bucket, reflecting how its water level and flow rates change over time in response to the input conditions. By storing these results in `bucket_dictionary`, we create an accessible record of each bucket's behavior, forming the basis for training and evaluating our LSTM model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_dictionary = {}\n",
    "for ibuc in range(n_buckets):\n",
    "    bucket_dictionary[ibuc] = run_bucket_simulation(ibuc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.7 Visualizing a Sample of the Bucket Fluxes\n",
    "\n",
    "After running our simulations, we visually inspect the results to ensure that the model is behaving as expected. Visualization helps us confirm the presence of fluxes and validate the realism of the generated values, particularly for spigot (channel flow) and overflow (flooding) scenarios.\n",
    "\n",
    "The function `viz_simulation` is designed to plot the model simulations for a given bucket. This visualization includes:\n",
    "\n",
    "- **Model Inputs**: A plot of the input variables (such as precipitation and evapotranspiration) over time. This helps us understand the driving forces behind the bucket's hydrological behavior.\n",
    "- **Model Outputs**: A plot of the output variables (such as overflow and spigot flow) along with the water level in the bucket over time. This gives us insight into the system's response to the inputs.\n",
    "\n",
    "By examining these plots, we can assess whether the fluxes and water levels are within realistic ranges and behaving as expected under different conditions. This step is crucial for verifying the accuracy and reliability of our simulations before proceeding to model training.\n",
    "\n",
    "The following code block defines the `viz_simulation` function and demonstrates how to visualize the simulation results for a selected bucket. This visual inspection is an integral part of the modeling process, ensuring that our data is suitable for training the LSTM model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viz_simulation(ibuc):\n",
    "    fig = plt.figure(figsize=(12, 3))\n",
    "    ax1 = fig.add_subplot(1, 2, 1)\n",
    "    ax2 = fig.add_subplot(1, 2, 2)\n",
    "    print('Bucket:', ibuc)\n",
    "    print(\"Overflow mean:\", np.round(bucket_dictionary[ibuc].q_overflow.mean(),2))\n",
    "    print(\"Overflow max:\", np.round(bucket_dictionary[ibuc].q_overflow.max(),2))\n",
    "    bucket_dictionary[ibuc].loc[:100,input_vars].plot(ax=ax1)\n",
    "    bucket_dictionary[ibuc].loc[:100,output_vars].plot(ax=ax2)\n",
    "    bucket_dictionary[ibuc].loc[:100,\"h_bucket\"].plot(ax=ax2)\n",
    "    ax1.set_title('Model inputs')\n",
    "    ax2.set_title('Model outputs')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's visualize the simulation results for the buckets in our validation set. This will help us assess the model's performance on different bucket configurations and ensure that the simulated fluxes are realistic and consistent with our expectations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for ibuc in buckets_for_val:\n",
    "    viz_simulation(ibuc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Deep Learning Model\n",
    "\n",
    "In this section, we will set up our deep learning model and outline the training procedure. The model will be trained on the \"ground truth\" data generated from our bucket simulations, learning to predict hydrological behaviors based on various inputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Defining the Neural Network Model\n",
    "\n",
    "Now we use the data from our previous simulations to train a deep learning model. This process involves learning from the generated data and making predictions based on that learning. Our goal is to use the simulations to extract valuable insights and apply them in our learning and prediction processes.\n",
    "\n",
    "We define a class called `LSTM1`, which is a PyTorch module for a single-layer Long Short-Term Memory (LSTM) network. Here's a brief overview of its structure and functionality:\n",
    "\n",
    "- **Input**: The module takes a tensor `x` of shape `(batch_size, seq_length, input_size)`, representing a sequence of `batch_size` samples, each of length `seq_length`, with `input_size` features at each time step.\n",
    "- **LSTM Layer**: Defined using the `nn.LSTM` class, it specifies the size of the input layer (`input_size`) and the hidden state (`hidden_size`). The `batch_first=True` parameter indicates that the first dimension of the input tensor is the batch size.\n",
    "- **Activation Functions**: The output of the LSTM layer passes through a ReLU activation function for non-linearity.\n",
    "- **Fully Connected Layer**: A linear layer (`nn.Linear`) with `num_classes` output units to generate the final predictions.\n",
    "- **Forward Method**: This method processes the input tensor `x`, along with an optional tuple `init_states` representing the initial hidden and internal states of the LSTM layer, and returns the output tensor prediction. If `init_states` is not provided, it initializes as a tensor of zeros.\n",
    "\n",
    "This LSTM model is designed to capture the temporal dynamics of our hydrological system, learning to predict future states based on past observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM1(nn.Module):\n",
    "    def __init__(self, num_classes, input_size, hidden_size, num_layers, batch_size, seq_length):\n",
    "        super(LSTM1, self).__init__()\n",
    "        self.num_classes = num_classes\n",
    "        self.num_layers = num_layers\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.seq_length = seq_length \n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, batch_first=True)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.tanh = nn.Tanh()\n",
    "        self.fc_1 =  nn.Linear(hidden_size, num_classes) #fully connected 1\n",
    "   \n",
    "    def forward(self, x, init_states=None):\n",
    "\n",
    "        if init_states is None:\n",
    "            h_t = Variable(torch.zeros(1, x.size(0), self.hidden_size, device=device)) # hidden state\n",
    "            c_t = Variable(torch.zeros(1, x.size(0), self.hidden_size, device=device)) # internal state\n",
    "            init_states = (h_t, c_t)\n",
    "           \n",
    "        out, _ = self.lstm(x, init_states)\n",
    "        out = self.relu(out)\n",
    "        prediction = self.fc_1(out) # Dense, fully connected layer\n",
    "        \n",
    "        return prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Defining a Procedure for Model Validation\n",
    "\n",
    "Model validation helps us verify that the model is functioning as expected, especially when we make changes to hyperparameters. Validation ensures that our model can accurately predict outcomes based on the input data and that it generalizes well to unseen data.\n",
    "\n",
    "In this part of the notebook, we define a function to validate and test the LSTM model. This function also checks the water balance of the system, an important aspect of ensuring the model's accuracy in simulating hydrological processes. Here's a brief overview of what the function does:\n",
    "\n",
    "- **Model Prediction**: It uses the pre-defined LSTM model to make predictions on the validation data.\n",
    "- **Performance Metrics**: The function computes two key metrics, the Nash-Sutcliffe Efficiency (NSE), for the `spigot_out` and `mass_overflow` columns of the dataframe. These metrics help us understand how well the model's predictions match the actual data.\n",
    "- **Visualization**: We plot the actual `spigot_out` and `mass_overflow` values against their corresponding LSTM predictions. This visual comparison is crucial for assessing the model's performance.\n",
    "- **Water Balance Check**: The function sums up the input (precipitation), evapotranspiration, mass overflow, spigot outflow, and the last recorded water level, comparing this to the total mass out of or left in the system. This check ensures that the model accurately represents the water balance in the system.\n",
    "- **Mass Residual**: We calculate and print out the percent mass residual as a measure of the system's balance.\n",
    "\n",
    "This validation procedure is essential for ensuring that our model is reliable and effective in simulating the complex dynamics of hydrological systems.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_validation_period(lstm, np_val_seq_X, ibuc, n_plot=100):\n",
    "    \n",
    "    def __make_prediction():\n",
    "        lstm_output_val = lstm(torch.Tensor(np_val_seq_X[ibuc]).to(device=device))\n",
    "        val_spigot_prediction = []\n",
    "        val_overflow_prediction = []\n",
    "        for i in range(lstm_output_val.shape[0]):\n",
    "            val_spigot_prediction.append((lstm_output_val[i,-1,1].cpu().detach().numpy() * \\\n",
    "                                    np.std(df.loc[train_start:train_end,'q_spigot'])) + \\\n",
    "                                   np.mean(df.loc[train_start:train_end,'q_spigot']))\n",
    "\n",
    "            val_overflow_prediction.append((lstm_output_val[i,-1,0].cpu().detach().numpy() * \\\n",
    "                                    np.std(df.loc[train_start:train_end,'q_overflow'])) + \\\n",
    "                                   np.mean(df.loc[train_start:train_end,'q_overflow']))\n",
    "        return val_spigot_prediction, val_overflow_prediction\n",
    "    \n",
    "    def __compute_nse():\n",
    "        spigot_out = df.loc[val_start:val_end, 'q_spigot']\n",
    "        spigot_mean = np.mean(spigot_out)\n",
    "        spigot_pred_variance = 0\n",
    "        spigot_obs_variance = 0\n",
    "\n",
    "        overflow_out = df.loc[val_start:val_end, 'q_overflow']\n",
    "        overflow_mean = np.mean(overflow_out)\n",
    "        overflow_pred_variance = 0\n",
    "        overflow_obs_variance = 0\n",
    "\n",
    "        for i, pred_spigot in enumerate(val_spigot_prediction):\n",
    "            t = i + seq_length - 1\n",
    "            spigot_pred_variance += np.power(( pred_spigot          - spigot_out.values[t]), 2)\n",
    "            spigot_obs_variance  += np.power(( spigot_mean          - spigot_out.values[t]), 2)\n",
    "\n",
    "        for i, pred_overflow in enumerate(val_overflow_prediction):\n",
    "            t = i + seq_length - 1\n",
    "            overflow_pred_variance += np.power((pred_overflow          - overflow_out.values[t]), 2)\n",
    "            overflow_obs_variance  += np.power((overflow_mean          - overflow_out.values[t]), 2)\n",
    "        spigot_nse = np.round( 1 - ( spigot_pred_variance / spigot_obs_variance   ), 4)\n",
    "        overland_flow_nse = np.round( 1 - ( overflow_pred_variance / overflow_obs_variance ), 4)\n",
    "        return spigot_nse, overland_flow_nse\n",
    "\n",
    "    def __compute_mass_balance():\n",
    "        mass_in = df.sum()['precip']\n",
    "        mass_out = df.sum()['et'] + \\\n",
    "                   df.sum()['q_overflow'] + \\\n",
    "                   df.sum()['q_spigot'] + \\\n",
    "                   df.loc[num_records - 1, 'h_bucket']\n",
    "        return mass_in, mass_out\n",
    "\n",
    "        \n",
    "    df = bucket_dictionary[ibuc]\n",
    "    val_spigot_prediction, val_overflow_prediction = __make_prediction()\n",
    "    spigot_nse, overland_flow_nse = __compute_nse()\n",
    "    mass_in, mass_out = __compute_mass_balance()\n",
    "        \n",
    "    print(\"Spigot NSE\", spigot_nse)\n",
    "    print(\"Overflow NSE\", overland_flow_nse)\n",
    "    print(\"Mass into the system: \", mass_in)\n",
    "    print(\"Mass out or left over:\", mass_out)\n",
    "    print(\"percent mass resudual: {:.0%}\".format((mass_in - mass_out) /mass_in))\n",
    "    \n",
    "    fig = plt.figure(figsize=(12, 3))\n",
    "    ax1 = fig.add_subplot(1, 2, 1)\n",
    "    ax2 = fig.add_subplot(1, 2, 2)\n",
    "    ax1.plot(df.loc[val_start + seq_length-1:val_start+n_plot+seq_length - 1,'q_spigot'].values, label=\"Spigot out\")\n",
    "    ax1.plot(val_spigot_prediction[:n_plot], label=\"LSTM spigot out\")\n",
    "    ax1.legend()\n",
    "    ax2.plot(df.loc[val_start+seq_length-1:val_start+n_plot+seq_length-1,'q_overflow'].values, label=\"Overflow\")\n",
    "    ax2.plot(val_overflow_prediction[:n_plot], label=\"LSTM Overflow\")\n",
    "    ax2.legend()\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Instantiating the Neural Network Model (LSTM)\n",
    "\n",
    "Now that we have defined our LSTM model's architecture and identified the necessary hyperparameters in Section 1.2, the next step is to instantiate the LSTM model. This involves creating a specific instance of the model with the parameters tailored to our dataset and simulation requirements.\n",
    "\n",
    "The process of instantiation includes:\n",
    "\n",
    "- Setting the number of output classes (`num_classes`) based on our model's output requirements.\n",
    "- Defining the size of the input layer (`input_size`) to match the number of features in our dataset.\n",
    "- Specifying the size of the hidden state (`hidden_state_size`) and the number of LSTM layers (`num_layers`), which determine the model's capacity to learn and remember information over time.\n",
    "- Setting the batch size (`batch_size`) and sequence length (`seq_length`), which are important for the training process.\n",
    "\n",
    "We also set the model to run on the specified device (CPU or CUDA/GPU) for efficient computation. The following code block creates this LSTM model instance, setting it up for the subsequent training and evaluation steps. This step is crucial in bringing our deep learning model to life, ready to learn from the data we have prepared.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1)\n",
    "lstm = LSTM1(num_classes=n_output,  \n",
    "             input_size=n_input,    \n",
    "             hidden_size=hidden_state_size, \n",
    "             num_layers=num_layers, \n",
    "             batch_size=batch_size, \n",
    "             seq_length=seq_length).to(device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Setting Up the Data to Feed into the Model\n",
    "\n",
    "We will set up data for three key purposes:  \n",
    "- **Training**: This data is used to calculate the loss, which is then backpropagated through the model to adjust its parameters.\n",
    "- **Validation**: After training, we use the validation data to get predictions from the model and assess its performance.\n",
    "- **Testing**: This is the data we ultimately use to report the LSTM's performance. It's crucial that testing is the final step; revisiting validation after testing would lead to P-hacking, compromising the integrity of our model evaluation.\n",
    "\n",
    "Each of these datasets plays a vital role in developing a robust and accurate deep learning model.\n",
    "\n",
    "Note: testing is the last thing we would do, if we go back to validation after this step, we would be P-hacking."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Fitting a Scaler to the Training Set to Transform All the Data**\n",
    "\n",
    "Normalization and standardization of data are key steps in preparing it for a deep learning model like LSTM. Here, we fit a scaler to the training set, which allows us to transform both input and output variables to a normalized and standardized scale. This process is crucial for several reasons:\n",
    "\n",
    "- It helps in training the model more effectively, as normalized data typically leads to faster convergence and better performance.\n",
    "- It ensures consistency across different datasets (training, validation, and testing), making the model's learning and predictions more reliable.\n",
    "\n",
    "The scaler is fitted only on the training data to prevent information leakage from the validation and testing sets. Once fitted, this scaler will be used to transform all our datasets accordingly, ensuring that our LSTM model receives data in the most suitable format for learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_scaler():\n",
    "    frames = [bucket_dictionary[ibuc].loc[train_start:train_end, input_vars] for ibuc in buckets_for_training]\n",
    "    df_in = pd.concat(frames)    \n",
    "    scaler_in = StandardScaler()\n",
    "    _ = scaler_in.fit_transform(df_in)\n",
    "\n",
    "    frames = [bucket_dictionary[ibuc].loc[train_start:train_end, output_vars] for ibuc in buckets_for_training]\n",
    "    df_out = pd.concat(frames)    \n",
    "    scaler_out = StandardScaler()\n",
    "    _ = scaler_out.fit_transform(df_out)\n",
    "    return scaler_in, scaler_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's execute the `fit_scaler` function to obtain the scalers for input and output data. These scalers, `scaler_in` and `scaler_out`, will be used to normalize our datasets, ensuring that our LSTM model receives data in a consistent and standardized format.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaler_in, scaler_out = fit_scaler()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Function to Create Data Loader for Each Data Split**\n",
    "\n",
    "Proper data preparation and organization are crucial for the efficient training of our deep learning model. In this step, we focus on creating data loaders, which are essential tools in handling batch processing and shuffling of the data during training.\n",
    "\n",
    "The data loaders not only streamline the process of feeding data into the neural network but also incorporate important preprocessing steps. This includes scaling the input and output variables using the scalers we previously fitted. Such preprocessing ensures that the data is in the optimal format for the LSTM model to process and learn from.\n",
    "\n",
    "The function `make_data_loader` is designed to:\n",
    "\n",
    "- Transform and scale the data for each bucket using the fitted scalers.\n",
    "- Organize the data into sequences of inputs (`np_seq_X`) and corresponding outputs (`np_seq_y`), which are crucial for time-series prediction tasks like ours.\n",
    "- Create a TensorDataset and DataLoader for each bucket. These loaders will manage the batching and shuffling of data, making the training process more efficient and effective.\n",
    "\n",
    "This function is a key component in our data pipeline, ensuring that each split of data (training, validation, and testing) is correctly prepared and ready for use in training our LSTM model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_data_loader(start, end, bucket_list):\n",
    "    loader = {}\n",
    "    np_seq_X = {}\n",
    "    np_seq_y = {}\n",
    "    for ibuc in bucket_list:\n",
    "        df = bucket_dictionary[ibuc]\n",
    "        scaler_in_i = scaler_in.transform(df.loc[start:end, input_vars])\n",
    "        scaler_out_i = scaler_out.transform(df.loc[start:end, output_vars])\n",
    "        np_seq_X[ibuc] = np.zeros((scaler_in_i.shape[0] - seq_length, seq_length, n_input))\n",
    "        np_seq_y[ibuc] = np.zeros((scaler_out_i.shape[0] - seq_length, seq_length, n_output))\n",
    "        for i in range(0, scaler_in_i.shape[0] - seq_length):\n",
    "            t = i+seq_length\n",
    "            np_seq_X[ibuc][i, :, :] = scaler_in_i[i:t,:]\n",
    "            np_seq_y[ibuc][i, :, :] = scaler_out_i[i:t,:]\n",
    "\n",
    "        ds = torch.utils.data.TensorDataset(torch.Tensor(np_seq_X[ibuc]), \n",
    "                                                  torch.Tensor(np_seq_y[ibuc]))\n",
    "        loader[ibuc] = torch.utils.data.DataLoader(ds, batch_size=batch_size, shuffle=True)\n",
    "    return loader, np_seq_X, np_seq_y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have defined the `make_data_loader` function, we will use it along with the parameters we set earlier in the notebook to generate data loaders for our training, validation, and testing sets. These data loaders are crucial for efficiently managing the flow of data through our LSTM model during each phase of the training and evaluation process.\n",
    "\n",
    "- **Training Data Loader**: This loader will handle the data designated for training the model. It ensures that the model is exposed to various scenarios and learns to predict the system's behavior accurately.\n",
    "- **Validation Data Loader**: The validation loader is used to assess the model's performance on data it hasn't seen during training. This step is for tuning the model and ensuring it generalizes well.\n",
    "- **Testing Data Loader**: Finally, the testing loader will provide the data used to evaluate the model's final performance. These data are only used after the model has been trained and validated.\n",
    "\n",
    "The following code block executes the `make_data_loader` function for each of these datasets, creating the respective loaders and organizing the data into the required format for the LSTM model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader, np_train_seq_X, np_train_seq_y = make_data_loader(train_start, train_end, buckets_for_training)\n",
    "val_loader, np_val_seq_X, np_val_seq_y = make_data_loader(val_start, val_end, buckets_for_val)\n",
    "test_loader, np_test_seq_X, np_test_seq_y = make_data_loader(test_start, test_end, buckets_for_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5 Training the Model: Learning the General Response of the Example Dynamic \"Hydrologic\" System\n",
    "\n",
    "Now we arrive at a pivotal moment in our project - training the model. All the steps we've taken so far have been in preparation for this crucial phase. Here, we define a function to train our LSTM neural network model, utilizing the `nn.MSELoss()` loss function and the Adam optimizer, along with the hyperparameters we defined earlier.\n",
    "\n",
    "**Key Aspects of the Training Process**:\n",
    "\n",
    "- **Epochs**: The training is conducted over a specified number of epochs. Each epoch represents a complete pass through the entire training dataset.\n",
    "- **Batch Processing**: Within each epoch, the training data is divided into batches. This is handled efficiently by our data loaders.\n",
    "- **Model Forward Pass**: For each batch, the data is fed through the LSTM model to generate predictions.\n",
    "- **Loss Calculation**: The model's predictions are compared with the target values using the Mean Squared Error (MSE) loss function.\n",
    "- **Backward Pass and Optimization**: The gradients are calculated, and the Adam optimizer updates the model's weights.\n",
    "- **Progress Tracking**: We use the `tqdm` library to visually track the progress of training, providing insights into the loss and RMSE (Root Mean Squared Error) at each step.\n",
    "- **RMSE Calculation**: After each epoch, we calculate the average RMSE, which gives us a measure of the model's prediction accuracy.\n",
    "\n",
    "This function encapsulates the entire training process, systematically updating our model's parameters to learn the general response of our hydrologic system. The outcome will be a trained LSTM model capable of making accurate predictions based on the patterns it has learned from the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(lstm, train_loader, buckets_for_training):\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = optim.Adam(lstm.parameters(), lr=learning_rate[0])\n",
    "    epoch_bar = tqdm(range(num_epochs),desc=\"Training\", position=0, total=num_epochs)\n",
    "    \n",
    "    # Create a dictionary to store the results\n",
    "    results = {}\n",
    "    for epoch in epoch_bar:\n",
    "\n",
    "        for ibuc in buckets_for_training:\n",
    "\n",
    "            batch_bar = tqdm(enumerate(train_loader[ibuc]),\n",
    "                             desc=\"Bucket: {}, Epoch: {}\".format(str(ibuc),str(epoch)),\n",
    "                             position=1,\n",
    "                             total=len(train_loader[ibuc]), leave=False, disable=True)\n",
    "\n",
    "            for i, (data, targets) in batch_bar:\n",
    "\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                optimizer = optim.Adam(lstm.parameters(), lr=learning_rate[epoch])\n",
    "\n",
    "                data = data.to(device=device)\n",
    "                targets = targets.to(device=device)\n",
    "\n",
    "                # Forward\n",
    "                lstm_output = lstm(data) \n",
    "                loss = criterion(lstm_output, targets)\n",
    "\n",
    "                #backward\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "\n",
    "                # gradient descent or adam step\n",
    "                optimizer.step()\n",
    "\n",
    "                batch_bar.set_postfix(loss=loss.to(device).item(),\n",
    "                                      RMSE=\"{:.2f}\".format(loss**(1/2)),\n",
    "                                      epoch=epoch)\n",
    "                batch_bar.update()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                rmse_list = []\n",
    "                for i, (data_, targets_) in enumerate(train_loader[ibuc]):\n",
    "                    data_ = data_.to(device=device)\n",
    "                    targets_ = targets_.to(device=device)\n",
    "                    lstm_output_ = lstm(data_)\n",
    "                    MSE_ = criterion(lstm_output_, targets_)\n",
    "                    rmse_list.append(MSE_**(1/2))\n",
    "\n",
    "            meanrmse = np.mean(np.array(torch.Tensor(rmse_list)))\n",
    "            epoch_bar.set_postfix(loss=loss.cpu().item(),\n",
    "                                  RMSE=\"{:.2f}\".format(meanrmse),\n",
    "                                  epoch=epoch)\n",
    "            \n",
    "            if ibuc not in results:\n",
    "                results[ibuc] = {\"loss\": [], \"RMSE\": []}\n",
    "            results[ibuc][\"loss\"].append(loss.cpu().item())\n",
    "            results[ibuc][\"RMSE\"].append(meanrmse)\n",
    "            #....todo..? also add IT metrics in results in this function ...\n",
    "            batch_bar.update()\n",
    "\n",
    "        𝗰𝗹𝗲𝗮𝗿_𝗼𝘂𝘁𝗽𝘂𝘁\n",
    "        \n",
    "    return lstm, results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to put our training function into action. We will run the `train_model` function, specifying the buckets to use for training and the LSTM model instance we previously instantiated. This step will engage the model in the learning process, using the training data to adjust and refine its parameters for accurate prediction of hydrologic behaviors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lstm, results = train_model(lstm, train_loader, buckets_for_training)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.6 Visualizing the Learning Curves\n",
    "\n",
    "After training our LSTM model, it's important to assess its learning progress and convergence. To do this, we visualize the learning curves, plotting the loss and Root Mean Square Error (RMSE) metrics for each epoch. These curves provide valuable insights into the model's learning behavior over time.\n",
    "\n",
    "- **Loss Curve**: This graph shows the loss value (calculated using the MSE loss function) at each epoch. A decreasing trend in the loss curve indicates that the model is learning effectively.\n",
    "- **RMSE Curve**: The RMSE curve displays the root mean square error for each epoch. Like the loss curve, a downward trend in RMSE suggests improvement in the model's prediction accuracy.\n",
    "\n",
    "By examining these curves, we can determine if the model fitting has converged, which is indicated by a plateau in both loss and RMSE values. This visualization is crucial for understanding the model's learning dynamics and for making any necessary adjustments to the training process.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viz_learning_curve(results):\n",
    "    fig = plt.figure(figsize=(12, 3))\n",
    "    ax1 = fig.add_subplot(1, 2, 1)\n",
    "    ax2 = fig.add_subplot(1, 2, 2)\n",
    "    for ibuc in buckets_for_training:\n",
    "        ax1.plot(range(num_epochs), results[ibuc]['loss'])\n",
    "        ax2.plot(range(num_epochs), results[ibuc]['RMSE'])\n",
    "    ax1.set_ylabel('Loss')\n",
    "    ax2.set_ylabel('RMSE')\n",
    "    ax1.set_xlabel('Epoch')\n",
    "    ax2.set_xlabel('Epoch')\n",
    "    plt.suptitle(\"Learning curves for each bucket\") \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now visualize the learning curves using the `viz_learning_curve` function. This will provide us with graphical representations of the loss and RMSE over each training epoch, offering insights into the model's learning progress and convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_learning_curve(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.6 Checking that the Model Works on the Validation Data\n",
    "\n",
    "With our LSTM model now trained, the next crucial step is to evaluate its performance on the validation data. This phase is essential for assessing how well the model generalizes to data it hasn't seen during training.\n",
    "\n",
    "By testing the model on our validation split, we can:\n",
    "\n",
    "- Verify the model's predictive accuracy against known outcomes.\n",
    "- Identify any potential issues, such as overfitting to the training data.\n",
    "- Make any necessary adjustments before the final evaluation on the test data.\n",
    "\n",
    "The following code block will apply the `check_validation_period` function to each bucket in our validation set. This function not only tests the model's predictions but also checks the water balance of the system, ensuring that our model is providing realistic and reliable outputs.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for ibuc in buckets_for_val:\n",
    "    check_validation_period(lstm, np_val_seq_X, ibuc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Clean up the code, add modularity, for easy use and modification\n",
    "\n",
    "## 3.1 src\n",
    "Contains the clean, modular code for the project, divided into several scripts:  \n",
    "  - `data_generation.py`: Generates synthetic data simulating a leaking bucket hydrologic response.  \n",
    "  - `lstm.py`: Implements the LSTM deep learning model.  \n",
    "  - `model_controller.py`: Manages model functions such as training and data normalization.  \n",
    "  - `validation.py`: Executes model validation and computes performance statistics.  \n",
    "  - `vizualization.py`: Provides functionality for visualizing data and predictions.  \n",
    "  \n",
    "## 3.2 run\n",
    "\n",
    "Includes Jupyter notebooks and scripts for direct execution of the model training and simulation:\n",
    "  - `run_deep_bucket_lab_with_graphics.ipynb`: A Jupyter notebook with comprehensive visual outputs.\n",
    "  - `run_deep_bucket_lab.py`: A script for running the model from the command line.\n",
    "\n",
    "## 3.3 configuration\n",
    "\n",
    "Stores configuration files that dictate model parameters and settings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Experimentation\n",
    "\n",
    "With our general leaky bucket model setup and training/validation framework in place, we now enter the experimentation phase. This stage allows us to explore various scenarios and understand how different configurations and settings impact the model's behavior and predictions.\n",
    "\n",
    "Experimentation is key to gaining deeper insights into the dynamics of our hydrologic system and the capabilities of our LSTM model. It enables us to test hypotheses, explore 'what-if' scenarios, and refine our understanding of the system's responses under varied conditions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this phase, you're encouraged to experiment with variations in the bucket characteristics and modeling setup. Here are some ideas to get you started:\n",
    "\n",
    "- **Data and system uncertainty**:\n",
    "  - `./experiments/noise_iteration.py` Hosts a script designed to run a specific experiment testing model sensitivity to input noise.\n",
    "- **Modifying Ground Truth Data**:\n",
    "  - To simulate a more \"flashy\" system, you might reduce the probability of heavy precipitation and increase its magnitude.\n",
    "  - For simulating smaller buckets, try reducing the size of the bucket attributes.\n",
    "  - To add more variability, consider increasing the noise level from 0.1 to different values.\n",
    "\n",
    "- **Adjusting the Modeling Setup**:\n",
    "  - Experiment with the number of training buckets. Increasing or decreasing this number can show how the model performs with more or less training data.\n",
    "  - Alter the length of the time series. This can help you understand the impact of short-term versus long-term data on the model's learning and predictions.\n",
    "\n",
    "These modifications can lead to new insights and help you understand the robustness and adaptability of your model. Feel free to try out different combinations and observe how they affect the model's performance and outputs.\n",
    "\n",
    "Head over to `deep_bucket_lab/experiments`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}