{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple Forward Simulation of Burgers Equation with Φ<sub>Flow</sub>\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, with Φ<sub>Flow</sub>.\n",
    "\n",
    "Burgers equations is a very simple, yet non-linear and non-trivial, model equation, and hence a very good starting point for experiments. It contains an advection term (motion / transport) and a diffusion term (dissipation due to the second law of thermodynamics). Burgers equation can already lead to interesting shock formations, and is given by:\n",
    "$\n",
    "    \\frac{\\partial u}{\\partial{t}} + u \\nabla u =\n",
    "    \\nu \\nabla\\cdot \\nabla u\n",
    "$ \n",
    "\n",
    "Let's get some preliminaries out of the way: we'll import the Φ<sub>Flow</sub> library, and define our simulation domain with $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 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`, which we'll use to initialize the velocity $u$ in the simulation in the next cell.\n",
    "\n",
    "**Note:** Below, the first command with a \"!\" prefix installs the [Φ<sub>Flow</sub> Python package from GitHub](https://github.com/tum-pbs/PhiFlow) via `pip` in your python environment. (You can skip or modify this command if necessary, depending on your system.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using phiflow version: 2.0.0rc0\n"
     ]
    }
   ],
   "source": [
    "#!pip install --upgrade --quiet git+https://github.com/tum-pbs/PhiFlow@develop\n",
    "#!pip install --upgrade --quiet phiflow\n",
    "\n",
    "from phi.flow import *\n",
    "import phi\n",
    "print(\"Using phiflow version: {}\".format(phi.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've loaded phiflow, we can start defining some globals and constants. The `initial` array contains the aforementioned sine wave that we'll use to produce a nice shock in the center of our domain.\n",
    "\n",
    "First we'll initialize some constants (denote by upper-case names):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 128\n",
    "STEPS = 32\n",
    "DT = 1./STEPS\n",
    "NU = 0.01/np.pi\n",
    "\n",
    "# initialization of velocities\n",
    "INITIAL = np.asarray( [-np.sin(np.pi * x) * 1. for x in np.linspace(-1,1,N)] ) # 1D array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we can define the `domain` to contain $\\Omega$, as outlined above periodically for $[-1,1]$. We also initialize a 1D `velocity` from the `initial` array.\n",
    "\n",
    "Just to illustrate, we'll also print some info about the velocity object: it's a `phiflow.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 the first five entries via the conversion to a numpy array with `numpy()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Velocity tensor shape: (x=128)\n",
      "Velocity tensor type: <class 'phi.math._tensors.NativeTensor'>\n",
      "Velocity tensor content: [1.2246469e-16 4.9453720e-02 9.8786421e-02 1.4787737e-01 1.9660644e-01]\n"
     ]
    }
   ],
   "source": [
    "DOMAIN = Domain(x=N, boundaries=PERIODIC, bounds=Box[-1:1])\n",
    "velocity = DOMAIN.scalar_grid(INITIAL)\n",
    "#velocity = DOMAIN.grid(Noise(vector=1)) # random init instead of sine curve\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 content: \" + format( velocity.values.numpy()[0:5] ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the simulation\n",
    "\n",
    "Now we're ready to run the simulation itself. To compute the diffusion and advection components of our model equation we can simply call the existing `diffusion` and `semi_lagrangian` operators in phiflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New velocity content at t=1.0: [0.00274862 0.01272991 0.02360343 0.03478042 0.0460869 ]\n"
     ]
    }
   ],
   "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] ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we're actually collecting all time steps in the list `vels`. 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 shows 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. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7ef47d8160>]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 34844 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "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')\n",
    "fig.plot(np.linspace(-1,1,len(vels[10].flatten())), vels[10].flatten(), lw=2, color='green')\n",
    "fig.plot(np.linspace(-1,1,len(vels[20].flatten())), vels[20].flatten(), lw=2, color='cyan')\n",
    "fig.plot(np.linspace(-1,1,len(vels[32].flatten())), vels[32].flatten(), lw=2, color='purple')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 purely optional, of course, but makes it easier to see what's happening in our Burgers simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "resulting image size(128, 528)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31316 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 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",
    "    # (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) \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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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(vector=1)` to get more chaotic results.\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."
   ]
  },
  {
   "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
}
