{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HykKFEeAoXan"
   },
   "source": [
    "# Burgers Optimization with a Physics-Informed NN\n",
    "\n",
    "To illustrate how the physics-informed losses work, let's consider a reconstruction task\n",
    "as an inverse problem example.\n",
    "We'll use Burgers equation $\\frac{\\partial u}{\\partial{t}} + u \\nabla u = \\nu \\nabla \\cdot \\nabla u$ as a simple yet non-linear equation in 1D, for which we have a series of _observations_ at time $t=0.5$. \n",
    "The solution should fulfill the residual formulation for Burgers equation and match the observations.\n",
    "In addition, let's impose Dirichlet boundary conditions $u=0$\n",
    "at the sides of our computational domain, and define the solution in\n",
    "the time interval $t \\in [0,1]$.\n",
    "\n",
    "Note that similar to the previous forward simulation example, \n",
    "we will still be sampling the solution with 128 points ($n=128$), but now we have a discretization via the NN. So we could also sample points inbetween without having to explicitly choose a basis function for interpolation. The discretization via the NN now internally determines how to use its degrees of freedom to arrange the activation functions as basis functions. So we have no direct control over the reconstruction.\n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/physicalloss-code.ipynb)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g9TYIr2loXar"
   },
   "source": [
    "## Formulation\n",
    "\n",
    "In terms of notation from {doc}`overview-equations` and the previous section, this reconstruction problem means we are solving\n",
    "\n",
    "$$\n",
    "\\text{arg min}_{\\theta} \\sum_i |f(x_i ; \\theta)-y^*_i|^2 + R(x_i) ,\n",
    "$$\n",
    "\n",
    "where $x$ and $y^*$ are solutions at different locations in space and time, i.e.  $x,y^* \\in \\mathbb{R}$.\n",
    "Together, they represent two-dimensional solutions\n",
    "$x(p_i,t_i)$ and $y^*(p_i,t_i)$ for a spatial coordinate $p_i$ and a time $t_i$, where the index $i$ sums over a set of chosen $p_i,t_i$\n",
    "locations. The residual function $R$ above collects additional evaluations of $f$ and its derivatives to formulate the residual for $\\mathcal{P}$. I.e., $R$ should simply converge to zero above, and we've omitted scaling factors in the objective function for simplicity.\n",
    "\n",
    "Note that, effectively, we're only dealing with individual point samples of a single solution for $\\mathcal{P}$ \n",
    "here.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3mh1Lf_XoXas"
   },
   "source": [
    "## Preliminaries\n",
    "\n",
    "Let's just load phiflow with the tensorflow backend for now, and initialize the random sampling. (_Note: this example uses an older version of phiflow (1.5.1)._)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "L30jgXowoXat",
    "outputId": "bcceefc0-fb33-44c7-959c-0b742e5a1606"
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@1.5.1\n",
    "\n",
    "from phi.tf.flow import *\n",
    "import numpy as np\n",
    "\n",
    "#rnd = TF_BACKEND  # for phiflow: sample different points in the domain each iteration\n",
    "rnd = math.choose_backend(1)  # use same random points for all iterations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JvVRcVfUoXat"
   },
   "source": [
    "We're importing phiflow here, but we won't use it to compute a solution to the PDE as in {doc}`overview-burgers-forw`. Instead, we'll use the  derivatives of an NN (as explained in the previous section) to set up a loss formulation for training.\n",
    "\n",
    "Next, we set up a simple NN with 8 fully connected layers and `tanh` activations with 20 units each. \n",
    "\n",
    "We'll also define the `boundary_tx` function which gives an array of constraints for the solution (all for $=0.5$ in this example), and the `open_boundary` function which stores constraints for $x= \\pm1$ being 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "AwWNtbm_oXau"
   },
   "outputs": [],
   "source": [
    "def network(x, t):\n",
    "    \"\"\" Dense neural network with 8 hidden layers and 3021 parameters in total.\n",
    "        Parameters will only be allocated once (auto reuse).\n",
    "    \"\"\"\n",
    "    y = math.stack([x, t], axis=-1)\n",
    "    for i in range(8):\n",
    "        y = tf.layers.dense(y, 20, activation=tf.math.tanh, name='layer%d' % i, reuse=tf.AUTO_REUSE)\n",
    "    return tf.layers.dense(y, 1, activation=None, name='layer_out', reuse=tf.AUTO_REUSE)\n",
    "\n",
    "def boundary_tx(N):\n",
    "    x = np.linspace(-1,1,128)\n",
    "    # precomputed solution from forward simulation:\n",
    "    u = 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",
    "    t = np.asarray(rnd.ones_like(x)) * 0.5\n",
    "    perm = np.random.permutation(128) \n",
    "    return (x[perm])[0:N], (t[perm])[0:N], (u[perm])[0:N]\n",
    "\n",
    "def _ALT_t0(N): # alternative, impose original initial state at t=0\n",
    "    x = rnd.random_uniform([N], -1, 1)\n",
    "    t = rnd.zeros_like(x)\n",
    "    u = - math.sin(np.pi * x)\n",
    "    return x, t, u\n",
    "\n",
    "def open_boundary(N):\n",
    "    t = rnd.random_uniform([N], 0, 1)\n",
    "    x = math.concat([math.zeros([N//2]) + 1, math.zeros([N//2]) - 1], axis=0)\n",
    "    u = math.zeros([N])\n",
    "    return x, t, u\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4qpLH7XEoXaw"
   },
   "source": [
    "Most importantly, we can now also construct the residual loss function `f` that we'd like to minimize in order to guide the NN to retrieve a solution for our model equation. As can be seen in the equation at the top, we need derivatives w.r.t. $t$, $x$ and a second derivative for $x$. The first three lines of `f` below do just that.\n",
    "\n",
    "Afterwards, we simply combine the derivates according to Burgers equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "E2LrmmT4oXaw"
   },
   "outputs": [],
   "source": [
    "def f(u, x, t):\n",
    "    \"\"\" Physics-based loss function with Burgers equation \"\"\"\n",
    "    u_t = gradients(u, t)\n",
    "    u_x = gradients(u, x)\n",
    "    u_xx = gradients(u_x, x)\n",
    "    return u_t + u*u_x - (0.01 / np.pi) * u_xx\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TWR6OvKboXax"
   },
   "source": [
    "Next, let's set up the sampling points in the inner domain, such that we can compare the solution with the previous forward simulation in phiflow. \n",
    "\n",
    "The next cell allocates two tensors: `grid_x` will cover the size of our domain, i.e., the -1 to 1 range, with 128 cells, while `grid_t` will sample the time interval $[0,1]$ with 33 time stamps.\n",
    "\n",
    "The last `math.expand_dims()` call simply adds another `batch` dimension, so that the resulting tensor is compatible with the following examples.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xxLRbUQzoXay",
    "outputId": "a2504aae-26dc-48a3-a5fa-12c6539f3f62"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-2-f2ebc84042fb>:7: dense (from tensorflow.python.keras.legacy_tf_layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.Dense instead.\n",
      "WARNING:tensorflow:From /Users/thuerey/miniconda3/envs/tf/lib/python3.8/site-packages/tensorflow/python/keras/legacy_tf_layers/core.py:187: Layer.apply (from tensorflow.python.keras.engine.base_layer_v1) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.__call__` method instead.\n"
     ]
    }
   ],
   "source": [
    "N=128\n",
    "grids_xt = np.meshgrid(np.linspace(-1, 1, N), np.linspace(0, 1, 33), indexing='ij')\n",
    "grid_x, grid_t = [tf.convert_to_tensor(t, tf.float32) for t in grids_xt]\n",
    "\n",
    "# create 4D tensor with batch and channel dimensions in addition to space and time\n",
    "# in this case gives shape=(1, N, 33, 1)\n",
    "grid_u = math.expand_dims(network(grid_x, grid_t))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3iPGedVcoXaz"
   },
   "source": [
    "Now, `grid_u` contains a full graph to evaluate our NN at $128 \\times 33$ positions, and returns the results in a $[1,128,33,1]$ array once we run it through `session.run`. Let's give this a try: we can initialize a TF session, evaluate `grid_u` and show it in an image, just like the phiflow solution we computed previously. \n",
    "\n",
    "(Note, we'll use the `show_state` as in {doc}`overview-burgers-forw`. Hence, the x axis does not show actual simulation time, but is showing 32 steps \"blown\" up by a factor of 16 to make the changes over time easier to see in the image.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 356
    },
    "id": "MHYmnuiEoXaz",
    "outputId": "7f48e8b7-97f3-442c-bf67-e1deb3aff6ea",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of grid_u: (1, 128, 33, 1)\n",
      "Randomly initialized network state:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26372 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab as plt\n",
    "print(\"Size of grid_u: \"+format(grid_u.shape))\n",
    "\n",
    "session = Session(None)\n",
    "session.initialize_variables()\n",
    "\n",
    "def show_state(a, title):\n",
    "    for i in range(4): a = np.concatenate( [a,a] , axis=3)\n",
    "    a = np.reshape( a, [a.shape[1],a.shape[2]*a.shape[3]] )\n",
    "    fig, axes = plt.subplots(1, 1, figsize=(16, 5))\n",
    "    im = axes.imshow(a, origin='upper', cmap='inferno')\n",
    "    plt.colorbar(im) ; plt.xlabel('time'); plt.ylabel('x'); plt.title(title)\n",
    "    \n",
    "print(\"Randomly initialized network state:\")\n",
    "show_state(session.run(grid_u),\"Uninitialized NN\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "amnB9mIBoXa0"
   },
   "source": [
    "This visualization already shows a smooth transition over space and time. So far, this is purely the random initialization of the NN that we're sampling here. So it has nothing to do with a solution of our PDE-based model up to now.\n",
    "\n",
    "The next steps will actually evaluate the constraints in terms of data (from the `boundary` functions), and the model constraints from `f` to retrieve an actual solution to the PDE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-PHuAh0UoXa1"
   },
   "source": [
    "## Loss function and training\n",
    "\n",
    "As objective for the learning process we can now combine the _direct_ constraints, i.e., the solution at $t=0.5$ and the Dirchlet $u=0$ boundary conditions with the loss from the PDE residuals. For both boundary constraints we'll use 100 points below, and then sample the solution in the inner region with an additional 1000 points.\n",
    "\n",
    "The direct constraints are evaluated via `network(x, t)[:, 0] - u`, where `x` and `t` are the space-time location where we'd like to sample the solution, and `u` provides the corresponding ground truth value.\n",
    "\n",
    "For the physical loss points, we have no ground truth solutions, but we'll only evaluate the PDE residual via the NN derivatives, to see whether the solution satisfies PDE model. If not, this directly gives us an error to be reduced via a update step in the optimization. The corresponding expression is of the form  `f(network(x, t)[:, 0], x, t)` below. Note that for both data and physics terms the `network()[:, 0]` expressions don't remove any data from the $L^2$ evaluation, they simply discard the last size-1 dimension of the $(n,1)$ tensor returned by the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bfZvwnFooXa1",
    "outputId": "f0188a14-5922-4faa-b1aa-fb6482bd81f3"
   },
   "outputs": [],
   "source": [
    "# Boundary loss\n",
    "N_SAMPLE_POINTS_BND = 100\n",
    "x_bc, t_bc, u_bc = [math.concat([v_t0, v_x], axis=0) for v_t0, v_x in zip(boundary_tx(N_SAMPLE_POINTS_BND), open_boundary(N_SAMPLE_POINTS_BND))]\n",
    "x_bc, t_bc, u_bc = np.asarray(x_bc,dtype=np.float32), np.asarray(t_bc,dtype=np.float32) ,np.asarray(u_bc,dtype=np.float32)\n",
    "#with app.model_scope():\n",
    "loss_u = math.l2_loss(network(x_bc, t_bc)[:, 0] - u_bc)  # normalizes by first dimension, N_bc\n",
    "\n",
    "# Physics loss inside of domain\n",
    "N_SAMPLE_POINTS_INNER = 1000\n",
    "x_ph, t_ph = tf.convert_to_tensor(rnd.random_uniform([N_SAMPLE_POINTS_INNER], -1, 1)), tf.convert_to_tensor(rnd.random_uniform([N_SAMPLE_POINTS_INNER], 0, 1))\n",
    "loss_ph = math.l2_loss(f(network(x_ph, t_ph)[:, 0], x_ph, t_ph))  # normalizes by first dimension, N_ph\n",
    "\n",
    "# Combine\n",
    "ph_factor = 1.\n",
    "loss = loss_u + ph_factor * loss_ph # allows us to control the relative influence of loss_ph \n",
    "\n",
    "optim = tf.train.GradientDescentOptimizer(learning_rate=0.02).minimize(loss)\n",
    "#optim = tf.train.AdamOptimizer(learning_rate=0.001).minimize(loss) # alternative, but not much benefit here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rJBkiuECoXa2"
   },
   "source": [
    "The code above just initializes the evaluation of the loss, we still didn't do any optimization steps, but we're finally in a good position to get started with this.\n",
    "\n",
    "Despite the simple equation, the convergence is typicaly very slow. The iterations themselves are fast to compute, but this setup needs a _lot_ of iterations. To keep the runtime in a reasonable range, we only do 10k iterations by default below (`ITERS`). You can increase this value to get better results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "64g5NVHAoXa3",
    "outputId": "b0425572-351c-4877-9d49-a83ae9a91a91"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, loss: 0.108430\n",
      "Step 1, loss: 0.103907\n",
      "Step 2, loss: 0.100539\n",
      "Step 1000, loss: 0.057177\n",
      "Step 2000, loss: 0.053230\n",
      "Step 3000, loss: 0.049326\n",
      "Step 4000, loss: 0.046316\n",
      "Step 5000, loss: 0.044153\n",
      "Step 6000, loss: 0.041827\n",
      "Step 7000, loss: 0.039517\n",
      "Step 8000, loss: 0.037237\n",
      "Step 9000, loss: 0.035256\n",
      "Step 10000, loss: 0.033931\n",
      "Runtime 110.15s\n"
     ]
    }
   ],
   "source": [
    "session.initialize_variables()\n",
    "\n",
    "import time\n",
    "start = time.time()\n",
    "\n",
    "ITERS = 10000\n",
    "for optim_step in range(ITERS+1):\n",
    "  _, loss_value = session.run([optim, loss])\n",
    "  if optim_step<3 or optim_step%1000==0: \n",
    "        print('Step %d, loss: %f' % (optim_step,loss_value))\n",
    "        #show_state(grid_u)\n",
    "        \n",
    "end = time.time()\n",
    "print(\"Runtime {:.2f}s\".format(end-start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5NLrym59oXa4"
   },
   "source": [
    "This training can take a significant amount of time, around 1 minute on a typical notebook, but at least the error goes down significantly (roughly from around 0.2 to ca. 0.03), and the network seems to successfully converge to a solution.\n",
    "\n",
    "Let's show the reconstruction of the network, by evaluating the network at the centers of a regular grid, so that we can show the solution as an image. Note that this is actually fairly expensive, we have to run through the whole network with a few thousand weights for all of the $128 \\times 32$ sampling points in the grid.\n",
    "\n",
    "It looks pretty good on first sight, though. There's been a very noticeable change compared to the random initialization shown above:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 323
    },
    "id": "JXcISF40oXa4",
    "outputId": "afa97e6c-4d63-433a-aa45-fa29c2c7b17e"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28216 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_state(session.run(grid_u),\"After Training\") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JD-WnU-1oXa5"
   },
   "source": [
    "---\n",
    "\n",
    "## Evaluation\n",
    "\n",
    "Let's compare solution in a bit more detail. Here are the actual sample points used for constraining the solution (at time step 16, $t=1/2$) shown in gray, versus the reconstructed solution in blue:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 332
    },
    "id": "A95iq_2_oXa6",
    "outputId": "59c9373e-b56b-47dd-ed7f-de7830a6c43e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f84721c3fd0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 20096 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "u = session.run(grid_u)\n",
    "\n",
    "# solution is imposed at t=1/2 , which is 16 in the array\n",
    "BC_TX = 16 \n",
    "uT = u[0,:,BC_TX,0]\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "fig.plot(np.linspace(-1,1,len(uT)), uT, lw=2, color='blue', label=\"NN\")\n",
    "fig.scatter(x_bc[0:100], u_bc[0:100], color='gray', label=\"Reference\")\n",
    "plt.title(\"Comparison at t=1/2\")\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dBbSjnJ8oXa6"
   },
   "source": [
    "Not too bad at the sides of the domain (the Dirichlet boundary conditions $u=0$ are fulfilled), but the shock in the center (at $x=0$) is not well represented.\n",
    "\n",
    "Let's check how well the initial state at $t=0$ was reconstructed. That's the most interesting, and toughest part of the problem (the rest basically follows from the model equation and boundary conditions, given the first state).\n",
    "\n",
    "It turns out, the accuracy of the initial state is actually not that great: the blue curve from the PINN is quite far away from the constraints via the reference data (shown in gray)... The solution will get better with larger number of iterations, but it requires a surprisingly large number of them for this fairly simple case. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 332
    },
    "id": "4n5dAMhBoXa7",
    "outputId": "e6c0e4a2-b5d3-476f-f349-f2b1ca75917a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f847223da60>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 25640 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ground truth solution at t0\n",
    "t0gt = np.asarray( [ [-math.sin(np.pi * x) * 1.] for x in np.linspace(-1,1,N)] )\n",
    "velP0 = u[0,:,0,0]\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "fig.plot(np.linspace(-1,1,len(velP0)), velP0, lw=2, color='blue', label=\"NN\")\n",
    "fig.plot(np.linspace(-1,1,len(t0gt)), t0gt, lw=2, color='gray', label=\"Reference\") \n",
    "plt.title(\"Comparison at t=0\")\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7m382fpVoXa8"
   },
   "source": [
    "Especially the maximum / minimum at $x=\\pm 1/2$ are far off, and the boundaries at $x=\\pm 1$ are not fulfilled: the solution is not at zero.\n",
    "\n",
    "We have the forward simulator for this simulation, so we can use the $t=0$ solution of the network to \n",
    "evaluate how well the temporal evoluation was reconstructed. This measures how well the temporal evolution of the model equation was captured via the soft constraints of the PINN loss.\n",
    "\n",
    "The graph below shows the initial state in blue, and two evolved states at $t=8/32$ and $t=15/32$. Note that this is all from the simulated version, we'll show the learned version next. \n",
    "\n",
    "(Note: The code segments below also have some optional code to show the states at `[STEPS//4]`. It's commented out by default, you can uncomment or add additional ones to visualize more of the time evolution if you like.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 348
    },
    "id": "v9uHDCL0oXa8",
    "outputId": "a52cca99-fc5b-44de-aadb-e66c02ea3388"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 128, 1)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f84722de760>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 29792 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# re-simulate with phiflow from solution at t=0\n",
    "DT = 1./32.\n",
    "STEPS = 32-BC_TX # depends on where BCs were imposed\n",
    "INITIAL = u[...,BC_TX:(BC_TX+1),0] # np.reshape(u0, [1,len(u0),1]) \n",
    "print(INITIAL.shape)\n",
    "\n",
    "DOMAIN = Domain([N], boundaries=PERIODIC, box=box[-1:1])\n",
    "state = [BurgersVelocity(DOMAIN, velocity=INITIAL, viscosity=0.01/np.pi)]\n",
    "physics = Burgers()\n",
    "\n",
    "for i in range(STEPS):\n",
    "    state.append( physics.step(state[-1],dt=DT) )\n",
    "\n",
    "# we only need \"velocity.data\" from each phiflow state\n",
    "vel_resim = [x.velocity.data for x in state]\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "pltx = np.linspace(-1,1,len(vel_resim[0].flatten()))\n",
    "fig.plot(pltx, vel_resim[ 0].flatten(),       lw=2, color='blue',  label=\"t=0\")\n",
    "#fig.plot(pltx, vel_resim[STEPS//4].flatten(), lw=2, color='green', label=\"t=0.125\")\n",
    "fig.plot(pltx, vel_resim[STEPS//2].flatten(), lw=2, color='cyan',  label=\"t=0.25\")\n",
    "fig.plot(pltx, vel_resim[STEPS-1].flatten(),  lw=2, color='purple',label=\"t=0.5\")\n",
    "#fig.plot(pltx, t0gt, lw=2, color='gray', label=\"t=0 Reference\") # optionally show GT, compare to blue\n",
    "plt.title(\"Resimulated u from solution at t=0\")\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "71n_TaJkoXa9"
   },
   "source": [
    "And here is the PINN output from `u` at the same time steps:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 348
    },
    "id": "tetNobgfoXa9",
    "outputId": "c385a596-c24d-4164-8b4b-8d39014bf5ce"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(128,)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f8470cfb1f0>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 27744 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "velP = [u[0,:,x,0] for x in range(33)]\n",
    "print(velP[0].shape)\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "fig.plot(pltx, velP[BC_TX+ 0].flatten(),       lw=2, color='blue',  label=\"t=0\")\n",
    "#fig.plot(pltx, velP[BC_TX+STEPS//4].flatten(), lw=2, color='green', label=\"t=0.125\")\n",
    "fig.plot(pltx, velP[BC_TX+STEPS//2].flatten(), lw=2, color='cyan',  label=\"t=0.25\")\n",
    "fig.plot(pltx, velP[BC_TX+STEPS-1].flatten(),  lw=2, color='purple',label=\"t=0.5\")\n",
    "plt.title(\"NN Output\")\n",
    "plt.xlabel('x'); plt.ylabel('u'); plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qx6rOskWoXa-"
   },
   "source": [
    "Judging via eyeball norm, these two versions of $u$ look quite similar, but not surprisingly the errors grow over time and there are significant differences. Especially the steepening of the solution near the shock at $x=0$ is not \"captured\" well. It's a bit difficult to see in these two graphs, though, let's quantify the error and show the actual difference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 348
    },
    "id": "lWP_Rf2aoXa_",
    "outputId": "e6bb27b6-9020-44d1-f499-e1a8547b11cf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean absolute error for re-simulation across 16 steps: 0.01067\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f8472550e50>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24292 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "error = np.sum( np.abs( np.asarray(vel_resim[0:16]).flatten() - np.asarray(velP[BC_TX:BC_TX+STEPS]).flatten() )) / (STEPS*N)\n",
    "print(\"Mean absolute error for re-simulation across {} steps: {:7.5f}\".format(STEPS,error))\n",
    "\n",
    "fig = plt.figure().gca()\n",
    "fig.plot(pltx, (vel_resim[0       ].flatten()-velP[BC_TX         ].flatten()), lw=2, color='blue',  label=\"t=5\")\n",
    "fig.plot(pltx, (vel_resim[STEPS//4].flatten()-velP[BC_TX+STEPS//4].flatten()), lw=2, color='green', label=\"t=0.625\")\n",
    "fig.plot(pltx, (vel_resim[STEPS//2].flatten()-velP[BC_TX+STEPS//2].flatten()), lw=2, color='cyan',  label=\"t=0.75\")\n",
    "fig.plot(pltx, (vel_resim[STEPS-1 ].flatten()-velP[BC_TX+STEPS-1 ].flatten()), lw=2, color='purple',label=\"t=1\")\n",
    "plt.title(\"u Error\")\n",
    "plt.xlabel('x'); plt.ylabel('MAE'); plt.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XDI_7eK8oXbA"
   },
   "source": [
    "The code above will compute a mean absolute error of ca. $1.5 \\cdot 10^{-2}$ between ground truth re-simulation and the PINN evolution, which is significant for the value range of the simulation.\n",
    "\n",
    "And for comparison with the forward simulation and following cases, here are also all steps over time with a color map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 370
    },
    "id": "2Trf5FSOoXbB",
    "outputId": "ec10bceb-fca2-4e96-d66d-609feed5e913",
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28304 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# show re-simulated solution again as full image over time\n",
    "sn = np.concatenate(vel_resim, axis=-1)\n",
    "sn = np.reshape(sn, list(sn.shape)+[1] ) # print(sn.shape)\n",
    "show_state(sn,\"Re-simulated u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EKQ8qe4FoXbC"
   },
   "source": [
    "Next, we'll store the full solution over the course of the $t=0 \\dots 1$ time interval, so that we can compare it later on to the full solution from a regular forward solve and compare it to the differential physics solution.\n",
    "\n",
    "Thus, stay tuned for the full evaluation and the comparison. This will follow in {doc}`diffphys-code-tf`, after we've discussed the details of how to run the differential physics optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BY0ltx8uoXbC",
    "outputId": "f4a0333b-0250-40bf-fd9d-e65c35f3e1ac"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vels array shape: (128, 33)\n"
     ]
    }
   ],
   "source": [
    "vels = session.run(grid_u) # special for showing NN results, run through TF \n",
    "vels = np.reshape( vels, [vels.shape[1],vels.shape[2]] )\n",
    "\n",
    "# save for comparison with other methods\n",
    "import os; os.makedirs(\"./temp\",exist_ok=True)\n",
    "np.savez_compressed(\"./temp/burgers-pinn-solution.npz\",vels) ; print(\"Vels array shape: \"+format(vels.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "50WLIo4voXbD"
   },
   "source": [
    "---\n",
    "\n",
    "## Next steps\n",
    "\n",
    "This setup is just a starting point for PINNs and physical soft-constraints, of course. The parameters of the setup were chosen to run relatively quickly. As we'll show in the next sections, the behavior of such an inverse solve can be improved substantially by a tighter integration of solver and learning. \n",
    "\n",
    "The solution of the PINN setup above can also directly be improved, however. E.g., try to:\n",
    "\n",
    "* Adjust parameters of the training to further decrease the error without making the solution diverge\n",
    "* Adapt the NN architecture for further improvements (keep track of the weight count, though).\n",
    "* Activate a different optimizer, and observe the change in behavior (this typically requires adjusting the learning rate). Note that the more complex optimizers don't necessarily do better in this relatively simple example.\n",
    "* Or modify the setup to make the test case more interesting: e.g., move the boundary conditions to a later point in simulation time, to give the reconstruction a larger time interval to reconstruct."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "physicalloss-code-may17b.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
}
