daml-part2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "%%HTML\n",
    "<link rel=\"stylesheet\" href=\"rise.css\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# <span style=\"color:teal\"> Data assimilation, machine learning and dynamical systems - Part II </span>\n",
    "#### Marc Bocquet¹ [marc.bocquet@enpc.fr](mailto:marc.bocquet@enpc.fr) and Alban Farchi¹ [alban.farchi@enpc.fr](mailto:alban.farchi@enpc.fr)\n",
    "#### (1) CEREA, École des Ponts and EdF R&D, IPSL, Île-de-France, France\n",
    "\n",
    "During this session, we will apply standard machine learning methods to learn the dynamics of the Lorenz 1996 model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:green\"> The Lorenz 1996 model </span>\n",
    "\n",
    "The Lorenz 1996 (L96, [Lorenz and Emanuel 1998](https://journals.ametsoc.org/view/journals/atsc/55/3/1520-0469_1998_055_0399_osfswo_2.0.co_2.xml)) is a low-order chaotic model commonly used in data assimilation to asses the performance of new algorithms. It represents the evolution of some unspecified scalar meteorological quantity (perhaps vorticity or temperature) over a latitude circle.\n",
    "\n",
    "The model **dynamics** is driven by the following set of ordinary differential equations (ODEs):\n",
    "$$\n",
    "    \\forall n \\in [1, N_{x}], \\quad \\frac{\\mathrm{d}x_{n}}{\\mathrm{d}t} =\n",
    "    (x_{n+1}-x_{n-2})x_{n-1}-x_{n}+F,\n",
    "$$\n",
    "where the indices are periodic: $x_{-1}=x_{N_{x}-1}$, $x_{0}=x_{N_{x}}$, and $x_{1}=x_{N_{x}+1}$, and where the system size $N_{x}$ can take arbitrary values.\n",
    "\n",
    "In the standard configuration, $N_{x}=40$ and the forcing coefficient is $F=8$. The ODEs are integrated using a fourth-order Runge-Kutta scheme with a time step of $0.05$ model time unit (MTU). The resulting dynamics is **chaotic** with a doubling time of errors around $0.42$ MTU (the corresponding Lyapunov is hence $0.61$ MTU). For comparison, $0.05$ MTU represent six hours of real time and correspond to an average autocorrelation around $0.967$. Finally, the model variability (spatial average of the time standard deviation per variable) is $3.64$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## <span style=\"color:green\"> The true model dynamics </span>\n",
    "\n",
    "In this series of experiments, we will try to emulate the dynamics of the L96 model using artificial neural networks (NN).\n",
    "1. We start by running the **true model** to build a training dataset.\n",
    "2. We build and **train a NN** using this dataset.\n",
    "3. We evaluate the **forecast skill** of the surrogate model (the NN)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Importing all modules </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tqdm.auto import tqdm, trange\n",
    "\n",
    "import utils\n",
    "\n",
    "utils.set_style()\n",
    "seeds = [3, 31, 314, 3141, 31415, 314159, 3141592, 31415926]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Defining the neural network model </span>\n",
    "\n",
    "In the following cell, we define the true Lorenz 1996 model using standard values: \n",
    "- the number of variables $N_{x}$ is set to `Nx=40`;\n",
    "- the forcing coefficient $F$ is set to `F=8`;\n",
    "- the integration time step is set to `dt=0.05`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create model\n",
    "true_model = utils.Lorenz1996Model(Nx=40, dt=0.05, F=8)\n",
    "\n",
    "# save some statistics about the model\n",
    "true_model.model_var = 3.64\n",
    "true_model.doubling_time = 0.42\n",
    "true_model.lyap_time = 0.61"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Short model integration </span>\n",
    "\n",
    "In the following cells, we perform a rather short model integration, in order to illustrate the model dynamics. The initial condition is a random field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define rng\n",
    "rng = np.random.default_rng(seed=seeds.pop(0))\n",
    "\n",
    "# allocate memory\n",
    "Nt_plot = 500\n",
    "xt_plot = np.zeros((Nt_plot+1, true_model.Nx))\n",
    "\n",
    "# initialisation and integrate\n",
    "xt_plot[0] = rng.normal(loc=3, scale=1, size=true_model.Nx)\n",
    "for t in trange(Nt_plot, desc='model integration'):\n",
    "    xt_plot[t+1] = true_model.forward(xt_plot[t])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_traj(\n",
    "    xt_plot, \n",
    "    true_model,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see first a spin-up period of about $1$ MTU, where the initial condition is progressively forgotten and the trajectory progressively gets back to the model attractor. After this spin-up period, the dynamics is characterised by waves moving slowly towards the east (i.e. decreasing variable index). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:green\"> Prepare the dataset </span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> A long model integration for the training data</span>\n",
    "\n",
    "We now use a true model trajectory to make the **training dataset**. This trajectory starts from a random field (different than for the plotting trajectory) and we discard the first $100$ time steps to get rid of the spin-up process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define rng\n",
    "rng = np.random.default_rng(seed=seeds.pop(0))\n",
    "\n",
    "# allocate memory\n",
    "Nt_train = 10000\n",
    "Nt_spinup = 100\n",
    "xt_train = np.zeros((Nt_train+1, true_model.Nx))\n",
    "\n",
    "# initialisation and spin-up\n",
    "xt_train[0] = rng.normal(loc=3, scale=1, size=true_model.Nx)\n",
    "for t in trange(Nt_spinup, desc='spin-up integration'):\n",
    "    xt_train[0] = true_model.forward(xt_train[0])\n",
    "\n",
    "# model integration\n",
    "for t in trange(Nt_train, desc='model integration (train)'):\n",
    "    xt_train[t+1] = true_model.forward(xt_train[t])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Preprocess the training data </span>\n",
    "\n",
    "The training dataset is made of input/output pairs, where the input is the state at a given time, and the output is the state at the following time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make input/output pairs: input = xt[t], output = xt[t+1]\n",
    "x_train_raw = xt_train[:-1]\n",
    "y_train_raw = xt_train[1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training dataset is then **normalised** and **split** into training and validation data. For our experiments, we keep one tenth of the data for validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalise the training data using numpy's broadcasting rules\n",
    "x_mean = x_train_raw.mean(axis=0)\n",
    "y_mean = y_train_raw.mean(axis=0)\n",
    "x_std = x_train_raw.std(axis=0)\n",
    "y_std = y_train_raw.std(axis=0)\n",
    "def normalise_x(x):\n",
    "    return (x - x_mean)/x_std\n",
    "def normalise_y(y):\n",
    "    return (y - y_mean)/y_std\n",
    "def denormalise_x(x_norm):\n",
    "    return x_norm*x_std + x_mean\n",
    "def denormalise_y(y_norm):\n",
    "    return y_norm*y_std + y_mean\n",
    "x_train_raw_norm = normalise_x(x_train_raw)\n",
    "y_train_raw_norm = normalise_y(y_train_raw)\n",
    "\n",
    "# split into training / validation\n",
    "index_train = np.array([i for i in range(Nt_train) if i%10])\n",
    "index_valid = np.array([i for i in range(Nt_train) if not i%10])\n",
    "x_train_norm = x_train_raw_norm[index_train]\n",
    "y_train_norm = y_train_raw_norm[index_train]\n",
    "x_valid_norm = x_train_raw_norm[index_valid]\n",
    "y_valid_norm = y_train_raw_norm[index_valid]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> A shorter model integration for the test data</span>\n",
    "\n",
    "We repeat the same process to make the **test dataset**. In this case, the trajectory starts from another random field (and we still get rid of the spin-up process) and can be somewhat shorter, but the normalisation must be the same as for the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define rng\n",
    "rng = np.random.default_rng(seed=seeds.pop(0))\n",
    "\n",
    "# allocate memory\n",
    "Nt_test = 1000\n",
    "Nt_spinup = 100\n",
    "xt_test = np.zeros((Nt_test+1, true_model.Nx))\n",
    "\n",
    "# initialisation and spin-up\n",
    "xt_test[0] = rng.normal(loc=3, scale=1, size=true_model.Nx)\n",
    "for t in trange(Nt_spinup, desc='spin-up integration'):\n",
    "    xt_test[0] = true_model.forward(xt_test[0])\n",
    "\n",
    "# model integration\n",
    "for t in trange(Nt_test, desc='model integration (test)'):\n",
    "    xt_test[t+1] = true_model.forward(xt_test[t])\n",
    "\n",
    "# make input/output pairs: input = xt[t], output = xt[t+1]\n",
    "x_test = xt_test[:-1]\n",
    "y_test = xt_test[1:]\n",
    "\n",
    "# normalise the test data using numpy's broadcasting rules\n",
    "x_test_norm = normalise_x(x_test)\n",
    "y_test_norm = normalise_y(y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> An ensemble model integration for the forecast skill data</span>\n",
    "\n",
    "In order to assess the forecast skill of the surrogate model, we will use a different test dataset, in which we record an ensemble of **trajectories** (instead of an ensemble of input/output pairs). This will allow us to measure the accuracy of the forecast for longer integration times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define rng\n",
    "rng = np.random.default_rng(seed=seeds.pop(0))\n",
    "\n",
    "# allocate memory\n",
    "Nt_fs = 400\n",
    "Nt_spinup = 100\n",
    "Ne_fs = 512\n",
    "xt_fs = np.zeros((Nt_fs+1, Ne_fs, true_model.Nx))\n",
    "\n",
    "# initialisation and spin-up\n",
    "xt_fs[0] = rng.normal(loc=3, scale=1, size=(Ne_fs, true_model.Nx))\n",
    "for t in trange(Nt_spinup, desc='spin-up integration'):\n",
    "    xt_fs[0] = true_model.forward(xt_fs[0])\n",
    "    \n",
    "# model integration\n",
    "for t in trange(Nt_fs, desc='model integration (ensemble)'):\n",
    "    xt_fs[t+1] = true_model.forward(xt_fs[t])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## <span style=\"color:green\"> The baseline model: persistence </span>\n",
    "\n",
    "In this first test series, we use **persistence** as surrogate model. This will provide baselines for our NN results. Persistence is defined as the model for which there is no time evolution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Evaluate the model</span>\n",
    "\n",
    "The mean square error (MSE) is the loss function that we will use to train our NNs later. Therefore, the test MSE is a measure of the efficiency of the learning/training process.\n",
    "\n",
    "The test MSE of persistence is a number whose absolute value is not that important per se (because the input and output data have been normalised) but it will be useful to normalise the test MSE of our trained NNs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# compute test MSE\n",
    "test_mse_baseline = np.mean(np.square(y_test_norm - x_test_norm))\n",
    "\n",
    "# compute forecast skill\n",
    "fs_baseline = np.sqrt(np.mean(np.square(xt_fs-xt_fs[0]), axis=2))\n",
    "\n",
    "# show test MSE\n",
    "print(f'test mse of persistence = {test_mse_baseline}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### <span style=\"color:blue\"> Example of surrogate model integration</span>\n",
    "\n",
    "In the following cell, we show one example of model integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_compare_traj(\n",
    "    xt_fs[:, 0],\n",
    "    np.broadcast_to(xt_fs[0, 0], shape=xt_fs[:, 0].shape),\n",
    "    true_model,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Forecast skill</span>\n",
    "\n",
    "In the following cell, we plot the average forecast skill, normalised by the model variability. The shadow delimits the 90% confidence interval (percentiles 5 and 95)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_forecast_skill(\n",
    "    dict(\n",
    "        persistence=fs_baseline,\n",
    "    ),\n",
    "    true_model,\n",
    "    p1=5,\n",
    "    p2=95,\n",
    "    xmax=4,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The error rapidly grows as time evolves. After about $1$ Lyapunov time, the error oscillates around $\\sqrt{2}$, which is the theoretical asymptotic value due to the normalisation and which is consistent with the wave behaviour of the dynamics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## <span style=\"color:green\"> A naive ML model </span>\n",
    "\n",
    "### <span style=\"color:blue\"> Construct and train the model</span>\n",
    "\n",
    "In this second test series, we train and evaluate a dense NN (sequential NN with only dense layers). In order to create this model, we use the sequential API of tensorflow as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_sequential_network(num_layers, num_nodes, activation):\n",
    "    # create a sequential network\n",
    "    network = tf.keras.models.Sequential()\n",
    "    # add the input layers\n",
    "    network.add(tf.keras.Input(shape=(true_model.Nx,)))\n",
    "    # add the internal layers\n",
    "    for i in range(num_layers):\n",
    "        network.add(tf.keras.layers.Dense(num_nodes, activation=activation))\n",
    "    # add the output layer without activation\n",
    "    network.add(tf.keras.layers.Dense(true_model.Nx))\n",
    "    return network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell, we actually build a dense NN with $4$ internal layers and $128$ nodes per layer. The total number of parameters of this model is $59944$. This is actually quite large for a $40$-variable system. This is because the dense architecture is rather \"inefficient\" in terms of parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set seed\n",
    "tf.keras.utils.set_random_seed(seeds.pop(0))\n",
    "\n",
    "# define the NN\n",
    "num_layers = 4\n",
    "num_nodes = 128\n",
    "activation = 'relu'\n",
    "naive_network = make_sequential_network(num_layers, num_nodes, activation)\n",
    "\n",
    "# compilation\n",
    "naive_network.compile(loss='mse', optimizer='adam')\n",
    "\n",
    "# print short summary\n",
    "naive_network.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell, we actually train the model for $256$ epochs. Going beyond part I, we use an EarlyStopping callback to end the training when the validation loss stops improving. This should avoid overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the ML model\n",
    "tf.keras.utils.set_random_seed(seeds.pop(0))\n",
    "num_epochs = 256\n",
    "tqdm_callback = utils.tqdm_callback(num_epochs, 'naive NN training')\n",
    "early_stopping_callback = tf.keras.callbacks.EarlyStopping(\n",
    "    monitor='val_loss',\n",
    "    patience=16,\n",
    "    verbose=0,\n",
    "    restore_best_weights=True,\n",
    ")\n",
    "fit_naive = naive_network.fit(\n",
    "    x_train_norm, \n",
    "    y_train_norm,\n",
    "    verbose=0,\n",
    "    epochs=num_epochs, \n",
    "    validation_data=(x_valid_norm, y_valid_norm),\n",
    "    callbacks=[tqdm_callback, early_stopping_callback],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell we plot the training history, that is, the evolution of the training MSE (the `loss`) and the validation MSE (the `val_loss`) as a function of the number of epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_learning_curve(\n",
    "    fit_naive.history['loss'],\n",
    "    fit_naive.history['val_loss'],\n",
    "    title='Naive NN training',\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both values are visually closely related. The validation MSE is more noisy than the training MSE, which is expected because the training data is nine times as large as the validation data. After several epochs, the validation MSE gets a bit higher than the training MSE. This is explained by the fact that this data is not used in the gradient descent algorithm. Finally, at the end the validation MSE stops improving. This is the sign that the model is starting to overfit the training data and that we should stop the training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Evaluate the model</span>\n",
    "\n",
    "We now compute the test MSE to evaluate our surrogate model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute test MSE\n",
    "test_mse_naive = naive_network.evaluate(x_test_norm, y_test_norm, verbose=0, batch_size=Nt_test)\n",
    "\n",
    "# compute forecast skill\n",
    "xt_naive = np.zeros(xt_fs.shape)\n",
    "xt_naive[0] = xt_fs[0]\n",
    "for t in trange(xt_naive.shape[0]-1, desc='naive surrogate model integration'):\n",
    "    x_norm = normalise_x(xt_naive[t])\n",
    "    y_norm = naive_network.predict(x_norm, batch_size=Ne_fs, verbose=0)\n",
    "    xt_naive[t+1] = denormalise_y(y_norm)\n",
    "fs_naive = np.sqrt(np.mean(np.square(xt_fs-xt_naive), axis=2))\n",
    "\n",
    "# show test MSE\n",
    "print(f'test mse of persistence = {test_mse_baseline}')\n",
    "print(f'test mse of naive model = {test_mse_naive}')\n",
    "print()\n",
    "print(f'relative test mse of naive model = {test_mse_naive/test_mse_baseline}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We obtain a reduction of about 80%, which is already quite good, but we will see later that it is possible to do much better."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### <span style=\"color:blue\"> Example of surrogate model integration</span>\n",
    "\n",
    "In the following cell, we show once again one example of model integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_compare_traj(\n",
    "    xt_fs[:, 0],\n",
    "    xt_naive[:, 0],\n",
    "    true_model,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The error is lower than in the first test series, but only during the first few integration steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### <span style=\"color:blue\"> Forecast skill</span>\n",
    "\n",
    "In the following cell, we plot once again the average forecast skill, normalised by the model variability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_forecast_skill(\n",
    "    dict(\n",
    "        persistence=fs_baseline,\n",
    "        naive=fs_naive,\n",
    "    ),\n",
    "    true_model,\n",
    "    p1=5,\n",
    "    p2=95,\n",
    "    xmax=4,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This curve confirms that the naive surrogate model is more accurate than persistence for one integration step, and that it remains more accurate until about $2$ Lyapunov times."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span style=\"color:green\"> A smart ML model </span>\n",
    "\n",
    "### <span style=\"color:blue\"> Build and train the model</span>\n",
    "\n",
    "In this third and last test series, we train and evaluate a smart NN. This NN uses a sparse architecture with convolutional NN and controlled nonlinearity to reproduce the **model tendencies**, as well as a Runge-Kutta integration scheme to **emulate the dynamics**. In order to implement this NN, we use both the functional API (for the model tendency) and the subclassing API (for the integration scheme) of tensorflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SmartNetwork(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self, num_filters, kernel_size, dt=0.05, **kwargs):\n",
    "        super(SmartNetwork, self).__init__(**kwargs)\n",
    "        self.dt = dt\n",
    "        \n",
    "        # reshape layers\n",
    "        reshape_input = tf.keras.layers.Reshape((true_model.Nx, 1))\n",
    "        reshape_output = tf.keras.layers.Reshape((true_model.Nx,))\n",
    "        \n",
    "        # padding layer\n",
    "        border = kernel_size//2\n",
    "        def apply_padding(x):\n",
    "            x_left = x[..., -border:, :]\n",
    "            x_right = x[..., :border, :]\n",
    "            return tf.concat([x_left, x, x_right], axis=-2)\n",
    "        padding_layer = tf.keras.layers.Lambda(apply_padding)\n",
    "        \n",
    "        # convolutional layers\n",
    "        conv_layer_1 = tf.keras.layers.Conv1D(num_filters, kernel_size)\n",
    "        conv_layer_2 = tf.keras.layers.Conv1D(1, 1)\n",
    "        \n",
    "        # network for the model tendencies\n",
    "        x_in = tf.keras.Input(shape=(true_model.Nx,))\n",
    "        x = reshape_input(x_in)\n",
    "        x = padding_layer(x)\n",
    "        x1 = conv_layer_1(x)\n",
    "        x2 = x1 * x1\n",
    "        x3 = tf.concat([x1, x2], axis=-1)\n",
    "        x_out = conv_layer_2(x3)\n",
    "        x_out = reshape_output(x_out)\n",
    "        self.tendency = tf.keras.Model(inputs=x_in, outputs=x_out)\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self, x):\n",
    "        dx_dt_0 = self.tendency(x)\n",
    "        dx_dt_1 = self.tendency(x+0.5*self.dt*dx_dt_0)\n",
    "        dx_dt_2 = self.tendency(x+0.5*self.dt*dx_dt_1)\n",
    "        dx_dt_3 = self.tendency(x+self.dt*dx_dt_2)\n",
    "        dx_dt =  (dx_dt_0 + 2*dx_dt_1 + 2*dx_dt_2 + dx_dt_3)/6\n",
    "        return x + self.dt*dx_dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set seed\n",
    "tf.keras.utils.set_random_seed(seeds.pop(0))\n",
    "\n",
    "# define the NN\n",
    "num_filters = 6\n",
    "kernel_size = 5\n",
    "smart_network = SmartNetwork(num_filters, kernel_size, dt=true_model.dt)\n",
    "\n",
    "# compilation\n",
    "smart_network.compile(loss='mse', optimizer='adam')\n",
    "\n",
    "# print short summary\n",
    "smart_network.tendency.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total number of parameters is only $49$. Furthermore in this case, with well-chosen parameters it is possible to reproduce the true dynamics up to machine precision: the model is said to be **identifiable**. Also note that this network is built in such a way that we don't need the input and output data to be normalised.\n",
    "\n",
    "In the following cell, we actually train the model for up to $128$ epochs. Once again, we use an EarlyStopping callback to end the training when the validation loss stops improving in order to avoid overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the ML model\n",
    "tf.keras.utils.set_random_seed(seeds.pop(0))\n",
    "num_epochs = 128\n",
    "tqdm_callback = utils.tqdm_callback(num_epochs, 'smart NN training')\n",
    "early_stopping_callback = tf.keras.callbacks.EarlyStopping(\n",
    "    monitor='val_loss',\n",
    "    patience=16,\n",
    "    verbose=0,\n",
    "    restore_best_weights=True)\n",
    "fit_smart = smart_network.fit(\n",
    "    denormalise_x(x_train_norm),\n",
    "    denormalise_y(y_train_norm),\n",
    "    verbose=0,\n",
    "    epochs=num_epochs, \n",
    "    validation_data=(denormalise_x(x_valid_norm), denormalise_y(y_valid_norm)),\n",
    "    callbacks=[tqdm_callback, early_stopping_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell we plot the training history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_learning_curve(\n",
    "    fit_smart.history['loss'],\n",
    "    fit_smart.history['val_loss'],\n",
    "    title='Smart NN training',\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once again, the training and validation MSE are visually closely related. However, by contrast with the previous test series, after about $30$ epochs, the MSEs have decreased to $10^{-12}$, which should be very close to the numerical precision zero (tensorflow is working on simple precision for real numbers). Passed $30$ epochs, the MSEs oscillate at very low values. This behaviour can be considered as numerical noise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span style=\"color:blue\"> Evaluate the model</span>\n",
    "\n",
    "We now compute the test MSE to evaluate our surrogate model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute test MSE\n",
    "y_test_smart = smart_network.predict(denormalise_x(x_test_norm), batch_size=Nt_test, verbose=0)\n",
    "test_mse_smart = np.mean(np.square(normalise_y(y_test_smart)-y_test_norm))\n",
    "\n",
    "# compute forecast skill\n",
    "xt_smart = np.zeros(xt_fs.shape)\n",
    "xt_smart[0] = xt_fs[0]\n",
    "for t in trange(xt_smart.shape[0]-1, desc='smart surrogate model integration'):\n",
    "    xt_smart[t+1] = smart_network.predict(xt_smart[t], batch_size=Ne_fs, verbose=0)\n",
    "fs_smart = np.sqrt(np.mean(np.square(xt_fs-xt_smart), axis=2))\n",
    "\n",
    "# show test MSE\n",
    "print(f'test mse of persistence = {test_mse_baseline}')\n",
    "print(f'test mse of naive model = {test_mse_naive}')\n",
    "print(f'test mse of smart model = {test_mse_smart}')\n",
    "print()\n",
    "print(f'relative test mse of naive model = {test_mse_naive/test_mse_baseline}')\n",
    "print(f'relative test mse of smart model = {test_mse_smart/test_mse_baseline}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The test MSE is sufficiently close to zero so that we can consider that our surrogate model reproduces the true model dynamics up to numerical precision."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### <span style=\"color:blue\"> Example of surrogate model integration</span>\n",
    "\n",
    "In the following cell, we show once again one example of model integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_compare_traj(\n",
    "    xt_fs[:, 0],\n",
    "    xt_smart[:, 0],\n",
    "    true_model,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, the error is so low that it is not visible until about $6$ MTU. At that time, the true model trajectory and the surrogate model trajectory diverge from each other. Indeed, the two models are equivalent up to numerical precision, but they are not bit-wise equivalent, which means that this divergence is unavoidable because of the chaotic nature of the dynamics. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### <span style=\"color:blue\"> Forecast skill</span>\n",
    "\n",
    "In the following cell, we plot once again the average forecast skill, normalised by the model variability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.plot_l96_forecast_skill(\n",
    "    dict(\n",
    "        persistence=fs_baseline,\n",
    "        naive=fs_naive,\n",
    "        smart=fs_smart,\n",
    "    ),\n",
    "    true_model,\n",
    "    p1=5,\n",
    "    p2=95,\n",
    "    xmax=30,\n",
    "    linewidth=18,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This curve confirms that the the smart surrogate model is equivalent to the true model up to numerical precision. The numerical divergence between the true and surrogate model happens on average after about $10$ Lyapunov times."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}