{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P9P3fJaa30da"
   },
   "source": [
    "# Solving Inverse Problems with NNs\n",
    "\n",
    "Inverse problems encompass a large class of practical applications in science. In general, the goal here is not to directly compute a physical field like the velocity at a future time (this is the typical scenario for a _forward_ solve), but instead more generically compute one or more parameters in the model equations such that certain constraints are fulfilled. A very common objective is to find the optimal setting for a single parameter given some constraints. E.g., this could be the global diffusion constant for an advection-diffusion model such that it fits measured data as accurately as possible. Inverse problems are encountered for any model parameter adjusted via observations, or the reconstruction of initial conditions, e.g., for particle imaging velocimetry (PIV). More complex cases aim for computing boundary geometries w.r.t. optmal conditions, e.g. to obtain a shape with minimal drag in a fluid flow.\n",
    "\n",
    "A key aspect below will be that we're not aiming for solving only a _single instance_ of an inverse problem, but we'd like to use deep learning to solve a _larger collection_ of inverse problems. Thus, unlike the physics-informed example of {doc}`physicalloss-code` or the differentiable physics (DP) optimization of {doc}`diffphys-code-ns`, where we've solved an optimization problem for specific instances of inverse problems, we now aim for training an NN that learns to solve a larger class of inverse problems, i.e., a whole solution manifold. Nonetheless, we of course need to rely on a certain degree of similarity for these problems, otherwise there's nothing to learn (and the implied assumption of continuity in the solution manifold breaks down).\n",
    "\n",
    "Below we will run a very challenging test case as a representative of these inverse problems: we will aim for computing a high dimensional control function that exerts forces over the full course of an incompressible fluid simulation in order to reach a desired goal state for a passively advected marker in the fluid. This means we only have very indirect constraints to be fulfilled (a single state at the end of a sequence), and a large number of degrees of freedom (the control force function is a space-time function with the same degrees of freedom as the flow field itself).\n",
    "\n",
    "The _long-term_ nature of the control is one of the aspects which makes this a tough inverse problem: any changes to the state of the physical system can lead to large change later on in time, and hence a controller needs to anticipate how the system will behave when it is influenced. This means an NN also needs to learn how the underlying physics evolve and change, and this is exactly where the gradients from the DP training come in to guide the learning task towards solution that can reach the goal.\n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/diffphys-control.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gaCwYgN_L-We"
   },
   "source": [
    "## Formulation\n",
    "\n",
    "With the notation from {doc}`overview-equations` this gives the minmization problem \n",
    "\n",
    "$$\n",
    "\\text{arg min}_{\\theta} \\sum_m \\sum_i (f(x_{m,i} ; \\theta)-y^*_{m,i})^2 , \n",
    "$$\n",
    "\n",
    "where $y^*_{m,i}$ denotes the samples of the target state of the marker field, \n",
    "and $x_{m,i}$ denotes the simulated state of the marker density.\n",
    "As before, the index $i$ samples our solution at different spatial locations (typically all grid cells), while the index $m$ here indicates a large collection of different target states.\n",
    "\n",
    "\n",
    "Our goal is to train two networks $\\mathrm{OP}$ and $\\mathrm{CFE}$ with weights\n",
    "$\\theta_{\\mathrm{OP}}$ and $\\theta_{\\mathrm{CFE}}$ such that a sequence \n",
    "\n",
    "$$\n",
    "\\newcommand{\\pde}{\\mathcal{P}}\n",
    "\\newcommand{\\net}{\\mathrm{CFE}}\n",
    "\\mathbf{u}_{n},d_{n} = \\pde(\\net(~\\pde(\\net(\\cdots \\pde(\\net( \\mathbf{u}_0,d_0, d_{OP} ))\\cdots)))) = (\\pde~\\net)^n ( \\mathbf{u}_0,d_0, d_{OP} ) .\n",
    "$$\n",
    "\n",
    "minimizes the loss above. The $\\mathrm{OP}$ network is a predictor that determines the state $d_{OP}$ that the action of the $\\mathrm{CFE}$ should aim for, i.e., it does the longer term planning from which to determine the action. Given the target $d^*$, it computes \n",
    "$d_{OP} = \\mathrm{OP}(d,d^*)= f_{\\mathrm{OP}}(d,d^{*};\\theta_{\\mathrm{OP}})$.\n",
    "The $\\mathrm{CFE}$ acts additively on the velocity field by computing\n",
    "$\\mathbf{u} + f_{\\mathrm{CFE}}(\\mathbf{u},d, f_{\\mathrm{OP}}(d,d^{*};\\theta_{\\mathrm{OP}}) ;\\theta_{\\mathrm{CFE}}) $, where we've used $f_{\\mathrm{OP}}$ and $f_{\\mathrm{CFE}}$ to denote the NN representations of $\\mathrm{OP}$ and $\\mathrm{CFE}$, respectively, and $d^{*}$ to denote the target density state. $\\theta_{\\mathrm{OP}}$ and $\\theta_{\\mathrm{CFE}}$ denote the corresponding network weights.\n",
    "\n",
    "For this problem, the model PDE $\\mathcal{P}$ contains a discretized version of the incompressible Navier-Stokes equations in two dimensions for a velocity $\\mathbf{u}$:\n",
    "\n",
    "$$\\begin{aligned}\n",
    "  \\frac{\\partial u_x}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla u_x &= - \\frac{1}{\\rho} \\nabla p \n",
    "  \\\\\n",
    "  \\frac{\\partial u_y}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla u_y &= - \\frac{1}{\\rho} \\nabla p \n",
    "  \\\\\n",
    "  \\text{s.t.} \\quad \\nabla \\cdot \\mathbf{u} &= 0,\n",
    "\\end{aligned}$$\n",
    "\n",
    "without explicit viscosity, and with an additional transport equation for the marker density $d$ given by \n",
    "$\\frac{\\partial d}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla d = 0$.\n",
    "\n",
    "To summarize, we have a predictor $\\mathrm{OP}$ that gives us a direction, an actor $\\mathrm{CFE}$ that exerts a force on a physical model $\\mathcal{P}$. They all need to play hand in hand to reach a given target after $n$ iterations of the simulation. As apparent from this formulation, it's not a simple inverse problem, especially due to the fact that all three functions are non-linear. This is exactly why the gradients from the DP approach are so important. (The viewpoint above also indicates that _reinforcement learning_ is a potential option. In {doc}`reinflearn-code.ipynb` we'll compare DP with these alternatives.)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tFb5WYgzL-Wf"
   },
   "source": [
    "\n",
    "---\n",
    "\n",
    "## Control of incompressible fluids \n",
    "\n",
    "The next sections will walk you through all the necessary steps from data generation to network training using [Φ<sub>Flow</sub>](https://github.com/tum-pbs/PhiFlow). Due to the complexity of the control problem,\n",
    "we'll start with a supervised initialization of the networks, before switching to a more accurate end-to-end training with DP.\n",
    "(_Note: this example uses an older version 1.4.1 of Φ<sub>Flow</sub>._)\n",
    "\n",
    "The code below replicates an inverse problem example (the shape transitions experiment) from [Learning to Control PDEs with Differentiable Physics](https://ge.in.tum.de/publications/2020-iclr-holl/) {cite}`holl2019pdecontrol`, further details can be found in section D.2 of the paper's [appendix](https://openreview.net/pdf?id=HyeSin4FPB).\n",
    "\n",
    "First we need to load phiflow and check out the _PDE-Control_ git repository, which also contains some numpy arrays with intial shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pwVPXx_Y30dd"
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@1.4.1\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from phi.flow import *\n",
    "\n",
    "if not os.path.isdir('PDE-Control'):\n",
    "  print(\"Cloning, PDE-Control repo, this can take a moment\")\n",
    "  os.system(\"git clone --recursive https://github.com/holl-/PDE-Control.git\")\n",
    "    \n",
    "# now we can load the necessary phiflow libraries and helper functions\n",
    "import sys; sys.path.append('PDE-Control/src')\n",
    "from shape_utils import load_shapes, distribute_random_shape\n",
    "from control.pde.incompressible_flow import IncompressibleFluidPDE\n",
    "from control.control_training import ControlTraining\n",
    "from control.sequences import StaggeredSequence, RefinedSequence\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vQCiicZv30de"
   },
   "source": [
    "## Data generation\n",
    "\n",
    "Before starting the training, we have to generate a data set to train with, i.e., a set of ground truth time sequences $u^*$. Due to the complexity of the training below, we'll use a staged approach that pre-trains a supervised network as a rough initialization, and then refines it to learn control looking further and further ahead into the future. (This will be realized by training specialized models that deal with longer and longer sequences.) \n",
    "\n",
    "First, let's set up a domain and basic parameters of the data generation step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EthZeAV030df"
   },
   "outputs": [],
   "source": [
    "domain = Domain([64, 64])  # 1D Grid resolution and physical size\n",
    "step_count = 16  # how many solver steps to perform\n",
    "dt = 1.0  # Time increment per solver step\n",
    "example_count = 1000\n",
    "batch_size = 100\n",
    "data_path = 'shape-transitions'\n",
    "pretrain_data_path = 'moving-squares'\n",
    "shape_library = load_shapes('PDE-Control/notebooks/shapes')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uDvxnTavZQk2"
   },
   "source": [
    "The `shape_library` in the last line contains ten different shapes that we'll use to intialize a marker density with at random positions.\n",
    "\n",
    "This is what the shapes look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 142
    },
    "id": "9m8HFEPzQD70",
    "outputId": "010df794-01e9-49b5-b4a2-492bb4050311"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12088 characters]\n",
      "text/plain": [
       "<Figure size 1224x360 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab\n",
    "pylab.subplots(1, len(shape_library), figsize=(17, 5))\n",
    "for t in range(len(shape_library)):\n",
    "    pylab.subplot(1, len(shape_library), t+1)\n",
    "    pylab.imshow(shape_library[t], origin='lower')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eQdzVAJH30dg"
   },
   "source": [
    "The following cell uses these shapes to create the dataset we want to train our model on.\n",
    "Each example consists of a start and target (end) frame which are generated by placing a random shape from the `shape_library` somewhere within the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MaAzExNW30dh",
    "outputId": "14de4d69-177c-4f0e-c603-5f041ce72cce"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape-transitions/sim_000000\n",
      "shape-transitions/sim_000100\n",
      "shape-transitions/sim_000200\n",
      "shape-transitions/sim_000300\n",
      "shape-transitions/sim_000400\n",
      "shape-transitions/sim_000500\n",
      "shape-transitions/sim_000600\n",
      "shape-transitions/sim_000700\n",
      "shape-transitions/sim_000800\n",
      "shape-transitions/sim_000900\n"
     ]
    }
   ],
   "source": [
    "for scene in Scene.list(data_path): scene.remove()\n",
    "\n",
    "for _ in range(example_count // batch_size):\n",
    "    scene = Scene.create(data_path, count=batch_size, copy_calling_script=False)\n",
    "    print(scene)\n",
    "    start = distribute_random_shape(domain.resolution, batch_size, shape_library)\n",
    "    end__ = distribute_random_shape(domain.resolution, batch_size, shape_library)\n",
    "    [scene.write_sim_frame([start], ['density'], frame=f) for f in range(step_count)]\n",
    "    scene.write_sim_frame([end__], ['density'], frame=step_count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AIh2z-be30di"
   },
   "source": [
    "Since this dataset does not contain any intermediate frames, it does not allow for supervised pretraining. This is because to pre-train a CFE network, two consecutive frames are required while to pretrain an $\\mathrm{OP}_n$ network, three frames with a distance of $n/2$ are needed.\n",
    "\n",
    "Instead, we create a second dataset which contains these intermediate frames. This does not need to be very close to the actual dataset since it's only used for network initialization via pretraining. Here, we linearly move a rectangle around the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A0QlQ-AL30di",
    "outputId": "5f464216-98ba-4c10-c660-887d3394ba43"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "moving-squares/sim_000000\n",
      "moving-squares/sim_000100\n",
      "moving-squares/sim_000200\n",
      "moving-squares/sim_000300\n",
      "moving-squares/sim_000400\n",
      "moving-squares/sim_000500\n",
      "moving-squares/sim_000600\n",
      "moving-squares/sim_000700\n",
      "moving-squares/sim_000800\n",
      "moving-squares/sim_000900\n"
     ]
    }
   ],
   "source": [
    "for scene in Scene.list(pretrain_data_path): scene.remove()\n",
    "\n",
    "for scene_index in range(example_count // batch_size):\n",
    "    scene = Scene.create(pretrain_data_path, count=batch_size, copy_calling_script=False)\n",
    "    print(scene)\n",
    "    pos0 = np.random.randint(10, 56, (batch_size, 2))  # start position\n",
    "    pose = np.random.randint(10, 56, (batch_size, 2))  # end position\n",
    "    size = np.random.randint(6,  10,  (batch_size, 2))\n",
    "    for frame in range(step_count+1):\n",
    "        time = frame / float(step_count + 1)\n",
    "        pos = np.round(pos0 * (1 - time) + pose * time).astype(np.int)\n",
    "        density = AABox(lower=pos-size//2, upper=pos-size//2+size).value_at(domain.center_points())\n",
    "        scene.write_sim_frame([density], ['density'], frame=frame)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cwrv1oHs30dj"
   },
   "source": [
    "## Supervised initialization\n",
    "\n",
    "First we define a split of the 1000 data samples into 100 test, 100 validation, and 800 training samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZcbTw9cF30dk"
   },
   "outputs": [],
   "source": [
    "test_range = range(100)\n",
    "val_range = range(100, 200)\n",
    "train_range = range(200, 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "22TTd6xJ30dk"
   },
   "source": [
    "The following cell trains all $\\mathrm{OP}_n \\,\\, \\forall n\\in\\{2,4,8,16\\}$. Here the $n$ indicates the number of time steps for which the network predicts the target. In order to cover longer time horizons, we're using factors of two here to hierarchically divide the time intervals during which the physical system should be controlled.\n",
    "\n",
    "The `ControlTraining` class is used to set up the corresponding optimization problem.\n",
    "The loss for the supervised initialization is defined as the observation loss in terms of velocity at the center frame:\n",
    "\n",
    "$$\n",
    "L_o^\\textrm{sup} = \\left| \\mathrm{OP}\\Big(d_{t_i},d_{t_j}\\Big) - d^*_{(t_i+t_j)/2} \\right|^2 .\n",
    "$$\n",
    "\n",
    "Consequently, no sequence needs to be simulated (`sequence_class=None`) and an observation loss is required at frame $\\frac n 2$ (`obs_loss_frames=[n // 2]`).\n",
    "The pretrained network checkpoints are stored in `supervised_checkpoints`.\n",
    "\n",
    "*Note: The next cell will run for some time. The PDE-Control git repo comes with a set of pre-trained networks. So if you want to focus on the evaluation, you can skip the training and load the pretrained networks instead by commenting out the training cells, and uncommenting the cells for loading below.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b5ksGZwO30dl"
   },
   "outputs": [],
   "source": [
    "supervised_checkpoints = {}\n",
    "\n",
    "for n in [2, 4, 8, 16]:\n",
    "    app = ControlTraining(n, IncompressibleFluidPDE(domain, dt),\n",
    "                          datapath=pretrain_data_path, val_range=val_range, train_range=train_range, trace_to_channel=lambda _: 'density',\n",
    "                          obs_loss_frames=[n//2], trainable_networks=['OP%d' % n],\n",
    "                          sequence_class=None).prepare()\n",
    "    for i in range(1000):\n",
    "        app.progress()  # Run Optimization for one batch\n",
    "    supervised_checkpoints['OP%d' % n] = app.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_A9-V1Vg30dl",
    "outputId": "994c7fdc-a5aa-4769-eab8-1f27f29ad082"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'OP16': '/root/phi/model/control-training/sim_000003/checkpoint_00001000',\n",
       " 'OP2': '/root/phi/model/control-training/sim_000000/checkpoint_00001000',\n",
       " 'OP4': '/root/phi/model/control-training/sim_000001/checkpoint_00001000',\n",
       " 'OP8': '/root/phi/model/control-training/sim_000002/checkpoint_00001000'}"
      ]
     },
     "execution_count": 8,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "supervised_checkpoints # this is where the checkpoints end up when re-training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jD7nKXCv30dl"
   },
   "outputs": [],
   "source": [
    "# supervised_checkpoints = {'OP%d' % n: 'PDE-Control/networks/shapes/supervised/OP%d_1000' % n for n in [2, 4, 8, 16]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "imuKMrRG30dm"
   },
   "source": [
    "\n",
    "This concludes the pretraining of the OP networks. This makes it possible to at least perform a rough planning of the motions, which will be refined via end-to-end training below. Before this, we'll initialize the $\\mathrm{CFE}$ networks such that we can perform _actions_, i.e., apply forces to the simulation. This is completely decoupled from the $\\mathrm{OP}$ networks.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bklXI_o130dm"
   },
   "source": [
    "## CFE pretraining with differentiable physics\n",
    "\n",
    "To pretrain the $\\mathrm{CFE}$ networks, we set up a simulation with a single step of the differentiable solver.\n",
    "\n",
    "The following cell trains the $\\mathrm{CFE}$ network from scratch. If you have a pretrained model at hand, you can skip the training an load the checkpoint by running the cell after."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kCtOr3ed30dn"
   },
   "outputs": [],
   "source": [
    "app = ControlTraining(1, IncompressibleFluidPDE(domain, dt),\n",
    "                      datapath=pretrain_data_path, val_range=val_range, train_range=train_range, trace_to_channel=lambda _: 'density',\n",
    "                      obs_loss_frames=[1], trainable_networks=['CFE']).prepare()\n",
    "for i in range(1000):\n",
    "    app.progress()  # Run Optimization for one batch\n",
    "supervised_checkpoints['CFE'] = app.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-KOcgr5M30dn"
   },
   "outputs": [],
   "source": [
    "# supervised_checkpoints['CFE'] = 'PDE-Control/networks/shapes/CFE/CFE_2000'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "34EA3S98L-Wr"
   },
   "source": [
    "Note that we have not actually set up a simulation for the training, as the $\\mathrm{CFE}$ network only infers forces between pairs of states. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kQZLbOolI4-2"
   },
   "outputs": [],
   "source": [
    "# [TODO, show preview of CFE only?]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g9HcP2oK30do"
   },
   "source": [
    "## End-to-end training with differentiable physics\n",
    "\n",
    "Now that first versions of both network types exist, we can initiate the most important step of the setup at hand: the coupled end-to-end training of both networks via the differentiable fluid solver. While the pretraining stages relied on supervised training, the next step will yield a significantly improved quality for the control. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NipTaC2S30do"
   },
   "source": [
    "To initiate an end-to-end training of the $\\mathrm{CFE}$ and all $\\mathrm{OP}_n$ networks with the differentiable physics loss in phiflow, we create a new `ControlTraining` instance with the staggered execution scheme.\n",
    "\n",
    "The following cell builds the computational graph with `step_count` solver steps without initializing the network weights.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iJjrTrl530dp",
    "outputId": "69c6407e-efb2-41eb-df62-e472e9442834"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "App created. Scene directory is /root/phi/model/control-training/sim_000005 (INFO), 2021-04-09 00:41:17,299n\n",
      "\n",
      "Sequence class: <class 'control.sequences.StaggeredSequence'> (INFO), 2021-04-09 00:41:17,305n\n",
      "\n",
      "Partition length 16 sequence (from 0 to 16) at frame 8\n",
      "Partition length 8 sequence (from 0 to 8) at frame 4\n",
      "Partition length 4 sequence (from 0 to 4) at frame 2\n",
      "Partition length 2 sequence (from 0 to 2) at frame 1\n",
      "Execute -> 1\n",
      "Execute -> 2\n",
      "Partition length 2 sequence (from 2 to 4) at frame 3\n",
      "Execute -> 3\n",
      "Execute -> 4\n",
      "Partition length 4 sequence (from 4 to 8) at frame 6\n",
      "Partition length 2 sequence (from 4 to 6) at frame 5\n",
      "Execute -> 5\n",
      "Execute -> 6\n",
      "Partition length 2 sequence (from 6 to 8) at frame 7\n",
      "Execute -> 7\n",
      "Execute -> 8\n",
      "Partition length 8 sequence (from 8 to 16) at frame 12\n",
      "Partition length 4 sequence (from 8 to 12) at frame 10\n",
      "Partition length 2 sequence (from 8 to 10) at frame 9\n",
      "Execute -> 9\n",
      "Execute -> 10\n",
      "Partition length 2 sequence (from 10 to 12) at frame 11\n",
      "Execute -> 11\n",
      "Execute -> 12\n",
      "Partition length 4 sequence (from 12 to 16) at frame 14\n",
      "Partition length 2 sequence (from 12 to 14) at frame 13\n",
      "Execute -> 13\n",
      "Execute -> 14\n",
      "Partition length 2 sequence (from 14 to 16) at frame 15\n",
      "Execute -> 15\n",
      "Execute -> 16\n",
      "Target loss: Tensor(\"truediv_16:0\", shape=(), dtype=float32) (INFO), 2021-04-09 00:41:44,654n\n",
      "\n",
      "Force loss: Tensor(\"truediv_107:0\", shape=(), dtype=float32) (INFO), 2021-04-09 00:41:51,312n\n",
      "\n",
      "Supervised loss at frame 16: Tensor(\"truediv_108:0\", shape=(), dtype=float32) (INFO), 2021-04-09 00:41:51,332n\n",
      "\n",
      "Setting up loss (INFO), 2021-04-09 00:41:51,338n\n",
      "\n",
      "Preparing data (INFO), 2021-04-09 00:42:32,417n\n",
      "\n",
      "Initializing variables (INFO), 2021-04-09 00:42:32,443n\n",
      "\n",
      "Model variables contain 0 total parameters. (INFO), 2021-04-09 00:42:36,418n\n",
      "\n",
      "Validation (000000): Learning_Rate: 0.0005, GT_obs_16: 399498.75, Loss_reg_unscaled: 1.2424506, Loss_reg_scale: 1.0, Loss: 798997.5 (INFO), 2021-04-09 00:42:59,618n\n",
      "\n"
     ]
    }
   ],
   "source": [
    "staggered_app = ControlTraining(step_count, IncompressibleFluidPDE(domain, dt),\n",
    "                                datapath=data_path, val_range=val_range, train_range=train_range, trace_to_channel=lambda _: 'density',\n",
    "                                obs_loss_frames=[step_count], trainable_networks=['CFE', 'OP2', 'OP4', 'OP8', 'OP16'],\n",
    "                                sequence_class=StaggeredSequence, learning_rate=5e-4).prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tUzt45Xv30dp"
   },
   "source": [
    "The next cell initializes the networks using the supervised checkpoints and then trains all networks jointly. You can increase the number of optimization steps or execute the next cell multiple times to further increase performance.\n",
    "\n",
    "*Note: The next cell will run for some time. Optionally, you can skip this cell and load a pretrained networks instead with code in the cell below.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oMTXnLqR30dq"
   },
   "outputs": [],
   "source": [
    "staggered_app.load_checkpoints(supervised_checkpoints)\n",
    "for i in range(1000):\n",
    "    staggered_app.progress()  # run staggered Optimization for one batch\n",
    "staggered_checkpoint = staggered_app.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xpLvDj5-30dq"
   },
   "outputs": [],
   "source": [
    "# staggered_checkpoint = {net: 'PDE-Control/networks/shapes/staggered/all_53750' for net in ['CFE', 'OP2', 'OP4', 'OP8', 'OP16']}\n",
    "# staggered_app.load_checkpoints(staggered_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KkQxqh1u30dr"
   },
   "source": [
    "Now that the network is trained, we can infer some trajectories from the test set.\n",
    "(This corresponds to Fig 5b and 18b of the [original paper](https://openreview.net/pdf?id=HyeSin4FPB).)\n",
    "\n",
    "The following cell takes the first one hundred configurations, i.e. our test set as defined by `test_range`, and let's the network infer solutions for the corresponding inverse problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8AgfnvE830dr"
   },
   "outputs": [],
   "source": [
    "states = staggered_app.infer_all_frames(test_range)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rDEPL8E9fiFt"
   },
   "source": [
    "Via the index list `batches` below, you can choose to display some of the solutions. Each row show a temporal sequence starting with the initial condition, and evolving the simulation with the NN control forces for 16 time steps. The last step, at $t=16$ should match the target shown in the image on the far right."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 413
    },
    "id": "dJvecqOY30dr",
    "outputId": "27fbfd47-9d1f-4ded-8238-cadd905e3429"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53572 characters]\n",
      "text/plain": [
       "<Figure size 1008x432 with 30 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light",
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "batches = [0,1,2]\n",
    "\n",
    "pylab.subplots(len(batches), 10, sharey='row', sharex='col', figsize=(14, 6))\n",
    "pylab.tight_layout(w_pad=0)\n",
    "\n",
    "# solutions\n",
    "for i, batch in enumerate(batches):\n",
    "    for t in range(9):\n",
    "        pylab.subplot(len(batches), 10, t + 1 + i * 10)\n",
    "        pylab.title('t=%d' % (t * 2))\n",
    "        pylab.imshow(states[t * 2].density.data[batch, ..., 0], origin='lower')\n",
    "\n",
    "# add targets\n",
    "testset = BatchReader(Dataset.load(staggered_app.data_path,test_range), staggered_app._channel_struct)[test_range]\n",
    "for i, batch in enumerate(batches):\n",
    "        pylab.subplot(len(batches), 10, i * 10 + 10)\n",
    "        pylab.title('target')\n",
    "        pylab.imshow(testset[1][i,...,0], origin='lower')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UCU0AjuUMk-K"
   },
   "source": [
    "As you can see in the two right-most columns, the network does a very good job at solving these inverse problems: the fluid marker is pushed to the right spot and deformed in the right way to match the target.\n",
    "\n",
    "What looks fairly simple here is actually a tricky task for a neural network: \n",
    "it needs to guide a full 2D Navier-Stokes simulation over the course of 16 time integration steps. Hence, if the applied forces are slightly off or incoherent, the fluid can start swirling and moving chaotically. However, the network has learned to keep the motion together, and guide the  marker density to the target location.\n",
    "\n",
    "Next, we quantify the achieved error rate by comparing the mean absolute error in terms of the final density configuration relative to the initial density. With the standard training setup above, the next cell should give a relative residual error of 5-6%. Vice versa, this means that more than ca. 94% of the marker density ends up in the right spot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BOgtZlJrM-0N",
    "outputId": "6beb52e6-08b3-4976-96d4-59b660dd91a8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative MAE: 0.05450168251991272\n"
     ]
    }
   ],
   "source": [
    "errors = []\n",
    "for batch in enumerate(test_range):\n",
    "  initial = np.mean( np.abs( states[0].density.data[batch, ..., 0] - testset[1][batch,...,0] )) \n",
    "  solution = np.mean( np.abs( states[16].density.data[batch, ..., 0] - testset[1][batch,...,0] )) \n",
    "  errors.append( solution/initial )\n",
    "print(\"Relative MAE: \"+format(np.mean(errors)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XKzyhAGjL-Wv"
   },
   "source": [
    "## Next steps\n",
    "\n",
    "For further experiments with this source code, you can, e.g.:\n",
    "\n",
    "- Change the `test_range` indices to look at different examples, or test the generalization of the trained controller networks by using new shapes as targets.\n",
    "- Try using a `RefinedSequence` (instead of a `StaggeredSequence`) to train with the prediction refinement scheme. This will yield a further improved control and reduced density error."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "diffphys-control.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
}
