{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MPAGS Spectroscopy assignment\n",
    "\n",
    "Code by M. Brogi, modified by M. Lafarga Magro\n",
    "\n",
    "\n",
    "In this assignment you are going to extract the long-slit spectra of two stars using DS9 and this python code. Below you have:\n",
    "1. The instructions to open and draw regions in DS9\n",
    "2. Some python code to read in the regions drawn with DS9 and extract the spectra\n",
    "3. Some questions for you to answer\n",
    "\n",
    "Please mail marina.lafarga-magro@warwick.ac.uk your answers. Please make the header of the email \"MPAGS spectroscopy [your name]\". If you have any problems with the assignment let me know. Remember: there is no mark associated with the assignment, as the aim is that you engage with the topic. We will just keep note of whether you tried these questions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Questions\n",
    "\n",
    "Here you have all the questions you should try to answer. These questions are repeated again below in their corresponding sections.\n",
    "\n",
    "After opening the science image with DS9 (see below), try to answer the following questions:\n",
    "\n",
    "1. How many sources do you see in the science image?\n",
    "\n",
    "2. What are the spectral and the spatial directions?\n",
    "\n",
    "3. How was the slit oriented with respect to the image coordinates?\n",
    "\n",
    "After extracting the spectrum with the python code (again, see below), try to answer the following questions:\n",
    "\n",
    "4. Plot the 1D extracted spectrum of the sources and the sky.\n",
    "\n",
    "5. Look at the extracted sky spectra. Contrary to what we saw in the lecture, the sky emission lines are not clearly visible. The sky spectrum looks smooth and low-resolution. What does this tell you about the slit?\n",
    "\n",
    "6. Look at the extracted spectrum of the two sources. Even without having performed any flux calibration, could you tell which one of the stars is hotter? Why?\n",
    "\n",
    "7. Among the methods discussed for wavelength calibration, which one(s) would be applicable to these data?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DS9 instructions\n",
    "\n",
    "Open the science spectrum `science_spec.fits` with DS9. \n",
    "\n",
    "Try to answer the following questions:\n",
    "\n",
    "1. How many sources do you see in the science image?\n",
    "\n",
    "2. What are the spectral and the spatial directions?\n",
    "\n",
    "3. How was the slit oriented with respect to the image coordinates?\n",
    "\n",
    "We will now use DS9 to drag lines tracing the position of the spectral trails. These will be subsequently used to extract the 1D spectrum. Make sure `Edit` -> `Region` is selected from the scroll-down menu at the top, and adjust the zoom and color scaling parameters. Select `Region` -> `Shape` -> `Line` and draw a line on top of each spectral trace. Lines can be selected and moved / resized with one click. Double-clicking allows you to manually specify start and end points. When you are happy with the result, save the region file with `Region` -> `Save Regions`. In the pop-up window, make sure to select your current folder, and use `trails.reg` as file name. In the last pop-up window, select \"DS9\" as \"Format\" and \"Physical\" as \"Coordinate System\".\n",
    "\n",
    "This should result in a file named `trails.reg` in the current directory. That's it for DS9, we will now continue in this notebook.\n",
    "\n",
    "\n",
    "The following code will read in the trails you have just drawn, compute and remove the sky around each of the spectral trails, and extract the 1D spectra by co-adding the trails along the spatial dimension. After this, try plotting the extracted spectrum and answer the questions you will find below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load science data\n",
    "\n",
    "Read in the science observation `science_spec.npy`. This is the same as the `science_spec.fits` but it is ready to be read in python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = np.load('science_spec.npy')\n",
    "\n",
    "ny = img.shape[0]       # Number of pixel along spectral direction\n",
    "nx = img.shape[1]       # Number of pixel along spatial direction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can have a look at the image here. This should result in something similar as what we saw in DS9."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.imshow(img, vmax=6000)\n",
    "ax.set_title('science_spec')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load trails from DS9 and reformat\n",
    "\n",
    "Read in the region file you created with DS9 containing the trail positions `trails.reg` and reformat into pixel coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read\n",
    "with open('trails.reg','r') as f:\n",
    "    flist = f.read()\n",
    "\n",
    "# Reformat\n",
    "nlines = 2  # Two sources\n",
    "yy = np.zeros((nlines,2))\n",
    "xx = np.zeros((nlines,2))\n",
    "\n",
    "# Reformatting DS9 into pixel coordinates - skipping first 3 lines\n",
    "for iline in range(nlines):\n",
    "    line1 = flist.split('\\n')[iline+3]\n",
    "    line1 = line1.split(',')\n",
    "    xx[iline] = np.array([float(line1[0].split(\"(\")[1]), float(line1[2])], dtype='int')\n",
    "    yy[iline] = np.array([float(line1[1]), float(line1[3].split(\")\")[0])], dtype='int')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get trail position along the full spectral direction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "win = 25\n",
    "hwin = (win-1)//2\n",
    "ntrails=2\n",
    "\n",
    "trail = np.zeros((ntrails,ny))\n",
    "for itrail in range(ntrails):\n",
    "    pfit = np.poly1d(np.polyfit(yy[itrail,], xx[itrail,], 1))\n",
    "    trail[itrail,] = pfit(np.arange(ny))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract the spectra\n",
    "\n",
    "Use the trails to extract the spectra from the science image (we will do an unweighted extraction)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare arrays\n",
    "spec = np.zeros((ntrails,ny))\n",
    "sky = np.zeros((ntrails,ny))\n",
    "\n",
    "# Extract the spectra\n",
    "for itrail in range(ntrails):\n",
    "    for iy in range(ny):\n",
    "        xpos = int(trail[itrail,iy])\n",
    "        sky_i = np.nanmedian(np.array([img[iy,xpos-win:xpos-hwin].flatten(), img[iy,xpos+hwin:xpos+win].flatten()]))\n",
    "        spec_i = img[iy,xpos-hwin:xpos+hwin]\n",
    "        sky[itrail,iy] = sky_i\n",
    "        spec[itrail,iy] = (spec_i - sky_i).sum()\n",
    "\n",
    "# Save extracted data\n",
    "np.save('spec1.npy',spec[0,])\n",
    "np.save('spec2.npy',spec[1,])\n",
    "np.save('sky1.npy',sky[0,])\n",
    "np.save('sky2.npy',sky[1,])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the extracted data\n",
    "\n",
    "Try plotting the extracted spectra from the 2 sources and the extracted sky spectra. You have some example code commented out below to plot that, but try having a go yourself first. Take into account that we have not performed any wavelength calibration, so in the x-axis you will only have \"pixels\", and not \"wavelength\". Similarly with the flux, in the y-axis you will only have the number of counts observed, since we have not performed any flux calibration.\n",
    "\n",
    "Once you have the plot, try to answer the following questions (question number 4 is simply the plot):\n",
    "\n",
    "4. Plot the 1D extracted spectrum of the sources and the sky.\n",
    "\n",
    "5. Look at the extracted sky spectra. Contrary to what we saw in the lecture, the sky emission lines are not clearly visible. The sky spectrum looks smooth and low-resolution. What does this tell you about the slit?\n",
    "\n",
    "6. Look at the extracted spectrum of the two sources. Even without having performed any flux calibration, could you tell which one of the stars is hotter? Why?\n",
    "\n",
    "7. Among the methods discussed for wavelength calibration, which one(s) would be applicable to these data?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Example plots\n",
    "\n",
    "# # Plot source\n",
    "# fig, ax = plt.subplots(figsize=(10,6))\n",
    "# ax.plot(spec[0,])\n",
    "# ax.plot(spec[1,])\n",
    "# ax.set_xlabel('Pixel along spectral direction')\n",
    "# ax.set_ylabel('Total counts')\n",
    "# ax.set_title('Source')\n",
    "# plt.tight_layout()\n",
    "# plt.savefig('spec.pdf')\n",
    "# plt.show()\n",
    "\n",
    "# # Plot sky\n",
    "# fig, ax = plt.subplots(figsize=(10,6))\n",
    "# ax.plot(sky[0,])\n",
    "# ax.plot(sky[1,])\n",
    "# ax.set_xlabel('Pixel along spectral direction')\n",
    "# ax.set_ylabel('Total counts')\n",
    "# ax.set_title('Sky')\n",
    "# plt.tight_layout()\n",
    "# plt.savefig('spec.pdf')\n",
    "# plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "popurri",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}