{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Simple Forward Simulation of Burgers Equation with phiflow\n",
    "\n",
    "This chapter will give an introduction for how to run _forward_, i.e., regular simulations starting with a given initial state and approximating a later state numerically, and introduce the Φ<sub>Flow</sub> framework. Φ<sub>Flow</sub> provides a set of differentiable building blocks that directly interface with deep learning frameworks, and hence is a very good basis for the topics of this book. Before going for deeper and more complicated integrations, this notebook (and the next one), will show how regular simulations can be done with Φ<sub>Flow</sub>. Later on, we'll show that these simulations can be easily coupled with neural networks.\n",
    "\n",
    "The main repository for Φ<sub>Flow</sub> (in the following \"phiflow\") is [https://github.com/tum-pbs/PhiFlow](https://github.com/tum-pbs/PhiFlow), and additional API documentation and examples can be found at [https://tum-pbs.github.io/PhiFlow/](https://tum-pbs.github.io/PhiFlow/).\n",
    "\n",
    "For this jupyter notebook (and all following ones), you can find a _\"[run in colab]\"_ link at the end of the first paragraph. This link will load the latest version from the PBDL github repo in a colab notebook that you can execute on the spot: \n",
    "[[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/overview-burgers-forw.ipynb)\n",
    "\n",
    "## Model\n",
    "\n",
    "As physical model we'll use Burgers equation.\n",
    "This equation is a very simple, yet non-linear and non-trivial, model equation that can lead to interesting shock formations. Hence, it's a very good starting point for experiments, and it's 1D version (from equation {eq}`model-burgers1d`) is given by:\n",
    "\n",
    "$$\n",
    "    \\frac{\\partial u}{\\partial{t}} + u \\nabla u =\n",
    "    \\nu \\nabla\\cdot \\nabla u\n",
    "$$ \n",
    "\n"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Importing and loading phiflow\n",
    "\n",
    "Let's get some preliminaries out of the way: first we'll import the phiflow library, more specifically the `numpy` operators for fluid flow simulations: `phi.flow` (differentiable versions for a DL framework _X_ are loaded via `phi.X.flow` instead).\n",
    "\n",
    "**Note:** Below, the first command with a \"!\" prefix will install the [phiflow python package from GitHub](https://github.com/tum-pbs/PhiFlow) via `pip` in your python environment once you uncomment it. We've assumed that phiflow isn't installed, but if you have already done so, just comment out the first line (the same will hold for all following notebooks)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "#!pip install --upgrade --quiet phiflow\n",
    "!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@develop\n",
    "\n",
    "from phi.flow import *\n",
    "\n",
    "from phi import __version__\n",
    "print(\"Using phiflow version: {}\".format(phi.__version__))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Using phiflow version: 2.0.0rc2\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next we can define and initialize the necessary constants (denoted by upper-case names):\n",
    "our simulation domain will have `N=128` cells as discretization points for the 1D velocity $u$ in a periodic domain $\\Omega$ for the interval $[-1,1]$. We'll use 32 time `STEPS` for a time interval of 1, giving us `DT=1/32`. Additionally, we'll use a viscosity `NU` of $\\nu=0.01/\\pi$.\n",
    "\n",
    "We'll also define an initial state given by $-\\text{sin}(\\pi x)$ in the numpy array `INITIAL_NUMPY`, which we'll use to initialize the velocity $u$ in the simulation in the next cell. This initialization will produce a nice shock in the center of our domain. \n",
    "\n",
    "Phiflow is object-oriented and centered around field data in the form of grids (internally represented by a tensor object). I.e. you assemble your simulation by constructing a number of grids, and updating them over the course of time steps. \n",
    "\n",
    "Phiflow internally works with tensors that have named dimensions. This will be especially handy later on for 2D simulations with additional batch and channel dimensions, but for now we'll simply convert the 1D array into a phiflow tensor that has a single spatial dimension `'x'`.\n"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "N = 128\n",
    "STEPS = 32\n",
    "DT = 1./STEPS\n",
    "NU = 0.01/np.pi\n",
    "\n",
    "# initialization of velocities\n",
    "INITIAL_NUMPY = np.asarray( [-np.sin(np.pi * x) * 1. for x in np.linspace(-1,1,N)] ) # 1D numpy array\n",
    "\n",
    "INITIAL = math.tensor(INITIAL_NUMPY, spatial('x') ) # convert to phiflow tensor\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "\n",
    "Next, we initialize a 1D `velocity` grid from the `INITIAL` numpy array that was converted into a tensor.\n",
    "The extent of our domain $\\Omega$ is specifiied via the `bounds` parameter $[-1,1]$, and the grid uses periodic boundary conditions (`extrapolation.PERIODIC`). These two properties are the main difference between a tensor and a grid: the latter has boundary conditions and a physical extent.\n",
    "\n",
    "Just to illustrate, we'll also print some info about the velocity object: it's a `phi.math` tensor with a size of 128. Note that the actual grid content is contained in the `values` of the grid. Below we're printing five entries by using the `numpy()` function to convert the content of the phiflow tensor into a numpy array. For tensors with more dimensions, we'd need to specify the order here, e.g., `'y,x,vector'` for a 2D velocity field. (If we'd call `numpy('x,vector')` below, this would convert the 1D array into one with an additional dimension for each 1D velocity component.)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "velocity = CenteredGrid(INITIAL, extrapolation.PERIODIC, x=N, bounds=Box[-1:1])\n",
    "#velocity = CenteredGrid(Noise(), extrapolation.PERIODIC, x=N, bounds=Box[-1:1]) # random init\n",
    "\n",
    "print(\"Velocity tensor shape: \"   + format( velocity.shape )) # == velocity.values.shape\n",
    "print(\"Velocity tensor type: \"    + format( type(velocity.values) ))\n",
    "print(\"Velocity tensor entries 10 to 14: \" + format( velocity.values.numpy()[10:15] ))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Velocity tensor shape: (xˢ=128)\n",
      "Velocity tensor type: <class 'phi.math._tensors.CollapsedTensor'>\n",
      "Velocity tensor entries 10 to 14: [0.47480196 0.51774486 0.55942075 0.59972764 0.6385669 ]\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Running the simulation\n",
    "\n",
    "Now we're ready to run the simulation itself. To ccompute the diffusion and advection components of our model equation we can simply call the existing `diffusion` and `semi_lagrangian` operators in phiflow: `diffuse.explicit(u,...)` computes an explicit diffusion step via central differences for the term $\\nu \\nabla\\cdot \\nabla u$ of our model. Next, `advect.semi_lagrangian(f,u)` is used for a stable first-order approximation of the transport of an arbitrary field `f` by a velocity `u`. In our model we have $\\partial u / \\partial{t} + u \\nabla f$, hence we use the `semi_lagrangian` function to transport the velocity with itself in the implementation:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "velocities = [velocity]\n",
    "age = 0.\n",
    "for i in range(STEPS):\n",
    "    v1 = diffuse.explicit(velocities[-1], NU, DT)\n",
    "    v2 = advect.semi_lagrangian(v1, v1, DT)\n",
    "    age += DT\n",
    "    velocities.append(v2)\n",
    "\n",
    "print(\"New velocity content at t={}: {}\".format( age, velocities[-1].values.numpy()[0:5] ))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "New velocity content at t=1.0: [0.00274862 0.01272991 0.02360343 0.03478042 0.0460869 ]\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Here we're actually collecting all time steps in the list `velocities`. This is not necessary in general (and could consume lots of memory for long-running sims), but useful here to plot the evolution of the velocity states later on.\n",
    "\n",
    "The print statements print a few of the velocity entries, and already show that something is happening in our simulation, but it's difficult to get an intuition for the behavior of the PDE just from these numbers. Hence, let's visualize the states over time to show what is happening.\n",
    "\n",
    "## Visualization\n",
    "\n",
    "We can visualize this 1D case easily in a graph: the following code shows the initial state in blue, and then times $10/32, 20/32, 1$ in green, cyan and purple. "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "# get \"velocity.values\" from each phiflow state with a channel dimensions, i.e. \"vector\"\n",
    "vels = [v.values.numpy('x,vector') for v in velocities] # gives a list of 2D arrays \n",
    "\n",
    "import pylab\n",
    "\n",
    "fig = pylab.figure().gca()\n",
    "fig.plot(np.linspace(-1,1,len(vels[ 0].flatten())), vels[ 0].flatten(), lw=2, color='blue',  label=\"t=0\")\n",
    "fig.plot(np.linspace(-1,1,len(vels[10].flatten())), vels[10].flatten(), lw=2, color='green', label=\"t=0.3125\")\n",
    "fig.plot(np.linspace(-1,1,len(vels[20].flatten())), vels[20].flatten(), lw=2, color='cyan',  label=\"t=0.625\")\n",
    "fig.plot(np.linspace(-1,1,len(vels[32].flatten())), vels[32].flatten(), lw=2, color='purple',label=\"t=1\")\n",
    "pylab.xlabel('x'); pylab.ylabel('u'); pylab.legend()\n"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f9ca2d4ef70>"
      ]
     },
     "metadata": {},
     "execution_count": 5
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ],
      "image/png": "[encoded data, 37736 characters]"
     },
     "metadata": {
      "needs_background": "light"
     }
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "This nicely shows the shock developing in the center of our domain, which forms from the collision of the two initial velocity \"bumps\", the positive one on left (moving right) and the negative one right of the center (moving left).\n",
    "\n",
    "As these lines can overlap quite a bit we'll also use a different visualization in the following chapters that shows the evolution over the course of all time steps in a 2D image. Our 1D domain will be shown along the Y-axis, and each point along X will represent one time step.\n",
    "\n",
    "The code below converts our collection of velocity states into a 2D array, repeating individual time steps 8 times to make the image a bit wider. This is purely optional, of course, but makes it easier to see what's happening in our Burgers simulation."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "def show_state(a, title):\n",
    "    # we only have 33 time steps, blow up by a factor of 2^4 to make it easier to see\n",
    "    # (could also be done with more evaluations of network)\n",
    "    a=np.expand_dims(a, axis=2)\n",
    "    for i in range(4):\n",
    "        a = np.concatenate( [a,a] , axis=2)\n",
    "\n",
    "    a = np.reshape( a, [a.shape[0],a.shape[1]*a.shape[2]] )\n",
    "    #print(\"Resulting image size\" +format(a.shape))\n",
    "\n",
    "    fig, axes = pylab.subplots(1, 1, figsize=(16, 5))\n",
    "    im = axes.imshow(a, origin='upper', cmap='inferno')\n",
    "    pylab.colorbar(im) ; pylab.xlabel('time'); pylab.ylabel('x'); pylab.title(title)\n",
    "        \n",
    "vels_img = np.asarray( np.concatenate(vels, axis=-1), dtype=np.float32 ) \n",
    "\n",
    "# save for comparison with reconstructions later on\n",
    "import os; os.makedirs(\"./temp\",exist_ok=True)\n",
    "np.savez_compressed(\"./temp/burgers-groundtruth-solution.npz\", np.reshape(vels_img,[N,STEPS+1])) # remove batch & channel dimension\n",
    "\n",
    "show_state(vels_img, \"Velocity\")"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ],
      "image/png": "[encoded data, 32764 characters]"
     },
     "metadata": {
      "needs_background": "light"
     }
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "This concludes a first simulation in phiflow. It's not overly complex, but because of that it's a good starting point for evaluating and comparing different physics-based deep learning approaches in the next chapter. But before that, we'll target a more complex simulation type in the next section."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Next steps\n",
    "\n",
    "Some things to try based on this simulation setup:\n",
    "\n",
    "- Feel free to experiment - the setup above is very simple, you can change the simulation parameters, or the initialization. E.g., you can use a noise field via `Noise()` to get more chaotic results (cf. the comment in the `velocity` cell above).\n",
    "\n",
    "- A bit more complicated: extend the simulation to 2D (or higher). This will require changes throughout, but all operators above support higher dimensions. Before trying this, you probably will want to check out the next example, which covers a 2D Navier-Stokes case."
   ],
   "metadata": {}
  }
 ],
 "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
}