{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wavelength calibration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to obtain the wavelength calibration for an arc spectrum, provided we already have a list (the 'master list') with the wavelengths of all the possible lines, and a rough idea of the wavelength range covered by that spectrum. \n",
    "\n",
    "Here I am going to explain how this wavelength calibration can be carried out using for that purpose a pattern matching based on line triplets. Note that a triplet is characterized by the relative location of the middle line relative to the other two lines. In this way, two triplets can look similar to each other if that relative location is not very different. But as soon as the middle line appears in a different relative position, the two triplets can readily be discarded as being similar. \n",
    "\n",
    "The idea behind the method here described is the following:\n",
    "\n",
    "- The potential wavelengths of the arc lines must be available in a master list, which should cover the full wavelength range of the observed arc spectra.\n",
    "\n",
    "- Given a particular observed arc spectrum, the coordinates of the line peaks (in pixel detector units) can be easily determined. Using these line peak coordinates, line triplets built from consecutive lines are going to be considered. Note that the lines must be consecutive in order to guanrantee that lineary deviations are not important when computing the relative position of the central linea of each line triplet. This means that, if `nlines_arc` is the observed number of arc lines in the spectrum, `nlines_arc-2` is the total number of line triplets considered.\n",
    "\n",
    "- Since each triplet from the observed arc spectrum is characterized by the relative location of the central line, triplets built from lines in the master list, with a similar relative location of the central line, are potential matches. Note that in the case of the master list, the triplets are not restricted to consecutive lines, since this master list can contain faint lines that may be undetected in the observed spectrum. As expected, the potential number of triplets built from the master list can be very large. So, to each triplet in the observed arc spectrum we can associate a bunch of triplets from the master list.\n",
    "\n",
    "- Each triplet from the master list provides a solution CRVAL1, CDELT1 for the wavelength calibration. The method is based on the assumption that, after plotting all the possible solutions in a diagram CDELT1 - CRVAL1, the erroneous solutions will be scattered in this diagram, while one should find an accumulation of points around the correct values.\n",
    "\n",
    "For the code of this notebook to be self-contained, we start with the required initialization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import sys\n",
    "import numpy as np\n",
    "from numpy.polynomial import polynomial\n",
    "from scipy.stats import norm\n",
    "import scipy.misc\n",
    "import matplotlib.pyplot as plt\n",
    "import itertools\n",
    "\n",
    "ldebug = True  # display intermediate information\n",
    "lplot = True   # display intermediate plots\n",
    "lpause = False  # pause execution waiting for the user\n",
    "\n",
    "def pause(lpause):\n",
    "    if lpause:\n",
    "        raw_input('press <RETURN> to continue...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a master list and generating a list of potential triplets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this notebook, we are going to create a synthetic master list covering a wavelength interval ranging from `wv_ini_master` to `wv_end_master`, containing a total of `nlines_master` lines. The resulting collection of master lines will be stored in a numpy array named `wv_master`. An auxiliary function `simulate_master_table` is defined to help in this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_master_table(my_seed, wv_ini_master, wv_end_master, nlines_master,\n",
    "                          ldebug=False):\n",
    "    \"\"\"Generates a simulated master table of wavelengths.\n",
    "\n",
    "    The location of the lines follows a random uniform distribution\n",
    "    between `wv_ini_master` and `wv_end_master`.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    my_seed : int\n",
    "        Seed to re-initialize random number generation.\n",
    "    wv_ini_master : float\n",
    "        Minimum wavelength in master table.\n",
    "    wv_end_master : float\n",
    "        Maximum wavelength in master table.\n",
    "    nlines_master : int\n",
    "        Total number of lines in master table.\n",
    "    ldebug : bool\n",
    "        If True intermediate results are displayed.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    wv_master : 1d numpy array, float\n",
    "        Array with wavelengths corresponding to the master table (Angstroms).\n",
    "    \"\"\"\n",
    "    if my_seed is not None:\n",
    "        np.random.seed(my_seed)\n",
    "        \n",
    "    if wv_end_master < wv_ini_master:\n",
    "        raise ValueError('wv_ini_master=' + str(wv_ini_master) +\n",
    "                         ' must be <= wv_end_master=' + str(wv_end_master))\n",
    " \n",
    "    wv_master = np.random.uniform(low=wv_ini_master,\n",
    "                                  high=wv_end_master,\n",
    "                                  size=nlines_master)\n",
    "    wv_master.sort()  # in-place sort\n",
    "\n",
    "    if ldebug:\n",
    "        print('>>> Master table:')\n",
    "        for val in zip(range(nlines_master), wv_master):\n",
    "            print(val)\n",
    "        pause(lpause)\n",
    "\n",
    "    return wv_master"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The master list with wavelengths is immediately defined using this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# master list parameters\n",
    "my_seed = 432\n",
    "wv_ini_master = 3000\n",
    "wv_end_master = 7000\n",
    "nlines_master = 120\n",
    "\n",
    "# generate line wavelengths, following a random uniform distribution\n",
    "wv_master = simulate_master_table(my_seed, wv_ini_master, wv_end_master, nlines_master,\n",
    "                                  ldebug=ldebug)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once this master list is generated, it is necessary to determine all the possible triplets that can be built from that list. In addition, the relative position of the central line of each triplet must also be computed. This is done by the following auxiliaty function. Note that given a particular master list, this function should be called only once. This is important when calibrating an image containing several spectra of the same arc (e.g. EMIR or MEGARA)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_triplets_master(wv_master, ldebug=False, lplot=False):\n",
    "    \"\"\"Compute information associated to triplets in master table.\n",
    "\n",
    "    Determine all the possible triplets that can be generated from the\n",
    "    array `wv_master`. In addition, the relative position of the central\n",
    "    line of each triplet is also computed.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    wv_master : 1d numpy array, float\n",
    "        Array with wavelengths corresponding to the master table (Angstroms).\n",
    "    ldebug : bool\n",
    "        If True intermediate results are displayed.\n",
    "    lplot : bool\n",
    "        If True intermediate plots are displayed.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    ntriplets_master : int\n",
    "        Number of triplets built from master table.\n",
    "    ratios_master_sorted : 1d numpy array, float\n",
    "        Array with values of the relative position of the central line of each\n",
    "        triplet, sorted in ascending order.\n",
    "    triplets_master_sorted_list : list of tuples\n",
    "        List with tuples of three numbers, corresponding to the three line\n",
    "        indices in the master table. The list is sorted to be in correspondence\n",
    "        with `ratios_master_sorted`.\n",
    "       \n",
    "    \"\"\"\n",
    "\n",
    "    nlines_master = wv_master.size\n",
    "\n",
    "    #---\n",
    "    # Generate all the possible triplets with the numbers of the lines in the\n",
    "    # master table. Each triplet is defined as a tuple of three numbers\n",
    "    # corresponding to the three line indices in the master table. The\n",
    "    # collection of tuples is stored in an ordinary python list.\n",
    "    iter_comb_triplets = itertools.combinations(range(nlines_master), 3)\n",
    "    triplets_master_list = []\n",
    "    for val in iter_comb_triplets:\n",
    "        triplets_master_list.append(val)\n",
    "\n",
    "    # Verify that the number of triplets coincides with the expected value.\n",
    "    ntriplets_master = len(triplets_master_list)\n",
    "    if ntriplets_master == scipy.misc.comb(nlines_master, 3, exact=True):\n",
    "        if ldebug:\n",
    "            print('>>> Total number of lines in master table:', \n",
    "                  nlines_master)\n",
    "            print('>>> Number of triplets in master table...:', \n",
    "                  ntriplets_master)\n",
    "    else:\n",
    "        raise ValueError('Invalid number of combinations')\n",
    "\n",
    "    # For each triplet, compute the relative position of the central line.\n",
    "    ratios_master = np.zeros(ntriplets_master)\n",
    "    for i_tupla in range(ntriplets_master):\n",
    "        i1, i2, i3 = triplets_master_list[i_tupla]\n",
    "        delta1 = wv_master[i2]-wv_master[i1]\n",
    "        delta2 = wv_master[i3]-wv_master[i1]\n",
    "        ratios_master[i_tupla] = delta1/delta2\n",
    "\n",
    "    # Compute the array of indices that index the above ratios in sorted order.\n",
    "    isort_ratios_master = np.argsort(ratios_master)\n",
    "\n",
    "    # Simultaneous sort of position ratios and triplets.\n",
    "    ratios_master_sorted = ratios_master[isort_ratios_master]\n",
    "    triplets_master_sorted_list = \\\n",
    "      [triplets_master_list[i] for i in isort_ratios_master]\n",
    "\n",
    "    if lplot:\n",
    "        # Compute and plot histogram with position ratios\n",
    "        bins_in = np.linspace(0.0, 1.0, 41)\n",
    "        hist, bins_out = np.histogram(ratios_master, bins=bins_in)\n",
    "        #\n",
    "        fig = plt.figure()\n",
    "        ax = fig.add_subplot(111)\n",
    "        width_hist = 0.8*(bins_out[1]-bins_out[0])\n",
    "        center = (bins_out[:-1]+bins_out[1:])/2\n",
    "        ax.bar(center, hist, align='center', width=width_hist)\n",
    "        ax.set_xlabel('distance ratio in each triplet')\n",
    "        ax.set_ylabel('Number of triplets')\n",
    "        plt.show()\n",
    "        pause(lpause)\n",
    "\n",
    "    return ntriplets_master, ratios_master_sorted, triplets_master_sorted_list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ntriplets_master, ratios_master_sorted, triplets_master_sorted_list = \\\n",
    "  gen_triplets_master(wv_master, ldebug=ldebug, lplot=lplot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating the observation of an arc spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are going to generate a simulated arc spectrum by assuming a particular wavelength range (defined by `wv_ini_arc` and `wv_end_arc`), the number of pixels (`naxis1_arc`), a probability for the lines in the master list to be detected in the arc (`prob_line_master_in_arc`), a minimum distance between arc lines (`delta_xpos_min_arc`), the maximum deviation from linearity of the calibration polynomial (`delta_lambda` in angstrom) an error in the location of the arc lines (`error_xpos_arc` in pixels), a polynomial degree for the wavelength calibration (`poly_degree`), and a fraction of lines in the arc that are not going to be present in the master list (`fraction_unknown_lines`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wv_ini_arc = 4000\n",
    "wv_end_arc = 5000\n",
    "naxis1_arc = 1024\n",
    "\n",
    "prob_line_master_in_arc = 0.80\n",
    "delta_xpos_min_arc = 4.0\n",
    "delta_lambda = 5.0\n",
    "error_xpos_arc = 0.3\n",
    "poly_degree = 2\n",
    "fraction_unknown_lines = 0.20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The actual generation of the arc lines is performed by the auxiliary function `simulate_arc`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_arc(wv_ini_master,\n",
    "                 wv_end_master,\n",
    "                 wv_master,\n",
    "                 wv_ini_arc,\n",
    "                 wv_end_arc,\n",
    "                 naxis1_arc,\n",
    "                 prob_line_master_in_arc,\n",
    "                 delta_xpos_min_arc,\n",
    "                 delta_lambda,\n",
    "                 error_xpos_arc,\n",
    "                 poly_degree,\n",
    "                 fraction_unknown_lines,\n",
    "                 ldebug=False, \n",
    "                 lplot=False,\n",
    "                 lpause=False):\n",
    "    \"\"\"Generates simulated input for arc calibration.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    wv_ini_master : float\n",
    "        Minimum wavelength in master table.\n",
    "    wv_end_master : float\n",
    "        Maximum wavelength in master table.\n",
    "    wv_master : 1d numpy array, float\n",
    "        Array with wavelengths corresponding to the master table (Angstroms).\n",
    "    wv_ini_arc : float\n",
    "        Minimum wavelength in arc spectrum.\n",
    "    wv_end_arc : float\n",
    "        Maximum wavelength in arc spectrum.\n",
    "    naxis1_arc : int\n",
    "        NAXIS1 of arc spectrum.\n",
    "    prob_line_master_in_arc : float\n",
    "        Probability that a master table line appears in the arc spectrum.\n",
    "    delta_xpos_min_arc : float\n",
    "        Minimum distance (pixels) between lines in arc spectrum.\n",
    "    delta_lambda : float\n",
    "        Maximum deviation (Angstroms) from linearity in arc calibration \n",
    "        polynomial.\n",
    "    error_xpos_arc : float\n",
    "        Standard deviation (pixels) employed to introduce noise in the arc\n",
    "        spectrum lines. The initial lines are shifted from their original\n",
    "        location following a Normal distribution with mean iqual to zero and\n",
    "        sigma equal to this parameter.\n",
    "    poly_degree : int\n",
    "        Polynomial degree corresponding to the original wavelength calibration\n",
    "        function.\n",
    "    fraction_unknown_lines : float\n",
    "        Fraction of lines that on average will be unknown (i.e., lines that\n",
    "        appear in the arc spectrum that are not present in the master table).\n",
    "    ldebug : bool\n",
    "        If True intermediate results are displayed.\n",
    "    lplot : bool\n",
    "        If True intermediate plots are displayed.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    nlines_arc : int\n",
    "        Number of arc lines\n",
    "    xpos_arc : 1d numpy array, float\n",
    "        Location of arc lines (pixels).\n",
    "    crval1_arc : float\n",
    "        CRVAL1 for arc spectrum (linear approximation).\n",
    "    cdelt1_arc : float\n",
    "        CDELT1 for arc spectrum (linear approximation).\n",
    "    c0_arc, c1_arc, c2_arc : floats\n",
    "        Coefficients of the second order polynomial.\n",
    "    ipos_wv_arc : 1d numpy array, int\n",
    "        Number of line in master table corresponding to each arc line. Unknown\n",
    "        lines (i.e. those that are not present in the master table) are\n",
    "        assigned to -1.\n",
    "    coeff_original : 1d numpy array, float\n",
    "        Polynomial coefficients ordered from low to high, corresponding to the\n",
    "        fit to the arc lines before the inclusion of unknown lines.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    if (wv_ini_arc < wv_ini_master) or (wv_ini_arc > wv_end_master):\n",
    "        print('wv_ini_master:', wv_ini_master)\n",
    "        print('wv_end_master:', wv_end_master)\n",
    "        print('wv_ini_arc...:', wv_ini_arc)\n",
    "        raise ValueError('wavelength_ini_arc outside valid range')\n",
    "\n",
    "    if (wv_end_arc < wv_ini_master) or (wv_end_arc > wv_end_master):\n",
    "        print('wv_ini_master:', wv_ini_master)\n",
    "        print('wv_end_master:', wv_end_master)\n",
    "        print('wv_end_arc...:', wv_end_arc)\n",
    "        raise ValueError('wavelength_ini_arc outside valid range')\n",
    "\n",
    "    if wv_end_arc < wv_ini_arc:\n",
    "        raise ValueError('wv_ini_arc=' + str(wv_ini_arc) +\n",
    "                         ' must be <= wv_end_arc=' + str(wv_end_arc))\n",
    "\n",
    "    #---\n",
    "\n",
    "    nlines_master = wv_master.size\n",
    "\n",
    "    crval1_arc = wv_ini_arc\n",
    "    cdelt1_arc = (wv_end_arc-wv_ini_arc)/float(naxis1_arc-1)\n",
    "    crpix1_arc = 1.0\n",
    "    if ldebug:\n",
    "        print('>>> CRVAL1, CDELT1, CRPIX1....:', crval1_arc, cdelt1_arc, crpix1_arc)\n",
    "\n",
    "    #---\n",
    "    # The arc lines constitute a subset of the master lines in the considered\n",
    "    # wavelength range.\n",
    "    i1_master = np.searchsorted(wv_master, wv_ini_arc)\n",
    "    i2_master = np.searchsorted(wv_master, wv_end_arc)\n",
    "    nlines_temp = i2_master-i1_master\n",
    "    nlines_arc_ini = int(round(nlines_temp*prob_line_master_in_arc))\n",
    "    ipos_wv_arc_ini = np.random.choice(range(i1_master,i2_master),\n",
    "                                       size = nlines_arc_ini,\n",
    "                                       replace = False)\n",
    "    ipos_wv_arc_ini.sort() # in-place sort\n",
    "    wv_arc_ini = wv_master[ipos_wv_arc_ini]\n",
    "    if ldebug:\n",
    "        print('>>> Number of master lines in arc region.:', nlines_temp)\n",
    "        print('>>> Initial number of arc lines..........:', nlines_arc_ini)\n",
    "        print('>>> Initial selection of master list lines for arc:')\n",
    "        print(ipos_wv_arc_ini)\n",
    "    # Remove too close lines.\n",
    "    ipos_wv_arc = np.copy(ipos_wv_arc_ini[0:1])\n",
    "    wv_arc = np.copy(wv_arc_ini[0:1])\n",
    "    i_last = 0\n",
    "    for i in range(1,nlines_arc_ini):\n",
    "        delta_xpos = (wv_arc_ini[i]-wv_arc_ini[i_last])/cdelt1_arc\n",
    "        if delta_xpos > delta_xpos_min_arc:\n",
    "            ipos_wv_arc = np.append(ipos_wv_arc, ipos_wv_arc_ini[i])\n",
    "            wv_arc = np.append(wv_arc, wv_arc_ini[i])\n",
    "            i_last = i\n",
    "        else:\n",
    "            if ldebug:\n",
    "                print('--> skipping line #',i,'. Too close to line #',i_last)\n",
    "    nlines_arc = len (wv_arc)\n",
    "    if ldebug:\n",
    "        print('>>> Intermediate number of arc lines.....:', nlines_arc)\n",
    "        print('>>> Intermediate selection of master list lines for arc:')\n",
    "        print(ipos_wv_arc)\n",
    "\n",
    "    # Generate pixel location of the arc lines.\n",
    "    if delta_lambda == 0.0:\n",
    "        # linear solution\n",
    "        xpos_arc = (wv_arc - crval1_arc)/cdelt1_arc + 1.0\n",
    "        c0_arc = wv_ini_arc\n",
    "        c1_arc = cdelt1_arc\n",
    "        c2_arc = 0.0\n",
    "    else:\n",
    "        # polynomial solution\n",
    "        c0_arc = wv_ini_arc\n",
    "        c1_arc = (wv_end_arc-wv_ini_arc-4*delta_lambda)/float(naxis1_arc-1)\n",
    "        c2_arc = 4*delta_lambda/float(naxis1_arc-1)**2\n",
    "        xpos_arc = (-c1_arc+np.sqrt(c1_arc**2-4*c2_arc*(c0_arc-wv_arc)))\n",
    "        xpos_arc /= 2*c2_arc\n",
    "        xpos_arc += 1 # convert from 0,...,(NAXIS1-1) to 1,...,NAXIS1\n",
    "\n",
    "    # Introduce noise in arc line positions.\n",
    "    if error_xpos_arc > 0:\n",
    "        xpos_arc += np.random.normal(loc=0.0, \n",
    "                                     scale=error_xpos_arc,\n",
    "                                     size=nlines_arc)\n",
    "    # Check that the order of the lines has not been modified.\n",
    "    xpos_arc_copy = np.copy(xpos_arc)\n",
    "    xpos_arc_copy.sort() # in-place sort\n",
    "    if sum(xpos_arc == xpos_arc_copy) != len(xpos_arc):\n",
    "        raise ValueError('FATAL ERROR: arc line switch after introducing noise')\n",
    "\n",
    "    if lplot:\n",
    "        fig = plt.figure()\n",
    "        ax = fig.add_subplot(111)\n",
    "        ax.set_xlim([1,naxis1_arc])\n",
    "        ax.set_ylim([wv_ini_arc,wv_end_arc])\n",
    "        ax.plot(xpos_arc, wv_arc, 'ro')\n",
    "        xp = np.array([1,naxis1_arc])\n",
    "        yp = np.array([wv_ini_arc,wv_end_arc])\n",
    "        ax.plot(xp,yp,'b-')\n",
    "        xp = np.arange(1,naxis1_arc+1)\n",
    "        yp = c0_arc+c1_arc*(xp-1)+c2_arc*(xp-1)**2\n",
    "        ax.plot(xp,yp,'g:')\n",
    "        ax.set_xlabel('pixel position in arc spectrum')\n",
    "        ax.set_ylabel('wavelength (Angstrom)')\n",
    "        plt.show()\n",
    "        pause(lpause)\n",
    "\n",
    "    # Unweighted polynomial fit.\n",
    "    coeff_original = polynomial.polyfit(xpos_arc, wv_arc, poly_degree)\n",
    "    poly_original = polynomial.Polynomial(coeff_original)\n",
    "    \n",
    "    if lplot:\n",
    "        fig = plt.figure()\n",
    "        ax = fig.add_subplot(111)\n",
    "        ax.set_xlim([1,naxis1_arc])\n",
    "        if delta_lambda == 0.0:\n",
    "            if error_xpos_arc > 0:\n",
    "                ax.set_ylim([-4*error_xpos_arc*cdelt1_arc,\n",
    "                              4*error_xpos_arc*cdelt1_arc])\n",
    "            else:\n",
    "                ax.set_ylim([-1.1,1.1])\n",
    "        else:\n",
    "            ax.set_ylim([-delta_lambda*1.5,delta_lambda*1.5])\n",
    "        yp = wv_arc - (crval1_arc+(xpos_arc-1)*cdelt1_arc)\n",
    "        ax.plot(xpos_arc, yp, 'ro')\n",
    "        xp = np.array([1,naxis1_arc])\n",
    "        yp = np.array([0,0])\n",
    "        ax.plot(xp,yp,'b-')\n",
    "        xp = np.arange(1,naxis1_arc+1)\n",
    "        yp = c0_arc+c1_arc*(xp-1)+c2_arc*(xp-1)**2\n",
    "        yp -= crval1_arc+cdelt1_arc*(xp-1)\n",
    "        ax.plot(xp,yp,'g:')\n",
    "        yp = poly_original(xp)\n",
    "        yp -= crval1_arc+cdelt1_arc*(xp-1)\n",
    "        ax.plot(xp,yp,'m-')\n",
    "        ax.set_xlabel('pixel position in arc spectrum')\n",
    "        ax.set_ylabel('residuals (Angstrom)')\n",
    "        plt.show()\n",
    "        pause(lpause)\n",
    "\n",
    "    #---\n",
    "    # Include unknown lines (lines that do not appear in the master table).\n",
    "    nunknown_lines = int(round(fraction_unknown_lines * float(nlines_arc)))\n",
    "    if ldebug:\n",
    "        print('>>> Number of unknown arc lines..........:', nunknown_lines)\n",
    "    for i in range(nunknown_lines):\n",
    "        iiter = 0\n",
    "        while True:\n",
    "            iiter += 1\n",
    "            if iiter > 1000:\n",
    "                raise ValueError('iiter > 1000 while adding unknown lines')\n",
    "            xpos_dum = np.random.uniform(low = 1.0,\n",
    "                                         high = float(naxis1_arc),\n",
    "                                         size = 1)\n",
    "            isort = np.searchsorted(xpos_arc, xpos_dum)\n",
    "            newlineok = False\n",
    "            if isort == 0:\n",
    "                dxpos1 = abs(xpos_arc[isort]-xpos_dum)\n",
    "                if dxpos1 > delta_xpos_min_arc:\n",
    "                    newlineok = True\n",
    "            elif isort == nlines_arc:\n",
    "                dxpos2 = abs(xpos_arc[isort-1]-xpos_dum)\n",
    "                if dxpos2 > delta_xpos_min_arc:\n",
    "                    newlineok = True\n",
    "            else:\n",
    "                dxpos1 = abs(xpos_arc[isort]-xpos_dum)\n",
    "                dxpos2 = abs(xpos_arc[isort-1]-xpos_dum)\n",
    "                if (dxpos1 > delta_xpos_min_arc) and \\\n",
    "                   (dxpos2 > delta_xpos_min_arc):\n",
    "                    newlineok = True\n",
    "            if newlineok:\n",
    "                xpos_arc = np.insert(xpos_arc, isort, xpos_dum)\n",
    "                ipos_wv_arc = np.insert(ipos_wv_arc, isort, -1)\n",
    "                nlines_arc += 1\n",
    "                if ldebug:\n",
    "                    print('--> adding unknown line at pixel:',xpos_dum)\n",
    "                break\n",
    "    if ldebug:\n",
    "        print('>>> Final number of arc lines............:', nlines_arc)\n",
    "        for val in zip(range(nlines_arc), ipos_wv_arc, xpos_arc):\n",
    "            print(val)\n",
    "        pause(lpause)\n",
    "\n",
    "    if lplot:\n",
    "        fig = plt.figure()\n",
    "        ax = fig.add_subplot(111)\n",
    "        ax.set_ylim([0.0,3.0])\n",
    "        ax.vlines(wv_master, 0.0, 1.0)\n",
    "        ax.vlines(wv_arc, 1.0, 2.0, colors='r', linestyle=':')\n",
    "        ax.vlines(wv_ini_arc, 0.0, 3.0, colors='m', linewidth=3.0)\n",
    "        ax.vlines(wv_end_arc, 0.0, 3.0, colors='m', linewidth=3.0)\n",
    "        ax.set_xlabel('wavelength')\n",
    "        axbis = ax.twiny()\n",
    "        axbis.vlines(xpos_arc, 2.0, 3.0, colors='g')\n",
    "        xmin_xpos_master = (wv_ini_master - crval1_arc)/cdelt1_arc + 1.0\n",
    "        xmax_xpos_master = (wv_end_master - crval1_arc)/cdelt1_arc + 1.0\n",
    "        axbis.set_xlim([xmin_xpos_master, xmax_xpos_master])\n",
    "        axbis.set_xlabel('pixel position in arc spectrum')\n",
    "        plt.show()\n",
    "        pause(lpause)\n",
    "\n",
    "    return nlines_arc, xpos_arc, crval1_arc, cdelt1_arc, \\\n",
    "           c0_arc, c1_arc, c2_arc, \\\n",
    "           ipos_wv_arc, coeff_original"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the help of this function it is immediate to simulate the arc spectrum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlines_arc, xpos_arc, crval1_arc, cdelt1_arc, \\\n",
    "  c0_arc, c1_arc, c2_arc, ipos_wv_arc, coeff_original = \\\n",
    "  simulate_arc(wv_ini_master,\n",
    "               wv_end_master,\n",
    "               wv_master,\n",
    "               wv_ini_arc,\n",
    "               wv_end_arc,\n",
    "               naxis1_arc,\n",
    "               prob_line_master_in_arc,\n",
    "               delta_xpos_min_arc,\n",
    "               delta_lambda,\n",
    "               error_xpos_arc,\n",
    "               poly_degree,\n",
    "               fraction_unknown_lines,\n",
    "               ldebug=ldebug, \n",
    "               lplot=lplot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last plot represents all the lines of the master list (black vertical lines at the bottom) and the arc lines corresponding to the simulated spectrum (green vertical lines at the top). The intermediate dotted red lines try to connect the master lines with the arc lines. However, since we are introducing a distortion in the wavelength scale, the green lines do not align perfectly with the master lines. In addition, some green lines have been artificially added to the arc spectrum that do not appear in the master list. These two effects are used to check how robust the wavelength calibration procedure is when confronting these problems. \n",
    "\n",
    "At this point all the required information is already available to perform the wavelength calibration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The wavelength calibration procedure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlines_master = wv_master.size\n",
    "\n",
    "nlines_arc = xpos_arc.size\n",
    "if nlines_arc < 5:\n",
    "    raise ValueError('Insufficient arc lines=' + str(nlines_arc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code should be incorporated into a function. We start by defining some parameters which are going to be useful later. The meaning of those parameters will be explained when they are used for the first time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters for wavelength calibration\n",
    "\n",
    "wv_ini_search = wv_ini_master  # minimum wavelength to be considered\n",
    "wv_end_search = wv_end_master  # maximum wavelength to be considered\n",
    "times_sigma_r = 3.0  # maximum allowed uncertainty in arc line location\n",
    "frac_triplets_for_sum = 0.50  # fraction of layers to be used when computing the cost function for each initial solution\n",
    "times_sigma_theil_sen = 10.0  # times sigma to remove deviant lines using the Theil-Sen method\n",
    "poly_degree_wfit = 2  # degree for polynomial fit to wavelength calibration\n",
    "times_sigma_polfilt = 10.0  # times sigma to reject deviant lines using a polynomial fit\n",
    "times_sigma_inclusion = 5.0  # times sigma to incorporate new lines that remain unidentified after the main procedure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is time now to generate all the triplets that can be built using consecutive arc lines. This number is equal to the total number of arc lines minus 2. \n",
    "\n",
    "A given arc triplet will be defined by the location of three lines $x_1$, $x_2$ and $x_3$ (in pixel units). This triplet will be characterized by the relative location of the central line (`ratio_arc`), defined as $$r \\equiv \\frac{x_2-x_1}{x_3-x_1}.$$\n",
    "This number will always verify $r \\in [0.1]$. In addition, we are assuming that those line locations have been measured with an uncertainty $\\Delta x$ (in pixel units; the same for each $x_i$, with $i=1, 2, 3$) given by `error_xpos_arc`. An important expression is the relation between this $\\Delta x$ and $\\Delta r$. Using the law of propagation of errors, it is easy to show that $$\\Delta r = \\frac{\\sqrt{2} \\, \\Delta x}{x_3-x_1} \\sqrt{r(r-1)+1}.$$\n",
    "\n",
    "For each arc triplet, characterized by its $r$ value ($r_{arc-triplet}$), compatible triplets from the master table are sought. By compatible triplet we consider a master line triplet which $r$ value $r_{master-triplet}$ is close to $r_{arc-triplet}$ ($\\pm$ a number of times the value of $\\Delta r$; this number of times is determined by `times_sigma_r`).\n",
    "\n",
    "Each compatible triplet from the master table provides an estimate for CRVAL1 and CDELT1. By definition, $\\lambda = \\lambda_{ini} + (N_{pixel}-1) \\Theta$, where $\\lambda_{ini}$=CRVAL1 (the central wavelength of the first pixel, in Angstrom), $\\Theta$=CDELT1 (the dispersion, in Angstrom/pixel), and $N_{pixel}$ is the pixel number. Since the triplet from the master table provides the wavelengths for three lines ($\\lambda_i$, for $i=1,2,3$), the previous expression leads to $$\\Theta = \\displaystyle\\frac{\\lambda_3-\\lambda_1}{x_3-x_1}$$ and $$\\lambda_{ini} = \\lambda_2 -(x_2-1) \\Theta.$$ \n",
    "\n",
    "From here is also possible to derive the uncertainties in these two parameters (which are correlated!): $$\\Delta\\Theta = \\sqrt{2}\\,\\Theta\\,\\frac{\\Delta x}{x_3-x_1}$$ and $$\\Delta\\lambda_{ini} = \\theta \\,\\Delta x \\sqrt{1+2 \\frac{(x_2-1)^2}{(x_3-x_1)^2}}.$$\n",
    "\n",
    "As an additional constraint, the only valid solutions are those for which the initial and the final wavelengths for the arc are restricted to a predefined wavelength interval (defined by `wv_ini_search` and `wv_end_search`). **Note that this interval must be generous enough to accommodate the real solution. On the contrary, the algorithm can fail to identify the triplets located just at the border of the wavelength interval.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crval1_search = np.array([])\n",
    "cdelt1_search = np.array([])\n",
    "error_crval1_search = np.array([])\n",
    "error_cdelt1_search = np.array([])\n",
    "itriplet_search = np.array([])\n",
    "clabel_search = []\n",
    "\n",
    "ntriplets_arc = nlines_arc - 2\n",
    "if ldebug:\n",
    "    print('>>> Total number of arc lines............:', nlines_arc)\n",
    "    print('>>> Total number of arc triplets.........:', ntriplets_arc)\n",
    "\n",
    "# Loop in all the arc line triplets. Note that only triplets built from\n",
    "# consecutive arc lines are considered.\n",
    "for i in range(ntriplets_arc):\n",
    "    i1, i2, i3 = i, i+1, i+2\n",
    "\n",
    "    dist12 = xpos_arc[i2]-xpos_arc[i1]\n",
    "    dist13 = xpos_arc[i3]-xpos_arc[i1]\n",
    "    ratio_arc = dist12/dist13\n",
    "\n",
    "    pol_r = ratio_arc*(ratio_arc-1)+1\n",
    "    error_ratio_arc = np.sqrt(2)*error_xpos_arc/dist13*np.sqrt(pol_r)\n",
    "\n",
    "    ratio_arc_min = max(0.0, ratio_arc-times_sigma_r*error_ratio_arc)\n",
    "    ratio_arc_max = min(1.0, ratio_arc+times_sigma_r*error_ratio_arc)\n",
    "\n",
    "    # determine compatible triplets from the master list\n",
    "    j_loc_min = np.searchsorted(ratios_master_sorted, ratio_arc_min)-1\n",
    "    j_loc_max = np.searchsorted(ratios_master_sorted, ratio_arc_max)+1\n",
    "\n",
    "    if j_loc_min < 0:\n",
    "        j_loc_min = 0\n",
    "    if j_loc_max > ntriplets_master:\n",
    "        j_loc_max = ntriplets_master\n",
    "\n",
    "    if ldebug:\n",
    "        print(i, ratio_arc_min, ratio_arc, ratio_arc_max, \n",
    "              j_loc_min, j_loc_max)\n",
    "\n",
    "    # each triplet from the master list provides a potential solution\n",
    "    # for CRVAL1 and CDELT1\n",
    "    for j_loc in range(j_loc_min, j_loc_max):\n",
    "        j1, j2, j3 = triplets_master_sorted_list[j_loc]\n",
    "        # initial solutions for CDELT1, CRVAL1 and CRVALN\n",
    "        cdelt1_temp = (wv_master[j3]-wv_master[j1])/dist13\n",
    "        crval1_temp = wv_master[j2]-(xpos_arc[i2]-1)*cdelt1_temp\n",
    "        crvaln_temp = crval1_temp + float(naxis1_arc-1)*cdelt1_temp\n",
    "        # check that CRVAL1 and CRVALN are within the valid limits\n",
    "        if wv_ini_search <= crval1_temp <= wv_end_search:\n",
    "            if wv_ini_search <= crvaln_temp <= wv_end_search:\n",
    "                # Compute errors\n",
    "                error_crval1_temp = cdelt1_temp*error_xpos_arc* \\\n",
    "                  np.sqrt(1+2*((xpos_arc[i2]-1)**2)/(dist13**2))\n",
    "                error_cdelt1_temp = np.sqrt(2)*cdelt1_temp* \\\n",
    "                  error_xpos_arc/dist13\n",
    "                # Store values and errors\n",
    "                crval1_search = np.append(crval1_search, crval1_temp)\n",
    "                cdelt1_search = np.append(cdelt1_search, cdelt1_temp)\n",
    "                error_crval1_search = np.append(error_crval1_search, \n",
    "                                                error_crval1_temp)\n",
    "                error_cdelt1_search = np.append(error_cdelt1_search, \n",
    "                                                error_cdelt1_temp)\n",
    "                # Store additional information about the triplets\n",
    "                itriplet_search = np.append(itriplet_search, i)\n",
    "                clabel_search.append((j1, j2, j3))\n",
    "         \n",
    "# Maximum allowed value for CDELT1\n",
    "cdelt1_max = (wv_end_search-wv_ini_search)/float(naxis1_arc-1)\n",
    "\n",
    "if lplot:\n",
    "    # CDELT1 vs CRVAL1 diagram (original coordinates)\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlabel('cdelt1 (Angstroms/pixel)')\n",
    "    ax.set_ylabel('crval1 (Angstroms)')\n",
    "    ax.scatter(cdelt1_search, crval1_search, s=200, alpha=0.1)\n",
    "    xmin = 0.0\n",
    "    xmax = cdelt1_max\n",
    "    dx = xmax-xmin\n",
    "    xmin -= dx/20\n",
    "    xmax += dx/20\n",
    "    ax.set_xlim([xmin,xmax])\n",
    "    ymin = wv_ini_search\n",
    "    ymax = wv_end_search\n",
    "    dy = ymax-ymin\n",
    "    ymin -= dy/20\n",
    "    ymax += dy/20\n",
    "    ax.set_ylim([ymin,ymax])\n",
    "    xp_limits = np.array([0., cdelt1_max])\n",
    "    yp_limits = wv_end_search-float(naxis1_arc-1)*xp_limits\n",
    "    xp_limits = np.concatenate((xp_limits,[xp_limits[0],xp_limits[0]]))\n",
    "    yp_limits = np.concatenate((yp_limits,[yp_limits[1],yp_limits[0]]))\n",
    "    ax.plot(xp_limits, yp_limits, linestyle='-', color='red')\n",
    "    plt.show()\n",
    "    print('Number of points in last plot:', len(cdelt1_search))\n",
    "    pause(lpause) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last plot shows the initial solutions plotted in the CDELT1-CRVAL1 diagram. Note that there is a concentration of solutions overplotted around CDELT1=1.0 and CRVAL1=4000 Angstroms (the right solution). However, there are some regions of this diagram with also high concentration of points. The limits corresponding to the valid area in this diagram are displayed with the red lines. Outside that region, no valid solutions can be found (the previous code has discarded them).\n",
    "\n",
    "Since we are going to measure distances in this diagram, it is useful to normalized the ranges in both axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalize the values of CDELT1 and CRVAL1 to the interval [0,1] in each\n",
    "# case.\n",
    "cdelt1_search_norm = cdelt1_search/cdelt1_max\n",
    "error_cdelt1_search_norm = error_cdelt1_search/cdelt1_max\n",
    "#\n",
    "crval1_search_norm = (crval1_search-wv_ini_search)\n",
    "crval1_search_norm /= (wv_end_search-wv_ini_search)\n",
    "error_crval1_search_norm = error_crval1_search/(wv_ini_search-wv_end_search)\n",
    "\n",
    "# Intermediate plots\n",
    "if lplot:\n",
    "    # CDELT1 vs CRVAL1 diagram (normalized coordinates)\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlabel('normalized cdelt1')\n",
    "    ax.set_ylabel('normalized crval1')\n",
    "    ax.scatter(cdelt1_search_norm, crval1_search_norm, s=200, alpha=0.1)\n",
    "    xmin = -0.05\n",
    "    xmax =  1.05\n",
    "    ymin = -0.05\n",
    "    ymax =  1.05\n",
    "    ax.set_xlim([xmin,xmax])\n",
    "    ax.set_ylim([ymin,ymax])\n",
    "    xp_limits = np.array([0.,1.,0.,0.])\n",
    "    yp_limits = np.array([1.,0.,0.,1.])\n",
    "    ax.plot(xp_limits, yp_limits, linestyle='-', color='red')\n",
    "    plt.show()\n",
    "    print('Number of points in last plot:', len(cdelt1_search))\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This plot can be repeated but using different colors for solutions derived from each arc triplet. In addition, the number of arc triplet is also overplotted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if lplot:\n",
    "    # CDELT1 vs CRVAL1 diagram (normalized coordinates)\n",
    "    # with different color for each arc triplet and overplotting \n",
    "    # the arc triplet number\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlabel('normalized cdelt1')\n",
    "    ax.set_ylabel('normalized crval1')\n",
    "    ax.scatter(cdelt1_search_norm, crval1_search_norm, s=200, alpha=0.1,\n",
    "               c=itriplet_search)\n",
    "    for i in range(len(itriplet_search)):\n",
    "        ax.text(cdelt1_search_norm[i], crval1_search_norm[i], \n",
    "                str(int(itriplet_search[i])), fontsize=6)\n",
    "    ax.set_xlim([xmin,xmax])\n",
    "    ax.set_ylim([ymin,ymax])\n",
    "    ax.plot(xp_limits, yp_limits, linestyle='-', color='red')\n",
    "    plt.show()\n",
    "    pause(lpause)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same plot but overplotting triplet numbers (the number of three lines within the master list)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if lplot:\n",
    "    # CDELT1 vs CRVAL1 diagram (normalized coordinates)\n",
    "    # including triplet numbers\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlabel('normalized cdelt1')\n",
    "    ax.set_ylabel('normalized crval1')\n",
    "    ax.scatter(cdelt1_search_norm, crval1_search_norm, s=200, alpha=0.1,\n",
    "               c=itriplet_search)\n",
    "    for i in range(len(clabel_search)):\n",
    "        ax.text(cdelt1_search_norm[i], crval1_search_norm[i], \n",
    "                clabel_search[i], fontsize=6)\n",
    "    ax.set_xlim([xmin,xmax])\n",
    "    ax.set_ylim([ymin,ymax])\n",
    "    ax.plot(xp_limits, yp_limits, linestyle='-', color='red')\n",
    "    plt.show()\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the same plot but including only the error bars (remember that those error bars are correlated!)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if lplot:\n",
    "    # CDELT1 vs CRVAL1 diagram (normalized coordinates)\n",
    "    # with error bars (note that errors in this plot are highly\n",
    "    # correlated)\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlabel('normalized cdelt1')\n",
    "    ax.set_ylabel('normalized crval1')\n",
    "    ax.errorbar(cdelt1_search_norm, crval1_search_norm, \n",
    "                xerr=error_cdelt1_search_norm,\n",
    "                yerr=error_crval1_search_norm,\n",
    "                fmt='none')\n",
    "    ax.set_xlim([xmin,xmax])\n",
    "    ax.set_ylim([ymin,ymax])\n",
    "    ax.plot(xp_limits, yp_limits, linestyle='-', color='red')\n",
    "    plt.show()\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step consists in segregating the different solutions by arc triplet. In this way, we have different plots like the previous ones (one for each arc triplet). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---\n",
    "# Segregate the different solutions (normalized to [0,1]) by triplet. In\n",
    "# this way the solutions are saved in different layers (a layer for each\n",
    "# triplet). The solutions will be stored as python lists of numpy arrays.\n",
    "ntriplets_layered_list = []\n",
    "cdelt1_layered_list = []\n",
    "error_cdelt1_layered_list = []\n",
    "crval1_layered_list = []\n",
    "error_crval1_layered_list = []\n",
    "itriplet_layered_list = []\n",
    "clabel_layered_list = []\n",
    "for i in range(ntriplets_arc):\n",
    "    ldum = (itriplet_search == i)\n",
    "    ntriplets_layered_list.append(ldum.sum())\n",
    "    #\n",
    "    cdelt1_dum = cdelt1_search_norm[ldum]\n",
    "    cdelt1_layered_list.append(cdelt1_dum)\n",
    "    error_cdelt1_dum = error_cdelt1_search_norm[ldum]\n",
    "    error_cdelt1_layered_list.append(error_cdelt1_dum)\n",
    "    #\n",
    "    crval1_dum = crval1_search_norm[ldum]\n",
    "    crval1_layered_list.append(crval1_dum)\n",
    "    error_crval1_dum = error_crval1_search_norm[ldum]\n",
    "    error_crval1_layered_list.append(error_crval1_dum)\n",
    "    #\n",
    "    itriplet_dum = itriplet_search[ldum]\n",
    "    itriplet_layered_list.append(itriplet_dum)\n",
    "    #\n",
    "    clabel_dum = [i for (i, v) in zip(clabel_search, ldum) if v]\n",
    "    clabel_layered_list.append(clabel_dum)\n",
    "\n",
    "if ldebug:\n",
    "    print('>>> List with no. of solutions/triplet...:',\n",
    "          ntriplets_layered_list)\n",
    "    print('>>> Total number of potential solutions..:',\n",
    "          sum(ntriplets_layered_list))\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once all the initial solutions have been segregated in different layers (one for each arc triplet), we can start computing the cost function for each solution.\n",
    "\n",
    "For each solution, we find the nearest solution in each of the the remaining `ntriplets_arc`-1 layers. After sorting all these distances, the cost function is the coaddition of a fraction of them (the fraction is given by `frac_triplets_for_sum`). After the computation of the cost function, this function is normalized and segregated by arc triplet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---\n",
    "# Computation of the cost function.\n",
    "#\n",
    "# For each solution, corresponding to a particular triplet, find the\n",
    "# nearest solution in each of the remaining ntriplets_arc-1 layers.\n",
    "# Compute the distance (in normalized coordinates) to those closest\n",
    "# solutions, and obtain the sum of distances considering only a fraction\n",
    "# of them (after sorting them in ascending order).\n",
    "ntriplets_for_sum = \\\n",
    "  max(1, int(round(frac_triplets_for_sum*float(ntriplets_arc))))\n",
    "funcost_search = np.zeros(len(itriplet_search))\n",
    "for k in range(len(itriplet_search)):\n",
    "    itriplet_local = itriplet_search[k]\n",
    "    x0 = cdelt1_search_norm[k]\n",
    "    y0 = crval1_search_norm[k]\n",
    "    dist_to_layers = np.array([])\n",
    "    for i in range(ntriplets_arc):\n",
    "        if i != itriplet_local:\n",
    "            if ntriplets_layered_list[i] > 0:\n",
    "                x1 = cdelt1_layered_list[i]\n",
    "                y1 = crval1_layered_list[i]\n",
    "                dist2 = (x0-x1)**2 + (y0-y1)**2\n",
    "                dist_to_layers = np.append(dist_to_layers, dist2.min())\n",
    "            else:\n",
    "                dist_to_layers = np.append(dist_to_layers, np.inf)\n",
    "    dist_to_layers.sort() # in-place sort\n",
    "    funcost_search[k] = dist_to_layers[range(ntriplets_for_sum)].sum()\n",
    "    \n",
    "# Normalize the cost function\n",
    "funcost_min = funcost_search.min()\n",
    "if ldebug:\n",
    "    print('funcost_min:',funcost_min)\n",
    "funcost_search /= funcost_min\n",
    "\n",
    "# Segregate the cost function by arc triplet.\n",
    "funcost_layered_list = []\n",
    "for i in range(ntriplets_arc):\n",
    "    ldum = (itriplet_search == i)\n",
    "    funcost_dum = funcost_search[ldum]\n",
    "    funcost_layered_list.append(funcost_dum)\n",
    "if ldebug:\n",
    "    for i in range(ntriplets_arc):\n",
    "        jdum = funcost_layered_list[i].argmin()\n",
    "        print('>>>',i,funcost_layered_list[i][jdum],\n",
    "              clabel_layered_list[i][jdum],\n",
    "              cdelt1_layered_list[i][jdum], \n",
    "              crval1_layered_list[i][jdum])\n",
    "    pause(lpause)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The tabulated data indicate, for each arc line triplet, the corresponding cost function, the number of the three lines within the master list, and the solutions for CDELT1 and CRVAL1 (in normalized units). Since the arc line triplets are consecutive, the above table provides, for each individual arc line, three identifications in the master list (except for the first and last arc lines, for which only one solution are available, and the second and penultimate arc lines, for which two solutions are possible).\n",
    "\n",
    "After the previous calculations, we can replot the CDELT1-CRVAL1 diagram, but now with symbol sizes proportional to the inverse of the cost function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if lplot:\n",
    "    # CDELT1 vs CRVAL1 diagram (normalized coordinates)\n",
    "    # with symbol size proportional to the inverse of the cost function\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlabel('normalized cdelt1')\n",
    "    ax.set_ylabel('normalized crval1')\n",
    "    ax.scatter(cdelt1_search_norm, crval1_search_norm, \n",
    "               s=2000/funcost_search, c=itriplet_search, alpha=0.2)\n",
    "    ax.set_xlim([xmin,xmax])\n",
    "    ax.set_ylim([ymin,ymax])\n",
    "    ax.plot(xp_limits, yp_limits, linestyle='-', color='red')\n",
    "    plt.show()\n",
    "    print('Number of points in last plot:', len(cdelt1_search))\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The largets symbols indicate the solutions with higher probabilities of being the correct ones.\n",
    "\n",
    "Then, we proceed to compile, for each arc line, the three possible identifications within the master list (for the first and last arc line, only one identification is available; for the second and penultiname arc line, two identifications are derived)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We store the identifications of each line in a python list of lists\n",
    "# named diagonal_ids (which grows as the different triplets are\n",
    "# considered). A similar list of lists is also employed to store the\n",
    "# corresponding cost functions.\n",
    "for i in range(ntriplets_arc):\n",
    "    jdum = funcost_layered_list[i].argmin()\n",
    "    k1, k2, k3 = clabel_layered_list[i][jdum]\n",
    "    funcost_dum = funcost_layered_list[i][jdum]\n",
    "    if i == 0:\n",
    "        diagonal_ids = [[k1],[k2],[k3]]\n",
    "        diagonal_funcost = [[funcost_dum],[funcost_dum], [funcost_dum]]\n",
    "    else:\n",
    "        diagonal_ids[i].append(k1)\n",
    "        diagonal_ids[i+1].append(k2)\n",
    "        diagonal_ids.append([k3])\n",
    "        diagonal_funcost[i].append(funcost_dum)\n",
    "        diagonal_funcost[i+1].append(funcost_dum)\n",
    "        diagonal_funcost.append([funcost_dum])\n",
    "\n",
    "if ldebug:\n",
    "    for i in range(nlines_arc):\n",
    "        print(i, diagonal_ids[i],diagonal_funcost[i])\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is the line identification. Several scenarios are considered:\n",
    "\n",
    "- Lines with three identifications:\n",
    "\n",
    "  - Type A: the three identifications are identical. We keep the lowest value of the three cost functions.\n",
    "  \n",
    "  - Type B: two identifications are identical and one is different. We keep the line with two identifications and the lowest of the corresponding two cost functions.\n",
    "  \n",
    "  - Type C: the three identifications are different. We keep the one which is closest to a previously identified type B line (it can not be next to a Type A line). We take the corresponding cost function.\n",
    "  \n",
    "- Lines with two identifications (second and penultiname lines):\n",
    "\n",
    "  - Type D: the two identifications are identical. We keep the lowest cost function value.\n",
    "\n",
    "- Lines with one identification (first and last lines):\n",
    "\n",
    "  - Type E: the two lines next (or previous) to the considered line have been identified. We keep its cost function.\n",
    "  \n",
    "Note that after this procedure, some lines can remain unidentified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The solutions are stored in a list of dictionaries. The dictionaries\n",
    "# contain the following elements:\n",
    "# - lineok: bool, indicates whether the line has been properly identified\n",
    "# - type: 'A','B','C','D','E',...\n",
    "# - id: index of the line in the master table\n",
    "# - funcost: cost function associated the the line identification\n",
    "# First, initialize solution.\n",
    "solution = []\n",
    "for i in range(nlines_arc):\n",
    "    solution.append({'lineok': False, \n",
    "                     'type': None,\n",
    "                     'id': None, \n",
    "                     'funcost': None})\n",
    "\n",
    "# Type A lines.\n",
    "for i in range(2,nlines_arc-2):\n",
    "    j1,j2,j3 = diagonal_ids[i]\n",
    "    if j1 == j2 == j3:\n",
    "        solution[i]['lineok'] = True\n",
    "        solution[i]['type'] = 'A'\n",
    "        solution[i]['id'] = j1\n",
    "        solution[i]['funcost'] = min(diagonal_funcost[i])\n",
    "\n",
    "if ldebug:\n",
    "    print('\\n* Including Type A lines:')\n",
    "    for i in range(nlines_arc):\n",
    "        print(i, solution[i])\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type B lines.\n",
    "for i in range(2,nlines_arc-2):\n",
    "    if solution[i]['lineok'] == False:\n",
    "        j1,j2,j3 = diagonal_ids[i]\n",
    "        f1,f2,f3 = diagonal_funcost[i]\n",
    "        if j1 == j2:\n",
    "            if max(f1,f2) < f3:\n",
    "                solution[i]['lineok'] = True\n",
    "                solution[i]['type'] = 'B'\n",
    "                solution[i]['id'] = j1\n",
    "                solution[i]['funcost'] = min(f1,f2)\n",
    "        elif j1 == j3:\n",
    "            if max(f1,f3) < f2:\n",
    "                solution[i]['lineok'] = True\n",
    "                solution[i]['type'] = 'B'\n",
    "                solution[i]['id'] = j1\n",
    "                solution[i]['funcost'] = min(f1,f3)\n",
    "        elif j2 == j3:\n",
    "            if max(f2,f3) < f1:\n",
    "                solution[i]['lineok'] = True\n",
    "                solution[i]['type'] = 'B'\n",
    "                solution[i]['id'] = j2\n",
    "                solution[i]['funcost'] = min(f2,f3)\n",
    "\n",
    "if ldebug:\n",
    "    print('\\n* Including Type B lines:')\n",
    "    for i in range(nlines_arc):\n",
    "        print(i, solution[i])\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type C lines.\n",
    "for i in range(2,nlines_arc-2):\n",
    "    if solution[i]['lineok'] == False:\n",
    "        j1,j2,j3 = diagonal_ids[i]\n",
    "        f1,f2,f3 = diagonal_funcost[i]\n",
    "        if solution[i-1]['type'] == 'B':\n",
    "            if min(f2,f3) > f1:\n",
    "                solution[i]['lineok'] = True\n",
    "                solution[i]['type'] = 'C'\n",
    "                solution[i]['id'] = j1\n",
    "                solution[i]['funcost'] = f1\n",
    "        elif solution[i+1]['type'] == 'B':\n",
    "            if min(f1,f2) > f3:\n",
    "                solution[i]['lineok'] = True\n",
    "                solution[i]['type'] = 'C'\n",
    "                solution[i]['id'] = j3\n",
    "                solution[i]['funcost'] = f3\n",
    "\n",
    "if ldebug:\n",
    "    print('\\n* Including Type C lines:')\n",
    "    for i in range(nlines_arc):\n",
    "        print(i, solution[i])\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type D lines.\n",
    "for i in [1,nlines_arc-2]:\n",
    "    j1,j2 = diagonal_ids[i]\n",
    "    if j1 == j2:\n",
    "        f1,f2 = diagonal_funcost[i]\n",
    "        solution[i]['lineok'] = True\n",
    "        solution[i]['type'] = 'D'\n",
    "        solution[i]['id'] = j1\n",
    "        solution[i]['funcost'] = min(f1,f2)\n",
    "\n",
    "if ldebug:\n",
    "    print('\\n* Including Type D lines:')\n",
    "    for i in range(nlines_arc):\n",
    "        print(i, solution[i])\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Type E lines.\n",
    "i = 0\n",
    "if solution[i+1]['lineok'] and solution[i+2]['lineok']:\n",
    "    solution[i]['lineok'] = True\n",
    "    solution[i]['type'] = 'E'\n",
    "    solution[i]['id'] = diagonal_ids[i][0]\n",
    "    solution[i]['funcost'] = diagonal_funcost[i][0]\n",
    "i = nlines_arc-1\n",
    "if solution[i-2]['lineok'] and solution[i-1]['lineok']:\n",
    "    solution[i]['lineok'] = True\n",
    "    solution[i]['type'] = 'E'\n",
    "    solution[i]['id'] = diagonal_ids[i][0]\n",
    "    solution[i]['funcost'] = diagonal_funcost[i][0]\n",
    "\n",
    "if ldebug:\n",
    "    print('\\n* Including Type E lines:')\n",
    "    for i in range(nlines_arc):\n",
    "        print(i, solution[i])\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that the solutions do not contain duplicated values. If they are present (probably due to the influence of an unknown line that unfortunately falls too close to a real line in the master table), we keep the solution with the lowest cost function. The removed lines are labelled as type='R'. The procedure is repeated several times in case a line appears more than twice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lduplicated = True\n",
    "nduplicated = 0\n",
    "while lduplicated:\n",
    "    lduplicated = False\n",
    "    for i1 in range(nlines_arc):\n",
    "        if solution[i1]['lineok']:\n",
    "            j1 = solution[i1]['id']\n",
    "            for i2 in range(i1+1,nlines_arc):\n",
    "                if solution[i2]['lineok']:\n",
    "                    j2 = solution[i2]['id']\n",
    "                    if j1 == j2:\n",
    "                        lduplicated = True\n",
    "                        nduplicated += 1\n",
    "                        f1 = solution[i1]['funcost']\n",
    "                        f2 = solution[i2]['funcost']\n",
    "                        if f1 < f2:\n",
    "                            solution[i2]['lineok'] = False\n",
    "                            solution[i2]['type'] = 'R'\n",
    "                        else:\n",
    "                            solution[i1]['lineok'] = False\n",
    "                            solution[i1]['type'] = 'R'\n",
    "\n",
    "if ldebug:\n",
    "    if nduplicated > 0:\n",
    "        print('\\n* Removing Type R lines:')\n",
    "        for i in range(nlines_arc):\n",
    "            print(i, solution[i])\n",
    "    else:\n",
    "        print('\\n* No duplicated Type R lines have been found')\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we expect the wavelength calibration to be a smooth function with a small deviation from linearity, we can easily remove wrong lines which are going to appear very far from a robust linear fit (computed using the Theil-Sen method). An important parameter in this step is `times_sigma_theil_sen`, the number of times sigma that is employed to determine that a point is too far from the linear fit.\n",
    "\n",
    "For this procedure we define three auxiliary functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_data_for_fit(wv_master, xpos_arc, solution):\n",
    "    \"\"\"Select information from valid arc lines to facilitate posterior fits.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    wv_master : 1d numpy array, float\n",
    "        Array with wavelengths corresponding to the master table (Angstroms).\n",
    "    xpos_arc : 1d numpy array, float\n",
    "        Location of arc lines (pixels).\n",
    "    solution : ouput from previous call to arccalibration\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    nfit : int\n",
    "        Number of valid points for posterior fits.\n",
    "    ifit : list of int\n",
    "        List of indices corresponding to the arc lines which coordinates are\n",
    "        going to be employed in the posterior fits.\n",
    "    xfit : 1d numpy aray\n",
    "        X coordinate of points for posterior fits.\n",
    "    yfit : 1d numpy array\n",
    "        Y coordinate of points for posterior fits.\n",
    "    wfit : 1d numpy array\n",
    "        Cost function of points for posterior fits. The inverse of these values\n",
    "        can be employed for weighted fits.\n",
    "    \"\"\"\n",
    "\n",
    "    nlines_arc = len(solution)\n",
    "\n",
    "    if nlines_arc != xpos_arc.size:\n",
    "        raise ValueError('invalid nlines_arc=' + str(nlines_arc))\n",
    "\n",
    "    nfit = 0\n",
    "    ifit = []\n",
    "    xfit = np.array([])\n",
    "    yfit = np.array([])\n",
    "    wfit = np.array([])\n",
    "    for i in range(nlines_arc):\n",
    "        if solution[i]['lineok']:\n",
    "            ifit.append(i)\n",
    "            xfit = np.append(xfit, xpos_arc[i])\n",
    "            yfit = np.append(yfit, wv_master[solution[i]['id']])\n",
    "            wfit = np.append(wfit, solution[i]['funcost'])\n",
    "            nfit += 1\n",
    "    return nfit, ifit, xfit, yfit, wfit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_theil_sen(x, y):\n",
    "    \"\"\"Compute a robust linear fit using the Theil-Sen method.\n",
    "\n",
    "    See http://en.wikipedia.org/wiki/Theil%E2%80%93Sen_estimator for details.\n",
    "    This function \"pairs up sample points by the rank of their x-coordinates\n",
    "    (the point with the smallest coordinate being paired with the first point\n",
    "    above the median coordinate, etc.) and computes the median of the slopes of\n",
    "    the lines determined by these pairs of points\".\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : 1d numpy array, float\n",
    "        X coordinate.\n",
    "    y : 1d numpy array, float\n",
    "        Y coordinate.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    intercept : float\n",
    "        Intercept of the linear fit.\n",
    "    slope : float\n",
    "        Slope of the linear fit.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    if x.ndim == y.ndim == 1:\n",
    "        n = x.size\n",
    "        if n == y.size:\n",
    "            if n < 5:\n",
    "                raise ValueError('n=' + str(n) +' is < 5')\n",
    "            result = []  # python list\n",
    "            if (n % 2) == 0:\n",
    "                iextra = 0\n",
    "            else:\n",
    "                iextra = 1\n",
    "            for i in range(n//2):\n",
    "                ii = i + n//2 + iextra\n",
    "                deltax = x[ii]-x[i]\n",
    "                deltay = y[ii]-y[i]\n",
    "                result.append(deltay/deltax)\n",
    "            slope = np.median(result)\n",
    "            result = y - slope*x  # numpy array\n",
    "            intercept = np.median(result)\n",
    "            return intercept, slope\n",
    "        else:\n",
    "            raise ValueError('Invalid input sizes')\n",
    "    else:\n",
    "        raise ValueError('Invalid input dimensions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmaG(x):\n",
    "    \"\"\"Compute a robust estimator of the standard deviation.\n",
    "\n",
    "    See Eq. 3.36 (page 84) in Statistics, Data Mining, and Machine\n",
    "    in Astronomy, by Ivezic, Connolly, VanderPlas & Gray\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : 1d numpy array, float\n",
    "        Array of input values which standard deviation is requested.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    sigmag : float\n",
    "        Robust estimator of the standar deviation\n",
    "    \"\"\"\n",
    "\n",
    "    q25, q75 = np.percentile(x, [25.0, 75.0])\n",
    "    sigmag = 0.7413*(q75-q25)\n",
    "    return sigmag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the previous definitions of these functions, we can proceed with the fit and the removal of deviant points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ldebug:\n",
    "    print('\\n>>> Theil-Sen filtering...')\n",
    "nfit, ifit, xfit, yfit, wfit = \\\n",
    "  select_data_for_fit(wv_master, xpos_arc, solution)\n",
    "intercept, slope = fit_theil_sen(xfit, yfit)\n",
    "rfit = abs(yfit - (intercept + slope*xfit))\n",
    "if ldebug:\n",
    "    print('rfit:', rfit)\n",
    "sigma_rfit = sigmaG(rfit)\n",
    "if ldebug:\n",
    "    print('sigmaG:',sigma_rfit)\n",
    "nremoved = 0\n",
    "for i in range(nfit):\n",
    "    if rfit[i] > times_sigma_theil_sen*sigma_rfit:\n",
    "        solution[ifit[i]]['lineok'] = False\n",
    "        solution[ifit[i]]['type'] = 'T'\n",
    "        nremoved += 1\n",
    "\n",
    "if ldebug:\n",
    "    if nremoved > 0:\n",
    "        print('\\n* Removing Type T lines:')\n",
    "        for i in range(nlines_arc):\n",
    "            print(i, solution[i])\n",
    "    else:\n",
    "        print('\\nNo Type T lines have been found and removed')\n",
    "    pause(lpause)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the previous filtering, we continue by removing points that deviate from a polynomial fit. The filtered lines are labelled as type='P'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ldebug:\n",
    "    print('\\n>>> Polynomial filtering...')\n",
    "nfit, ifit, xfit, yfit, wfit = \\\n",
    "  select_data_for_fit(wv_master, xpos_arc, solution)\n",
    "coeff_fit = polynomial.polyfit(xfit, yfit, poly_degree, w=1/wfit)\n",
    "poly = polynomial.Polynomial(coeff_fit)\n",
    "rfit = abs(yfit - poly(xfit))\n",
    "if ldebug:\n",
    "    print('rfit:',rfit)\n",
    "sigma_rfit = sigmaG(rfit)\n",
    "if ldebug:\n",
    "    print('sigmaG:',sigma_rfit)\n",
    "nremoved = 0\n",
    "for i in range(nfit):\n",
    "    if rfit[i] > times_sigma_polfilt*sigma_rfit:\n",
    "        solution[ifit[i]]['lineok'] = False\n",
    "        solution[ifit[i]]['type'] = 'P'\n",
    "        nremoved += 1\n",
    "\n",
    "if ldebug:\n",
    "    if nremoved > 0:\n",
    "        print('\\n* Removing Type P lines:')\n",
    "        for i in range(nlines_arc):\n",
    "            print(i, solution[i])\n",
    "    else:\n",
    "        print('\\nNo type P lines have been found and removed')\n",
    "    pause(lpause)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After removing deviant lines, we can now try to incorporate identifications for those lines that remain unidentified after the previous procedure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ldebug:\n",
    "    print('\\n>>> Polynomial prediction of unknown lines...')\n",
    "nfit, ifit, xfit, yfit, wfit = \\\n",
    "  select_data_for_fit(wv_master, xpos_arc, solution)\n",
    "coeff_fit = polynomial.polyfit(xfit, yfit, poly_degree_wfit, w=1/wfit)\n",
    "poly = polynomial.Polynomial(coeff_fit)\n",
    "rfit = abs(yfit - poly(xfit))\n",
    "if ldebug:\n",
    "    print('rfit:',rfit)\n",
    "sigma_rfit = sigmaG(rfit)\n",
    "if ldebug:\n",
    "    print('sigmaG:',sigma_rfit)\n",
    "\n",
    "list_id_already_found = []\n",
    "list_funcost_already_found = []\n",
    "for i in range(nlines_arc):\n",
    "    if solution[i]['lineok']:\n",
    "        list_id_already_found.append(solution[i]['id'])\n",
    "        list_funcost_already_found.append(solution[i]['funcost'])\n",
    "\n",
    "nnewlines = 0\n",
    "for i in range(nlines_arc):\n",
    "    if not solution[i]['lineok']:\n",
    "        zfit = poly(xpos_arc[i]) # predicted wavelength\n",
    "        isort = np.searchsorted(wv_master, zfit)\n",
    "        if isort == 0:\n",
    "            ifound = 0\n",
    "            dlambda = wv_master[ifound]-zfit\n",
    "        elif isort == nlines_master:\n",
    "            ifound = isort - 1\n",
    "            dlambda = zfit - wv_master[ifound]\n",
    "        else:\n",
    "            dlambda1 = zfit-wv_master[isort-1]\n",
    "            dlambda2 = wv_master[isort]-zfit\n",
    "            if dlambda1 < dlambda2:\n",
    "                ifound = isort - 1\n",
    "                dlambda = dlambda1\n",
    "            else:\n",
    "                ifound = isort\n",
    "                dlambda = dlambda2\n",
    "        if ldebug:\n",
    "            print(i,zfit,ifound,dlambda)\n",
    "        if ifound not in list_id_already_found:  # unused line\n",
    "            if dlambda < times_sigma_inclusion*sigma_rfit:\n",
    "                list_id_already_found.append(ifound)\n",
    "                solution[i]['lineok'] = True\n",
    "                solution[i]['type'] = 'I'\n",
    "                solution[i]['id'] = ifound\n",
    "                solution[i]['funcost'] = max(list_funcost_already_found)\n",
    "                nnewlines += 1\n",
    "\n",
    "if ldebug:\n",
    "    if nnewlines > 0:\n",
    "        print('\\n* Including Type I lines:')\n",
    "        for i in range(nlines_arc):\n",
    "            print(i, solution[i])\n",
    "    else:\n",
    "        print('\\nNo Type I lines have been found and added')\n",
    "    pause(lpause)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the solution has been found, we can finally determine the coefficients of the final wavelength calibration fit. For this purpose, we use an auxiliary function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_solution(wv_master,\n",
    "                 xpos_arc,\n",
    "                 solution,\n",
    "                 naxis1_arc,\n",
    "                 poly_degree_wfit,\n",
    "                 weighted,\n",
    "                 ldebug=False,\n",
    "                 lplot=False,\n",
    "                 lpause=False):\n",
    "    \"\"\"Fit polynomial to arc calibration solution.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    wv_master : 1d numpy array, float\n",
    "        Array with wavelengths corresponding to the master table (Angstroms).\n",
    "    xpos_arc : 1d numpy array, float\n",
    "        Location of arc lines (pixels).\n",
    "    solution : list of dictionaries\n",
    "        A list of size equal to the number of arc lines, which elements are\n",
    "        dictionaries containing all the relevant information concerning the\n",
    "        line identification.\n",
    "    naxis1_arc : int\n",
    "        NAXIS1 of arc spectrum.\n",
    "    poly_degree_wfit : int\n",
    "        Polynomial degree corresponding to the wavelength calibration function\n",
    "        to be fitted.\n",
    "    weighted : bool\n",
    "        Determines whether the polynomial fit is weighted or not, using as\n",
    "        weights the values of the cost function obtained in the line\n",
    "        identification.\n",
    "    ldebug : bool\n",
    "        If True intermediate results are displayed.\n",
    "    lplot : bool\n",
    "        If True intermediate plots are displayed.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    coeff : 1d numpy array, float\n",
    "        Coefficients of the polynomial fit.\n",
    "    crval1_approx : float\n",
    "        Approximate CRVAL1 value.\n",
    "    cdetl1_approx : float\n",
    "        Approximate CDELT1 value.\n",
    "    \"\"\"\n",
    "\n",
    "    nlines_arc = len(solution)\n",
    "\n",
    "    if nlines_arc != xpos_arc.size:\n",
    "        raise ValueError('Invalid nlines_arc=' + str(nlines_arc))\n",
    "\n",
    "    # Select information from valid lines.\n",
    "    nfit, ifit, xfit, yfit, wfit = \\\n",
    "      select_data_for_fit(wv_master, xpos_arc, solution)\n",
    "\n",
    "    # Select list of filtered out and unidentified lines\n",
    "    list_R = []\n",
    "    list_T = []\n",
    "    list_P = []\n",
    "    list_unidentified = []\n",
    "    for i in range(nlines_arc):\n",
    "        if not solution[i]['lineok']:\n",
    "            if solution[i]['type'] is None:\n",
    "                list_unidentified.append(i)\n",
    "            elif solution[i]['type'] == 'R':\n",
    "                list_R.append(i)\n",
    "            elif solution[i]['type'] == 'T':\n",
    "                list_T.append(i)\n",
    "            elif solution[i]['type'] == 'P':\n",
    "                list_P.append(i)\n",
    "            else:\n",
    "                raise ValueError('Unexpected \"type\"')\n",
    "\n",
    "    # Obtain approximate linear fit using the robust Theil-Sen method.\n",
    "    intercept, slope = fit_theil_sen(xfit, yfit)\n",
    "    cdelt1_approx = slope\n",
    "    crval1_approx = intercept + cdelt1_approx\n",
    "    if ldebug:\n",
    "        print('>>> Approximate CRVAL1 :',crval1_approx)\n",
    "        print('>>> Approximate CDELT1 :',cdelt1_approx)\n",
    "\n",
    "    # Polynomial fit.\n",
    "    if weighted:\n",
    "        coeff = polynomial.polyfit(xfit, yfit, poly_degree_wfit, w=1/wfit)\n",
    "    else:\n",
    "        coeff = polynomial.polyfit(xfit, yfit, poly_degree_wfit)\n",
    "\n",
    "    if ldebug:\n",
    "        print('>>> Fitted coefficients:', coeff)\n",
    "    xpol = np.linspace(1,naxis1_arc,naxis1_arc)\n",
    "    poly = polynomial.Polynomial(coeff)\n",
    "    ypol = poly(xpol)-(crval1_approx+(xpol-1)*cdelt1_approx)\n",
    "\n",
    "    if lplot:\n",
    "        fig = plt.figure()\n",
    "        ax = fig.add_subplot(111)\n",
    "        # identified lines\n",
    "        xp = np.copy(xfit)\n",
    "        yp = yfit-(crval1_approx+(xp-1)*cdelt1_approx)\n",
    "        ax.set_xlim([1-0.05*naxis1_arc,naxis1_arc+0.05*naxis1_arc])\n",
    "        ax.set_xlabel('pixel position in arc spectrum')\n",
    "        ax.set_ylabel('residuals (Angstrom)')\n",
    "        ax.plot(xp, yp, 'go', label=\"identified lines\")\n",
    "        for i in range(nfit):\n",
    "            ax.text(xp[i], yp[i], solution[ifit[i]]['type'], fontsize=15)\n",
    "        # polynomial fit\n",
    "        ax.plot(xpol,ypol,'c-', label=\"polynomial fit\")\n",
    "        # unidentified lines\n",
    "        if len(list_unidentified) > 0:\n",
    "            ymin = np.concatenate((yp,ypol)).min()\n",
    "            ymax = np.concatenate((yp,ypol)).max()\n",
    "            for i in list_unidentified:\n",
    "                xp = np.array([xpos_arc[i], xpos_arc[i]])\n",
    "                yp = np.array([ymin, ymax])\n",
    "                if i == list_unidentified[0]:\n",
    "                    ax.plot(xp,yp,'r--',label='unidentified lines')\n",
    "                else:\n",
    "                    ax.plot(xp,yp,'r--')\n",
    "        # R, T and P lines\n",
    "        for val in zip([\"R\",\"T\",\"P\"],[list_R, list_T, list_P],['red','blue','magenta']):\n",
    "            list_X = val[1]\n",
    "            if len(list_X) > 0:\n",
    "                xp = np.array([])\n",
    "                yp = np.array([])\n",
    "                for i in list_X:\n",
    "                    xp = np.append(xp, xpos_arc[i])\n",
    "                    yp = np.append(yp, wv_master[solution[i]['id']])\n",
    "                yp -= crval1_approx+(xp-1)*cdelt1_approx\n",
    "                ax.plot(xp, yp, marker='x', markersize = 15, c=val[2],\n",
    "                        linewidth=0, label='removed line ('+val[0]+')')\n",
    "\n",
    "        # shrink current axis by 20% and put a legend to the right\n",
    "        box = ax.get_position()\n",
    "        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])\n",
    "        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "                \n",
    "        plt.show()\n",
    "        pause(lpause)\n",
    "\n",
    "    return coeff, crval1_approx, cdelt1_approx\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coeff, crval1_approx, cdelt1_approx = \\\n",
    "  fit_solution(wv_master,\n",
    "               xpos_arc,\n",
    "               solution,\n",
    "               naxis1_arc,\n",
    "               poly_degree_wfit = 2,\n",
    "               weighted = False,\n",
    "               ldebug=ldebug,\n",
    "               lplot=lplot,\n",
    "               lpause=lpause)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('>>> original crval1, cdelt1...:', crval1_arc, cdelt1_arc)\n",
    "print('>>> original c0, c1, c2.......:', c0_arc, c1_arc, c2_arc)\n",
    "print('>>> approximate crval1, cdelt1:', crval1_approx, cdelt1_approx)\n",
    "print('>>> fitted coefficients.......:', coeff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
