{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Molecular dynamics \n",
    "\n",
    "* *simple visualization* of live MD simulation visualisation \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run if you need extra packages: numba and scipy\n",
    "!pip install scipy numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "\n",
    "import numpy as np\n",
    "from MD_utils import lennardjones, simple_molecule_vis\n",
    "from scipy.constants import N_A\n",
    "epsilon = 0.9959\n",
    "sigma = 0.3405\n",
    "k_B = 0.008311\n",
    "M = 39.948"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### initial condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt = simple_molecule_vis(bs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot.display()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy  as np\n",
    "# bs jest rozmiarem kostki w której znajdują się cząstki.  \n",
    "nx = 8\n",
    "dx = sigma *1.2\n",
    "bs = dx*nx \n",
    "box = np.array([bs,bs,bs])\n",
    "try:\n",
    "    plt.update_box(bs=bs)\n",
    "    print(\"no plot object\")\n",
    "except:\n",
    "    pass\n",
    "\n",
    "Nparticles = nx**3\n",
    "print(Nparticles)\n",
    "\n",
    "U = np.zeros((3,Nparticles))\n",
    "l = 0\n",
    "for i in range(nx):\n",
    "    for j in range(nx):\n",
    "        for k in range(nx):\n",
    "            U[:,l] = (dx*i-dx*(nx/2-0.50),dx*j-dx*(nx/2-0.5),dx*k-dx*(nx/2-0.5))\n",
    "            l+=1\n",
    "V = np.zeros_like(U)\n",
    "plt.pkts.positions = U.T.copy()\n",
    "plt.pkts.point_size = .3\n",
    "plt.pkts.colors = 0xff0000*np.ones(Nparticles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "box = np.array([bs,bs,bs])\n",
    "%time Epot, F, Vir = lennardjones(U, box,sigma = 0.3405, epsilon=0.9959)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = 0.1*(np.random.randn(3,Nparticles)-0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Start simulation\n",
    "\n",
    "import time \n",
    "n_steps = 1520\n",
    "dt = 0.005\n",
    "N = Nparticles\n",
    "\n",
    "box = np.array([bs,bs,bs])\n",
    "plt.update_box(bs=bs)\n",
    "\n",
    "(epot,F,Vir) = lennardjones(U,box,sigma = 0.3405, epsilon=0.9959)\n",
    "traj = []\n",
    "for i in range(n_steps):\n",
    "    U += V * dt + 0.5 * F/M * dt * dt\n",
    "    F0 = F[:]\n",
    "\n",
    "    (epot, F, Vir) = lennardjones(U, box,sigma = 0.3405, epsilon=0.9959)\n",
    "    \n",
    "    V += 0.5 * (F + F0)/M * dt\n",
    "    U -= bs*np.rint(U/bs)\n",
    "    \n",
    "    traj.append(U[:,233].copy())\n",
    "    \n",
    "    if i%10==0:\n",
    "        T = M*np.sum(V**2)/(k_B*(3*N-6))  \n",
    "        P = 1/bs**3*(  N*k_B*T + 1/3.*Vir )\n",
    "        Ek = 0.5*M*np.sum(V**2)\n",
    "        plt.pkts.positions = U.T.copy()\n",
    "        plt.pkts.colors = (np.sum((V**2),axis=0)/10.0* 0xffffff).astype(np.int64)\n",
    "        plt.update_box(bs=bs)\n",
    "        clear_output(wait=True)\n",
    "     \n",
    "        print(i,epot,Ek,epot+Ek,\"T=\",T,\"P=\",P)\n",
    "        print(\"Vir:\",Vir)\n",
    "        print(\"bs:\",bs)\n",
    "    \n",
    "traj = np.array(traj)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The trajectory of a single atom "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import k3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_traj = k3d.points(traj.copy(), color=0xffff00, point_size=.05)\n",
    "plt.plot += plt_traj"
   ]
  },
  {
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}