{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Burgers Optimization with a Differentiable Physics Gradient\n",
    "\n",
    "To illustrate the process of computing gradients in a _differentiable physics_ (DP) setting, we target the same inverse problem (the reconstruction task) used for the PINN example  in {doc}`physicalloss-code`. The choice of DP as a method has some immediate implications: we start with a discretized PDE, and the evolution of the  system is now fully determined by the resulting numerical solver. 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. Instead, we can rely on our discretized model. \n",
    "\n",
    "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 phiflow, as shown in the overview section on _Burgers forward simulations_. \n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/diffphys-code-burgers.ipynb)\n",
    "\n",
    "\n",
    "## Initialization\n",
    "\n",
    "phiflow directly gives us a sequence of differentiable operations, provided that we don't use the _numpy_ backend.\n",
    "The important step here is to include `phi.tf.flow` instad of `phi.flow` (for _pytorch_ you could use `phi.torch.flow`).\n",
    "\n",
    "So, as a first step, let's set up some constants, and intialize our domain with 128 points. We'll directly allocate a `velocity` field, and our constraint at $t=0.5$ (step 16), now as a `scalar_grid` in phiflow.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@develop\n",
    "#!pip install --upgrade --quiet phiflow\n",
    "from phi.tf.flow import *\n",
    "\n",
    "N = 128\n",
    "STEPS = 32\n",
    "DT = 1./STEPS\n",
    "NU = 0.2/np.pi\n",
    "\n",
    "DOMAIN = Domain(x=N, boundaries=PERIODIC, bounds=Box[-1:1])\n",
    "\n",
    "# initialization of velocities\n",
    "velocity = DOMAIN.scalar_grid() # start from zero\n",
    "\n",
    "SOLUTION_T16 = DOMAIN.scalar_grid(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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can verify that the fields of our simulation are now backed by TensorFlow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensorflow.python.framework.ops.EagerTensor"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(velocity.values.native())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradients\n",
    "\n",
    "The `record_gradients` function of phiflow triggers the generation of a gradient tape to compute gradients of a simulation via `math.gradients(loss, values)`.\n",
    "\n",
    "To use it for the Burgers case we need to specify a loss function: we want the solution at $t=0.5$ to match the reference data. Thus we simply compute an $L^2$ difference between step number 16 and our constraint array as `loss`. Afterwards, we evaluate the gradient of the initial velocity state `velocity` with respect to this loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/thuerey/miniconda3/envs/tf/lib/python3.8/site-packages/tensorflow/python/ops/math_grad.py:297: setdiff1d (from tensorflow.python.ops.array_ops) is deprecated and will be removed after 2018-11-30.\n",
      "Instructions for updating:\n",
      "This op will be removed after the deprecation date. Please switch to tf.sets.difference().\n",
      "Loss: 0.382915\n"
     ]
    }
   ],
   "source": [
    "velocities = [velocity]\n",
    "\n",
    "with math.record_gradients(velocity.values):\n",
    "\n",
    "    for time_step in range(STEPS):\n",
    "        v1 = diffuse.explicit(1.0*velocities[-1], NU, DT, substeps=1)\n",
    "        v2 = advect.semi_lagrangian(v1, v1, DT)\n",
    "        velocities.append(v2)\n",
    "\n",
    "    loss = field.l2_loss(velocities[16] - SOLUTION_T16)*2./N # MSE\n",
    "\n",
    "    grad = math.gradients(loss, velocity.values)\n",
    "\n",
    "print('Loss: %f' % (loss))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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. However, for fairness regarding a comparison with the previous PINN case, we include them.\n",
    "\n",
    "Note that we've done a lot of calcuations here: first the 32 steps of our simulation, and then another 16 steps backwards from the loss. They were recorded by the gradient tape, and used to backpropagate the loss to the initial state of the simulation.\n",
    "\n",
    "Not surprisingly, because we're starting from zero, there's also a significant intial error of ca. 0.38 for the 16th simulation step.\n",
    "\n",
    "So what do we get as a gradient here? It has the same dimensions as the velocity, and we can easily visualize it:\n",
    "Starting from the zero state for `velocity` (shown in blue), the first gradient is shown as a green line below. 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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19296 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab as plt\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "pltx = np.linspace(-1,1,N)\n",
    "\n",
    "# first gradient\n",
    "fig.plot(pltx, grad.numpy('x')           , lw=2, color='green',      label=\"Gradient\")  \n",
    "fig.plot(pltx, velocity.values.numpy('x'), lw=2, color='mediumblue', label=\"u at t=0\")\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend();\n",
    "\n",
    "# some (optional) other fields to plot:\n",
    "#fig.plot(pltx,  (velocities[16]).values.numpy('x')   , lw=2, color='cyan', label=\"u at t=0.5\")  \n",
    "#fig.plot(pltx,  (SOLUTION_T16).values.numpy('x')   , lw=2, color='red', label=\"solution at t=0.5\")  \n",
    "#fig.plot(pltx,  (velocities[16] - SOLUTION_T16).values.numpy('x')   , lw=2, color='blue', label=\"difference at t=0.5\")  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This gives us a \"search direction\" for each velocity variable. Based on a linear approximation, the gradient tells us how to change each of them to increase the loss function (gradients _always_ point \"upwards\"). Thus, we can use the gradient to run an optimization and find an initial state `velocity` that minimizes our loss.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization \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). We're using a learning rate `LR=5` for this step. We're also re-evaluating the loss for the updated state to check how we did. \n",
    "\n",
    "In the following code block, 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 for `velocity`.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization step 0, loss: 0.382915\n",
      "Optimization step 1, loss: 0.329677\n",
      "Optimization step 2, loss: 0.284684\n",
      "Optimization step 3, loss: 0.246604\n",
      "Optimization step 4, loss: 0.214305\n"
     ]
    }
   ],
   "source": [
    "LR = 5.\n",
    "\n",
    "grads=[]\n",
    "for optim_step in range(5):\n",
    "    velocities = [velocity]\n",
    "    with math.record_gradients(velocity.values):\n",
    "        for time_step in range(STEPS):\n",
    "            v1 = diffuse.explicit(1.0*velocities[-1], NU, DT)\n",
    "            v2 = advect.semi_lagrangian(v1, v1, DT)\n",
    "            velocities.append(v2)\n",
    "\n",
    "        loss = field.l2_loss(velocities[16] - SOLUTION_T16)*2./N # MSE\n",
    "        print('Optimization step %d, loss: %f' % (optim_step,loss))\n",
    "\n",
    "        grads.append( math.gradients(loss, velocity.values) )\n",
    "\n",
    "    velocity = velocity - LR * grads[-1]\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Now we can check well the 16th state of the simulation actually matches 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 `velocity` (which we updated five times via the gradient by now) was updated 16 times by the solver. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26164 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, SOLUTION_T16.values.numpy('x'), lw=2, color='forestgreen',  label=\"Reference\") \n",
    "\n",
    "# optimized state of our simulation after 16 steps\n",
    "fig.plot(pltx, velocities[16].values.numpy('x'), lw=2, color='mediumblue', label=\"Simulated velocity\")\n",
    "\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend(); plt.title(\"After 5 Optimization Steps at t=0.5\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This seems to be going in the right direction! It's defintely not perfect, but we've only computed 5 GD update 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 that it's 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. In the example above, we're backpropagating through all 16 steps of the simulation, and we could easily enlarge this \"look-ahead\" of the optimization with minor changes to the code.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More optimization steps\n",
    "\n",
    "Before moving on to more complex physics simulations, or involving NNs, let's finish the optimization task at hand, and run more steps to get a better solution.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization step 0, loss: 0.186823\n",
      "Optimization step 5, loss: 0.098101\n",
      "Optimization step 10, loss: 0.055802\n",
      "Optimization step 15, loss: 0.033452\n",
      "Optimization step 20, loss: 0.020761\n",
      "Optimization step 25, loss: 0.013248\n",
      "Optimization step 30, loss: 0.008643\n",
      "Optimization step 35, loss: 0.005733\n",
      "Optimization step 40, loss: 0.003848\n",
      "Runtime 47.50s\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "\n",
    "for optim_step in range(45):\n",
    "    velocities = [velocity]\n",
    "    with math.record_gradients(velocity.values):\n",
    "        for time_step in range(STEPS):\n",
    "            v1 = diffuse.explicit(1.0*velocities[-1], NU, DT)\n",
    "            v2 = advect.semi_lagrangian(v1, v1, DT)\n",
    "            velocities.append(v2)\n",
    "\n",
    "        loss = field.l2_loss(velocities[16] - SOLUTION_T16)*2./N # MSE\n",
    "        if optim_step%5==0: \n",
    "            print('Optimization step %d, loss: %f' % (optim_step,loss))\n",
    "\n",
    "        grad = math.gradients(loss, velocity.values)\n",
    "\n",
    "    velocity = velocity - LR * grad\n",
    "\n",
    "end = time.time()\n",
    "print(\"Runtime {:.2f}s\".format(end-start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thinking back to the PINN version from {doc}`diffphys-code-burgers`, we have a much lower error here after only 50 steps (by ca. an order of magnitude), and the overall runtime is also much lower (ca. a factor of 4, depending on your hardware).\n",
    "\n",
    "Let's plot again how well our solution at $t=0.5$ (blue) matches the constraints (green) now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28584 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure().gca()\n",
    "fig.plot(pltx, SOLUTION_T16.values.numpy('x'), lw=2, color='forestgreen',  label=\"Reference\") \n",
    "fig.plot(pltx, velocities[16].values.numpy('x'), lw=2, color='mediumblue', label=\"Simulated velocity\")\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend(); plt.title(\"After 50 Optimization Steps at t=0.5\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not bad. But how well is the initial state recovered via backpropagation through the 16 simulation steps? This is what we're changing, and because it's only indirectly constrained via the observation later in time there is more room to deviate from a desired or expected solution.\n",
    "\n",
    "This is shown in the next plot: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33244 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",
    "INITIAL_GT = np.asarray( [-np.sin(np.pi * x) * 1. for x in np.linspace(-1,1,N)] ) # 1D array\n",
    "fig.plot(pltx, INITIAL_GT.flatten()      , lw=2, color='forestgreen', label=\"Ground truth initial state\")  # ground truth initial state of sim\n",
    "fig.plot(pltx, velocity.values.numpy('x'), lw=2, color='mediumblue',  label=\"Optimized initial state\") # manual\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend(); plt.title(\"Initial State After 50 Optimization Steps\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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 solutions that exactly or numerically very closely satisfy the constraints. Hence, our minimizer not necessarily finds the exact state we started from. However, it's still quite close in this Burgers scenario.\n",
    "\n",
    "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": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31664 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab\n",
    "\n",
    "def show_state(a):\n",
    "    a=np.expand_dims(a, axis=2)\n",
    "    for i in range(4):\n",
    "        a = np.concatenate( [a,a] , axis=2)\n",
    "    a = np.reshape( a, [a.shape[0],a.shape[1]*a.shape[2]] )\n",
    "    fig, axes = pylab.subplots(1, 1, figsize=(16, 5))\n",
    "    im = axes.imshow(a, origin='upper', cmap='inferno')\n",
    "    pylab.colorbar(im) \n",
    "        \n",
    "# get numpy versions of all states \n",
    "vels = [ x.values.numpy('x,vector') for x in velocities] \n",
    "# concatenate along vector/features dimension\n",
    "vels = np.concatenate(vels, axis=-1) \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, the one with the PINN, and the DP version, 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Solutions Ground Truth (top), PINN (middle) , DiffPhys (bottom):\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37348 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 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\"] \n",
    "sol_pi=npfile=np.load(\"./temp/burgers-pinn-solution.npz\")[\"arr_0\"] \n",
    "sol_dp=npfile=np.load(\"./temp/burgers-diffphys-solution.npz\")[\"arr_0\"] \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,[N*3+20,33,1]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's quite clearly visible here that the PINN solution (in the middle) 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 (at the bottom) 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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAE PINN: 0.22991 \n",
      "MAE DP:   0.05811\n",
      "\n",
      "Error GT to PINN (top) , DiffPhys (bottom):\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29672 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 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: {:7.5f} \\nMAE DP:   {: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) ,[N*2+10,33,1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's a pretty clear result: the PINN error is almost 4 times higher than the one from 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.\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",
    "Now we have a first example to show similarities and differences of the two approaches. In the next section, we'll present a discussion of the findings so far, before moving to more complex cases in the following chapter.\n",
    "\n",
    "\n",
    "## Next steps\n",
    "\n",
    "As with the PINN version, there's variety of things that can be improved and experimented with using the code above:\n",
    "\n",
    "* You can try to adjust the training parameters to further improve the reconstruction.\n",
    "* As for the PINN case, you can activate a different optimizer, and observe the changing (not necessarily improved) convergence behavior.\n",
    "* Vary the number of steps, or the resolution of the simulation and reconstruction.\n"
   ]
  }
 ],
 "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
}
