{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Burgers Optimization with a (Manual) Differentiable Physics Gradient\n",
    "\n",
    "To illustrate the process of compute gradients in a differentiable physics setting, we first demonstrate a process that isn't recommended for more complex real-world cases, but more clearly shows the way we can get gradient information from a physical simulation in a DL framework. We'll target tensorflow in the following.\n",
    "\n",
    "We target the same reconstruction task like for the PINN example. This has some immediate implications: the evolution of the  system is now fully determined by our PDE formulatio. Hence, the only real unknown is the initial state! We will still need to re-compute all the states betwwen the initial and target state many times, just now we won't need an NN for this step, we can rely on our discretized model. Also, as we choose an initial discretization for the DP approach, the unknown initial state consists of the sampling points of the involved physical fields, and we can simply represent these unknowns as floating point variables. Hence, even for the initial state we do not need to set up an NN. Thus, our Burgers reconstruction problem reduces to a gradient-based opitmization without any NN when solving it with DP. Nonetheless, it's a very good starting point to illustrate the process.\n",
    "\n",
    "First, we'll set up our discretized simulation. Here we can employ the phiflow solver, as shown in the overview section on _Burgers forward simulations_. phiflow directly gives us a computational graph in tensorflow. So, as a first step, let's set up our grid with 128 points, and construct a tensorflow graph for 32 steps of the simulation. (As usual, you can ignore the warnings when loading TF and phiflow.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Each velocity is a phiflow grid like this one: Grid[128(1), size=[2.], <class 'numpy.float32'>]\n"
     ]
    }
   ],
   "source": [
    "from phi.tf.flow import *\n",
    "\n",
    "# run with phiflow\n",
    "n = 128\n",
    "dt = 1./32.\n",
    "initial = np.zeros([n,1]) # start from 0\n",
    "\n",
    "domain = Domain([n], boundaries=PERIODIC, box=box[-1:1])\n",
    "state0 = BurgersVelocity(domain, velocity=initial, viscosity=0.01/np.pi)\n",
    "\n",
    "# start with a state0, to be modified\n",
    "state_in = state0.copied_with(velocity=placeholder)\n",
    "states = [state_in]\n",
    "\n",
    "for i in range(32):\n",
    "    states.append( Burgers().step(states[-1],dt=dt) )\n",
    "\n",
    "print(\"Note: each velocity is a phiflow grid like this one \" + format(states[-1].velocity) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**A lot** has happened in these first steps, but we didnt run a single calculation for our simulation. In contrast to the first example, which used numpy as phiflow backend, we now have used the tensorflow backend (the trick here was importing `phi.tf.flow` instad of `phi.flow`). This gave us a tensorflow graph, but didn't run any real simulation steps. That will only happen once we run a tensorflow session and pass some data through this graph. So, in a way _nothing_ has happened so far in terms of our simulation! \n",
    "\n",
    "Next, let's set up the loss, such that we can start the optimization. That's actually pretty simple, we want the solution at $t=0.5$, i.e. step number 16, to respect the data in terms of an L2 norm. Let's also directly initialize the TF variables, and see what loss we get from the initial zero'ed initialization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
      "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",
      "Initial loss 0.3829152584075928\n"
     ]
    }
   ],
   "source": [
    "sess = Session(None) \n",
    "\n",
    "# enforce constraints at time t=0.5 (step 16)\n",
    "solution_t16 = np.asarray( [0.008612174447657694, 0.02584669669548606, 0.043136357266407785, 0.060491074685516746, 0.07793926183951633, 0.0954779141740818, 0.11311894389663882, 0.1308497114054023, 0.14867023658641343, 0.1665634396808965, 0.18452263429574314, 0.20253084411376132, 0.22057828799835133, 0.23865132431365316, 0.25673879161339097, 0.27483167307082423, 0.2929182325574904, 0.3109944766354339, 0.3290477753208284, 0.34707880794585116, 0.36507311960102307, 0.38303584302507954, 0.40094962955534186, 0.4188235294008765, 0.4366357052408043, 0.45439856841363885, 0.4720845505219581, 0.4897081943759776, 0.5072391070000235, 0.5247011051514834, 0.542067187709797, 0.5593576751669057, 0.5765465453632126, 0.5936507311857876, 0.6106452944663003, 0.6275435911624945, 0.6443221318186165, 0.6609900633731869, 0.67752574922899, 0.6939334022562877, 0.7101938106059631, 0.7263049537163667, 0.7422506131457406, 0.7580207366534812, 0.7736033721649875, 0.7889776974379873, 0.8041371279965555, 0.8190465276590387, 0.8337064887158392, 0.8480617965162781, 0.8621229412131242, 0.8758057344502199, 0.8891341984763013, 0.9019806505391214, 0.9143881632159129, 0.9261597966464793, 0.9373647624856912, 0.9476871303793314, 0.9572273019669029, 0.9654367940878237, 0.9724097482283165, 0.9767381835635638, 0.9669484658390122, 0.659083299684951, -0.659083180712816, -0.9669485121167052, -0.9767382069792288, -0.9724097635533602, -0.9654367970450167, -0.9572273263645859, -0.9476871280825523, -0.9373647681120841, -0.9261598056102645, -0.9143881718456056, -0.9019807055316369, -0.8891341634240081, -0.8758057205293912, -0.8621229450911845, -0.8480618138204272, -0.833706571569058, -0.8190466131476127, -0.8041372124868691, -0.7889777195422356, -0.7736033858767385, -0.758020740007683, -0.7422507481169578, -0.7263049162371344, -0.7101938950789042, -0.6939334061553678, -0.677525822052029, -0.6609901538934517, -0.6443222327338847, -0.6275436932970322, -0.6106454472814152, -0.5936507836778451, -0.5765466491708988, -0.5593578078967361, -0.5420672759411125, -0.5247011730988912, -0.5072391580614087, -0.4897082914472909, -0.47208460952428394, -0.4543985995006753, -0.4366355580500639, -0.41882350871539187, -0.40094955631843376, -0.38303594105786365, -0.36507302109186685, -0.3470786936847069, -0.3290476440540586, -0.31099441589505206, -0.2929180880304103, -0.27483158663081614, -0.2567388003912687, -0.2386513127155433, -0.22057831776499126, -0.20253089403524566, -0.18452269630486776, -0.1665634500729787, -0.14867027528284874, -0.13084990929476334, -0.1131191325854089, -0.09547794429803691, -0.07793928430794522, -0.06049114408297565, -0.0431364527809777, -0.025846763281087953, -0.00861212501518312] )\n",
    "target = state0.copied_with(velocity=np.reshape(solution_t16,[n,1]))\n",
    "\n",
    "loss = math.sum( (states[16].velocity.data - target.velocity.data)**2 ) / n\n",
    "sess.initialize_variables()\n",
    "\n",
    "# compute initial loss\n",
    "s,l = sess.run([states[-1],loss], feed_dict={state_in: state0})\n",
    "print(\"Initial loss \"+format(l))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not surprisingly, because we're starting from zero, there's a significant intial error of ca. 0.38 for the 16th simulation step.\n",
    "\n",
    "Note that because we're only constraining timestep 16, we could actually omit steps 17 to 31 in this setup. They don't have any degrees of freedom and are not constrained in any way. For fairness and for direct comparison with the previous PINN case, let's include them.\n",
    "\n",
    "Now we have our simulation graph in TF, we can use TF to give us a gradient for the initial state for the loss. All we need to do is run `tf.gradients(loss, [state_in.velocity.data]`, which will give us a direction for each velocity variable. Based on a linear approximation, this direction tells us how to change each of them to increase the loss function (gradients _always_ point \"upwards\"). In the following code snipped, we're additionally saving all these gradients in a list called `grads`, such that we can visualize them later on. (Normally, we could discard each gradient after performing an update step.)\n",
    "\n",
    "Based on the gradient, we can now take a step in the opposite direction to bring the loss down (instead of increasing it). Below we're using a learning rate `LR=5` for this step. Afterwards, we're re-evaluating the loss for the updated state to check how we did. As these updates aren't exactly fast, we're only computing five of them here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "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",
      "Step 0, loss = 0.3829152584075928\n",
      "Step 1, loss = 0.3334265947341919\n",
      "Step 2, loss = 0.2905231714248657\n",
      "Step 3, loss = 0.25346457958221436\n",
      "Step 4, loss = 0.2215319573879242\n"
     ]
    }
   ],
   "source": [
    "LR = 5.\n",
    "stateN = state0\n",
    "grads = []\n",
    "for i in range(5):\n",
    "    grads.append( tf.gradients(loss, [state_in.velocity.data] ) ) # use TF\n",
    "    grads.append( gradients(loss, state_in.velocity) ) # phiflow wrapper also works\n",
    "\n",
    "    state_updated = state_in.copied_with( velocity=(state_in.velocity - LR*grads[-1]) )\n",
    "    stateN,l = sess.run([state_updated,loss], feed_dict={state_in: stateN})\n",
    "\n",
    "    #s,l = sess.run([states[-1],loss], feed_dict={state_in: stateN})\n",
    "    print( \"Step {}, loss = {}\".format(i,l) )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the reasons for the bad performance here is that we're querying states via `sess.run()` multiple times, instead of computing the update in one go. Nonetheless, this simple example already clearly shows that the loss is nicely going down: from above 0.38 to ca. 0.22. This is a good sign!\n",
    "\n",
    "Before improving on the performance, let's visualize the reconstruction.\n",
    "\n",
    "Starting from the zero state `state0` (shown in blue), the first gradient is shown as a red line here. If you compare it with the solution it points in the opposite direction, as expected. The solution is much larger in magnitude, so we omit it here (see the next graph).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x13fa0ae90>]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16968 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "pltx = np.linspace(-1,1,n)\n",
    "\n",
    "# first gradient\n",
    "[grad0,s0] = sess.run([grads[0],state0], feed_dict={state_in: stateN}) #.velocity.data\n",
    "fig.plot(pltx, grad0[0].flatten()             , lw=2, color='red')  \n",
    "\n",
    "fig.plot(pltx, state0.velocity.data.flatten(), lw=2, color='mediumblue')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So the gradient looks good (feel free to plot the others from the five iterations above).\n",
    "\n",
    "How well does the 16th state of the simulation actually match the target after the 5 update steps? This is what the loss measures, after all. The next graph shows the constraints (i.e. the solution we'd like to obtain) in green, and the reconstructed state after the initial state was updated 16 times by the solver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x13fbd43d0>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 20584 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure().gca()\n",
    "\n",
    "# target constraint at t=0.5\n",
    "fig.plot(pltx, target.velocity.data.flatten(), lw=2, color='forestgreen') \n",
    "\n",
    "# constrained state of simulation\n",
    "contrained_state = sess.run(states[16], feed_dict={state_in: stateN}).velocity.data\n",
    "fig.plot(pltx, contrained_state.flatten(), lw=2, color='mediumblue')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This seems to be going in the right direction! It's not there yet, but we've also only computed 5 steps so far. The two peaks with a positive velocity on the left side of the shock and the negative peak on the right side are starting to show.\n",
    "\n",
    "This is a good indicator that the backpropagation of gradients through all of our 16 simulated steps is behaving correctly, and driving the solution in the right direction. This hints at how powerful this setup is: the gradient that we obtain from each of the simulation steps (and each operation within them) can easily be chained together into more complex sequences. Here, we're backpropagation through all 16 steps of the simulation, and we could easily enlarge this \"look-ahead\" of the optimization without any additional effort.\n",
    "\n",
    "A simple direct reconstruction problem like this one is always a good initial test for a DP solver, e.g., before moving to more complex setups like coupling it with an NN. If the direct optimization does not converge, there's probably still something fundamentally wrong, and there's no point involving an NN. \n",
    "\n",
    "Our Burgers example here looks good - before including an NN, let's address the performance issue so: here we've updated our solutoin via a gradient that is \"manually\" computed with multiple TF session run() calls. TF's own optimizers do this internally, more flexibly and more efficiently. Let's start over, and switch to a `tf.train.Optimizer()` instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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": 4
}
