{
 "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",
    "from astropy.visualization import quantity_support\n",
    "quantity_support()\n",
    "\n",
    "import utils\n",
    "import utils.logging_config\n",
    "utils.logging_config.set_log_level(\"info\")\n",
    "import logging\n",
    "logger = logging.getLogger(\"task2 (mesh)\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path('data')\n",
    "# DATA_NAME = 'data0.txt'\n",
    "# DATA_NAME = 'data1.txt'\n",
    "# DATA_NAME = 'data0_noise.txt'\n",
    "DATA_NAME = 'data1_noise.txt'\n",
    "CACHE_ROOT = Path('.cache')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:10 - utils.load - Loaded 9913 rows and 9 columns from data/data1_noise.txt\n",
      "11:03:10 - task2 (mesh) - Softening length: 0.037\n"
     ]
    }
   ],
   "source": [
    "points, columns = utils.load_data(DATA_ROOT / DATA_NAME)\n",
    "logger.debug(f\"Fetched {points.shape[0]} points, columns: {columns}\")\n",
    "# points = points[1:100, ...]\n",
    "# points = points[::2]\n",
    "# TODO remove\n",
    "softenings = points[:, -1]\n",
    "assert np.all(softenings == softenings[0]), \"Softening lengths are not uniform\"\n",
    "SOFTENING = softenings[0]\n",
    "logger.info(f\"Softening length: {SOFTENING}\")\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, 91732 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38772 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "utils.particles_plot_3d(particles[:, :3], particles[:, 3], title=\"Particle distribution\")\n",
    "# Note: colormap corresponds to the mass of the particles\n",
    "utils.particles_plot_2d(particles)\n",
    "# Note: colormap corresponds to the order of the particles in the array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12892 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19544 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Also consider the velocity distribution\n",
    "velocities = points[:, [5, 6, 7]]\n",
    "r = np.linalg.norm(particles[..., :3], axis=-1)\n",
    "v = np.linalg.norm(velocities, axis=-1)\n",
    "plt.figure()\n",
    "plt.plot(r, v, 'o')\n",
    "plt.xlabel('r')\n",
    "plt.ylabel('v')\n",
    "plt.show()\n",
    "\n",
    "## Check the velocity direction\n",
    "radial_velocities = np.zeros_like(v)\n",
    "circular_velocities = np.zeros_like(v)\n",
    "for i in range(particles.shape[0]):\n",
    "    if r[i] > 0:\n",
    "        radial_velocities[i] = np.abs(np.dot(velocities[i], particles[i, :3]) / r[i])\n",
    "        circular_velocities[i] = np.linalg.norm(np.cross(particles[i, :3], velocities[i])) / r[i]\n",
    "    else:\n",
    "        radial_velocities[i] = 0\n",
    "        circular_velocities[i] = 0\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(r, radial_velocities, 'o', label=\"Radial velocity\")\n",
    "plt.plot(r, circular_velocities, 'o', label=\"Circular velocity\")\n",
    "plt.xlabel('r')\n",
    "plt.ylabel('v')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### System characteristics\n",
    "- central black hole with mass $\\sim 10 \\%$\n",
    "- particles orbit circularly in the equatorial plane\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choice of units\n",
    "Recap of the particle properties:\n",
    "- $\\sim 10^4$ particles\n",
    "- around 1 black hole (10% of the mass)\n",
    "\n",
    "$\\implies$ ???"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:12 - task2 (mesh) - Considering a globular cluster - total mass of stars: 9.999999999999998, maximum radius of particles: 0.9096731180509756\n",
      "11:03:12 - utils.units - Set scales: M_SCALE = 1e+03 solMass, R_SCALE = 1.1 pc\n",
      "11:03:12 - task2 (mesh) - Black hole mass: 1e+03 solMass\n"
     ]
    }
   ],
   "source": [
    "# Set G = 1\n",
    "G = 1\n",
    "\n",
    "# from the particle number we can estimate the total (stellar) mass, excluding the BH\n",
    "M_TOT = 1e4 * u.M_sun\n",
    "# the radius aound the black hole follows from ??? # TODO\n",
    "R_TOT = 1 * u.pc\n",
    "\n",
    "# Rescale the units of the particles - considering only the orbiting stars\n",
    "M_particles = particles[:,3].sum() - 1\n",
    "R_particles = np.max(np.linalg.norm(particles[:, :3], axis=1))\n",
    "\n",
    "logger.info(f\"Considering a globular cluster - total mass of stars: {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",
    "logger.info(f\"Black hole mass: {utils.apply_units(particles[0, -1], \"mass\"):.2g}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "tags": [
     "plot_particle_distribution"
    ]
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84608 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 43092 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Plot the particles with units applied\n",
    "\n",
    "positions = utils.apply_units(particles[:, :3], \"position\")\n",
    "masses = utils.apply_units(particles[:, 3], \"mass\")\n",
    "utils.particles_plot_3d(positions, masses, title=\"Particle distribution\")\n",
    "\n",
    "utils.particles_plot_2d(positions)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:12 - utils.particles - Found mean interparticle distance: 0.0262396757880128\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/n_squared_forces__n_9913__softening_multiplier_0.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/n_squared_forces__n_9913__softening_multiplier_1.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/n_squared_forces__n_9913__softening_multiplier_10.npy\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_20__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_75__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_50__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_100__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_125__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_150__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_175__mapping_particle_to_cells_nn.npy\n",
      "11:03:12 - task2 (mesh) - Loaded forces from .cache/mesh_forces__n_9913__mesh_size_200__mapping_particle_to_cells_nn.npy\n"
     ]
    }
   ],
   "source": [
    "### Direct N body force computation\n",
    "epsilon = utils.mean_interparticle_distance(particles)\n",
    "\n",
    "epsilon_range = np.logspace(-2, 2, 5)\n",
    "epsilon_range = np.logspace(-1, 1, 3)\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",
    "    if cache_file.exists():\n",
    "        f = np.load(cache_file)\n",
    "        logger.info(f\"Loaded forces from {cache_file}\")\n",
    "    else:\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 = [20, 75, 50, 100, 125, 150, 175, 200]\n",
    "mapping = utils.particle_to_cells_nn\n",
    "\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",
    "    if cache_file.exists():\n",
    "        f = np.load(cache_file)\n",
    "        logger.info(f\"Loaded forces from {cache_file}\")\n",
    "    else:\n",
    "        f = utils.mesh_forces(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",
    "\n",
    "\n",
    "# TODO: compare computation time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "tags": [
     "plot_force_radial"
    ]
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 248076 characters]",
      "text/plain": [
       "<Figure size 640x480 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",
    "# many of the particles have the same distance from the origin, so we skip some of them\n",
    "SKIP_N = 5\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f'Radial force dependence (every {SKIP_N}th particle)')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('$r$')\n",
    "plt.ylabel('$|F(r)|$')\n",
    "\n",
    "\n",
    "colors = plt.cm.Greys(np.linspace(0.5, 1, len(epsilon_range)))\n",
    "for f, e, c in zip(n_squared_forces, epsilon_range, colors):\n",
    "    # remove the black hole:\n",
    "    plt.plot(r[1::SKIP_N], np.linalg.norm(f, axis=1)[1::SKIP_N], 'o', label=f\"$N^2 - {e:.1g} \\\\cdot \\\\epsilon$\", color=c, zorder=10)\n",
    "\n",
    "colors = plt.cm.coolwarm(np.linspace(0, 1, len(mesh_size_range)))\n",
    "for f, s, c in zip(mesh_forces, mesh_size_range, colors):\n",
    "    # remove the black hole:\n",
    "    plt.plot(r[1::SKIP_N], np.linalg.norm(f, axis=1)[1::SKIP_N], '.', label=f\"Mesh - N={s}\", alpha=0.2, color=c)\n",
    "\n",
    "# put the legend outside\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "plot_force_radial_single"
    ]
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61000 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Zoom in on a particularly clear case\n",
    "plt.figure()\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "\n",
    "ith = 1\n",
    "f = n_squared_forces[ith]\n",
    "f_val = np.linalg.norm(f, axis=1)\n",
    "\n",
    "plt.plot(r, f_val, 'o', label=f'$N^2 + {epsilon_range[ith]:.2g} \\\\cdot \\\\epsilon$', color=\"black\", zorder=10)\n",
    "\n",
    "f = mesh_forces[1]\n",
    "f_val = np.linalg.norm(f, axis=1)\n",
    "\n",
    "plt.plot(r, f_val, '.', label=f'Mesh [n={mesh_size_range[1]}]', color=plt.cm.Blues(0.5))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discussion\n",
    "- Using the baseline of $N^2 + 1 \\varepsilon$ softening we can see that already a 20 x 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 => noisy data even for the non-noisy particle distributions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.integrate as spi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:13 - task2 (mesh) - Considering 9913 particles\n",
      "11:03:13 - task2 (mesh) - Total mass: 10.999999999999998\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",
    "\n",
    "logger.info(f\"Considering {p0.shape[0]} particles\")\n",
    "logger.info(f\"Total mass: {np.sum(p0[:,6])}\")\n",
    "\n",
    "if logger.isEnabledFor(logging.DEBUG):\n",
    "    # assert that the ODE reshaping is consistent\n",
    "    p0_ref = p0.copy()\n",
    "    y0, _ = utils.ode_setup(p0, None)\n",
    "    logger.debug(y0[0:7])\n",
    "    p0_reconstructed = utils.to_particles(y0)\n",
    "    logger.debug(f\"{p0_ref[0]} -> {p0_reconstructed[0]}\")\n",
    "    logger.debug(f\"{p0_ref[1]} -> {p0_reconstructed[1]}\")\n",
    "\n",
    "    assert np.allclose(p0_ref, p0_reconstructed)\n",
    "    logger.debug(\"Consistency check passed\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "kjhfsdf"
    ]
   },
   "outputs": [],
   "source": [
    "def integrate(method: str, force_function: callable, p0: np.ndarray, t_range: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Integrate the gravitational movement of the particles, using the specified method\n",
    "    - method: the integration method to use (\"scipy\" or \"rk4\")\n",
    "    - force_function: the function that computes the forces acting on the particles\n",
    "    - p0: the initial conditions of the particles (n, 7) array, unflattened\n",
    "    - t_range: the time range to integrate over\n",
    "    Returns: the integrated positions and velocities of the particles in a 'flattened' array (time_steps, nx7)\n",
    "    \"\"\"\n",
    "    y0, y_prime = utils.ode_setup(p0, force_function)\n",
    "    \n",
    "    if method == \"scipy\":\n",
    "        sol = spi.odeint(y_prime, y0, t_range, rtol=1e-2)\n",
    "    elif method == \"rk4\":\n",
    "        sol = np.zeros((t_range.shape[0], y0.shape[0]))\n",
    "        sol[0] = y0\n",
    "        dt = t_range[1] - t_range[0]\n",
    "        for i in range(1, t_range.shape[0]):\n",
    "            t = t_range[i]\n",
    "            sol[i,...] = utils.runge_kutta_4(sol[i-1], t, y_prime, dt)\n",
    "\n",
    "\n",
    "    logger.info(f\"Integration done, shape: {sol.shape}\")\n",
    "    return sol\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:13 - utils.particles - Found mean interparticle distance: 0.0262396757880128\n",
      "11:03:13 - task2 (mesh) - Mean velocity: 0.001109653989710481, timestep: 0.023646718735142812\n"
     ]
    }
   ],
   "source": [
    "# Determine the integration timesteps\n",
    "# let's first compute the crossing time\n",
    "v = np.linalg.norm(particles[:, 3:6], axis=1)\n",
    "v_mean = np.mean(v)\n",
    "# a timestep should result in a small displacement, wrt. to the mean interparticle distance\n",
    "r_inter = utils.mean_interparticle_distance(particles)\n",
    "\n",
    "dt = r_inter / v_mean * 1e-3\n",
    "logger.info(f\"Mean velocity: {v_mean}, timestep: {dt}\")\n",
    "\n",
    "if np.isnan(dt):\n",
    "    raise ValueError(\"Invalid timestep\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11:03:13 - task2 (mesh) - Integration range: 0.0 -> 0.6857548433191415, n_steps: 30\n",
      "11:03:13 - utils.particles - Found mean interparticle distance: 0.0262396757880128\n",
      "11:03:13 - task2 (mesh) - Given softening length 0.037 does not match the mean interparticle distance 0.0262396757880128\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'to_particles' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[15], line 22\u001b[0m\n\u001b[1;32m     15\u001b[0m force_function \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mlambda\u001b[39;00m x: utils\u001b[38;5;241m.\u001b[39mn_body_forces(x, G, epsilon)\n\u001b[1;32m     16\u001b[0m \u001b[38;5;66;03m# force_function = lambda x: 0\u001b[39;00m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# force_function = lambda x: utils.n_body_forces_basic(x, G, epsilon)\u001b[39;00m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# force_function = lambda x: utils.analytical_forces(x)\u001b[39;00m\n\u001b[1;32m     19\u001b[0m \u001b[38;5;66;03m# force_function = lambda x: utils.mesh_forces_v2(x, G, 135, utils.particle_to_cells_nn)\u001b[39;00m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;66;03m# force_function = lambda x: utils.mesh_forces(x, G, 135, utils.particle_to_cells_nn)\u001b[39;00m\n\u001b[0;32m---> 22\u001b[0m sol \u001b[38;5;241m=\u001b[39m \u001b[43mintegrate\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mrk4\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mforce_function\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mp0\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt_range\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[13], line 20\u001b[0m, in \u001b[0;36mintegrate\u001b[0;34m(method, force_function, p0, t_range)\u001b[0m\n\u001b[1;32m     18\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(\u001b[38;5;241m1\u001b[39m, t_range\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m]):\n\u001b[1;32m     19\u001b[0m         t \u001b[38;5;241m=\u001b[39m t_range[i]\n\u001b[0;32m---> 20\u001b[0m         sol[i,\u001b[38;5;241m.\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;241m.\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[43mutils\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunge_kutta_4\u001b[49m\u001b[43m(\u001b[49m\u001b[43msol\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my_prime\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdt\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     23\u001b[0m logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIntegration done, shape: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00msol\u001b[38;5;241m.\u001b[39mshape\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m sol\n",
      "File \u001b[0;32m~/Documents/Uni/HS24/Computational Astrophysics/projects/nbody/utils/integrate.py:88\u001b[0m, in \u001b[0;36mrunge_kutta_4\u001b[0;34m(y, t, f, dt)\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mrunge_kutta_4\u001b[39m(y: np\u001b[38;5;241m.\u001b[39mndarray, t: \u001b[38;5;28mfloat\u001b[39m, f: \u001b[38;5;28mcallable\u001b[39m, dt: \u001b[38;5;28mfloat\u001b[39m):\n\u001b[1;32m     85\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     86\u001b[0m \u001b[38;5;124;03m    Runge-Kutta 4th order integrator.\u001b[39;00m\n\u001b[1;32m     87\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 88\u001b[0m     k1 \u001b[38;5;241m=\u001b[39m \u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[43my\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m     k2 \u001b[38;5;241m=\u001b[39m f(y \u001b[38;5;241m+\u001b[39m k1\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m dt, t \u001b[38;5;241m+\u001b[39m dt\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m2\u001b[39m)\n\u001b[1;32m     90\u001b[0m     k3 \u001b[38;5;241m=\u001b[39m f(y \u001b[38;5;241m+\u001b[39m k2\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m dt, t \u001b[38;5;241m+\u001b[39m dt\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m2\u001b[39m)\n",
      "File \u001b[0;32m~/Documents/Uni/HS24/Computational Astrophysics/projects/nbody/utils/solver_mesh.py:123\u001b[0m, in \u001b[0;36mode_setup.<locals>.f\u001b[0;34m(y, t)\u001b[0m\n\u001b[1;32m    118\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mf\u001b[39m(y, t):\n\u001b[1;32m    119\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    120\u001b[0m \u001b[38;5;124;03m    Computes the right hand side of the ODE system.\u001b[39;00m\n\u001b[1;32m    121\u001b[0m \u001b[38;5;124;03m    The ODE system is linearized around the current positions and velocities.\u001b[39;00m\n\u001b[1;32m    122\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 123\u001b[0m     p \u001b[38;5;241m=\u001b[39m \u001b[43mto_particles\u001b[49m(y)\n\u001b[1;32m    124\u001b[0m     \u001b[38;5;66;03m# this is explicitly a copy, which has shape (n, 7)\u001b[39;00m\n\u001b[1;32m    125\u001b[0m     \u001b[38;5;66;03m# columns x, y, z, vx, vy, vz, m\u001b[39;00m\n\u001b[1;32m    126\u001b[0m     \u001b[38;5;66;03m# (need to keep y intact since integrators make multiple function calls)\u001b[39;00m\n\u001b[1;32m    128\u001b[0m     forces \u001b[38;5;241m=\u001b[39m force_function(p[:, [\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]])\n",
      "\u001b[0;31mNameError\u001b[0m: name 'to_particles' is not defined"
     ]
    }
   ],
   "source": [
    "## Integration setup - use the n_squared forces for a few timesteps only, to see if the orbits are stable\n",
    "t_orbit = 2 * np.pi * r_inter / v_mean\n",
    "n_steps = int(t_orbit / dt * 5)\n",
    "n_steps = 30\n",
    "t_range = np.arange(0, n_steps*dt, dt)\n",
    "assert t_range.shape[0] == n_steps\n",
    "logger.info(f\"Integration range: {t_range[0]} -> {t_range[-1]}, n_steps: {n_steps}\")\n",
    "\n",
    "# The force function can be interchanged\n",
    "epsilon = utils.mean_interparticle_distance(particles)\n",
    "if epsilon != SOFTENING:\n",
    "    logger.warning(f\"Given softening length {SOFTENING} does not match the mean interparticle distance {epsilon}\")\n",
    "    epsilon = SOFTENING\n",
    "\n",
    "force_function = lambda x: utils.n_body_forces(x, G, epsilon)\n",
    "# force_function = lambda x: 0\n",
    "# force_function = lambda x: utils.n_body_forces_basic(x, G, epsilon)\n",
    "# force_function = lambda x: utils.analytical_forces(x)\n",
    "# force_function = lambda x: utils.mesh_forces(x, G, 135, utils.particle_to_cells_nn)\n",
    "\n",
    "sol = integrate(\"rk4\", force_function, p0, t_range)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "plot_system_evolution"
    ]
   },
   "outputs": [],
   "source": [
    "### Show the evolution\n",
    "particles_in_time = utils.to_particles_3d(sol)\n",
    "\n",
    "# First show the evolution of the particles in a series of 2D heatmaps:\n",
    "# reduce the font size\n",
    "fig, axs = plt.subplots(4, 6, figsize=(20, 12))\n",
    "fig.suptitle(\"Particle evolution (top view)\")\n",
    "\n",
    "print(t_range.shape)\n",
    "# make sure we have enough time steps to show\n",
    "diff = axs.size - particles_in_time.shape[0]\n",
    "if diff > 0:\n",
    "    logger.debug(f\"Adding dummy time steps: {diff=} -> {axs.size=}\")\n",
    "    plot_t_range = np.concatenate([t_range, np.zeros(diff)])\n",
    "    plot_particles_in_time = particles_in_time\n",
    "elif diff < 0:\n",
    "    logger.debug(f\"Too many steps to plot - reducing: {particles_in_time.shape[0]} -> {axs.size}\")\n",
    "    # skip some of the time steps\n",
    "    plot_t_range = []\n",
    "    plot_particles_in_time = []\n",
    "    for i in range(axs.size):\n",
    "        idx = int(i / axs.size * particles_in_time.shape[0])\n",
    "        # make sure we have the first and last time step are included\n",
    "        if i == 0:\n",
    "            idx = 0\n",
    "        elif i == axs.size - 1:\n",
    "            idx = particles_in_time.shape[0] - 1\n",
    "\n",
    "        plot_t_range.append(t_range[idx])\n",
    "        plot_particles_in_time.append(particles_in_time[idx])\n",
    "# othwerwise, we have exactly the right number of time steps\n",
    "\n",
    "for p, t, a in zip(plot_particles_in_time, plot_t_range, axs.flat):\n",
    "    a.set_title(f\"t={t:.2g}\")\n",
    "    utils.particles_plot_2d(p, ax=a)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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(n_steps)\n",
    "for i in range(n_steps):\n",
    "    p = particles_in_time[i]\n",
    "    ke_per_particle = 0.5 * p[:, 6] * np.linalg.norm(p[:,3:6], axis=1)**2\n",
    "    # logger.debug(f\"KE: {ke_per_particle.shape}\")\n",
    "    k_e = np.sum(ke_per_particle)\n",
    "    energies[i] = k_e\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(t_range, energies)\n",
    "plt.title('Kinetic energy')\n",
    "plt.xlabel('Integration time')\n",
    "plt.ylabel('Energy')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# radial extrema of the particles - disk surface\n",
    "r_mins = np.zeros(n_steps)\n",
    "r_maxs = np.zeros(n_steps)\n",
    "for i in range(n_steps):\n",
    "    p = particles_in_time[i][1:,...] # remove the black hole\n",
    "    r = np.linalg.norm(p[:,:3], axis=1)\n",
    "    r_mins[i] = np.min(r)\n",
    "    r_maxs[i] = np.max(r)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(t_range, r_mins, label='r_min')\n",
    "plt.plot(t_range, r_maxs, label='r_max')\n",
    "plt.title('Radial extrema')\n",
    "plt.xlabel('Integration time')\n",
    "plt.ylabel('r')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full PM solver\n",
    "We now have all the tools to implement the full PM solver:\n",
    "- force computation using mesh\n",
    "- integrator with RK4\n",
    "- estimate for good timesteps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "22:37:37 - task2 (mesh) - Integration range: 0.0 -> 0.2128204686162853, n_steps: 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "22:38:43 - task2 (mesh) - Integration done, shape: (10, 69391)\n"
     ]
    }
   ],
   "source": [
    "t_orbit = 2 * np.pi * r_inter / v_mean\n",
    "n_steps = int(t_orbit / dt * 5)\n",
    "n_steps = 10\n",
    "t_range = np.arange(0, n_steps*dt, dt)\n",
    "logger.info(f\"Integration range: {t_range[0]} -> {t_range[-1]}, n_steps: {n_steps}\")\n",
    "\n",
    "\n",
    "mesh_size = 125 # as per the previous discussion\n",
    "r = np.linalg.norm(particles[:,:3], axis=1)\n",
    "r_max = np.max(r)\n",
    "mesh_bound = (-2*r_max, 2*r_max)\n",
    "\n",
    "# The force function is now more \"rigid\" => the mesh is fixed and boundary conditions are enforced\n",
    "force_function = lambda x: utils.mesh_solver(x, G, utils.particle_to_cells_nn, mesh_size, mesh_bound, \"vanishing\")\n",
    "sol = integrate(\"rk4\", force_function, p0, t_range)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10,)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n",
      "0\n",
      "(9913, 7)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 204816 characters]",
      "text/plain": [
       "<Figure size 2000x1200 with 24 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Show some results\n",
    "particles_in_time = utils.to_particles_3d(sol)\n",
    "\n",
    "# First show the evolution of the particles in a series of 2D heatmaps:\n",
    "# reduce the font size\n",
    "fig, axs = plt.subplots(4, 6, figsize=(20, 12))\n",
    "fig.suptitle(\"Particle evolution (top view)\")\n",
    "\n",
    "print(t_range.shape)\n",
    "# make sure we have enough time steps to show\n",
    "diff = axs.size - particles_in_time.shape[0]\n",
    "if diff > 0:\n",
    "    logger.debug(f\"Adding dummy time steps: {diff=} -> {axs.size=}\")\n",
    "    plot_t_range = np.concatenate([t_range, np.zeros(diff)])\n",
    "    plot_particles_in_time = particles_in_time\n",
    "elif diff < 0:\n",
    "    logger.debug(f\"Too many steps to plot - reducing: {particles_in_time.shape[0]} -> {axs.size}\")\n",
    "    # skip some of the time steps\n",
    "    plot_t_range = []\n",
    "    plot_particles_in_time = []\n",
    "    for i in range(axs.size):\n",
    "        idx = int(i / axs.size * particles_in_time.shape[0])\n",
    "        # make sure we have the first and last time step are included\n",
    "        if i == 0:\n",
    "            idx = 0\n",
    "        elif i == axs.size - 1:\n",
    "            idx = particles_in_time.shape[0] - 1\n",
    "\n",
    "        plot_t_range.append(t_range[idx])\n",
    "        plot_particles_in_time.append(particles_in_time[idx])\n",
    "# othwerwise, we have exactly the right number of time steps\n",
    "\n",
    "for p, t, a in zip(plot_particles_in_time, plot_t_range, axs.flat):\n",
    "    a.set_title(f\"t={t:.2g}\")\n",
    "    # remove nan values\n",
    "    print(np.sum(np.isnan(p)))\n",
    "    p = p[~np.isnan(p).any(axis=1)]\n",
    "    print(p.shape)\n",
    "    utils.particles_plot_2d(p, ax=a)\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "function_runge_kutta"
    ]
   },
   "outputs": [],
   "source": [
    "print(inspect.getsource(utils.runge_kutta_4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "function_mesh_force"
    ]
   },
   "outputs": [],
   "source": [
    "print(inspect.getsource(utils.mesh_forces))\n",
    "print(inspect.getsource(utils.mesh_poisson))"
   ]
  }
 ],
 "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
}
