{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy as sc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Global parameters\n",
    "N = 100\n",
    "space_range = [-1, 1]\n",
    "center = np.array([0, 0, 0])\n",
    "r0 = 0.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_mesh(n_side, space_range) -> tuple[np.ndarray, tuple[np.ndarray, np.ndarray, np.ndarray]]:\n",
    "    x = np.linspace(space_range[0], space_range[1], n_side)\n",
    "    y = np.linspace(space_range[0], space_range[1], n_side)\n",
    "    z = np.linspace(space_range[0], space_range[1], n_side)\n",
    "    mesh = np.zeros((n_side, n_side, n_side))\n",
    "    return mesh, np.meshgrid(x, y, z)\n",
    "\n",
    "def show_mesh(grid, mesh, title='Mesh'):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "    ax.scatter(grid[0], grid[1], grid[2], c=mesh, alpha=0.2, s=1)\n",
    "    ax.set_title(title)\n",
    "    plt.show()\n",
    "\n",
    "def show_mesh_2d(mesh, title='Mesh'):\n",
    "    fig = plt.figure()\n",
    "    axs = fig.subplots(1, 3)\n",
    "    im = axs[0].imshow(np.sum(mesh, axis=0))\n",
    "    axs[0].imshow(np.sum(mesh, axis=0))\n",
    "    axs[1].imshow(np.sum(mesh, axis=1))\n",
    "    axs[2].imshow(np.sum(mesh, axis=2))\n",
    "    fig.colorbar(im, ax=axs, orientation='horizontal', fraction=0.02, pad=0.1)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def density_distribution(grid, mesh, center):\n",
    "    radial_distance = np.linalg.norm(grid - center[:, None, None, None], axis=0)\n",
    "    # radial_distance = np.sqrt((grid[0] - center[0])**2 + (grid[1] - center[1])**2 + (grid[2] - center[2])**2)\n",
    "    n = mesh.shape[0]//3\n",
    "    distance_steps = np.linspace(0, np.max(radial_distance), n)\n",
    "    density = np.zeros(n)\n",
    "\n",
    "    for i in range(1, n):\n",
    "        d_inner = distance_steps[i-1]\n",
    "        d_outer = distance_steps[i]\n",
    "        mask_outer = radial_distance < d_outer\n",
    "        mask_inner = radial_distance > d_inner\n",
    "        volume = 4/3 * np.pi * (d_outer**3 - d_inner**3)\n",
    "        # print(f'Calculating density for distance {d_inner}-{d_outer}')\n",
    "        # show_mesh(grid, mask_outer & mask_inner, title=f'{d_inner}-{d_outer}')\n",
    "        density[i-1] = np.sum(mesh[mask_outer & mask_inner]) / volume\n",
    "\n",
    "    return distance_steps, density\n",
    "\n",
    "def plot_distribution(x, y, title='Density distribution'):\n",
    "    plt.plot(x, y)\n",
    "    plt.vlines(r0, np.min(y), np.max(y), colors='r', linestyles='dashed')\n",
    "    plt.xlabel('Distance')\n",
    "    plt.title(title)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25608 characters]",
      "text/plain": [
       "<Figure size 640x480 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30916 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mesh, grid = create_mesh(N, space_range)\n",
    "# show_mesh(grid, mesh, 'Initial mesh')\n",
    "\n",
    "# linalg.norm does not work here, because we want to do it matrix-wise\n",
    "radial_distance = np.sqrt((grid[0] - center[0])**2 + (grid[1] - center[1])**2 + (grid[2] - center[2])**2)\n",
    "\n",
    "mesh[radial_distance < r0] = 1\n",
    "\n",
    "# show_mesh(grid, mesh, 'Mesh after applying the condition')\n",
    "show_mesh_2d(mesh)\n",
    "\n",
    "\n",
    "distance_steps, density = density_distribution(grid, mesh, center)\n",
    "plot_distribution(distance_steps, density, 'Density distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 130136 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38860 characters]",
      "text/plain": [
       "<Figure size 640x480 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30120 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Construct the mesh and grid of the fourier space mapped quantities\n",
    "rho_k = np.fft.fftn(mesh)\n",
    "k = np.fft.fftfreq(N, d=(space_range[1] - space_range[0]) / N)\n",
    "kx, ky, kz = np.meshgrid(k, k, k)\n",
    "\n",
    "# print(rho_k.shape, kx.shape, ky.shape, kz.shape)\n",
    "\n",
    "# apply the poisson equation in fourier space\n",
    "k_abs = np.sqrt(kx**2 + ky**2 + kz**2)\n",
    "# intermediately set this to one and later on set the result back to zero\n",
    "k_corrected = np.where(k_abs == 0, 1, k_abs)\n",
    "phi_k = - rho_k * (4 * np.pi / k_corrected**2)\n",
    "phi_k[k_abs == 0] = 0\n",
    "\n",
    "# transform back to real space\n",
    "phi = np.fft.ifftn(phi_k)\n",
    "show_mesh(grid, np.real(phi), 'Potential')\n",
    "show_mesh_2d(np.real(phi), 'Potential')\n",
    "\n",
    "distance_steps, phi_d = density_distribution(grid, np.real(phi), center)\n",
    "plot_distribution(distance_steps, phi_d, 'Potential distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31092 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# analytical solution\n",
    "def phi_analytical(r, r0):\n",
    "    phi = np.zeros_like(r)\n",
    "    phi[r < r0] = np.abs(r[r < r0]/r0)**2 - 3\n",
    "    phi[r >= r0] = -2 * np.abs(r0/r[r >= r0])\n",
    "    phi *= (4 * np.pi / 6) * r0**2\n",
    "    return phi\n",
    "\n",
    "distance_steps = np.linspace(0, np.max(radial_distance), N)\n",
    "phi_a = phi_analytical(distance_steps, r0)\n",
    "plot_distribution(distance_steps, phi_a, 'Analytical potential distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def acceleration(phi, grid, center, end):\n",
    "    # line between center and end\n",
    "    # todo center and end are wrong\n",
    "    direction = (end - center / np.linalg.norm(end - center)).astype(int)\n",
    "    print(direction)\n",
    "    points = [center + i * direction for i in range(1, phi.shape[0]//2)]\n",
    "    phi_values = phi[points]\n",
    "    print(phi_values)\n",
    "\n",
    "    a = - np.gradient(phi_values)\n",
    "    return a\n",
    "\n",
    "a = acceleration(np.real(phi), grid, center, np.array([1, 1, 1]))\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
}
