{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "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(\"task1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path('data')\n",
    "DATA_NAME = 'data.txt'\n",
    "NBINS = 30\n",
    "CACHE_ROOT = Path('.cache')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12:58:01 - utils.load - Loaded 50010 rows and 10 columns from data/data.txt\n",
      "12:58:01 - task1 - Fetched 50010 points, columns: ['ID', 'MASS', 'X', 'Y', 'Z', 'VX', 'VY', 'VZ', 'SOFTENING', '?']\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",
    "# particles = utils.remove_outliers(particles)\n",
    "# logger.warning(f\"Reduced selection to {particles.shape[0]} particles\")\n",
    "# TODO: make sure this is not executed in the final version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 87048 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": "markdown",
   "metadata": {},
   "source": [
    "### Choice of units\n",
    "Recap of the particle properties:\n",
    "- $\\sim 10^4$ particles\n",
    "- spherical distribution\n",
    "\n",
    "$\\implies$ this seems to be a **globular cluster**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12:58:03 - utils.units - Set scales: M_SCALE = 0.021634629254180957 solMass, R_SCALE = 0.02759802045521005 pc\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4622219.258999999\n"
     ]
    }
   ],
   "source": [
    "# Set G = 1\n",
    "G = 1\n",
    "\n",
    "\n",
    "# Since we have a globular cluster, we can use typical values\n",
    "M_TOT = 1e5 * u.M_sun\n",
    "R_TOT = 20 * u.pc\n",
    "\n",
    "# Rescale the units of the particles\n",
    "M_particles = particles[:,3].sum()\n",
    "print(M_particles.sum())\n",
    "R_particles = np.max(np.linalg.norm(particles[:, :3], axis=1))\n",
    "m_scale = M_TOT / M_particles\n",
    "r_scale = R_TOT / R_particles\n",
    "utils.seed_scales(r_scale, m_scale)\n",
    "\n",
    "\n",
    "\n",
    "# ### Plot again with scales\n",
    "# positions = utils.apply_units(particles[:, :3], \"position\")\n",
    "# masses = utils.apply_units(particles[:, 3], \"mass\")\n",
    "# fig = plt.figure()\n",
    "# ax = fig.add_subplot(111, projection='3d')\n",
    "# sc = ax.scatter(positions[:,0], positions[:,1], positions[:,2], cmap='viridis', c=masses)\n",
    "#\n",
    "# ax.set_xlabel(f'x [{positions.unit}]')\n",
    "# ax.set_ylabel(f'y [{positions.unit}]')\n",
    "# ax.set_zlabel(f'z [{positions.unit}]')\n",
    "# cbar = plt.colorbar(sc, ax=ax, pad=0.1)\n",
    "# cbar.set_label(f'Mass [{masses.unit}]')\n",
    "#\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12:58:03 - utils.particles - Half mass radius: 0.1887071270700712 for 25004th particle of 50010\n",
      "12:58:03 - task1 - Scale length: a=0.078, total mass: m=4.6e+06\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34368 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.log10(np.min(all_r)), np.log10(np.max(all_r)), NBINS, endpoint=True)\n",
    "\n",
    "## Using the Hernquist model with parameters M and a\n",
    "m = M_particles\n",
    "a = utils.half_mass_radius(particles) / (1 + np.sqrt(2))\n",
    "logger.debug(f\"Scale length: {a=:.2g}, total mass: {m=:.2g}\")\n",
    "rho_t = utils.model_density_distribution(rbins, m, a)\n",
    "\n",
    "## Numerical density distribution\n",
    "rho, err = utils.density_distribution(rbins, particles, ret_error=True)\n",
    "\n",
    "## Add the units\n",
    "rho = utils.apply_units(rho, \"density\")\n",
    "err = utils.apply_units(err, \"density\")\n",
    "rho_t = utils.apply_units(rho_t, \"density\")\n",
    "rbins = utils.apply_units(rbins, \"position\")\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "plt.title('Density distribution')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel(f'r [{rbins.unit}]')\n",
    "plt.ylabel(f'$\\\\rho$ [{rho.unit}]')\n",
    "plt.errorbar(rbins.value, rho.value, yerr=err.value, fmt='o', label=\"Particle data\")\n",
    "plt.plot(rbins.value, rho_t.value, 'o', label=\"Analytic model\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "## Note that by construction, the first shell contains no particles\n",
    "## => the numerical density is zero there\n",
    "## Having more bins means to have shells that are nearly empty\n",
    "## => the error is large, NBINS = 30 is a good compromise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "13:14:03 - utils.particles - Half mass radius: 0.1887071270700712 for 25004th particle of 50010\n",
      "13:14:03 - utils.particles - Number of particles within half mass radius: 25004 of 50010\n",
      "13:14:03 - utils.particles - Found mean interparticle distance: 0.010402746349924056\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "13:14:03 - utils.forces_basic - Computing forces for 50010 particles using spherical approximation\n",
      "13:14:03 - utils.forces_basic - Particle 0 done\n",
      "13:14:07 - utils.forces_basic - Particle 5000 done\n",
      "13:14:11 - utils.forces_basic - Particle 10000 done\n",
      "13:14:15 - utils.forces_basic - Particle 15000 done\n",
      "13:14:18 - utils.forces_basic - Particle 20000 done\n",
      "13:14:22 - utils.forces_basic - Particle 25000 done\n",
      "13:14:26 - utils.forces_basic - Particle 30000 done\n",
      "13:14:30 - utils.forces_basic - Particle 35000 done\n",
      "13:14:34 - utils.forces_basic - Particle 40000 done\n",
      "13:14:38 - utils.forces_basic - Particle 45000 done\n",
      "13:14:41 - utils.forces_basic - Particle 50000 done\n"
     ]
    }
   ],
   "source": [
    "### Direct N body force computation and comparison with analytical model\n",
    "epsilon = utils.mean_interparticle_distance(particles)\n",
    "\n",
    "\n",
    "# TODO: compute this in a threaded way\n",
    "epsilon_range = np.logspace(-3, 3, 7)\n",
    "n_squared_forces = []\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",
    "        np.save(cache_file, f)\n",
    "        logger.debug(f\"Saved forces to {cache_file}\")\n",
    "    n_squared_forces.append(f)\n",
    "\n",
    "analytical_force = utils.analytical_forces(particles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"## 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\\nf_diff = f_nsquare_2e - f_analytical\\ndiff_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\\nfig, ax = plt.subplots(4, sharex=True)\\nax[0].hist(diff_mag, bins=NBINS)\\nax[0].set_title('Magnitude of the force difference')\\nax[0].set_yscale('log')\\n\\nax[1].hist(f_diff[:,0], bins=NBINS)\\nax[1].set_title('X component of the force difference')\\nax[1].set_yscale('log')\\n\\nax[2].hist(f_diff[:,1], bins=NBINS)\\nax[2].set_title('Y component of the force difference')\\nax[2].set_yscale('log')\\n\\nax[3].hist(f_diff[:,2], bins=NBINS)\\nax[3].set_title('Z component of the force difference')\\nax[3].set_yscale('log')\\n\\nplt.title('Error in forces')\\nplt.show()\\n\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61300 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "SKIP_N = 20\n",
    "## Plot the radial force behaviour\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",
    "plt.plot(r[::SKIP_N], np.linalg.norm(analytical_force[::SKIP_N], axis=1), '.', label=\"Analytical\", zorder=10, alpha=0.5)\n",
    "\n",
    "for e, f in zip(epsilon_range, n_squared_forces):\n",
    "    r_red = r[::SKIP_N]\n",
    "    f_red = f[::SKIP_N]\n",
    "    plt.plot(r, np.linalg.norm(f, axis=1), '.', label=f\"$N^2$ - {e:.1g} * $\\\\epsilon$\")\n",
    "\n",
    "\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion\n",
    "- basic $N^2$ matches analytical solution without dropoff. but: noisy data from \"bad\" samples\n",
    "- $N^2$ with softening matches analytical solution but has a dropoff. No noisy data.\n",
    "    - softening $\\approx 1 \\varepsilon$ is a sweet spot since the dropoff is \"late\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Relaxation time\n",
    "\n",
    "For an estimate of the crossing time we use the half mass radius and the circular velocity at the half mass radius.\n",
    "\n",
    "$\\implies$ the crossing time is of order radius / velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "13:10:27 - utils.particles - Half mass radius: 0.1887071270700712 for 25004th particle of 50010\n",
      "13:10:27 - task1 - Crossing time for half mass system: 5.392276839004401e-05\n",
      "13:10:27 - task1 - Direct estimate of the relaxation timescale: 0.025\n"
     ]
    }
   ],
   "source": [
    "## Computes the relaxation timescale of a set of particles using the velocity at the half mass radius.\n",
    "\n",
    "# enclosed mass at half mass radius\n",
    "m_half = np.sum(particles[:, 3]) / 2\n",
    "r_half = utils.half_mass_radius(particles)\n",
    "\n",
    "v_c = np.sqrt(G * m_half / r_half)\n",
    "\n",
    "t_c = r_half / v_c\n",
    "# TODO: what is the unit of the crossing time?\n",
    "# we expressed the mass in solar masses, the radius in pc and set G = 1\n",
    "# => ???\n",
    "logger.info(f\"Crossing time for half mass system: {t_c}\")\n",
    "\n",
    "\n",
    "## Using the derived formula for the relaxation timescale\n",
    "# We can directly compute the relaxation timescale through the estimate\n",
    "# t_relax = t_c * N / (10 * log(N))\n",
    "n = particles.shape[0]\n",
    "t_rel = t_c * n / (10 * np.log(n))\n",
    "logger.info(f\"Direct estimate of the relaxation timescale: {t_rel:.2g}\")\n",
    "\n",
    "\n",
    "## Manual computation of the relaxation timescale by summing up 2 body interactions\n",
    "# From the crossing time, we can estimate the number of interactions per crossing time\n",
    "# We then multiply by the number of particles to get the relaxation timescale\n",
    "n_interactions = n * np.log(n)\n",
    "# TODO\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion\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.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
