{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o4JZ84moBKMr"
   },
   "source": [
    "# Differentiable Fluid Simulations\n",
    "\n",
    "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n",
    "\n",
    "**test version, rough... not newest one**\n",
    "\n",
    "\n",
    "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n",
    "\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": 1,
   "metadata": {
    "id": "da1uZcDXdVcF"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v2/phi/physics/_boundaries.py:396: SyntaxWarning: \"is\" with a literal. Did you mean \"==\"?\n",
      "  return self.velocity is 0 and self.angular_velocity is 0\n",
      "/Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v2/phi/physics/_boundaries.py:396: SyntaxWarning: \"is\" with a literal. Did you mean \"==\"?\n",
      "  return self.velocity is 0 and self.angular_velocity is 0\n",
      "/Users/thuerey/miniconda3/envs/tf/lib/python3.8/_collections_abc.py:743: MatplotlibDeprecationWarning: The global colormaps dictionary is no longer considered public API.\n",
      "  for key in self._mapping:\n",
      "/Users/thuerey/miniconda3/envs/tf/lib/python3.8/_collections_abc.py:744: MatplotlibDeprecationWarning: The global colormaps dictionary is no longer considered public API.\n",
      "  yield (key, self._mapping[key])\n"
     ]
    }
   ],
   "source": [
    "#!pip install --upgrade --quiet phiflow\n",
    "#!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@develop\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. \n",
    "\n",
    "While the simulation setup is very similar to before, we'll define an additional batch dimension in addition to the spatial dimensions `x` and `y`. Below, this additional dimension is triggered by inflow tensor, which contains an named `nbatch` dimension, in addition to the spatial ones (the name `batch` is alredy taken). The two `print` statements below illustrate this behavior: they print the tensor shapes with all dimensions, and next only the spatial ones.\n",
    "\n",
    "TODO xxx **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",
    "TODO ,explain expansion!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "WrA3IXDxv31P"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inflow shape: (nbatch=4, x=32, y=40)\n",
      "Inflow shape, spatial only: (x=32, y=40)\n",
      "Initial velocity shape: (x=32, y=40, vector=2) \n",
      "\n",
      "Velocity shape after 1 step: (nbatch=4, x=32, y=40, vector=2)\n"
     ]
    }
   ],
   "source": [
    "domain = Domain(x=32, y=40, boundaries=CLOSED, bounds=Box[0:32, 0:40])\n",
    "\n",
    "smoke    = domain.scalar_grid(0)  # sampled at cell centers\n",
    "velocity = domain.staggered_grid(0)  # sampled in staggered form at face centers \n",
    "\n",
    "sphere_centers = [(10,5), (12,5), (19,5), (16,5)]\n",
    "inflow_location = math.tensor(sphere_centers, names='nbatch,vector')\n",
    "#inflow_location = math.tensor([(4, 5), (8, 5), (12, 5), (16, 5)], names='inflow_loc,vector')\n",
    "inflow = domain.grid(Sphere(center=inflow_location, radius=3)) * 0.6\n",
    "\n",
    "\n",
    "#print(f\"DEBUG : {velocity.shape.nbatch}\")\n",
    "\n",
    "print(f\"Inflow shape: {inflow.shape}\")\n",
    "print(f\"Inflow shape, spatial only: {inflow.shape.spatial}\")\n",
    "print(f\"Initial velocity shape: {velocity.shape} \\n\")\n",
    "\n",
    "# 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",
    "  #world.step(dt=1.5)\n",
    "  #for _ in range(20):\n",
    "\n",
    "for frame in range(20):\n",
    "  smoke = advect.mac_cormack(smoke, velocity, dt=1) + inflow\n",
    "  buoyancy_force = smoke * (0, 0.5) >> velocity\n",
    "  velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force\n",
    "  velocity, _, _, _ = fluid.make_incompressible(velocity, domain)\n",
    "  if frame==0:\n",
    "    print(f\"Velocity shape after 1 step: {velocity.shape}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO** explain nbatch\n",
    "phiflow automatically expands dimensions...\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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15336 characters]\n",
      "text/plain": [
       "<Figure size 1440x288 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#pylab.imshow(np.concatenate(fluid.density.data[...,0], axis=1), origin='lower', cmap='magma')\n",
    "fig, axes = pylab.subplots(1, 4, figsize=(20, 4))\n",
    "for i in range(inflow.shape.nbatch):\n",
    "  axes[i].imshow(smoke.values.nbatch[i].numpy('y,x'), origin='lower', cmap='magma')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**TODO** add labels...\n",
    "\n",
    "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",
    "This can be achieved by changing the import statement to [`from phi.tf.flow import *`](https://tum-pbs.github.io/PhiFlow/phi/tf/flow.html) (or [`phi.torch.flow`](https://tum-pbs.github.io/PhiFlow/phi/torch/flow.html) for PyTorch).\n",
    "Tensors created after this import will be allocated using TensorFlow and operations on these will be executed with TensorFlow.\n",
    "Note that this will make use of a GPU through CUDA if your TensorFlow version supports it, and potentially custom CUDA operators of phiflow (if you dont have those compiled you might see a warning below, which you can safely ignore)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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-v2/phi/tf/cuda/build/resample.so. Run \"python setup.py tf_cuda\" to compile them\n"
     ]
    }
   ],
   "source": [
    "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.flow`, let's set up the simulation just like before. But now, the initial velocity state that we start from will be a variable quantity that we can modify. 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=20$. Thus, after 20 steps 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 and interesting here -- 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'll run the full chain of steps for the simulation over time in TensorFlow, and track the gradients for the initial velocity state at $t=0$ to be modified in the optimization.\n",
    "\n",
    "**REMOVE 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": 5,
   "metadata": {
    "id": "NlJMJikaHOL6"
   },
   "outputs": [],
   "source": [
    "domain = Domain(x=32, y=40, boundaries=CLOSED, bounds=Box[0:32, 0:40])\n",
    "sphere_centers = [(10.,5.), (12.,5.), (19.,5.), (16.,5.)]\n",
    "inflow_location = math.tensor(sphere_centers, names='inflow_loc,vector', convert=True)\n",
    "inflow = domain.grid(Sphere(center=inflow_location, radius=3)) * 0.6\n",
    "\n",
    "# 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": {},
   "source": [
    "We can verify that tensors are now backed by TensorFlow with the following `print` statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensorflow.python.framework.ops.EagerTensor"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(inflow.values.native())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vSdGIEzCgq5-"
   },
   "source": [
    "Note that tensors created with NumPy will keep on using NumPy/SciPy operations unless a TensorFlow tensor is also passed to the same operation.\n",
    "\n",
    "**CONTINUE**\n",
    "\n",
    "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": 7,
   "metadata": {
    "id": "wSrIezfWHjcQ"
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'world' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-70b40c8f4c48>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mworld\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\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      2\u001b[0m \u001b[0mpylab\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimshow\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconcatenate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msession\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfluid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdensity\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morigin\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'lower'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcmap\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'magma'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'world' is not defined"
     ]
    }
   ],
   "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": null,
   "metadata": {
    "id": "b9xHtdDQRrjL"
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "id": "TA6Ibs-mXsTc"
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "id": "7KPpyIwjYETi"
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "id": "pvvF6xqmaRLX"
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "id": "SQBtCmhZaYYj"
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "id": "i7ZahlUudex8"
   },
   "outputs": [],
   "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": null,
   "metadata": {
    "id": "Pqw5BDxmdkut"
   },
   "outputs": [],
   "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
}
