single_rss.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6589599",
   "metadata": {},
   "source": [
    "# Quality control of WEAVE RSS files\n",
    "\n",
    "Relative wavelength calibration and throughput test for individual fibres\n",
    "\n",
    "Based on sky emission lines identified in Row-stacked spectra (RSS)\n",
    "\n",
    "# 0. Initialisation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64b59e2e-ea23-4687-bca8-0be5c0ead95c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5629f632",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib ipympl\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import colors\n",
    "from matplotlib.ticker import AutoMinorLocator\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import glob\n",
    "from time import time\n",
    "from scipy import ndimage\n",
    "\n",
    "from astropy.io import fits\n",
    "from astropy.wcs import WCS\n",
    "from astropy.coordinates import SkyCoord\n",
    "from astropy import units as u\n",
    "from astropy import constants as c\n",
    "from astropy.table import Table\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "451c1859-6722-423c-b4fa-82c06f057f98",
   "metadata": {},
   "source": [
    "Plotting functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "317e8dac-7df3-47cf-ba2f-3b41dbd9dae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def new_figure(fig_name, figsize=None, nrows=1, ncols=1, sharex='col', sharey='row', gridspec_kw={'hspace': 0, 'wspace': 0}):\n",
    "    if figsize is None:\n",
    "        figsize = (9 + ncols, 4 + nrows)\n",
    "        \n",
    "    plt.close(fig_name)\n",
    "    fig = plt.figure(fig_name, figsize=figsize)\n",
    "    axes = fig.subplots(nrows=nrows, ncols=ncols, squeeze=False,\n",
    "                        sharex=sharex, sharey=sharey,\n",
    "                        gridspec_kw=gridspec_kw\n",
    "                       )\n",
    "    fig.set_tight_layout(True)\n",
    "    for ax in axes.flat:\n",
    "        ax.xaxis.set_minor_locator(AutoMinorLocator())\n",
    "        ax.yaxis.set_minor_locator(AutoMinorLocator())\n",
    "        ax.tick_params(which='both', bottom=True, top=True, left=True, right=True)\n",
    "        ax.tick_params(which='major', direction='inout', length=8, grid_alpha=.3)\n",
    "        ax.tick_params(which='minor', direction='in', length=2, grid_alpha=.1)\n",
    "        ax.grid(True, which='both')\n",
    "\n",
    "    #fig.suptitle(f'{self.filename} {fig_name}')\n",
    "    fig.suptitle(fig_name)\n",
    "    \n",
    "    return fig, axes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ba4437-730a-41c6-866b-bdae5a708b36",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_cmap = plt.get_cmap(\"gist_earth\").copy()\n",
    "default_cmap.set_bad('gray')\n",
    "\n",
    "\n",
    "def colour_map(fig, ax, cblabel, data, cmap=default_cmap, norm=None, xlabel=None, x=None, ylabel=None, y=None):\n",
    "    \n",
    "    if norm is None:\n",
    "        percentiles = np.array([1, 16, 50, 84, 99])\n",
    "        ticks = np.nanpercentile(data, percentiles)\n",
    "        linthresh = np.median(data[data > 0])\n",
    "        norm = colors.SymLogNorm(vmin=ticks[0], vmax=ticks[-1], linthresh=linthresh)\n",
    "    else:\n",
    "        ticks = None\n",
    "    if y is None:\n",
    "        y = np.arange(data.shape[0])\n",
    "    if x is None:\n",
    "        x = np.arange(data.shape[1])\n",
    "\n",
    "    im = ax.imshow(data,\n",
    "                   extent=(x[0]-(x[1]-x[0])/2, x[-1]+(x[-1]-x[-2])/2, y[0]-(y[1]-y[0])/2, y[-1]+(y[-1]-y[-2])/2),\n",
    "                   interpolation='nearest', origin='lower',\n",
    "                   cmap=cmap,\n",
    "                   norm=norm,\n",
    "                  )\n",
    "    ax.set_aspect('auto')\n",
    "    if xlabel is not None:\n",
    "        ax.set_xlabel(xlabel)\n",
    "    if ylabel is not None:\n",
    "        ax.set_ylabel(ylabel)\n",
    "\n",
    "    cb = fig.colorbar(im, ax=ax, orientation='vertical', shrink=.9)\n",
    "    cb.ax.set_ylabel(cblabel)\n",
    "    if ticks is not None:\n",
    "        cb.ax.set_yticks(ticks=ticks, labels=[f'{value:.3g} ({percent}%)' for value, percent in zip(ticks, percentiles)])\n",
    "    cb.ax.tick_params(labelsize='small')\n",
    "    \n",
    "    return im, cb\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ba31853-36ba-4e15-90ac-1dd7a812d086",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Ancillary data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d21d052-4a64-4bfa-a0dc-adbd8418ea1e",
   "metadata": {},
   "source": [
    "UVES sky emission atlas: <https://www.eso.org/observing/dfo/quality/UVES/pipeline/sky_spectrum.html>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969b85c4-d75e-4824-8a98-0fad62e52ddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "wave_flux = np.empty((0, 2))\n",
    "filenames = glob.glob('sky/UVES_sky_emission_atlas/gident_*.dat')\n",
    "filenames.sort()\n",
    "for filename in filenames:\n",
    "    print(filename)\n",
    "    wave_flux = np.concatenate((wave_flux, np.loadtxt(filename, usecols=(1, 4), skiprows=3, comments=['#', '--------'])), axis=0)\n",
    "UVES_atlas = Table(wave_flux, names=('wavelength', 'flux'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c3def7-dd49-4227-98cf-168b0caa2ed8",
   "metadata": {},
   "outputs": [],
   "source": [
    "UVES_atlas"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6daf8566-4ef0-49f1-8bb8-45355e2a335e",
   "metadata": {},
   "source": [
    "## RSS files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d88171eb-9a70-4f79-b50e-c2ffff481cd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_lower_envelope(x, y, min_separation):\n",
    "    '''\n",
    "    Fit lower envelope of a single spectrum:\n",
    "    1) Find local minima, with a minimum separation `min_separation`.\n",
    "    2) Interpolate linearly between them.\n",
    "    '''\n",
    "    valleys = []\n",
    "    y[np.isnan(y)] = np.inf\n",
    "    for i in range(min_separation, y.size-min_separation-1):\n",
    "        if np.argmin(y[i-min_separation:i+min_separation+1]) == min_separation:\n",
    "            valleys.append(i)\n",
    "    y[~np.isfinite(y)] = np.nan\n",
    "    return np.fmin(y, np.interp(x, x[valleys], y[valleys]))\n",
    "\n",
    "\n",
    "def boxcar(data, h):\n",
    "    cumulative = np.nancumsum(data)\n",
    "    left = np.clip(np.arange(data.size) - h, 0, data.size - 1)\n",
    "    right = np.clip(np.arange(data.size) + h, 0, data.size - 1)\n",
    "    return (cumulative[right] - cumulative[left]) / (right - left)\n",
    "\n",
    "\n",
    "def find_continuum(x, y, h):\n",
    "    continuum = find_lower_envelope(x, y, h)\n",
    "    #continuum = boxcar(np.fmin(continuum, boxcar(continuum, h)), h)\n",
    "\n",
    "    offset = y - continuum\n",
    "    offset = offset[offset < 2*np.nanmedian(offset)]\n",
    "    ##continuum += .5*np.nanmedian(offset)\n",
    "    continuum += np.nanpercentile(offset, 50)\n",
    "    #continuum = find_lower_envelope(x, continuum, h)\n",
    "\n",
    "    #return lower_envelope_original + np.nanmedian(offset)\n",
    "    #continuum = boxcar(lower_envelope_original + np.nanmedian(offset), 10*h)\n",
    "    #continuum = boxcar(spectrum, h)\n",
    "    #continuum = find_lower_envelope(x, continuum, h)\n",
    "    #continuum = boxcar(continuum, 2*h)\n",
    "    #continuum = np.fmax(lower_envelope_original, continuum)\n",
    "    \n",
    "    return continuum\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a8a68d-24bf-4b16-8e32-d8725e94103d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Emission_line(object):\n",
    "    \n",
    "    def __init__(self, left=0, right=-1):\n",
    "        '''Simple class to handle emission lines (object or sky)'''\n",
    "        self.blue_index = left\n",
    "        self.red_index = right\n",
    "        self.observed_wavelength = None\n",
    "        self.reference_wavelength = None\n",
    "        self.mean_intensity = None\n",
    "        self.reference_intensity = None\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf756b5-5f45-4d30-abe1-e870dfecb6dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WEAVE_RSS(object):\n",
    "    \n",
    "    def __init__(self, filename):\n",
    "        '''Read a WEAVE \"single exposure\" file (i.e. row-stacked spectra for just one arm)'''\n",
    "        self.filename = filename\n",
    "        self.hdu = fits.open(filename)\n",
    "        self.wcs = WCS(self.hdu[1].header)\n",
    "        pixels = np.arange(self.hdu[1].data.shape[1])\n",
    "        self.wavelength = self.wcs.spectral.array_index_to_world(pixels).to_value(u.Angstrom)\n",
    "        self.counts = self.hdu[3].data\n",
    "        self.counts_error = np.where(self.hdu[4].data > 0, 1/np.sqrt(self.hdu[4].data), np.nan)\n",
    "        self.sky_counts = self.hdu[3].data - self.hdu[1].data\n",
    "        self.sensitivity_function = self.hdu[5].data\n",
    "        self.flux = self.hdu[1].data*self.sensitivity_function\n",
    "        self.sky_sub_ivar = self.hdu[2].data\n",
    "        \n",
    "        bad = np.isnan(self.counts_error).nonzero()\n",
    "        self.counts[bad] = np.nan\n",
    "        self.sky_counts[bad] = np.nan\n",
    "        self.sensitivity_function[bad] = np.nan\n",
    "        self.flux[bad] = np.nan\n",
    "        self.sky_sub_ivar[bad] = np.nan\n",
    "        \n",
    "        self.fibtable = self.hdu[6].data\n",
    "        #self.sky_fibres = np.where(self.fibtable['TARGUSE'] != 'Patata')\n",
    "        self.sky_fibres = np.where(self.fibtable['TARGUSE'] == 'S')\n",
    "        self.target_fibres = np.where(self.fibtable['TARGUSE'] == 'T')\n",
    "        self.n_sky_fibres = self.sky_fibres[0].size\n",
    "        self.n_fibres = self.counts.shape[0]\n",
    "\n",
    "        heliocentric_correction = (1 + np.nanmean(self.fibtable['Helio_cor'])/3e5)\n",
    "        wave = UVES_atlas['wavelength'] * heliocentric_correction\n",
    "        inside = np.where((wave > self.wavelength[0]) & (wave < self.wavelength[-1]))\n",
    "        self.sky_lines = UVES_atlas[inside].copy()\n",
    "        self.sky_lines['wavelength'] *= heliocentric_correction\n",
    "        \n",
    "        self.intensity = self.counts.copy()\n",
    "        self.continuum, self.strong_sky_lines = self.find_sky_lines()\n",
    "        self.f_sky = self.estimate_sky()\n",
    "        self.line_offset, self.line_throughput = self.trace_sky_lines()\n",
    "\n",
    "        self.original_offset = self.line_offset.copy()\n",
    "        self.original_throughput = self.line_throughput.copy()\n",
    "        self.original_continuum = self.continuum.copy()\n",
    "\n",
    "        for i in range(0):\n",
    "            line_weight = self.strong_sky_lines['counts']\n",
    "            total_weight = np.nansum(line_weight)\n",
    "            # TODO: flux-conserving interpolation\n",
    "            wavelength_correction = np.nansum(self.line_offset*line_weight[:, np.newaxis], axis=0) / total_weight\n",
    "            #wavelength_correction /= 2\n",
    "            for fibre in range(self.n_fibres):\n",
    "                self.intensity[fibre] = np.interp(self.wavelength,\n",
    "                                                         self.wavelength - wavelength_correction[fibre], self.intensity[fibre])        \n",
    "                self.continuum[fibre] = np.interp(self.wavelength,\n",
    "                                                         self.wavelength - wavelength_correction[fibre], self.continuum[fibre])        \n",
    "            fibre_throughput = np.nansum(self.line_throughput*line_weight[:, np.newaxis], axis=0) / total_weight\n",
    "            self.intensity /= fibre_throughput[:, np.newaxis]\n",
    "            fluxcal = np.nanmean(self.counts) / np.nanmean(self.intensity)\n",
    "            print(fluxcal, fibre_throughput[302])\n",
    "            self.intensity *= fluxcal\n",
    "            self.continuum *= fluxcal\n",
    "            self.continuum /= fibre_throughput[:, np.newaxis]\n",
    "            self.f_sky = self.estimate_sky()\n",
    "            self.line_offset, self.line_throughput = self.trace_sky_lines()\n",
    "\n",
    "        for i in range(2):\n",
    "            # TODO: flux-conserving interpolation\n",
    "            wavelength_correction = np.nanmedian(self.line_offset, axis=0)\n",
    "            for fibre in range(self.n_fibres):\n",
    "                self.intensity[fibre] = np.interp(self.wavelength,\n",
    "                                                         self.wavelength - wavelength_correction[fibre], self.intensity[fibre])        \n",
    "                self.continuum[fibre] = np.interp(self.wavelength,\n",
    "                                                         self.wavelength - wavelength_correction[fibre], self.continuum[fibre])        \n",
    "            fibre_throughput = np.nanmedian(self.line_throughput, axis=0)\n",
    "            self.intensity /= fibre_throughput[:, np.newaxis]\n",
    "            fluxcal = np.nanmean(self.counts) / np.nanmean(self.intensity)\n",
    "            print(fluxcal, fibre_throughput[302])\n",
    "            self.intensity *= fluxcal\n",
    "            self.continuum *= fluxcal\n",
    "            self.continuum /= fibre_throughput[:, np.newaxis]\n",
    "            self.f_sky = self.estimate_sky()\n",
    "            self.line_offset, self.line_throughput = self.trace_sky_lines()\n",
    "\n",
    "\n",
    "    def find_sky_lines(self, min_separation=5):\n",
    "        '''Find continuum and identify strong sky emission lines'''\n",
    "        \n",
    "        # Find continuum for all fibres:\n",
    "        print(f\"> Find continuum for {self.n_fibres} fibres:\")\n",
    "        t0 = time()\n",
    "        continuum = np.empty_like(self.counts)\n",
    "        for i, spectrum in enumerate(self.intensity):\n",
    "            continuum[i] = find_continuum(self.wavelength, spectrum, min_separation)\n",
    "        print(f\"  Done ({time()-t0:.3g} s)\")\n",
    "\n",
    "        # Identify strong sky emission lines:\n",
    "        #line_fraction = (self.intensity - continuum) / (self.intensity + continuum)\n",
    "        line_fraction = 1 - continuum/self.intensity\n",
    "        p16, p50 = np.nanpercentile(line_fraction, [16, 50])\n",
    "        line_threshold = p50 + 3*(p50-p16)\n",
    "\n",
    "        #line_mask = np.all(line_fraction > line_threshold, axis=0)\n",
    "        line_mask = np.nanmedian(line_fraction, axis=0) > line_threshold\n",
    "        line_mask[0] = False\n",
    "        line_mask[-1] = False\n",
    "\n",
    "        line_left = np.where(~line_mask[:-1] & line_mask[1:])[0]\n",
    "        line_right = np.where(line_mask[:-1] & ~line_mask[1:])[0]\n",
    "        line_right += 1\n",
    "        print(f'  {line_left.size} strong sky lines ({np.count_nonzero(line_mask)} out of {self.wavelength.size} wavelengths with line fraction > {line_threshold:.3f})')\n",
    "        min_line_width = int(2 * np.nanpercentile(line_right-line_left, 50))\n",
    "        line_left = np.fmin(line_left, (line_left + line_right - min_line_width + 1) // 2).clip(0, self.wavelength.size)\n",
    "        line_right = np.fmax(line_right, (line_left + line_right + min_line_width + 1) // 2).clip(0, self.wavelength.size)\n",
    "\n",
    "        t = Table((line_left, line_right), names=('left', 'right'))\n",
    "        t.add_column(0., name='wavelength')\n",
    "        t.add_column(0., name='counts')\n",
    "        t.add_column(0., name='reference_wavelength')\n",
    "        t.add_column(0., name='reference_flux')\n",
    "        for line in t:\n",
    "            wavelength = self.wavelength[line['left']:line['right']]\n",
    "            inside = np.where((self.sky_lines['wavelength'] >= wavelength[0])\n",
    "                              & (self.sky_lines['wavelength'] < wavelength[-1]))\n",
    "            line['reference_flux'] = np.nansum(self.sky_lines[inside]['flux'])\n",
    "            line['reference_wavelength'] = np.nansum(self.sky_lines[inside]['flux'] * self.sky_lines[inside]['wavelength'])\n",
    "            line['reference_wavelength'] /= line['reference_flux']\n",
    "        return continuum, t[t['reference_flux'] > 0]\n",
    "        \n",
    "        \n",
    "    def estimate_sky(self, max_bins=101):\n",
    "        '''Linear fit at every wavelength for the spaxels with fainter flux'''\n",
    "        print(f\"> Estimating sky:\")\n",
    "        t0 = time()\n",
    "\n",
    "        flux = np.nanmean(self.intensity, axis=1)\n",
    "        flux -= np.nanmin(flux)\n",
    "        faint = np.where(flux < 2*np.nanmedian(flux))[0]\n",
    "\n",
    "        f_sky = np.zeros_like(self.wavelength)\n",
    "        for i in range(self.wavelength.size):\n",
    "            valid = np.isfinite(self.intensity[faint, i])\n",
    "            if np.count_nonzero(valid):\n",
    "                slope, f_sky[i] = np.polyfit(flux[faint][valid], self.intensity[faint, i][valid], 1)\n",
    "                if np.isnan(f_sky[i]) or slope < 0:\n",
    "                    f_sky[i] = np.nanmean(self.intensity[faint, i])\n",
    "            else:\n",
    "                f_sky[i] = np.nan\n",
    "\n",
    "        print(f\"  Done ({time()-t0:.3g} s)\")\n",
    "        return f_sky\n",
    "\n",
    "\n",
    "    def trace_sky_lines(self, min_separation=5):\n",
    "        '''Trace the wavelength and intensity of sky emission lines'''\n",
    "\n",
    "        reference_spectrum = self.f_sky.copy()\n",
    "        reference_spectrum -= find_continuum(self.wavelength, reference_spectrum, min_separation)\n",
    "        \n",
    "        for line in self.strong_sky_lines:\n",
    "            wavelength = self.wavelength[line['left']:line['right']]\n",
    "            spectrum = reference_spectrum[line['left']:line['right']]\n",
    "            weight = spectrum\n",
    "            #weight = spectrum**2\n",
    "            #weight = spectrum**3\n",
    "            line['wavelength'] = np.nansum(weight*wavelength) / np.nansum(weight)\n",
    "            line['counts'] = np.nanmean(spectrum)\n",
    "            #line['counts'] = np.nanstd(spectrum)\n",
    "\n",
    "        #line_weight = self.strong_sky_lines['counts']\n",
    "        #total_weight = np.nansum(line_weight)\n",
    "        \n",
    "        # Trace line wavelengths for every fibre to compare with the reference spectrum:\n",
    "\n",
    "        line_fibre_wavelength = []\n",
    "        line_fibre_intensity = []\n",
    "        for line in self.strong_sky_lines:\n",
    "            left = line['left']\n",
    "            right = line['right']\n",
    "            weight = (self.intensity[:, left:right] - self.continuum[:, left:right])\n",
    "            #weight = self.intensity[:, left:right]**2\n",
    "            line_fibre_wavelength.append(np.nansum(weight * self.wavelength[np.newaxis, left:right], axis=1) / np.nansum(weight, axis=1))\n",
    "            line_fibre_intensity.append(np.nanmean(self.intensity[:, left:right] - self.continuum[:, left:right], axis=1))\n",
    "            #line_fibre_intensity.append(np.nanstd(self.intensity[:, left:right] - self.continuum[:, left:right], axis=1))\n",
    "            \n",
    "        line_offset = np.array(line_fibre_wavelength) - self.strong_sky_lines['wavelength'][:, np.newaxis]\n",
    "        line_offset -= np.nanmedian(line_offset, axis=1)[:, np.newaxis]\n",
    "        #line_offset -= np.nanmean(line_offset, axis=1)[:, np.newaxis]\n",
    "        #line_offset -= (np.nansum(line_offset*line_weight[:, np.newaxis], axis=1)/total_weight)[:, np.newaxis]\n",
    "\n",
    "        line_throughput = np.array(line_fibre_intensity) / self.strong_sky_lines['counts'][:, np.newaxis]\n",
    "        line_throughput /= np.nanmedian(line_throughput)\n",
    "        #line_throughput /= np.nanmedian(line_throughput, axis=1)[:, np.newaxis]\n",
    "        #line_throughput /= np.nanmean(line_throughput, axis=1)[:, np.newaxis]\n",
    "        #line_throughput /= (np.nansum(line_throughput*line_weight[:, np.newaxis], axis=1)/total_weight)[:, np.newaxis]\n",
    "        \n",
    "        return line_offset, line_throughput\n",
    "        \n",
    "        \n",
    "#rss = WEAVE_RSS('data/v3/casu.ast.cam.ac.uk/weavedata/1500016316/L1/single_3039342.fit')\n",
    "#rss = WEAVE_RSS('data/v3/casu.ast.cam.ac.uk/weavedata/1500016658/L1/single_3042890.fit')\n",
    "rss = WEAVE_RSS('data/v3/casu.ast.cam.ac.uk/weavedata/1500016520/L1/single_3045973.fit')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d364558f-6986-4b6b-8c63-a8fa5b08882b",
   "metadata": {},
   "source": [
    "# QC tests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b9ee3ce-3a9d-4462-95cd-c9eb689ddca1",
   "metadata": {},
   "source": [
    "## Wavelength correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb775d99-b574-4ecc-96a5-06d93b898032",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_relative_wavelength_correction(self):\n",
    "    '''Relative line offset, before and after correction'''\n",
    "    fig, axes = new_figure('relative_wavelength_correction', nrows=4)\n",
    "    \n",
    "    \n",
    "    ax = axes[0, 0]\n",
    "    ax.set_ylabel('original $\\Delta\\lambda$ [$\\AA$]')\n",
    "    ax.set_ylim(-.25, .25)\n",
    "\n",
    "    p16, p50, p84 = np.nanpercentile(self.original_offset, [16, 50, 84], axis=0)\n",
    "    ax.plot(p50, 'k-', alpha=.5)\n",
    "    ax.fill_between(np.arange(self.n_fibres), p16, p84, color='k', alpha=.1)\n",
    "    \n",
    "    #for line in self.sky_fibres[0]:\n",
    "    #    ax.axvline(line, c='b', ls='-', alpha=.2)\n",
    "    cb = fig.colorbar(None, ax=ax)\n",
    "    cb.remove()\n",
    "\n",
    "    \n",
    "    ax = axes[1, 0]\n",
    "    ax.set_ylabel('new $\\Delta\\lambda$ [$\\AA$]')\n",
    "    ax.set_ylim(-.25, .25)\n",
    "    ax.get_shared_y_axes().join(ax, axes[0, 0])\n",
    "\n",
    "    p16, p50, p84 = np.nanpercentile(self.line_offset, [16, 50, 84], axis=0)\n",
    "    ax.plot(p50, 'k-', alpha=.5)\n",
    "    ax.fill_between(np.arange(self.n_fibres), p16, p84, color='k', alpha=.1)\n",
    "    cb = fig.colorbar(None, ax=ax)\n",
    "    cb.remove()\n",
    "\n",
    "    \n",
    "    ax = axes[2, 0]\n",
    "    ax.set_ylim(-1.5, len(self.strong_sky_lines)+.5)\n",
    "    im, cb = colour_map(fig, ax, 'original $\\Delta\\lambda$ [$\\AA$]', self.original_offset,\n",
    "                        xlabel='fibre', ylabel='line ID (increasing $\\lambda$)', cmap='turbo', norm=colors.Normalize(vmin=-.2, vmax=.2))\n",
    "\n",
    "    ax = axes[3, 0]\n",
    "    ax.set_ylim(-1.5, len(self.strong_sky_lines)+.5)\n",
    "    ax.get_shared_y_axes().join(ax, axes[2, 0])\n",
    "    im, cb = colour_map(fig, ax, 'new $\\Delta\\lambda$ [$\\AA$]', self.line_offset,\n",
    "                        xlabel='fibre', ylabel='line ID (increasing $\\lambda$)', cmap='turbo', norm=colors.Normalize(vmin=-.2, vmax=.2))\n",
    "\n",
    "    plt.savefig(f'{self.filename}-{fig.get_label()}.pdf')\n",
    "\n",
    "\n",
    "test_relative_wavelength_correction(rss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a7818b-d453-450a-894a-37e2800a9fe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_absolute_wavelength_calibration(self):\n",
    "    '''Compare measured wavelength of sky lines with UVES atlas'''\n",
    "    fig, axes = new_figure('absolute_wavelength_calibration', ncols=2)\n",
    "\n",
    "    \n",
    "    ax = axes[0, 0]\n",
    "    ax.set_ylabel(r'$\\Delta\\lambda\\ [\\AA]$')\n",
    "    ax.set_xlabel(r'wavelength [$\\AA$]')\n",
    "    \n",
    "    ratio = self.strong_sky_lines['wavelength'] - self.strong_sky_lines['reference_wavelength']\n",
    "    really_strong = np.where(self.strong_sky_lines['reference_flux'] > np.nanmean(self.strong_sky_lines['reference_flux']))\n",
    "    ax.plot(self.strong_sky_lines['wavelength'][really_strong], ratio[really_strong], 'co')\n",
    "    ax.plot(self.strong_sky_lines['wavelength'], ratio, 'k+')\n",
    "    \n",
    "    #p16_o, p50_o, p84_o = np.nanpercentile(self.line_offset, [16, 50, 84], axis=1)\n",
    "    #ax.errorbar(self.strong_sky_lines['wavelength'], p50_o, yerr=(np.abs(p16_o), p84_o), fmt='rx')\n",
    "    #ax.fill_between(self.strong_sky_lines['wavelength'], p16_o, p84_o, color='k', alpha=.1)\n",
    "\n",
    "    p16, p50, p84 = np.nanpercentile(ratio, [16, 50, 84])\n",
    "    ax.axhline(p50, c='k', ls=':', alpha=.5)\n",
    "    ax.fill_between(self.strong_sky_lines['wavelength'], p16, p84, color='k', alpha=.1)\n",
    "\n",
    "    \n",
    "    ax = axes[0, 1]\n",
    "    ax.set_xlabel(r'counts')\n",
    "    \n",
    "    ax.plot(self.strong_sky_lines['counts'][really_strong], ratio[really_strong], 'co')\n",
    "    ax.plot(self.strong_sky_lines['counts'], ratio, 'k+')\n",
    "    #ax.errorbar(self.strong_sky_lines['counts'], p50_o, yerr=(np.abs(p16_o), p84_o), fmt='rx')\n",
    "\n",
    "    ax.axhline(p50, c='k', ls=':', alpha=.5, label=f'{p50:.2f}')\n",
    "    ax.fill_between(np.sort(self.strong_sky_lines['counts']), p16, p84, color='k', alpha=.1, label=f'{p16:.2f} - {p84:.2f}')\n",
    "    ax.legend()\n",
    "    \n",
    "test_absolute_wavelength_calibration(rss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3066df4a-a0b7-4a90-adc8-d84ae29b7c54",
   "metadata": {},
   "source": [
    "## Fibre throughput"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d3518b-63eb-4511-870e-f52cdad2f62f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_fibre_throughput(self):\n",
    "    '''Fibre throughput, before and after correction'''\n",
    "    fig, axes = new_figure('fibre_throughput', nrows=4)\n",
    "\n",
    "    \n",
    "    ax = axes[0, 0]\n",
    "    ax.set_ylabel('original throughput')\n",
    "    ax.set_ylim(.75, 1.25)\n",
    "\n",
    "    p16, p50, p84 = np.nanpercentile(self.original_throughput, [16, 50, 84], axis=0)\n",
    "    ax.plot(p50, 'k-', alpha=.5)\n",
    "    ax.fill_between(np.arange(self.n_fibres), p16, p84, color='k', alpha=.1)\n",
    "    \n",
    "    #for line in self.sky_fibres[0]:\n",
    "    #    ax.axvline(line, c='b', ls='-', alpha=.2)\n",
    "    cb = fig.colorbar(None, ax=ax)\n",
    "    cb.remove()\n",
    "\n",
    "    \n",
    "    ax = axes[1, 0]\n",
    "    ax.set_ylabel('new throughput')\n",
    "    ax.set_ylim(.75, 1.25)\n",
    "    ax.get_shared_y_axes().join(ax, axes[0, 0])\n",
    "\n",
    "    p16, p50, p84 = np.nanpercentile(self.line_throughput, [16, 50, 84], axis=0)\n",
    "    ax.plot(p50, 'k-', alpha=.5)\n",
    "    ax.fill_between(np.arange(self.n_fibres), p16, p84, color='k', alpha=.1)\n",
    "    cb = fig.colorbar(None, ax=ax)\n",
    "    cb.remove()\n",
    "\n",
    "    \n",
    "    ax = axes[2, 0]\n",
    "    ax.set_ylim(-1.5, len(self.strong_sky_lines)+.5)\n",
    "    im, cb = colour_map(fig, ax, 'original throughput', self.original_throughput,\n",
    "                        xlabel='fibre', ylabel='line ID (increasing $\\lambda$)', cmap='turbo', norm=colors.Normalize(vmin=.8, vmax=1.2))\n",
    "\n",
    "    ax = axes[3, 0]\n",
    "    ax.set_ylim(-1.5, len(self.strong_sky_lines)+.5)\n",
    "    ax.get_shared_y_axes().join(ax, axes[2, 0])\n",
    "    im, cb = colour_map(fig, ax, 'new throughput', self.line_throughput,\n",
    "                        xlabel='fibre', ylabel='line ID (increasing $\\lambda$)', cmap='turbo', norm=colors.Normalize(vmin=.8, vmax=1.2))\n",
    "\n",
    "    plt.savefig(f'{self.filename}-{fig.get_label()}.pdf')\n",
    "\n",
    "\n",
    "test_fibre_throughput(rss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7819d956-7adf-4ec8-a9b1-35b17d0e96c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_throughput_map(self):\n",
    "    '''Spatial distribution and dependence on total flux'''\n",
    "    flux = np.nanmean(self.intensity, axis=1)\n",
    "    fig, axes = new_figure('throughput_map', nrows=2)\n",
    "    \n",
    "    ax = axes[0, 0]\n",
    "    ax.set_ylabel(r'original throughput')\n",
    "    ax.set_ylim(.65, 1.55)\n",
    "    p16, p50, p84 = np.nanpercentile(self.original_throughput, [16, 50, 84], axis=0)\n",
    "    ax.errorbar(flux, p50, yerr=(p50-p16, p84-p50), fmt='k+', alpha=.2)\n",
    "    \n",
    "    ax = axes[1, 0]\n",
    "    ax.set_ylabel(r'new throughput')\n",
    "    ax.set_ylim(.65, 1.55)\n",
    "    #ax.set_ylim(.8, 1.2)\n",
    "    #ax.set_yscale('log')\n",
    "    p16, p50, p84 = np.nanpercentile(self.line_throughput, [16, 50, 84], axis=0)\n",
    "    ax.errorbar(flux, p50, yerr=(p50-p16, p84-p50), fmt='k+', alpha=.2)\n",
    "\n",
    "    ax.set_xlabel(r'mean intensity')\n",
    "    ax.set_xscale('log')\n",
    "\n",
    "\n",
    "test_throughput_map(rss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e20ac0-09fe-4f6c-a715-d69771efafe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#rss.strong_sky_lines.show_in_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d214971-1dc9-403a-b545-d572d7385d1c",
   "metadata": {},
   "source": [
    "## Sky subtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "267305ee-f7bb-4255-a55d-1f4dee854cba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_sky_spectrum(self):\n",
    "    '''Compare different methods to estimate the sky spectrum. Show average subtracted spectrum, as well as fibre with maximuum signal.'''\n",
    "    fig, axes = new_figure('sky_spectrum', nrows=3)#, figsize=(10, 8))\n",
    "\n",
    "    ax = axes[0, 0]\n",
    "    ax.set_ylabel('sky counts')\n",
    "    #ax.set_yscale('log')\n",
    "    ax.plot(self.wavelength, self.sky_counts[0], 'r-', alpha=.5, label='original sky')\n",
    "    ax.plot(self.wavelength, self.f_sky, 'k-', alpha=1, label='linear fit')\n",
    "    ax.legend()\n",
    "\n",
    "    ax = axes[1, 0]\n",
    "    ax.set_ylabel('mean spectrum')\n",
    "    ax.plot(self.wavelength, np.nanmean(self.counts, axis=0) - self.sky_counts[0], 'r-', alpha=.5, label='original mean spectrum')\n",
    "    ax.plot(self.wavelength, np.nanmean(self.intensity, axis=0) - self.f_sky, 'k-', alpha=1, label='new mean spectrum')\n",
    "\n",
    "    ax = axes[2, 0]\n",
    "    peak = np.nanargmax(np.nanmean(self.intensity, axis=1))\n",
    "    ax.set_ylabel(f'fibre {peak} (peak)')\n",
    "    ax.plot(self.wavelength, self.counts[peak] - self.sky_counts[0], 'r-', alpha=.5, label=f'fibre {peak} (original)')\n",
    "    ax.plot(self.wavelength, self.intensity[peak] - self.f_sky, 'k-', alpha=1, label=f'fibre {peak} (new)')\n",
    "\n",
    "    ax.set_xlabel(r'wavelength [$\\AA$]')\n",
    "    for ax in axes.ravel():\n",
    "        for line in self.strong_sky_lines:\n",
    "            ax.axvspan(self.wavelength[line['left']], self.wavelength[line['right']], color='b', alpha=.1)\n",
    "\n",
    "    plt.savefig(f'{self.filename}-{fig.get_label()}.pdf')\n",
    "\n",
    "\n",
    "test_sky_spectrum(rss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d438007d-6b8a-49b7-8db3-9073e439a75f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_sky_map(self):\n",
    "    '''Compare different sky estimates'''\n",
    "    fig, axes = new_figure('sky_map', nrows=5)\n",
    "\n",
    "    ax = axes[0, 0]\n",
    "    ax.set_ylabel('fibre')\n",
    "    im, cb = colour_map(fig, ax, 'original counts', self.counts, xlabel='', x=self.wavelength)\n",
    "\n",
    "    ax = axes[1, 0]\n",
    "    ax.set_ylabel('fibre')\n",
    "    im, cb = colour_map(fig, ax, 'original subtraction', self.counts-self.sky_counts[0], xlabel='', x=self.wavelength, norm=cb.norm)\n",
    "    ax.get_shared_y_axes().join(ax, axes[0, 0])\n",
    "\n",
    "    ax = axes[2, 0]\n",
    "    ax.set_ylabel('fibre')\n",
    "    im, cb = colour_map(fig, ax, 'new subtraction', self.intensity-self.f_sky, xlabel='', x=self.wavelength, norm=cb.norm)\n",
    "    ax.get_shared_y_axes().join(ax, axes[0, 0])\n",
    "\n",
    "    ax = axes[3, 0]\n",
    "    ax.set_ylabel('fibre')\n",
    "    im, cb = colour_map(fig, ax, 'new counts', self.intensity, xlabel='', x=self.wavelength, norm=cb.norm)\n",
    "    ax.get_shared_y_axes().join(ax, axes[0, 0])\n",
    "\n",
    "    ax = axes[4, 0]\n",
    "    ax.set_ylabel('fibre')\n",
    "    im, cb = colour_map(fig, ax, 'ratio', self.intensity/self.counts, xlabel='', x=self.wavelength, cmap='turbo')\n",
    "    ax.get_shared_y_axes().join(ax, axes[0, 0])\n",
    "\n",
    "    ax.set_xlabel(r'wavelength [$\\AA$]')\n",
    "    plt.savefig(f'{self.filename}-{fig.get_label()}.pdf')\n",
    "\n",
    "\n",
    "test_sky_map(rss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b879932-51e1-4dca-82c4-c1ca90a94fca",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Illustrative examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "170a0d22-b98d-4f9b-a67e-d35b162147f7",
   "metadata": {},
   "source": [
    "## Line-continuum separation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bb4088e-1505-467d-bf27-2195cc7e21a8",
   "metadata": {},
   "source": [
    "Use the sky spectrum as an illustrative example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b00ec8c3-563f-4de1-84e7-b6b103d0db50",
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrum = rss.f_sky\n",
    "fibre = 497\n",
    "spectrum = rss.counts[fibre]\n",
    "\n",
    "\n",
    "#lower_envelope_smooth = boxcar(np.fmax(lower_envelope_original, find_lower_envelope(rss.wavelength, smooth, h)), h)\n",
    "\n",
    "#continuum = boxcar(np.fmax(lower_envelope_original, find_lower_envelope(rss.wavelength, np.fmax(lower_envelope_original, lower_envelope_smooth), h)), h)\n",
    "#continuum = (lower_envelope_original + lower_envelope_smooth) / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e408b5af-2227-4092-bef1-3ba72f896f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = new_figure('line_continuum_separation')\n",
    "\n",
    "ax = axes[0, 0]\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(30, 3e4)\n",
    "\n",
    "ax.plot(rss.wavelength, spectrum, 'b-', alpha=.2, label='spectrum')\n",
    "#ax.plot(rss.wavelength, smooth, 'k:', alpha=1, label='boxcar')\n",
    "#ax.plot(rss.wavelength, lower_envelope_smooth, 'r--', alpha=1, label='lower envelope smooth')\n",
    "#ax.plot(rss.wavelength, lower_envelope_original, 'k--', alpha=1, label='lower envelope original')\n",
    "#ax.plot(rss.wavelength, spectrum - continuum, 'k-', alpha=.2, label='line emission')\n",
    "\n",
    "ax.plot(rss.wavelength, find_continuum(rss.wavelength, spectrum, 5), 'k-', alpha=1, label='continuum')\n",
    "ax.plot(rss.wavelength, find_lower_envelope(rss.wavelength, spectrum, 5), 'k--', alpha=1, label='lower envelope')\n",
    "\n",
    "'''\n",
    "ax.plot(rss.wavelength, rss.counts[fibre], 'r:', alpha=1, label='old intensity')\n",
    "ax.plot(rss.wavelength, rss.intensity[fibre], 'r-', alpha=1, label='new intensity')\n",
    "ax.plot(rss.wavelength, rss.original_continuum[fibre], 'r:', alpha=1, label='old continuum')\n",
    "ax.plot(rss.wavelength, rss.continuum[fibre], 'r--', alpha=1, label='new continuum')\n",
    "'''\n",
    "\n",
    "ax.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe1afab-8b32-40d3-98ba-d6c38e7a1511",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = new_figure('line_emission_map', nrows=4, sharey=True)\n",
    "\n",
    "total_cont = np.sum(rss.continuum, axis=1)\n",
    "x=rss.wavelength\n",
    "x=np.arange(rss.wavelength.size)\n",
    "\n",
    "ax = axes[0, 0]\n",
    "im, cb = colour_map(fig, ax, 'continuum', rss.continuum, x=x)\n",
    "\n",
    "ax = axes[1, 0]\n",
    "im, cb = colour_map(fig, ax, 'counts', rss.counts, x=x)\n",
    "\n",
    "ax = axes[2, 0]\n",
    "im, cb = colour_map(fig, ax, 'line emission', rss.intensity-rss.continuum, x=x)\n",
    "\n",
    "ax = axes[3, 0]\n",
    "im, cb = colour_map(fig, ax, 'line ratio', 1-rss.continuum/rss.intensity, x=x)\n",
    "for line in rss.strong_sky_lines:\n",
    "    ax.axvline(line['left'], c='w', ls='--')\n",
    "    ax.axvline(line['right'], c='k', ls='--')    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf8567e7-ad9c-4de5-969f-45b645b536c7",
   "metadata": {},
   "source": [
    "## Sky estimation from linear fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1432833c-44f8-4f08-a05f-cab9b7917926",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_cont = np.nansum(rss.intensity, axis=1)\n",
    "total_cont -= np.nanmin(total_cont)\n",
    "f_sky = np.zeros_like(rss.wavelength)\n",
    "f_obj = np.zeros_like(rss.wavelength)\n",
    "\n",
    "faint = np.where(total_cont < 2*np.nanmedian(total_cont))[0]\n",
    "for i in range(rss.wavelength.size):\n",
    "    valid = np.isfinite(rss.intensity[faint, i])\n",
    "    if np.count_nonzero(valid):\n",
    "        f_obj[i], f_sky[i] = np.polyfit(total_cont[faint][valid], rss.intensity[faint, i][valid], 1)\n",
    "        if np.isnan(f_sky[i]) or f_obj[i] < 0:\n",
    "            #print(np.count_nonzero(np.isnan(total_cont[faint])), np.count_nonzero(np.isnan(rss.intensity[faint, i])), ' are nan')\n",
    "            f_sky[i] = np.nanmedian(rss.intensity[faint, i])\n",
    "            f_obj[i] = 0\n",
    "    else:\n",
    "        f_sky[i] = np.nan\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a2e12e-e259-4d8b-8c57-b4ddad1abfad",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = new_figure('linear fit', nrows=1)\n",
    "\n",
    "ax = axes[0, 0]\n",
    "pix = 0\n",
    "ax.plot(total_cont, rss.counts[:, pix], 'c.')\n",
    "ax.plot(total_cont[faint], rss.counts[faint, pix], 'r.')\n",
    "ax.plot(total_cont[faint], f_sky[pix] + total_cont[faint]*f_obj[pix], 'k+')\n",
    "\n",
    "plt.show()\n",
    "print(rss.wavelength[pix], f_sky[pix])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c6247a7-997c-4ada-90d4-96d9a15cea4c",
   "metadata": {},
   "source": [
    "## Line Spread Function (LSF)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a4320bc-fe83-4ff3-83ed-b26a700ef04c",
   "metadata": {},
   "source": [
    "Parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01e601b8-5ed3-4f1f-b3da-e69750d365fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "LSF_requested_resolution = 0.01  # Angstrom\n",
    "LSF_wavelength_range = 20  # Angstrom\n",
    "LSF_dlambda = np.arange(-LSF_wavelength_range, LSF_wavelength_range+.5*LSF_requested_resolution, LSF_requested_resolution)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6402fe73-432e-417d-8d90-02edf3c34870",
   "metadata": {},
   "source": [
    "Function definitions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e3b6f51-4e2f-4b24-8dd0-0d2fee9af5d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalise(x):\n",
    "    x -= np.median(x)\n",
    "    norm = x[x.size//2]\n",
    "    if norm > 0:\n",
    "        x /= norm\n",
    "    else:\n",
    "        x *= np.nan\n",
    "    return x\n",
    "\n",
    "\n",
    "def find_LSF(delta_l, spectrum):\n",
    "    median_skyline = np.zeros((rss.sky_lines['wavelength'].size, delta_l.size))\n",
    "    for i, line in enumerate(rss.sky_lines['wavelength']):\n",
    "        sed = np.interp(line+delta_l, rss.wavelength, spectrum)\n",
    "        sed = normalise(sed)\n",
    "        median_skyline[i] = sed\n",
    "    return normalise(np.nanmedian(median_skyline, axis=0))\n",
    "\n",
    "\n",
    "def find_FWHM(delta_l, profile):\n",
    "    threshold = np.max(profile)/2\n",
    "    left = np.max(delta_l[(delta_l < 0) & (profile < threshold)])\n",
    "    right = np.min(delta_l[(delta_l > 0) & (profile < threshold)])\n",
    "    return right-left\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c5a7e2-3bec-4d0c-be42-b32677241c77",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_profile(x, mu=0, sigma=1, norm=True):\n",
    "    g = np.exp(-.5 * ((x-mu) / sigma)**2)\n",
    "    if norm:\n",
    "        g /= np.sqrt(2*np.pi) * sigma\n",
    "    return  g\n",
    "\n",
    "\n",
    "def refine_Gaussian(x, I, mu0, sigma0):\n",
    "\n",
    "    good = np.where(np.isfinite(I))\n",
    "    weight = np.exp(-.5 * ((x[good]-mu0) / sigma0)**2) * (I[good] - np.min(I[good]))\n",
    "    total_weight = np.sum(weight)\n",
    "\n",
    "    ivar0 = 1 / sigma0**2\n",
    "    mu = np.sum(weight * x[good]) / total_weight\n",
    "    ivar = total_weight / np.sum(weight * (x[good] - mu)**2)\n",
    "\n",
    "    ivar1 = ivar - ivar0\n",
    "    mu1 = (mu*ivar - mu0*ivar0) / ivar1\n",
    "    return mu1, 1/np.sqrt(ivar1)\n",
    "\n",
    "\n",
    "def fit_Gaussian(x, I, mu0=None, sigma0=None):\n",
    "\n",
    "    if mu0 is None:\n",
    "        mu0 = np.nanmean(x)\n",
    "    if sigma0 is None:\n",
    "        sigma0 = np.nanstd(x)\n",
    "    \n",
    "    delta2 = np.inf\n",
    "    while delta2 > LSF_requested_resolution**2:\n",
    "        mu1, sigma1 = refine_Gaussian(x, I, mu0, sigma0)\n",
    "        delta2 = (mu1 - mu0)**2 + (sigma1 - sigma0)**2\n",
    "        mu0 = mu1\n",
    "        sigma0 = sigma1\n",
    "        #print(mu0, sigma0)\n",
    "    \n",
    "    return mu0, sigma0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3da3769-dda7-45ef-ae5a-646daa7e637f",
   "metadata": {},
   "source": [
    "FWHM of mode-based sky:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55057686-31b1-4733-b998-9ae96b62181d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_LSF = find_LSF(LSF_dlambda, rss.f_sky)\n",
    "sky_FWHM = find_FWHM(LSF_dlambda, sky_LSF)\n",
    "#LSF_sigma = sky_FWHM / np.sqrt(8*np.log(2))\n",
    "\n",
    "LSF_mu, LSF_sigma = fit_Gaussian(LSF_dlambda, sky_LSF, 0, sky_FWHM / np.sqrt(8*np.log(2)))\n",
    "\n",
    "print(f'FWHM of mode-based sky = {sky_FWHM:.3f} (sigma = {sky_FWHM / np.sqrt(8*np.log(2)):.3f}) Angstrom')\n",
    "print(f'Moments (mu, sigma) = ({LSF_mu:.3f}, {LSF_sigma:.3f}) Angstrom')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf083024-8f61-4276-b14c-1afc94186aaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = new_figure('sky_LSF')\n",
    "\n",
    "    \n",
    "ax = axes[0, 0]\n",
    "ax.set_ylabel(r'Median Line Spread Function (LSF)')\n",
    "#ax.set_yscale('log')\n",
    "#ax.set_ylim(5e-4, 2)\n",
    "\n",
    "\n",
    "#ax.plot(LSF_dlambda, np.nancumsum(sky_LSF), 'k--', alpha=.5)\n",
    "#ax.plot(LSF_dlambda, np.nancumsum(sky_LSF*LSF_dlambda)*np.nancumsum(sky_LSF), 'b-', alpha=.5)\n",
    "ax.plot(LSF_dlambda, sky_LSF, 'k-', alpha=1, label=f'FWHM = {sky_FWHM:.3f} (sigma = {sky_FWHM / np.sqrt(8*np.log(2)):.3f}) $\\\\AA$ ')\n",
    "ax.plot(LSF_dlambda, gaussian_profile(LSF_dlambda, LSF_mu, LSF_sigma, False), 'k--', alpha=1, label=f'Gaussian ($\\\\mu, \\\\sigma$) = ({LSF_mu:.3f}, {LSF_sigma:.3f}) $\\\\AA$ ')\n",
    "ax.axvline(LSF_mu, c='k', ls=':')\n",
    "ax.legend()\n",
    "\n",
    "ax = axes[-1, 0]\n",
    "ax.set_xlabel(r'$(\\lambda - \\lambda_0)$ [$\\AA$]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ffc864-e650-4c9d-b14e-6884c7788764",
   "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.8.10"
  },
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}