{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qT_RWmTEugu9"
   },
   "source": [
    "# Supervised training for RANS flows around airfoils\n",
    "\n",
    "## Overview \n",
    "\n",
    "Let's directly look at an example with a fairly complicated context:\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 $\\mathbf{u}$ and a pressure field $p$ in a computational domain $\\Omega$\n",
    "around the airfoil in the center of $\\Omega$.\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 know aim for training a neural network that completely bypasses the numerical solver,\n",
    "and produces the solution in terms of $\\mathbf{u}$ and $p$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
    "$\\text{arg min}_{\\theta} \\sum_i |f(x_i ; \\theta)-y^*_i|^2$ , \n",
    "\n",
    "where $x$ and $y^*$ each consist of a set of physical fields (pressure, x velocity, and y velocity), each with a dimension of $128^2$, i.e.  $x,y^* \\in \\mathbb{R}^{3\\times128^2}$, and the index $i$ evaluates the difference across all discretization points in our data sets.\n",
    "\n",
    "Thus, the only point to keep in mind is that our quantities of interest 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. This implies that we need to be careful with simple summation (as above), and that we should take care to normalize the data.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code coming up...\n",
    "\n",
    "Let's get started with the implementation. Note that we'll skip the data generation process here. This example is adapted from [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": "59b8aa53-0858-49d1-c8d5-edb92f0cba9c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Torch version 1.7.1\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",
    "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",
    "\n",
    "npfile=np.load(\"data-airfoils.npz\")\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": {
    "id": "RY1F4kdWPLNG"
   },
   "source": [
    "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": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 307
    },
    "id": "6y6YGxMeNECD",
    "outputId": "7f8d669e-3023-4c42-8673-21b780f5620d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Here are all 3 inputs are shown at the top (mask,in x, in y) \n",
      "Side by side with the 3 output channels (p,vx,vy) at the bottom:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32972 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x12C4E0D10>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from PIL import Image\n",
    "from matplotlib import cm\n",
    "\n",
    "def show(a, stats=False): # helper to show three target channels, normalized with colormap\n",
    "  for i in range(3):\n",
    "    b = np.flipud(a[i].transpose())\n",
    "    min,mean,max = np.min(b),np.mean(b),np.max(b); \n",
    "    if stats: print([b[0,0],min,mean,max])\n",
    "    b -= min; b /= (max-min)\n",
    "    if stats: print(\"{}\".format([\"p:\",\"ux:\",\"uy:\"][i])) # note, for inputs the first channel is the mask...\n",
    "    display(Image.fromarray( cm.magma(b, bytes=True) ))\n",
    "\n",
    "def showSbs(a1,a2, stats=False): # helper to show three target channels, normalized with colormap side by side\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([b[0,0],min,mean,max])\n",
    "    b -= min; b /= (max-min)\n",
    "    c.append(b)\n",
    "  c = np.concatenate(c,axis=1)\n",
    "  display(Image.fromarray( cm.magma(c, bytes=True) ))\n",
    "\n",
    "NUM=72\n",
    "print(\"\\nHere are all 3 inputs are shown at the top (mask,in x, in y) \\nSide by side with the 3 output channels (p,vx,vy) at the bottom:\")\n",
    "showSbs(npfile[\"inputs\"][NUM],npfile[\"targets\"][NUM])\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",
    "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. (_Learning rate decay_ would be better, i.e. give slightly faster training, but is omitted here for clarity.) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BGN4GqxkIueM",
    "outputId": "bde9ee69-5d3b-43b3-ccaa-9d18a85ab9bc"
   },
   "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. It uses the typical stack of convolutions across different spatial resolutions. The main trick is to introduce _skip connection_ from the encoder to the decoder part, to ensure no information is lots 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 an own unit.)\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/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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, 'layer2', transposed=False, bn=True,  relu=False, dropout=dropout )\n",
    "        self.layer3 = blockUNet(channels*2, channels*2, 'layer3', transposed=False, bn=True,  relu=False, dropout=dropout )\n",
    "        self.layer4 = blockUNet(channels*2, channels*4, 'layer4', transposed=False, bn=True,  relu=False, dropout=dropout )\n",
    "        self.layer5 = blockUNet(channels*4, channels*8, 'layer5', transposed=False, bn=True,  relu=False, dropout=dropout ) \n",
    "        self.layer6 = blockUNet(channels*8, channels*8, 'layer6', transposed=False, bn=True,  relu=False, dropout=dropout , size=2,pad=0)\n",
    "        self.layer7 = blockUNet(channels*8, channels*8, 'layer7', transposed=False, bn=False, 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, 'dlayer7', transposed=True, bn=True, relu=True, dropout=dropout , size=2,pad=0)\n",
    "        self.dlayer6 = blockUNet(channels*16,channels*8, 'dlayer6', transposed=True, bn=True, relu=True, dropout=dropout , size=2,pad=0)\n",
    "        self.dlayer5 = blockUNet(channels*16,channels*4, 'dlayer5', transposed=True, bn=True, relu=True, dropout=dropout ) \n",
    "        self.dlayer4 = blockUNet(channels*8, channels*2, 'dlayer4', transposed=True, bn=True, relu=True, dropout=dropout )\n",
    "        self.dlayer3 = blockUNet(channels*4, channels*2, 'dlayer3', transposed=True, bn=True, relu=True, dropout=dropout )\n",
    "        self.dlayer2 = blockUNet(channels*4, channels  , 'dlayer2', transposed=True, bn=True, relu=True, dropout=dropout )\n",
    "\n",
    "        self.dlayer1 = nn.Sequential()\n",
    "        self.dlayer1.add_module('dlayer1_relu', nn.ReLU(inplace=True))\n",
    "        self.dlayer1.add_module('dlayer1_tconv', nn.ConvTranspose2d(channels*2, 3, 4, 2, 1, bias=True))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # note, this Unet stack could be automated of course - but it's a bit clearer this way...\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",
    "Here, 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 for $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": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "N1uV8k5VIoqT",
    "outputId": "f085f33f-30d1-4fe4-bb41-e1d1496ba655"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DfpNet(\n",
      "  (layer1): Sequential(\n",
      "    (layer1): Conv2d(3, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (layer2_leakyrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    (layer2_conv): Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (layer2_bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (layer3_leakyrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    (layer3_conv): Conv2d(16, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (layer3_bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (layer4): Sequential(\n",
      "    (layer4_leakyrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    (layer4_conv): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (layer4_bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (layer5): Sequential(\n",
      "    (layer5_leakyrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    (layer5_conv): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "    (layer5_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (layer6): Sequential(\n",
      "    (layer6_leakyrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    (layer6_conv): Conv2d(64, 64, kernel_size=(2, 2), stride=(2, 2))\n",
      "    (layer6_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (layer7): Sequential(\n",
      "    (layer7_leakyrelu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
      "    (layer7_conv): Conv2d(64, 64, kernel_size=(2, 2), stride=(2, 2))\n",
      "  )\n",
      "  (dlayer7): Sequential(\n",
      "    (dlayer7_relu): ReLU(inplace=True)\n",
      "    (dlayer7_upsam): Upsample(scale_factor=2.0, mode=nearest)\n",
      "    (dlayer7_tconv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (dlayer7_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (dlayer6): Sequential(\n",
      "    (dlayer6_relu): ReLU(inplace=True)\n",
      "    (dlayer6_upsam): Upsample(scale_factor=2.0, mode=nearest)\n",
      "    (dlayer6_tconv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (dlayer6_bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (dlayer5): Sequential(\n",
      "    (dlayer5_relu): ReLU(inplace=True)\n",
      "    (dlayer5_upsam): Upsample(scale_factor=2.0, mode=nearest)\n",
      "    (dlayer5_tconv): Conv2d(128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (dlayer5_bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (dlayer4): Sequential(\n",
      "    (dlayer4_relu): ReLU(inplace=True)\n",
      "    (dlayer4_upsam): Upsample(scale_factor=2.0, mode=nearest)\n",
      "    (dlayer4_tconv): Conv2d(64, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (dlayer4_bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (dlayer3): Sequential(\n",
      "    (dlayer3_relu): ReLU(inplace=True)\n",
      "    (dlayer3_upsam): Upsample(scale_factor=2.0, mode=nearest)\n",
      "    (dlayer3_tconv): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (dlayer3_bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (dlayer2): Sequential(\n",
      "    (dlayer2_relu): ReLU(inplace=True)\n",
      "    (dlayer2_upsam): Upsample(scale_factor=2.0, mode=nearest)\n",
      "    (dlayer2_tconv): Conv2d(32, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (dlayer2_bn): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (dlayer1): Sequential(\n",
      "    (dlayer1_relu): ReLU(inplace=True)\n",
      "    (dlayer1_tconv): ConvTranspose2d(16, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
      "  )\n",
      ")\n",
      "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": "ErLQ6UYlQT8-"
   },
   "source": [
    "## Training\n",
    "\n",
    "Finally, we can train the model. This step can take a while, as we'll go over all 320 samples 100 times, and continually evaluate the validation samples to keep track of how well we're doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kkOYvwxqKa8n",
    "outputId": "e71149b6-8888-44d0-bce9-d5288b1d9fb8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training from scratch\n",
      "Epoch: 0, L1 train: 0.23608, L1 vali: 0.22773\n",
      "Epoch: 1, L1 train: 0.22596, L1 vali: 0.22119\n",
      "Epoch: 2, L1 train: 0.21803, L1 vali: 0.21487\n",
      "Epoch: 3, L1 train: 0.21162, L1 vali: 0.20664\n",
      "Epoch: 4, L1 train: 0.20456, L1 vali: 0.20017\n",
      "Epoch: 5, L1 train: 0.19872, L1 vali: 0.19361\n",
      "Epoch: 6, L1 train: 0.19307, L1 vali: 0.18902\n",
      "Epoch: 7, L1 train: 0.18764, L1 vali: 0.18313\n",
      "Epoch: 8, L1 train: 0.18274, L1 vali: 0.17782\n",
      "Epoch: 9, L1 train: 0.17719, L1 vali: 0.17278\n",
      "Epoch: 10, L1 train: 0.17280, L1 vali: 0.16770\n",
      "Epoch: 11, L1 train: 0.16881, L1 vali: 0.16246\n",
      "Epoch: 12, L1 train: 0.16275, L1 vali: 0.15712\n",
      "Epoch: 13, L1 train: 0.15851, L1 vali: 0.15384\n",
      "Epoch: 14, L1 train: 0.15592, L1 vali: 0.14751\n",
      "Epoch: 15, L1 train: 0.14877, L1 vali: 0.14329\n",
      "Epoch: 16, L1 train: 0.14498, L1 vali: 0.13888\n",
      "Epoch: 17, L1 train: 0.13987, L1 vali: 0.13402\n",
      "Epoch: 18, L1 train: 0.13579, L1 vali: 0.13007\n",
      "Epoch: 19, L1 train: 0.13181, L1 vali: 0.12809\n",
      "Epoch: 20, L1 train: 0.12681, L1 vali: 0.12327\n",
      "Epoch: 21, L1 train: 0.12324, L1 vali: 0.11969\n",
      "Epoch: 22, L1 train: 0.11602, L1 vali: 0.11148\n",
      "Epoch: 23, L1 train: 0.11144, L1 vali: 0.10759\n",
      "Epoch: 24, L1 train: 0.10772, L1 vali: 0.10140\n",
      "Epoch: 25, L1 train: 0.10440, L1 vali: 0.10212\n",
      "Epoch: 26, L1 train: 0.09807, L1 vali: 0.09606\n",
      "Epoch: 27, L1 train: 0.09329, L1 vali: 0.09090\n",
      "Epoch: 28, L1 train: 0.08888, L1 vali: 0.08182\n",
      "Epoch: 29, L1 train: 0.08625, L1 vali: 0.08044\n",
      "Epoch: 30, L1 train: 0.08248, L1 vali: 0.08004\n",
      "Epoch: 31, L1 train: 0.07684, L1 vali: 0.07337\n",
      "Epoch: 32, L1 train: 0.07180, L1 vali: 0.06905\n",
      "Epoch: 33, L1 train: 0.07012, L1 vali: 0.06635\n",
      "Epoch: 34, L1 train: 0.06508, L1 vali: 0.06049\n",
      "Epoch: 35, L1 train: 0.06191, L1 vali: 0.06010\n",
      "Epoch: 36, L1 train: 0.05935, L1 vali: 0.05591\n",
      "Epoch: 37, L1 train: 0.05504, L1 vali: 0.05180\n",
      "Epoch: 38, L1 train: 0.05436, L1 vali: 0.04854\n",
      "Epoch: 39, L1 train: 0.05151, L1 vali: 0.04681\n",
      "Epoch: 40, L1 train: 0.04978, L1 vali: 0.04431\n",
      "Epoch: 41, L1 train: 0.04864, L1 vali: 0.04357\n",
      "Epoch: 42, L1 train: 0.04755, L1 vali: 0.04144\n",
      "Epoch: 43, L1 train: 0.04658, L1 vali: 0.04245\n",
      "Epoch: 44, L1 train: 0.04391, L1 vali: 0.04071\n",
      "Epoch: 45, L1 train: 0.04504, L1 vali: 0.03778\n",
      "Epoch: 46, L1 train: 0.04246, L1 vali: 0.03717\n",
      "Epoch: 47, L1 train: 0.04292, L1 vali: 0.03732\n",
      "Epoch: 48, L1 train: 0.04186, L1 vali: 0.03603\n",
      "Epoch: 49, L1 train: 0.04050, L1 vali: 0.03514\n",
      "Epoch: 50, L1 train: 0.04045, L1 vali: 0.03408\n",
      "Epoch: 51, L1 train: 0.04003, L1 vali: 0.03458\n",
      "Epoch: 52, L1 train: 0.03965, L1 vali: 0.03332\n",
      "Epoch: 53, L1 train: 0.03847, L1 vali: 0.03322\n",
      "Epoch: 54, L1 train: 0.03771, L1 vali: 0.03241\n",
      "Epoch: 55, L1 train: 0.03759, L1 vali: 0.03199\n",
      "Epoch: 56, L1 train: 0.03729, L1 vali: 0.03124\n",
      "Epoch: 57, L1 train: 0.03738, L1 vali: 0.03124\n",
      "Epoch: 58, L1 train: 0.03612, L1 vali: 0.03081\n",
      "Epoch: 59, L1 train: 0.03600, L1 vali: 0.03119\n",
      "Epoch: 60, L1 train: 0.03518, L1 vali: 0.03008\n",
      "Epoch: 61, L1 train: 0.03485, L1 vali: 0.02950\n",
      "Epoch: 62, L1 train: 0.03407, L1 vali: 0.02895\n",
      "Epoch: 63, L1 train: 0.03494, L1 vali: 0.02905\n",
      "Epoch: 64, L1 train: 0.03395, L1 vali: 0.02903\n",
      "Epoch: 65, L1 train: 0.03340, L1 vali: 0.02823\n",
      "Epoch: 66, L1 train: 0.03394, L1 vali: 0.02846\n",
      "Epoch: 67, L1 train: 0.03471, L1 vali: 0.02803\n",
      "Epoch: 68, L1 train: 0.03379, L1 vali: 0.02757\n",
      "Epoch: 69, L1 train: 0.03401, L1 vali: 0.02807\n",
      "Epoch: 70, L1 train: 0.03418, L1 vali: 0.02765\n",
      "Epoch: 71, L1 train: 0.03405, L1 vali: 0.02842\n",
      "Epoch: 72, L1 train: 0.03239, L1 vali: 0.02720\n",
      "Epoch: 73, L1 train: 0.03253, L1 vali: 0.02760\n",
      "Epoch: 74, L1 train: 0.03298, L1 vali: 0.02694\n",
      "Epoch: 75, L1 train: 0.03166, L1 vali: 0.02704\n",
      "Epoch: 76, L1 train: 0.03235, L1 vali: 0.02675\n",
      "Epoch: 77, L1 train: 0.03201, L1 vali: 0.02712\n",
      "Epoch: 78, L1 train: 0.03069, L1 vali: 0.02600\n",
      "Epoch: 79, L1 train: 0.03228, L1 vali: 0.02637\n",
      "Epoch: 80, L1 train: 0.03006, L1 vali: 0.02546\n",
      "Epoch: 81, L1 train: 0.03180, L1 vali: 0.02574\n",
      "Epoch: 82, L1 train: 0.03017, L1 vali: 0.02540\n",
      "Epoch: 83, L1 train: 0.03192, L1 vali: 0.02602\n",
      "Epoch: 84, L1 train: 0.03004, L1 vali: 0.02592\n",
      "Epoch: 85, L1 train: 0.02985, L1 vali: 0.02506\n",
      "Epoch: 86, L1 train: 0.03043, L1 vali: 0.02560\n",
      "Epoch: 87, L1 train: 0.03020, L1 vali: 0.02498\n",
      "Epoch: 88, L1 train: 0.02907, L1 vali: 0.02516\n",
      "Epoch: 89, L1 train: 0.02957, L1 vali: 0.02496\n",
      "Epoch: 90, L1 train: 0.02943, L1 vali: 0.02447\n",
      "Epoch: 91, L1 train: 0.02832, L1 vali: 0.02472\n",
      "Epoch: 92, L1 train: 0.02832, L1 vali: 0.02467\n",
      "Epoch: 93, L1 train: 0.02914, L1 vali: 0.02492\n",
      "Epoch: 94, L1 train: 0.02912, L1 vali: 0.02387\n",
      "Epoch: 95, L1 train: 0.02861, L1 vali: 0.02467\n",
      "Epoch: 96, L1 train: 0.02931, L1 vali: 0.02387\n",
      "Epoch: 97, L1 train: 0.02951, L1 vali: 0.02465\n",
      "Epoch: 98, L1 train: 0.02995, L1 vali: 0.02439\n",
      "Epoch: 99, L1 train: 0.02881, L1 vali: 0.02360\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": [
    "Yay, the model is trained...! The losses nicely went 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 trained 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": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(100,)]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 18152 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "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": {},
   "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",
    "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. So 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 model with training data! That won't tell you much because overfitting is a very common problem. At least use validation data, and if that looks good, try some more different inputs, i.e., _test data_. (That's coming up next below.) \n",
    "\n",
    "First, let's visualize some validation samples. All visualizations in the following show 2 rows: the target at the top, the model output at the bottom. Each row shows pressure, x velocity, and y velocity (from left to right). Typically, especially the _pressure_ in the first column is interesting, due to it's cubic scaling w.r.t. the varying velocities, it is the toughest quantity to infer for the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "N6OONK0bL4ev",
    "outputId": "6ab05f74-2cb2-406a-e06b-e29beba74e97"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation sample 0, target (top) & model output (bottom)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32864 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x12A5B9D90>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation sample 10, target (top) & model output (bottom)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28812 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x11738EA10>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation sample 20, target (top) & model output (bottom)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31344 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x11738EA10>"
      ]
     },
     "metadata": {},
     "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<3: # only show first 3 of each batch here, enlarge to show all\n",
    "        print(\"Validation sample {}, target (top) & model output (bottom)\".format(i * batch_size) ) # only show each first one\n",
    "        showSbs(targets_curr[0] , outputs_curr[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4cIO36A_FhdE"
   },
   "source": [
    "## Test evaluation\n",
    "\n",
    "Now let's look at **proper** test samples, i.e. OOD samples, for which in this case we'll use new airfoil shapes. 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 new shapes.\n",
    "\n",
    "As these samples are at least slightly OOD, we can draw conclusions about how well the model generalizes, which the validation data would not really tell us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "dD5O58L9HLE-",
    "outputId": "64bd417a-7fad-42c0-822c-8075eba1ba06"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 10/10 test samples\n",
      "Test sample 0, loss 0.16995319724082947, target (top) & model output (bottom)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37712 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x12C246D10>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test sample 1, loss 0.1963399201631546, target (top) & model output (bottom)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34976 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x12C246E50>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test sample 2, loss 0.18705692887306213, target (top) & model output (bottom)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32852 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=384x256 at 0x12C246BD0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test sample 3, loss 0.15921331942081451, target (top) & model output (bottom)\n",
      "Test sample 4, loss 0.17740412056446075, target (top) & model output (bottom)\n",
      "Test sample 5, loss 0.20241335034370422, target (top) & model output (bottom)\n",
      "Test sample 6, loss 0.16636453568935394, target (top) & model output (bottom)\n",
      "Test sample 7, loss 0.21479596197605133, target (top) & model output (bottom)\n",
      "Test sample 8, loss 0.16884702444076538, target (top) & model output (bottom)\n",
      "Test sample 9, loss 0.17588728666305542, target (top) & model output (bottom)\n",
      "\n",
      "Average test error: 0.1818275645375252\n"
     ]
    }
   ],
   "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\".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",
    "\n",
    "    print(\"Test sample {}, loss {}, target (top) & model output (bottom)\".format(i, lossL1t.item()) ) \n",
    "    if i<3: # again only show first 3 of each batch here, enlarge value to show all\n",
    "        showSbs(targets_curr[0] , outputs_curr[0])\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. This is not too bad for new shapes, but clearly leaves room for improvement! There are many obvious things to try here (see the suggestions below), e.g. longer training, larger data sets, larger networks etc. \n",
    "\n",
    "---\n",
    "\n",
    "## Nex steps\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",
    "* 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",
    "**(TODO, us: for exercise, provide data with \"errors\" (nan & huge neg number in 1 cell), filter out to make model train...)**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "AIAA-airfoils.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
}
