{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "tough-notification",
   "metadata": {},
   "source": [
    "3: Writing and including custom Effects\n",
    "=======================================\n",
    "\n",
    "In this tutorial, we will load the model of MICADO (including Armazones, ELT, MORFEO) and then turn off all effect that modify the spatial extent of the stars. The purpose here is to see in detail what happens to the **distribution of the stars flux on a sub-pixel level** when we add a plug-in astrometric Effect to the optical system.\n",
    "\n",
    "For real simulation, we will obviously leave all normal MICADO effects turned on, while still adding the plug-in Effect. Hopefully this tutorial will serve as a refernce for those who want to see **how to create Plug-ins** and how to manipulate the effects in the MICADO optical train model.\n",
    "\n",
    "\n",
    "Create and optical model for MICADO and the ELT\n",
    "-----------------------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "constant-weekly",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tempfile import TemporaryDirectory\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "import scopesim as sim\n",
    "from scopesim_templates.stellar import stars, star_grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40fabcee",
   "metadata": {},
   "source": [
    "Scopesim works by using so-called instrument packages, which have to be downloaded separately. For normal use, you would set the package directory (a local folder path, `local_package_folder` in this example), download the required packages *once*, and then **remove the download command**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "661ea82b",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_package_folder = \"./inst_pkgs\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1350c51d",
   "metadata": {},
   "source": [
    "However, to be able to run this example on the *Readthedocs* page, we need to include a temporary directory.\n",
    "\n",
    "**Do not** copy and run this code locally, it is **only** needed to set things up for *Readthedocs*!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d33b08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tempfile import TemporaryDirectory\n",
    "local_package_folder = TemporaryDirectory().name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acute-calculator",
   "metadata": {},
   "source": [
    "Download the required instrument packages for an observation with MICADO at the ELT.\n",
    "\n",
    "Again, you would only need to do this **once**, not every time you run the rest of the script, assuming you set a (permanent) instrument package folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "gorgeous-blond",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.download_packages([\"Armazones\", \"ELT\", \"MICADO\", \"MORFEO\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "loving-skill",
   "metadata": {},
   "source": [
    "We can see which Effects are already included by calling ``micado.effects``:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "celtic-fluid",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmd = sim.UserCommands(use_instrument=\"MICADO\", set_modes=[\"SCAO\", \"IMG_1.5mas\"])\n",
    "micado = sim.OpticalTrain(cmd)\n",
    "\n",
    "micado.effects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "suburban-inquiry",
   "metadata": {},
   "source": [
    "Now we turn off all ``Effects`` that cause spatial aberrations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bound-literature",
   "metadata": {},
   "outputs": [],
   "source": [
    "for effect_name in [\"full_detector_array\", \"micado_adc_3D_shift\", \n",
    "                    \"micado_ncpas_psf\", \"relay_psf\"]:\n",
    "    micado[effect_name].include = False\n",
    "    print(micado[effect_name])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "narrow-bulletin",
   "metadata": {},
   "source": [
    "The normal detector window is set to 1024 pixels square.\n",
    "Let's reduce the size of the detector readout window:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "allied-matrix",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado[\"detector_window\"].data[\"x_cen\"] = 0          # [mm] distance from optical axis on the focal plane\n",
    "micado[\"detector_window\"].data[\"y_cen\"] = 0\n",
    "micado[\"detector_window\"].data[\"x_size\"] = 64          # [pixel] width of detector\n",
    "micado[\"detector_window\"].data[\"y_size\"] = 64"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "original-sugar",
   "metadata": {},
   "source": [
    "By default ScopeSim works on the whole pixel level for saving computation time.\n",
    "However it is capable of integrating sub pixel shift.\n",
    "For this we need to turn on the sub-pixel mode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aerial-warehouse",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.cmds[\"!SIM.sub_pixel.flag\"] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "anonymous-chorus",
   "metadata": {},
   "source": [
    "We can test what's happening by making a grid of stars and observing them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "indoor-norway",
   "metadata": {},
   "outputs": [],
   "source": [
    "src = star_grid(n=9, mmin=20, mmax=20.0001, separation=0.0015 * 15)\n",
    "src.fields[0][\"x\"] -= 0.00075\n",
    "src.fields[0][\"y\"] -= 0.00075\n",
    "\n",
    "micado.observe(src, update=True)\n",
    "\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.imshow(micado.image_planes[0].data, origin=\"lower\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lightweight-louisiana",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado[\"detector_window\"].data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eastern-psychology",
   "metadata": {},
   "source": [
    "Writing a custom Effect object\n",
    "------------------------------\n",
    "\n",
    "The following code snippet creates a new ``Effect`` class.\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "weighted-mortgage",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from astropy.table import Table\n",
    "\n",
    "from scopesim.effects import Effect\n",
    "from scopesim.base_classes import SourceBase\n",
    "\n",
    "\n",
    "class PointSourceJitter(Effect):\n",
    "    def __init__(self, **kwargs):\n",
    "        super(PointSourceJitter, self).__init__(**kwargs)   # initialise the underlying Effect class object\n",
    "        self.meta[\"z_order\"] = [500]                        # z_order number for knowing when and how to apply the Effect\n",
    "        self.meta[\"max_jitter\"] = 0.001                     # [arcsec] - a parameter needed by the effect\n",
    "        self.meta.update(kwargs)                            # add any extra parameters passed when initialising\n",
    "\n",
    "    def apply_to(self, obj):                                # the function that does the work\n",
    "        if isinstance(obj, SourceBase):\n",
    "            for field in obj.fields:\n",
    "                if isinstance(field, Table):\n",
    "                    dx, dy = 2 * (np.random.random(size=(2, len(field))) - 0.5)\n",
    "                    field[\"x\"] += dx * self.meta[\"max_jitter\"]\n",
    "                    field[\"y\"] += dy * self.meta[\"max_jitter\"]\n",
    "\n",
    "        return obj"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "drawn-vacation",
   "metadata": {},
   "source": [
    "Lets break it down a bit (**THIS IS JUST A STEP-BY-STEP EXPLANATION OF THE CODE ABOVE, NOT SOMETHING NEW!**):\n",
    "\n",
    "    class PointSourceJitter(Effect):\n",
    "        ...\n",
    "\n",
    "Here we are subclassing the ``Effect`` object from ScopeSim.\n",
    "This has the basic functionality for reading in ASCII and FITS files, and for communicating with the ``OpticsManager`` class in ScopeSim.\n",
    "\n",
    "The initialisation function looks like this:\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        super(PointSourceJitter, self).__init__(**kwargs)   # initialise the underlying Effect class object\n",
    "        self.meta[\"z_order\"] = [500]\n",
    "\n",
    "Here we make sure to activate the underlying Effect object.\n",
    "The ``z_order`` keyword in the meta dictionary is used by ScopeSim to determine when and where this Effect should be applied during a simulations run.\n",
    "The exact z-order numbers are described in [insert link here].\n",
    "\n",
    "The main function of any Effect is the ``apply_to`` method:\n",
    "\n",
    "    def apply_to(self, obj):\n",
    "        if isinstance(obj, SourceBase):\n",
    "            ...\n",
    "\n",
    "        return obj\n",
    "\n",
    "It should be noted that what is passed in via (``obj``) must be returned in the same format. The contents of the ``obj`` can change, but the ``obj`` object must be returned.\n",
    "\n",
    "All the code which enacts the results of the physical effect are contained in this method.\n",
    "For example, if we are writing a redshifting Effect, we could write the code to shift the wavelength array of a ``Source`` object by ``z+1`` here.\n",
    "\n",
    "There are 4 main classes that are cycled through during an observation run:\n",
    "* ``SourceBase``: contains the original 2+1D distribtion of light,\n",
    "* ``FieldOfViewBase``: contains a (quasi-)monochromatic cutout from the Source object,\n",
    "* ``ImagePlaneBase``: contains the expectation flux image on the detector plane\n",
    "* ``DetectorBase``: contains the electronic readout image\n",
    "\n",
    "An ``Effect`` object can be applied to any number of objects based on one or more of these base classes.\n",
    "Just remember to segregate the base-class-specific code with ``if`` statements.\n",
    "\n",
    "One further method should be mentioned: ``def fov_grid()``.\n",
    "This method is used by ``FOVManager`` to estimate how many ``FieldOfView`` objects to generate in order to best simulation the observation.\n",
    "If your Effect object might alter this estimate, then you should include this method in your class. See the code base for further details.\n",
    "\n",
    "**Note**: The ``fov_grid`` method will be depreciated in a future release of ScopeSim.\n",
    "    It will most likely be replaced by a ``FOVSetupBase`` class that will be cycled through the ``apply_to`` function.\n",
    "    However this is not yet 100% certain, so please bear with us."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "shaped-priority",
   "metadata": {},
   "source": [
    "Including a custom Effect\n",
    "-------------------------\n",
    "\n",
    "First we need to initialise an instance of the Effect object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "empirical-skill",
   "metadata": {},
   "outputs": [],
   "source": [
    "jitter_effect = PointSourceJitter(max_jitter=0.001, name=\"random_jitter\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "southern-gothic",
   "metadata": {},
   "source": [
    "Then we can add it to the optical model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "considerable-factory",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.optics_manager.add_effect(jitter_effect)\n",
    "\n",
    "micado.effects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "forty-statistics",
   "metadata": {},
   "source": [
    "When we want to observe, we need to include the ``update=True`` flag so that the optical model is updated to include the instance of our new ``Effect``:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exempt-purse",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado.observe(src, update=True)\n",
    "\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.imshow(micado.image_planes[0].data, origin=\"lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interstate-qatar",
   "metadata": {},
   "source": [
    "We can update the parameters of the object on-the-fly by accessing the meta dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sound-preference",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado[\"random_jitter\"].meta[\"max_jitter\"] = 0.005\n",
    "\n",
    "micado.observe(src, update=True)\n",
    "\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.imshow(micado.image_planes[0].data, origin=\"lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fluid-canon",
   "metadata": {},
   "source": [
    "Here we can see that there is a certain amount of sub-pixel jitter being introduced into each observation.\n",
    "However this bare-bones approach is not very realistic.\n",
    "We should therefore turn the PSF back on to get a more realistic observation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "future-approval",
   "metadata": {},
   "outputs": [],
   "source": [
    "micado[\"relay_psf\"].include = True\n",
    "\n",
    "micado.observe(src, update=True)\n",
    "hdus = micado.readout()\n",
    "\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.imshow(hdus[0][1].data, origin=\"lower\", norm=LogNorm())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fossil-range",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.16"
  },
  "nbsphinx": {
   "execute": "auto"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}