{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# automatically reflect changes in imported modules\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.units as u\n",
    "\n",
    "import utils\n",
    "import utils.logging_config\n",
    "import logging\n",
    "logger = logging.getLogger(\"task2 (mesh)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path('data')\n",
    "DATA_NAME = 'data0.txt'\n",
    "# DATA_NAME = 'data1.txt'\n",
    "NBINS = 30\n",
    "CACHE_ROOT = Path('.cache')\n",
    "\n",
    "G = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17:02:59 - utils.load - Loaded 1008 rows and 9 columns from data/data0.txt\n"
     ]
    }
   ],
   "source": [
    "points, columns = utils.load_data(DATA_ROOT / DATA_NAME)\n",
    "logger.debug(f\"Fetched {points.shape[0]} points, columns: {columns}\")\n",
    "# reorder the columns to match the expected order (x, y, z, mass)\n",
    "particles = points[:, [2, 3, 4, 1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 236832 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot the distribution of the particles\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.scatter(particles[:,0], particles[:,1], particles[:,2], cmap='viridis', c=particles[:,3])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17:03:00 - utils.particles - Found mean interparticle distance: 0.05692997700471936\n"
     ]
    }
   ],
   "source": [
    "### Direct N body force computation\n",
    "epsilon = utils.mean_interparticle_distance(particles)\n",
    "\n",
    "epsilon_range = np.logspace(-2, 2, 5)\n",
    "n_squared_forces = []\n",
    "\n",
    "SAVE_FORCES = False\n",
    "\n",
    "for e in epsilon_range:\n",
    "    n_particles = particles.shape[0]\n",
    "    cache_file = CACHE_ROOT / f\"n_squared_forces__n_{n_particles}__softening_multiplier_{e:.0f}.npy\"\n",
    "    try:\n",
    "        f = np.load(cache_file)\n",
    "    except FileNotFoundError:\n",
    "        f = utils.n_body_forces(particles, G, e * epsilon)\n",
    "        if SAVE_FORCES:\n",
    "            np.save(cache_file, f)\n",
    "            logger.debug(f\"Saved forces to {cache_file}\")\n",
    "    n_squared_forces.append(f)\n",
    "\n",
    "### Mesh based force computation\n",
    "mesh_size_range = [10, 20, 50, 100, 150, 200]\n",
    "mapping = utils.particle_to_cells_nn\n",
    "mesh_forces = []\n",
    "for mesh_size in mesh_size_range:\n",
    "    cache_file = CACHE_ROOT / f\"mesh_forces__n_{n_particles}__mesh_size_{mesh_size}__mapping_{mapping.__name__}.npy\"\n",
    "    try:\n",
    "        f = np.load(cache_file)\n",
    "    except FileNotFoundError:\n",
    "        f = utils.mesh_forces_v2(particles, G, mesh_size, mapping)\n",
    "        if SAVE_FORCES:\n",
    "            np.save(cache_file, f)\n",
    "            logger.debug(f\"Saved forces to {cache_file}\")\n",
    "    mesh_forces.append(f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 83636 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Compare the mesh computation with the direct summation\n",
    "r = np.linalg.norm(particles[:,:3], axis=1)\n",
    "\n",
    "plt.figure()\n",
    "plt.title('Radial force')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('$r$')\n",
    "plt.ylabel('$F_r(r)$')\n",
    "\n",
    "# many of the particles have the same distance from the origin, so we skip some of them\n",
    "SKIP_N = 20\n",
    "\n",
    "for f,e in zip(n_squared_forces, epsilon_range):\n",
    "    plt.plot(r[::SKIP_N], np.linalg.norm(f, axis=1)[::SKIP_N], 'o', label=f\"$N^2$ - {e:.1g} * $\\\\epsilon$\", alpha=0.3)\n",
    "for f, s in zip(mesh_forces, mesh_size_range):\n",
    "    plt.plot(r[::SKIP_N], np.linalg.norm(f, axis=1)[::SKIP_N], '.', label=f\"Mesh - N={s}\")\n",
    "\n",
    "plt.ylim([5e-4, 1e2])\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# TODO: compare computation time\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion\n",
    "- Using the baseline of N^2 and 1 epsilon softening we can see that already a 20 x 20 grid provides good accuracy but the mapping breaks down at small distances (dip)\n",
    "- Larger grids are more stable, especially at small distances => 50 x 50 x 50 already seems to be a good choice\n",
    "- very large grids show overdiscretization\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.integrate as spi\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17:03:01 - task2 (mesh) - Considering 1008 particles\n",
      "17:03:01 - task2 (mesh) - Total mass: 11.0\n"
     ]
    }
   ],
   "source": [
    "# load the particles in the format [x, y, z, vx, vy, vz, mass]\n",
    "p0 = points[:, [2, 3, 4, 5, 6, 7, 1]]\n",
    "# p0 = p0[::10] # only take a subset for now\n",
    "\n",
    "logger.info(f\"Considering {p0.shape[0]} particles\")\n",
    "logger.info(f\"Total mass: {np.sum(p0[:,6])}\")\n",
    "\n",
    "if logger.level >= logging.DEBUG:\n",
    "    # assert that the ODE reshaping is consistent\n",
    "    logger.debug(f\"{p0[0]}, {p0[1]}\")\n",
    "    y0, _ = utils.ode_setup(p0, None)\n",
    "    logger.debug(y0[0:7])\n",
    "    p0_reconstructed = utils.to_particles(y0)\n",
    "    logger.debug(p0_reconstructed[0])\n",
    "    assert np.allclose(p0, p0_reconstructed)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17:17:03 - task2 (mesh) - Integration done, shape: (50, 7056)\n"
     ]
    }
   ],
   "source": [
    "# we can now integrate this system for a few steps\n",
    "dt = 0.01\n",
    "n_steps = 50\n",
    "t_range = np.arange(0, n_steps*dt, dt)\n",
    "\n",
    "## The force function can be interchanged\n",
    "force_function = lambda x: utils.n_body_forces(x, G, epsilon) # epsilon was computed above\n",
    "# force_function = lambda x: utils.mesh_forces_v2(x, G, 50, utils.particle_to_cells_nn)\n",
    "\n",
    "y0, y_prime = utils.ode_setup(p0, force_function)\n",
    "\n",
    "\n",
    "# integrate the system\n",
    "sol = spi.odeint(y_prime, y0, t_range, rtol=1e-3)\n",
    "logger.info(f\"Integration done, shape: {sol.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 164620 characters]",
      "text/plain": [
       "<Figure size 1800x1200 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "### Plot a fixed number of states\n",
    "N_STATES = 10\n",
    "SKIP_NTH_PARTICLE = 10\n",
    "fig, axs = plt.subplots(1, N_STATES, subplot_kw={'projection': '3d'})\n",
    "for i, ax in enumerate(axs.flat):\n",
    "    p = utils.to_particles(sol[i])[::SKIP_NTH_PARTICLE]\n",
    "    ax.scatter(p[:,0], p[:,1], p[:,2], cmap='viridis', c=range(p.shape[0]))\n",
    "    ax.set_title(f\"t={t_range[i]:.2f}\")\n",
    "# set size\n",
    "fig.set_size_inches(18, 12)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# # Also show the velocities\n",
    "# scale = 1e-2\n",
    "# fig = plt.figure()\n",
    "# ax = fig.add_subplot(111, projection='3d')\n",
    "# p_reduced = plast[::10]\n",
    "# ax.quiver(p_reduced[:,0], p_reduced[:,1], p_reduced[:,2], p_reduced[:,3], p_reduced[:,4], p_reduced[:,5], length=0.1, normalize=True)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38752 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#### Plot some key quantities of the system as a whole\n",
    "# sol has the shape (n_steps, n_particles*6) where the first 3*n are the positions and the last 3*n are the velocities\n",
    "\n",
    "# kinetic energy\n",
    "energies = np.zeros(sol.shape[0])\n",
    "for i in range(sol.shape[0]):\n",
    "    p = utils.to_particles(sol[i])\n",
    "    k_e = 0.5 * np.sum(p[:,6] * np.linalg.norm(p[:,3:6], axis=1)**2)\n",
    "    energies[i] = k_e\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(t_range, energies)\n",
    "plt.title('Kinetic energy')\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Energy')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# to"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "projects-X-9bmgL6",
   "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.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
