{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o4JZ84moBKMr"
   },
   "source": [
    "# Differentiable Fluid Simulations\n",
    "\n",
    "**(Note - this example is not final. TODO: update from diffphys-code-ns-v1-syncWith2.ipynb, diff display for target, adjust params so that the loss properly goes down)**\n",
    "\n",
    "Next, we'll target a more complex example with the Navier-Stokes equations as physical model. In line with {doc}`overview-ns-forw`, we'll target a 2D case.\n",
    "\n",
    "As optimization objective we'll consider a more difficult variant of the previous Burgers 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}_0$ 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",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/diffphys-code-ns.ipynb)\n",
    "\n",
    "## Physical Model\n",
    "\n",
    "We'll use a Navier-Stokes model 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. For the velocity this gives:\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",
    "\\end{aligned}$\n",
    "\n",
    "With an additional transport equation for the passively advected marker density $s$:\n",
    "\n",
    "$\\begin{aligned}\n",
    "  \\frac{\\partial s}{\\partial{t}} + \\mathbf{u} \\cdot \\nabla s &= 0 \n",
    "\\end{aligned}$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formulation\n",
    "\n",
    "\n",
    "With the notation from {doc}`overview-equations` the inverse problem outlined above can be formulated as a  minimization problem \n",
    "\n",
    "$$\n",
    "\\text{arg min}_{\\mathbf{u}_{0}} \\sum_i |f(x_{t_e,i} ; \\theta)-y^*_{t_e,i}|^2 ,\n",
    "$$\n",
    "\n",
    "where $y^*_{t_e,i}$ are samples of the reference solution at a targeted time $t_e$,\n",
    "and $x_{t_e,i}$ denotes the estimate of our simulator at the same\n",
    "sampling locations and time. \n",
    "The index $i$ here runs over all discretized, spatial degrees of freedom in our fluid solver (we'll have $32\\times40$ below).\n",
    "\n",
    "In contrast to before, we are not dealing with pre-computed quantities anymore, but now $x_{t_e,i}$ is a complex, non-linear function itself.\n",
    "More specifically, the simulator starts with the initial velocity $\\mathbf{u}_0$ and density $s_0$ to compute the $x_{t_e,i}$, by $n$ evaluations of the discretized PDE $\\mathcal{P}$.\n",
    "This gives as simulated final state \n",
    "$y_{t_e,i} = s_{t_e} = \\mathcal{P}^n(\\mathbf{u}_0,s_0)$, where we will leave $s_0$ fixed in the following, and focus on $\\mathbf{u}_0$ as our degrees of freedom.\n",
    "Hence, the optimization can only change $\\mathbf{u}_0$ to align $y_{t_e,i}$ with the references $y^*_{t_e,i}$ as closely as possible.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Starting the Implementation\n",
    "\n",
    "First, let's get the loading of python modules out of the way. By importing `phi.torch.flow`, we get fluid simulation functions that work within pytorch graphs and can provide gradients (`phi.tf.flow` would be the alternative for tensorflow).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "da1uZcDXdVcF"
   },
   "outputs": [],
   "source": [
    "#!pip install --upgrade --quiet phiflow\n",
    "#!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@develop\n",
    "\n",
    "from phi.torch.flow import *  \n",
    "import pylab # for visualizations later on\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rdSTbMoaS0Uz"
   },
   "source": [
    "## Batched simulations\n",
    "\n",
    "Now we can set up the simulation, which will work in line with the previous \"regular\" simulation example from the {doc}`overview-ns-forw`. However, now we'll directly include an additional dimension, in line with a mini-batch used for NN training. For this, we'll include an additional named dimension called `inflow_loc`. This dimension will exist \"above\" the previous spatial dimensions `y`, `x` and the channel dimensions `vector`. As indicated by the name `inflow_loc`, the main differences for this dimension will lie in different locations of the inflow, in order to obtain different flow simulations. The named dimensions in phiflow make it very convenient to broadcast information across matching dimensions in different tensors.\n",
    "\n",
    "The `Domain` object is allocated just like before, but the `INFLOW_LOCATION` tensor now receives a string\n",
    "`'inflow_loc,vector'` that indicates the names of the two dimensions. This leads to the creation of an `inflow_loc` dimensions in addition to the two spatial dimensions (the `vector` part).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "mphMP0sYIOz-"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inflow tensor shape and dimensions: (inflow_loc=4, x=32, y=40)\n"
     ]
    }
   ],
   "source": [
    "DOMAIN = Domain(x=32, y=40, boundaries=CLOSED, bounds=Box[0:32, 0:40])\n",
    "\n",
    "INFLOW_LOCATION = math.tensor([(10, 5), (12, 5), (14, 5), (16, 5)], 'inflow_loc,vector', convert=True)\n",
    "INFLOW = DOMAIN.grid(Sphere(center=INFLOW_LOCATION, radius=3)) * (1./3.)\n",
    "print(f\"Inflow tensor shape and dimensions: {INFLOW.shape}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last `print` statement verifies that our `INFLOW` grid likewise has an `inflow_loc` dimension in addition to the spatial dimenions `x` and `y`. You can test for the existence of a tensor dimension in phiflow with the `.exists` boolean, which can be evaluated for any dimension name. E.g., above `INFLOW.inflow_loc.exists` will give `True`, while `INFLOW.some_unknown_dim.exists` will give `False`.\n",
    "\n",
    "Phiflow tensors are automatically broadcast to new dimensions via their names, and hence typically no-reshaping operations are required. E.g., you can easily add or multiply tensors with differing dimensions. Below we'll multiply a staggered grid with a tensor of ones along the `inflow_loc` dimension to get a staggered velocity that has `x,y,inflow_loc` as dimensions via `DOMAIN.staggered_grid(0) * math.ones(inflow_loc=4)`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can easily simulate a few steps now starting with these different initial conditions. Thanks to the broadcasting, the exact same code we used for the single forward simulation in the overview chapter will produce four simulations with different smoke inflow positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21136 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "smoke = DOMAIN.scalar_grid(0)        # sampled at cell centers\n",
    "velocity = DOMAIN.staggered_grid(0)  # sampled in staggered form at face centers \n",
    "\n",
    "def step(smoke, velocity):\n",
    "  smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "  buoyancy_force = smoke * (0, 1) >> velocity\n",
    "  velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force\n",
    "  velocity, _ = fluid.make_incompressible(velocity, DOMAIN)\n",
    "  return smoke, velocity\n",
    "\n",
    "for _ in range(20):\n",
    "    smoke,velocity = step(smoke,velocity)\n",
    "\n",
    "smoke_final = smoke # final state, without optimization\n",
    "fig, axes = pylab.subplots(1, 4, figsize=(16, 5))\n",
    "for i in range(INFLOW.shape.inflow_loc):\n",
    "  axes[i].imshow(smoke_final.values.inflow_loc[i].numpy('y,x'), origin='lower', cmap='magma')\n",
    "  axes[i].set_title(f\"Inflow {INFLOW_LOCATION.inflow_loc[i].vector.unstack_spatial('x,y')}\" + (\", Reference\" if i==3 else \"\"))\n",
    "pylab.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last image shows the state of the advected smoke fields after 20 simulation steps. The final smoke shape of simulation `[3]` with an inflow at `(16,5)`, with the straight plume on the far right, will be our **reference state** below. The initial velocity of the other three will be modified in the optimization procedure below to match this reference.\n",
    "\n",
    "(As a small side note: phiflow tensors will keep track of their chain of operations using the backend they were created for. E.g. a tensor created with NumPy will keep using NumPy/SciPy operations unless a PyTorch or TensorFlow tensor is also passed to the same operation. Thus, it is a good idea to verify that tensors have the right `type` once in a while.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradients\n",
    "\n",
    "Let's look at how to get gradients from our simulation. The first trivial step taken care of above was to include `phi.torch.flow` to import differentiable operators from which to build our simulator.\n",
    "\n",
    "Now we want to optimize the initial velocities so that all simulations arrive at a final state that is similar to the simulation no the right, where the inflow is located at `(16, 5)`, i.e. centered along `x`.\n",
    "To achieve this, we record the gradients during the simulation and define a simple $L^2$ based loss function. The loss function we'll use is given by $L = | s_{t_e} - s_{t_e}^* |^2$, where $s_{t_e}$ denotes the smoke density, and $s_{t_e}^*$\n",
    "denotes the reference state from the fourth simulation in our batch (both evaluated at the last time step $t_e$).\n",
    "When evaluating the loss function we treat the reference state as an external constant via `field.stop_gradient()`.\n",
    "As outlined at the top, $s$ is a function of $\\mathbf{u}$ (via the advection equation), which in turn is given by the Navier-Stokes equations. Thus, via a chain of multiple time steps $s$ depends in the initial velocity state $\\mathbf{u}_0$.\n",
    "\n",
    "\n",
    "It is important that our initial velocity has the `inflow_loc` dimension before we record the gradients, such that we have the full \"mini-batch\" of four versions of our velocity (three of which will be updated via gradients in our optimization later on). To get the approprate velocity tensor, we multiply a `staggered_grid` tensor from our domain object, with a tensor full of ones along the `inflow_loc` dimension.\n",
    "\n",
    "Phiflow provides a unified API for gradients across different platforms by using functions that need to return a loss values, in addition to optional state values. Below the `simulate` function computes the $L^2$ error outlined above and returns the evolved `smoke` and `velocity` states after 20 simulation steps.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Velocity tensor shape: (inflow_loc=4, x=32, y=40, vector=2)\n"
     ]
    }
   ],
   "source": [
    "initial_smoke = DOMAIN.grid(math.zeros(inflow_loc=4))\n",
    "initial_velocity = DOMAIN.staggered_grid(0) * math.ones(inflow_loc=4)\n",
    "\n",
    "print(\"Velocity tensor shape: \"+format(initial_velocity.shape))\n",
    "\n",
    "def simulate(smoke: CenteredGrid, velocity: StaggeredGrid):\n",
    "    for _ in range(20):\n",
    "        smoke,velocity = step(smoke,velocity)\n",
    "        \n",
    "    loss = field.l2_loss(smoke - field.stop_gradient(smoke.inflow_loc[-1]) )\n",
    "    # optionally, use smoother loss with diffusion steps - no difference here, but can be useful for more complex cases\n",
    "    #loss = field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10))\n",
    "    \n",
    "    return loss, smoke, velocity\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Phiflow's `field.functional_gradient()` function is the central function to compute gradients. Next, we'll use it to obtain the gradient with respect to the initial velocity. Since the velocity is the second argument of the `simulate()` function, we pass `wrt=[1]`. (Note that phiflow also has a `field.spatal_gradient` function which instead computes derivatives of tensors along spatial dimensions like `x,y`.)\n",
    "\n",
    "`functional_gradient` in generates a gradient function, which we evaluate one term for test purposes with the initial states for smoke and velocity. The last statement prints a summary of one part of the tensor:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "grad can be implicitly created only for scalar outputs",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-e357726b3491>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;31m#[velocity_grad] = sim_grad(initial_smoke, initial_velocity)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;31m#(velocity_grad,) = sim_grad(initial_smoke, initial_velocity)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m \u001b[0msim_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minitial_smoke\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minitial_velocity\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      9\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Gradient info \"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvelocity_grad\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/phiflow-v2/phi/math/_functional.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    300\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mkey\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrads\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    301\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrads\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_trace_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwrt_natives\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 302\u001b[0;31m         \u001b[0mnative_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrads\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mnatives\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    303\u001b[0m         \u001b[0moutput_key\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmatch_output_signature\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecorded_mappings\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    304\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_output\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-v2/phi/torch/_torch_backend.py\u001b[0m in \u001b[0;36meval_grad\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    635\u001b[0m             \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\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[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    636\u001b[0m             \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maux\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput\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[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mtuple\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\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[0;32m--> 637\u001b[0;31m             \u001b[0mgrads\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwrt_args\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# TODO this does not work during jit for some reason\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    638\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mget_output\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    639\u001b[0m                 \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\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~/miniconda3/envs/tf/lib/python3.8/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mgrad\u001b[0;34m(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)\u001b[0m\n\u001b[1;32m    195\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    196\u001b[0m     \u001b[0mgrad_outputs_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_tensor_or_tensors_to_tuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrad_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\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[0m\n\u001b[0;32m--> 197\u001b[0;31m     \u001b[0mgrad_outputs_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_make_grads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_outputs_\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    198\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    199\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mretain_graph\u001b[0m \u001b[0;32mis\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[0;32m~/miniconda3/envs/tf/lib/python3.8/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36m_make_grads\u001b[0;34m(outputs, grads)\u001b[0m\n\u001b[1;32m     48\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_grad\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     49\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\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---> 50\u001b[0;31m                     \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"grad can be implicitly created only for scalar outputs\"\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     51\u001b[0m                 \u001b[0mnew_grads\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmemory_format\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpreserve_format\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     52\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;31mRuntimeError\u001b[0m: grad can be implicitly created only for scalar outputs"
     ]
    }
   ],
   "source": [
    "sim_grad = field.functional_gradient(simulate, wrt=[1], get_output=False)\n",
    "\n",
    "TODO fix\n",
    "#[velocity_grad] = sim_grad(initial_smoke, initial_velocity)\n",
    "#(velocity_grad,) = sim_grad(initial_smoke, initial_velocity)\n",
    "sim_grad(initial_smoke, initial_velocity)\n",
    "\n",
    "print(\"Gradient info \" + format(velocity_grad))\n",
    "velocity_grad.inflow_loc[0].values.vector[0] # one example, location 0, x component, automatically prints size & content range\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last two lines just print some information about the resulting gradient field. Naturally, it has the same shape as the velocity itself: it's a staggered grid with four inflow locations. The last line shows how to access the x-components of one of the gradients.\n",
    "\n",
    "We can directly use this to take a look at the content of the computed gradient. The following images show the `x` components (via `vector[0].values.numpy()`) of the staggered grid gradient in `velocity_grad`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = pylab.subplots(1, 4, figsize=(16, 5))\n",
    "for i in range(INFLOW.shape.inflow_loc):\n",
    "  im = axes[i].imshow(velocity_grad.inflow_loc[i].values.vector[0].numpy('y,x'), origin='lower', cmap='bwr')\n",
    "  pylab.colorbar(im, ax=axes[i])\n",
    "  axes[i].set_title(f\"Gradient for x vel - In {math.to_int(INFLOW_LOCATION).inflow_loc[i].vector.unstack_spatial('x,y')}\")\n",
    "pylab.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not surprisingly, the fourth gradient on the left is zero (it's already matching the reference). The other three gradients have detected variations for the intial round inflow posiitons shown as positive and negative regions around the circular shape.\n",
    "\n",
    "**TODO, double check after code update**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3mpyowRYUSS4"
   },
   "source": [
    "\n",
    "\n",
    "## Optimization\n",
    "\n",
    "The gradient visualized above is just the linearized change that points in the direction of an increasing loss. Now we can proceed by updating the initial velocities in the opposite direction to minimize the loss, and iterate to find a minimizer.\n",
    "\n",
    "This is a difficult task: the simulation is producing different dynamics due to the differing initial spatial density configuration, and an optimizer should now find a single initial velocity state that gives the same state as the reference simulation at $t=20$. Thus, after non-linear 20 update steps the simulation should reproduce the desired marker density state. It would be much easier to simply change the position of the marker inflow to arrive at this goal, but -- to make things more difficult and interesting here -- the inflow is _not_ a degree of freedom. The optimizer can only change the initial velocity $\\mathbf{u}_0$.\n",
    "\n",
    "The following cell re-evalutes the gradient function, and iterates several times to optimize $\\mathbf{u}_0$ with a learning rate (step size) `LR`.\n",
    "\n",
    "`field.functional_gradient` has a parameter `get_output` that determines whether the original results of the function (`simulate` here) are returned, or only the gradient. As it's interesting to track how the loss evolves over the course of the iterations, let's redefine the gradient function with `get_output=True`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_velocity = DOMAIN.staggered_grid(0) * math.ones(inflow_loc=4) # reset\n",
    "LR = 1e-02\n",
    "\n",
    "sim_grad_wloss = field.functional_gradient(simulate, wrt=[1], get_output=True)\n",
    "\n",
    "for optim_step in range(20):\n",
    "    initial_smoke = DOMAIN.grid(math.zeros(inflow_loc=4))\n",
    "    \n",
    "    # sim_grad__wloss returns: [loss,s,v,velocity _grad] \n",
    "    sim_grad = field.functional_gradient(simulate, wrt=[1], get_output=True)\n",
    "    [loss,s,v, velocity_grad] = sim_grad_wloss(initial_smoke, initial_velocity)\n",
    "\n",
    "    print('Optimization step %d, loss: %f' % (optim_step,loss))\n",
    "\n",
    "    initial_velocity = initial_velocity - LR * velocity_grad\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss should have gone down significantly, from above 60 to below 40, and now we can also visualize the initial velocities that were obtained in the optimization. **TODO, double check after code update**\n",
    "\n",
    "The following images show the resulting three initial velocities (we're skipping the fourth set with `inflow_loc[0]` because it only contains zeros).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = pylab.subplots(1, 3, figsize=(12, 5))\n",
    "for i in range(INFLOW.shape.inflow_loc-1):\n",
    "  im = axes[i].imshow(initial_velocity.inflow_loc[i].values.vector[0].numpy('y,x'), origin='lower', cmap='bwr')\n",
    "  pylab.colorbar(im, ax=axes[i])\n",
    "  axes[i].set_title(f\"Initial velocity x for {math.to_int(INFLOW_LOCATION).inflow_loc[i].vector.unstack_spatial('x,y')}\")\n",
    "pylab.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = pylab.subplots(1, 3, figsize=(12, 5))\n",
    "for i in range(INFLOW.shape.inflow_loc-1):\n",
    "  im = axes[i].imshow(initial_velocity.inflow_loc[i].values.vector[1].numpy('y,x'), origin='lower', cmap='bwr')\n",
    "  pylab.colorbar(im, ax=axes[i])\n",
    "  axes[i].set_title(f\"Initial velocity y for {math.to_int(INFLOW_LOCATION).inflow_loc[i].vector.unstack_spatial('x,y')}\")\n",
    "pylab.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Re-simulation\n",
    "\n",
    "We can also visualize how the full simulation over the course of 20 steps turns out, given the new initial velocity conditions for each of the inflow locations. This is what happened internally at optimization time, and what was measured by our loss function. Hence, it's good to get an intuition for which solutions the optimization has found.\n",
    "\n",
    "Below, we re-run the forward simulation with the new initial conditions from `initial_velocity`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smoke = initial_smoke \n",
    "velocity = initial_velocity\n",
    "\n",
    "for _ in range(20):\n",
    "    smoke,velocity = step(smoke,velocity)\n",
    "\n",
    "fig, axes = pylab.subplots(1, 4, figsize=(16, 5))\n",
    "for i in range(INFLOW.shape.inflow_loc):\n",
    "  axes[i].imshow(smoke.values.inflow_loc[i].numpy('y,x'), origin='lower', cmap='magma')\n",
    "  axes[i].set_title(f\"Solution for {math.to_int(INFLOW_LOCATION).inflow_loc[i].vector.unstack_spatial('x,y')}\" + (\", Reference\" if i==3 else \"\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Naturally, the image on the right is the same (this is the reference), and the other three simulations now exhibit a noticeable shift towards the right. This is especially visible for the far left simulation (0), which was strongly curved in its original state, and now has a shape that tries to match the straighter reference on the right.  **TODO, double check after code update**\n",
    "\n",
    "Note that all three simulations need to \"work\" with a fixed inflow, hence they cannot simply \"produce\" marker density out of the blue to match the target. Also each simulation 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 difficult, and it is not possible to exactly satisfy the constraints to match the reference simulation in this scenario. This is exactly what makes it an interesting test case, though.\n",
    "\n",
    "The next set of images show the difference before and after optimization by subtracting the final state of the unmodified simulation (`smoke_final` from above) from the optimized solutions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = pylab.subplots(1, 3, figsize=(16, 5))\n",
    "for i in range(3):\n",
    "  axes[i].imshow( smoke.values.inflow_loc[i].numpy('y,x') - smoke_final.values.inflow_loc[i].numpy('y,x') , origin='lower', cmap='magma')\n",
    "  axes[i].set_title(f\"Diff {math.to_int(INFLOW_LOCATION).inflow_loc[i].vector.unstack_spatial('x,y')}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conclude... **TODO, double check after code update**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Conclusions\n",
    "\n",
    "This example illustrated how the differentiable physics approach can easily be extended towards significanlty more \n",
    "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 jointly solved a single inverse problem. There were three cases to solve, of course, but nonetheless the setup remains a direct optimization. Thus, in line with the PINN example of {doc}`physicalloss-code` we've not really dealt with a _machine learning_ task here. However, that will be the topic of the next chapters.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Based on the code example above, we can recommend experimenting with the following:\n",
    "\n",
    "- Modify the setup of the simulation to differ more strongly across the four instances, run longer, or use a finer spatial discretization (i.e. larger grid size). Note that this will make the optimization problem tougher, and hence might not converge directly with this simple setup.\n",
    "\n",
    "- For more complex cases, add a multi-resolution optimization, that first solves with a coarse discretization, and then uses this solution as an initial guess for a finer discretization. (Ideally, implement it with a loop that continuously refines the discretization.)\n"
   ]
  }
 ],
 "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
}
