{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simulation import Lattice, Geometry\n",
    "from symbolic.generator import LBM\n",
    "import symbolic.D2Q9 as D2Q9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "lbm = LBM(D2Q9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cavity(geometry, x, y):\n",
    "    if x == 1 or y == 1 or x == geometry.size_x-2:\n",
    "        return 2\n",
    "    elif y == geometry.size_y-2:\n",
    "        return 3\n",
    "    else:\n",
    "        return 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "boundary = \"\"\"\n",
    "    if ( m == 2 ) {\n",
    "        u_0 = 0.0;\n",
    "        u_1 = 0.0;\n",
    "    }\n",
    "    if ( m == 3 ) {\n",
    "        u_0 = 0.1;\n",
    "        u_1 = 0.0;\n",
    "    }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_moment_plot(lattice, m):\n",
    "    velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_size())))\n",
    "    for x, y in lattice.geometry.inner_cells():\n",
    "        velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.gid(x,y)]**2 + m[2,lattice.gid(x,y)]**2)\n",
    "    plt.figure(figsize=(8,8))\n",
    "    plt.imshow(velocity, origin='lower', vmin=0.0, cmap=plt.get_cmap('seismic'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(nX=64, nY=64, nSteps=1000, tau=0.6):\n",
    "    lattice = Lattice(\n",
    "        descriptor = D2Q9,\n",
    "        geometry   = Geometry(nX, nY),\n",
    "\n",
    "        moments = lbm.moments(optimize = False),\n",
    "        collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = tau),\n",
    "\n",
    "        boundary_src = boundary)\n",
    "    lattice.setup_geometry(cavity)\n",
    "    \n",
    "    for i in range(0,nSteps):\n",
    "        lattice.evolve()\n",
    "    generate_moment_plot(lattice, lattice.get_moments())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "from ipywidgets import interactive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba655c48ce2b4f4898ca9fcb61e693a5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=64, description='nX', max=1024, min=32, step=32), IntSlider(value=64, de…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "interactive(test, nX=(32,1024,32), nY=(32,1024,32), nSteps=(0,100000,500), tau=(0.515,1.0,0.01))"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}