{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o4JZ84moBKMr"
   },
   "source": [
    "# Differentiable Fluid Simulations\n",
    "\n",
    "\n",
    "**(old phiflow1.x version)**\n",
    "\n",
    "\n",
    "Next, we'll target a more complex example with the Navier-Stokes equations as model. We'll target a 2D case with velocity $\\mathbf{u}$, no explicit viscosity term, and a marker density $d$ that drives a simple Boussinesq buoyancy term $\\eta d$ adding a force along the y dimension:\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 + \\eta d\n",
    "  \\\\\n",
    "  \\text{s.t.} \\quad \\nabla \\cdot \\mathbf{u} &= 0,\n",
    "  \\\\\n",
    "  \\frac{\\partial d}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla d &= 0 \n",
    "\\end{aligned}$\n",
    "\n",
    "As optimization objective we'll consider a more difficult variant of the previous example: the state of the observed density $d$ should match a given target after $n=20$ steps of simulation. In contrast to before, the marker $d$ cannot be modified in any way, but only the initial state of the velocity $\\mathbf{u}$ at $t=0$. This gives us a split between observable quantities for the loss formulation, and quantities that we can interact with during the optimization (or later on via NNs).\n",
    "\n",
    "First, let's get the loading of python modules out of the way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "da1uZcDXdVcF"
   },
   "outputs": [],
   "source": [
    "#!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@1.5\n",
    "# # %tensorflow_version 1.x  # colab\n",
    "\n",
    "#!pip install --upgrade --quiet phiflow\n",
    "\n",
    "from phi.flow import *  # The Dash GUI is not supported on Google Colab, ignore the warning\n",
    "import pylab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BVV1IKVqDfLl"
   },
   "source": [
    "## Setting up the simulation\n",
    "\n",
    "To make things a bit more interesting - and to move a bit closer to a NN training process - let's set up of four fluid simulations that run in parallel, i.e. a mini batch similar to DL training. In phiflow we can directly pass a `batch_size=4` parameter to the `Fluid` object. Each fluid simulation is fully independent. In this case they differ by having circular Inflows at different locations.\n",
    "\n",
    "Like before, let's plot the marker density after a few steps of simulation (each call to `step()` now updates all four simulations). Note that the boundaries between the four simulations are not visible in the image, but it shows four completely separate density states. The different inflow positions in conjunction with the solid wall boundaries (zero Dirichlet for velocity, and Neumann for pressure), result in four  different end states of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "WrA3IXDxv31P"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1283dad10>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 11516 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "world = World()\n",
    "fluid = world.add(Fluid(Domain([40, 32], boundaries=CLOSED), buoyancy_factor=0.05, batch_size=4), physics=IncompressibleFlow())\n",
    "centers = [[5,10], [5,12], [5,14], [5,16]]\n",
    "world.add(Inflow(Sphere(center=centers, radius=3), rate=0.2));\n",
    "\n",
    "for frame in range(20):\n",
    "  world.step(dt=1.5)\n",
    "\n",
    "pylab.imshow(np.concatenate(fluid.density.data[...,0], axis=1), origin='lower', cmap='magma')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we see four simulations, (0) to (3). This final density of simulation (0), with the curved plume on the far left, will be our **reference state**, while the initial velocity of the other three will be modified in the optimization procedure below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rdSTbMoaS0Uz"
   },
   "source": [
    "## Differentiation\n",
    "\n",
    "The simulation we just computed was using purely (non-differentiable) operations from numpy.\n",
    "To enable differentiability, we need to build a TensorFlow graph that computes this result.\n",
    "\n",
    "(Note, the first line of the next block, `%tensorflow_version 1.x` is only necessary when running in environments that by default have newer tensorflow versions installed, e.g., `colab`. Uncomment if you're running this notebook there.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "mphMP0sYIOz-"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Could not load resample cuda libraries: CUDA binaries not found at /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/cuda/build/resample.so. Run \"python setup.py cuda\" to compile them\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/util.py:119: The name tf.AUTO_REUSE is deprecated. Please use tf.compat.v1.AUTO_REUSE instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/profiling.py:12: The name tf.RunOptions is deprecated. Please use tf.compat.v1.RunOptions instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/profiling.py:13: The name tf.RunMetadata is deprecated. Please use tf.compat.v1.RunMetadata instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/session.py:17: 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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/flow.py:15: UserWarning: TensorFlow-CUDA solver is not available. To compile it, download phiflow sources and run\n",
      "$ python setup.py tf_cuda\n",
      "before reinstalling phiflow.\n",
      "  warnings.warn(\"TensorFlow-CUDA solver is not available. To compile it, download phiflow sources and run\\n$ python setup.py tf_cuda\\nbefore reinstalling phiflow.\")\n"
     ]
    }
   ],
   "source": [
    "# %tensorflow_version 1.x\n",
    "from phi.tf.flow import *  # Causes deprecation warnings with TF 1.15\n",
    "import pylab\n",
    "session = Session(None)  # Used to run the TensorFlow graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3mpyowRYUSS4"
   },
   "source": [
    "Now that we have imported `phi.tf.fluid`, let's set up the simulation just like before. But now, we want to start from a velocity that we can modify, i.e. a variable. Then we can optimize these initial velocities so that all simulations arrive at a final state that is similar to the first simulation from the previous example. I.e., the state shown in the left-most image above.\n",
    "\n",
    "This is a fairly tough task: we're producing diffent dynamics by changing the boundary conditions (the marker inflow position), and an optimizer should now find a single initial velocity state, that gives the same state as simulation `0` above at $t=30$. Thus, after 20 steps with $\\Delta t=1.5$ the simulation should reproduce a different set of boundary conditions from the velocity state. It would be much easier to simply change the position of the marker inflow to arrive at this goal, but -- to make things a bit more difficult -- the inflow is _not_ a degree of freedom. The optimizer can only change the velocity $\\mathbf{u}$ at time $t=0$.\n",
    "\n",
    "To achieve this, we create a TensorFlow variable for the velocity at t=0.\n",
    "It is initialized with zeros (like with the NumPy simulation above) and can later be used as a target for optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "NlJMJikaHOL6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/session.py:28: The name tf.global_variables_initializer is deprecated. Please use tf.compat.v1.global_variables_initializer instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/session.py:29: The name tf.train.Saver is deprecated. Please use tf.compat.v1.train.Saver instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "world = World()\n",
    "fluid = world.add(Fluid(Domain([40, 32], boundaries=CLOSED), buoyancy_factor=0.05, batch_size=4), physics=IncompressibleFlow())\n",
    "world.add(Inflow(Sphere(center=centers, radius=3), rate=0.2));\n",
    "fluid.velocity = variable(fluid.velocity)  # create TensorFlow variable\n",
    "initial_state = fluid.state  # Remember the state at t=0 for later visualization\n",
    "session.initialize_variables()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vSdGIEzCgq5-"
   },
   "source": [
    "The simulation now contains variables in the initial state.\n",
    "Since all later states depend on the value of the variable, the `step` method cannot directly compute concrete state values.\n",
    "Instead, `world.step` will extend the TensorFlow graph by the operations needed to perform the step.\n",
    "\n",
    "To execute the graph with actual data, we can use `session.run`, just like with regular TensorFlow 1.x. While `run` would usually be used to infer predictions from a learning model, it now executes the graph of simulation steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "wSrIezfWHjcQ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/tf_backend.py:167: The name tf.sparse_tensor_dense_matmul is deprecated. Please use tf.sparse.sparse_dense_matmul instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/tf_backend.py:61: The name tf.div_no_nan is deprecated. Please use tf.math.divide_no_nan instead.\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x143f6c290>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAACICAYAAAD+r7D/AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8li6FKAAAK10lEQVR4nO3db4wcd33H8fendhL+SbUDyHLttHGFBTKhOJGFHBVVbQjCoRTnQVSlRaqrRrIqUTVUSChpnhSpUotaQanUEllJGreKEmhIGzeif4zrij4oBhtMSGJSO6UUR04MAgOlEmDz7YOZSzbnvdze3e7e/S7vl3TamdnZnd+f1edmfzOzk6pCktSen1juAkiSFscAl6RGGeCS1CgDXJIaZYBLUqMMcElq1NppbiyJ5yxK0sJ9s6peO3vhVAO8s2b6m5Skpl342rClDqFIUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjZo3wJO8LMnnknwpyeNJPtgv35LkSJJTST6e5NLJF1eSNGOUPfAfANdV1ZuB7cCuJDuBDwEfqarXAd8GbplcMSVJs80b4NX53372kv6vgOuAB/vl+4EbJ1JCSdJQI42BJ1mT5DhwFjgIPAWcq6rz/SqngU2TKaIkaZiRAryqLlTVdmAz8BbgDaNuIMneJEeTHF1kGSVJQyzoLJSqOgccBq4F1iVZ2z+1GXh6jtfsq6odVbVjSSWVJL3AKGehvDbJun765cDbgRN0QX5Tv9oe4OFJFVKSdLG186/CRmB/kjV0gf+JqnokyRPAA0n+EPgicPcEyylJmiVVNb2NJQVrprY9SVodLhwbNgztlZiS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUqFHuiXlFksNJnkjyeJJb++WXJzmY5GT/uH7yxZUkzRhlD/w88P6q2gbsBN6bZBtwG3CoqrYCh/p5SdKUzBvgVXWmqr7QT3+P7o70m4DdwP5+tf3AjZMqpCTpYgsaA09yJXA1cATYUFVn+qeeATaMtWSSpBe1dtQVk7wK+CTwvqr6bpLnnquq6u44P/R1e4G9Sy2oJOmFRtoDT3IJXXjfV1UP9YufTbKxf34jcHbYa6tqX1XtqKod4yiwJKkzylkoAe4GTlTVhweeOgDs6af3AA+Pv3iSpLmkaujIx/MrJG8F/h34MvDjfvHv042DfwL4aeBrwK9W1bfmea+CNUstsyS9xFw4NmwUY94AHycDXJIWY3iAeyWmJDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNWqUe2Lek+RskscGll2e5GCSk/3j+skWU5I02yh74PcCu2Ytuw04VFVbgUP9vCRpiuYN8Kr6DDD7ZsW7gf399H7gxjGXS5I0j8WOgW+oqjP99DPAhjGVR5I0orVLfYOqqu5u88Ml2QvsXep2JEkvtNg98GeTbAToH8/OtWJV7auqHVW1Y5HbkiQNsdgAPwDs6af3AA+PpziSpFGNchrh/cB/AK9PcjrJLcAfA29PchK4vp+XJE1RquYcvh7/xpKCNRN7/03rfumiZU+fOzyx7b0U2caTZxtP1rD2hZXexheODRuG9kpMSWqUAS5JjWpyCGXwK9A/bL/quentv33xex+/8wIAv3L8uV8CWOFflVYG23jybOPJG9bGw9oXVnobO4QiSauKAS5JjVrylZjLYfDr5ps+vfu56QtD1n3TTf1rrn9+2TX/tlK+Fq1ctvHk2caTN6yNh7UvtNnG7oFLUqOa2gOfOSAxeBBirv+msw2+ZtPx5w9srJyDFCuDbTx5tvHkvVTa2D1wSWqUAS5JjTLAJalRBrgkNcoAl6RGNXUWysxR4ON3DpzbedNor525THbwfXQx23jybOPJe6m0sXvgktSopvbAZwz+2MzgVVOj/giQ5mcbT55tPHnD2nghP2a10rkHLkmNMsAlqVFL+j3wJLuAj9L9yPddVfWi98b0lmrts40nzzaeLG+pBiRZA/wFcAOwDfi1JNsWX0BJ0kIseg88ybXAH1TVO/r52wGq6o9e5DUT3QOXpNVp/Hfk2QR8fWD+dL/sBZLsTXI0ydElbEuSNMvETyOsqn3APpjZA5ckjcNSAvxp4IqB+c39shfzTbjw/e5x1XgNq6s+sPrqZH1WvtVWp3HX52eGLVzKGPha4D+Bt9EF9+eBX6+qx+d53dFhYzmtWm31gdVXJ+uz8q22Ok2rPoveA6+q80l+B/hnuiOT98wX3pKk8VnSGHhVfQr41JjKIklagOW4EnPfMmxzklZbfWD11cn6rHyrrU5Tqc+SrsSUJC0ffwtFkho11QBPsivJk0lOJbltmtsehyRXJDmc5Ikkjye5tV9+eZKDSU72j+uXu6wLkWRNki8meaSf35LkSN9PH09y6XKXcSGSrEvyYJKvJDmR5NqW+yjJ7/Wft8eS3J/kZS31UZJ7kpxN8tjAsqH9kc6f9/V6NMk1y1fyuc1Rpz/pP3OPJvm7JOsGnru9r9OTSd4xrnJMLcBXyW+nnAfeX1XbgJ3Ae/s63AYcqqqtwKF+viW3AicG5j8EfKSqXgd8G7hlWUq1eB8F/qmq3gC8ma5uTfZRkk3A7wI7quoqujO+bqatProX2DVr2Vz9cQOwtf/bC3xsSmVcqHu5uE4Hgauq6ufoTrG+HaDPiJuBN/av+cs+D5dsmnvgbwFOVdV/VdUPgQeA3VPc/pJV1Zmq+kI//T26YNhEV4/9/Wr7gRuXp4QLl2Qz8MvAXf18gOuAB/tVWqvPTwK/ANwNUFU/rKpzNNxHdGeLvby/9uIVwBka6qOq+gzwrVmL5+qP3cBfV+ezwLokG6dT0tENq1NV/UtVne9nP0t3cSN0dXqgqn5QVV8FTtHl4ZJNM8BH+u2UViS5ErgaOAJsqKoz/VPPABuWqViL8WfAB4Af9/OvBs4NfBBb66ctwDeAv+qHhe5K8koa7aOqehr4U+B/6IL7O8Ax2u4jmLs/VktO/Bbwj/30xOrkQcxFSPIq4JPA+6rqu4PPVXdaTxOn9iR5F3C2qo4td1nGaC1wDfCxqroa+D6zhksa66P1dHtwW4CfAl7JxV/dm9ZSf4wiyR10w633TXpb0wzwxfx2yoqT5BK68L6vqh7qFz878zWvfzy7XOVboJ8H3p3kv+mGtK6jGz9e139dh/b66TRwuqqO9PMP0gV6q310PfDVqvpGVf0IeIiu31ruI5i7P5rOiSS/CbwLeE89f472xOo0zQD/PLC1P3p+Kd2g/oEpbn/J+vHhu4ETVfXhgacOAHv66T3Aw9Mu22JU1e1VtbmqrqTrj3+tqvcAh4Gb+tWaqQ9AVT0DfD3J6/tFbwOeoNE+ohs62ZnkFf3nb6Y+zfZRb67+OAD8Rn82yk7gOwNDLStaujuUfQB4d1X938BTB4Cbk1yWZAvdAdrPjWWjVTW1P+CddEdnnwLumOa2x1T+t9J91XsUON7/vZNu3PgQcBL4NHD5cpd1EXX7ReCRfvpn+w/YKeBvgcuWu3wLrMt24GjfT38PrG+5j4APAl8BHgP+BrispT4C7qcbv/8R3TekW+bqDyB0Z6s9BXyZ7uybZa/DiHU6RTfWPZMNdw6sf0dfpyeBG8ZVDq/ElKRGeRBTkhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1Kj/B+WGLvmcG7MHAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "world.step()\n",
    "pylab.imshow(np.concatenate(session.run(fluid.density).data[...,0], axis=1), origin='lower', cmap='magma')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iJc6UdYHhtOH"
   },
   "source": [
    "Let's build a graph for the full simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "b9xHtdDQRrjL"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building graph for time step 0\n",
      "Building graph for time step 1\n",
      "Building graph for time step 2\n",
      "Building graph for time step 3\n",
      "Building graph for time step 4\n",
      "Building graph for time step 5\n",
      "Building graph for time step 6\n",
      "Building graph for time step 7\n",
      "Building graph for time step 8\n",
      "Building graph for time step 9\n",
      "Building graph for time step 10\n",
      "Building graph for time step 11\n",
      "Building graph for time step 12\n",
      "Building graph for time step 13\n",
      "Building graph for time step 14\n",
      "Building graph for time step 15\n",
      "Building graph for time step 16\n",
      "Building graph for time step 17\n",
      "Building graph for time step 18\n",
      "Building graph for time step 19\n"
     ]
    }
   ],
   "source": [
    "for step in range(20):\n",
    "  print('Building graph for time step %d' % step)\n",
    "  world.step(dt=1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2VQ92g2rs6wM"
   },
   "source": [
    "When calling `session.run` now, the full simulation is evaluated using TensorFlow operations.\n",
    "This will take advantage of your GPU, if available.\n",
    "If you compile Φ<sub>Flow</sub> with [CUDA support](https://github.com/tum-pbs/PhiFlow/blob/master/documentation/Installation_Instructions.md), the TensorFlow graph will also use optimized operators for efficient simulation and training runs.\n",
    "\n",
    "The `session.run()` call of the following code block will now retrieve the final fluid density, and for that it actually needs to process all 20 simulation steps of the TF graph that we just constructed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "TA6Ibs-mXsTc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computing frames...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x147e73190>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 11864 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "print('Computing frames...')\n",
    "original_state = np.concatenate(session.run(fluid.density).data[...,0], axis=1)\n",
    "pylab.imshow(original_state, origin='lower', cmap='magma')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IClfRMfoyGUa"
   },
   "source": [
    "Next, we define the *loss* function. This is the value we want to decrease via optimization.\n",
    "For this example, we want the marker densities of all final simulation states to match the left-most one, called `target`, in terms of an $L^2$ norm.\n",
    "\n",
    "For the optimizer, we choose again gradient descent for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "7KPpyIwjYETi"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.7/site-packages/tensorflow_core/python/ops/math_grad.py:1375: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "Initial loss: 65.417114\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x14b36c750>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 11864 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "target = session.run(fluid.density).data[0,...]\n",
    "loss = math.l2_loss(fluid.density.data[1:,...] - target)\n",
    "optim = tf.train.GradientDescentOptimizer(learning_rate=0.01).minimize(loss)\n",
    "session.initialize_variables()\n",
    "\n",
    "print('Initial loss: %f' % session.run(loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AALD66-N0U5F"
   },
   "source": [
    "With the loss and optimizer set up, all that's left is to run the actual optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "pvvF6xqmaRLX"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization step 0, loss 65.417122 . The first step sets up the adjoint graph.\n",
      "Optimization step 1, loss 61.038647 \n",
      "Optimization step 2, loss 55.527832 \n",
      "Optimization step 10, loss 43.914444 \n",
      "Optimization step 20, loss 40.739956 \n",
      "Optimization step 30, loss 39.523136 \n",
      "Optimization step 40, loss 39.361752 \n",
      "Optimization step 50, loss 39.348873 \n",
      "Optimization step 60, loss 37.156372 \n",
      "Optimization step 70, loss 37.729256 \n",
      "Optimization step 80, loss 35.622345 \n",
      "Optimization step 90, loss 35.464912 \n"
     ]
    }
   ],
   "source": [
    "for optim_step in range(100):\n",
    "    _, loss_value = session.run([optim, loss])\n",
    "    if optim_step<3 or optim_step%10==0: \n",
    "        print('Optimization step %d, loss %f %s' % (optim_step, loss_value, '' if optim_step else '. The first step sets up the adjoint graph.'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss should have gone down significantly, from above 60 to below 40, and now we can visualize how well the reconstruction for the last frame turned out for simulations 1 to 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "SQBtCmhZaYYj"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x14b4ad890>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 26504 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "optimized_state = np.concatenate(session.run(fluid.density).data[...,0], axis=1)\n",
    "pylab.title(\"Original state at t=20 (top), and optimized version (bottom):\")\n",
    "pylab.imshow( np.concatenate([optimized_state, original_state],axis=0), origin='lower', cmap='magma')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HP7aDQfpKifp"
   },
   "source": [
    "Naturally, the image on the left is the same (this is simulation (0), our reference), and the other three simulations now exhibit a noticeable curved shape towards the left. This is especially visible for the far right simulation (3), which was moving straight up in its original state, and now has an off-center plume shape that tries to match the reference on the left. \n",
    "\n",
    "Note that this simulation (like both others) needs to \"work\" with a fixed inflow, hence it cannot simply \"produce\" marker density out of the blue to match the target. Also it needs to take into account how the non-linear model equations change the state of the system over the course of 20 time steps. So the optimization goal is quite tough, and it is not possible to exactly satisfy the constraints to match simulation (0) in this scenario. This is exactly what makes it an interesting test case, though.\n",
    "\n",
    "Now we can also have a look at the now-optimized initial velocity field. Here we see the y-components (from `data[...,0]`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "i7ZahlUudex8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x14b5363d0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 13804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "optimized_velocity_field = session.run(initial_state.velocity).at_centers()\n",
    "\n",
    "pylab.title('Initial y-velocity (optimized)')\n",
    "pylab.imshow(np.concatenate(optimized_velocity_field.data[...,0], axis=1), origin='lower')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This nicely shows how the initial perturbation of the velocity field grows for the larger inflow deviations towards the right:\n",
    "\n",
    "For comlpeteness, here are the x-components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "Pqw5BDxmdkut"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x14b3efa50>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 13436 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "pylab.title('Initial x-velocity (optimized)')\n",
    "pylab.imshow(np.concatenate(optimized_velocity_field.data[...,1], axis=1), origin='lower')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ooqVxCPM8PXl"
   },
   "source": [
    "This example illustrates how the differentiable physics approach can easily be extended towards significanlty more complex PDEs. Above, we've optimized for a mini-batch of 20 steps of a full Navier-Stokes solver.\n",
    "\n",
    "This is a powerful basis to bring NNs into the picture! As you might have noticed, our degrees of freedom were still a regular grid, and we've solved a single inverse problem. There was more than one case to solve, but nonetheless no question of generalization and no neural network structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Differentiable Physics with Fluid Simulations.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
}
