Weighted mask tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Weighted mask cross-correlation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why a weighted mask?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have seen (or we will see!) the power of the likelihood frameworks developed by Brogi & Line (2019), Gibson et al. (2020). They provide the highest constraining power and can get precise confidence intervals for atmospheric parameters. So why even bothering with alternative methods?\n",
    "\n",
    "I see two reasons why you should consider a weighted mask (sometimes called \"binary\") approach.\n",
    "1. It is the most **model independent** method I am aware of to detect a molecule/atom;\n",
    "2. It allows you to visualize what you are fitting.\n",
    "\n",
    "It is possible to perform a statistically sound fit also with the binary mask method (Pino et al., 2020). It seems to deliver a lower precision on model parameters. However, this is an avenue that I have not fully explored and that could provide complementary information to the cross-correlation-likelihood mapping methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is a weighted binary mask?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A set of apertures that collect all the flux that falls within them, and exclude everything outside of them. We may want to give weights that account for the amount of \"information\" in each aperture.\n",
    "\n",
    "The only input information we need is the position of the lines we want to include in the mask (which we can get from line lists such as VALD, ExoMOL, ...), and, if we want, a set of weights. As you can see, these are only a subset of the information that we need to include in an atmospheric model that we use for cross-correlating. In this sense, this approach is less model dependent (if you use weight, it still is)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"Binary_mask.gif\" width=\"750\" align=\"center\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is also called \"stacking\", and is a method used in extragalctic studies to stack faint lines from many different galaxies to build up SNR. Here, we cumulate many lines belonging to the same species, but on a single planet. As an example, here is a model (in red) of a water + clouds transmission spectrum of HD189733b, and in black is the \"noisy\" version where I have added simulated ESPRESSO white noise corresponding to one single transit of the planet."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\" Noisy_model.png\" width=\"750\" align=\"center\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we apply the binary mask stacking on a specific water band, by taking from the strongest to the weakest line. With 100 - 1000 lines, we get a strong and clear detection (of course, this is a somewhat ideal case. This only works with **white noise**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table><tr>\n",
    "<td> <img src=\"Binary_mask_application_simulated_data.gif\" width=\"500\" align=\"left\">\n",
    "<td> <img src=\"Binary_mask_application_simulated_data_2.gif\" width=\"500\" align=\"right\">\n",
    "</tr></table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can have different kinds of masks, including specific to search some atoms or molecules or built starting from stellar data. David will cover more of the binary mask concept later.\n",
    "\n",
    "Now, I give you my implementation of this concept."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Weighted mask implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definitions, and preparing the ESPRESSO pipeline weighted mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mask has the form of a list of wavelengths and the corresponding weights. There are two additional \"meta-parameters\" that we need to decide on, but we have good arguments for their choice and in my opinion they do not introduce additional model dependence.\n",
    "1. Every individual aperture has a size in wavelength (or better radial velocity) which we need to specify. This controls the amount of flux that is collected in that aperture;\n",
    "2. We do not know in advance the position of the rest frame of the exoplanet atmosphere signal. Thus, we will need to slide the mask on the spectrum, to identify that. We will do this in velocity space, and we need to decide on which grid.\n",
    "To determine the best values for these parameters, we need to be familiar with the difference between \"resolving power\" and \"resolution\". Note that in the literature there is often some confusion on these terms. I believe that this is the correct nomenclature from the spectroscopy community, but you may often find that the terms are used interchangeably. Always take care on what is meant precisely! Here are my definitions.\n",
    "\n",
    "*Resolving power*: c/FWHM_LSF [km/s]. This controls the size of the LSF, and thus determines how far can two spectrscopic lines be in order for us to resolve them, and be able to individually characterize them.\n",
    "*Resolution*: c/pixel_size [km/s]. This controls the size of a single pixel. Why is it different compared to the resolving power? For instruments like ESPRESSO, we aim to a precision in RV which is much much higher compared to the resolving power. We achieve that with several tricks which rely on the stability of the instrument, i.e., the LSF does not chang in time during the night. We want to be able to monitor this, and to know the shape of the LSF very precisely. Thus, we oversample it with several pixels.\n",
    "\n",
    "For ESPRESSO, we see that R = 140.000 is the resolving power. It corresponds to 2.14 km/s. The LSF is oversample 4.5 times, so the individual pixel size is 0.475 km/s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Aperture**: The optimal choice collects the largest amount of flux from photons which are independent from each other (to avoid having correlations). Thus, the unit to which we need to compare is the resolution. We chose each aperture to be exactly 0.475 km/s wide. A larger aperture will result in unnecesssary correlation. A smaller aperture will result in the loss of some flux, and thus higher photon noise.\n",
    "\n",
    "**RV grid**: How far do we go? We need to have an idea of the maximum orbital speed of the planet. We can get that from the period and semi-major axis of the planet. For WASP-7b:\n",
    "\n",
    "$K_\\mathrm{p} = 2\\pi a/P \\sim 200 \\mathrm{km\\ s^{-1}}$\n",
    "\n",
    "Step size: This is not crucial at this point, but if we want to perform a fit on this weighted binary mask (see Pino et al., 2020), we want every point to be statistically independent. The safe choice is to make sure that each point comes from an individually resolved spectral line, i.e. we step by one resolving power element (you probably could do a bit better, but careful with correlations!). Here, I put a smaller step just to show you a well sampled CCF, but consider that there may be some correlation between neighboring points in the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from astropy.io import fits\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from astropy import visualization\n",
    "\n",
    "aperture_width = 0.4759\n",
    "RV_table = np.arange(-210., 210., 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first load the final products of our previous pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here is where we were at.\n",
    "\n",
    "S2D_wavelengths_vacuum_laboratory = np.load('S2D_wavelengths_vacuum_laboratory.npy')\n",
    "S2D_spectra_stellar_rest_frame_normalized = np.load('S2D_spectra_barycentric_normalized.npy')\n",
    "\n",
    "# (exposures, orders, pixels)\n",
    "np.shape(S2D_wavelengths_vacuum_laboratory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mask is independent of exposure (lines are always the same), but we will have to build a separate mask for every order. Thus, the shape of our mask will be (n_lines, n_RV_points, n_orders). Since n_lines differs from order to order, this is no numpy array.\n",
    "\n",
    "To start with, we will use the masks providede by the ESPRESSO pipeline (when you install it). Let's read one of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binary_mask_directory = '/data2/pino/Data/binary_masks/'\n",
    "name = 'G2' # Try also: F9, G8, G9, K2, K6, M0, M2, M3, M4, M5\n",
    "\n",
    "mask_file = fits.open(binary_mask_directory + 'ESPRESSO_pipeline/' + 'ESPRESSO_{}.fits'.format(name))\n",
    "\n",
    "mask_wavelengths_air = mask_file[1].data['lambda'] # Wavelengths are provided in air.\n",
    "mask_line_contrast = mask_file[1].data['contrast'] # Built on an observed template stellar spectum\n",
    "\n",
    "print(np.shape(mask_wavelengths_air))\n",
    "print(np.shape(mask_line_contrast))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is 1D a list of wavelengths and weights. It is the input we need, but we need to rearrange it to apply it to the S2D spectra we have. First, let's convert the wavelengths to vacuum, where the observed spectra are given."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def n_air(mask_wavelengths_air):\n",
    "    t=15. \n",
    "    p=760.\n",
    "    \n",
    "    # wavelength in A, t=temperature in C, p=pression in millibar\n",
    "    n_air = 1e-6 * p * (1 + (1.049-0.0157*t)*1e-6*p) / 720.883 / (1 + 0.003661*t) * (64.328 + 29498.1/(146-(1e3/mask_wavelengths_air/10.)**2) + 255.4/(41-(1e3/mask_wavelengths_air/10.)**2))\n",
    "    return n_air + 1\n",
    "\n",
    "mask_wavelengths_vac = mask_wavelengths_air * n_air(mask_wavelengths_air)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now rearrange by selecting only wavelengths that are fully contained in each order for all RV_shifts that we consider. We do not want lines to be only present in a subset of exposures, otherwise we would combine different information across exposures. You could in principle improve the SNR by taking all lines, but then you should be careful with the interpretation of your operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To account for the fact that the blue CCD has less useful pixels, let's put nan in place of the wavelength value in the parts of the CCD we do not want to include.\n",
    "# Since the wavelength array is in common across spectra, we only build this for S2D_wavelengths_vacuum_laboratory[0] \n",
    "wave_appo = np.where(S2D_spectra_stellar_rest_frame_normalized[0]==0, np.nan, S2D_wavelengths_vacuum_laboratory[0])\n",
    "wave_min = np.nanmin(wave_appo, axis=1) # One for each order. Shape: (n_exposures, n_pixels)\n",
    "wave_max = np.nanmax(wave_appo, axis=1)\n",
    "\n",
    "# The bluest Doppler shift is given by the minimum RV considered in our mask, while the maximum Doppler shift is given by the maximum one.\n",
    "# This function thus searches for the position of wave_min and wave_max inside the 1D mask provided by the ESPRESSO pipeline.\n",
    "indexes_start_order = np.searchsorted(mask_wavelengths_vac, wave_min*(1.-RV_table[0]/2.99792458e5)+wave_min*aperture_width/2./2.99792458e5)\n",
    "indexes_end_order = np.searchsorted(mask_wavelengths_vac, wave_max*(1.-RV_table[-1]/2.99792458e5)-wave_max*aperture_width/2./2.99792458e5)\n",
    "\n",
    "mask_wavelengths_vac_S2D = []\n",
    "mask_line_contrast_S2D = []\n",
    "\n",
    "# We use the indexes to build the 2D mask.\n",
    "for n, (i, j) in enumerate(zip(indexes_start_order, indexes_end_order)):\n",
    "    mask_wavelengths_vac_S2D.append(mask_wavelengths_vac[i:j])\n",
    "    mask_line_contrast_S2D.append(mask_line_contrast[i:j])\n",
    "    \n",
    "n_lines_per_order = []    \n",
    "for order in range(len(mask_wavelengths_vac_S2D)):\n",
    "    n_lines_per_order.append(len(mask_wavelengths_vac_S2D[order]))\n",
    "    \n",
    "plt.plot(n_lines_per_order)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A practical way to make the computation faster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We follow Pino et al. (2020) Appendix B and, rather than cycling over the mask apertures, we build a weight for every pixel based on whether it falls within a mask or not. This is convenient for several reasons.\n",
    "1. It allows matrix calculation that is much faster;\n",
    "2. It allows us to store the mask for future use;\n",
    "3. It makes it much simpler to implement other weights in our formalism (see later).\n",
    "\n",
    "In this formalism, for every pixel we build the following weight:\n",
    "$$w_\\mathrm{pixel} = w_\\mathrm{mask}\\cdot\\overline{\\Delta \\lambda}\\ ,$$\n",
    "where $w_\\mathrm{mask}$ is the weight of the mask in that pixel (could be 0 if no aperture falls on that pixel), and $\\overline{\\Delta \\lambda}$ could be the entire wavelength range spun by the pixel if it falls entirely in a mask, or the corresponding section.\n",
    "\n",
    "<img src=\"Binary_mask_implementation.png\" width=\"750\" align=\"center\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_pixel_by_pixel_mask_single_order(l,mask_ll,mask_W,RV_table,bin_width):\n",
    "    \"\"\"Project a list of weights for a binary mask on a list of pixels. \n",
    "\n",
    "    Takes into account:\n",
    "    - Non-constant pixel size in wavelength dimensions\n",
    "    - Partial overlap between binary mask and pixel\n",
    "\n",
    "    Input\n",
    "    -------\n",
    "\n",
    "    l: floats[n]\n",
    "        List of central wavelengths of the pixels in the detector\n",
    "\n",
    "    mask_ll, mask_W: floats[m]\n",
    "        Lists containing:\n",
    "        - The central wavelenghts of the lines in the binary mask\n",
    "        - The weights corresponding to those lines\n",
    "\n",
    "    RV_table: floats[n_rv]\n",
    "        List containing the array of radial velocities over which the CCF will be computed.\n",
    "        Every time we shift the binary mask, a new set of weights needs to be computed.\n",
    "\n",
    "    bin_width: float\n",
    "        - Size of the binary mask windows in units of km s-1\n",
    "\n",
    "\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    \n",
    "    weights_pixels: floats[n_rv, n]\n",
    "        Weights to be attributed to single pixels\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "\n",
    "    Below you find a version of the routine that computes weights as mask_weight/\\overbar{delta_lambda}_n (notation Pino et al. 2019).\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "\n",
    "    # Compute mask width in wavelength units, accounting for shift\n",
    "    mask_ll_shift = np.outer(mask_ll,(1.+RV_table/2.99792458e5))\n",
    "    mask_width = mask_ll_shift*bin_width/2.99792458e5\n",
    "\n",
    "\n",
    "    #\n",
    "    # Preparation, everything in matrix form\n",
    "    #\n",
    "\n",
    "    # wavelengths corresponding to the beginning and the ending of the mask bins.\n",
    "    lstart = mask_ll_shift - mask_width/2.\n",
    "    lstop = mask_ll_shift + mask_width/2.\n",
    "\n",
    "\n",
    "    av_band = np.stack([lstart,lstop], axis=2)\n",
    "\n",
    "    # Compute extrema of pixels in wavelength space\n",
    "\n",
    "    wave_bins = [ [(l[i-1] + l[i])/2.,(l[i] + l[i+1])/2.] for i in range(1,len(l)-1)]\n",
    "    wave_bins.insert(0, [l[0] - (l[1] - l[0])/2., l[0] + (l[1] - l[0])/2. ])\n",
    "    wave_bins.append([l[-1] - (l[-1] - l[-2])/2., l[-1] + (l[-1] - l[-2])/2. ])\n",
    "    wave_bins = np.array(wave_bins)\n",
    "    \n",
    "    # Indexes of pixels within which the binary mask starts and stops ([n_rv, m])\n",
    "    i_start = np.searchsorted(wave_bins[:,1], av_band[:,:,0])\n",
    "    i_end = np.searchsorted(wave_bins[:,1], av_band[:,:,1])\n",
    "\n",
    "    # Create a mask that weights every pixel based on:\n",
    "    # (1) How much of the pixel is illuminated;\n",
    "    # (2) The weight of the binary mask\n",
    "\n",
    "    # In each pixel, the weight is computed as mask_weight * \\overbar{delta_lambda} (notation Pino et al. 2020), where \\overbar{delta_lambda}_n can be:\n",
    "    # - delta_lambda_n, i.e the full pixel. This is the first case below.\n",
    "    # - Only a part of a pixel. The other if/else clauses deal with that\n",
    "    weights_pixels = [] \n",
    "    for n_rv in range(len(i_start[1])):\n",
    "        pixel_interm = []\n",
    "        mask_pixel_single_rv = np.zeros(len(wave_bins))\n",
    "        for m, (i_s, i_e) in enumerate(zip(i_start[:, n_rv], i_end[:,n_rv])):\n",
    "\n",
    "            # Pixels fully contained within a binary mask aperture\n",
    "            if any(range(i_s+1, i_e)):\n",
    "                mask_pixel_single_rv[range(i_s+1, i_e)] = mask_W[m]*np.diff(wave_bins[range(i_s+1, i_e)], axis=1)[:,0]\n",
    "\n",
    "            # The binary mask aperture is fully contained within a pixel\n",
    "            if i_s == i_e:\n",
    "                mask_pixel_single_rv[i_s] = mask_W[m]*np.diff(av_band[m, n_rv])\n",
    "\n",
    "            # Binary mask partially covers the pixels found at the start and end\n",
    "            else:\n",
    "                mask_pixel_single_rv[i_s] = mask_W[m]*(wave_bins[i_s,1]-av_band[m, n_rv,0])\n",
    "                mask_pixel_single_rv[i_e] = mask_W[m]*(av_band[m, n_rv,1] - wave_bins[i_e,0])              \n",
    "        weights_pixels.append(mask_pixel_single_rv)\n",
    "    weights_pixels = np.array(weights_pixels)\n",
    "\n",
    "    return weights_pixels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_by_pixel_mask = np.zeros((len(S2D_wavelengths_vacuum_laboratory[0]), len(RV_table), len(S2D_wavelengths_vacuum_laboratory[0,0]))) # (n_orders, n_RVs, n_pizels)\n",
    "\n",
    "for order, (w, ww_ll, weights) in enumerate(zip(S2D_wavelengths_vacuum_laboratory[0], mask_wavelengths_vac_S2D, mask_line_contrast_S2D)):\n",
    "    print(\"Generating mask for order {}\".format(order))\n",
    "    if any(weights != 0.): \n",
    "        pixel_by_pixel_mask[order] = build_pixel_by_pixel_mask_single_order(w, ww_ll, weights, RV_table, aperture_width)\n",
    "    else: \n",
    "        pixel_by_pixel_mask[order] = np.zeros( (len(RV_table), len(w)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('pixel_by_pixel_mask', pixel_by_pixel_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for w, s, m in zip(S2D_wavelengths_vacuum_laboratory[50], S2D_spectra_stellar_rest_frame_normalized[50], pixel_by_pixel_mask[:, 400]):\n",
    "    plt.plot(w, s, 'k', w, 100.*(m)+1, 'r')\n",
    "\n",
    "plt.xlim(4110, 4150)\n",
    "plt.ylim(0.8, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to calculate the cross-correlation as a simple matrix product! However, let's consider for a second if we are done or not. Let's have a look at the spectra again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for w, o in zip(S2D_wavelengths_vacuum_laboratory[50], S2D_spectra_stellar_rest_frame_normalized[50]):\n",
    "    plt.plot(w, o, 'k')\n",
    "\n",
    "# plt.xlim(4000., 4010)\n",
    "plt.ylim(0.5, 1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for w, o in zip(S2D_wavelengths_vacuum_bary[50], S2D_spectra_stellar_rest_frame_normalized[50]):\n",
    "    plt.plot(w, o, 'k')\n",
    "\n",
    "plt.xlim(4000., 4040)\n",
    "plt.ylim(0.5, 1.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clearly, if we simply apply the weighted mask, we will be dominated by the noise picked in the bad pixels. We thus apply two additional sets of weights:\n",
    "1. We weight orders with respect to each other, so that orders with overall higher SNR are weighted more.\n",
    "2. Within a single order, we need to downweight the borders where we receive very little flux (crucial in the blue CCD!)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weights on orders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the signal-to-noise registered in each order that is found in the headers of the ESPRESSO. Note that, due to color effects, the relative weights of orders could change throughout the night and we need to account for that. So we calculate weights on an exposure-by-exposure basis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_sorted_hdus(input_data_directory, data_type, target_name):\n",
    "    \n",
    "    # Get all data in folder of the specified type\n",
    "    import glob\n",
    "    data_files = np.array(glob.glob(input_data_directory + '*_{}_*'.format(data_type)))\n",
    "    \n",
    "    # Read date and time from the filename (standardized)\n",
    "    date_time_observations = get_date_time_observations(data_files, data_type=data_type)\n",
    "    \n",
    "    # Sort exposures in time\n",
    "    time_sorted_exposures_mask = compute_time_sorted_exposures_mask(date_time_observations)\n",
    "    data_files_sorted = data_files[time_sorted_exposures_mask]\n",
    "\n",
    "    # Open the fits and return them. Not the cheapest in terms of memory, but we'll use some of the info in the headers\n",
    "    hdus = []\n",
    "    for file_name in data_files_sorted:\n",
    "        hdu = fits.open(file_name)\n",
    "        if hdu[0].header['HIERARCH ESO OBS TARG NAME'] == target_name:\n",
    "            hdus.append(hdu)    \n",
    "\n",
    "    return hdus\n",
    "\n",
    "def get_date_time_observations(data_files, data_type):\n",
    "    \"\"\" Extract the date and time of an expsure from its name.\n",
    "    \"\"\"\n",
    "    import re\n",
    "    re.compile(data_type)\n",
    "\n",
    "    date_time_observations = [] \n",
    "\n",
    "    for file in data_files:\n",
    "        start = re.search(data_type + '_', file).end()\n",
    "\n",
    "        date = file[start:start+10]\n",
    "        time = file[start+11:start+23]\n",
    "\n",
    "        date_time_observations.append([date, time])\n",
    "\n",
    "    return date_time_observations\n",
    "\n",
    "def compute_time_sorted_exposures_mask(date_time_observations):\n",
    "    \"\"\" From strings to datetime objects\n",
    "    \"\"\"\n",
    "\n",
    "    date_time_observations_datetime=[]\n",
    "    from datetime import datetime\n",
    "    for dt in date_time_observations:\n",
    "        date_time_observations_datetime.append(datetime.strptime(dt[0] + ' ' + dt[1][:-4], \"%Y-%m-%d %H:%M:%S\"))    \n",
    "\n",
    "    return np.argsort(date_time_observations_datetime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data_directory = '/data2/pino/Data/observations/ESPRESSO/W76b_2018_10_31/'\n",
    "data_type='S2D_SKYSUB_A'\n",
    "target_name='BD+01   316' # From the header of ESPRESSO files\n",
    "\n",
    "hdus_S2D = read_sorted_hdus(input_data_directory = input_data_directory, data_type=data_type, target_name=target_name)\n",
    "\n",
    "SNR_per_exposure_per_order = []\n",
    "for n_exp, hdu in enumerate(hdus_S2D):\n",
    "    SNR_per_exposure_per_order.append([])\n",
    "    for order in range(1, 171):\n",
    "        SNR_per_exposure_per_order[n_exp].append(hdu[0].header['HIERARCH ESO QC ORDER{} SNR'.format(order)])\n",
    "\n",
    "SNR_per_exposure_per_order = np.array(SNR_per_exposure_per_order)\n",
    "\n",
    "# We place an arbitrary cut to exclude low SNR orders. In these orders, we have no guarantee that we are dominated by gaussian noise.\n",
    "order_weights = np.where(SNR_per_exposure_per_order>25., SNR_per_exposure_per_order, 0.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We remove a significant amount of orders. This is an arbitrary choice, so it's up to our experience and tests to understand what is best here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "100.*np.sum(order_weights == 0.)/(np.shape(S2D_spectra_stellar_rest_frame_normalized)[0]*np.shape(S2D_spectra_stellar_rest_frame_normalized)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weights within orders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even within a single order, we can clearly see that the noise is a function of wavelength. In some spectrographs (e.g., HARPS-N), this could be a negligible effect. However, ESPRESSO has large orders for which this effect is significant. As we have seen in the previous notebook, these oscillations are due to the variable throughput from the blaze function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data_directory = '/data2/pino/Data/observations/ESPRESSO/W76b_2018_10_31/'\n",
    "\n",
    "import glob\n",
    "data_files_A_blaze = glob.glob(input_data_directory + '*BLAZE_A_*')\n",
    "blaze_profiles = fits.open(data_files_A_blaze[0])[1].data\n",
    "\n",
    "# np.einsum is a clear way to perform matrix multiplication in any dimension. It is also pretty optimized.\n",
    "#\n",
    "# Syntax: (input indexes, output indexes. Reapeated indexes not appearing in the output are contracted according to tensor operations contraction of indexes)\n",
    "#\n",
    "# Here: j = n_orders, l = n_radial_velocities, m = n_pixels. \n",
    "blaze_weighted_pixel_mask = np.einsum('jlm,jm->jlm', pixel_by_pixel_mask, blaze_profiles)\n",
    "np.save('blaze_weighted_pixel_mask', blaze_weighted_pixel_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for w, s, m in zip(S2D_wavelengths_vacuum_laboratory[50], S2D_spectra_stellar_rest_frame_normalized[50], blaze_weighted_pixel_mask[:, 400]):\n",
    "    plt.plot(w, s, 'k', w, 150.*(m)+1, 'r')\n",
    "\n",
    "plt.xlim(4110, 4150)\n",
    "plt.ylim(0.8, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculation of the CCF as matrix multiplication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# i = n_exposures, j = n_orders, l = n_radial_velocity, m = n_pixel.\n",
    "#\n",
    "# We sum on all indexes excpet exposures and RV. So in the end we have one CCF for every exposure, and for every shift we selected\n",
    "ccfs_summed = np.einsum('ij,jlm,ijm->il', order_weights, blaze_weighted_pixel_mask, S2D_spectra_stellar_rest_frame_normalized)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we may want to preserve the real contrast of lines, to have a more direct interpretation of what we see. Pino et al. (2020) demonstrates that this can be achieved by normalizing by all the weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalization_ccfs = np.einsum('ij,jlm->il', order_weights, pixel_by_pixel_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ccfs_summed = ccfs_summed[:-1]\n",
    "normalization_ccfs = normalization_ccfs[:-1]\n",
    "\n",
    "ccfs_summed_normalized = ccfs_summed/normalization_ccfs\n",
    "\n",
    "for c in ccfs_summed_normalized:\n",
    "    plt.plot(c)\n",
    "    \n",
    "# plt.xlim(380, 440)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "norm = visualization.ImageNormalize(ccfs_summed_normalized, interval=visualization.ZScaleInterval())\n",
    "imshow = ax.imshow(ccfs_summed_normalized, norm=norm, origin='lower', aspect='auto', cmap= matplotlib.cm.get_cmap('Blues'), interpolation='none')\n",
    "ax.set_xlabel('RV point')\n",
    "ax.set_ylabel('Exposure')\n",
    "fig.colorbar(imshow)\n",
    "fig.tight_layout()\n",
    "\n",
    "# plt.xlim(400, 500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The slope that we see is an artificial product of the missing color-correction and of the pipeline. We also see some differences in the continua, possibly due to the fact that the planet + Doppler shadow signal is extremely strong in this planet. We thus normalize the CCFs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CCF post-processing. What are we seeing?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's calculate at which radial velocity we are expecting the planet. To do that, we need to calculate the planet phases that we are observing, and use the known orbital solution from the literature. We need to take care of the time conversions: the known ephemerid is given in BJD TDB, so we need to convert our exposure times to that same time scale and format. We use astropy for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy import time\n",
    "from astropy.coordinates import SkyCoord, EarthLocation\n",
    "\n",
    "target_coordinates = SkyCoord.from_name(target_name)\n",
    "site = EarthLocation.of_site('paranal')\n",
    "\n",
    "time_observations = []\n",
    "\n",
    "for hdu in hdus_S2D:\n",
    "    exp_start=hdu[0].header['MJD-OBS']\n",
    "    exp_dur=hdu[0].header['EXPTIME']\n",
    "    time_observations.append(exp_start + exp_dur/(60.*60.*24.))\n",
    "\n",
    "time_observations=time.Time(np.array(time_observations), format='mjd', scale='utc', location=site)\n",
    "\n",
    "# We compute the time travel difference to the barycenter of the Solar System\n",
    "dt_bary = time_observations.light_travel_time(target_coordinates)\n",
    "\n",
    "# For the barycentric times, which are of our interest, the correct scale is TDB. (https://docs.astropy.org/en/stable/time/index.html#barycentric-and-heliocentric-light-travel-time-corrections).\n",
    "bjd_time_observations = time_observations.tdb + dt_bary  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T0 =   2457273.4191 # BJD TDB; https://ui.adsabs.harvard.edu/abs/2022ApJS..258...40K/abstract\n",
    "P_days = 1.8098806\n",
    "r_AU = 0.0330\n",
    "\n",
    "phases = ((bjd_time_observations.jd-T0)/P_days)%1\n",
    "phases = np.where(phases < 0.5, 1+phases, phases) # To have them as a monotic sequence\n",
    "\n",
    "r_km = r_AU*1.496e8\n",
    "P_sec = P_days*24.*60.*60.\n",
    "v_max = 2*np.pi*r_km/P_sec\n",
    "velocities_spun_during_transit = v_max*np.sin(phases*2.*np.pi)\n",
    "\n",
    "k_star = 0.1193 # km/s\n",
    "stellar_phases = 0.5+phases\n",
    "stellar_velocities = k_star*np.sin(stellar_phases*2.*np.pi)\n",
    "\n",
    "stellar_velocities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig = plt.figure()\n",
    "# ax = fig.add_subplot(111)\n",
    "\n",
    "# norm = visualization.ImageNormalize(ccfs_summed/normalization_ccfs, interval=visualization.ZScaleInterval())\n",
    "# imshow = ax.imshow(ccfs_summed_normalized, norm=norm, origin='lower', aspect='auto', cmap= matplotlib.cm.get_cmap('Blues'), interpolation='none', extent=[np.min(RV_table), np.max(RV_table), np.min(phases), np.max(phases)])\n",
    "# ax.set_xlabel('RV point')\n",
    "# ax.set_ylabel('Exposure')\n",
    "# fig.colorbar(imshow)\n",
    "# fig.tight_layout()\n",
    "\n",
    "# plt.plot(velocities_spun_during_transit, phases, 'r')\n",
    "# plt.plot(velocities_spun_during_transit-5, phases, 'r--')\n",
    "\n",
    "# # plt.xlim(-20, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to exclude the region where the planet + RML signal come from.\n",
    "\n",
    "approximate_systemic_velocity = -1.1\n",
    "vsini = 3.3\n",
    "max_velocity_during_transit = np.max(abs(velocities_spun_during_transit))\n",
    "\n",
    "plot_debug = False\n",
    "\n",
    "i1 = np.searchsorted(RV_table, approximate_systemic_velocity-vsini - 30.)\n",
    "i2 = np.searchsorted(RV_table, approximate_systemic_velocity-vsini+vsini+30.)\n",
    "\n",
    "ccfs_summed_normalized_continuum_normalized = []    \n",
    "\n",
    "\n",
    "for c in ccfs_summed_normalized:\n",
    "\n",
    "    fit = np.polyfit(np.concatenate((RV_table[:i1], RV_table[i2:])), np.concatenate((c[:i1], c[i2:])),2)\n",
    "    fit_fn = np.poly1d(fit) \n",
    "\n",
    "    if plot_debug:\n",
    "        import matplotlib.pyplot as plt\n",
    "        plt.plot(RV_table[i2:], c[i2:], 'k', RV_table[i2:], fit_fn(RV_table[i2:]), 'r')\n",
    "        plt.plot(np.concatenate((RV_table[:i1], RV_table[i2:])), np.concatenate((c[:i1], c[i2:])), 'k', RV_table, fit_fn(RV_table), 'r')\n",
    "        plt.show()\n",
    "\n",
    "    ccfs_summed_normalized_continuum_normalized.append(c/fit_fn(RV_table))\n",
    "\n",
    "ccfs_summed_normalized_continuum_normalized = np.array(ccfs_summed_normalized_continuum_normalized)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "norm = visualization.ImageNormalize(ccfs_summed_normalized_continuum_normalized, interval=visualization.ZScaleInterval())\n",
    "imshow = ax.imshow(ccfs_summed_normalized_continuum_normalized, norm=norm, origin='lower', aspect='auto', cmap= matplotlib.cm.get_cmap('Blues'), interpolation='none', extent=[np.min(RV_table), np.max(RV_table), np.min(phases), np.max(phases)])\n",
    "ax.set_xlabel('RV point')\n",
    "ax.set_ylabel('Exposure')\n",
    "fig.colorbar(imshow)\n",
    "fig.tight_layout()\n",
    "\n",
    "# plt.plot(velocities_spun_during_transit, phases, 'r')\n",
    "# plt.plot(velocities_spun_during_transit-5, phases, 'r--')\n",
    "# plt.xlim(-20, 10)"
   ]
  },
  {
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}