Virtual spectrometer SCS Viking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6386344d-b7ac-440d-9926-f03af4ff9d6f",
   "metadata": {},
   "source": [
    "# Training the Virtual Spectrometer with Viking and PES data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1711c3b9-5065-4a44-8b1b-a3e861b92bc5",
   "metadata": {},
   "source": [
    "The objective here is to use the Viking detector to train the Virtual Spectrometer. This means that we will fit (\"train\") a model, which maps the PES measurements with the Viking measurements and use their correlation to interpolate in cases where the Viking is not available.\n",
    "\n",
    "The following conditions must be satisfied for this to be possible:\n",
    "* The PES settings are the same in the \"training\" run and interesting run.\n",
    "* The photon energies of the beam in \"training\" and in the interesting run are similar.\n",
    "* The beam intensities are similar.\n",
    "* The sample between PES and Viking is transparent.\n",
    "* 1 pulse trains in \"training\".\n",
    "\n",
    "The following software implements:\n",
    "1. retrieve data and calibrate Viking using the SCS toolbox;\n",
    "2. the Virtual Spectrometer training excluding the last 10 trains avalable so that we can use them for validation;\n",
    "3. the Virtual Spectrometer resolution function plotting;\n",
    "4. comparison of the Virtual spectrometer results in a selected set in which the Viking data was available.\n",
    "\n",
    "Finally, the model is applied in data without the grating. This last part may be applied independently from the rest if the modal has been written in a `joblib` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4a627555-522a-4c9d-b6b2-6ff77148eaab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "# replace this\n",
    "sys.path.append('/home/danilo/scratch/karabo/devices/pes_to_spec')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "78bbc433-ac5e-44c3-8740-3e93800c4532",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cupy is not installed in this environment, no access to the GPU\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import dask.array as da\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from pes_to_spec.model import Model\n",
    "\n",
    "import toolbox_scs as tb\n",
    "from euxfel_bunch_pattern import indices_at_sase\n",
    "\n",
    "from scipy.signal import fftconvolve"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7609899-5bc0-4211-ae97-010b3edcf676",
   "metadata": {},
   "source": [
    "# Training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "95da5231-e454-4f7f-a1ce-eef7e52fe457",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pes channel names to be used for reference later\n",
    "pes_map = dict(channel_1_A=\"PES_S_raw\",\n",
    "                channel_1_B=\"PES_SSW_raw\",\n",
    "                channel_1_C=\"PES_SW_raw\",\n",
    "                channel_1_D=\"PES_WSW_raw\",\n",
    "                channel_2_A=\"PES_W_raw\",\n",
    "                channel_2_B=\"PES_WNW_raw\",\n",
    "                channel_2_C=\"PES_NW_raw\",\n",
    "                channel_2_D=\"PES_NNW_raw\",\n",
    "                channel_3_A=\"PES_E_raw\",\n",
    "                channel_3_B=\"PES_ESE_raw\",\n",
    "                channel_3_C=\"PES_SE_raw\",\n",
    "                channel_3_D=\"PES_SSE_raw\",\n",
    "                channel_4_A=\"PES_N_raw\",\n",
    "                channel_4_B=\"PES_NNE_raw\",\n",
    "                channel_4_C=\"PES_NE_raw\",\n",
    "                channel_4_D=\"PES_ENE_raw\",\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48bb4c8c-04ad-44d5-b123-643ce3253ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "proposal = 2953\n",
    "runTrain = 322  # run containing the data without sample\n",
    "darkNB = 375  # dark run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0a467b2f-5f99-4ed8-bb1d-cb429454d3ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "newton: only 50.0% of trains (629 out of 1259) contain data.\n"
     ]
    }
   ],
   "source": [
    "v = tb.Viking(proposal)\n",
    "fields = ['XTD10_SA3',\n",
    "          *list(pes_map.values()) # add PES\n",
    "         ]\n",
    "v.FIELDS += fields\n",
    "v.X_RANGE = slice(0, 1500) # define the dispersive axis range of interest (in pixels)\n",
    "v.Y_RANGE = slice(29, 82) # define the non-dispersive axis range of interest (in pixels)\n",
    "v.ENERGY_CALIB = [1.47802667e-06, 2.30600328e-02, 5.15884589e+02] # energy calibration, see further below\n",
    "v.BL_POLY_DEG = 1 # define the polynomial degree for baseline subtraction\n",
    "v.BL_SIGNAL_RANGE = [500, 545] # define the range containing the signal, to be excluded for baseline subtraction\n",
    "\n",
    "v.load_dark(darkNB)  # load a dark image (averaged over the dark run number)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f6124d9-8c1b-44f8-a078-07475a9674fc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_train = v.from_run(runTrain)  # load refNB. The `newton` variable contains the CCD images.\n",
    "v.integrate(data_train)  # integrate over the non-dispersive dimension \n",
    "v.removePolyBaseline(data_train)  # remove baseline\n",
    "data_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "294b5f3a-1d59-444e-80ab-4834d26d62dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform PES data into the format expected\n",
    "pes_data = {k: da.from_array(data_train[item].to_numpy())\n",
    "            for k, item in pes_map.items() if item in data_train}\n",
    "xgm = data_train.XTD10_SA3.isel(sa3_pId=0).to_numpy()[:, np.newaxis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b477bf49-f5ca-4df0-b6ed-a270ee35cd28",
   "metadata": {},
   "outputs": [],
   "source": [
    "channels = tuple(pes_data.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f154e38-d208-477e-9d9c-ef2a632514c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy = data_train.newt_x.to_numpy()"
   ]
  },
  {
   "cell_type": "code",