{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qT_RWmTEugu9"
   },
   "source": [
    "# Supervised training for RANS flows around airfoils\n",
    "\n",
    "## Overview \n",
    "\n",
    "For this example of supervised training\n",
    "we have a turbulent airflow around wing profiles, and we'd like to know the average motion\n",
    "and pressure distribution around this airfoil for different Reynolds numbers and angles of attack.\n",
    "Thus, given an airfoil shape, Reynolds numbers, and angle of attack, we'd like to obtain\n",
    "a velocity field and a pressure field around the airfoil.\n",
    "\n",
    "This is classically approximated with _Reynolds-Averaged Navier Stokes_ (RANS) models, and this\n",
    "setting is still one of the most widely used applications of Navier-Stokes solver in industry.\n",
    "However, instead of relying on traditional numerical methods to solve the RANS equations,\n",
    "we now aim for training a surrogate model via a neural network that completely bypasses the numerical solver,\n",
    "and produces the solution in terms of velocity and pressure.\n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/supervised-airfoils.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2ZaC0_MdJMTH"
   },
   "source": [
    "## Formulation\n",
    "\n",
    "\n",
    "With the supervised formulation from {doc}`supervised`, our learning task is pretty straight-forward, and can be written as \n",
    "\n",
    "$$\\begin{aligned}\n",
    "\\text{arg min}_{\\theta} \\sum_i |f(x_i ; \\theta)-y^*_i|^2 ,\n",
    "\\end{aligned}$$\n",
    "\n",
    "where $x$ and $y^*$ each consist of a set of physical fields,\n",
    "and the index $i$ evaluates the difference across all discretization points in our data sets.\n",
    "\n",
    "The goal is to infer velocity $\\mathbf{u} = u_x,u_y$ and a pressure field $p$ in a computational domain $\\Omega$\n",
    "around the airfoil in the center of $\\Omega$. \n",
    "$u_x,u_y$ and $p$ each have a dimension of $128^2$.\n",
    "As inputs we have the Reynolds number $\\text{Re} \\in \\mathbb{R}$, the angle of attack\n",
    "$\\alpha \\in \\mathbb{R}$, and the airfoil shape $\\mathbf{s}$ encoded as a rasterized grid with $128^2$.\n",
    "Both constant, scalar inputs $\\text{Re}$ and $\\alpha$ are likewise extended to a size of $128^2$.\n",
    "Thus, put together, both input and output have the same dimensions: $x,y^* \\in \\mathbb{R}^{3\\times128\\times128}$.\n",
    "This is exactly what we'll specify as input and output dimensions for the NN below.\n",
    "\n",
    "A point to keep in mind here is that our quantities of interest in $y^*$ contain three different physical fields. While the two velocity components are quite similar in spirit, the pressure field typically has a different behavior with an approximately squared scaling with respect to the velocity (cf. [Bernoulli](https://en.wikipedia.org/wiki/Bernoulli%27s_principle)). This implies that we need to be careful with simple summations (as in the minimization problem above), and that we should take care to normalize the data.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ewKoLbFCJMTJ"
   },
   "source": [
    "## Code coming up...\n",
    "\n",
    "Let's get started with the implementation. Note that we'll skip the data generation process here. The code below is adapted from {cite}`thuerey2020dfp` and [this codebase](https://github.com/thunil/Deep-Flow-Prediction), which you can check out for details. Here, we'll simply download a small set of training data generated with a Spalart-Almaras RANS simulation in [OpenFOAM](https://openfoam.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JwZudtWauiGa",
    "outputId": "802dc68e-ac0d-4a0e-ffb0-30153433b98b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Torch version 1.8.1+cu101\n",
      "Downloading training data (300MB), this can take a few minutes the first time...\n",
      "Loaded data, 320 training, 80 validation samples\n",
      "Size of the inputs array: (320, 3, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import os.path, random\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "print(\"Torch version {}\".format(torch.__version__))\n",
    "\n",
    "# get training data\n",
    "if True:\n",
    "    # download\n",
    "    if not os.path.isfile('data-airfoils.npz'):\n",
    "        import urllib.request\n",
    "        url=\"https://ge.in.tum.de/download/2019-deepFlowPred/data.npz\"\n",
    "        print(\"Downloading training data (300MB), this can take a few minutes the first time...\")\n",
    "        urllib.request.urlretrieve(url, 'data-airfoils.npz')\n",
    "    npfile=np.load('data-airfoils.npz')\n",
    "\n",
    "else:\n",
    "    # alternative: load from google drive (upload there beforehand):\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/gdrive')\n",
    "    npfile=np.load('gdrive/My Drive/data-airfoils.npz')\n",
    "\n",
    "\n",
    "print(\"Loaded data, {} training, {} validation samples\".format(len(npfile[\"inputs\"]),len(npfile[\"vinputs\"])))\n",
    "\n",
    "print(\"Size of the inputs array: \"+format(npfile[\"inputs\"].shape))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you run this notebook in colab, the `else` statement above (which is deactivated by default) might be interesting for you: instead of downloading the training data anew every time, you can manually download it once and store it in your google drive. We assume it's stored in the root directory as `data-airfoils.npz`. Afterwards, you can use the code above to load the file from your google drive, which is typically much faster. This is highly recommended if you want to experiment more extensively via colab."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RY1F4kdWPLNG"
   },
   "source": [
    "## RANS training data\n",
    "\n",
    "Now we have some training data. In general it's very important to understand the data we're working with as much as possible (for any ML task the _garbage-in-gargabe-out_ principle definitely holds). We should at least understand the data in terms of dimensions and rough statistics, but ideally also in terms of content. Otherwise we'll have a very hard time interpreting the results of a training run. And despite all the DL magic: if you can't make out any patterns in your data, NNs surely won't find any useful ones.\n",
    "\n",
    "Hence, let's look at one of the training samples... The following is just some helper code to show images side by side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 353
    },
    "id": "6y6YGxMeNECD",
    "outputId": "91a5b3c7-4b87-4e31-d6e8-105d09d08f7e"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50064 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab\n",
    "\n",
    "# helper to show three target channels: normalized, with colormap, side by side\n",
    "def showSbs(a1,a2, stats=False, bottom=\"NN Output\", top=\"Reference\", title=None): \n",
    "  c=[]\n",
    "  for i in range(3):\n",
    "    b = np.flipud( np.concatenate((a2[i],a1[i]),axis=1).transpose())\n",
    "    min, mean, max = np.min(b), np.mean(b), np.max(b); \n",
    "    if stats: print(\"Stats %d: \"%i + format([min,mean,max]))\n",
    "    b -= min; b /= (max-min)\n",
    "    c.append(b)\n",
    "  fig, axes = pylab.subplots(1, 1, figsize=(16, 5))\n",
    "  axes.set_xticks([]); axes.set_yticks([]); \n",
    "  im = axes.imshow(np.concatenate(c,axis=1), origin='upper', cmap='magma')\n",
    "\n",
    "  pylab.colorbar(im); pylab.xlabel('p, ux, uy'); pylab.ylabel('%s           %s'%(bottom,top))\n",
    "  if title is not None: pylab.title(title)\n",
    "\n",
    "NUM=72\n",
    "showSbs(npfile[\"inputs\"][NUM],npfile[\"targets\"][NUM], stats=False, bottom=\"Target Output\", top=\"Inputs\", title=\"3 inputs are shown at the top (mask, in-ux, in-uy), with the 3 output channels (p,ux,uy) at the bottom\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TVHhr8zCPUfN"
   },
   "source": [
    "Next, let's define a small helper class `DfpDataset` to organize inputs and targets. We'll transfer the corresponding data to the pytorch `DataLoader` class. \n",
    "\n",
    "We also set up some globals to control training parameters, maybe most importantly: the learning rate `LR`, i.e. $\\eta$ from the previous setions. When your training run doesn't converge this is the first parameter to experiment with.\n",
    "\n",
    "Here, we'll keep it relatively small throughout. (Using _learning rate decay_ would be better, i.e. potentially give an improed convergence, but is omitted here for clarity.) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BGN4GqxkIueM",
    "outputId": "b97ce299-8c12-402a-d90a-c583224c8971"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training & validation batches: 32 , 8\n"
     ]
    }
   ],
   "source": [
    "# some global training constants\n",
    "\n",
    "# number of training epochs\n",
    "EPOCHS = 100\n",
    "# batch size\n",
    "BATCH_SIZE = 10\n",
    "# learning rate\n",
    "LR = 0.00002\n",
    "\n",
    "class DfpDataset():\n",
    "    def __init__(self, inputs,targets): \n",
    "        self.inputs  = inputs\n",
    "        self.targets = targets\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.inputs)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.inputs[idx], self.targets[idx]\n",
    "\n",
    "tdata = DfpDataset(npfile[\"inputs\"],npfile[\"targets\"])\n",
    "vdata = DfpDataset(npfile[\"vinputs\"],npfile[\"vtargets\"])\n",
    "\n",
    "trainLoader = torch.utils.data.DataLoader(tdata, batch_size=BATCH_SIZE, shuffle=True , drop_last=True) \n",
    "valiLoader  = torch.utils.data.DataLoader(vdata, batch_size=BATCH_SIZE, shuffle=False, drop_last=True) \n",
    "\n",
    "print(\"Training & validation batches: {} , {}\".format(len(trainLoader),len(valiLoader) ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9ys7KZu_P5NB"
   },
   "source": [
    "## Network setup\n",
    "\n",
    "Now we can set up the architecture of our neural network, we'll use a fully convolutional U-net. This is a widely used architecture that uses a stack of convolutions across different spatial resolutions. The main deviation from a regular conv-net is to introduce _skip connection_ from the encoder to the decoder part. This ensures that no information is lost during feature extraction. (Note that this only works if the network is to be used as a whole. It doesn't work in situations where we'd, e.g., want to use the decoder as a standalone component.)\n",
    "\n",
    "Here's a overview of the architecure:\n",
    "\n",
    "![An overview of the U-net we're using for this learning task](resources/supervised-airfoils-unet.jpg)\n",
    "\n",
    "First, we'll define a helper to set up a convolutional block in the network, `blockUNet`. Note, we don't use any pooling! Instead we use strides and transpose convolutions (these need to be symmetric for the decoder part, i.e. have an uneven kernel size), following [best practices](https://distill.pub/2016/deconv-checkerboard/). The full pytroch neural network is managed via the `DfpNet` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "PVNjz4FjGhdY"
   },
   "outputs": [],
   "source": [
    "import os, sys, random\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.autograd \n",
    "import torch.utils.data \n",
    "\n",
    "def blockUNet(in_c, out_c, name, transposed=False, bn=True, relu=True, size=4, pad=1, dropout=0.):\n",
    "    block = nn.Sequential()\n",
    "\n",
    "    if relu:\n",
    "        block.add_module('%s_relu' % name, nn.ReLU(inplace=True))\n",
    "    else:\n",
    "        block.add_module('%s_leakyrelu' % name, nn.LeakyReLU(0.2, inplace=True))\n",
    "\n",
    "    if not transposed:\n",
    "        block.add_module('%s_conv' % name, nn.Conv2d(in_c, out_c, kernel_size=size, stride=2, padding=pad, bias=True))\n",
    "    else:\n",
    "        block.add_module('%s_upsam' % name, nn.Upsample(scale_factor=2, mode='bilinear'))\n",
    "        # reduce kernel size by one for the upsampling (ie decoder part)\n",
    "        block.add_module('%s_tconv' % name, nn.Conv2d(in_c, out_c, kernel_size=(size-1), stride=1, padding=pad, bias=True))\n",
    "\n",
    "    if bn:\n",
    "        block.add_module('%s_bn' % name, nn.BatchNorm2d(out_c))\n",
    "    if dropout>0.:\n",
    "        block.add_module('%s_dropout' % name, nn.Dropout2d( dropout, inplace=True))\n",
    "\n",
    "    return block\n",
    "    \n",
    "class DfpNet(nn.Module):\n",
    "    def __init__(self, channelExponent=6, dropout=0.):\n",
    "        super(DfpNet, self).__init__()\n",
    "        channels = int(2 ** channelExponent + 0.5)\n",
    "\n",
    "        self.layer1 = nn.Sequential()\n",
    "        self.layer1.add_module('layer1', nn.Conv2d(3, channels, 4, 2, 1, bias=True))\n",
    "\n",
    "        self.layer2 = blockUNet(channels  , channels*2, 'enc_layer2', transposed=False, bn=True, relu=False, dropout=dropout )\n",
    "        self.layer3 = blockUNet(channels*2, channels*2, 'enc_layer3', transposed=False, bn=True, relu=False, dropout=dropout )\n",
    "        self.layer4 = blockUNet(channels*2, channels*4, 'enc_layer4', transposed=False, bn=True, relu=False, dropout=dropout )\n",
    "        self.layer5 = blockUNet(channels*4, channels*8, 'enc_layer5', transposed=False, bn=True, relu=False, dropout=dropout ) \n",
    "        self.layer6 = blockUNet(channels*8, channels*8, 'enc_layer6', transposed=False, bn=True, relu=False, dropout=dropout , size=2,pad=0)\n",
    "        self.layer7 = blockUNet(channels*8, channels*8, 'enc_layer7', transposed=False, bn=True, relu=False, dropout=dropout , size=2,pad=0)\n",
    "     \n",
    "        # note, kernel size is internally reduced by one for the decoder part\n",
    "        self.dlayer7 = blockUNet(channels*8, channels*8, 'dec_layer7', transposed=True, bn=True, relu=True, dropout=dropout , size=2,pad=0)\n",
    "        self.dlayer6 = blockUNet(channels*16,channels*8, 'dec_layer6', transposed=True, bn=True, relu=True, dropout=dropout , size=2,pad=0)\n",
    "        self.dlayer5 = blockUNet(channels*16,channels*4, 'dec_layer5', transposed=True, bn=True, relu=True, dropout=dropout ) \n",
    "        self.dlayer4 = blockUNet(channels*8, channels*2, 'dec_layer4', transposed=True, bn=True, relu=True, dropout=dropout )\n",
    "        self.dlayer3 = blockUNet(channels*4, channels*2, 'dec_layer3', transposed=True, bn=True, relu=True, dropout=dropout )\n",
    "        self.dlayer2 = blockUNet(channels*4, channels  , 'dec_layer2', transposed=True, bn=True, relu=True, dropout=dropout )\n",
    "\n",
    "        self.dlayer1 = nn.Sequential()\n",
    "        self.dlayer1.add_module('dec_layer1_relu', nn.ReLU(inplace=True))\n",
    "        self.dlayer1.add_module('dec_layer1_tconv', nn.ConvTranspose2d(channels*2, 3, 4, 2, 1, bias=True))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # note, this Unet stack could be allocated with a loop, of course... \n",
    "        out1 = self.layer1(x)\n",
    "        out2 = self.layer2(out1)\n",
    "        out3 = self.layer3(out2)\n",
    "        out4 = self.layer4(out3)\n",
    "        out5 = self.layer5(out4)\n",
    "        out6 = self.layer6(out5)\n",
    "        out7 = self.layer7(out6)\n",
    "        # ... bottleneck ...\n",
    "        dout6 = self.dlayer7(out7)\n",
    "        dout6_out6 = torch.cat([dout6, out6], 1)\n",
    "        dout6 = self.dlayer6(dout6_out6)\n",
    "        dout6_out5 = torch.cat([dout6, out5], 1)\n",
    "        dout5 = self.dlayer5(dout6_out5)\n",
    "        dout5_out4 = torch.cat([dout5, out4], 1)\n",
    "        dout4 = self.dlayer4(dout5_out4)\n",
    "        dout4_out3 = torch.cat([dout4, out3], 1)\n",
    "        dout3 = self.dlayer3(dout4_out3)\n",
    "        dout3_out2 = torch.cat([dout3, out2], 1)\n",
    "        dout2 = self.dlayer2(dout3_out2)\n",
    "        dout2_out1 = torch.cat([dout2, out1], 1)\n",
    "        dout1 = self.dlayer1(dout2_out1)\n",
    "        return dout1\n",
    "\n",
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        m.weight.data.normal_(0.0, 0.02)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        m.weight.data.normal_(1.0, 0.02)\n",
    "        m.bias.data.fill_(0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QAl3VgKVQSI3"
   },
   "source": [
    "Next, we can initialize an instance of the `DfpNet`.\n",
    "\n",
    "Below, the `EXPO` parameter here controls the exponent for the feature maps of our Unet: this directly scales the network size (3 gives a model with ca. 150k parameters). This is relatively small for a generative model with $3 \\times 128^2 = \\text{ca. }49k$ outputs, but yields fast training times and prevents overfitting given the relatively small data set we're using here. Hence it's a good starting point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "N1uV8k5VIoqT",
    "outputId": "ceb8fb79-6154-423a-9d17-d6e4dd0c2718"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainable params: 147555   -> crucial! always keep in view... \n"
     ]
    }
   ],
   "source": [
    "# channel exponent to control network size\n",
    "EXPO = 3\n",
    "\n",
    "# setup network\n",
    "net = DfpNet(channelExponent=EXPO)\n",
    "#print(net) # to double check the details...\n",
    "\n",
    "model_parameters = filter(lambda p: p.requires_grad, net.parameters())\n",
    "params = sum([np.prod(p.size()) for p in model_parameters])\n",
    "\n",
    "# crucial parameter to keep in view: how many parameters do we have?\n",
    "print(\"Trainable params: {}   -> crucial! always keep in view... \".format(params)) \n",
    "\n",
    "net.apply(weights_init)\n",
    "\n",
    "criterionL1 = nn.L1Loss()\n",
    "optimizerG = optim.Adam(net.parameters(), lr=LR, betas=(0.5, 0.999), weight_decay=0.0)\n",
    "\n",
    "targets = torch.autograd.Variable(torch.FloatTensor(BATCH_SIZE, 3, 128, 128))\n",
    "inputs  = torch.autograd.Variable(torch.FloatTensor(BATCH_SIZE, 3, 128, 128))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UNjBAvfWJMTR"
   },
   "source": [
    "With an exponent of 3, this network has 147555 trainable parameters. As the subtle hint in the print statement indicates, this is a crucial number to always have in view when training NNs. It's easy to change settings, and get a network that has millions of parameters, and as a result probably all kinds of convergence and overfitting problems. The number of parrameters definitely has to be matched with the amount of training data, and should also scale with the depth of the network. How these three relate to each other exactly is problem dependent, though."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ErLQ6UYlQT8-"
   },
   "source": [
    "## Training\n",
    "\n",
    "Finally, we can train the model. This step can take a while, as the training runs over all 320 samples 100 times, and continually evaluate the validation samples to keep track of how well the current state of the NN is doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kkOYvwxqKa8n",
    "outputId": "a973331c-9e8f-401f-dfd7-9b7177faca0f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training from scratch\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:3458: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.\n",
      "  \"See the documentation of nn.Upsample for details.\".format(mode)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, L1 train: 0.20192, L1 vali: 0.17916\n",
      "Epoch: 1, L1 train: 0.18912, L1 vali: 0.18013\n",
      "Epoch: 2, L1 train: 0.17917, L1 vali: 0.17145\n",
      "Epoch: 20, L1 train: 0.08637, L1 vali: 0.08432\n",
      "Epoch: 40, L1 train: 0.03682, L1 vali: 0.03222\n",
      "Epoch: 60, L1 train: 0.03176, L1 vali: 0.02710\n",
      "Epoch: 80, L1 train: 0.02772, L1 vali: 0.02522\n",
      "Training done, saved model\n"
     ]
    }
   ],
   "source": [
    "history_L1 = []\n",
    "history_L1val = []\n",
    "\n",
    "if os.path.isfile(\"model\"):\n",
    "  print(\"Found existing model, loading & skipping training\")\n",
    "  net.load_state_dict(torch.load(\"model\")) # optionally, load existing model\n",
    "\n",
    "else:\n",
    "  print(\"Training from scratch\")\n",
    "  for epoch in range(EPOCHS):\n",
    "      net.train()\n",
    "      L1_accum = 0.0\n",
    "      for i, traindata in enumerate(trainLoader, 0):\n",
    "          inputs_curr, targets_curr = traindata\n",
    "          inputs.data.copy_(inputs_curr.float())\n",
    "          targets.data.copy_(targets_curr.float())\n",
    "\n",
    "          net.zero_grad()\n",
    "          gen_out = net(inputs)\n",
    "\n",
    "          lossL1 = criterionL1(gen_out, targets)\n",
    "          lossL1.backward()\n",
    "          optimizerG.step()\n",
    "          L1_accum += lossL1.item()\n",
    "\n",
    "      # validation\n",
    "      net.eval()\n",
    "      L1val_accum = 0.0\n",
    "      for i, validata in enumerate(valiLoader, 0):\n",
    "          inputs_curr, targets_curr = validata\n",
    "          inputs.data.copy_(inputs_curr.float())\n",
    "          targets.data.copy_(targets_curr.float())\n",
    "\n",
    "          outputs = net(inputs)\n",
    "          outputs_curr = outputs.data.cpu().numpy()\n",
    "\n",
    "          lossL1val = criterionL1(outputs, targets)\n",
    "          L1val_accum += lossL1val.item()\n",
    "\n",
    "      # data for graph plotting\n",
    "      history_L1.append( L1_accum / len(trainLoader) )\n",
    "      history_L1val.append( L1val_accum / len(valiLoader) )\n",
    "\n",
    "      if epoch<3 or epoch%20==0:\n",
    "          print( \"Epoch: {}, L1 train: {:7.5f}, L1 vali: {:7.5f}\".format(epoch, history_L1[-1], history_L1val[-1]) )\n",
    "\n",
    "  torch.save(net.state_dict(), \"model\" )\n",
    "  print(\"Training done, saved model\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4KuUpJsSL3Jv"
   },
   "source": [
    "The model is finally trained! The losses should have nicely gone down in terms of absolute values: With the standard settings from an initial value of around 0.2 for the validation  loss, to ca. 0.02 after 100 epochs. \n",
    "\n",
    "Let's look at the graphs to get some intuition for how the training progressed over time. This is typically important to identify longer-term trends in the training. In practice it's tricky to spot whether the overall trend of 100 or so noisy numbers in a command line log is going slightly up or down - this is much easier to spot in a visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 286
    },
    "id": "Hlki3bU8JMTT",
    "outputId": "2a36d013-5af2-4580-f782-86b65f6f8c11"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20564 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "l1train = np.asarray(history_L1)\n",
    "l1vali  = np.asarray(history_L1val)\n",
    "\n",
    "plt.plot(np.arange(l1train.shape[0]),l1train,'b',label='Training loss')\n",
    "plt.plot(np.arange(l1vali.shape[0] ),l1vali ,'g',label='Validation loss')\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3Vj1E9ZdJMTU"
   },
   "source": [
    "You should see a curve that goes down for ca. 40 epochs, and then starts to flatten out. In the last part, it's still slowly decreasing, and most importantly, the validation loss is not increasing. This would be a certain sign of overfitting, and something that we should avoid. (Try decreasing the amount of training data artificially, then you should be able to intentionally cause overfitting.)\n",
    "\n",
    "## Training progress and validation\n",
    "\n",
    "If you look closely at this graph, you should spot something peculiar:\n",
    "_Why is the validation loss lower than the training loss_?\n",
    "The data is similar to the training data of course, but in a way it's slightly \"tougher\", because the model certainly never received any validation samples during training. It is natural that the validation loss slightly deviates from the training loss, but how can the L1 loss be _lower_ for these inputs?\n",
    "\n",
    "This is a subtlety of the training loop above: it runs a training step first, and the loss for each point in the graph is measured with the evolving state of the network in an epoch. The network is updated, and afterwards runs through the validation samples. Thus all validation samples are using a state that is definitely different (and hopefully a bit better) than the initial states of the epoch. Hence, the validation loss can be slightly lower.\n",
    "\n",
    "A general word of caution here: never evaluate your network with training data! That won't tell you much because overfitting is a very common problem. At least use data the network hasn't seen before, i.e. validation data, and if that looks good, try some more different (at least slightly out-of-distribution) inputs, i.e., _test data_. The next cell runs the trained network over the validation data, and displays one of them with the `showSbs` function.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 386
    },
    "id": "N6OONK0bL4ev",
    "outputId": "240ad666-e266-42d1-c13f-3573eddc7221",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:3458: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.\n",
      "  \"See the documentation of nn.Upsample for details.\".format(mode)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 113380 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "net.eval()\n",
    "for i, validata in enumerate(valiLoader, 0):\n",
    "    inputs_curr, targets_curr = validata\n",
    "    inputs.data.copy_(inputs_curr.float())\n",
    "    targets.data.copy_(targets_curr.float())\n",
    "    \n",
    "    outputs = net(inputs)\n",
    "    outputs_curr = outputs.data.cpu().numpy()\n",
    "    if i<1: showSbs(targets_curr[0] , outputs_curr[0], title=\"Validation sample %d\"%(i*BATCH_SIZE))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TJtxXgGoJMTV"
   },
   "source": [
    "Visually, there should at least be a rough resemblance here between input out network output. We'll save the more detailed evaluation for the test data, though."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4cIO36A_FhdE"
   },
   "source": [
    "## Test evaluation\n",
    "\n",
    "Now let's look at actual test samples: In this case we'll use new airfoil shapes as out-of-distribution (OOD) data. These are shapes that the network never saw in any training samples, and hence it tells us a bit about how well the model generalizes to unseen inputs (the validation data wouldn't suffice to draw conclusions about generalization).\n",
    "\n",
    "We'll use the same visualization as before, and as indicated by the Bernoulli equation, especially the _pressure_ in the first column is a challenging quantity for the network. Due to it's cubic scaling w.r.t. the input freestream velocity and localized peaks, it is the toughest quantity to infer for the network.\n",
    "\n",
    "The cell below first downloads a smaller archive with these test data samples, and then runs them through the network. The evaluation loop also computes the accumulated L1 error such that we can quantify how well the network does on the test samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "dD5O58L9HLE-",
    "outputId": "9ea61f54-2bc2-4881-f6bf-e37502e48383"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading test data, this should be fast...\n",
      "Loaded 10/10 test samples\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:3458: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.\n",
      "  \"See the documentation of nn.Upsample for details.\".format(mode)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Average test error: 0.029000501427799464\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 117196 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 123640 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 129928 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "if not os.path.isfile('data-airfoils-test.npz'):\n",
    "  import urllib.request\n",
    "  url=\"https://ge.in.tum.de/download/2019-deepFlowPred/data_test.npz\"\n",
    "  print(\"Downloading test data, this should be fast...\")\n",
    "  urllib.request.urlretrieve(url, 'data-airfoils-test.npz')\n",
    "\n",
    "nptfile=np.load('data-airfoils-test.npz')\n",
    "print(\"Loaded {}/{} test samples\\n\".format(len(nptfile[\"test_inputs\"]),len(nptfile[\"test_targets\"])))\n",
    "\n",
    "testdata = DfpDataset(nptfile[\"test_inputs\"],nptfile[\"test_targets\"])\n",
    "testLoader  = torch.utils.data.DataLoader(testdata, batch_size=1, shuffle=False, drop_last=True) \n",
    "\n",
    "net.eval()\n",
    "L1t_accum = 0.\n",
    "for i, validata in enumerate(testLoader, 0):\n",
    "    inputs_curr, targets_curr = validata\n",
    "    inputs.data.copy_(inputs_curr.float())\n",
    "    targets.data.copy_(targets_curr.float())\n",
    "\n",
    "    outputs = net(inputs)\n",
    "    outputs_curr = outputs.data.cpu().numpy()\n",
    "\n",
    "    lossL1t = criterionL1(outputs, targets)\n",
    "    L1t_accum += lossL1t.item()\n",
    "    if i<3: showSbs(targets_curr[0] , outputs_curr[0],  title=\"Test sample %d\"%(i))\n",
    "\n",
    "print(\"\\nAverage test error: {}\".format( L1t_accum/len(testLoader) ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "By722sgwFbwG"
   },
   "source": [
    "The average test error with the default settings should be ca. 0.03. As the inputs are normalized, this means the average error across all three fields is 3% w.r.t. the maxima of each quantity. This is not too bad for new shapes, but clearly leaves room for improvement.\n",
    "\n",
    "Looking at the visualizations, you'll notice that especially high-pressure peaks and pockets of larger y-velocities are missing in the outputs. This is primarily caused by the small network, which does not have enough resources to reconstruct details.\n",
    "\n",
    "Nonetheless, we have successfully replaced a fairly sophisticated RANS solver with a very small and fast neural network model. It has GPU support \"out-of-the-box\" (via pytorch), is differentiable, and introduces an error of only a few per-cent.\n",
    "\n",
    "---\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vhH-rUZ-JMTX"
   },
   "source": [
    "## Next steps\n",
    "\n",
    "There are many obvious things to try here (see the suggestions below), e.g. longer training, larger data sets, larger networks etc. \n",
    "\n",
    "* Experiment with learning rate, dropout, and model size to reduce the error on the test set. How small can you make it with the given training data?\n",
    "\n",
    "* The setup above uses normalized data [the original fields by undoing the normalization](https://github.com/thunil/Deep-Flow-Prediction), and check how well the network does w.r.t. the original \n",
    "\n",
    "* As you'll see, it's a bit limited here what you can get out of this dataset, head over to [the main github repo of this project](https://github.com/thunil/Deep-Flow-Prediction) to download larger data sets, or generate own data\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "supervised-airfoils-may16a.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
}
