{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import utils\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import logging\n",
    "# logging.basicConfig(level=logging.INFO)\n",
    "logging.basicConfig(level=logging.DEBUG)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# silence some debug messages\n",
    "logging.getLogger('matplotlib.font_manager').setLevel(logging.WARNING)\n",
    "logging.getLogger('matplotlib.ticker').setLevel(logging.WARNING)\n",
    "logging.getLogger('matplotlib.pyplot').setLevel(logging.WARNING)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path('data')\n",
    "# DATA_NAME = 'data.txt'\n",
    "DATA_NAME = 'data0.txt'\n",
    "# DATA_NAME = 'data1.txt'\n",
    "NBINS = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:utils.load:Loaded 1008 rows and 9 columns from data/data0.txt\n",
      "DEBUG:__main__:Fetched 1008 points, columns: ['ID', 'M', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'eps']\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]]\n",
    "\n",
    "particles = particles[::10, ...]\n",
    "# particles = particles[:5000, ...]\n",
    "# TODO: remove this\n",
    "\n",
    "particles = utils.remove_outliers(particles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 87124 characters]",
      "text/plain": [
       "<Figure size 640x480 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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29604 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_r = utils.r_distribution(particles)\n",
    "rbins = np.logspace(np.min(all_r), np.max(all_r), NBINS)\n",
    "rho, err = utils.density_distribution(rbins, particles, ret_error=True)\n",
    "rho_t = utils.model_density_distribution(rbins)\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "plt.title('Density distribution')\n",
    "# set the scale to log-log\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('r')\n",
    "plt.ylabel('rho')\n",
    "plt.errorbar(rbins, rho, yerr=err, fmt='o', label=\"Particle data\")\n",
    "plt.plot(rbins, rho_t, 'r-', label=\"Analytic model\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choice of units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set G = 1\n",
    "G = 1\n",
    "\n",
    "# Since we have an elliptical galaxy, we can use typical values\n",
    "M_TOT = 1e12 # Msol\n",
    "R_TOT = 100 # kpc\n",
    "\n",
    "# For this system the virial theorem gives the typical velocity\n",
    "V_MEAN = np.sqrt(G * M_TOT / R_TOT)\n",
    "# For units of time we use the dynamical time or crossing time\n",
    "T_MEAN = R_TOT / V_MEAN\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:utils.particles:Half mass radius: 0.0 for 0th particle of 101\n",
      "DEBUG:utils.particles:Number of particles within half mass radius: 0 of 101\n",
      "/home/remy/Documents/Uni/HS24/Computational Astrophysics/projects/nbody/utils/particles.py:73: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  rho = n_half_mass / (4/3 * np.pi * r_half_mass**3)\n",
      "INFO:__main__:Mean interparticle distance: nan\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Computing forces for 101 particles using spherical approximation\n",
      "/home/remy/Documents/Uni/HS24/Computational Astrophysics/projects/nbody/utils/forces.py:71: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  f = - m_current * m_enclosed / r_current**2\n",
      "DEBUG:utils.forces:Particle 0 done\n"
     ]
    }
   ],
   "source": [
    "### Direct N body force computation\n",
    "epsilon = utils.mean_interparticle_distance(particles)\n",
    "logger.info(f\"Mean interparticle distance: {epsilon}\")\n",
    "\n",
    "f_nsquare = utils.n_body_forces(particles, G, 0)\n",
    "f_nsquare_1e = utils.n_body_forces(particles, G, epsilon)\n",
    "f_nsquare_2e = utils.n_body_forces(particles, G, 2*epsilon)\n",
    "f_nsquare_5e = utils.n_body_forces(particles, G, 5*epsilon)\n",
    "f_nsquare_05e = utils.n_body_forces(particles, G, 0.5*epsilon)\n",
    "f_nsquare_01e = utils.n_body_forces(particles, G, 0.1*epsilon)\n",
    "f_nsquare_001e = utils.n_body_forces(particles, G, 0.01*epsilon)\n",
    "\n",
    "f_analytical = utils.analytical_forces(particles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/remy/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/matplotlib/axes/_axes.py:6973: RuntimeWarning: All-NaN slice encountered\n",
      "  xmin = min(xmin, np.nanmin(xi))\n",
      "/home/remy/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/matplotlib/axes/_axes.py:6974: RuntimeWarning: All-NaN slice encountered\n",
      "  xmax = max(xmax, np.nanmax(xi))\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "autodetected range of [nan, nan] is not finite",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 13\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# plot the distribution of the error\u001b[39;00m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# create 4 stacked histograms, sharing the same x axis\u001b[39;00m\n\u001b[1;32m     12\u001b[0m fig, ax \u001b[38;5;241m=\u001b[39m plt\u001b[38;5;241m.\u001b[39msubplots(\u001b[38;5;241m4\u001b[39m, sharex\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m---> 13\u001b[0m \u001b[43max\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhist\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdiff_mag\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbins\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mNBINS\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     14\u001b[0m ax[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mset_title(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mMagnitude of the force difference\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     15\u001b[0m ax[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mset_yscale(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlog\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "File \u001b[0;32m~/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/matplotlib/__init__.py:1476\u001b[0m, in \u001b[0;36m_preprocess_data.<locals>.inner\u001b[0;34m(ax, data, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1473\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m   1474\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(ax, \u001b[38;5;241m*\u001b[39margs, data\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m   1475\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m data \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m-> 1476\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1477\u001b[0m \u001b[43m            \u001b[49m\u001b[43max\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1478\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mmap\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43msanitize_sequence\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1479\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m{\u001b[49m\u001b[43mk\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43msanitize_sequence\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1481\u001b[0m     bound \u001b[38;5;241m=\u001b[39m new_sig\u001b[38;5;241m.\u001b[39mbind(ax, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m   1482\u001b[0m     auto_label \u001b[38;5;241m=\u001b[39m (bound\u001b[38;5;241m.\u001b[39marguments\u001b[38;5;241m.\u001b[39mget(label_namer)\n\u001b[1;32m   1483\u001b[0m                   \u001b[38;5;129;01mor\u001b[39;00m bound\u001b[38;5;241m.\u001b[39mkwargs\u001b[38;5;241m.\u001b[39mget(label_namer))\n",
      "File \u001b[0;32m~/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/matplotlib/axes/_axes.py:7001\u001b[0m, in \u001b[0;36mAxes.hist\u001b[0;34m(self, x, bins, range, density, weights, cumulative, bottom, histtype, align, orientation, rwidth, log, color, label, stacked, **kwargs)\u001b[0m\n\u001b[1;32m   6997\u001b[0m \u001b[38;5;66;03m# Loop through datasets\u001b[39;00m\n\u001b[1;32m   6998\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(nx):\n\u001b[1;32m   6999\u001b[0m     \u001b[38;5;66;03m# this will automatically overwrite bins,\u001b[39;00m\n\u001b[1;32m   7000\u001b[0m     \u001b[38;5;66;03m# so that each histogram uses the same bins\u001b[39;00m\n\u001b[0;32m-> 7001\u001b[0m     m, bins \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhistogram\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbins\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweights\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mw\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mhist_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   7002\u001b[0m     tops\u001b[38;5;241m.\u001b[39mappend(m)\n\u001b[1;32m   7003\u001b[0m tops \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(tops, \u001b[38;5;28mfloat\u001b[39m)  \u001b[38;5;66;03m# causes problems later if it's an int\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/numpy/lib/_histograms_impl.py:797\u001b[0m, in \u001b[0;36mhistogram\u001b[0;34m(a, bins, range, density, weights)\u001b[0m\n\u001b[1;32m    688\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    689\u001b[0m \u001b[38;5;124;03mCompute the histogram of a dataset.\u001b[39;00m\n\u001b[1;32m    690\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    793\u001b[0m \n\u001b[1;32m    794\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    795\u001b[0m a, weights \u001b[38;5;241m=\u001b[39m _ravel_and_check_weights(a, weights)\n\u001b[0;32m--> 797\u001b[0m bin_edges, uniform_bins \u001b[38;5;241m=\u001b[39m \u001b[43m_get_bin_edges\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbins\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mrange\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweights\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    799\u001b[0m \u001b[38;5;66;03m# Histogram is an integer or a float array depending on the weights.\u001b[39;00m\n\u001b[1;32m    800\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m weights \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/numpy/lib/_histograms_impl.py:430\u001b[0m, in \u001b[0;36m_get_bin_edges\u001b[0;34m(a, bins, range, weights)\u001b[0m\n\u001b[1;32m    427\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m n_equal_bins \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    428\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m`bins` must be positive, when an integer\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m--> 430\u001b[0m     first_edge, last_edge \u001b[38;5;241m=\u001b[39m \u001b[43m_get_outer_edges\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mrange\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    432\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m np\u001b[38;5;241m.\u001b[39mndim(bins) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    433\u001b[0m     bin_edges \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39masarray(bins)\n",
      "File \u001b[0;32m~/.local/share/virtualenvs/projects-X-9bmgL6/lib/python3.12/site-packages/numpy/lib/_histograms_impl.py:323\u001b[0m, in \u001b[0;36m_get_outer_edges\u001b[0;34m(a, range)\u001b[0m\n\u001b[1;32m    321\u001b[0m     first_edge, last_edge \u001b[38;5;241m=\u001b[39m a\u001b[38;5;241m.\u001b[39mmin(), a\u001b[38;5;241m.\u001b[39mmax()\n\u001b[1;32m    322\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (np\u001b[38;5;241m.\u001b[39misfinite(first_edge) \u001b[38;5;129;01mand\u001b[39;00m np\u001b[38;5;241m.\u001b[39misfinite(last_edge)):\n\u001b[0;32m--> 323\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    324\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mautodetected range of [\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m, \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m] is not finite\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(first_edge, last_edge))\n\u001b[1;32m    326\u001b[0m \u001b[38;5;66;03m# expand empty range to avoid divide by zero\u001b[39;00m\n\u001b[1;32m    327\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m first_edge \u001b[38;5;241m==\u001b[39m last_edge:\n",
      "\u001b[0;31mValueError\u001b[0m: autodetected range of [nan, nan] is not finite"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 14180 characters]",
      "text/plain": [
       "<Figure size 640x480 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## compare the two force calculations\n",
    "# since the forces were computed for each particle, rather than comparing them directly we compare the relative error in the magnitude and direction of the forces\n",
    "\n",
    "\n",
    "# f_diff = f_nsquare_1e - f_analytical\n",
    "f_diff = f_nsquare_2e - f_analytical\n",
    "diff_mag = np.linalg.norm(f_diff, axis=1)\n",
    "\n",
    "\n",
    "# plot the distribution of the error\n",
    "# create 4 stacked histograms, sharing the same x axis\n",
    "fig, ax = plt.subplots(4, sharex=True)\n",
    "ax[0].hist(diff_mag, bins=NBINS)\n",
    "ax[0].set_title('Magnitude of the force difference')\n",
    "ax[0].set_yscale('log')\n",
    "\n",
    "ax[1].hist(f_diff[:,0], bins=NBINS)\n",
    "ax[1].set_title('X component of the force difference')\n",
    "ax[1].set_yscale('log')\n",
    "\n",
    "ax[2].hist(f_diff[:,1], bins=NBINS)\n",
    "ax[2].set_title('Y component of the force difference')\n",
    "ax[2].set_yscale('log')\n",
    "\n",
    "ax[3].hist(f_diff[:,2], bins=NBINS)\n",
    "ax[3].set_title('Z component of the force difference')\n",
    "ax[3].set_yscale('log')\n",
    "\n",
    "plt.title('Error in forces')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32248 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Plot the radial force behaviour\n",
    "r = np.linalg.norm(particles[:,:3], axis=1)\n",
    "r_sort = np.argsort(r)\n",
    "# r = r[r_sort]\n",
    "# f_nsquare_1e = f_nsquare_1e[r_sort]\n",
    "# f_nsquare_2e = f_nsquare_2e[r_sort]\n",
    "# f_analytical = f_analytical[r_sort]\n",
    "\n",
    "plt.figure()\n",
    "plt.title('Radial force behaviour')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('r')\n",
    "plt.ylabel('f')\n",
    "\n",
    "plt.plot(r, np.linalg.norm(f_analytical, axis=1), 'o', label=\"Analytical\")\n",
    "\n",
    "plt.plot(r, np.linalg.norm(f_nsquare, axis=1), 'o', label=\"N^2\")\n",
    "plt.plot(r, np.linalg.norm(f_nsquare_1e, axis=1), 'o', label=\"N^2 - 1epsilon\")\n",
    "plt.plot(r, np.linalg.norm(f_nsquare_2e, axis=1), 'o', label=\"N^2 - 2epsilon\")\n",
    "plt.plot(r, np.linalg.norm(f_nsquare_5e, axis=1), 'o', label=\"N^2 - 5epsilon\")\n",
    "# plt.plot(r, np.linalg.norm(f_nsquare_05e, axis=1), 'o', label=\"N^2 - 0.5epsilon\")\n",
    "# plt.plot(r, np.linalg.norm(f_nsquare_01e, axis=1), 'o', label=\"N^2 - 0.1epsilon\")\n",
    "# plt.plot(r, np.linalg.norm(f_nsquare_001e, axis=1), 'o', label=\"N^2 - 0.01epsilon\")\n",
    "\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task 2 - Particle mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:utils.integrate:Reshaped (1008, 7) to y0 with (6048,) and masses with (1008,)\n",
      "DEBUG:utils.integrate:y with shape (6048,)\n",
      "DEBUG:utils.integrate:Unstacked y into x with shape (1008, 3) and v with shape (1008, 3)\n",
      "DEBUG:utils.forces:Computing forces for 1008 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:utils.forces:Particle 1000 done\n",
      "DEBUG:utils.integrate:y with shape (6048,)\n",
      "DEBUG:utils.integrate:Unstacked y into x with shape (1008, 3) and v with shape (1008, 3)\n",
      "DEBUG:utils.forces:Computing forces for 1008 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Particle 1000 done\n",
      "DEBUG:utils.integrate:y with shape (6048,)\n",
      "DEBUG:utils.integrate:Unstacked y into x with shape (1008, 3) and v with shape (1008, 3)\n",
      "DEBUG:utils.forces:Computing forces for 1008 particles using n^2 algorithm\n",
      "DEBUG:utils.forces:Particle 0 done\n",
      "DEBUG:utils.forces:Particle 1000 done\n",
      "/tmp/ipykernel_36731/3498856150.py:13: ODEintWarning: Excess accuracy requested (tolerances too small). Run with full_output = 1 to get quantitative information.\n",
      "  sol = spi.odeint(y_prime, y0, np.arange(0, n_steps*dt, dt), atol=1e3,rtol=1e2)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " lsoda--  at t (=r1), too much accuracy requested  \u0000\u0000\n",
      "       for precision of machine..  see tolsf (=r2) \u0000\u0000\n",
      "      in above,  r1 =  0.1000000000000D+00   r2 =                  NaN\n"
     ]
    }
   ],
   "source": [
    "# load the particles in the format [x, y, z, vx, vy, vz, mass]\n",
    "particles = points[:, [2, 3, 4, 5, 6, 7, 1]]\n",
    "\n",
    "# we can now integrate this system for a few steps\n",
    "dt = 0.1\n",
    "n_steps = 10\n",
    "force_function = lambda x: utils.n_body_forces(x, G, epsilon)\n",
    "y0, y_prime = utils.ode_setup(particles, force_function)\n",
    "\n",
    "import scipy.integrate as spi\n",
    "\n",
    "# integrate the system\n",
    "sol = spi.odeint(y_prime, y0, np.arange(0, n_steps*dt, dt), atol=1e3,rtol=1e2)\n"
   ]
  }
 ],
 "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
