SPH-interactive.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "from numba import jit\n",
    "from numpy import linalg as LA\n",
    "\n",
    "# %matplotlib inline\n",
    "GRAVITY = 0.025\n",
    "RANGE = 10.0\n",
    "RANGE2 = RANGE * RANGE\n",
    "PRESSURE = 1.0\n",
    "VISCOSITY = 0.05\n",
    "DENSITY = 0.2\n",
    "N = 144"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def move(x, v, f):\n",
    "    v[:] += f[:]\n",
    "    v[:,1] -= GRAVITY\n",
    "    x[:] += v[:]\n",
    "    \n",
    "    cond = x < 5.0\n",
    "    v[cond] += (5.0 - x[cond]) * 0.5 - v[cond] * 0.5\n",
    "    cond = x > 460.0\n",
    "    v[cond] += (460.0 - x[cond]) * 0.5 - v[cond] * 0.5\n",
    "\n",
    "@jit\n",
    "def find_neighbors(x):\n",
    "    neighbors = []\n",
    "    \n",
    "    for i in range(N):\n",
    "        for j in range(N):\n",
    "            if i == j:\n",
    "                continue\n",
    "                \n",
    "            dist = (x[i,0] - x[j,0]) ** 2 + (x[i,1] - x[j,1]) ** 2\n",
    "            if dist < RANGE2:\n",
    "                neighbors.append((i, j))\n",
    "                \n",
    "    return neighbors\n",
    "\n",
    "@jit\n",
    "def calc_pressure(neighbors, x, pressure, rho):\n",
    "    for i, j in neighbors:\n",
    "        diff = x[i] - x[j]\n",
    "        distance = LA.norm(diff)\n",
    "        weight = 1.0 - distance / RANGE\n",
    "        temp = weight ** 3\n",
    "        rho[i] += temp\n",
    "        rho[j] += temp\n",
    "        \n",
    "    rho[rho < DENSITY] = DENSITY\n",
    "    pressure[:] = rho[:] - DENSITY\n",
    "\n",
    "@jit\n",
    "def calc_force(neighbors, x, f, pressure, rho):\n",
    "    for i, j in neighbors:\n",
    "        diff = x[i] - x[j]\n",
    "        distance = LA.norm(diff)\n",
    "        weight = 1.0 - distance / RANGE\n",
    "\n",
    "        p_weight = weight * (pressure[i] + pressure[j]) / (rho[i] + rho[j]) * PRESSURE\n",
    "        v_weight = weight / (rho[i] + rho[j]) * VISCOSITY\n",
    "\n",
    "        rv = v[j] - v[i]\n",
    "        n = diff / distance\n",
    "        f[i] += n * p_weight + rv * v_weight\n",
    "        f[j] -= n * p_weight - rv * v_weight\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# %%prun -s cumulative\n",
    "xi = np.linspace(50, 150, int(np.sqrt(N)))\n",
    "yi = np.linspace(50, 150, int(np.sqrt(N)))\n",
    "xx, yy = np.meshgrid(xi, yi)\n",
    "x = np.c_[xx.ravel(), yy.ravel()]\n",
    "\n",
    "v = np.zeros_like(x)\n",
    "f = np.zeros_like(x)\n",
    "rho = np.zeros(N)\n",
    "pressure = np.zeros(N)\n",
    "\n",
    "fig = plt.figure(figsize = (8, 8))\n",
    "plt.xlim(0, 465.0)\n",
    "plt.ylim(0, 465.0)\n",
    "\n",
    "anim = []\n",
    "for t in range(1200):\n",
    "    f[:] = rho[:] = 0.0\n",
    "    neighbors = find_neighbors(x)\n",
    "    calc_pressure(neighbors, x, pressure, rho)\n",
    "    calc_force(neighbors, x, f, pressure, rho)\n",
    "    move(x, v, f)\n",
    "    img = plt.plot(x[:,0], x[:,1], 'go')\n",
    "    anim.append(img)\n",
    "    \n",
    "plt.rcParams['animation.html'] = 'html5'\n",
    "a = animation.ArtistAnimation(fig, anim, interval=16)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}