{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qT_RWmTEugu9"
   },
   "source": [
    "# Reducing Numerical Errors with Deep Learning\n",
    "\n",
    "First, we'll target numerical errors that arise in the discretization of a continuous PDE $\\mathcal P^*$, i.e. when we formulate $\\mathcal P$. This approach will demonstrate that, despite the lack of closed-form descriptions, discretization errors often are functions with regular and repeating structures and, thus, can be learned by a neural network. Once the network is trained, it can be evaluated locally to improve the solution of a PDE-solver, i.e., to reduce its numerical error. The resulting method is a hybrid one: it will always run (a coarse) PDE solver, and then improve it at runtime with corrections inferred by an NN.\n",
    "\n",
    " \n",
    "Pretty much all numerical methods contain some form of iterative process: repeated updates over time for explicit solvers, or within a single update step for implicit solvers. \n",
    "An example for the second case could be found [here](https://github.com/tum-pbs/CG-Solver-in-the-Loop),\n",
    "but below we'll target the first case, i.e. iterations over time.\n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/diffphys-code-sol.ipynb)\n",
    "\n",
    "\n",
    "## Problem formulation\n",
    "\n",
    "In the context of reducing errors, it's crucial to have a _differentiable physics solver_, so that the learning process can take the reaction of the solver into account. This interaction is not possible with supervised learning or PINN training. Even small inference errors of a supervised NN can accumulate over time, and lead to a data distribution that differs from the distribution of the pre-computed data. This distribution shift can lead to sub-optimal results, or even cause blow-ups of the solver.\n",
    "\n",
    "In order to learn the error function, we'll consider two different discretizations of the same PDE $\\mathcal P^*$: \n",
    "a _reference_ version, which we assume to be accurate, with a discretized version \n",
    "$\\mathcal P_r$, and solutions $\\mathbf r \\in \\mathscr R$, where $\\mathscr R$ denotes the manifold of solutions of $\\mathcal P_r$.\n",
    "In parallel to this, we have a less accurate approximation of the same PDE, which we'll refer to as the _source_ version, as this will be the solver that our NN should later on interact with. Analogously,\n",
    "we have $\\mathcal P_s$ with solutions $\\mathbf s \\in \\mathscr S$.\n",
    "After training, we'll obtain a _hybrid_ solver that uses $\\mathcal P_s$ in conjunction with a trained network to obtain improved solutions, i.e., solutions that are closer to the ones produced by $\\mathcal P_r$.\n",
    "\n",
    "```{figure} resources/diffphys-sol-manifolds.jpeg\n",
    "---\n",
    "height: 150px\n",
    "name: diffphys-sol-manifolds\n",
    "---\n",
    "Visual overview of coarse and reference manifolds\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tayrJa7_ZzS_"
   },
   "source": [
    "\n",
    "Let's assume $\\mathcal{P}$ advances a solution by a time step $\\Delta t$, and let's denote $n$ consecutive steps by a superscript:\n",
    "$\n",
    "\\newcommand{\\pde}{\\mathcal{P}}\n",
    "\\newcommand{\\pdec}{\\pde_{s}}\n",
    "\\newcommand{\\vc}[1]{\\mathbf{s}_{#1}} \n",
    "\\newcommand{\\vr}[1]{\\mathbf{r}_{#1}} \n",
    "\\newcommand{\\vcN}{\\vs}          \n",
    "\\newcommand{\\project}{\\mathcal{T}}   \n",
    "\\pdec^n ( \\mathcal{T} \\vr{t} ) = \\pdec(\\pdec(\\cdots \\pdec( \\mathcal{T} \\vr{t}  )\\cdots)) .\n",
    "$ \n",
    "The corresponding state of the simulation is\n",
    "$\n",
    "\\mathbf{s}_{t+n} = \\mathcal{P}^n ( \\mathcal{T} \\mathbf{r}_{t} ) .\n",
    "$\n",
    "Here we assume a mapping operator $\\mathcal{T}$ exists that transfers a reference solution to the source manifold. This could, e.g., be a simple downsampling operation.\n",
    "Especially for longer sequences, i.e. larger $n$, the source state \n",
    "$\\newcommand{\\vc}[1]{\\mathbf{s}_{#1}} \\vc{t+n}$\n",
    "will deviate from a corresponding reference state\n",
    "$\\newcommand{\\vr}[1]{\\mathbf{r}_{#1}} \\vr{t+n}$. \n",
    "This is what we will address with an NN in the following.\n",
    "\n",
    "As before, we'll use an $L^2$-norm to quantify the deviations, i.e., \n",
    "an error function $\\newcommand{\\loss}{e} \n",
    "\\newcommand{\\corr}{\\mathcal{C}} \n",
    "\\newcommand{\\vc}[1]{\\mathbf{s}_{#1}} \n",
    "\\newcommand{\\vr}[1]{\\mathbf{r}_{#1}} \n",
    "\\loss (\\vc{t},\\mathcal{T} \\vr{t})=\\Vert\\vc{t}-\\mathcal{T} \\vr{t}\\Vert_2$. \n",
    "Our learning goal is to train at a correction operator \n",
    "$\\mathcal{C} ( \\mathbf{s} )$ such that \n",
    "a solution to which the correction is applied has a lower error than the original unmodified (source) \n",
    "solution: $\\newcommand{\\loss}{e} \n",
    "\\newcommand{\\corr}{\\mathcal{C}} \n",
    "\\newcommand{\\vr}[1]{\\mathbf{r}_{#1}} \n",
    "\\loss ( \\mathcal{P}_{s}( \\corr (\\mathcal{T} \\vr{t_0}) ) , \\mathcal{T} \\vr{t_1}) < \\loss ( \\mathcal{P}_{s}( \\mathcal{T} \\vr{t_0} ), \\mathcal{T} \\vr{t_1})$. \n",
    "\n",
    "The correction function \n",
    "$\\newcommand{\\vcN}{\\mathbf{s}} \\newcommand{\\corr}{\\mathcal{C}} \\corr (\\vcN | \\theta)$ \n",
    "is represented as a deep neural network with weights $\\theta$\n",
    "and receives the state $\\mathbf{s}$ to infer an additive correction field with the same dimension.\n",
    "To distinguish the original states $\\mathbf{s}$ from the corrected ones, we'll denote the latter with an added tilde $\\tilde{\\mathbf{s}}$.\n",
    "The overall learning goal now becomes\n",
    "\n",
    "$$\n",
    "\\newcommand{\\corr}{\\mathcal{C}}  \n",
    "\\newcommand{\\vr}[1]{\\mathbf{r}_{#1}} \n",
    "\\text{argmin}_\\theta | ( \\mathcal{P}_{s} \\corr )^n ( \\mathcal{T} \\vr{t} ) - \\mathcal{T} \\vr{t}|^2\n",
    "$$\n",
    "\n",
    "To simplify the notation, we've dropped the sum over different samples here (the $i$ from previous versions).\n",
    "A crucial bit that's easy to overlook in the equation above, is that the correction depends on the modified states, i.e.\n",
    "it is a function of\n",
    "$\\tilde{\\mathbf{s}}$, so we have \n",
    "$\\newcommand{\\vctN}{\\tilde{\\mathbf{s}}} \\newcommand{\\corr}{\\mathcal{C}} \\corr (\\vctN | \\theta)$.\n",
    "These states actually evolve over time when training. They don't exist beforehand.\n",
    "\n",
    "**TL;DR**:\n",
    "We'll train a network $\\mathcal{C}$ to reduce the numerical errors of a simulator with a more accurate reference. It's crucial to have the _source_ solver realized as a differential physics operator, such that it can give gradients for an improved training of $\\mathcal{C}$.\n",
    "\n",
    "<br>\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting started with the implementation\n",
    "\n",
    "The following replicates an experiment) from [Solver-in-the-loop: learning from differentiable physics to interact with iterative pde-solvers](https://ge.in.tum.de/publications/2020-um-solver-in-the-loop/) {cite}`holl2019pdecontrol`, further details can be found in section B.1 of the [appendix](https://arxiv.org/pdf/2007.00016.pdf) of the paper.\n",
    "\n",
    "First, let's download the prepared data set (for details on generation & loading cf. https://github.com/tum-pbs/Solver-in-the-Loop), and let's get the data handling out of the way, so that we can focus on the _interesting_ parts..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JwZudtWauiGa",
    "outputId": "bd3a4a4d-706f-4210-ee4e-ca4e370b762d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded data, 6 training sims\n"
     ]
    }
   ],
   "source": [
    "import os, sys, logging, argparse, pickle, glob, random, distutils.dir_util\n",
    "\n",
    "if not os.path.isfile('sol-data-karman-2d-train.pickle'):\n",
    "  import urllib.request\n",
    "  url=\"https://ge.in.tum.de/download/2020-solver-in-the-loop/sol-karman-2d-data.pickle\"\n",
    "  print(\"Downloading training data (73MB), this can take a moment the first time...\")\n",
    "  urllib.request.urlretrieve(url, 'sol-data-karman-2d-train.pickle')\n",
    "\n",
    "with open('sol-data-karman-2d-train.pickle', 'rb') as f: dataPreloaded = pickle.load(f)\n",
    "print(\"Loaded data, {} training sims\".format(len(dataPreloaded)) )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RY1F4kdWPLNG"
   },
   "source": [
    "Also let's get installing / importing all the necessary libraries out of the way. And while we're at it, we can set the random seed - obviously, 42 is the ultimate choice here 🙂"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BGN4GqxkIueM",
    "outputId": "095adbf8-1ef6-41fd-938e-6cafcf0fdfdc"
   },
   "outputs": [],
   "source": [
    "#!pip install --upgrade --quiet phiflow\n",
    "#%tensorflow_version 1.x\n",
    "\n",
    "from phi.tf.flow import *\n",
    "import phi.tf.util\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "\n",
    "random.seed(42)\n",
    "np.random.seed(42)\n",
    "tf.compat.v1.set_random_seed(42)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OhnzPdoww11P"
   },
   "source": [
    "## Simulation setup\n",
    "\n",
    "Now we can set up the _source_ simulation $\\mathcal{P}_{s}$. \n",
    "Note that we won't deal with \n",
    "$\\mathcal{P}_{r}$\n",
    "below: the downsampled reference data is contained in the training data set. It was generated with a four times finer discretization. Below we're focusing on the interaction of the source solver and the NN. \n",
    "\n",
    "This code block and the next ones will define lots of functions, that will be used later on for training.\n",
    "\n",
    "The `KarmanFlow` solver below simulates a relatively standard wake flow case with a spherical obstacle in a rectangular domain, and an explicit viscosity solve to obtain different Reynolds numbers. This is the geometry of the setup:\n",
    "\n",
    "```{figure} resources/diffphys-sol-domain.png\n",
    "---\n",
    "height: 200px\n",
    "name: diffphys-sol-domain\n",
    "---\n",
    "Domain setup for the wake flow case.\n",
    "```\n",
    "\n",
    "The solver applies inflow boundary conditions for the y-velocity with a pre-multiplied mask (`velBCy,velBCyMask`), to set the y components at the bottom of the domai. It then calls `super().step()` to run the _regular_ phiflow fluid solving step.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "6WNMcdWUw4EP"
   },
   "outputs": [],
   "source": [
    "class KarmanFlow(IncompressibleFlow):\n",
    "    def __init__(self, pressure_solver=None, make_input_divfree=False, make_output_divfree=True):\n",
    "        IncompressibleFlow.__init__(self, pressure_solver, make_input_divfree, make_output_divfree)\n",
    "\n",
    "        self.infl = Inflow(box[5:10, 25:75])\n",
    "        self.obst = Obstacle(Sphere([50, 50], 10))\n",
    "\n",
    "    def step(self, fluid, re, res, velBCy, velBCyMask, dt=1.0, gravity=Gravity()):\n",
    "        # apply viscosity\n",
    "        alpha = 1.0/math.reshape(re, [fluid._batch_size, 1, 1, 1]) * dt * res * res\n",
    "\n",
    "        cy = diffuse(CenteredGrid(fluid.velocity.data[0].data), alpha)\n",
    "        cx = diffuse(CenteredGrid(fluid.velocity.data[1].data), alpha)\n",
    "\n",
    "        # apply velocity BCs, only for y velocity for now. note: content of velBCy should be pre-multiplied\n",
    "        cy = cy*(1.0 - velBCyMask) + velBCy\n",
    "\n",
    "        fluid = fluid.copied_with(velocity=StaggeredGrid([cy.data, cx.data], fluid.velocity.box))\n",
    "\n",
    "        return super().step(fluid=fluid, dt=dt, obstacles=[self.obst], gravity=gravity, density_effects=[self.infl], velocity_effects=())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RYFUGICgxk0K"
   },
   "source": [
    "## Network architecture\n",
    "\n",
    "We'll also define two alternative versions of a neural networks to represent \n",
    "$\\newcommand{\\vcN}{\\mathbf{s}} \\newcommand{\\corr}{\\mathcal{C}} \\corr$. In both cases we'll use fully convolutional networks, i.e. networks without any fully-connected layers. We'll use Keras within tensorflow to define the layers of the network (mostly via `Conv2D`), typically activated via ReLU and LeakyReLU functions, respectively.\n",
    "The inputs to the network are: \n",
    "- 2 fields with x,y velocity\n",
    "- plus the Reynolds number as constant channel.\n",
    "\n",
    "The output is: \n",
    "- a 2 component field containing the x,y velocity\n",
    "\n",
    "First, let's define a minimal network consisting only of three convolutional layers with ReLU activations (we're also using keras here for simplicity). The input channel dimension is defined via the `tensor_in`, then we'll go to 32 and 64 features, before reducing to 2 channels in the output. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "qIrWYTy6xscA"
   },
   "outputs": [],
   "source": [
    "def network_small(tensor_in):\n",
    "    return keras.Sequential([\n",
    "        keras.layers.Input(tensor=tensor_in),\n",
    "        keras.layers.Conv2D(filters=32, kernel_size=5, padding='same', activation=tf.nn.relu),\n",
    "        keras.layers.Conv2D(filters=64, kernel_size=5, padding='same', activation=tf.nn.relu),\n",
    "        keras.layers.Conv2D(filters=2,  kernel_size=5, padding='same', activation=None),  # u, v\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YfHvdI7yxtdj"
   },
   "source": [
    "For flexibility (and larger-scale tests later on), let's define a _proper_ ResNet with a few more layers. This architecture is the one from the original paper, and will give a fairly good performance (`network_small` above will train faster, but give a sub-optimal performance at inference time)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "TyfpA7Fbx0ro"
   },
   "outputs": [],
   "source": [
    "def network_medium(tensor_in):\n",
    "    l_input = keras.layers.Input(tensor=tensor_in)\n",
    "    block_0 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(l_input)\n",
    "    block_0 = keras.layers.LeakyReLU()(block_0)\n",
    "\n",
    "    l_conv1 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(block_0)\n",
    "    l_conv1 = keras.layers.LeakyReLU()(l_conv1)\n",
    "    l_conv2 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(l_conv1)\n",
    "    l_skip1 = keras.layers.add([block_0, l_conv2])\n",
    "    block_1 = keras.layers.LeakyReLU()(l_skip1)\n",
    "\n",
    "    l_conv3 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(block_1)\n",
    "    l_conv3 = keras.layers.LeakyReLU()(l_conv3)\n",
    "    l_conv4 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(l_conv3)\n",
    "    l_skip2 = keras.layers.add([block_1, l_conv4])\n",
    "    block_2 = keras.layers.LeakyReLU()(l_skip2)\n",
    "\n",
    "    l_conv5 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(block_2)\n",
    "    l_conv5 = keras.layers.LeakyReLU()(l_conv5)\n",
    "    l_conv6 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(l_conv5)\n",
    "    l_skip3 = keras.layers.add([block_2, l_conv6])\n",
    "    block_3 = keras.layers.LeakyReLU()(l_skip3)\n",
    "\n",
    "    l_conv7 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(block_3)\n",
    "    l_conv7 = keras.layers.LeakyReLU()(l_conv7)\n",
    "    l_conv8 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(l_conv7)\n",
    "    l_skip4 = keras.layers.add([block_3, l_conv8])\n",
    "    block_4 = keras.layers.LeakyReLU()(l_skip4)\n",
    "\n",
    "    l_conv9 = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(block_4)\n",
    "    l_conv9 = keras.layers.LeakyReLU()(l_conv9)\n",
    "    l_convA = keras.layers.Conv2D(filters=32, kernel_size=5, padding='same')(l_conv9)\n",
    "    l_skip5 = keras.layers.add([block_4, l_convA])\n",
    "    block_5 = keras.layers.LeakyReLU()(l_skip5)\n",
    "\n",
    "    l_output = keras.layers.Conv2D(filters=2,  kernel_size=5, padding='same')(block_5)\n",
    "    return keras.models.Model(inputs=l_input, outputs=l_output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ew-MgPSlyLW-"
   },
   "source": [
    "Next, we're coming to two functions which are pretty important: they transform the simulation state into an input tensor for the network, and vice versa. Hence, they're the interface between _keras/tensorflow_ and _phiflow_.\n",
    "\n",
    "The `to_keras` function uses the `staggered_tensor` from phiflow (a 2 component tensor instead of 2 separate grids), from which we discard the outermost layer. We then add a constant channel via `tf.constant` that is multiplied by the desired Reynolds number. The resulting stack of grids is provided as a tensor as input to the neural network. \n",
    "\n",
    "After network evaluation, we transform the output tensor back into a phiflow grid via the `to_staggered` function. \n",
    "It converts the 2-component tensor that is returned by the network into a phiflow staggered grid object, so that it is compatible with the velocity field of the fluid simulation.\n",
    "(Note: these are two _centered_ grids with different sizes, so we leave the work to the`unstack_staggered_tensor` function in `StaggeredGrid()` constructor)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "hhGFpTjGyRyg"
   },
   "outputs": [],
   "source": [
    "def to_keras(fluidstate, ext_const_channel):\n",
    "  # drop the unused edges of the staggered velocity grid making its size the same as the centered grid\n",
    "  return math.concat(\n",
    "      [\n",
    "          fluidstate.velocity.staggered_tensor()[:, :-1:, :-1:, 0:2],\n",
    "          tf.constant(shape=fluidstate.density.data.shape, value=1.0)*math.reshape(value=ext_const_channel, shape=[fluidstate._batch_size, 1, 1, 1]),\n",
    "      ],\n",
    "      axis=-1\n",
    "  )\n",
    "\n",
    "def to_staggered(tensor_cen, box):\n",
    "  return StaggeredGrid(math.pad(tensor_cen, ((0,0), (0,1), (0,1), (0,0))), box=box)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VngMwN_9y00S"
   },
   "source": [
    "---\n",
    "\n",
    "## Data handling\n",
    "\n",
    "So far so good - we also need to take care of a few more mundane tasks, e.g. the some data handling and randomization. Below we define a `Dataset` class that stores all \"ground truth\" reference data (already downsampled).\n",
    "\n",
    "We actually have a lot of data dimensions: multiple simulations, with many time steps, each with different fields. This makes the code below a bit more difficult to read.\n",
    "\n",
    "The data format for the numpy array `dataPreloaded`: is  `['sim_name', frame, field (dens & vel)]`, where each field has dimension `[batch-size, y-size, x-size, channels]` (this is the standard for a phiflow export)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "tjywcdD2y20t"
   },
   "outputs": [],
   "source": [
    "class Dataset():\n",
    "    def __init__(self, data_preloaded, num_frames, num_sims=None, batch_size=1, is_testset=False):\n",
    "        self.epoch         = None\n",
    "        self.epochIdx      = 0\n",
    "        self.batch         = None\n",
    "        self.batchIdx      = 0\n",
    "        self.step          = None\n",
    "        self.stepIdx       = 0\n",
    "\n",
    "        self.dataPreloaded = data_preloaded\n",
    "        self.batchSize     = batch_size\n",
    "\n",
    "        self.numSims       = num_sims\n",
    "        self.numBatches    = num_sims//batch_size\n",
    "        self.numFrames     = num_frames\n",
    "        self.numSteps      = num_frames\n",
    "        \n",
    "        # initialize directory keys (using naming scheme from SoL codebase)\n",
    "        # constant additional per-sim channel: Reynolds numbers from data generation\n",
    "        # hard coded for training and test data here\n",
    "        if not is_testset:\n",
    "            self.dataSims = ['karman-fdt-hires-set/sim_%06d'%i for i in range(num_sims) ]\n",
    "            ReNrs = [160000.0, 320000.0, 640000.0,  1280000.0,  2560000.0,  5120000.0]\n",
    "            self.extConstChannelPerSim = { self.dataSims[i]:[ReNrs[i]] for i in range(num_sims) }\n",
    "        else:\n",
    "            self.dataSims = ['karman-fdt-hires-testset-reduced/sim_%06d'%i for i in range(num_sims) ]\n",
    "            ReNrs = [240000.0, 960000.0, 3840000.0] \n",
    "            self.extConstChannelPerSim = { self.dataSims[i]:[ReNrs[i]] for i in range(num_sims) }\n",
    "\n",
    "        self.dataFrms = [ np.arange(num_frames) for _ in self.dataSims ]  \n",
    "\n",
    "        #print(format(self.dataPreloaded[self.dataSims[0]][0][0].shape )) # debugging example: check shape of a single marker density field\n",
    "\n",
    "        # the data has the following shape ['sim', frame, field (dens/vel)] where each field is [batch-size, y-size, x-size, channels]\n",
    "        self.resolution = self.dataPreloaded[self.dataSims[0]][0][0].shape[1:3]  \n",
    "\n",
    "        # compute data statistics for normalization\n",
    "        self.dataStats = {\n",
    "            'std': (\n",
    "                np.std(np.concatenate([np.absolute(self.dataPreloaded[asim][i][0].reshape(-1)) for asim in self.dataSims for i in range(num_frames)], axis=-1)),  # marker density\n",
    "                (\n",
    "                    np.std(np.concatenate([np.absolute(self.dataPreloaded[asim][i][1][...,0].reshape(-1)) for asim in self.dataSims for i in range(num_frames)])),  # vel[0]\n",
    "                    np.std(np.concatenate([np.absolute(self.dataPreloaded[asim][i][1][...,1].reshape(-1)) for asim in self.dataSims for i in range(num_frames)])),  # vel[1]\n",
    "                )\n",
    "            ),\n",
    "            'ext.std': [ np.std([np.absolute(self.extConstChannelPerSim[asim][0]) for asim in self.dataSims]) ] # Re\n",
    "        }\n",
    "\n",
    "        if not is_testset:\n",
    "            print(\"Data stats: \"+format(self.dataStats))\n",
    "\n",
    "    # re-shuffle data for next epoch\n",
    "    def newEpoch(self, exclude_tail=0, shuffle_data=True):\n",
    "        self.numSteps = self.numFrames - exclude_tail\n",
    "        simSteps = [ (asim, self.dataFrms[i][0:(len(self.dataFrms[i])-exclude_tail)]) for i,asim in enumerate(self.dataSims) ]\n",
    "        sim_step_pair = []\n",
    "        for i,_ in enumerate(simSteps):\n",
    "            sim_step_pair += [ (i, astep) for astep in simSteps[i][1] ]  # (sim_idx, step) ...\n",
    "\n",
    "        if shuffle_data: random.shuffle(sim_step_pair)\n",
    "        self.epoch = [ list(sim_step_pair[i*self.numSteps:(i+1)*self.numSteps]) for i in range(self.batchSize*self.numBatches) ]\n",
    "        self.epochIdx += 1\n",
    "        self.batchIdx = 0\n",
    "        self.stepIdx = 0\n",
    "\n",
    "    def nextBatch(self):  \n",
    "        self.batchIdx += self.batchSize\n",
    "        self.stepIdx = 0\n",
    "\n",
    "    def nextStep(self):\n",
    "        self.stepIdx += 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "twIMJ3V0N1FX"
   },
   "source": [
    "The `nextEpoch`, `nextBatch`, and `nextStep` functions will be called at training time to randomize the training data.\n",
    "\n",
    "Now we need one more function that compiles the data for a mini batch to train with, called `getData` below. It returns batches of the desired size in terms of marker density, velocity, and Reynolds number.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "Dfwd4TnqN1Tn"
   },
   "outputs": [],
   "source": [
    "    # for class Dataset():\n",
    "\n",
    "    # get one mini batch of data: [marker density, velocity, Reynolds number] all from ground truth\n",
    "    def getData(self, consecutive_frames, with_skip=1):\n",
    "        marker_dens = [\n",
    "            math.concat([\n",
    "                self.dataPreloaded[\n",
    "                    self.dataSims[self.epoch[self.batchIdx+i][self.stepIdx][0]]  # sim_key\n",
    "                ][\n",
    "                    self.epoch[self.batchIdx+i][self.stepIdx][1]+j*with_skip  # steps\n",
    "                ][0]\n",
    "                for i in range(self.batchSize)\n",
    "            ], axis=0) for j in range(consecutive_frames+1)\n",
    "        ]\n",
    "        velocity = [\n",
    "            math.concat([\n",
    "                self.dataPreloaded[\n",
    "                    self.dataSims[self.epoch[self.batchIdx+i][self.stepIdx][0]]  # sim_key\n",
    "                ][\n",
    "                    self.epoch[self.batchIdx+i][self.stepIdx][1]+j*with_skip  # steps\n",
    "                ][1]\n",
    "                for i in range(self.batchSize)\n",
    "            ], axis=0) for j in range(consecutive_frames+1)\n",
    "        ]\n",
    "        ext = [\n",
    "            self.extConstChannelPerSim[\n",
    "                self.dataSims[self.epoch[self.batchIdx+i][self.stepIdx][0]]\n",
    "            ][0] for i in range(self.batchSize)\n",
    "        ]\n",
    "        return [marker_dens, velocity, ext]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bIWnyPYlz8q7"
   },
   "source": [
    "Note that the `density` here denotes a passively advected marker field, and not the density of the fluid. Below we'll be focusing on the velocity only, the marker density is tracked purely for visualization purposes.\n",
    "\n",
    "After all the definitions we can finally run some code. We can define the dataset object with the downloaded data from the first cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "59EBdEdj0QR2",
    "outputId": "8043f090-4e7b-4178-d2d2-513981e3905b"
   },
   "outputs": [],
   "source": [
    "nsims = 6\n",
    "batch_size = 3\n",
    "simsteps = 500\n",
    "\n",
    "dataset = Dataset( data_preloaded=dataPreloaded, num_frames=simsteps, num_sims=nsims, batch_size=batch_size )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0N92RooWPzeA"
   },
   "source": [
    "Additionally, we've defined several global variables to control the training and the simulation in the next code cells.\n",
    "\n",
    "The most important and interesting one is `msteps`. It defines the number of simulation steps that are unrolled at each training iteration. This directly influences the runtime of each training step, as we first have to simulate all steps forward, and then backpropagate the gradient through all `msteps` simulation steps interleaved with the NN evaluations. However, this is where we'll receive important feedback in terms of gradients how the inferred corrections actually influence a running simulation. Hence, larger `msteps` are typically better.\n",
    "\n",
    "In addition we define the `source` and `reference` simulations below. The reference is just a placeholder for data, only the `source` simulation is actually executed. We also define the actual NN `network`. All three are initialized with the size given in the data set `dataset.resolution`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EjgkdCzKP2Ip",
    "outputId": "6b21bd54-15aa-4440-b274-c3a68ab244f8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "conv2d (Conv2D)              (3, 64, 32, 32)           2432      \n",
      "_________________________________________________________________\n",
      "conv2d_1 (Conv2D)            (3, 64, 32, 64)           51264     \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (3, 64, 32, 2)            3202      \n",
      "=================================================================\n",
      "Total params: 56,898\n",
      "Trainable params: 56,898\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thuerey/Library/Python/3.7/lib/python/site-packages/ipykernel_launcher.py:23: DeprecationWarning: placeholder_like may not respect the batch dimension. For State objects, use placeholder(state.shape) instead.\n",
      "/Users/thuerey/Library/Python/3.7/lib/python/site-packages/ipykernel_launcher.py:24: DeprecationWarning: placeholder_like may not respect the batch dimension. For State objects, use placeholder(state.shape) instead.\n"
     ]
    }
   ],
   "source": [
    "# one of the most crucial! how many simulation steps to look into the future while training\n",
    "msteps = 4\n",
    "\n",
    "# this is the actual resolution in terms of cells\n",
    "source_res = list(dataset.resolution)\n",
    "# this is only a virtual size, in terms of abstract units for the bounding box of the domain (in practice it's important for conversions or when rescaling to physical units)\n",
    "sim_len = 100.\n",
    "\n",
    "source     =   Fluid(Domain(resolution=source_res, box=box[0:sim_len*2, 0:sim_len], boundaries=OPEN), buoyancy_factor=0, batch_size=batch_size)\n",
    "reference = [ Fluid(Domain(resolution=source_res, box=box[0:sim_len*2, 0:sim_len], boundaries=OPEN), buoyancy_factor=0, batch_size=batch_size) for _ in range(msteps) ]\n",
    "\n",
    "# velocity BC\n",
    "vn = np.zeros(source.velocity.data[0].data.shape)  # st.velocity.data[0] is considered as the velocity field in y axis!\n",
    "vn[..., 0:2, 0:vn.shape[2]-1, 0] = 1.0\n",
    "vn[..., 0:vn.shape[1], 0:1,   0] = 1.0\n",
    "vn[..., 0:vn.shape[1], -1:,   0] = 1.0\n",
    "velBCy = vn\n",
    "velBCyMask = np.copy(vn)  # warning, mask needs to be binary, 0..1, this only works if vel is also 1\n",
    "\n",
    "lr_in =   tf.placeholder(tf.float32, shape=[])  # learning rate\n",
    "Re_in =   tf.placeholder(tf.float32, shape=[batch_size])  # Reynolds numbers\n",
    "\n",
    "source_in =   phi.tf.util.placeholder_like(source)\n",
    "reference_in = [ phi.tf.util.placeholder_like(source) for _ in range(msteps) ]\n",
    "\n",
    "network = network_small(to_keras(source_in, Re_in)) # use small network for testing\n",
    "#network = network_medium(to_keras(source_in, Re_in)) # optionally switch to larger network\n",
    "network.summary() \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AbpNPzplQZMF"
   },
   "source": [
    "## Interleaving simulation and NN\n",
    "\n",
    "Now comes the **most crucial** step in the whole setup: we define the chain of simulation steps and network evaluations to be used at training time. After all the work defining helper functions, it's acutally pretty simple: we loop over `msteps`, call the simulator via `KarmanFlow.step` for an input state, and afterwards evaluate the correction via `network(to_keras(...))`. The NN correction is then added to the last simulation state in the `prediction` list (we're actually simply overwriting the last simulated step `prediction[-1]` with `velocity + correction[-1]`.\n",
    "\n",
    "One other important thing that's happening here is normalization: the inputs to the network are divided by the standard deviations in `dataset.dataStats`. This is slightly complicated as we have to append the scaling for the Reynolds numbers to the normalization for the velocity. After evaluating the `network`, we only have a velocity left, so we can simply multiply by the standard deviation of the velocity again (via `* dataset.dataStats['std'][1]`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D5NeMcLGQaxh",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prediction, correction = [], []\n",
    "for i in range(msteps):\n",
    "    prediction += [\n",
    "        KarmanFlow().step(\n",
    "            source_in if i==0 else prediction[-1],\n",
    "            re=Re_in,\n",
    "            res=source_res[-1],  # reference resolution is size in x direction\n",
    "            velBCy=velBCy,\n",
    "            velBCyMask=velBCyMask\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    correction += [\n",
    "        to_staggered(\n",
    "            network(\n",
    "                to_keras(prediction[-1], Re_in)/[\n",
    "                    *(dataset.dataStats['std'][1]),  # velocity\n",
    "                    dataset.dataStats['ext.std'][0]  # Re\n",
    "                ]\n",
    "            ) * dataset.dataStats['std'][1],\n",
    "            box=source.velocity.box\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    prediction[-1] = prediction[-1].copied_with(velocity=prediction[-1].velocity + correction[-1])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c4yLlDM3QfUR"
   },
   "source": [
    "We also need to define a loss function for training. Here, we can simply use an $L^2$ loss over the whole sequence, i.e. the iteration over `msteps`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "K2JcO3-QQgC9"
   },
   "outputs": [],
   "source": [
    "loss_steps = [\n",
    "    tf.nn.l2_loss(\n",
    "        (reference_in[i].velocity.staggered_tensor() - prediction[i].velocity.staggered_tensor())\n",
    "        /dataset.dataStats['std'][1]\n",
    "    )\n",
    "    for i in range(msteps)\n",
    "]\n",
    "loss = tf.reduce_sum(loss_steps)/msteps\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E6Vly1_0QhZ1"
   },
   "source": [
    "## Training\n",
    "\n",
    "For the training, we use a standard Adam optimizer, and only run 4 epochs by default. This should be increased for the larger network or to obtain more accurate results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PuljFamYQksW"
   },
   "outputs": [],
   "source": [
    "lr = 1e-4\n",
    "adapt_lr = True\n",
    "resume = 0 # set to 1 to load existing network\n",
    "epochs = 4\n",
    "\n",
    "opt = tf.compat.v1.train.AdamOptimizer(learning_rate=lr_in)\n",
    "train_step = opt.minimize(loss)\n",
    "\n",
    "tf_session = tf.Session() \n",
    "scene = Scene.create(\".\", count=batch_size, mkdir=False, copy_calling_script=False) # not really used here\n",
    "sess = Session(scene, session=tf_session)\n",
    "tf.compat.v1.keras.backend.set_session(tf_session)\n",
    "\n",
    "sess.initialize_variables()\n",
    "\n",
    "# optional, load existing network...\n",
    "if resume>0: \n",
    "    ld_network = keras.models.load_model('./nn_epoch{:04d}.h5'.format(resume))\n",
    "    network.set_weights(ld_network.get_weights())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p8hUXJDkRQST"
   },
   "source": [
    "As this setups supports an (optional) fairly accurate training with the medium sized network from above, we'll define a helper function for scheduling learning rate decay. This helps to make the network converge later on in the training. The steps below (after 10,15 etc. epochs) were determined heuristically from previous runs, and are not active by the default run with `epochs=4` (feel free to change and experiment to train more accurate networks)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "Am3hSdNgRPEh"
   },
   "outputs": [],
   "source": [
    "def lr_schedule(epoch, current_lr):\n",
    "    lr = current_lr\n",
    "    if   epoch == 25: lr *= 0.5\n",
    "    elif epoch == 20: lr *= 1e-1\n",
    "    elif epoch == 15: lr *= 1e-1\n",
    "    elif epoch == 10: lr *= 1e-1\n",
    "    return lr\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lrALctV1RWBO"
   },
   "source": [
    "Finally, we can start training the NN! This is very straight forward now, we simply loop over the desired number of iterations, get a batch each time via `getData`, feed it into the source simulation input `source_in`, and compare it in the loss with the `reference` data for the batch.\n",
    "\n",
    "The setup above will automatically take care that the differentiable physics solver used here provides the right gradient information, and provides it to the tensorflow network. Be warned: due to the complexity of the setup, this training run can take a while... (If you have a saved `final.h5` model from a previous run, you can potentially skip this block and load the previously trained model instead.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "m3Nd8YyHRVFQ",
    "outputId": "148d951b-7070-4a95-c6d7-0fd91d29606e"
   },
   "outputs": [],
   "source": [
    "current_lr = lr\n",
    "steps = 0\n",
    "for j in range(epochs):  # training\n",
    "    dataset.newEpoch(exclude_tail=msteps)\n",
    "    if j<resume:\n",
    "        print('resume: skipping {} epoch'.format(j+1))\n",
    "        steps += dataset.numSteps*dataset.numBatches\n",
    "        continue\n",
    "\n",
    "    current_lr = lr_schedule(j, current_lr) if adapt_lr else lr\n",
    "    for ib in range(dataset.numBatches):   \n",
    "        for i in range(dataset.numSteps): \n",
    "            batch = getData(dataset, consecutive_frames=msteps) # should be dataset.getData\n",
    "            re_nr = batch[2] # Reynolds numbers\n",
    "            source = source.copied_with(density=batch[0][0], velocity=batch[1][0])\n",
    "            reference = [ reference[k].copied_with(density=batch[0][k+1], velocity=batch[1][k+1]) for k in range(msteps) ]\n",
    "\n",
    "            my_feed_dict = { source_in: source, Re_in: re_nr, lr_in: current_lr }\n",
    "            my_feed_dict.update(zip(reference_in, reference))\n",
    "            _, l2 = sess.run([train_step, loss], my_feed_dict)\n",
    "            steps += 1\n",
    "\n",
    "            if (j==0 and i<3) or (j==0 and ib==0 and i%31==0) or (ib==0 and i%124==0):\n",
    "              print('epoch {:03d}/{:03d}, batch {:03d}/{:03d}, step {:04d}/{:04d}: loss={}'.format( j+1, epochs, ib+1, dataset.numBatches, i+1, dataset.numSteps, l2 ))\n",
    "            dataset.nextStep()\n",
    "\n",
    "        dataset.nextBatch()\n",
    "\n",
    "    if j%10==9: network.save('./nn_epoch{:04d}.h5'.format(j+1))\n",
    "\n",
    "# all done! save final version\n",
    "network.save('./final.h5')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "swG7GeDpWT_Z"
   },
   "source": [
    "The loss should go down from above 1000 initially to below 10. This is a good sign, but of course it's even more important to see how the resulting solver fares on new inputs.\n",
    "\n",
    "Note that with this training approach we've realized a hybrid solver, consisting of a regular _source_ simulator, and a network that was trained to specificially interact with this simulator for a chosen domain of simulation cases.\n",
    "\n",
    "Let's see how well this works by applying it to a set of test data inputs with new Reynolds numbers that were not part of the training data.\n",
    "\n",
    "To keep things somewhat simple, we won't aim for a high-performance version of our hybrid solver. For performance, please check out the external code base: the network trained here should be directly useable in [this apply script](https://github.com/tum-pbs/Solver-in-the-Loop/blob/master/karman-2d/karman_apply.py).\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation \n",
    "\n",
    "In order to evaluate the performance of our DL-powered solver, we essentially only need to repeat the inner loop of each training iteration for more steps. While we were limited to `msteps` evaluations at training time, we can now run our solver for arbitrary lengths. This is a good test for how well our solver has learned to keep the data within the desired distribution, and represents a generalization test for longer rollouts.\n",
    "\n",
    "We can reuse the solver code from above, but in the following, we will consider two simulated versions: for comparison, we'll run one reference simulation in the _source_ space (i.e., without any modifications). This version receives the regular outputs of each evaluation of the simulator, and ignores the learned correction (denoted as `sourcesim` below). The second version, `prediction`, repeatedly computes the source solver plus the learned correction, and advances this state in the solver.\n",
    "\n",
    "A subtle but important point: we still have to use the normalization from the original training data set here, i.e., the `dataset.dataStats['std']` values. Below we'll create a new test data set with it's own mean and standard deviation, but the model never saw this data before. It was trained with the data in `dataset` above, and hence we have to use the constants from there to make sure the model receives values that it can relate to the data it was trained with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RumKebW_05xp"
   },
   "outputs": [],
   "source": [
    "# similar to msteps from before\n",
    "evalsteps=1\n",
    "\n",
    "sourcesim_in = phi.tf.util.placeholder_like(source) # reuse source from above\n",
    "sourcesim, prediction, correction = [], [], []\n",
    "\n",
    "for i in range(evalsteps):\n",
    "    sourcesim += [\n",
    "        KarmanFlow().step(\n",
    "            sourcesim_in if i==0 else prediction[-1],\n",
    "            re=Re_in,\n",
    "            res=source_res[-1], \n",
    "            velBCy=velBCy,\n",
    "            velBCyMask=velBCyMask\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    # important: use normalization from _original_ dataset here, not the testdata below\n",
    "    correction += [\n",
    "        to_staggered(\n",
    "            network(\n",
    "                to_keras(sourcesim[-1], Re_in)/[\n",
    "                    *(dataset.dataStats['std'][1]),  # velocity\n",
    "                    dataset.dataStats['ext.std'][0]  # Re\n",
    "                ]\n",
    "            ) * dataset.dataStats['std'][1],\n",
    "            box=source.velocity.box\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    prediction += [sourcesim[-1].copied_with(velocity=sourcesim[-1].velocity + correction[-1])]\n",
    "   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we also need a set of new data to apply the learned solver to. Below, we'll download a new set of Reynolds numbers (inbetween the ones used for training), so that we can later on run the unmodified simulator and the DL-powered one on these cases.\n",
    "\n",
    "We're creating a new dataset object `dataset_test` here, which organizes the data. We're simply using the first batch of the unshuffled dataset, though."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded test data, 3 training sims\n",
      "dict_keys(['karman-fdt-hires-testset-reduced/sim_000000', 'karman-fdt-hires-testset-reduced/sim_000001', 'karman-fdt-hires-testset-reduced/sim_000002'])\n",
      "Reynolds numbers in test data set: [240000.0, 960000.0, 3840000.0]\n"
     ]
    }
   ],
   "source": [
    "if not os.path.isfile('sol-data-karman-2d-test.pickle'):\n",
    "  import urllib.request\n",
    "  url=\"https://ge.in.tum.de/download/2020-solver-in-the-loop/sol-data-karman-2d-test.pickle\"\n",
    "  print(\"Downloading test data (38MB), this can take a moment the first time...\")\n",
    "  urllib.request.urlretrieve(url, 'sol-data-karman-2d-test.pickle')\n",
    "\n",
    "with open('sol-data-karman-2d-test.pickle', 'rb') as f: data_test_preloaded = pickle.load(f)\n",
    "print(\"Loaded test data, {} training sims\".format(len(data_test_preloaded)) )\n",
    "\n",
    "print(format(data_test_preloaded.keys()))\n",
    "dataset_test = Dataset( data_preloaded=data_test_preloaded, is_testset=True, num_frames=simsteps, num_sims=3, batch_size=batch_size )\n",
    "\n",
    "# we only need 1 batch to start with\n",
    "dataset_test.newEpoch(shuffle_data=False)\n",
    "batch = getData(dataset_test, consecutive_frames=msteps) \n",
    "re_nr = batch[2] # Reynolds numbers\n",
    "source_test = source.copied_with(density=batch[0][0], velocity=batch[1][0])\n",
    "reference = [ reference[k].copied_with(density=batch[0][k+1], velocity=batch[1][k+1]) for k in range(msteps) ]\n",
    "\n",
    "print(\"Reynolds numbers in test data set: \"+format(re_nr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can evaluate the _source_ simulation, and the _prediction_ of our learned hybrid solver for 120 steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source simulation steps 121\n"
     ]
    }
   ],
   "source": [
    "steps_source = [source_test]\n",
    "\n",
    "for i in range(120):\n",
    "    my_feed_dict = { sourcesim_in: steps_source[-1], Re_in: re_nr, lr_in: current_lr }\n",
    "    sourcesim_out = sess.run([sourcesim[0]], my_feed_dict) \n",
    "    steps_source.append( sourcesim_out[0] )\n",
    "\n",
    "print(\"Source simulation steps \"+format(len(steps_source)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The way the evaluation is implemented above, we can only evalute a proper source simulation (i.e. one that is not modified by the NN) for `evalsteps=1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Steps with hybrid solver 121\n"
     ]
    }
   ],
   "source": [
    "steps_pred = [source_test]\n",
    "assert(evalsteps==1) # for other settings, deactivate the source evaluation below\n",
    "\n",
    "for i in range(120):\n",
    "    my_feed_dict = { sourcesim_in: steps_pred[-1], Re_in: re_nr, lr_in: current_lr }\n",
    "    prediction_out = sess.run([prediction[0]], my_feed_dict)  # w corr\n",
    "    steps_pred.append( prediction_out[0] )\n",
    "\n",
    "print(\"Steps with hybrid solver \"+format(len(steps_pred)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the differences of the two outputs by plotting the y component of the velocities over time. The two following code cells show six velocity snapshots for the batch index `b` in intervals of 20 time steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53140 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab\n",
    "b = 1 # batch index for the following comparisons\n",
    "\n",
    "fig, axes = pylab.subplots(1, 6, figsize=(16, 5))\n",
    "for i in range(1,7):\n",
    "  v = steps_source[i*20].velocity.data[0].data[b,:,:,0] \n",
    "  axes[i-1].imshow( v , origin='lower', cmap='magma')\n",
    "  axes[i-1].set_title(f\" Source simulation t={i*20} \")\n",
    "pylab.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57712 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = pylab.subplots(1, 6, figsize=(16, 5))\n",
    "for i in range(1,7):\n",
    "  v = steps_pred[i*20].velocity.data[0].data[b,:,:,0] \n",
    "  axes[i-1].imshow( v , origin='lower', cmap='magma')\n",
    "  axes[i-1].set_title(f\" Hybrid solver t={i*20} \")\n",
    "pylab.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "They both start out with the same initial state at $t=0$ (a downsampled solution from the reference solution manifold), and at $t=20$ the solutions still share similarities. Over time, the source version strongly diffuses the structures in the flow and looses momentum. The flow behind the obstacles becomes straight, and lacks clear vortices. \n",
    "\n",
    "The version produced by the hybrid solver does much better. It preserves the vortex shedding even after more than one hundred updates. Note that both outputs were produced by the same underlying solver. The second version just profits from the learned corrector which has learned to revert the overly strong dissipation of the solver. However, it also produces some visible axis-aligned structures, especially near the sides of the domain. (This could be alleviated with improved training setups, e.g., more look-ahead, and a larger model.)\n",
    "\n",
    "We can also compare and quantify how the models do w.r.t. reference data. The next cell plots one time step of the three versions: the reference data after 50 steps, and the re-simulated version of the source and our hybrid solver, together with a per-cell error of the two:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 51060 characters]\n",
      "text/plain": [
       "<Figure size 1008x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "index = 50 # time step index\n",
    "\n",
    "ref = [ math.concat( dataset_test.dataPreloaded[ dataset_test.dataSims[b] ][ j ][1] , axis=0) for j in range(121) ]\n",
    "\n",
    "fig, axes = pylab.subplots(1, 4, figsize=(14, 5))\n",
    "\n",
    "v = ref[index][:,:,0] \n",
    "axes[0].imshow( v , origin='lower', cmap='magma')\n",
    "axes[0].set_title(f\" Reference \")\n",
    "\n",
    "v = steps_source[index].velocity.data[0].data[b,:,:,0] \n",
    "axes[1].imshow( v , origin='lower', cmap='magma')\n",
    "axes[1].set_title(f\" Source \")\n",
    "\n",
    "v = steps_pred[index].velocity.data[0].data[b,:,:,0] \n",
    "axes[2].imshow( v , origin='lower', cmap='magma')\n",
    "axes[2].set_title(f\" Learned \")\n",
    "\n",
    "# show error side by side\n",
    "err_source = ref[index][:,1:-2,0] - steps_source[index].velocity.data[0].data[b,:,1:-1,0] \n",
    "err_pred   = ref[index][:,1:-2,0] - steps_pred[index].velocity.data[0].data[b,:,1:-1,0] \n",
    "v = np.concatenate([err_source,err_pred], axis=1)\n",
    "axes[3].imshow( v , origin='lower', cmap='magma')\n",
    "axes[3].set_title(f\" Errors: Source & Learned\")\n",
    "\n",
    "pylab.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows very clearly how the pure source simulation in the middle deviates from the reference on the left. Apart from the slight vertical streaks, the learned version stays much closer to the reference solution. \n",
    "\n",
    "The two per-cell error images on the right also illustrate this: the source version has much larger errors that show how it systematically underestimates the vortices that should form. The error for the learned version is much more evenly distributed and significantly smaller.\n",
    "\n",
    "Next, we can quantify these observations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAE for source: 0.05912385508418083 , and hybrid: 0.0402643121778965\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21180 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "errors_source, errors_pred = [], []\n",
    "for index in range(100):\n",
    "  vs = ref[index][:,1:-2,0] - steps_source[index].velocity.data[0].data[b,:,1:-1,0] \n",
    "  vp = ref[index][:,1:-2,0] - steps_pred[index].velocity.data[0].data[b,:,1:-1,0] \n",
    "  errors_source.append(np.mean(np.abs(vs))) \n",
    "  errors_pred.append(np.mean(np.abs(vp)))\n",
    "  #print(\"Errors for step \"+format(index)+\" source: \"+format(np.mean(np.abs(vs))) +\" , hybrid: \"+format(np.mean(np.abs(vp))))\n",
    "\n",
    "fig = pylab.figure().gca()\n",
    "pltx = np.linspace(0,99,100)\n",
    "fig.plot(pltx, errors_source, lw=2, color='mediumblue', label='Source')  \n",
    "fig.plot(pltx, errors_pred,   lw=2, color='green', label='Hybrid')\n",
    "pylab.xlabel('Time step'); pylab.ylabel('Error'); fig.legend()\n",
    "\n",
    "print(\"MAE for source: \"+format(np.mean(errors_source)) +\" , and hybrid: \"+format(np.mean(errors_pred)) )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The overall mean absolute error (MAE) is almost 50% larger for the regular simulation. The graph below shows this behavior over time: the error of the source version is clearly larger than the errors of the hybrid simulator.\n",
    "\n",
    "This concludes our evaluation. Note that the improved behavior of the hybrid solver is difficult to reliably measure with simple vector norms such as an MAE or $L^2$ norm. To improve this, we'd need to employ other, domain-specific metrics. In this case, metrics for fluids based on vorticity and turbulence properties of the flow would be applicable. However, in this text, we instead want to focus on DL-related topics and target another inverse problem with differentiable physics solvers in the next chapter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Next steps\n",
    "\n",
    "* Modify the training to further reduce the training error. With the _medium_ network you should be able to get the loss down to around 1.\n",
    "\n",
    "* Use the external github code to generate new test data, and run your model on these cases. You'll see that a reduced training error not always directly correlates with an improved test performance.\n",
    "\n",
    "* Turn off the differentiable physics training (by setting `msteps=1`), and compare it with the DP version.\n",
    "\n",
    "* Likewise, train a model with a larger `msteps` setting, e.g., 8 or 16. Note that due to the recurrent nature of the training, you'll probabaly have to load a pre-trained state to stabilize the first iterations."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "diffphys-code-sol.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
