{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "framed-clerk",
   "metadata": {},
   "source": [
    "# Wide-field observations of PSF variations with SCAO\n",
    "\n",
    "In this example we will show the effect of the off-axis AO correction on the shape of the PSF when observing with the SCAO mode.\n",
    "To make a semi-realistic observation, we will observe an open cluster created by the ``ScopeSim_Templates`` package.\n",
    "This cluster extends over the full MICADO field of view.\n",
    "As such it should show the full extend of the expected PSF variations over the field.\n",
    "\n",
    "It should be noted that adding PSF field variations is a computationally expensive process. ScopeSim does it's best to reduce the level of complexity by discretising the PSF variation into different zones over the field of view.\n",
    "Occasionally the borders between these zones is visible "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unlimited-cliff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from astropy import units as u\n",
    "\n",
    "import scopesim as sim\n",
    "import scopesim_templates as sim_tp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "about-darwin",
   "metadata": {},
   "source": [
    "The PSF field-variation is complex and computationally expensive.\n",
    "Hence this functionality is not inluded in the ``MICADO`` package.\n",
    "For use cases requireing field variations we must use the pipeline-oriented ``MICADO`` package.\n",
    "As we will be using the SCAO mode, we can ignore the ``MORFEO`` module.\n",
    "\n",
    "    sim.download_packages([\"Armazones\", \"ELT\", \"MORFEO\", \"MICADO\"])\n",
    "    \n",
    ".. note:: In ScopeSim v0.5 both download_packages (new format) and download_package (old format) exist.\n",
    "\n",
    "If we would like to keep the instrument packages in a separate directory, we can set the following config value:\n",
    "\n",
    "    sim.rc.__config__[\"!SIM.file.local_packages_path\"] = \"path/to/packages\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "formed-metabolism",
   "metadata": {},
   "source": [
    "## Set up an open cluster ``Source`` object\n",
    "\n",
    "The first step is to create a ``Source`` object. \n",
    "Here we take advantage of the ``cluster`` function from the ``ScopeSim_Templates.stellar.clusters`` submodule.\n",
    "The ``cluster`` function draws masses from an standard Kroupa IMF until the ``mass`` limit has been reached.\n",
    "It then generates randomly draw positions from a 2D gaussian distribution with a HWHM equal to the given ``core_radius`` in parsec.\n",
    "The ``distance`` parameter sets the distance in parsec to the centre of the cluster.\n",
    "This allows the final on-sky star positions in units of ``arcsec`` from the centre of the cluster to be calculated.\n",
    "Spectra for all the stars are taken from the Pickles (1998) catalogue and are imported from the ``Pyckles`` python packages.\n",
    "\n",
    "Here we choose a cluster with a mass of 10000 solar masses (~250000 stars) with a core radius of 0.5 parsec, located at a distance of 8 kpc from the Earth.\n",
    "This essentially gives us a cluster that covers the whole MICADO filed of view in wide-field mode (4mas / pixel)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "altered-paris",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = sim_tp.cluster(mass=10000, core_radius=0.5, distance=8000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3f68de-ca2a-470f-928e-6b2c95da4415",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "enhanced-strip",
   "metadata": {},
   "source": [
    "## Set up the MICADO optical system for SCAO and include a field-varying PSF\n",
    "\n",
    "The next step is the standard proceedure - set the instrument to ``MICADO`` and the modes to ``SCAO`` and ``IMG_4mas``, then create an ``Simulation`` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "continent-springer",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado = sim.Simulation(\"MICADO\", [\"SCAO\", \"IMG_4mas\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "identified-archive",
   "metadata": {},
   "source": [
    "By default the PSF model in the pipeline-oriented ``MICADO`` package is set to use a field-constant PSF cube.\n",
    "We need to change this to use a Field-Varying PSF model.\n",
    "\n",
    "**Note:** to view all the optical effects contained within the optical model, use the ``.effects`` attribute in the ``OpticalTrain`` class.\n",
    "\n",
    "Here we see the PSF model in the 4th last line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "peripheral-monday",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.optical_train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ruled-lounge",
   "metadata": {},
   "source": [
    "We can access the PSF model using the standard python dictionary notion using the name of the element.\n",
    "In this case it is called ``relay_psf`` as this refers to the internal MICADO relay optics which are connected to the SCAO system.\n",
    "\n",
    "We do not want to use the default ``relay_psf`` object, so we need to tell ScopeSim not to include it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "subjective-parallel",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.optical_train[\"relay_psf\"].include = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "threatened-steam",
   "metadata": {},
   "source": [
    "Next we want to add a ``FieldVaryingPSF`` object from the ScopeSim library of optical effect.\n",
    "Before we initialise the object though, we will need the correct PSF cube.\n",
    "\n",
    "A small library of precomputed PSFs can be found on the ScopeSim server: https://scopesim.univie.ac.at/InstPkgSvr/psfs/\n",
    "\n",
    "Once we have downloaded the PSF cube of our choice (in this case the cube corresponding to ESO median atmospheric conditions) we can pass the file path to the ``FieldVaryingPSF`` object.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">For the purpose of testing this notebook, we're using a much smaller test PSF file. If you'd like to see the real deal, simply uncomment the one below and comment out the test file instead.</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b0961b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pooch import retrieve\n",
    "# fname = retrieve(\n",
    "#     \"https://scopesim.univie.ac.at/InstPkgSvr/psfs/AnisoCADO_SCAO_FVPSF_4mas_EsoMedian_20190328.fits\",\n",
    "#     known_hash=\"a9bd309e5ac025f2aa5f8ded85323465578906e47025ce86501985c9b258474c\",\n",
    "#     fname=\"AnisoCADO_SCAO_FVPSF_4mas_EsoMedian_20190328.fits\",\n",
    "#     progressbar=True,\n",
    "# )\n",
    "fname = retrieve(\n",
    "    \"https://scopesim.univie.ac.at/InstPkgSvr/psfs/test_fvpsf.fits\",\n",
    "    known_hash=\"c13aba82e4faf42568b5996f3d32135ddb69b537a832c753da189b698a7ad48c\",\n",
    "    fname=\"test_fvpsf.fits\",\n",
    "    progressbar=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "improving-august",
   "metadata": {},
   "outputs": [],
   "source": [
    "fv_psf = sim.effects.psfs.FieldVaryingPSF(filename=fname, name=\"SCAO_FV_PSF\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decimal-productivity",
   "metadata": {},
   "source": [
    "We now need to add the new PSF to the MICADO relay optics element inside the ``MICADO`` optical train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "elegant-swimming",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.optical_train.optics_manager[\"default_ro\"].add_effect(fv_psf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "funded-given",
   "metadata": {},
   "source": [
    "Viewing an on-axis PSF is boring. Therefore we will move the 1024x1024 detector window off centre 10 arcseconds in both x and y directions.\n",
    "We will also increase the exposure time to 60 seconds.\n",
    "\n",
    "Now we are ready to observe and readout the MICADO detector window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "revised-paper",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.settings[\"!OBS.dit\"] = 600    # [s]\n",
    "micado.settings[\"!DET.x\"] = 10       # [arcsec]\n",
    "micado.settings[\"!DET.y\"] = 10\n",
    "\n",
    "hdus = micado(cluster)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "external-billion",
   "metadata": {},
   "source": [
    "The elongation of the PSF along the radial axis is quite visible."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "animated-multiple",
   "metadata": {},
   "source": [
    "## First look at limited FOV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prescription-certificate",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(figsize=(12, 12))\n",
    "img = hdus[1].data\n",
    "ax.imshow(img, origin=\"lower\", norm=\"log\", vmin=0.9*np.median(img), vmax=1.1*np.median(img));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wrong-excuse",
   "metadata": {},
   "source": [
    "## Larger view\n",
    "\n",
    "Alternatively we could re-centre the detector window, but expand the window size to 4096x4096 pixels (16x16 arcsecond) which is the equivalent of the central MICADO detector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "romantic-birth",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.settings[\"!DET.x\"] = 0     \n",
    "micado.settings[\"!DET.y\"] = 0\n",
    "micado.settings[\"!DET.width\"] = 4096\n",
    "micado.settings[\"!DET.height\"] = 4096\n",
    "\n",
    "hdus = micado(cluster)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "classical-ethiopia",
   "metadata": {},
   "source": [
    "It should be noted, the larger the detector size, the longer the simulation takes to run.\n",
    "Hence the option to simply shift the detector centre while keeping the detector size small.\n",
    "\n",
    "This option does however let us see the change in orientation of the PSF elongation around the on-axis position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "patient-thinking",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(figsize=(12, 12))\n",
    "img = hdus[1].data\n",
    "ax.imshow(img, origin=\"lower\", norm=\"log\", vmin=0.9*np.median(img), vmax=3*np.median(img));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "quality-scroll",
   "metadata": {},
   "source": [
    "## Full detector array\n",
    "\n",
    "<div class=\"alert alert-block alert-warning\">It is ONLY recommended to run the full MICADO detector configuration is you have >8 GB of RAM</div>\n",
    "\n",
    "While the ``MICADO`` package only has a customisable detector window for simulated observations, the ``MICADO`` package contains a full description of the 9 detectors in the MICADO detector array.\n",
    "\n",
    "This requires simulating a ~13000x13000 pixel field, which understandable takes a lot of computer resources and takes on the **order of 10 minutes to run.**\n",
    "\n",
    "\n",
    "\n",
    "Ideally the user would write a script using the commands in this notebook and run that script in the background.\n",
    "\n",
    "That said, swapping to the full detector array is a simple matter of setting the ``.include`` parameter to ``True`` and ``False`` for the two Detector options provided in the ``MICADO`` package (see ``micado.effects`` table above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collectible-actress",
   "metadata": {},
   "outputs": [],
   "source": [
    "# micado.optical_train[\"full_detector_array\"].include = True\n",
    "# micado.optical_train[\"detector_window\"].include = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "considered-institute",
   "metadata": {},
   "source": [
    "In this case, it is worth saving the final ``HDUList`` to disk as a multi-extension FITS file so the data are safe in case something happens to the Python session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mounted-burns",
   "metadata": {},
   "outputs": [],
   "source": [
    "# micado(cluster, filename=\"micado_full_detector.fits\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}