{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spaxel assignment\n",
    "\n",
    "This notebook shows the principle, how stellar particles or gas particles are assigned to the different spaxels. We show this here for squared spaxels.\n",
    "\n",
    "We start with two particles and assign them to the spatial matching spaxels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "from rubix.telescope.utils import square_spaxel_assignment\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "from jaxtyping import Float, Array \n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "\n",
    "# Define the particle coordinates\n",
    "coords = jnp.array([[0.5, 1.5], [2.5, 3.5]])\n",
    "print(\"coords\", coords)\n",
    "\n",
    "# Define the spatial bin edges\n",
    "spatial_bin_edges = jnp.array([0.0, 1.0, 2.0, 3.0, 4.0])\n",
    "\n",
    "# Compute the pixel assignments\n",
    "pixel_assignments = square_spaxel_assignment(coords, spatial_bin_edges)\n",
    "\n",
    "# Create a discrete colormap\n",
    "max_assignment = np.max(pixel_assignments)\n",
    "colors = plt.cm.viridis(np.linspace(0, 1, int(max_assignment) + 1))\n",
    "cmap = ListedColormap(colors)\n",
    "\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "# Plotting the particles with labels\n",
    "plt.subplot(1, 2, 1)\n",
    "scatter = plt.scatter(coords[:, 0], coords[:, 1], c=pixel_assignments, cmap=cmap, edgecolor='k')\n",
    "plt.colorbar(scatter, ticks=np.arange(0, max_assignment + 1))\n",
    "plt.title('Particle Coordinates and Pixel Assignments')\n",
    "plt.xlabel('X Coordinate')\n",
    "plt.ylabel('Y Coordinate')\n",
    "plt.xlim(spatial_bin_edges[0], spatial_bin_edges[-1])\n",
    "plt.ylim(spatial_bin_edges[0], spatial_bin_edges[-1])\n",
    "\n",
    "\n",
    "# Label each point with its pixel index\n",
    "for i, (x, y) in enumerate(coords[:, :2]):\n",
    "    plt.text(x, y, str(pixel_assignments[i]), color='red', fontsize=8)\n",
    "\n",
    "#create the bins\n",
    "for edge in spatial_bin_edges:\n",
    "    plt.axvline(edge, color='k', linestyle='--')\n",
    "    plt.axhline(edge, color='k', linestyle='--')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we do the same with a lot more random points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "#create random data\n",
    "n_stars = 1000\n",
    "coords = np.random.normal(2, 0.5, (n_stars, 3))\n",
    "coords = jnp.array(coords)\n",
    "\n",
    "# Compute the pixel assignments\n",
    "pixel_assignments = square_spaxel_assignment(coords, spatial_bin_edges)\n",
    "\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "# Plotting the particles with labels\n",
    "plt.subplot(1, 2, 1)\n",
    "scatter = plt.scatter(coords[:, 0], coords[:, 1], c=pixel_assignments, cmap=cmap, edgecolor='k')\n",
    "plt.colorbar(scatter, ticks=np.arange(0, max_assignment + 1))\n",
    "plt.title('Particle Coordinates and Pixel Assignments')\n",
    "plt.xlabel('X Coordinate')\n",
    "plt.ylabel('Y Coordinate')\n",
    "plt.xlim(spatial_bin_edges[0], spatial_bin_edges[-1])\n",
    "plt.ylim(spatial_bin_edges[0], spatial_bin_edges[-1])\n",
    "\n",
    "\n",
    "#create the bins\n",
    "for edge in spatial_bin_edges:\n",
    "    plt.axvline(edge, color='k', linestyle='--')\n",
    "    plt.axhline(edge, color='k', linestyle='--')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the last plot shows how many particles fall in each spaxel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "image = np.zeros((len(spatial_bin_edges) - 1, len(spatial_bin_edges) - 1))\n",
    "\n",
    "# Count the number of particles in each pixel\n",
    "for i in range(len(spatial_bin_edges) - 1):\n",
    "    for j in range(len(spatial_bin_edges) - 1):\n",
    "        image[i, j] = np.sum(pixel_assignments == (i + (len(spatial_bin_edges) - 1) * j))\n",
    "        \n",
    "        \n",
    "plt.imshow(image, cmap='viridis', origin='lower')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rubix",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}