{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Burgers Optimization with a DL-based DP\n",
    "\n",
    "As shown in the previous example, the differentable Burgers solver allows us to easily backpopragate through arbitrary sequences of time-updates. (If you haven't done so, please check out the \"manual\" example of the previous notebook - the setup below directly builds on that one.)\n",
    "\n",
    "Now, let's use a \"proper\" DL framework for this task, i.e. an actual optimizer from TensorFlow. This is more flexible as will also be applicable to NN training runs later on, and will make the optimization much faster. The first step, setting up the simulation is almost identical to before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "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/tf_backend.py:287: to_complex64 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.cast` instead.\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/tf_backend.py:375: The name tf.fft is deprecated. Please use tf.signal.fft instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/tf_backend.py:387: The name tf.ifft is deprecated. Please use tf.signal.ifft instead.\n",
      "\n",
      "WARNING:tensorflow:From /Users/thuerey/Dropbox/mbaDevelSelected/phiflow-v1.5/phi/tf/tf_backend.py:399: The name tf.real is deprecated. Please use tf.math.real 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": [
    "from phi.tf.flow import *\n",
    "\n",
    "# run with phiflow\n",
    "n = 128\n",
    "steps = 32\n",
    "dt = 1./steps\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",
    "# now use a real TF variable instead\n",
    "state_in = state0.copied_with(velocity= variable(state0.velocity) ) # <-- main change!\n",
    "states = [state_in]\n",
    "\n",
    "for i in range(steps):\n",
    "    states.append( Burgers().step(states[-1],dt=dt) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's easy to overlook, but there was an important change here: we've declared the initial state of the simulation to be a `tf.variable`. \n",
    "\n",
    "Now we don't need to keep track and pass along `state_in` as in the previous manual example, TF knows about the initial state as a variable we'd like to optimize. This is in contrast to the previous version of this optimization, where we explicitly stored and updated the `state_in` field containing the intial state.\n",
    "\n",
    "Nonetheless, this likewise didnt run anything! We only have a graph of operations in tensorflow so far. Let's set up the loss like before:"
   ]
  },
  {
   "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"
     ]
    }
   ],
   "source": [
    "sess = Session(None) \n",
    "\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the graph encodes the full temporal sequence, and we can compute and propagate derivatives for all intermediates, we can can now backpropagate a gradient for this loss to any earlier variable, and hence also the initial velocity state that we're interested in (and that we already declared as a `variable()` to TF above).\n",
    "\n",
    "Also, in contrast to the PINN example we don't need additional terms for enforcing the boundary conditions, these are already present in the simulation, and hence encoded in the execution graph for the `states` sequence above.\n",
    "\n",
    "Now we're ready to run the optimization with one of the optimizers provided in tensorflow, e.g., the simplest starting point is always _gradient descent_ (GD), i.e. the `tf.train.GradientDescentOptimizer` from TF:"
   ]
  },
  {
   "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",
      "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.382915\n",
      "Step 0, loss: 0.382915\n",
      "Step 5, loss: 0.194040\n",
      "Step 10, loss: 0.103610\n",
      "Step 15, loss: 0.059475\n",
      "Step 20, loss: 0.035991\n",
      "Step 25, loss: 0.022581\n",
      "Step 30, loss: 0.014490\n",
      "Step 35, loss: 0.009447\n",
      "Step 40, loss: 0.006238\n",
      "Step 45, loss: 0.004161\n",
      "Runtime 1.32s\n"
     ]
    }
   ],
   "source": [
    "LR = 5.\n",
    "\n",
    "optim = tf.train.GradientDescentOptimizer(learning_rate=LR).minimize(loss)\n",
    "#optim = tf.train.AdamOptimizer(learning_rate=0.05).minimize(loss) # overshoots?\n",
    "sess.initialize_variables() \n",
    "\n",
    "print('Initial loss: %f' % sess.run(loss))\n",
    "\n",
    "import time\n",
    "start = time.time()\n",
    "\n",
    "for optim_step in range(50):\n",
    "    _, loss_value = sess.run([optim, loss])\n",
    "    \n",
    "    if optim_step%5==0: \n",
    "        print('Optimization step %d, loss: %f' % (optim_step,loss_value))\n",
    "\n",
    "end = time.time()\n",
    "print(\"Runtime {:.2f}s\".format(end-start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code already includes a timer, and already ran 50 steps which should take around one second. This is much faster than our previous \"manual\" version. Also, the loss now then down to below 0.005.\n",
    "\n",
    "Let's visualize the reconstruction: first let's check how well the last state of the simulation matches the target:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x14384c190>]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22076 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",
    "# initial guess for optimization, this is where we started out...\n",
    "fig.plot(pltx, initial.flatten()             , lw=2, color='linen')  \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]).velocity.data\n",
    "fig.plot(pltx, contrained_state.flatten(), lw=2, color='mediumblue')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not bad. How well is the initial state recovered via backpropagation through the 16 simulation steps?\n",
    "\n",
    "This also matches quite well, as shown in the next plot. Naturally, this is a tougher task: the optimization receives direct feedback what the state at $t=0.5$ should look like, but due to the non-linear model equation, we typically have a large number of solution that exactly or numerically very closely satisfy the constraints. Hence, our minimizer not necessarily finds the exact state we started from. It's still quite close in this Burgers scenario:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x14386c5d0>]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24348 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure().gca()\n",
    "pltx = np.linspace(-1,1,n)\n",
    "\n",
    "# ground truth state at time=0 , move down\n",
    "t0gt = np.asarray( [ [-math.sin(np.pi * x) * 1.] for x in np.linspace(-1,1,n)] )\n",
    "fig.plot(pltx, t0gt.flatten()             , lw=2, color='forestgreen')  # ground truth initial state of sim\n",
    "\n",
    "#fig.plot(pltx, stateN.velocity.data.flatten(), lw=2, color='mediumblue') # manual\n",
    "fig.plot(pltx, sess.run(states[0]).velocity.data.flatten(), lw=2, color='mediumblue') # manual\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This was all very qualitative so far. Before measuring the overall error of the reconstruction, let's visualize the full evolution of our system over time as this also yields the solution in the form of a numpy array that we can compare to the other versions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vels array shape: (1, 128, 33, 1)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 10528 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def show_state(a):\n",
    "    # we only have 33 time steps, blow up by a factor of 2^4 to make it easier to see\n",
    "    for i in range(4):\n",
    "        a = np.concatenate( [a,a] , axis=3)\n",
    "\n",
    "    a = np.reshape( a, [a.shape[1],a.shape[2]*a.shape[3]] )\n",
    "    plt.imshow(a, origin='upper', cmap='magma')\n",
    "    \n",
    "# run all states through TF to get numpy arrays \n",
    "vels = sess.run(states) # convert all states in 1 sess.run() call\n",
    "vels = [ x.velocity.data for x in vels] \n",
    "#vels = np.asarray( [ sess.run(x, feed_dict={state_in: stateN}).velocity.data for x in states] ) # also works, but slower\n",
    "vels = np.concatenate(vels, axis=-1)\n",
    "vels = np.reshape(vels, list(vels.shape)+[1] ) ; print(\"Vels array shape: \"+format(vels.shape))\n",
    "\n",
    "# save for comparison with other methods\n",
    "np.savez_compressed(\"./temp/burgers-diffphys-solution.npz\", np.reshape(vels,[n,steps+1])) # remove batch & channel dimension\n",
    "\n",
    "show_state(vels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Physics-Informed vs. Differentiable Physics Reconstruction\n",
    "\n",
    "Now we have both versions, so let's compare both reconstructions in more detail.\n",
    "\n",
    "Let's first look at the solutions side by side. The code below generates an image with 3 versions, from top to bottom: the \"ground truth\" (GT) solution as given by the regular forward simulation, in the middle the PINN reconstruction, and at the bottom the differentiable physics version.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Solutions Ground Truth (top), PINN (middle) , DiffPhys (bottom):\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26008 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# note, this requires running the forward sim & PINN notebooks beforehand\n",
    "sol_gt=npfile=np.load(\"./temp/burgers-groundtruth-solution.npz\")[\"arr_0\"] ; #print(format(sol_gt.shape)) \n",
    "sol_pi=npfile=np.load(\"./temp/burgers-pinn-solution.npz\")[\"arr_0\"] ; #print(format(sol_pi.shape)) \n",
    "sol_dp=npfile=np.load(\"./temp/burgers-diffphys-solution.npz\")[\"arr_0\"] ; #print(format(sol_dp.shape)) \n",
    "\n",
    "divider = np.ones([10,33])*-1. # we'll sneak in a block of -1s to show a black divider in the image\n",
    "sbs = np.concatenate( [sol_gt, divider, sol_pi, divider, sol_dp], axis=0)\n",
    "\n",
    "print(\"\\nSolutions Ground Truth (top), PINN (middle) , DiffPhys (bottom):\")\n",
    "show_state(np.reshape(sbs,[1,n*3+20,33,1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's clealry visible here that the PINN solution recovers the overall shape of the solution, hence the temporal constraints are at least partially fulfilled. However, it doesn't manage to capture the amplitudes of the GT solution very well.\n",
    "\n",
    "The reconstruction from the optimization with a differentiable solver is much closer to the ground truth thanks to an improved flow of gradients over the whole course of the sequence. In addition, it can leverage the grid-based discretization for both forwards as well as backwards passes, and in this way provide a more accurate signal to the unknown initial state. It is nonetheless visible, that the reconstruction lacks certain \"sharper\" features of the GT version, e.g., visible in the bottom left corner of the solution image.\n",
    "\n",
    "Let's quantify these errors over the whole sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAE PINN: \t0.23629 \n",
      "MAE DiffPhys: \t0.04718\n",
      "\n",
      "Error GT to PINN (top) , DiffPhys (bottom):\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 16764 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "err_pi = np.sum( np.abs(sol_pi-sol_gt)) / (steps*n)\n",
    "err_dp = np.sum( np.abs(sol_dp-sol_gt)) / (steps*n)\n",
    "print(\"MAE PINN: \\t{:7.5f} \\nMAE DiffPhys: \\t{:7.5f}\".format(err_pi,err_dp))\n",
    "\n",
    "print(\"\\nError GT to PINN (top) , DiffPhys (bottom):\")\n",
    "show_state(np.reshape( np.concatenate([sol_pi-sol_gt, divider, sol_dp-sol_gt],axis=0) ,[1,n*2+10,33,1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's a pretty clear result: the PINN error is ca. 5 times higher than the Differentiable Physics (DP) reconstruction.\n",
    "\n",
    "This difference also shows clearly in the jointly visualized image at the bottom: the magnitudes of the errors of the DP reconstruction are much closer to zero, as indicated by the purple color above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a first example to show similarities and differences of the two approaches. In the section, we'll present a more in depth discussion of these findings, before moving to more complex cases.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Next steps\n",
    "\n",
    "As with the PINN version, there's variety of things that can be improved and experimented with here:\n",
    "\n",
    "* You can try to adjust the training parameters to improve the reconstruction.\n",
    "* As for the PINN case, you can activate a different optimizer, and observe the changing (not necessarily improved) behavior.\n",
    "* Vary the number of steps, or the resolution of the simulation and reconstruction."
   ]
  },
  {
   "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
}
