{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 92,
   "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",
    "from astropy.visualization import quantity_support\n",
    "quantity_support()\n",
    "\n",
    "\n",
    "import utils\n",
    "import utils.logging_config\n",
    "utils.logging_config.set_log_level(\"info\")\n",
    "import logging\n",
    "logger = logging.getLogger(\"task1\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "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": 94,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12:40:29 - utils.load - Loaded 50010 rows and 10 columns from data/data.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]]\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": 95,
   "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": 113,
   "metadata": {
    "tags": [
     "plot_particle_distribution"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "18:12:23 - task1 - Considering a globular cluster - total mass of particles: 4622219.258999999, maximum radius of particles: 724.689657812915\n",
      "18:12:23 - utils.units - Set scales: M_SCALE = 0.022 solMass, R_SCALE = 0.028 pc\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 134656 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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",
    "R_particles = np.max(np.linalg.norm(particles[:, :3], axis=1))\n",
    "logger.info(f\"Considering a globular cluster - total mass of particles: {M_particles}, maximum radius of particles: {R_particles}\")\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",
    "reduced = utils.remove_outliers(particles.copy())\n",
    "positions = utils.apply_units(reduced[:, :3], \"position\")\n",
    "masses = utils.apply_units(reduced[:, 3], \"mass\")\n",
    "utils.particles_plot_3d(positions, masses, title=\"Particle distribution\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "tags": [
     "plot_density_distribution"
    ]
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32720 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(\"R\")\n",
    "plt.xlabel(f'r [{rbins.unit:latex}]')\n",
    "plt.ylabel(f'$\\\\rho$ [{rho.unit:latex}]')\n",
    "plt.errorbar(rbins, rho, yerr=err, fmt='o', label=\"Particle data\")\n",
    "plt.plot(rbins, rho_t, '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": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12:40:33 - utils.particles - Found mean interparticle distance: 0.010402746349924056\n"
     ]
    }
   ],
   "source": [
    "### Direct N body force computation and comparison with analytical model\n",
    "epsilon = utils.mean_interparticle_distance(particles)\n",
    "\n",
    "\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": 99,
   "metadata": {
    "tags": [
     "plot_force_radial"
    ]
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57884 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "SKIP_N = 1\n",
    "\n",
    "## Plot the radial force behaviour\n",
    "r = np.linalg.norm(particles[:,:3], axis=1)\n",
    "r = utils.apply_units(r, \"position\")\n",
    "\n",
    "plt.figure()\n",
    "f_analytical = np.linalg.norm(analytical_force, axis=1)\n",
    "f_analytical = utils.apply_units(f_analytical, \"force\")\n",
    "\n",
    "plt.plot(r[::SKIP_N], f_analytical[::SKIP_N], '.', label=\"Analytical\", zorder=10, alpha=0.5, color='black')\n",
    "\n",
    "colors = plt.cm.Blues(np.linspace(0.3,1,len(epsilon_range)))\n",
    "for e, f, c in zip(epsilon_range, n_squared_forces, colors):\n",
    "    # print((np.log(e)+10) * 10)\n",
    "    f_abs = np.linalg.norm(f, axis=1)\n",
    "    f_abs = utils.apply_units(f_abs, \"force\")\n",
    "    plt.plot(r[::SKIP_N], f_abs[::SKIP_N], '.', label=f\"$N^2$ - {e:.1g} * $\\\\epsilon$\", color=c)\n",
    "\n",
    "\n",
    "plt.title('Radial force')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel(f'$r$ [{r.unit:latex}]')\n",
    "plt.ylabel(f'$|F(r)|$ [{f_analytical.unit:latex}]')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nparticles = utils.remove_outliers(particles)\\nr = np.linalg.norm(particles[:,:3], axis=1)\\nr = utils.apply_units(r, \"position\")\\n\\nplt.figure()\\nf_nsquare = utils.analytical_forces(particles)\\nf_abs = np.linalg.norm(f_nsquare, axis=1)\\nf_abs = utils.apply_units(f_abs, \"force\")\\nplt.plot(r.value[::SKIP_N], f_abs.value[::SKIP_N], \\'.\\', label=\"N^2\")\\n\\nf_mesh = utils.mesh_forces(particles, G, 405, utils.particle_to_cells_nn)\\nf_abs = np.linalg.norm(f_mesh, axis=1)\\nf_abs = utils.apply_units(f_abs, \"force\")\\nplt.plot(r.value[::SKIP_N], f_abs.value[::SKIP_N], \\'.\\', label=\"Mesh\")\\n\\nplt.title(\\'Radial force\\')\\nplt.xscale(\\'log\\')\\nplt.yscale(\\'log\\')\\nplt.xlabel(f\\'$r$ [{r.unit}]\\')\\nplt.ylabel(f\\'$F_r(r) [{f_analytical.unit}]$\\')\\nplt.legend()\\nplt.show()\\n'"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## For fun, do this with the mesh solver as well:\n",
    "\"\"\"\n",
    "particles = utils.remove_outliers(particles)\n",
    "r = np.linalg.norm(particles[:,:3], axis=1)\n",
    "r = utils.apply_units(r, \"position\")\n",
    "\n",
    "plt.figure()\n",
    "f_nsquare = utils.analytical_forces(particles)\n",
    "f_abs = np.linalg.norm(f_nsquare, axis=1)\n",
    "f_abs = utils.apply_units(f_abs, \"force\")\n",
    "plt.plot(r.value[::SKIP_N], f_abs.value[::SKIP_N], '.', label=\"N^2\")\n",
    "\n",
    "f_mesh = utils.mesh_forces(particles, G, 405, utils.particle_to_cells_nn)\n",
    "f_abs = np.linalg.norm(f_mesh, axis=1)\n",
    "f_abs = utils.apply_units(f_abs, \"force\")\n",
    "plt.plot(r.value[::SKIP_N], f_abs.value[::SKIP_N], '.', label=\"Mesh\")\n",
    "\n",
    "plt.title('Radial force')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel(f'$r$ [{r.unit}]')\n",
    "plt.ylabel(f'$F_r(r) [{f_analytical.unit}]$')\n",
    "plt.legend()\n",
    "plt.show()\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": 101,
   "metadata": {
    "tags": [
     "compute_relaxation_time"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12:40:50 - task1 - Central velocity @ HM 3098.502740746268 solMass(1/2) / pc(1/2)\n",
      "12:40:50 - task1 - Crossing time for half mass system: 1.7e-06 pc(3/2) / solMass(1/2)\n",
      "12:40:50 - task1 - Direct estimate of the relaxation timescale: 0.00078 pc(3/2) / solMass(1/2)\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",
    "# set the units\n",
    "m_half = utils.apply_units(m_half, \"mass\")\n",
    "r_half = utils.apply_units(r_half, \"position\")\n",
    "\n",
    "v_c = np.sqrt(G * m_half / r_half)\n",
    "logger.info(f\"Central velocity @ HM {v_c}\")\n",
    "\n",
    "t_c = r_half / v_c\n",
    "logger.info(f\"Crossing time for half mass system: {t_c:.2g}\")\n",
    "# TODO: how to treat this unit?\n",
    "\n",
    "## Using the derived formula for the relaxation timescale\n",
    "# Compute the relaxation timescale through the estimate\n",
    "# t_relax = t_c *n_relax = t_c * N / (10 * log(N))\n",
    "n = particles.shape[0]\n",
    "n_relax = n / (10 * np.log(n))\n",
    "t_rel = t_c * n_relax\n",
    "logger.info(f\"Direct estimate of the relaxation timescale: {t_rel:.2g}\")\n"
   ]
  },
  {
   "attachments": {
    "Pasted image 20250118155501.png": {
     "image/png": "[encoded data, 19132 characters]"
    }
   },
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion\n",
    "\n",
    "\n",
    "The estimate for $n_{relax}$ comes from the contribution of each star-star encounter to the velocity dispersion. This depends on the perpendicular force\n",
    "\n",
    "![Pasted image 20250118155501.png](<attachment:Pasted image 20250118155501.png>)\n",
    "\n",
    "The full integration of the force along the path of the star reduces to the main contribution at the closest approach (distance $b$):\n",
    "$$\n",
    "\\delta v \\approx \\frac{2 G m}{b} \\to \\delta v \\approx \\frac{2 G m}{b + \\varepsilon}\n",
    "$$\n",
    "$\\implies$ a bigger softening length leads to a smaller $\\delta v$.\n",
    "\n",
    "Using $n_{relax} = \\frac{\\delta v^2}{v^2}$, and knowing that the value of $v^2$ is derived from the Virial theorem (i.e. unaffected by the softening length), we can see that $n_{relax}$ scales like $\\varepsilon^{-2}$.\n",
    "\n",
    "##### Effect\n",
    "- The relaxation time decreases with increasing softening length\n",
    "- From the integration over all impact parameters $b$ even $b_{min}$ is chosen to be larger than $\\varepsilon$ $\\implies$ expect only a small effect on the relaxation time\n",
    "\n",
    "#TODO : The softening dampens the change of velocity => time to relax is longer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "function_density_distribution"
    ]
   },
   "outputs": [],
   "source": [
    "print(inspect.getsource(utils.density_distribution))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {
    "tags": [
     "function_analytical_forces"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def analytical_forces(particles: np.ndarray):\n",
      "    \"\"\"\n",
      "    Computes the interparticle forces without computing the n^2 interactions.\n",
      "    This is done by using newton's second theorem for a spherical mass distribution.\n",
      "    The force on a particle at radius r is simply the force exerted by a point mass with the enclosed mass.\n",
      "    Assumes that the particles array has the following columns: x, y, z, m.\n",
      "    \"\"\"\n",
      "\n",
      "    n = particles.shape[0]\n",
      "    forces = np.zeros((n, 3))\n",
      "\n",
      "    logger.debug(f\"Computing forces for {n} particles using spherical approximation\")\n",
      "\n",
      "    r_particles = np.linalg.norm(particles[:, :3], axis=1)\n",
      "    for i in range(n):\n",
      "        r_current = np.linalg.norm(particles[i, 0:3])\n",
      "        m_current = particles[i, 3]\n",
      "\n",
      "        m_enclosed = np.sum(particles[r_particles < r_current, 3])\n",
      "\n",
      "        # the force is the same as the force exerted by a point mass at the center\n",
      "        f = - m_current * m_enclosed / r_current**2\n",
      "        forces[i] = f\n",
      "\n",
      "        if i % 5000 == 0:\n",
      "            logger.debug(f\"Particle {i} done\")\n",
      "\n",
      "    return forces\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(utils.analytical_forces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "tags": [
     "function_n2_forces"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def n_body_forces(particles: np.ndarray, G: float, softening: float = 0):\n",
      "    \"\"\"\n",
      "    Computes the gravitational forces between a set of particles.\n",
      "    Assumes that the particles array has the following columns: x, y, z, m.\n",
      "    \"\"\"\n",
      "    if particles.shape[1] != 4:\n",
      "        raise ValueError(\"Particles array must have 4 columns: x, y, z, m\")\n",
      "\n",
      "    x_vec = particles[:, 0:3]\n",
      "    masses = particles[:, 3]\n",
      "\n",
      "    n = particles.shape[0]\n",
      "    forces = np.zeros((n, 3))\n",
      "    logger.debug(f\"Computing forces for {n} particles using n^2 algorithm (using {softening=:.2g})\")\n",
      "\n",
      "    for i in range(n):\n",
      "        # the current particle is at x_current\n",
      "        x_current = x_vec[i, :]\n",
      "        m_current = masses[i]\n",
      "\n",
      "        # first compute the displacement to all other particles (and its magnitude)\n",
      "        r_vec = x_vec - x_current\n",
      "        r = np.linalg.norm(r_vec, axis=1)\n",
      "\n",
      "        # add softening to the denominator\n",
      "        r_adjusted = r**2 + softening**2\n",
      "        # usually with a square root: r' = sqrt(r^2 + softening^2) and then cubed, but we combine that below\n",
      "\n",
      "        # the numerator is tricky:\n",
      "        # m is a list of scalars and r_vec is a list of vectors (2D array)\n",
      "        # we only want row_wise multiplication\n",
      "        num = G * (masses * r_vec.T).T\n",
      "        # a zero value is expected where we have the same particle\n",
      "        r_adjusted[i] = 1\n",
      "        num[i] = 0\n",
      "\n",
      "        f = - np.sum((num.T / r_adjusted**1.5).T, axis=0) * m_current\n",
      "        forces[i] = f\n",
      "\n",
      "        if i!= 0 and i % 5000 == 0:\n",
      "            logger.debug(f\"Particle {i} done\")\n",
      "\n",
      "    return forces\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(utils.n_body_forces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {
    "tags": [
     "function_interparticle_distance"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def mean_interparticle_distance(particles: np.ndarray):\n",
      "    \"\"\"\n",
      "    Computes the mean interparticle distance of a set of particles.\n",
      "    Assumes that the particles array has the following columns: x, y, z ...\n",
      "    \"\"\"\n",
      "    if particles.shape[1] < 3:\n",
      "        raise ValueError(\"Particles array must have at least 3 columns: x, y, z\")\n",
      "    \n",
      "\n",
      "    r_half_mass = half_mass_radius(particles)\n",
      "    r = np.linalg.norm(particles[:, :3], axis=1)\n",
      "\n",
      "    n_half_mass = np.sum(r < r_half_mass)\n",
      "    logger.debug(f\"Number of particles within half mass radius: {n_half_mass} of {particles.shape[0]}\")\n",
      "\n",
      "    rho = n_half_mass / (4/3 * np.pi * r_half_mass**3)\n",
      "    # the mean distance between particles is the inverse of the density\n",
      "\n",
      "    epsilon = (1 / rho)**(1/3)\n",
      "    logger.info(f\"Found mean interparticle distance: {epsilon}\")\n",
      "    return epsilon\n",
      "    # TODO: check if this is correct\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(utils.mean_interparticle_distance))"
   ]
  }
 ],
 "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
}
