{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P9P3fJaa30da"
   },
   "source": [
    "# Solving Inverse Problems with NNs\n",
    "\n",
    "**TODOs: 1) bottom: show targets in results;**\n",
    "\n",
    "\n",
    "Inverse problems encompass a large class of practical scenarios that appear 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 a parameter in the model equation 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 demonstrated 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 PINN example of {doc}`physicalloss-code` or the 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 exaclty where the gradients from the DP training come in to guide the learning task towards solution that can reach the goal.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formulation\n",
    "\n",
    "With the notation from {doc}`overview-equations` this gives the minmization problem \n",
    "\n",
    "$\\text{arg min}_{\\theta} \\sum_m \\sum_i (f(x_{m,i} ; \\theta)-y^*_{m,i})^2$\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 $n$ 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  ))\\cdots)))) = (\\pde\\net)^n ( \\mathbf{u}_0,d_0 ) .\n",
    "$\n",
    "\n",
    "minimizes the loss above. The $\\mathrm{OP}$ network is a predictor that determines the action of the $\\mathrm{CFE}$ network given the target $d^*$, i.e., $\\mathrm{OP}(\\mathbf{u},d,d^*)=d_{OP}$,\n",
    "and $\\mathrm{CFE}$ acts additively on the velocity field via\n",
    "$\\mathrm{CFE}(\\mathbf{u},d,d_{OP}) = \\mathbf{u} + f_{\\mathrm{OP}}(\\mathbf{u},d,d_{OP};\\theta_{\\mathrm{OP}})$,  \n",
    "where we've used $f_{\\mathrm{OP}}$ to denote the NN representation of $\\mathrm{CFE}$).\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",
    "with an additional transport equation for the marker density $d$:\n",
    "\n",
    "$\\begin{aligned}\n",
    "  \\frac{\\partial d}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla d &= 0 .\n",
    "\\end{aligned}$\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.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning to Control Incompressible Fluids \n",
    "\n",
    "This notebook will walk you through data generation, supervised network initialization and end-to-end training using our differentiable PDE solver, [Φ<sub>Flow</sub>](https://github.com/tum-pbs/PhiFlow). \n",
    "(_Note: this example uses an older version of Φ<sub>Flow</sub> (1.4.1)._)\n",
    "\n",
    "The code below replicates the shape transitions experiment from the ICLR 2020 paper by Holl et al. {cite}`holl2019pdecontrol`, [Learning to Control PDEs with Differentiable Physics](https://ge.in.tum.de/publications/2020-iclr-holl/), further details can be found in section D.2 of the [appendix](https://openreview.net/pdf?id=HyeSin4FPB).\n",
    "\n",
    "First we need to load phiflow and download the control code helpers (which will end up under `./src`) and some numpy arrays with intial shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pwVPXx_Y30dd",
    "outputId": "17eeffaa-9651-48ea-b3a9-d479fc53adab"
   },
   "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",
    "# # this essentially copies over the data and code from https://github.com/holl-/PDE-Control\n",
    "# if not os.path.isfile('shapes/Shape_000000.npz'):\n",
    "#   import urllib.request\n",
    "#   url=\"https://ge.in.tum.de/download/2020-iclr-holl/control.zip\"\n",
    "#   print(\"Downloading, this can take a few minutes the first time...\")\n",
    "#   urllib.request.urlretrieve(url, 'control.zip')\n",
    "#   import os\n",
    "#   os.system(\"unzip control.zip\")\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; \n",
    "#sys.path.append('./src')\n",
    "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., the goal is to pre-compute our 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 (i.e., being trained for longer simulation sequences). \n",
    "\n",
    "First, let's set up a domain and basic parameters of the data generation step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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": "eQdzVAJH30dg"
   },
   "source": [
    "The following cell creates 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 somewhere within the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MaAzExNW30dh",
    "outputId": "ae340096-5016-4314-97a8-edd89f617040"
   },
   "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, two consecutive frames are required and to pretrain an $OP_n$, three frames with distance $n/2$ are needed.\n",
    "\n",
    "Instead, we create a second dataset which contains intermediate frames. This does not need to look like the actual dataset since it's only used for network initialization. Here, we linearly move a rectangle around the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A0QlQ-AL30di",
    "outputId": "1c8a2743-c846-44a2-fe74-f07123166df4"
   },
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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 OP$_n \\,\\, \\forall n\\in\\{2,4,8,16\\}$.\n",
    "The `ControlTraining` class is used to set up the optimization problem.\n",
    "\n",
    "The loss for the supervised initialization is defined as the observation loss at the center frame.\n",
    "\n",
    "$\\boldsymbol L_o^\\textrm{sup} = \\left|\\mathrm{OP}(o(t_i),o(t_j)) - u^*\\left(\\frac{t_i+t_j}{2}\\right)\\right|^2.$\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. If you have a set of pretrained networks, you can skip it and load the pretrained networks instead (see instructions below).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "b5ksGZwO30dl"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From PDE-Control/src/control/control_training.py:31: The name tf.reset_default_graph is deprecated. Please use tf.compat.v1.reset_default_graph instead.\n",
      "\n",
      "App created. Scene directory is /Users/thuerey/phi/model/control-training/sim_000000 (INFO), 2021-04-07 15:04:48,590n\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/app.py:23: The name tf.ConfigProto is deprecated. Please use tf.compat.v1.ConfigProto instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/app.py:25: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/session.py:18: The name tf.get_default_graph is deprecated. Please use tf.compat.v1.get_default_graph instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/app.py:51: The name tf.placeholder is deprecated. Please use tf.compat.v1.placeholder instead.\n",
      "\n",
      "Sequence class: <class 'control.sequences.SkipSequence'> (INFO), 2021-04-07 15:04:48,733n\n",
      "\n",
      "Partition length 2 sequence (from 0 to 2) at frame 1\n",
      "WARNING:tensorflow:From PDE-Control/src/control/pde/incompressible_flow.py:19: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/grid_layers.py:9: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.keras.layers.Conv2D` instead.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.7/site-packages/tensorflow_core/python/layers/convolutional.py:424: Layer.apply (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.__call__` method instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/struct/context.py:18: UserWarning: struct.unsafe() is deprecated. Use map() with new_type argument to avoid validation.\n",
      "  warnings.warn(\"struct.unsafe() is deprecated. Use map() with new_type argument to avoid validation.\")\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Dimensions must be equal, but are 35 and 37 for 'OP2/add_1' (op: 'AddV2') with input shapes: [?,35,35,16], [?,37,37,16].",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mInvalidArgumentError\u001b[0m                      Traceback (most recent call last)",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py\u001b[0m in \u001b[0;36m_create_c_op\u001b[0;34m(graph, node_def, inputs, control_inputs)\u001b[0m\n\u001b[1;32m   1606\u001b[0m   \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1607\u001b[0;31m     \u001b[0mc_op\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mc_api\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTF_FinishOperation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop_desc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1608\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mInvalidArgumentError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mInvalidArgumentError\u001b[0m: Dimensions must be equal, but are 35 and 37 for 'OP2/add_1' (op: 'AddV2') with input shapes: [?,35,35,16], [?,37,37,16].",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-2da40b2caa44>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m                           \u001b[0mdatapath\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpretrain_data_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mval_range\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mval_range\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_range\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrain_range\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrace_to_channel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'density'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m                           \u001b[0mobs_loss_frames\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m//\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable_networks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'OP%d'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m                           sequence_class=None).prepare()\n\u001b[0m\u001b[1;32m      8\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1000\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m         \u001b[0mapp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprogress\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Run Optimization for one batch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/control_training.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, n, pde, datapath, val_range, train_range, trace_to_channel, obs_loss_frames, trainable_networks, sequence_class, batch_size, view_size, learning_rate, learning_rate_half_life, dt, new_graph)\u001b[0m\n\u001b[1;32m     61\u001b[0m         \u001b[0mexecutor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecutor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPDEExecutor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mworld\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpde\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable_networks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     62\u001b[0m         \u001b[0msequence\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msequence\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msequence_class\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexecutor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 63\u001b[0;31m         \u001b[0msequence\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     64\u001b[0m         \u001b[0mall_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mall_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworldstate\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mframe\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msequence\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mframe\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     65\u001b[0m         \u001b[0;31m# --- Loss ---\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/sequences.py\u001b[0m in \u001b[0;36mexecute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     47\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     48\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 49\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpartition_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep_count\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     50\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mpartition_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_frame_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/sequences.py\u001b[0m in \u001b[0;36mpartition_execute\u001b[0;34m(self, n, start_frame_index, **kwargs)\u001b[0m\n\u001b[1;32m     53\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaf_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_frames\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstart_frame_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_frames\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstart_frame_index\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 55\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbranch_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_frame_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     56\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mleaf_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_frame\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mend_frame\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/sequences.py\u001b[0m in \u001b[0;36mbranch_execute\u001b[0;34m(self, n, start_frame_index, **kwargs)\u001b[0m\n\u001b[1;32m    109\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    110\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mbranch_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_frame_index\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 111\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpartition\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_frame_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    112\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    113\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/sequences.py\u001b[0m in \u001b[0;36mpartition\u001b[0;34m(self, n, start_frame_index)\u001b[0m\n\u001b[1;32m     63\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mpartition\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstart_frame_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m         self.executor.partition(n, self._frames[start_frame_index], self._frames[start_frame_index + n],\n\u001b[0;32m---> 65\u001b[0;31m                                 self._frames[start_frame_index + n // 2])\n\u001b[0m\u001b[1;32m     66\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     67\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__getitem__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mitem\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/hierarchy.py\u001b[0m in \u001b[0;36mpartition\u001b[0;34m(self, n, initial_frame, target_frame, center_frame)\u001b[0m\n\u001b[1;32m     56\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mpartition\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minitial_frame\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_frame\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcenter_frame\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m         \u001b[0mPartitioningExecutor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpartition\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minitial_frame\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_frame\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcenter_frame\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 58\u001b[0;31m         \u001b[0mcenter_frame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworldstate\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpde\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minitial_frame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworldstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_frame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworldstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'OP%d'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mn\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainable_networks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     59\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     60\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mcenter_frame\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindex\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworldsteps\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/pbdl-book/PDE-Control/src/control/pde/incompressible_flow.py\u001b[0m in \u001b[0;36mpredict\u001b[0;34m(self, n, initial_worldstate, target_worldstate, trainable)\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0md1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minitial_worldstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdensity\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_worldstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdensity\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvariable_scope\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'OP%d'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 20\u001b[0;31m             \u001b[0mdensity_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mu_net\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdomain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0md1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlevels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrainable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreuse\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAUTO_REUSE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopied_with\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mage\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mage\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0md2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mage\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;36m2.\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopied_with\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0md1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     21\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minitial_worldstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate_replaced\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdensity_pred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/standard_networks.py\u001b[0m in \u001b[0;36mu_net\u001b[0;34m(domain, input_fields, output_field, levels, filters, blocks_per_level, skip_combine, training, trainable, reuse)\u001b[0m\n\u001b[1;32m     43\u001b[0m         \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconv_layer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresolutions\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlevel_filters\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstrides\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mactivation\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'valid'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'down_convolution_%d'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mlevel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrainable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreuse\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mreuse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     44\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mblocks_per_level\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 45\u001b[0;31m             \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mresidual_block\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlevel_filters\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'down_res_block_%d_%d'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mlevel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtraining\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrainable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreuse\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mreuse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     46\u001b[0m         \u001b[0mresolutions\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minsert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     47\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/grid_layers.py\u001b[0m in \u001b[0;36mresidual_block\u001b[0;34m(grid, nb_channels, kernel_size, stride, activation, _project_shortcut, padding, name, training, trainable, reuse)\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mresidual_block\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrid\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnb_channels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mactivation\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mleaky_relu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_project_shortcut\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"SYMMETRIC\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreuse\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAUTO_REUSE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mgrid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrank\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_residual_block_2d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrid\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnb_channels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mactivation\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_project_shortcut\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreuse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     26\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mgrid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrank\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_residual_block_1d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrid\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnb_channels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mactivation\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_project_shortcut\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreuse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/grid_layers.py\u001b[0m in \u001b[0;36m_residual_block_2d\u001b[0;34m(grid, nb_channels, kernel_size, stride, activation, _project_shortcut, padding, name, training, trainable, reuse)\u001b[0m\n\u001b[1;32m     86\u001b[0m         \u001b[0;31m# shortcut = tf.layers.batch_normalization(shortcut, name=None if name is None else name+\"/normid\", training=training, trainable=trainable, reuse=reuse)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m     \u001b[0my\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mshortcut\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m     \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mactivation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mCenteredGrid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbox\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgrid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbox\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mextrapolation\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgrid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextrapolation\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py\u001b[0m in \u001b[0;36mbinary_op_wrapper\u001b[0;34m(x, y)\u001b[0m\n\u001b[1;32m    897\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname_scope\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mop_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    898\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 899\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    900\u001b[0m       \u001b[0;32melif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msparse_tensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSparseTensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    901\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py\u001b[0m in \u001b[0;36m_add_dispatch\u001b[0;34m(x, y, name)\u001b[0m\n\u001b[1;32m   1195\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mgen_math_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1197\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mgen_math_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_v2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1198\u001b[0m   \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1199\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mgen_math_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py\u001b[0m in \u001b[0;36madd_v2\u001b[0;34m(x, y, name)\u001b[0m\n\u001b[1;32m    547\u001b[0m   \u001b[0;31m# Add nodes to the TensorFlow graph.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    548\u001b[0m   _, _, _op = _op_def_lib._apply_op_helper(\n\u001b[0;32m--> 549\u001b[0;31m         \"AddV2\", x=x, y=y, name=name)\n\u001b[0m\u001b[1;32m    550\u001b[0m   \u001b[0m_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    551\u001b[0m   \u001b[0m_inputs_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/op_def_library.py\u001b[0m in \u001b[0;36m_apply_op_helper\u001b[0;34m(self, op_type_name, name, **keywords)\u001b[0m\n\u001b[1;32m    792\u001b[0m         op = g.create_op(op_type_name, inputs, dtypes=None, name=scope,\n\u001b[1;32m    793\u001b[0m                          \u001b[0minput_types\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minput_types\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattrs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mattr_protos\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 794\u001b[0;31m                          op_def=op_def)\n\u001b[0m\u001b[1;32m    795\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    796\u001b[0m       \u001b[0;31m# Conditionally invoke tfdbg v2's op callback(s).\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/util/deprecation.py\u001b[0m in \u001b[0;36mnew_func\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    505\u001b[0m                 \u001b[0;34m'in a future version'\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mdate\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'after %s'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mdate\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    506\u001b[0m                 instructions)\n\u001b[0;32m--> 507\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    508\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    509\u001b[0m     doc = _add_deprecated_arg_notice_to_docstring(\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py\u001b[0m in \u001b[0;36mcreate_op\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m   3355\u001b[0m         \u001b[0;32mraise\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input #%d is not a tensor: %s\"\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3356\u001b[0m     return self._create_op_internal(op_type, inputs, dtypes, input_types, name,\n\u001b[0;32m-> 3357\u001b[0;31m                                     attrs, op_def, compute_device)\n\u001b[0m\u001b[1;32m   3358\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3359\u001b[0m   def _create_op_internal(\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py\u001b[0m in \u001b[0;36m_create_op_internal\u001b[0;34m(self, op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_device)\u001b[0m\n\u001b[1;32m   3424\u001b[0m           \u001b[0minput_types\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minput_types\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3425\u001b[0m           \u001b[0moriginal_op\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_default_original_op\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3426\u001b[0;31m           op_def=op_def)\n\u001b[0m\u001b[1;32m   3427\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_create_op_helper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mret\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcompute_device\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcompute_device\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3428\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mret\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, node_def, g, inputs, output_types, control_inputs, input_types, original_op, op_def)\u001b[0m\n\u001b[1;32m   1768\u001b[0m           op_def, inputs, node_def.attr)\n\u001b[1;32m   1769\u001b[0m       self._c_op = _create_c_op(self._graph, node_def, grouped_inputs,\n\u001b[0;32m-> 1770\u001b[0;31m                                 control_input_ops)\n\u001b[0m\u001b[1;32m   1771\u001b[0m     \u001b[0;31m# pylint: enable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1772\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py\u001b[0m in \u001b[0;36m_create_c_op\u001b[0;34m(graph, node_def, inputs, control_inputs)\u001b[0m\n\u001b[1;32m   1608\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mInvalidArgumentError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1609\u001b[0m     \u001b[0;31m# Convert to ValueError for backwards compatibility.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1610\u001b[0;31m     \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1611\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1612\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mc_op\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Dimensions must be equal, but are 35 and 37 for 'OP2/add_1' (op: 'AddV2') with input shapes: [?,35,35,16], [?,37,37,16]."
     ]
    }
   ],
   "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": "cf4ff331-e175-4350-928b-7172043c22f1"
   },
   "outputs": [],
   "source": [
    "supervised_checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jD7nKXCv30dl"
   },
   "outputs": [],
   "source": [
    "# supervised_checkpoints = {'OP%d' % n: '../networks/shapes/supervised/OP%d_1000' % n for n in [2, 4, 8, 16]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "imuKMrRG30dm"
   },
   "source": [
    "## CFE Pretraining with Differentiable Physics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bklXI_o130dm"
   },
   "source": [
    "To pretrain the CFE, we set up a simulation with a single step of the differentiable solver.\n",
    "\n",
    "The following cell trains the CFE network from scratch. You can skip it 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'] = '../networks/shapes/CFE/CFE_2000'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we have not actually set up a simulation for the training. To evaluate the quality of our current solutions, we can use the `ControlTraining` class again but this time passing a `sequence_class`.\n",
    "With a sequence class, `ControlTraining` sets up the necessary physics steps and networks.\n",
    "Example sequence classes are `StaggeredSequence` and `RefinedSequence`.\n",
    "\n",
    "**TODO** , add preview?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g9HcP2oK30do"
   },
   "source": [
    "## End-to-end Training with Differentiable Physics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NipTaC2S30do"
   },
   "source": [
    "To initiate an end-to-end training of the $CFE$ and all $OP_n$ networks with the differentiable physics loss , 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": "75fc502c-b769-4fe8-9332-a1401d6476df"
   },
   "outputs": [],
   "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 skip it and load the checkpoint by running the cell after. 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. You can skip the next two cells and load the pretrained networks instead.*"
   ]
  },
  {
   "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: '../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 from the [paper](https://openreview.net/pdf?id=HyeSin4FPB)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8AgfnvE830dr"
   },
   "outputs": [],
   "source": [
    "states = staggered_app.infer_all_frames(test_range)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 461
    },
    "id": "dJvecqOY30dr",
    "outputId": "d20b09e3-a7bb-4298-ffba-97a966779650"
   },
   "outputs": [],
   "source": [
    "import pylab\n",
    "batches = [0, 1, 2]\n",
    "pylab.subplots(len(batches), 9, sharey='row', sharex='col', figsize=(12, 7))\n",
    "pylab.tight_layout(w_pad=0)\n",
    "for i, batch in enumerate(batches):\n",
    "    for t in range(9):\n",
    "        pylab.subplot(len(batches), 9, t + 1 + i * 9)\n",
    "        pylab.title('t=%d' % t * 2)\n",
    "        pylab.imshow(states[t * 2].density.data[batch, ..., 0], origin='lower')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sQE9mfID30ds"
   },
   "source": [
    "\n",
    "**TODO, show targets; discuss**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "- Change the `test_range` indices to look at different examples, and test the generalization of the trained controller networks.\n",
    "- Try using a `RefinedSequence` (instead of a `StaggeredSequence`) to train with the prediction refinement scheme. This will yield a further improved control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Shape Transitions.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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
