in model.fit verbose 1 -> 2

[notebooks.git] / fmda / fmda_kf_rnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "khOTxJsYc91W"
   },
   "source": [
    "# Kalman Filtering and Recurrent Neural Networks for Fuel Moisture\n",
    "## Jan Mandel, University of Colorado Denver\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sXaqfI-EdCEk"
   },
   "source": [
    "## Abstract"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZbtSiYm4dF7B"
   },
   "source": [
    "''Fuel moisture is an important factor of the spread of wildland fires. Some weather stations have fuel moisture sensors and data are available online. We review a simple model of fuel moisture from atmospheric conditions, and show how to adjust the model using the weather station data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eZ6dfHlZ63j1"
   },
   "source": [
    "## Table of contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Snyr0_3sGxty"
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WHIgN2uZ689b"
   },
   "source": [
    "1 Introduction\n",
    "\n",
    "2 Background\n",
    "\n",
    "2.1 Imports\n",
    "\n",
    "2.2 Kalman filter\n",
    "\n",
    "2.2.1 Overview\n",
    "\n",
    "2.2.2 Formulation\n",
    "\n",
    "2.2.3 A Kalman filter tester\n",
    "\n",
    "2.3 Fuel moisture model\n",
    "\n",
    "2.3.1 A simple time lag model\n",
    "\n",
    "2.3.1 Fuel moisture model with drying equilibrium, wetting equilibrium, and rain\n",
    "\n",
    "3 Methods\n",
    "\n",
    "3.1 Kalman filter demonstration on the simple model\n",
    "\n",
    "3.1.1 Creating synthetic data\n",
    "\n",
    "3.1.2 Running the Kalman filter\n",
    "\n",
    "3.2 Acquisition and preprocessing of real data\n",
    "\n",
    "3.2.1 Acquisition of fuel moisture observations\n",
    "\n",
    "3.2.2 Acquisition of weather data\n",
    "\n",
    "3.2.3 Preprocessing and visualization of the weather data\n",
    "\n",
    "4 Results\n",
    "\n",
    "4.1 Kalman filter with fuel moisture observations, followed by forecasting\n",
    "\n",
    "4.2 Model with an augmented state\n",
    "\n",
    "4.3 Kalman filter on the augmented model\n",
    "\n",
    "4.4 A comment on the information flow in the Kalman filter and in neural networks\n",
    "\n",
    "5. Conclusion\n",
    "\n",
    "Contributions of Authors\n",
    "\n",
    "Acknowledgements\n",
    "\n",
    "References\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZFafUPCTO1N1"
   },
   "source": [
    "## 1 Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4_RcdWybPFks"
   },
   "source": [
    "The Kalman filter is at the foundation of many technologies in daily use, from GPS to weather forecasting. No model is completely accurate. Think space navigation: the movement of a Apollo 13 between the moon and the earth, subject to gravitational forces and propulsion, with the position ascertained by visual measurements. No matter how accurate the model of spacecraft motion is, the measurements are always burdened with noise. The idea of Kalman filter is to evolve a quantification of the of the state (here, positin and velocity of the spacecraft) in the form of a covariance matrix, and, using an estimate of the uncertainty of the data, adjust the state to split the difference every time measurements are taken. \n",
    "\n",
    "Here, we use the Kalman filter to estimate the evolution of fuel (dead wood) moisture content from a simple theoretical model, adjusting the state of the model hourly for measurements from fuel moisture a sensor in a wood stick exposed to the elements. This is needed for forecasting of wildfire progress; for this purpose, we also want to have the filter adjust the model from the data, so that it gives more accurate data for future when we only have hourly weather forecast but no actual data - because the future has not happened yet. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M2kbwDPBTB7A"
   },
   "source": [
    "## 2 Background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ar1BbXac49hO"
   },
   "source": [
    "In this section, we take care of preliminaries: we install some packages we need, and then proceed with the Kalman filter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_5F5CuRqc91X"
   },
   "source": [
    "### 2.1 Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K6sWUMf0c91Y"
   },
   "source": [
    "We may need the pygrib package to read weather data, but pygrib requires current numpy while google colab is using an old numpy version for compatibility with tensorflow. We will upgrade numpy and restart the runtime then the notebook will need to be run again. If numpy is current, we just download and import packages we need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X9rvlymMZdJg"
   },
   "source": [
    "### 2.2 Kalman filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x5E2UE3F5gf2"
   },
   "source": [
    "#### 2.2.1 Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NPgTHlCLAlA-"
   },
   "source": [
    "The Kalman filter provides an estimate $u$ of the time evolution of some unknown process, called \"nature\" or \"truth\". We do not know with certainty what the nature is, but we can observe it at regular intervals (steps) with some error. In each step, model $F$ advances the model state $u$ in time, $ u \\leftarrow F(u)$, and attempts to reconcile the state with an observation $d$ of the true state, so $u \\approx d$. The filter modifies the model state $u$ to balance the uncertainty in the model and the data (this is called *analysis*) and the cycle continues. For that purpose, the filter evolves also an estimate of the uncertainly of the model.\n",
    "\n",
    "More generally, instead of $u \\approx d$, only a part of the state is observed, and $Hu \\approx d$ where $H$ is a matrix, or observation function. Basically, $Hu$ is what the data would be if the model was completely accurate. \n",
    "\n",
    "See Kalman (1960) for the original publication, Kalnay (2003) for a gentle introduction, and the [Wikipedia article](https://en.wikipedia.org/wiki/Extended_Kalman_filter)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y6j34L5s5pEL"
   },
   "source": [
    "#### 2.2.2 Formulation\n",
    "\n",
    "---\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b3GZW5vP5_o8"
   },
   "source": [
    "We present the Kalman filter in perhaps the most used form, as extended to nonlinear models.\n",
    " Consider a discrete time model of some natural\n",
    "process. At time step $k$, the model has state $u_{k}\\in\\mathbb{R}^{n}$, which\n",
    "can be approximated from the previous step $u_{k-1}$ by applying the model\n",
    "$\\mathcal{M}$ to get a forecast $u_{k}^{f}=\\mathcal{M}\\left(  u_{k-1}\\right)\n",
    "$. We model uncertainty in the model itself by adding normally distributed\n",
    "noise with mean zero and covariance $Q$ to the uncertainty of $u_{k}^{f}$. We\n",
    "also need to estimate now the uncertainty in the previous state $u_{k-1}$\n",
    "propagates to the uncertainty of the forecast $u_{k}^{f}$. So, assume that the\n",
    "model is differentiable and quantify the uncertainty of the state by a\n",
    "covariance matrix. That is,  assume that at step $k-1$, the state has\n",
    "(approximately) normal distribution with mean $u_{k-1}$ and covariance\n",
    "$P_{k-1}$. Using the Taylor expansion of order $1$ of the model operator at\n",
    "$u_{k-1}$, $\\mathcal{M}\\left(  u\\right)  \\approx\\mathcal{M}\\left(\n",
    "u_{k-1}\\right)  +\\mathcal{M}^{\\prime}\\left(  u_{k-1}\\right)  \\left(\n",
    "u-u_{k-1}\\right)  $, where $\\mathcal{M}^{\\prime}\\left(  u_{k-1}\\right)  $ is\n",
    "the Jacobian matrix of $\\mathcal{M}$ at $u_{k-1}$. It can be shown that the\n",
    "forecast has then (approximately)\\ normal distribution with mean and\n",
    "covariance\n",
    "$$\n",
    "u_{k}^{f}=\\mathcal{M}\\left(  u_{k-1}\\right)  ,\\ P_{k}^{f}=\\mathcal{M}\\left(\n",
    "u_{k-1}\\right)  P_{k-1}\\mathcal{M}^{\\prime}\\left(  u_{k-1}\\right)  +Q\n",
    "$$\n",
    "At time $k$, we also have an observation $d_{k}\\approx Hu_{k}$, where $H$ is a\n",
    "given observation operator, and we want to find $u_{k}$ so that both\n",
    "$$\n",
    "u_{k}\\approx u_{k}^{f}\\text{ and }d_{k}\\approx Hu_{k}.\n",
    "$$\n",
    "We quantify the uncertainly of the error of observation $d_{k}$ by a covariance\n",
    "matrix $R$: assume that the observation error has normal probability\n",
    "distribution with a known covariance $R$. Then, the likelihood of state $u$ is\n",
    "proportional to $e^{-\\left\\Vert d_{k}-Hu\\right\\Vert _{R^{-1}}^{2}/2}$, where\n",
    "we used the notation for the norm $\\left\\Vert v\\right\\Vert _{A}%\n",
    "=\\left(v^{\\top}Av\\right)^{1/2}$ induced by a positive definite matrix $A$. Similarly, we quantify the\n",
    "uncertainty of the state by a covariance matrix $P_{k}$. That is, the forecast\n",
    "state has (approximately) normal distribution with mean $u_{k}^{f}$  and covariance\n",
    "$P_{k}^{f}$. From the Bayes theorem of statistics, the probability distribution\n",
    "of the state after taking the data into account has density\n",
    "$$\n",
    "p_{k}\\left(  u\\right) \\propto e^\\frac{-\\left\\Vert d_{k}\n",
    "-Hu\\right\\Vert_{R^{-1}}^{2}}{2}e^\\frac{-\\left\\Vert u-u_{k}^{f}\\right\\Vert _{\n",
    "{P_{k}^f}^{-1}  }^{2}}{2}%\n",
    "$$\n",
    "where $\\propto$ means proportional.\n",
    "Note that the probability density at $u$ is maximal when $\\left\\Vert\n",
    "d_{k}-Hu\\right\\Vert _{R^{-1}}^{2}+\\left\\Vert u-u_{k}\\right\\Vert _{{P_{k}^{f}}^{-1}}^{2}$\n",
    " is minimal, which quantifies the statement that $d_{k}\\approx\n",
    "Hu_{k}$ and $u\\approx u_{k}^{f}$.  By a direct computation completing the\n",
    "square and using the Sherman-Morrison-Woodbury formula, \n",
    "$$p_{k}\\left(\n",
    "\t\tu\n",
    "\t   \\right) \\propto \n",
    "e^{-\\frac{\n",
    "\t\\left\\Vert u-u_{k\n",
    "\t         }\n",
    "\t\\right\\Vert_\n",
    "\t\t{P_{k\n",
    "\t\t      }^{-1}\n",
    "\t\t}^{2}\n",
    "\t}\n",
    "\t{2}},\n",
    "$$ \n",
    "which is the density of the normal distribution with the mean\n",
    "$$\n",
    "u_{k}^{f}=u_{k}^{f}+K_{k}(d-Hu_{k}^{f}),\\ \\text{where }K_{k}=P_{k}%\n",
    "^{f}H^{\\mathrm{T}}(HP_{k}^{f}H^{\\mathrm{T}}+R)^{-1}%\n",
    "$$\n",
    "and covariance\n",
    "$$\n",
    "P_{k}=\\left(  \\left(  P_{k}^{f}\\right)  ^{-1}+H^{\\mathrm{T}}R^{-1}H\\right)\n",
    "^{-1}=(I-KH)P_{k}^{f}.\n",
    "$$\n",
    "\n",
    "These are the equations of the extended Kalman filter. The original Kalman (1960) filter was\n",
    "formulated for a linear process. The extension to the\n",
    "nonlinear case made broad array of applications possible, including the Apollo spacecraft naviation (McGee and Schmidt, 1966),  and is\n",
    "still a de-facto standard in navigation and GPS.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-bvUtJ_OLwQA"
   },
   "outputs": [],
   "source": [
    "# Import packages\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt \n",
    "import subprocess,os\n",
    "import json\n",
    "import pytz\n",
    "from datetime import datetime, timedelta, time\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "import tensorflow as tf\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Dense, SimpleRNN\n",
    "from keras.utils.vis_utils import plot_model\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.metrics import mean_squared_error\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import keras.backend as K\n",
    "from keras.utils.vis_utils import plot_model\n",
    "from scipy.interpolate import LinearNDInterpolator, interpn\n",
    "from scipy.optimize import root\n",
    "\n",
    "## Local Modules\n",
    "from moisture_models import ext_kf, model_decay, model_moisture, model_augmented, run_augmented_kf\n",
    "from moisture_models import create_RNN, create_RNN_2, staircase, seq2batches\n",
    "from data_funcs import create_synthetic_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Uvsbbv2XZ2Hd"
   },
   "source": [
    "#### 2.2.3 A Kalman filter tester"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gcmGBqPOU1e5"
   },
   "source": [
    "It is a very good idea to make write a simple tester for every piece of code. How else would we know it actually works, and that something basic did not get broken inadvertently, perhaps as a side effect of changing something else? A simple tester may save a great deal of time trying to debug cryptic errors later. And, what better place for a tester that right after the code it is testing so that it gets run every time?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OsOqvQk6ZXZV"
   },
   "outputs": [],
   "source": [
    "# a basic ext_kf test\n",
    "\n",
    "u = [1,\n",
    "     2]\n",
    "P = [[2 , -1],\n",
    "    [-1 , 2]]\n",
    "A = [ [1 ,2],\n",
    "      [3 ,4]]\n",
    "u = np.array(u)      \n",
    "Q = np.array([[1,0],[0,1]])\n",
    "A = np.array(A)\n",
    "def fun(u):\n",
    "  return A @ u, A\n",
    "F = lambda u: fun(u)\n",
    "H = [[1, 0],\n",
    "     [0, 1]]\n",
    "d = [2,\n",
    "    3]\n",
    "R = [[2, 0],\n",
    "    [0, 2]]\n",
    "H = np.array(H)      \n",
    "d = np.array(d)\n",
    "R = np.array(R)\n",
    "ua,Pa = ext_kf(u,P,F,Q)\n",
    "print('ua=',ua)\n",
    "print('Pa=',Pa)\n",
    "ua,Pa = ext_kf(u,P,F,Q,d,H,R)\n",
    "print('ua=',ua)\n",
    "print('Pa=',Pa)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A9ZpmNcdRpmp"
   },
   "source": [
    "### 2.3  Fuel moisture models\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eZL8gN7ISGVh"
   },
   "source": [
    "#### 2.3.1 A simple fuel moisture model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1XvOC4kYSQgH"
   },
   "source": [
    "First consider a simplified fuel moisture model without considering the effect of rain.\n",
    "The evolution of fuel moisture content $m(t)$ is modeled by the time-lag differential equation on interval $\\left[\n",
    "t_{0},t_{1}\\right]  $,\n",
    "$$\n",
    "\\frac{dm}{dt}=\\frac{E-m(t)}{T},\\quad m(t_{0})=m_{0}.\n",
    "$$\n",
    "where the initial fuel moisture content $m_{0}=m\\left(  t_{0}\\right)  $ is the\n",
    "input, and $m_{1}=m(t_{1})$ is the output. Tnus, $m_1=F(m_0)$. The parameters of the model are the\n",
    "fuel moisture equilibrium $E$, assumed to be constant over the interval $\\left[\n",
    "t_{0},t_{1}\\right]  $, NS the characteristic decay time $T$. \n",
    "\n",
    "We can build the general model later by calling this simple model with different\n",
    "equilibria and time constants (drying, wetting, rain).\n",
    "\n",
    "Since $E$ is constant in time, the solution can be found\n",
    "analytically,\n",
    "$$\n",
    "m\\left(  t\\right)  =E+\\left(  m_{0}-E\\right)  e^{-t/T}%\n",
    "$$\n",
    "For convenience, we use $T_{1}=1/T$ instead of $T$, and the model becomes\n",
    "$$\n",
    "m_{1}=E+\\left(  m_{0}-E\\right)  e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}%\n",
    "$$\n",
    "In the extended Kalman filter, we will need the partial derivatives of $m_{1}$\n",
    "with respect to the input and the parameters. Compute\n",
    "$$\n",
    "\\frac{dm_{1}}{d_{m0}}=e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
    "$$\n",
    "$$\n",
    "\\frac{dm_{1}}{dE}=1-e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
    "$$\n",
    "$$\n",
    "\\frac{dm_{1}}{dT_{1}}=-\\left(  m_{0}-E\\right)  \\left(  t_{1}-t_{0}\\right)\n",
    "e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
    "$$\n",
    "At the moment, we need only ${dm_{1}}/{dm_{0}}$ but we put in the code all partials for possible use in future.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dOARZlj-RUCi"
   },
   "source": [
    "#### 2.3.2 Fuel moisture model with drying equilibrium, wetting equilibrium, and rain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AJp6FTpTSx5B"
   },
   "source": [
    "Here is a little more realistic fuel moisture model from Mandel et al. (2004). A rain-wetting lag time $t_{\\mathrm{r}}$ is reached for heavy rain only\n",
    "asymptotically, when the rain intensity $r$ (mm/h) is\n",
    "large:\n",
    "$$\n",
    "\\frac{\\mathrm{d}m}{\\mathrm{d}t}=\\frac{S-m}{t_{\\mathrm{r}}}\\left(1-\\exp\\left(-\\frac{r-r_0}{r_{\\mathrm{s}}}\n",
    "\\right)  \\right),\\ \\text{if}\\ r>r_0, \n",
    "$$\n",
    "where $r_0$ is the threshold rain intensity below which no perceptible\n",
    "wetting occurs, and $r_{\\mathrm{s}}$ is the saturation rain\n",
    "intensity. At the saturation rain intensity, $1-1/e\\approx 0.63$ of\n",
    "the maximal rain-wetting rate is achieved. For 10h fuel, the model takes $S=250\\,{\\%}$,\n",
    "$t_{\\mathrm{r}}=14$h, $r_0=0.05$mm/h and\n",
    "$r_{\\mathrm{s}}=8$mm/h. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ITsKE0psRblG"
   },
   "outputs": [],
   "source": [
    "### Define model function with drying, wetting, and rain equilibria\n",
    "\n",
    "# Parameters\n",
    "r0 = 0.05                                   # threshold rainfall [mm/h]\n",
    "rs = 8.0                                    # saturation rain intensity [mm/h]\n",
    "Tr = 14.0                                   # time constant for rain wetting model [h]\n",
    "S = 250                                     # saturation intensity [dimensionless]\n",
    "T = 10.0                                    # time constant for wetting/drying"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JDLU3B_jV42l"
   },
   "source": [
    "## 3. Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hLPJT3FcA2a7"
   },
   "source": [
    "### 3.1 Kalman filter demonstration on the simple model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kIA3X8vluFdd"
   },
   "source": [
    "We demonstrate the Kalman filter for this model on a simple artificial example. The model is solving the differential equation for one hour. The equilibrium $E$ is constant during the hour, but it changes over the day so that it is higher at night and lower during the day, with a 24-hour period.  First, we create the \"truth\" by choosing the equilibrium $E$ and solving the differential aquation every hour, with a small additive noise. The synthetic data is obtained as the values of the \"truth\", with random noise to simulate observation error."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bBv10PTiChhm"
   },
   "source": [
    "#### 3.1.1 Creating synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GfRxLp4HkcVz"
   },
   "outputs": [],
   "source": [
    "E,m_f,data,hour,h2,DeltaE = create_synthetic_data(days=20,power=4,data_noise=0.01,process_noise=0.0,DeltaE=0.0) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize simulated data\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(hour,E,linestyle='--',c='r',label='Equilibrium')\n",
    "plt.plot(hour,m_f,linestyle='-',c='k',label='10-h fuel truth')\n",
    "plt.scatter(hour[:h2],data[:h2],c='b',label='10-h fuel data')\n",
    "plt.title('Synthetic data')\n",
    "plt.xlabel('Time (hours)')\n",
    "plt.ylabel('Fuel moisture content (%)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-3WLAEpD2yJ"
   },
   "source": [
    "#### 3.1.2 Running the Kalman filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T4g-RrrYAlBD"
   },
   "source": [
    "We have used the same code for model and for the truth, and run the Kalman filter for 10 days. The graph below shows that the model state was remarkably close to the truth, even if the model is fed only noisy observations. This is because the dynamics of the model and of the truth are the same. After 10 days, we let the model continue without any new data to simulate forecasting the future, and the agreement with the truth was still very good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_-CjONZkD18n"
   },
   "outputs": [],
   "source": [
    "# using global E, m_f\n",
    "\n",
    "def plot_m(m,Ec=None,title=None,):  # global hour\n",
    "    hours=hour.shape[0]\n",
    "    %matplotlib inline\n",
    "    plt.figure(figsize=(16,4))\n",
    "    plt.plot(hour,E,linestyle='--',c='r',label='E=Equilibrium data')\n",
    "    # print(len(hour),len(m_f))\n",
    "    plt.plot(hour,m_f,linestyle='-',c='b',label='m_f=10-h fuel truth')\n",
    "    plt.scatter(hour[:h2],data[:h2],c='b',label='data=10-h fuel data')\n",
    "    if m is not None:\n",
    "        plt.plot(hour[:h2],m[:h2],linestyle='-',c='k',label='m=filtered')\n",
    "        plt.plot(hour[h2:hours],m[h2:hours],linestyle='-',c='r',label='m=forecast')\n",
    "    if Ec is not None:\n",
    "        plt.plot(hour,Ec,linestyle='-',c='g',label='Ec=Equilibrium correction')\n",
    "    if title is not None:\n",
    "        plt.title(title) \n",
    "    else:\n",
    "        plt.title('Kalman filtering and forecast on artificial data')\n",
    "    plt.xlabel('Time (hours)') \n",
    "    plt.ylabel('Fuel moisture content (%)')\n",
    "    plt.legend()\n",
    "\n",
    "def kf_example(DeltaE):\n",
    "    hours=hour.shape[0]\n",
    "    m = np.zeros(hours)\n",
    "    m[0]=0.1             # background state  \n",
    "    P = np.zeros(hours)\n",
    "    P[0] = 0.03 # background state variance\n",
    "    Q = np.array([0.02]) # process noise variance\n",
    "    H = np.array([1.])   # all observed\n",
    "    R = np.array([0.02]) # data variance\n",
    "\n",
    "    for t in range(h2):\n",
    "        # use lambda construction to pass additional arguments to the model \n",
    "        m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_decay(u,E[t]+DeltaE,partials=1),Q,\n",
    "                    d=data[t],H=H,R=R)\n",
    "    for t in range(h2,hours - 1):\n",
    "        m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_decay(u,E[t]+DeltaE,partials=1))\n",
    "    return m, P\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d0EFhTPZAlBD",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "DeltaE = 0.0          # bias\n",
    "m, P = kf_example(DeltaE)\n",
    "plot_m(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vqyB2Yz3uCsD"
   },
   "source": [
    "We have recovered the fuel moisture from data with random noise - we **filtered** the noise out. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dl7pBZ9B3Nox"
   },
   "source": [
    "Let's have a look at the evolution of the filter's estimate of the variance $P$. A common problem with the Kalman filter is when the variance converges to zero over time, then, since the filter trusts the model too much, it ignores the observations. Of course, once we switch to forecasting mode, the variance is not of interest. We could keep evolving the variance to bridge over periods when there are no observations, but not in this simplified version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wRJgbmGLc91g"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(P,linestyle='-',c='b',label='Estimated state variance P')\n",
    "plt.title('Kalman filtering and forecast on artificial data')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Estimated variance of fuel moisture (%^2)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ccr-uKbmAlBE"
   },
   "source": [
    "Now what if the model is wrong - different from nature? That is always so in reality. Now suppose that the model and the truth are not the same. That is always the case in reality.  Consider a simple case when the model thinks that the equilibrium $E$ is too high."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "spMdGW8oAlBE"
   },
   "outputs": [],
   "source": [
    "DeltaE = -0.05\n",
    "m, P = kf_example(DeltaE)\n",
    "plot_m(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_INvPNGCmM2"
   },
   "outputs": [],
   "source": [
    "DeltaE = 0.05\n",
    "m, P = kf_example(DeltaE)\n",
    "plot_m(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DQeF7J8T4j2i"
   },
   "source": [
    "We have found a good estimate of the state $m$, while data is available. Also, the estimated state variance $P$ converges with time - we have *learned* the variance that balances the noise. But for forecasting fuel moisture, we need to continue the fuel moisture model into the future, and we can't have any measurements from future. We only have the equilibrium from weather forecast. And the forecast and the truth disagree - as soon as there is no data to attract the simulation, the model is doing its own thing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-WMWCDz4DX45"
   },
   "source": [
    "#### 3.2 Model with an augmented state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jivOYEhiXMi5"
   },
   "source": [
    "In reality, the equilibrium moisture $E$ computed from atmospheric conditions\n",
    "generally does not agree with the data. We want to add a correction $\\Delta\n",
    "E$ to $E$ constant in time, and identify the new parameter $\\Delta E$ from data. \n",
    "Because the Kalman filter identifies state, add the parameter to the state.\n",
    "Define augmented state $u=\\left[\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "m\\\\\n",
    "\\Delta E\n",
    "\\end{array}\n",
    "\\right]  .$ Since $\\Delta E$ is constant in time, it satisfies the\n",
    "differential equation $\\frac{d\\Delta E}{dt}=0.$ So, we want to estimate the\n",
    "state $u$ governed by the\n",
    "$$\n",
    "\\frac{d}{dt}\\left[\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "m\\\\\n",
    "\\Delta E\n",
    "\\end{array}\n",
    "\\right]  =\\left[\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "\\frac{E+\\Delta E-m(t)}{T}\\\\\n",
    "0\n",
    "\\end{array}\n",
    "\\right]  ,\n",
    "$$\n",
    "which we write as $\\frac{du}{dt}=F(u),$ where\n",
    "$$\n",
    "F(u)=\\left[\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "F_{1}\\left(  u\\right)  \\\\\n",
    "F_{2}\\left(  u\\right)\n",
    "\\end{array}\n",
    "\\right]  =F\\left(  \\left[\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "m\\\\\n",
    "\\Delta E\n",
    "\\end{array}\n",
    "\\right]  \\right)  =\\left[\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "\\left(  E+\\Delta E-m(t)\\right)  T_{1}\\\\\n",
    "0\n",
    "\\end{array}\n",
    "\\right]  ,\\quad T_{1}=\\frac{1}{T}.\n",
    "$$\n",
    "The Jacobian of $F$ is\n",
    "$$\n",
    "\\left[\n",
    "\\begin{array}\n",
    "[c]{cc}\n",
    "\\frac{\\partial F_{1}}{\\partial u_{1}} & \\frac{\\partial F_{1}}{\\partial u_{2}\n",
    "}\\\\\n",
    "\\frac{\\partial F_{2}}{\\partial u_{1}} & \\frac{\\partial F_{2}}{\\partial u_{2}}\n",
    "\\end{array}\n",
    "\\right]  =\\left[\n",
    "\\begin{array}\n",
    "[c]{cc}\n",
    "\\frac{\\partial m_{1}}{\\partial m_{0}} & \\frac{\\partial m_{1}}{\\partial E}\\\\\n",
    "\\frac{\\partial\\Delta E}{\\partial m_{0}} & \\frac{\\partial\\Delta E}\n",
    "{\\partial\\Delta E}\n",
    "\\end{array}\n",
    "\\right]  =\\left[\n",
    "\\begin{array}\n",
    "[c]{cc}\n",
    "\\frac{\\partial m_{1}}{\\partial m_{0}} & \\frac{\\partial m_{1}}{\\partial E}\\\\\n",
    "0 & 1\n",
    "\\end{array}\n",
    "\\right]\n",
    "$$\n",
    "Here is a function that implements the augmented model $F$. The input is\n",
    "$u_{0}$. The output is $u_{1}$ and the Jacobian $du_{1}/du_{0}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MJ1C_1Omc91s"
   },
   "source": [
    "\n",
    "Define augmented model function. Also, add use drying, wetting, and rain equilibria"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9vUv1Uc5OfPI"
   },
   "outputs": [],
   "source": [
    "def augmented_example(DeltaE):\n",
    "    hours=hour.shape[0]\n",
    "    h2 = int(hours/2)\n",
    "    m, Ec = run_augmented_kf(data,E+DeltaE,E+DeltaE,0*E,h2,hours)  # data, E, hours are global\n",
    "    return m, Ec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QQnRlereDHbf"
   },
   "outputs": [],
   "source": [
    "m, Ec=augmented_example(0.1)\n",
    "plot_m(m, Ec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZOC6bplsIRCx"
   },
   "outputs": [],
   "source": [
    "m, Ec=augmented_example(0.0)\n",
    "plot_m(m, Ec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xuMx41fXIejB"
   },
   "outputs": [],
   "source": [
    "m, Ec=augmented_example(-0.1)\n",
    "plot_m(m, Ec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ejlc1sbENm6R"
   },
   "source": [
    "## From Kalman filter to neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6zeGArbmNkAW"
   },
   "source": [
    "In the Kalman filter, at each time step $k$,\n",
    "\n",
    "* the input state is $u_{k-1}$ size $n$ and its covariance matrix $P_{k-1}$ size $n \\times n$.\n",
    "* the model is applied to external data $e_k$ and the input $u_{k-1},P_{k-1}$ produce the forecast $u_k^f$ and its covariance $P^f_k$\n",
    "* the new state $u_k$ is found by minimizing $|| u^f_k - u_k||^2_{P^f_k} + ||H u_k - d_k||^2_{R}$   \n",
    "* the new state covariance is $P_k = ( (P^f_k)^{-1} + H^\\top R^{-1} H)^{-1}$.\n",
    "\n",
    "Here, the state consists of \n",
    "* the fuel moisture and the adjustment to the equilibrium, dimension 2\n",
    "* the covariance matrix of vector of dimension 2, which is symmetric $2 \\times 2$ matrix, given by 3 numbers because it is symmetric\n",
    "Thus, the dimension of the state is 2 + 3 = 5. The first component of the state, the fuel moisture, is the quantity of interest, the rest are auxiliary.\n",
    "\n",
    "\n",
    "This can be understood as:\n",
    "\n",
    "* a mapping $M$ of the 5-dimensional hidden and external data state to a new hidden state:\n",
    "$$M:(u_{k-1},P_{k-1},e_k) \\mapsto (u_{k},P_{k})$$\n",
    "* retrieving the output (the quantity of interest) as the first component of the hiddent state\n",
    "* feeding the hiddent state back to the mapping $M$ for the next step $k+1$\n",
    "* training consists of fitting the hidden state to minimize a loss function\n",
    "$$\\ell(u_{k},P_{k},d_k,R_k) \\to \\min$$\n",
    "\n",
    "Note that in the augmented Kalman filter above, the mapping $M$ is fixed and it has a one component of the hidden state as a parameter. To get a better fit, we could increase the number of parameters, e.g., by modeling the moisture in multiple layers, as in van der Kamp et al. (2017) two-layer model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fk72YB2mjuGk"
   },
   "source": [
    "Building and evaluating RNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "svt8wUAsSA67"
   },
   "source": [
    "A recurrent neural network (RNN) has a similar information flow but it can be more flexible and look for the best model automatically, i.e., build the model from data. \n",
    "\n",
    "We'll start by how to evaluate the map, then actually create it later.\n",
    "\n",
    "Some of the code is from https://machinelearningmastery.com/understanding-simple-recurrent-neural-networks-in-keras/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lphzeucMfI7L"
   },
   "outputs": [],
   "source": [
    "# Demo example\n",
    "hidden=5\n",
    "features=2\n",
    "timesteps=3\n",
    "demo_model = create_RNN(hidden_units=hidden, dense_units=1, \n",
    "                        input_shape=(timesteps,features), \n",
    "                        activation=['linear', 'linear'])\n",
    "print(demo_model.summary())\n",
    "w = demo_model.get_weights()\n",
    "#print(len(w),' weight arrays:',w)\n",
    "wname=('wx','wh','bh','wy','by','wz','bz')\n",
    "for i in range(len(w)):\n",
    "  print(i,':',wname[i],'shape=',w[i].shape)\n",
    "wx, wh, bh, wy, by = w\n",
    "plot_model(demo_model, to_file='model_plot.png', \n",
    "  show_shapes=True, show_layer_names=True,\n",
    "  expand_nested=True,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uFVQdXL0fODX"
   },
   "source": [
    "The input layer here is just a formality. The input of the hidden layer `simple_rnn` consist of vector passed by the input layer, followed by its own output from the previous time step.\n",
    "\n",
    "Now let’s do a simple experiment to see how the layers from a SimpleRNN and Dense layer produce an output. Keep this figure in view.\n",
    "<img src=\"https://machinelearningmastery.com/wp-content/uploads/2021/09/rnnCode1.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TcxvQPo1hYip"
   },
   "source": [
    "We’ll input x for three time steps and let the network generate an output. The values of the hidden units at time steps 1, 2 and 3 will be computed. $h(0)$ is initialized to the zero vector. The output $o(3)$ is computed from $h(3)$ and $w(3)$. An activation function is linear, $f(x)=x$, so the update of  $h(k)$  and the output $o(k)$ are given by\n",
    "\\begin{align*}\n",
    "h\\left(  0\\right)  = &0  \\\\\n",
    "h\\left(  k+1\\right)  =& \n",
    "x\\left(  k\\right) w_{x}\n",
    "  +h(k) w_{h}  + b_{h}\\\\\n",
    "o(k+1)=& h(k+1)w_{y} + b_y\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fqL5TEfpml7q"
   },
   "outputs": [],
   "source": [
    "# Reshape the input to sample_size x time_steps x features \n",
    "samples=4   # number of samples\n",
    "x = tf.reshape(tf.range(samples*timesteps*features),[samples,timesteps,features]) \n",
    "print('test input x=',x)\n",
    "print('model.predict start')\n",
    "y_pred_model = demo_model.predict(x)\n",
    "print('model.predict end')\n",
    "\n",
    "o3=np.zeros([samples,1])\n",
    "for i in range(samples):\n",
    "  h_0 = np.zeros(hidden)\n",
    "  h_1 = np.dot(x[i,0,:], wx) + np.dot(h_0,wh) + bh\n",
    "  h_2 = np.dot(x[i,1,:], wx) + np.dot(h_1,wh) + bh\n",
    "  h_3 = np.dot(x[i,2,:], wx) + np.dot(h_2,wh) + bh\n",
    "  o3[i,0] = np.dot(h_3, wy) + by\n",
    "#print('h1 = ', h_1,'h2 = ', h_2,'h3 = ', h_3)\n",
    "\n",
    "print(\"Prediction from network \", y_pred_model)\n",
    "print(\"Prediction from our computation \", o3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6qdqOFCvhQL1"
   },
   "source": [
    "The result is the same."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AkyiGlZF0WrM"
   },
   "source": [
    "#### Training and forecasting with the RNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e265QFwlw22_"
   },
   "source": [
    "We are given a sequence `x` of inputs size `[train_steps+forecast_steps,features]` and want to train a model so that at step `i` in `range(train_steps)`, the model returns close to `features[i,:]`. The trained model then returns for `i` in `range(train_steps,train_steps+forecast_steps)` a forecast `features[i,:]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Kg7wSrkk-HrE"
   },
   "outputs": [],
   "source": [
    "print('test preprocessing for RNN')\n",
    "trainsteps=5\n",
    "features=1\n",
    "outputs=1\n",
    "timesteps=3\n",
    "x = tf.reshape(tf.range(trainsteps*features),[trainsteps,features])\n",
    "y = tf.reshape(tf.range(trainsteps*outputs),[trainsteps,outputs])\n",
    "print('x=',x)\n",
    "print('y=',y)\n",
    "x_train, y_train = staircase(x,y,timesteps,trainsteps)\n",
    "print('x_train=',x_train)\n",
    "print('y_train=',y_train)\n",
    "x_train, y_train = seq2batches(x,y,timesteps,trainsteps)\n",
    "print('x_train=',x_train)\n",
    "print('y_train=',y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rHWxqARpSO_f"
   },
   "outputs": [],
   "source": [
    "E,m_f,data,hour,h2,DeltaE = create_synthetic_data(days=20,power=4,data_noise=0.01,process_noise=0.0,DeltaE=0.1) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vfcxN9JCH5Ku"
   },
   "outputs": [],
   "source": [
    "scale=False\n",
    "# transform as 2D, (timesteps, features) and (timesteps, outputs)\n",
    "Et = np.reshape(E,[E.shape[0],1])\n",
    "datat = np.reshape(data,[data.shape[0],1])\n",
    "if scale:\n",
    "    scalerx = MinMaxScaler()\n",
    "    scalerx.fit(Et)\n",
    "    Et = scalerx.transform(Et)\n",
    "    scalery = MinMaxScaler()\n",
    "    scalery.fit(datat)\n",
    "    datat = scalery.transform(datat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZjJxHiEVL5sJ"
   },
   "outputs": [],
   "source": [
    "# split data\n",
    "return_sequences=False\n",
    "shift = 0.\n",
    "print('shifting inputs by',shift)\n",
    "x_train, y_train = staircase(Et+shift,datat+shift,timesteps=5,trainsteps=h2,\n",
    "                             return_sequences=return_sequences)\n",
    "print('x_train shape=',x_train.shape)\n",
    "samples, timesteps, features = x_train.shape\n",
    "print('y_train shape=',y_train.shape)\n",
    "# the simplest model possible\n",
    "activation=['linear','linear']\n",
    "hidden_units=1\n",
    "dense_units=1\n",
    "dense_layers=1\n",
    "features=1\n",
    "hours=Et.shape[0]\n",
    "h0 = tf.convert_to_tensor(datat[:samples],dtype=tf.float32)\n",
    "# print('initial state=',h0)\n",
    "# statefull model version for traning\n",
    "model_fit=create_RNN_2(hidden_units=hidden_units, \n",
    "                        dense_units=dense_units, \n",
    "                        batch_shape=(samples,timesteps,features),\n",
    "                        stateful=True,\n",
    "                        return_sequences=return_sequences,\n",
    "                        # initial_state=h0,\n",
    "                        activation=activation,\n",
    "                        dense_layers=dense_layers)\n",
    "# same model stateless for prediction on the entire dataset - to start onlg\n",
    "# the real application will switch to prediction after training data end\n",
    "# and start from the state there\n",
    "print('model_fit input shape',x_train.shape,'output shape',model_fit(x_train).shape)\n",
    "plot_model(model_fit, to_file='model_plot.png', \n",
    "           show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ipX9EJqz17Lr"
   },
   "outputs": [],
   "source": [
    "model_predict=create_RNN_2(hidden_units=hidden_units, dense_units=dense_units,  \n",
    "                        input_shape=(hours,features),stateful = False,\n",
    "                        return_sequences=True,\n",
    "                        activation=activation,dense_layers=dense_layers)\n",
    "# model_predict=create_RNN_sequences(hidden_units=1, dense_units=1, input_shape=(hours,1), \n",
    "#                        activation=['linear', 'linear'])\n",
    "print('model_predict input shape',Et.shape,'output shape',model_predict(Et).shape)\n",
    "print(model_predict.summary())\n",
    "plot_model(model_predict, to_file='model_plot.png', \n",
    "           show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dE0OHg0tGVE9"
   },
   "outputs": [],
   "source": [
    "# fitting\n",
    "w_exact=  [np.array([[1.-np.exp(-0.1)]]), np.array([[np.exp(-0.1)]]), np.array([0.]),np.array([[1.0]]),np.array([-1.*DeltaE])]\n",
    "w_initial=[np.array([[1.-np.exp(-0.1)]]), np.array([[np.exp(-0.1)]]), np.array([0.]),np.array([[1.0]]),np.array([0.*DeltaE])]\n",
    "model_fit.set_weights(w_initial)\n",
    "model_fit.fit(x_train, y_train, epochs=1000, verbose=0,batch_size=samples)\n",
    "w_fitted=model_fit.get_weights()\n",
    "for i in range(len(w)):\n",
    "    print('weight',i,' exact:',w_exact[i],':  initial:',w_initial[i],' fitted:',w_fitted[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_eval(w,title):\n",
    "    # prediction on the entire dataset from zero state\n",
    "    model_predict.set_weights(w)\n",
    "    hours=Et.shape[0]\n",
    "    print('Et.shape=',Et.shape,'hours=',hours)\n",
    "    x_input=np.reshape(Et,(1, hours, 1))\n",
    "    y_output = model_predict.predict(x_input)\n",
    "    print('x_input.shape=',x_input.shape,'y_output.shape=',y_output.shape)\n",
    "    m = np.reshape(y_output,hours) - shift\n",
    "    print('weights=',w)\n",
    "    if scale:\n",
    "        print('scaling')\n",
    "        m = scalery.inverse_transform(m)\n",
    "    m = np.reshape(m,hours)\n",
    "    plot_m(m,title=title)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R2jkoZlAIaSb"
   },
   "outputs": [],
   "source": [
    "m_fitted=model_eval(w_fitted,'RNN prediction with fitted weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bO1ewTj9gGvg"
   },
   "outputs": [],
   "source": [
    "m_exact=model_eval(w_exact,'RNN prediction with exact weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "henChC0cmbuy"
   },
   "outputs": [],
   "source": [
    "m_initial=model_eval(w_initial,'RNN prediction with initial weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PZw7DNQD4Inr"
   },
   "outputs": [],
   "source": [
    "out = np.empty((hours,1))\n",
    "w=w_exact\n",
    "h=0\n",
    "for i in range(Et.shape[0]):\n",
    "    h=np.dot(Et[i,0],w[0])+np.dot(h,w[1]) + w[2]\n",
    "    out[i]=np.dot(h,w[3]) + w[4]\n",
    "if scale:\n",
    "    print('scaling')\n",
    "    out = scalery.inverse_transform(out)\n",
    "out=np.reshape(out,hours)\n",
    "print('max abs diff',np.max(np.abs(m_exact-out)))\n",
    "plot_m(out,title='Hand computed RNN prediction with exact weights')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6uXVJj9koGF2"
   },
   "source": [
    "### 3.2 Acquisition and preprocessing of real data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q3BpOBuzc91i"
   },
   "source": [
    "Data assimilation for fuel moisture from Remote Automated Weather Stations (RAWS) was developed in Vejmelka et al. (2016). First, they use regression from all RAWS in a given area to extend the data spatially from RAWS to a grid in the whole area, then they run the extended Kalman filter at each grid node. Here, we are interested in a simplified problem: estimate future fuel moisture at a single RAWS location from weather data.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c8Y6bL1Yc91i"
   },
   "source": [
    "#### 3.2.1 Acquisition of fuel moisture observations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0CuXyWBFc91i"
   },
   "source": [
    "We try to load the data from a saved file first. If that fails, retrieve the fuel moisture data from sensors on weather stations in the Mesowest network. Get all stations with fuel moisture data in a spatial box within one hour, then pick one station and retrieve the whole time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LFrlbbMmc91i"
   },
   "outputs": [],
   "source": [
    "jfile = 'raws.json'; vars='fuel_moisture'; case = 1\n",
    "# jfile = 'raws2.json'; vars='fuel_moisture,precip_accum_one_hour'; case = 2\n",
    "def json_w(j,f):\n",
    "  print('writing json file',f)\n",
    "  json.dump(j,open(f,'w'),indent=4)\n",
    "try:\n",
    "    #! wget --no-clobber http://math.ucdenver.edu/~jmandel/data/math4779f21/raws.json\n",
    "    j = json.load(open(jfile,'r'))\n",
    "    print('loaded from ',jfile)\n",
    "    # Take the first station in the boulding box that has data between time_start and time_s2.\n",
    "    # Then retrieve data for that station between time_start and time_end\n",
    "    time_start = j['time_start']      # start of data time series\n",
    "    # time_s2    = j['time_s2']         # end of segment to read coordinates\n",
    "    time_end  = j['time_end']         # end of data time series\n",
    "    meso_ts  = j['meso_ts']           # get meso observations time series\n",
    "    obs_lon =   j['obs_lon']          # where we retrieved observations\n",
    "    obs_lat =   j['obs_lat']\n",
    "except:\n",
    "    print(\"can't read\",jfile,', creating')\n",
    "    # set up bounds\n",
    "    time_start = \"201806010800\"  # June 1 2018 08:00 in format yyyymmddHHMM\n",
    "    time_s2    = \"201806010900\"  # June 1 2018 09:00 in format yyyymmddHHMM \n",
    "    time_end   = \"201907200900\"  # June 20 2018 09:00 in format yyyymmddHHMM \n",
    "    #time_start=  \"201810230100\"\n",
    "    #time_s2=  \"201810230300\"\n",
    "    #time_end  =  \"201806022300\"\n",
    "    !pip install MesoPy\n",
    "    from MesoPy import Meso\n",
    "    bounding_box = \"-115, 38, -110, 40\"  # min longtitude, latitude\n",
    "    meso_token=\"b40cb52cbdef43ef81329b84e8fd874f\"       # you should get your own if you do more of this\n",
    "    m = Meso(meso_token)# create a Meso object\n",
    "    print('reading MesoWest fuel moisture data')\n",
    "    json_w(m.variables(),'variables.json')\n",
    "    meso_obss = m.timeseries(time_start, time_s2, bbox=bounding_box, \n",
    "                             showemptystations = '0', vars=vars)   # ask the object for data\n",
    "    json_w(meso_obss,'meso_obss.json')                        \n",
    "    # pick one station and retrieve the whole time series.\n",
    "    station=meso_obss['STATION'][0]\n",
    "    json_w(station,'station.json')\n",
    "    lon,lat = (float(station['LONGITUDE']),float(station['LATITUDE']))\n",
    "    print(station['NAME'],'station',station['STID'],'at',lon,lat)\n",
    "    e = 0.01   # tolerance\n",
    "    bb = '%s, %s, %s, %s' % (lon - e, lat - e, lon + e, lat + e)\n",
    "    print('bounding box',bb)\n",
    "    meso_ts = m.timeseries(time_start, time_end, bbox=bb, showemptystations = '0', vars=vars)   # ask the object for data\n",
    "    json_w(meso_ts,'meso_ts.json')                        \n",
    "    obs_lon, obs_lat = (lon, lat)   # remember station coordinates for later\n",
    "    j={'time_start':time_start,'time_s2':time_s2,'time_end':time_end,\n",
    "       'meso_ts':meso_ts,'obs_lon':obs_lon,'obs_lat':obs_lat}\n",
    "    json_w(j,jfile)\n",
    "    print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3bXopS3btyz0",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# process the data retrieved for this station\n",
    "# print(json.dumps(meso_ts['STATION'][0], indent=4))\n",
    "station = meso_ts['STATION'][0]\n",
    "time_str  = station['OBSERVATIONS']['date_time']\n",
    "obs_time = [datetime.strptime(t, '%Y-%m-%dT%H:%M:%SZ').replace(tzinfo=pytz.UTC) for t in time_str]\n",
    "start_time = obs_time[0].replace(minute=0)     # remember obs_time and start_time for later\n",
    "end_time = obs_time[-1]\n",
    "obs_data = np.array(station['OBSERVATIONS'][\"fuel_moisture_set_1\"])\n",
    "# obs_data = np.array(station['OBSERVATIONS'][\"fuel_moisture\"])\n",
    "# display the data retrieved\n",
    "#for o_time,o_data in zip (obs_time,obs_data):\n",
    "#    print(o_time,o_data)\n",
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(obs_data,linestyle='-',c='k',label='10-h fuel data')\n",
    "plt.title(station['STID'] + ' 10 h fuel moisture data')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Fuel moisture content (%)')\n",
    "plt.legend()\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lJQFB4rAYu9P"
   },
   "outputs": [],
   "source": [
    "# %debug"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pY4hPeATK9wZ"
   },
   "source": [
    "#### 3.2.2 Acquisition of weather data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xhyjXqxVN6B2"
   },
   "source": [
    "Our weather data are results from atmospheric models, with assimilated observations from weather stations, satellites, radars, etc. The models can be run in reanalysis mode (for the past, with data for the period modeled)  or in forecast mode (for the future, with only past data assimilated - because future data are not here yet). We use the Real-Time Mesoscale Analysis ([RTMA](https://www.nco.ncep.noaa.gov/pmb/products/rtma/)) interpolated to the RAWS location. RTMA is a real-time product, posted hourly, and available only for few days in the past. We have our own collection of selected RAWS data over past few years, obtained as a side effect of running the fuel moisture modeling software [WRFXPY](https://github.com/openwfm/wrfxpy).\n",
    "\n",
    "First try to read the data already extracted for this RAWS and staged for download."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WlqJRP8Vc91o"
   },
   "outputs": [],
   "source": [
    "jfile = 'rtma.json'\n",
    "try:\n",
    "    print('loading from ',jfile)\n",
    "    # ! wget --no-clobber http://math.ucdenver.edu/~jmandel/data/math4779f21/rtma.json\n",
    "    j = json.load(open(jfile,'r'))\n",
    "    print('loaded from ',jfile)\n",
    "    if j['obs_lat']!=obs_lat or j['obs_lon']!=obs_lon:\n",
    "      print('lon lat doesnot agree, need to load original RTMA files')\n",
    "      read_rtma=True\n",
    "    else:\n",
    "      read_rtma=False\n",
    "except:\n",
    "    print(\"can't read\",jfile,', creating')\n",
    "    read_rtma=True\n",
    "\n",
    "print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "THI6gElyHOOc"
   },
   "source": [
    "Next, functions to get the files, open as grib, and interpolate to the station coordinates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2iBNHQg5hPxB"
   },
   "source": [
    "####<font color=red>Note: If read_rtma==True, the notebook will say it crashed when run the first time. This is because it needs to install different version of some python packages and restart runtime. Simply run it again.</fonr>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "j.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mxZABVDxt0gd"
   },
   "outputs": [],
   "source": [
    "# Set up environment to read RTMA gribs\n",
    "# we will need current numpy for pygrib - needed on Colab, tensorflow is using numpy 1.19\\\n",
    "if read_rtma:\n",
    "  def load_rtma(path,file,reload=0):\n",
    "    url='http://math.ucdenver.edu/~jmandel/rtma/' + path \n",
    "    if os.path.exists(file):\n",
    "      if reload:\n",
    "        print(file + ' already exists, removing')\n",
    "        os.remove(file)\n",
    "      else:\n",
    "        print(file + ' already exists, exiting')\n",
    "        # add checking size here\n",
    "        return 0\n",
    "    try:\n",
    "      print('loading ' + url + ' as ' + file)\n",
    "      ret = subprocess.check_output(['wget','--no-clobber','--output-document='+ file, url,],stderr=subprocess.STDOUT).decode() # execute command from python strings\n",
    "      if os.path.exists(file):\n",
    "        print('loaded ' + url + ' as ' + file)\n",
    "        return 0\n",
    "      else: \n",
    "        print('file transfer completed, but the file is missing? ' + url)  \n",
    "      return 1\n",
    "    except:\n",
    "      print('file transfer failed: ' + url)\n",
    "      return 2\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dQ-uJI2sy6I3"
   },
   "source": [
    "Create a function to transfer RTMA files in GRIB2 format from the stash. The function returns zero if the file transfer succeeded. If the file is not available, it returns a nonzero value. Note: if needed, maybe in future add more sophisticated checks, check the return code of wget and if the file size is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PL3gxK67AlBI"
   },
   "outputs": [],
   "source": [
    "if read_rtma:\n",
    "  def rtma_grib(t,var):\n",
    "    tpath = '%4i%02i%02i/%02i' % (t.year, t.month, t.day, t.hour)  # remote path on server\n",
    "    tstr  = '%4i%02i%02i%02i_' % (t.year, t.month, t.day, t.hour)  # time string for local path\n",
    "    gribfile = os.path.join('data',tstr + var + '.grib')\n",
    "    remote = tpath + '/' + var + '.grib'\n",
    "    if load_rtma(remote,gribfile):\n",
    "        print('cannot load remote file',remote,'as',gribfile)\n",
    "        return []\n",
    "    else:\n",
    "        try:\n",
    "            gf=GribFile(gribfile)\n",
    "            v = np.array(gf[1].values())\n",
    "        except:\n",
    "            print('cannot read grib file',gribfile)\n",
    "            return []\n",
    "        print('loaded ',gribfile,' containing array shape ',v.shape)\n",
    "        return gf[1]   # grib message\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OY1oTYKlfd17"
   },
   "outputs": [],
   "source": [
    "if read_rtma:\n",
    "    times = pd.date_range(start=time_start,end=time_end,freq='1H')\n",
    "    varnames=['temp','td','precipa']\n",
    "    j =    read_interp_rtma(varnames,times,obs_lat,obs_lon)      # temperature\n",
    "    for varname in varnames:\n",
    "        j[varname]=j[varname].tolist() \n",
    "    j['obs_lat']=obs_lat\n",
    "    j['obs_lon']=obs_lon\n",
    "    json.dump(j,open('rtma.json','w'),indent=4)\n",
    "    print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ccp10kurAlBI"
   },
   "outputs": [],
   "source": [
    "def interp_to_lat_lon_slow(lats,lons,v,lat,lon): \n",
    "    # on mesh with coordinates lats and lons interpolate v to given lat lon\n",
    "    interp=LinearNDInterpolator(list(zip(lats.flatten(),lons.flatten())),v.flatten())\n",
    "    return interp(lat,lon)\n",
    "def interp_to_lat_lon(lats,lons,v,lat,lon):\n",
    "    # on mesh with coordinates lats and lons interpolate v to given lat lon\n",
    "    points=(np.array(range(lats.shape[0]),float),np.array(range(lats.shape[1]),float))  # uniform mesh\n",
    "    def res(ij):  # interpolation of lons lats on the uniform mesh, to noninteger coordinates   \n",
    "       return np.hstack((interpn(points,lats,ij)-lat, interpn(points,lons,ij)-lon))\n",
    "    # solve for xi,xj such that lats(xi,xj)=lat lons(xi,xj)=lon, then interpolate to (xi, xj) on uniform grid \n",
    "    result = root(res,(0,0)) # solve res(ij) = 0\n",
    "    if not result.success:\n",
    "        print(result.message)\n",
    "        exit(1)\n",
    "    return interpn(points,v,result.x) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jvnpq6S5AlBI"
   },
   "source": [
    "The interpolation function needs to  be tested."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NVMJBYI7AlBI"
   },
   "outputs": [],
   "source": [
    "def interp_to_lat_lon_test(lats,lons):\n",
    "    print('testing interp_to_lat_lon')\n",
    "    vx, vy = np.meshgrid(range(lats.shape[0]),range(lats.shape[1]),indexing='ij')\n",
    "    i, j = (1,2)\n",
    "    lat,lon = ((lats[i,j]+lats[i+1,j+1])/2,(lons[i,j]+lons[i+1,j+1])/2)\n",
    "    vi = interp_to_lat_lon(lats,lons,vx,lat,lon)\n",
    "    vj = interp_to_lat_lon(lats,lons,vy,lat,lon)\n",
    "    print(vi,vj,'should be about',i+0.5,j+0.5)\n",
    "    test_slow = 0\n",
    "    if test_slow:\n",
    "        print('Testing against the standard slow method scipy.interpolate.LinearNDInterpolator. Please wait...')\n",
    "        vi_slow = interp_to_lat_lon_slow(lats,lons,vx,lat,lon)\n",
    "        print(vi_slow)\n",
    "        vj_slow = interp_to_lat_lon_slow(lats,lons,vy,lat,lon)\n",
    "        print(vj_slow)\n",
    "        \n",
    "#gf = rtma_grib(start_time,'temp')      #  read the first grib file and use it to test interpolation\n",
    "#lats, lons = gf.latlons()\n",
    "#interp_to_lat_lon_test(lats,lons)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vt-Mk8fIc91m"
   },
   "outputs": [],
   "source": [
    "#%debug\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LQbWB_3GAlBI"
   },
   "source": [
    "Now we are ready for a function to read the RTMA files and interpolate to the station coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b3JJH3XPAlBI"
   },
   "outputs": [],
   "source": [
    "if read_rtma:\n",
    "  def read_interp_rtma(varnames,times,lat,lon):\n",
    "    # read RTMA from start_time to end_time and interpolate to obs_lat obs_lon\n",
    "    ntimes = len(times)\n",
    "    time_str = 'time_str'\n",
    "    j={time_str:times.strftime('%Y-%m-%d %H:%M').tolist()}\n",
    "    for varname in varnames:\n",
    "        j[varname]=np.full(ntimes,np.nan)  # initialize array of nans as list\n",
    "    n=0\n",
    "    for t in times:\n",
    "        tim=t.strftime('%Y-%m-%d %H:%M')\n",
    "        should_be = j[time_str][n]\n",
    "        if tim != should_be:\n",
    "            print('n=',n,'time',tim,'expected',should_be)\n",
    "            raise 'Invalid time' \n",
    "        for varname in varnames:\n",
    "            gf = rtma_grib(t,varname)   # read and create grib object, download if needed\n",
    "            if gf:\n",
    "                lats,lons = gf.latlons()    # coordinates\n",
    "                v = gf.values()\n",
    "                vi=interp_to_lat_lon(lats,lons,v,lat,lon) # append to array\n",
    "                print(varname,'at',t,'interpolated to',lat,lon,' value ',vi)\n",
    "                j[varname][n] = vi\n",
    "            else:\n",
    "                print(varname,'at',t,' could not be loaded')\n",
    "        n = n+1\n",
    "    return j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bMpYIZT6c91o"
   },
   "outputs": [],
   "source": [
    "# %debug\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KVXBjGA0CiXr"
   },
   "source": [
    "#### 3.2.3 Preprocessing and visualization of the weather data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fNA3Vbo1c91o"
   },
   "outputs": [],
   "source": [
    "rtma = j\n",
    "td = np.array(rtma['td'])\n",
    "t2 = np.array(rtma['temp'])\n",
    "rain=np.array(rtma['precipa'])\n",
    "# compute relative humidity\n",
    "rh = 100*np.exp(17.625*243.04*(td - t2) / (243.04 + t2 - 273.15) / (243.0 + td - 273.15))\n",
    "Ed = 0.924*rh**0.679 + 0.000499*np.exp(0.1*rh) + 0.18*(21.1 + 273.15 - t2)*(1 - np.exp(-0.115*rh))\n",
    "Ew = 0.618*rh**0.753 + 0.000454*np.exp(0.1*rh) + 0.18*(21.1 + 273.15 - t2)*(1 - np.exp(-0.115*rh))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tZIK59bJAlBJ"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(t2,linestyle='-',c='k',label='Temperature')\n",
    "plt.title(station['STID'] + ' Temperature')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Temperature (K)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LbyqcuXYAlBJ"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(td,linestyle='-',c='k',label='Dew point')\n",
    "plt.title(station['STID'] + ' Dew point (K)')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Dew point (K)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dfoOK2kSc91p"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(rh,linestyle='-',c='k',label='Dew point')\n",
    "plt.title(station['STID'] + ' relative humidity')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Relative humidity (%)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MWTJ5b2kc91p"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(Ed,linestyle='-',c='r',label='drying equilibrium')\n",
    "plt.plot(Ew,linestyle=':',c='b',label='wetting equilibrium')\n",
    "plt.title(station['STID'] + ' drying and wetting equilibria')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Fuel moisture contents (%)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jY3_eeBRc91p"
   },
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PQKSRvRSAlBJ"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(rain,linestyle='-',c='k',label='Precipitation')\n",
    "plt.title(station['STID'] + ' Precipitation' )\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Precipitation (mm/hour)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Dwbt4UXfro5x"
   },
   "outputs": [],
   "source": [
    "print(rain[1900:2000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_yRu_7WvHc6P"
   },
   "source": [
    "Precipitation from RTMA is in kg/m${}^2$. 1m water depth over 1m${}^2$ is 1m${}^3$ with mass 1000 kg thus 1 kg/m${}^2$ is the same as 1 mm of precipitation. RTMA values are accumulations over 1 h so these are values in mm/h. So 9999 mm/h = 10m/h makes no sense. Replace anything over 1m/h by nan and try again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XPYO_Iuvc91q"
   },
   "outputs": [],
   "source": [
    "rain[rain > 1000] = np.NaN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GYWTfbBBc91q",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(rain,linestyle='-',c='k',label='Precipitation')\n",
    "plt.title(station['STID'] + ' Precipitation' )\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Precipitation (mm/hour)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q_L0R2Njc91q"
   },
   "source": [
    "Fix some missing data, then we can use the data for up to 1942 hours until a biger gap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_tkU7UJic91q"
   },
   "outputs": [],
   "source": [
    "# fix isolated nans\n",
    "def fixnan(a,n):\n",
    "    for c in range(n):\n",
    "        for i in np.where(np.isnan(a)):\n",
    "            a[i]=0.5*(a[i-1]+a[i+1])\n",
    "        if not any(np.isnan(a)):\n",
    "            break\n",
    "    return a\n",
    "\n",
    "rain=fixnan(rain,2)\n",
    "t2=fixnan(t2,2)\n",
    "rh=fixnan(rh,2)\n",
    "obs_data=fixnan(obs_data,2)\n",
    "Ed=fixnan(Ed,2)\n",
    "Ew=fixnan(Ew,2)\n",
    "\n",
    "print(np.where(np.isnan(rain)))\n",
    "print(np.where(np.isnan(t2)))\n",
    "print(np.where(np.isnan(rh)))\n",
    "print(np.where(np.isnan(obs_data)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XqQYnyI9DIy1"
   },
   "source": [
    "## 4 Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2tIC_Tqnc91r"
   },
   "source": [
    "### 4.1 Kalman filter with fuel moisture observations, followed by forecasting\n",
    "We run the model first with Kalman filter for 150 hours. The observations are the RAWS data\n",
    "After 150 hours, we run in forecast mode - the RAWS data are no longer used, and we run the model from the weather data without the Kalman filter. The weather data are taken to be RTMA interpolated to one RAWS location.\n",
    "In a real forecasting application, the model would be run from weather forecast rather than data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aXnSQM7wc91r"
   },
   "outputs": [],
   "source": [
    "# run KF on an initial data seqment\n",
    "\n",
    "hours=1200 # total simulation\n",
    "h2 = 300\n",
    "m = np.zeros(hours) # preallocate\n",
    "m[0]= obs_data[0]             # initial state  \n",
    "P = np.zeros(hours)\n",
    "P[0] = 1e-3 # background state variance\n",
    "H = np.array([1.])   # all oQ = np.array([0.02]) # process noise variancebserved\n",
    "Q = np.array([1e-3]) # process noise variance\n",
    "R = np.array([1e-3]) # data variance\n",
    "for t in range(hours-1):\n",
    "    # using lambda construction to pass additional arguments to the model \n",
    "    if t < h2 and not np.isnan(obs_data[t]) and not np.isnan(Ew[t]) and not np.isnan(rain[t]): # advance model and run KF\n",
    "        m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_moisture(u,Ed[t],Ew[t],rain[t],t,partials=1),Q,\n",
    "                    d=obs_data[t],H=H,R=R)\n",
    "    else:  # just advance to next hour, no process noise\n",
    "        m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_moisture(u,Ed[t],Ew[t],rain[t],t,partials=1),Q*0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "peMi-OF3c91r",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(Ed[:hours],linestyle='--',c='r',label='Drying Equilibrium')\n",
    "plt.plot(Ew[:hours],linestyle='--',c='b',label='Wetting Equilibrium')\n",
    "plt.plot(obs_data[:hours],linestyle=':',c='k',label='RAWS data')\n",
    "plt.plot(m[:h2],linestyle='-',c='k',label='filtered')\n",
    "plt.plot(range(h2,hours),m[h2:hours],linestyle='-',c='r',label='forecast')\n",
    "plt.title(station['STID'] + ' Kalman filtering and forecast with real data')\n",
    "plt.xlabel('Time (hours)') \n",
    "plt.ylabel('Fuel moisture content (%)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3TnwXYcLc91r"
   },
   "source": [
    "Clearly, there is a problem - the forecast fuel moisture is too high. We need to assimilate also some parameters of the model, not just its output state. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8SuVNg8TsW4d"
   },
   "source": [
    "### 4.3 Kalman filter on the augmented model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FYAbWNCfk2wD"
   },
   "source": [
    "Run augmented filter and plot the result:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q3NHr3oBsDg6"
   },
   "outputs": [],
   "source": [
    "m,Ec = run_augmented_kf(obs_data,Ed,Ew,rain,h2,hours)  # extract from state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hlkby_oTlB_f"
   },
   "outputs": [],
   "source": [
    "def plot_moisture(hmin,hmax):\n",
    "  print('training from 0 to',h2,'plot from',hmin,'to',hmax)\n",
    "  plt.figure(figsize=(16,4))\n",
    "  plt.plot(range(hmin,hmax),Ed[hmin:hmax],linestyle='--',c='r',label='Drying Equilibrium (%)')\n",
    "  plt.plot(range(hmin,hmax),Ew[hmin:hmax],linestyle='--',c='b',label='Wetting Equilibrium (%)')\n",
    "  plt.plot(range(hmin,hmax),Ec[hmin:hmax],linestyle='--',c='g',label='Equilibrium Correction (%)')\n",
    "  plt.plot(range(hmin,hmax),m[hmin:hmax],linestyle='-',c='k',label='filtered')\n",
    "  plt.plot(range(hmin,hmax),obs_data[hmin:hmax],linestyle='-',c='b',label='RAWS data (%)')\n",
    "  plt.plot(range(hmin,hmax),rain[hmin:hmax],linestyle='-',c='b',label='RTMA rain (mm/h)')\n",
    "  if hmin>=h2:\n",
    "    plt.plot(m[hmin:h2],linestyle='-',c='k',label='Filtered')\n",
    "  h1 = np.maximum(hmin,h2)\n",
    "  plt.plot(range(h1,hmax),m[h1:hmax],linestyle='-',c='r',label='Forecast (%)')\n",
    "  plt.title(station['STID'] +' Kalman filtering and forecast with augmented state, real data. Training 0:%i hmax' % h2)\n",
    "  plt.xlabel('Time (hours)') \n",
    "  plt.ylabel('Fuel moisture content (%)')\n",
    "  plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q-h5omKgnow2"
   },
   "outputs": [],
   "source": [
    "plot_moisture(0,hours)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0w0YtHtqnza5"
   },
   "source": [
    "A detailed view of transition from training to forecast:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B7sXGUotc91s"
   },
   "outputs": [],
   "source": [
    "plot_moisture(0,600)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xy7sIs0z_Kk6"
   },
   "outputs": [],
   "source": [
    "plot_moisture(300,500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y-C6IRFVxGUR"
   },
   "outputs": [],
   "source": [
    "plot_moisture(300,800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TvlCtT0X2ejp"
   },
   "outputs": [],
   "source": [
    "plot_moisture(800,1200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7W03QTo3c91t"
   },
   "source": [
    "Filtering by extended Kalman filter using RAWS data until 150 hours, then forecasting mode - running the model from interpolated RTMA only. For the first 60 hours the forecast is good, the equilibium correction made the model quite close to data. But then the big spike in equilibrium moisture around 230 hours attracted the solution, and it took a while for it to get back. The spike in the RAWS measurement is there but much smaller. The model becomes inaccurate during periods when the fuel moisture equilibrium is large.\n",
    "\n",
    "Possible reasons include: 1. There was something in the data we do not know about - maybe it rained but RTMA did not tell us. Try comparing with data from the RAWS itself? 2. The model is too simple, assumes the whole depth of the wood stick is wetting and drying at the same time. Perhaps the moisture got stored in the inside layers of the measurement stick. Try a two-layer model as in van der Kamp (2017) and make the state larger? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "owEI4EtTo7Ek"
   },
   "source": [
    "A detailed view of rain episode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C_hoDjgtpMEJ"
   },
   "outputs": [],
   "source": [
    "plot_moisture(900,1100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DRraWhwdpSkV"
   },
   "source": [
    "It seems there is some rain that the model does not know about."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1STfnlT40rPX"
   },
   "source": [
    "## RNN for real data, no rain yet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3cwY43iSnQ0t"
   },
   "source": [
    "#### Linear modeling by RELU - potential for generalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MotzNBvOnFvC"
   },
   "outputs": [],
   "source": [
    "def RELU(x):\n",
    "  if x>0. :\n",
    "    return x\n",
    "  else:\n",
    "    return 0.\n",
    "\n",
    "# network computing z = a*x1 + b*x2 with offset c\n",
    "def linrelu(x,a,b,c):\n",
    "  y = np.dot(np.array([[a, b], [-a, -b] ]), x) + np.array([c, -c])\n",
    "  y[0]=RELU(y[0])\n",
    "  y[1]=RELU(y[1])\n",
    "  return(np.dot([1,-1],y))-c\n",
    "x = [1,2]\n",
    "a = 2\n",
    "b = 4\n",
    "c = 3\n",
    "print(a*x[0]+b*x[1])\n",
    "linrelu(x,a,b,c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-p6dcLua_udD"
   },
   "source": [
    "### Basic RNN on real data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set seed for reproducibility\n",
    "tf.random.set_seed(123)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gSmbDPZIHbTr"
   },
   "source": [
    "Try with E average between drying and wetting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ymhNMZkoHfCl"
   },
   "outputs": [],
   "source": [
    "E = (Ed + Ew)/2\n",
    "print(Ed.shape,Ew.shape,rain.shape)\n",
    "first_rain=np.nonzero(rain>0)[0][0]\n",
    "print(first_rain)\n",
    "hours=first_rain\n",
    "E=E[:hours]\n",
    "data=obs_data[:hours]\n",
    "scale=False\n",
    "\n",
    "# transform as 2D, (timesteps, features) and (timesteps, outputs)\n",
    "Et = np.reshape(E,[E.shape[0],1])\n",
    "datat = np.reshape(data,[data.shape[0],1])\n",
    "if scale:\n",
    "  scalerx = MinMaxScaler()\n",
    "  scalerx.fit(Et)\n",
    "  Et = scalerx.transform(Et)\n",
    "  scalery = MinMaxScaler()\n",
    "  scalery.fit(datat)\n",
    "  datat = scalery.transform(datat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DPcxv85XILdn"
   },
   "source": [
    "Create the model again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gEkbHZSqIOq1"
   },
   "outputs": [],
   "source": [
    "# split data\n",
    "return_sequences=False\n",
    "x_train, y_train = staircase(Et,datat,timesteps=5,trainsteps=h2,\n",
    "                             return_sequences=return_sequences)\n",
    "print('x_train shape=',x_train.shape)\n",
    "samples, timesteps, features = x_train.shape\n",
    "print('y_train shape=',y_train.shape)\n",
    "# the simplest model possible\n",
    "activation=['linear','linear']\n",
    "hidden_units=3\n",
    "dense_units=1\n",
    "dense_layers=1\n",
    "features=1\n",
    "hours=Et.shape[0]\n",
    "h0 = tf.convert_to_tensor(datat[:samples],dtype=tf.float32)\n",
    "# print('initial state=',h0)\n",
    "# statefull model version for traning\n",
    "model_fit=create_RNN_2(hidden_units=hidden_units, \n",
    "                        dense_units=dense_units, \n",
    "                        batch_shape=(samples,timesteps,features),\n",
    "                        stateful=True,\n",
    "                        return_sequences=return_sequences,\n",
    "                        # initial_state=h0,\n",
    "                        activation=activation,\n",
    "                        dense_layers=dense_layers)\n",
    "# same model stateless for prediction on the entire dataset - to start onlg\n",
    "# the real application will switch to prediction after training data end\n",
    "# and start from the state there\n",
    "print('model_fit input shape',x_train.shape,'output shape',model_fit(x_train).shape)\n",
    "\n",
    "plot_model(model_fit, to_file='model_plot.png', \n",
    "           show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jtFJQu33NqfL"
   },
   "outputs": [],
   "source": [
    "model_predict=create_RNN_2(hidden_units=hidden_units, dense_units=dense_units,  \n",
    "                        input_shape=(hours,features),stateful = False,\n",
    "                        return_sequences=True,\n",
    "                        activation=activation,dense_layers=dense_layers)\n",
    "# model_predict=create_RNN_sequences(hidden_units=1, dense_units=1, input_shape=(hours,1), \n",
    "#                        activation=['linear', 'linear'])\n",
    "print('model_predict input shape',Et.shape,'output shape',model_predict(Et).shape)\n",
    "print(model_predict.summary())\n",
    "from keras.utils.vis_utils import plot_model\n",
    "plot_model(model_predict, to_file='model_plot.png', \n",
    "           show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wuxh5pq0OMSa"
   },
   "outputs": [],
   "source": [
    "# fitting\n",
    "DeltaE = 0\n",
    "w_exact=  [np.array([[1.-np.exp(-0.1)]]), np.array([[np.exp(-0.1)]]), np.array([0.]),np.array([[1.0]]),np.array([-1.*DeltaE])]\n",
    "w_initial=[np.array([[1.-np.exp(-0.1)]]), np.array([[np.exp(-0.1)]]), np.array([0.]),np.array([[1.0]]),np.array([-1.0])]\n",
    "w=model_fit.get_weights()\n",
    "for i in range(len(w)):\n",
    "  print('weight',i,'shape',w[i].shape,'ndim',w[i].ndim,'given',w_initial[i].shape)\n",
    "  for j in range(w[i].shape[0]):\n",
    "    if w[i].ndim==2:\n",
    "      for k in range(w[i].shape[1]):\n",
    "        w[i][j][k]=w_initial[i][0][0]/w[i].shape[0]\n",
    "    else:\n",
    "      w[i][j]=w_initial[i][0]\n",
    "model_fit.set_weights(w)\n",
    "model_fit.fit(x_train, y_train, epochs=5000, verbose=0,batch_size=samples)\n",
    "w_fitted=model_fit.get_weights()\n",
    "for i in range(len(w)):\n",
    "  print('weight',i,' exact:',w_exact[i],':  initial:',w_initial[i],' fitted:',w_fitted[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uJz1EgPyRTEH"
   },
   "outputs": [],
   "source": [
    "# evaluate model\n",
    "model_predict.set_weights(w_fitted)\n",
    "x_input=np.reshape(Et,(1, hours, 1))\n",
    "y_output = model_predict.predict(x_input)\n",
    "print('x_input.shape=',x_input.shape,'y_output.shape=',y_output.shape)\n",
    "print(shift)\n",
    "m = np.reshape(y_output,hours)\n",
    "print('weights=',w)\n",
    "if scale:\n",
    "    print('scaling')\n",
    "    m = scalery.inverse_transform(m)\n",
    "m = np.reshape(m,hours)\n",
    "hour=np.array(range(hours))\n",
    "title=\"First RNN forecast\"\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(hour,E,linestyle='--',c='r',label='E=Equilibrium data')\n",
    "# print(len(hour),len(m_f))\n",
    "plt.scatter(hour,data,c='b',label='data=10-h fuel data')\n",
    "if m is not None:\n",
    "    plt.plot(hour[:h2],m[:h2],linestyle='-',c='k',label='m=filtered')\n",
    "    plt.plot(hour[h2:hours],m[h2:hours],linestyle='-',c='r',label='m=forecast')\n",
    "plt.title(title) \n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VSwtgKPJPnH4"
   },
   "outputs": [],
   "source": [
    "# plot subinterval only\n",
    "def plot_int(lb=0,ub=hours,title=\"RNN forecast\"):\n",
    "  hour=np.array(range(hours))\n",
    "  plt.figure(figsize=(16,4))\n",
    "  plt.plot(hour[lb:ub],E[lb:ub],linestyle='--',c='r',label='Equilibrium data')\n",
    "  plt.scatter(hour[lb:ub],data[lb:ub],c='b',label='data=10-h fuel data')\n",
    "  if lb <= h2:\n",
    "    ub1 = min(h2,ub)\n",
    "    plt.plot(hour[lb:ub1],m[lb:ub1],linestyle='-',c='k',label='filtered')\n",
    "  if ub >= h2:\n",
    "    lb1 = max(h2,lb)\n",
    "    plt.plot(hour[lb1:ub],m[lb1:ub],linestyle='-',c='r',label='forecast')\n",
    "  plt.title(title) \n",
    "  plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vCjk9hZtkFym"
   },
   "outputs": [],
   "source": [
    "plot_int()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Sd3fDOnvmmdp"
   },
   "outputs": [],
   "source": [
    "plot_int(0,300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vHkc4KHdkAJp"
   },
   "outputs": [],
   "source": [
    "plot_int(300,500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Km5VWhcJlyvV"
   },
   "outputs": [],
   "source": [
    "plot_int(500,800)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TBayRudFcZWP"
   },
   "source": [
    "Next step: two features - drying and wetting equilibria"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SGbgxOm_kEc4"
   },
   "outputs": [],
   "source": [
    "print(Ed.shape,Ew.shape,rain.shape)\n",
    "first_rain=np.nonzero(rain>0)[0][0]\n",
    "print(first_rain)\n",
    "hours=first_rain\n",
    "Ed=Ed[:hours]\n",
    "Ew=Ew[:hours]\n",
    "h2 = 300\n",
    "# print(Ed.shape,Ew.shape)\n",
    "# (timesteps, features)\n",
    "Et = np.vstack((Ed, Ew)).T\n",
    "print(E.shape)\n",
    "data=obs_data[:hours]\n",
    "\n",
    "scale=False\n",
    "\n",
    "# transform as 2D, (timesteps, features) and (timesteps, outputs)\n",
    "datat = np.reshape(data,[data.shape[0],1])\n",
    "if scale:\n",
    "  scalerx = MinMaxScaler()\n",
    "  scalerx.fit(Et)\n",
    "  Et = scalerx.transform(Et)\n",
    "  scalery = MinMaxScaler()\n",
    "  scalery.fit(datat)\n",
    "  datat = scalery.transform(datat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b6aJAvBEkEBl"
   },
   "outputs": [],
   "source": [
    "# split data\n",
    "return_sequences=False\n",
    "x_train, y_train = staircase(Et,datat,timesteps=5,trainsteps=h2,\n",
    "                             return_sequences=return_sequences)\n",
    "print('x_train shape=',x_train.shape)\n",
    "samples, timesteps, features = x_train.shape\n",
    "print('y_train shape=',y_train.shape)\n",
    "# the simplest model possible\n",
    "activation=['linear','linear']\n",
    "hidden_units=6\n",
    "dense_units=1\n",
    "dense_layers=1\n",
    "features=Et.shape[1]\n",
    "hours=Et.shape[0]\n",
    "h0 = tf.convert_to_tensor(datat[:samples],dtype=tf.float32)\n",
    "# print('initial state=',h0)\n",
    "# statefull model version for traning\n",
    "model_fit=create_RNN_2(hidden_units=hidden_units, \n",
    "                        dense_units=dense_units, \n",
    "                        batch_shape=(samples,timesteps,features),\n",
    "                        stateful=True,\n",
    "                        return_sequences=return_sequences,\n",
    "                        # initial_state=h0,\n",
    "                        activation=activation,\n",
    "                        dense_layers=dense_layers)\n",
    "# same model stateless for prediction on the entire dataset - to start onlg\n",
    "# the real application will switch to prediction after training data end\n",
    "# and start from the state there\n",
    "print('model_fit input shape',x_train.shape,'output shape',model_fit(x_train).shape)\n",
    "\n",
    "plot_model(model_fit, to_file='model_plot.png', \n",
    "           show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ClBMYe8Lqr7P"
   },
   "outputs": [],
   "source": [
    "model_predict=create_RNN_2(hidden_units=hidden_units, dense_units=dense_units,  \n",
    "                        input_shape=(hours,features),stateful = False,\n",
    "                        return_sequences=True,\n",
    "                        activation=activation,dense_layers=dense_layers)\n",
    "# model_predict=create_RNN_sequences(hidden_units=1, dense_units=1, input_shape=(hours,1), \n",
    "#                        activation=['linear', 'linear'])\n",
    "# print('model_predict input shape',Et.shape,'output shape',model_predict(Et).shape)\n",
    "print(model_predict.summary())\n",
    "\n",
    "plot_model(model_predict, to_file='model_plot.png', \n",
    "           show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4U0kTEiksNZs"
   },
   "outputs": [],
   "source": [
    "w_initial=[np.array([[1.-np.exp(-0.1)]]), np.array([[np.exp(-0.1)]]), np.array([0.]),np.array([[1.0]]),np.array([-1.0])]\n",
    "w=model_fit.get_weights()\n",
    "for i in range(len(w)):\n",
    "  print('weight',i,'shape',w[i].shape,'ndim',w[i].ndim,'given',w_initial[i].shape)\n",
    "  for j in range(w[i].shape[0]):\n",
    "    if w[i].ndim==2:\n",
    "      for k in range(w[i].shape[1]):\n",
    "        w[i][j][k]=w_initial[i][0][0]/w[i].shape[0]\n",
    "    else:\n",
    "      w[i][j]=w_initial[i][0]\n",
    "model_fit.set_weights(w)\n",
    "model_fit.fit(x_train, y_train, epochs=5000, verbose=0,batch_size=samples)\n",
    "w_fitted=model_fit.get_weights()\n",
    "for i in range(len(w)):\n",
    "  print('weight',i,' exact:',w_exact[i],':  initial:',w_initial[i],' fitted:',w_fitted[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o10lIOl4sndv"
   },
   "outputs": [],
   "source": [
    "# evaluate model\n",
    "model_predict.set_weights(w_fitted)\n",
    "x_input=np.reshape(Et,(1, hours, 2))\n",
    "y_output = model_predict(x_input)\n",
    "print('x_input.shape=',x_input.shape,'y_output.shape=',y_output.shape)\n",
    "print(shift)\n",
    "m = np.reshape(y_output,hours)\n",
    "print('weights=',w)\n",
    "if scale:\n",
    "    print('scaling')\n",
    "    m = scalery.inverse_transform(m)\n",
    "m = np.reshape(m,hours)\n",
    "hour=np.array(range(hours))\n",
    "title=\"First RNN forecast\"\n",
    "plt.figure(figsize=(16,4))\n",
    "plt.plot(hour,Ed,linestyle='--',c='r',label='Drying equilibrium')\n",
    "plt.plot(hour,Ew,linestyle='--',c='b',label='Wetting equilibrium')\n",
    "# print(len(hour),len(m_f))\n",
    "plt.scatter(hour,data,c='b',label='data=10-h fuel data')\n",
    "if m is not None:\n",
    "    plt.plot(hour[:h2],m[:h2],linestyle='-',c='k',label='m=filtered')\n",
    "    plt.plot(hour[h2:hours],m[h2:hours],linestyle='-',c='r',label='m=forecast')\n",
    "plt.title(title) \n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mrWioCJVuU-G"
   },
   "outputs": [],
   "source": [
    "# plot subinterval only\n",
    "def plot_int(lb=0,ub=hours,title=\"RNN Prediction\"):\n",
    "  hour=np.array(range(hours))\n",
    "  plt.figure(figsize=(16,4))\n",
    "  plt.plot(hour[lb:ub],Ed[lb:ub],linestyle='--',c='r',label='Drying equilibrium')\n",
    "  plt.plot(hour[lb:ub],Ew[lb:ub],linestyle='--',c='b',label='Wetting equilibrium')\n",
    "  plt.plot(hour[lb:ub],data[lb:ub],linestyle='-',c='b',label='RAWS fuel moisture data')\n",
    "  if lb <= h2:\n",
    "    ub1 = min(h2,ub)\n",
    "    plt.plot(hour[lb:ub1],m[lb:ub1],linestyle='-',c='k',label='Fuel moisture fitted')\n",
    "  if ub >= h2:\n",
    "    lb1 = max(h2,lb)\n",
    "    plt.plot(hour[lb1:ub],m[lb1:ub],linestyle='-',c='r',label='Fuel moisture prediction')\n",
    "  plt.title(title) \n",
    "  plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qmGPeG61uqGI"
   },
   "outputs": [],
   "source": [
    "plot_int(0,600,title='RNN fitting and prediction')  # again the whole thing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SwnOSJlOuvAA"
   },
   "outputs": [],
   "source": [
    "plot_int(0,300,title='RNN Fitting') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EqCZD7uCvDrS"
   },
   "outputs": [],
   "source": [
    "plot_int(300,500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hYgLAXpUvSLo"
   },
   "outputs": [],
   "source": [
    "plot_int(500,800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(m.size)\n",
    "print(m.mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gVQxv9Blc91t"
   },
   "source": [
    "### 4.4 A comment on the information flow in the Kalman filter and in neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_g_OTEg6ePb9"
   },
   "source": [
    "## 5. Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aNxw7xI3FqFt"
   },
   "source": [
    "We have shown how to combine a model and data for improved forecasting of fuel moisture from weather forecast using the Kalman filter. Augmenting the filter state by a model parameter and joint estimation of augmented state resulted in an improvement of the forecast."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IWpmDwUPGElR"
   },
   "source": [
    "## Contributions of authors "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jujW1VFgGOCn"
   },
   "source": [
    "Not applicable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HWslw7HmGZmP"
   },
   "source": [
    "## Acknowledgements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xubqDAV2GjkZ"
   },
   "source": [
    "This Math Clinic was sponsored by the team of investigators of the NASA grant no. 80NSSC19K1091 *Coupled Interactive Forecasting of Weather, Fire Behavior, and Smoke Impact for Improved Wildland Fire Decision Making* under the NASA ROSES18 Disasters program. The author would like to thank Brian Zhang from the Math Clinic class for bringing the reference van der Kamp et al. (2017) to his attention."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZsNZxOv7c91t"
   },
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vFY-iS1Wc91t"
   },
   "source": [
    "J. Mandel, S. Amram, J. D. Beezley, G. Kelman, A. K. Kochanski, V. Y. Kondratenko, B. H. Lynn, B. Regev, and M. Vejmelka. *Recent advances and applications of WRF-SFIRE.* Natural Hazards and Earth System Science, 14(10):2829–2845, 2014. [doi:10.5194/nhessd-2-1759-2014](https://doi.org/10.5194/nhessd-2-1759-2014)\n",
    "\n",
    "R. E. Kalman. *A new approach to linear filtering and prediction problems.* Transactions of the ASME – Journal of Basic Engineering, Series D, 82:35–45, 1960. [doi:10.1115/1.3662552](https://doi.org/10.1115/1.3662552)\n",
    "\n",
    "E. Kalnay. *Atmospheric Modeling, Data Assimilation and Predictability.* Cambridge University Press, 2003. [doi:10.1017/CBO9780511802270](https://doi.org/10.1017/CBO9780511802270)\n",
    "\n",
    "D. W. van der Kamp, R. D. Moore, and I. G. McKendry. *A model for simulating the moisture content of standardized fuel sticks of various sizes.* Agricultural and Forest Meteorology, 236:123–134, 2017. [doi:10.1016/j.agrformet.2017.01.013](https://doi.org/10.1016/j.agrformet.2017.01.013)\n",
    "\n",
    "S. F. Schmidt. *Application of state-space methods to navigation problems.* volume 3 of Advances in Control Systems, C. T.  Leondes, ed., pages 293–340. Elsevier, 1966. [doi:10.1016/B978-1-4831-6716-9.50011-4](https://doi.org/10.1016/B978-1-4831-6716-9.50011-4)\n",
    "\n",
    "M. Vejmelka, A. K. Kochanski, and J. Mandel. *Data assimilation of dead fuel moisture observations from remote automatic weather stations.* International Journal of Wildland Fire, 25:558– 568, 2016. [doi:10.1071/WF14085](https://doi.org/10.1071/WF14085)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "gpuClass": "standard",
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}