{
 "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}) ) , \\mathcal{T} \\vr{t+1}) < \\loss ( \\mathcal{P}_{s}( \\mathcal{T} \\vr{t} ), \\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{arg min}_\\theta \\big( ( \\mathcal{P}_{s} \\corr )^n ( \\mathcal{T} \\vr{t} ) - \\mathcal{T} \\vr{t+n} \\big)^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": {
    "id": "hPgwGkzYdIww"
   },
   "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": "30ab90f0-4b0c-4451-81da-f85887aeb7b9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading training data (73MB), this can take a moment the first time...\n",
      "Loaded data, 6 training sims\n"
     ]
    }
   ],
   "source": [
    "import os, sys, logging, argparse, pickle, glob, random, distutils.dir_util, urllib.request\n",
    "\n",
    "fname_train = 'sol-karman-2d-train.pickle'\n",
    "if not os.path.isfile(fname_train):\n",
    "  print(\"Downloading training data (73MB), this can take a moment the first time...\")\n",
    "  urllib.request.urlretrieve(\"https://physicsbaseddeeplearning.org/data/\"+fname_train, fname_train)\n",
    "\n",
    "with open(fname_train, 'rb') as f: data_preloaded = pickle.load(f)\n",
    "print(\"Loaded data, {} training sims\".format(len(data_preloaded)) )\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 \ud83d\ude42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BGN4GqxkIueM",
    "outputId": "d934bf06-b6b9-41ce-be11-d1d6d3561c89"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": []
    }
   ],
   "source": [
    "#!pip install --upgrade --quiet phiflow\n",
    "#!pip uninstall phiflow\n",
    "!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@develop\n",
    "\n",
    "from phi.tf.flow import *\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "\n",
    "random.seed(42) \n",
    "np.random.seed(42)\n",
    "tf.random.set_seed(42)"
   ]
  },
  {
   "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 (sizes in the imlpementation are using an additional factor of 100).\n",
    "```\n",
    "\n",
    "The solver applies inflow boundary conditions for the y-velocity with a pre-multiplied mask (`vel_BcMask`), to set the y components at the bottom of the domain during the simulation step. This mask is created with the `HardGeometryMask` from phiflow, which initializes the spatially shifted entries for the components of a staggered grid correctly. The simulation step is quite straight forward: it computes contributions for viscosity, inflow, advection and finally makes the resulting motion divergence free via an implicit pressure solve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "6WNMcdWUw4EP"
   },
   "outputs": [],
   "source": [
    "class KarmanFlow():\n",
    "    def __init__(self, domain):\n",
    "        self.domain = domain\n",
    "\n",
    "        self.vel_BcMask = self.domain.staggered_grid(HardGeometryMask(Box[:5, :]) )\n",
    "    \n",
    "        self.inflow = self.domain.scalar_grid(Box[5:10, 25:75])         # scale with domain if necessary!\n",
    "        self.obstacles = [Obstacle(Sphere(center=[50, 50], radius=10))] \n",
    "\n",
    "    def step(self, density_in, velocity_in, re, res, buoyancy_factor=0, dt=1.0):\n",
    "        velocity = velocity_in\n",
    "        density = density_in\n",
    "\n",
    "        # viscosity\n",
    "        velocity = phi.flow.diffuse.explicit(field=velocity, diffusivity=1.0/re*dt*res*res, dt=dt)\n",
    "        \n",
    "        # inflow boundary conditions\n",
    "        velocity = velocity*(1.0 - self.vel_BcMask) + self.vel_BcMask * (1,0)\n",
    "\n",
    "        # advection \n",
    "        density = advect.semi_lagrangian(density+self.inflow, velocity, dt=dt)\n",
    "        velocity = advected_velocity = advect.semi_lagrangian(velocity, velocity, dt=dt)\n",
    "\n",
    "        # mass conservation (pressure solve)\n",
    "        pressure = None\n",
    "        velocity, pressure = fluid.make_incompressible(velocity, self.obstacles)\n",
    "        self.solve_info = { 'pressure': pressure, 'advected_velocity': advected_velocity }\n",
    "        \n",
    "        return [density, velocity]\n",
    "\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",
    "- 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 small network consisting only of four convolutional layers with ReLU activations (we're also using keras here for simplicity). The input dimensions are determined from input tensor in the `inputs_dict` (it has three channels: u,v, and Re). Then we process the data via three conv layers with 32 features each, before reducing to 2 channels in the output. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "qIrWYTy6xscA"
   },
   "outputs": [],
   "source": [
    "def network_small(inputs_dict):\n",
    "    l_input = keras.layers.Input(**inputs_dict)\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",
    "    block_1 = keras.layers.LeakyReLU()(l_conv2)\n",
    "\n",
    "    l_output = keras.layers.Conv2D(filters=2,  kernel_size=5, padding='same')(block_1) # u, v\n",
    "    return keras.models.Model(inputs=l_input, outputs=l_output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YfHvdI7yxtdj"
   },
   "source": [
    "For flexibility (and larger-scale tests later on), let's also 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(inputs_dict):\n",
    "    l_input = keras.layers.Input(**inputs_dict)\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 two vector components via `vector['x']` and `vector['y']` to discard the outermost layer of the velocity field grids. This gives two tensors of equal size that can be combined. \n",
    "It then adds a constant channel via `math.ones` that is multiplied by the desired Reynolds number in `ext_const_channel`. The resulting stack of grids is stacked along the `channels` dimensions, and represents an  input to the neural network. \n",
    "\n",
    "After network evaluation, we transform the output tensor back into a phiflow grid via the `to_phiflow` 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 `domain.staggered_grid` function, which also sets physical size and boundary conditions as given by the domain object)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "hhGFpTjGyRyg"
   },
   "outputs": [],
   "source": [
    "\n",
    "def to_keras(dens_vel_grid_array, ext_const_channel):\n",
    "    # align the sides the staggered velocity grid making its size the same as the centered grid\n",
    "    return math.stack(\n",
    "        [\n",
    "            math.pad( dens_vel_grid_array[1].vector['x'].values, {'x':(0,1)} , math.extrapolation.ZERO),\n",
    "            dens_vel_grid_array[1].vector['y'].y[:-1].values,         # v\n",
    "            math.ones(dens_vel_grid_array[0].shape)*ext_const_channel # Re\n",
    "        ],\n",
    "        math.channel('channels')\n",
    "    )\n",
    "\n",
    "def to_phiflow(tf_tensor, domain):\n",
    "    return domain.staggered_grid(\n",
    "        math.stack(\n",
    "            [\n",
    "                math.tensor(tf.pad(tf_tensor[..., 1], [(0,0), (0,1), (0,0)]), math.batch('batch'), math.spatial('y, x')), # v\n",
    "                math.tensor( tf_tensor[...,:-1, 0], math.batch('batch'), math.spatial('y, x')), # u \n",
    "            ], math.channel('vector')\n",
    "        )\n",
    "    )\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., 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/sim_%06d'%i for i in range(num_sims) ]\n",
    "            ReNrs = [120000.0, 480000.0, 1920000.0, 7680000.0] \n",
    "            self.extConstChannelPerSim = { self.dataSims[i]:[ReNrs[i]] for i in range(num_sims) }\n",
    "\n",
    "        self.dataFrames = [ np.arange(num_frames) for _ in self.dataSims ]  \n",
    "\n",
    "        # debugging example, check shape of a single marker density field:\n",
    "        #print(format(self.dataPreloaded[self.dataSims[0]][0][0].shape )) \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)), # density\n",
    "                np.std(np.concatenate([np.absolute(self.dataPreloaded[asim][i][1].reshape(-1)) for asim in self.dataSims for i in range(num_frames)], axis=-1)), # x-velocity\n",
    "                np.std(np.concatenate([np.absolute(self.dataPreloaded[asim][i][2].reshape(-1)) for asim in self.dataSims for i in range(num_frames)], axis=-1)), # y-velocity\n",
    "            )\n",
    "        }\n",
    "        self.dataStats.update({\n",
    "            'ext.std': [ np.std([np.absolute(self.extConstChannelPerSim[asim][0]) for asim in self.dataSims]) ] # Reynolds Nr\n",
    "        })\n",
    "\n",
    "        \n",
    "        if not is_testset:\n",
    "            print(\"Data stats: \"+format(self.dataStats))\n",
    "\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.dataFrames[i][0:(len(self.dataFrames[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 order of 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",
    "def getData(self, consecutive_frames):\n",
    "    d_hi = [\n",
    "        np.concatenate([\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 # frames\n",
    "            ][0]\n",
    "            for i in range(self.batchSize)\n",
    "        ], axis=0) for j in range(consecutive_frames+1)\n",
    "    ]\n",
    "    u_hi = [\n",
    "        np.concatenate([\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 # frames\n",
    "            ][1]\n",
    "            for i in range(self.batchSize)\n",
    "        ], axis=0) for j in range(consecutive_frames+1)\n",
    "    ]\n",
    "    v_hi = [\n",
    "        np.concatenate([\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 # frames\n",
    "            ][2]\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 [d_hi, u_hi, v_hi, 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": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "59EBdEdj0QR2",
    "outputId": "d9282614-d514-47d8-b911-c262c81c252e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data stats: {'std': (2.6542656, 0.23155601, 0.3066732), 'ext.std': [1732512.6262166172]}\n"
     ]
    }
   ],
   "source": [
    "nsims = 6\n",
    "batch_size = 3\n",
    "simsteps = 500\n",
    "\n",
    "dataset = Dataset( data_preloaded=data_preloaded, num_frames=simsteps, num_sims=nsims, batch_size=batch_size )"
   ]
  },
  {
   "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 resolution of the simulation in `source_res`, and allocate the fluid solver object called `simulator`. In order to create grids, it requires access to a `Domain` object, which mostly exists for convenience purposes: it stores resolution, physical size in `bounds`, and boundary conditions of the domain. This information needs to be passed to every grid, and hence it's convenient to have it in one place in the form of the `Domain`. For the setup described above, we need different boundary conditions along x and y: closed walls, and free flow in and out of the domain, respecitvely.\n",
    "\n",
    "We also instantiate the actual NN `network` in the next cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EjgkdCzKP2Ip",
    "outputId": "2d4b34f6-2d40-4273-fc2c-1dac7fe786cb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_1 (InputLayer)         [(None, 64, 32, 3)]       0         \n",
      "_________________________________________________________________\n",
      "conv2d (Conv2D)              (None, 64, 32, 32)        2432      \n",
      "_________________________________________________________________\n",
      "leaky_re_lu (LeakyReLU)      (None, 64, 32, 32)        0         \n",
      "_________________________________________________________________\n",
      "conv2d_1 (Conv2D)            (None, 64, 32, 32)        25632     \n",
      "_________________________________________________________________\n",
      "leaky_re_lu_1 (LeakyReLU)    (None, 64, 32, 32)        0         \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (None, 64, 32, 32)        25632     \n",
      "_________________________________________________________________\n",
      "leaky_re_lu_2 (LeakyReLU)    (None, 64, 32, 32)        0         \n",
      "_________________________________________________________________\n",
      "conv2d_3 (Conv2D)            (None, 64, 32, 2)         1602      \n",
      "=================================================================\n",
      "Total params: 55,298\n",
      "Trainable params: 55,298\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\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 a virtual size, in terms of abstract units for the bounding box of the domain (it's important for conversions or when rescaling to physical units)\n",
    "simulation_length = 100.\n",
    "\n",
    "# for readability\n",
    "from phi.physics._boundaries import Domain, OPEN, STICKY as CLOSED\n",
    "\n",
    "boundary_conditions = {\n",
    "    'x':(phi.physics._boundaries.STICKY,phi.physics._boundaries.STICKY), \n",
    "    'y':(phi.physics._boundaries.OPEN,  phi.physics._boundaries.OPEN) }\n",
    "\n",
    "domain = Domain(y=source_res[0], x=source_res[1], bounds=Box[0:2*simulation_length, 0:simulation_length], boundaries=boundary_conditions)\n",
    "simulator = KarmanFlow(domain=domain)\n",
    "\n",
    "network = network_small(dict(shape=(source_res[0],source_res[1], 3)))\n",
    "network.summary()\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 a function that encapsulates the chain of simulation steps and network evaluations in each training step. After all the work defining helper functions, it's actually pretty simple: we create a gradient tape via `tf.GradientTape()` such that we can backpropagate later on. We then loop over `msteps`, call the simulator via `simulator.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 velocity `prediction[-1][1]` with `prediction[-1][1] + 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`. 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]` and `[2]`).\n",
    "\n",
    "The `training_step` function also directly evaluates and returns the loss. Here, we can simply use an $L^2$ loss over the whole sequence, i.e. the iteration over `msteps`. This is requiring a few lines of code because we separately loop over 'x' and 'y' components, in order to normalize and compare to the ground truth values from the training data set.\n",
    "\n",
    "The \"learning\" happens in the last two lines via `tape.gradient()` and `opt.apply_gradients()`, which then contain the aggregated information about how to change the NN weights to nudge the simulation closer to the reference for the full chain of simulation steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "D5NeMcLGQaxh",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def training_step(dens_gt, vel_gt, Re, i_step):\n",
    "    with tf.GradientTape() as tape:\n",
    "        prediction, correction = [ [dens_gt[0],vel_gt[0]] ], [0] # predicted states with correction, inferred velocity corrections\n",
    "\n",
    "        for i in range(msteps):\n",
    "            prediction += [\n",
    "                simulator.step(\n",
    "                    density_in=prediction[-1][0],\n",
    "                    velocity_in=prediction[-1][1],\n",
    "                    re=Re, res=source_res[1],\n",
    "                )\n",
    "            ]       # prediction: [[density1, velocity1], [density2, velocity2], ...]\n",
    "\n",
    "            model_input = to_keras(prediction[-1], Re)\n",
    "            model_input /= math.tensor([dataset.dataStats['std'][1], dataset.dataStats['std'][2], dataset.dataStats['ext.std'][0]], channel('channels')) # [u, v, Re]\n",
    "            model_out = network(model_input.native(['batch', 'y', 'x', 'channels']), training=True)\n",
    "            model_out *= [dataset.dataStats['std'][1], dataset.dataStats['std'][2]] # [u, v]\n",
    "            correction += [ to_phiflow(model_out, domain) ]                         # [velocity_correction1, velocity_correction2, ...]\n",
    "\n",
    "            prediction[-1][1] = prediction[-1][1] + correction[-1]\n",
    "            #prediction[-1][1] = correction[-1]\n",
    "\n",
    "        # evaluate loss\n",
    "        loss_steps_x = [\n",
    "            tf.nn.l2_loss(\n",
    "                (\n",
    "                    vel_gt[i].vector['x'].values.native(('batch', 'y', 'x'))\n",
    "                    - prediction[i][1].vector['x'].values.native(('batch', 'y', 'x'))\n",
    "                )/dataset.dataStats['std'][1]\n",
    "            )\n",
    "            for i in range(1,msteps+1)\n",
    "        ]\n",
    "        loss_steps_x_sum = tf.math.reduce_sum(loss_steps_x)\n",
    "\n",
    "        loss_steps_y = [\n",
    "            tf.nn.l2_loss(\n",
    "                (\n",
    "                    vel_gt[i].vector['y'].values.native(('batch', 'y', 'x'))\n",
    "                    - prediction[i][1].vector['y'].values.native(('batch', 'y', 'x'))\n",
    "                )/dataset.dataStats['std'][2]\n",
    "            )\n",
    "            for i in range(1,msteps+1)\n",
    "        ]\n",
    "        loss_steps_y_sum = tf.math.reduce_sum(loss_steps_y)\n",
    "\n",
    "        loss = (loss_steps_x_sum + loss_steps_y_sum)/msteps\n",
    "\n",
    "        gradients = tape.gradient(loss, network.trainable_variables)\n",
    "        opt.apply_gradients(zip(gradients, network.trainable_variables))\n",
    "\n",
    "        return math.tensor(loss)    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c4yLlDM3QfUR"
   },
   "source": [
    "Once defined, we can prepare this function for executing the training step by calling phiflow's `math.jit_compile()` function. It automatically maps to the correct pre-compilation step of the chosen backend. E.g., for TF this internally creates a computational graph, and optimizes the chain of operations. For JAX, it can even compile optimized GPU code (if JAX is set up correctly). Using the jit compilation can make a huge difference in terms of runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "K2JcO3-QQgC9"
   },
   "outputs": [],
   "source": [
    "\n",
    "training_step_jit = math.jit_compile(training_step)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E6Vly1_0QhZ1"
   },
   "source": [
    "## Training\n",
    "\n",
    "For the training, we use a standard Adam optimizer, and run 15 epochs by default. This should be increased for the larger network or to obtain more accurate results. For longer training runs, it would also be beneficial to decrease the learning rate over the course of the epochs, but for simplicity, we'll keep `LR` constant here.\n",
    "\n",
    "Optionally, this is also the right point to load a network state to resume training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "PuljFamYQksW"
   },
   "outputs": [],
   "source": [
    "LR = 1e-4\n",
    "EPOCHS = 15\n",
    "\n",
    "opt = tf.keras.optimizers.Adam(learning_rate=LR) \n",
    "\n",
    "# optional, load existing network...\n",
    "# set to epoch nr. to load existing network from there\n",
    "resume = 0\n",
    "if resume>0: \n",
    "    ld_network = keras.models.load_model('./nn_epoch{:04d}.h5'.format(resume)) \n",
    "    #ld_network = keras.models.load_model('./nn_final.h5') # or the last one\n",
    "    network.set_weights(ld_network.get_weights())\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 `nn_final.h5` network from a previous run, you can potentially skip this block and load the previously trained model instead via the cell above.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "m3Nd8YyHRVFQ",
    "outputId": "686a3419-d022-4889-c0de-66e4e02953d1",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 001/015, batch 001/002, step 0001/0496: loss=2605.340576171875\n",
      "epoch 001/015, batch 001/002, step 0002/0496: loss=1485.1646728515625\n",
      "epoch 001/015, batch 001/002, step 0003/0496: loss=790.8267211914062\n",
      "epoch 001/015, batch 001/002, step 0129/0496: loss=98.64994049072266\n",
      "epoch 001/015, batch 001/002, step 0257/0496: loss=75.3546142578125\n",
      "epoch 001/015, batch 001/002, step 0385/0496: loss=70.05519104003906\n",
      "epoch 002/015, batch 001/002, step 0401/0496: loss=19.126527786254883\n",
      "epoch 003/015, batch 001/002, step 0401/0496: loss=9.628664016723633\n",
      "epoch 004/015, batch 001/002, step 0401/0496: loss=7.898053169250488\n",
      "epoch 005/015, batch 001/002, step 0401/0496: loss=3.6936004161834717\n",
      "epoch 006/015, batch 001/002, step 0401/0496: loss=3.172729730606079\n",
      "epoch 007/015, batch 001/002, step 0401/0496: loss=2.8511123657226562\n",
      "epoch 008/015, batch 001/002, step 0401/0496: loss=3.4968295097351074\n",
      "epoch 009/015, batch 001/002, step 0401/0496: loss=1.6942076683044434\n",
      "epoch 010/015, batch 001/002, step 0401/0496: loss=1.6551270484924316\n",
      "epoch 011/015, batch 001/002, step 0401/0496: loss=1.9383186101913452\n",
      "epoch 012/015, batch 001/002, step 0401/0496: loss=2.0140795707702637\n",
      "epoch 013/015, batch 001/002, step 0401/0496: loss=1.4174892902374268\n",
      "epoch 014/015, batch 001/002, step 0401/0496: loss=1.2593278884887695\n",
      "epoch 015/015, batch 001/002, step 0401/0496: loss=1.250532627105713\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training done, saved NN\n"
     ]
    }
   ],
   "source": [
    "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",
    "    for ib in range(dataset.numBatches):   \n",
    "        for i in range(dataset.numSteps): \n",
    "\n",
    "            # batch: [[dens0, dens1, ...], [x-velo0, x-velo1, ...], [y-velo0, y-velo1, ...], [ReynoldsNr(s)]]            \n",
    "            batch = getData(dataset, consecutive_frames=msteps)\n",
    "            \n",
    "            dens_gt = [   # [density0:CenteredGrid, density1, ...]\n",
    "                domain.scalar_grid(\n",
    "                    math.tensor(batch[0][k], math.batch('batch'), math.spatial('y, x'))\n",
    "                ) for k in range(msteps+1)\n",
    "            ]\n",
    "\n",
    "            vel_gt = [   # [velocity0:StaggeredGrid, velocity1, ...]\n",
    "                domain.staggered_grid(\n",
    "                    math.stack(\n",
    "                        [\n",
    "                            math.tensor(batch[2][k], math.batch('batch'), math.spatial('y, x')),\n",
    "                            math.tensor(batch[1][k], math.batch('batch'), math.spatial('y, x')),\n",
    "                        ], math.channel('vector')\n",
    "                    )\n",
    "                ) for k in range(msteps+1)\n",
    "            ]\n",
    "            re_nr = math.tensor(batch[3], math.batch('batch'))\n",
    "\n",
    "            loss = training_step_jit(dens_gt, vel_gt, re_nr, math.tensor(steps)) \n",
    "            \n",
    "            steps += 1\n",
    "            if (j==0 and ib==0 and i<3) or (j==0 and ib==0 and i%128==0) or (j>0 and ib==0 and i==400): # reduce output \n",
    "              print('epoch {:03d}/{:03d}, batch {:03d}/{:03d}, step {:04d}/{:04d}: loss={}'.format( j+1, EPOCHS, ib+1, dataset.numBatches, i+1, dataset.numSteps, loss ))\n",
    "            \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('./nn_final.h5'); print(\"Training done, saved NN\")\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 NN-solver combination fares on new inputs. With this training approach we've realized a hybrid solver, consisting of a regular _source_ simulator, and a network that was trained to specifically 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": {
    "id": "0c38ne0UdIxV"
   },
   "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 (stored in `steps_source` below). The second version, repeatedly computes the source solver plus the learned correction, and advances this state in the solver (`steps_hybrid`).\n",
    "\n",
    "We also need a set of new data. Below, we'll download a new set of Reynolds numbers (in between the ones used for training), so that we can later on run the unmodified simulator and the DL-powered one on these cases.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RumKebW_05xp",
    "outputId": "b119bc05-2f9d-4289-c951-f9f12627c7fb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading test data (38MB), this can take a moment the first time...\n",
      "Loaded test data, 4 training sims\n"
     ]
    }
   ],
   "source": [
    "fname_test = 'sol-karman-2d-test.pickle'\n",
    "if not os.path.isfile(fname_test):\n",
    "  print(\"Downloading test data (38MB), this can take a moment the first time...\")\n",
    "  urllib.request.urlretrieve(\"https://physicsbaseddeeplearning.org/data/\"+fname_test, fname_test)\n",
    "\n",
    "with open(fname_test, 'rb') as f: data_test_preloaded = pickle.load(f)\n",
    "print(\"Loaded test data, {} training sims\".format(len(data_test_preloaded)) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rZ9h-gRddIxb"
   },
   "source": [
    "Next we create a new dataset object `dataset_test` that organizes the data. We're simply using the first batch of the unshuffled dataset, though.\n",
    "\n",
    "A subtle but important point: we still have to use the normalization from the original training data set: `dataset.dataStats['std']` values. The test data set has it's own mean and standard deviation, and so the trained NN never saw this data before. The NN was trained with the data in `dataset` above, and hence we have to use the constants from there for normalization to make sure the network receives values that it can relate to the data it was trained with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9OPruTGMdIxe",
    "outputId": "254e71e0-c471-4fec-df6f-f116227d12f3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reynolds numbers in test data set: (120000.0, 480000.0, 1920000.0, 7680000.0) along batch\u1d47\n"
     ]
    }
   ],
   "source": [
    "dataset_test = Dataset( data_preloaded=data_test_preloaded, is_testset=True, num_frames=simsteps, num_sims=4, batch_size=4 )\n",
    "\n",
    "# we only need 1 batch with t=0 states to initialize the test simulations with\n",
    "dataset_test.newEpoch(shuffle_data=False)\n",
    "batch = getData(dataset_test, consecutive_frames=0) \n",
    "\n",
    "re_nr_test = math.tensor(batch[3], math.batch('batch')) # Reynolds numbers\n",
    "print(\"Reynolds numbers in test data set: \"+format(re_nr_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sMqRPg2pdIxh"
   },
   "source": [
    "Next we construct a `math.tensor` as initial state for the centered marker fields, and a staggered grid from the next two indices of the test set batch. Similar to `to_phiflow` above, we can use `phi.math.stack()` to combine two fields of appropriate size as a staggered grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "xK1MEaPqdIxi"
   },
   "outputs": [],
   "source": [
    "source_dens_initial = math.tensor( batch[0][0], math.batch('batch'), math.spatial('y, x'))\n",
    "\n",
    "source_vel_initial = domain.staggered_grid(phi.math.stack([\n",
    "    math.tensor(batch[2][0], math.batch('batch'),math.spatial('y, x')),\n",
    "    math.tensor(batch[1][0], math.batch('batch'),math.spatial('y, x'))], channel('vector')))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KhGVceo6dIxl"
   },
   "source": [
    "Now we can first run the _source_ simulation for 120 steps as baseline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nbTTl15kdIxl",
    "outputId": "14521920-1966-41d6-e3a2-41db9fb2f69d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source simulation steps 121\n"
     ]
    }
   ],
   "source": [
    "source_dens_test, source_vel_test = source_dens_initial, source_vel_initial\n",
    "steps_source = [[source_dens_test,source_vel_test]]\n",
    "\n",
    "# note - math.jit_compile() not useful for numpy solve... hence not necessary\n",
    "for i in range(120):\n",
    "    [source_dens_test,source_vel_test] = simulator.step(\n",
    "        density_in=source_dens_test,\n",
    "        velocity_in=source_vel_test,\n",
    "        re=re_nr_test,\n",
    "        res=source_res[1],\n",
    "    )\n",
    "    steps_source.append( [source_dens_test,source_vel_test] )\n",
    "\n",
    "print(\"Source simulation steps \"+format(len(steps_source)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vQV0qV5pdIxm"
   },
   "source": [
    "Next, we compute the corresponding states of our learned hybrid solver. Here, we closely follow the training code, however, now without any gradient tapes or loss computations. We only evaluate the NN in a forward pass for each simulated state to compute a correction field:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fH5tFfh9dIxn",
    "outputId": "1a3c76f6-e401-479e-911d-4bd58f69dab1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Steps with hybrid solver 121\n"
     ]
    }
   ],
   "source": [
    "source_dens_test, source_vel_test = source_dens_initial, source_vel_initial\n",
    "steps_hybrid = [[source_dens_test,source_vel_test]]\n",
    "        \n",
    "for i in range(120):\n",
    "    [source_dens_test,source_vel_test] = simulator.step(\n",
    "        density_in=source_dens_test,\n",
    "        velocity_in=source_vel_test,\n",
    "        re=math.tensor(re_nr_test),\n",
    "        res=source_res[1],\n",
    "    )\n",
    "    model_input = to_keras([source_dens_test,source_vel_test], re_nr_test )\n",
    "    model_input /= math.tensor([dataset.dataStats['std'][1], dataset.dataStats['std'][2], dataset.dataStats['ext.std'][0]], channel('channels')) # [u, v, Re]\n",
    "    model_out = network(model_input.native(['batch', 'y', 'x', 'channels']), training=False)\n",
    "    model_out *= [dataset.dataStats['std'][1], dataset.dataStats['std'][2]] # [u, v]\n",
    "    correction =  to_phiflow(model_out, domain) \n",
    "    source_vel_test = source_vel_test+correction\n",
    "\n",
    "    steps_hybrid.append( [source_dens_test,source_vel_test+correction] )\n",
    "    \n",
    "print(\"Steps with hybrid solver \"+format(len(steps_hybrid)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tnHYeOfldIxp"
   },
   "source": [
    "Given the stored states, we quantify the improvements that the NN yields, and visualize the results. \n",
    "\n",
    "In the following cells, the index `b` chooses one of the four test simulations (by default index 0, the lowest Re outside the training data range), and computes the accumulated mean absolute error (MAE) over all time steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 318
    },
    "id": "bU-PwcCCdIxq",
    "outputId": "66956540-891f-4af7-bafe-22fbd11d8b47"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAE for source: 0.1363069713115692 , and hybrid: 0.05150971934199333\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 18792 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab\n",
    "b = 0 # batch index for the following comparisons\n",
    "\n",
    "errors_source, errors_pred = [], []\n",
    "for index in range(100):\n",
    "  vx_ref = dataset_test.dataPreloaded[ dataset_test.dataSims[b] ][ index ][1][0,...]\n",
    "  vy_ref = dataset_test.dataPreloaded[ dataset_test.dataSims[b] ][ index ][2][0,...]\n",
    "  vxs = vx_ref - steps_source[index][1].values.vector[1].numpy('batch,y,x')[b,...]\n",
    "  vxh = vx_ref - steps_hybrid[index][1].values.vector[1].numpy('batch,y,x')[b,...]\n",
    "  vys = vy_ref - steps_source[index][1].values.vector[0].numpy('batch,y,x')[b,...] \n",
    "  vyh = vy_ref - steps_hybrid[index][1].values.vector[0].numpy('batch,y,x')[b,...] \n",
    "  errors_source.append(np.mean(np.abs(vxs)) + np.mean(np.abs(vys))) \n",
    "  errors_pred.append(np.mean(np.abs(vxh)) + np.mean(np.abs(vyh)))\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)) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aOQP6iCBdIxs"
   },
   "source": [
    "Due to the complexity of the training, the performance can vary, but typically the overall MAE is ca. 160% larger for the regular simulation compared to the hybrid simulator. \n",
    "The gap is typically even bigger for other Reynolds numbers within the training data range. \n",
    "The graph above also shows this behavior over time.\n",
    "\n",
    "Let's also 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": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 302
    },
    "id": "_3f8uhIIdIxs",
    "outputId": "85675563-7fbb-4476-dc15-a33c2e75e990"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 58660 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "c = 0          # channel selector, x=1 or y=0 \n",
    "interval = 20  # time interval\n",
    "\n",
    "fig, axes = pylab.subplots(1, 6, figsize=(16, 5))    \n",
    "for i in range(0,6):\n",
    "  v = steps_source[i*interval][1].values.vector[c].numpy('batch,y,x')[b,...]\n",
    "  axes[i].imshow( v , origin='lower', cmap='magma')\n",
    "  axes[i].set_title(f\" Source simulation t={i*interval} \")\n",
    "\n",
    "pylab.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 313
    },
    "id": "v2d2WTGedIxt",
    "outputId": "888de2df-1934-4813-ba56-98321c8c37c9"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 58700 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(0,6):\n",
    "  v = steps_hybrid[i*interval][1].values.vector[c].numpy('batch,y,x')[b,...]\n",
    "  axes[i].imshow( v , origin='lower', cmap='magma')\n",
    "  axes[i].set_title(f\" Hybrid solver t={i*interval} \")\n",
    "pylab.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ivS0SUiYdIxt"
   },
   "source": [
    "They both start out with the same initial state at $t=0$ (the 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 manages to revert the numerical errors of the source solver, including its overly strong dissipation. \n",
    "\n",
    "We can also visually compare how the NN does 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": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 349
    },
    "id": "23yyfljqdIxu",
    "outputId": "9051bab9-59dc-4029-9a34-588ffc3aba8a"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50484 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",
    "vx_ref = dataset_test.dataPreloaded[ dataset_test.dataSims[b] ][ index ][1][0,...]\n",
    "vx_src = steps_source[index][1].values.vector[1].numpy('batch,y,x')[b,...]\n",
    "vx_hyb = steps_hybrid[index][1].values.vector[1].numpy('batch,y,x')[b,...]\n",
    "\n",
    "fig, axes = pylab.subplots(1, 4, figsize=(14, 5))\n",
    "\n",
    "axes[0].imshow( vx_ref , origin='lower', cmap='magma')\n",
    "axes[0].set_title(f\" Reference \")\n",
    "\n",
    "axes[1].imshow( vx_src , origin='lower', cmap='magma')\n",
    "axes[1].set_title(f\" Source \")\n",
    "\n",
    "axes[2].imshow( vx_hyb , origin='lower', cmap='magma')\n",
    "axes[2].set_title(f\" Learned \")\n",
    "\n",
    "# show error side by side\n",
    "err_source = vx_ref - vx_src \n",
    "err_hybrid = vx_ref - vx_hyb \n",
    "v = np.concatenate([err_source,err_hybrid], 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": {
    "id": "BZByQsAydIxv"
   },
   "source": [
    "This shows very clearly how the pure source simulation in the middle deviates from the reference on the left. 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 (i.e. brighter colors) 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 in magnitude.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UQTY8m6LdIxv"
   },
   "source": [
    "This concludes our evaluation. Note that the improved behavior of the hybrid solver can be 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": {
    "id": "Dl3vzF_XdIxv"
   },
   "source": [
    "## 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",
    "* Turn off the differentiable physics training (by setting `msteps=1`), and compare it with the DP version.\n",
    "\n",
    "* Likewise, train a network with a larger `msteps` setting, e.g., 8 or 16. Note that due to the recurrent nature of the training, you'll probably have to load a pre-trained state to stabilize the first iterations.\n",
    "\n",
    "* Use the external github code to generate new test data, and run your trained NN on these cases. You'll see that a reduced training error not always directly correlates with an improved test performance.\n",
    "\n"
   ]
  }
 ],
 "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": 0
}
