{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The point spread function\n",
    "\n",
    "The point spread functions blures the data in spatial direction and is an observational artefact. So far we have implemented a gaussian kernel. With that kernel the mock-observations can be convolved to mimic real observations. It is important that the sum of the kernel is 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NBVAL_SKIP\n",
    "from rubix.telescope.psf.kernels import gaussian_kernel_2d\n",
    "\n",
    "kernel = gaussian_kernel_2d(20,20,3.5)\n",
    "print(kernel.shape)\n",
    "print(kernel.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NBVAL_SKIP\n",
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(kernel, cmap='hot')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here an example for the PSF convolution, we have an artificial datacube of shape (50,50,300), which contains random numbers in the spatial dimension. Each layer in the wavelength dimension is convolved with the kernel. We plot one spaxel [10,10] along the wavelength range for the original random data and the psf smoothed random data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NBVAL_SKIP\n",
    "# Get an example Datacube\n",
    "from rubix.telescope.psf.psf import apply_psf\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "datacube = np.ones((50,50,300))\n",
    "# create random data\n",
    "for i in range(300):\n",
    "    datacube[:,:,i] = np.random.rand(50,50)\n",
    "\n",
    "datacube = jnp.array(datacube)\n",
    "\n",
    "convolved_datacube = apply_psf(datacube, kernel)\n",
    "print(convolved_datacube.shape)\n",
    "\n",
    "plt.plot(convolved_datacube[10,10,:], label='convolved')\n",
    "plt.plot(datacube[10,10,:], label='original')\n",
    "plt.legend()"
   ]
  }
 ],
 "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
}