Virtual spectrometer SCS Grating.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6386344d-b7ac-440d-9926-f03af4ff9d6f",
   "metadata": {},
   "source": [
    "# Training the Virtual Spectrometer with grating and PES data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1711c3b9-5065-4a44-8b1b-a3e861b92bc5",
   "metadata": {},
   "source": [
    "The objective here is to use the grating monochromator 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 grating 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",
    "* 1 pulse trains in \"training\".\n",
    "\n",
    "The following software implements:\n",
    "1. retrieve data;\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": [
    "## Get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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": 13,
   "id": "fd8dacae-c22e-4c20-9df9-8720a2814320",
   "metadata": {},
   "outputs": [],
   "source": [
    "proposal = 900331\n",
    "runTrain = 69"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "25000b87-246d-467b-b770-8cde527faec4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "navitar: only 71.9% of trains (7176 out of 9974) contain data.\n",
      "energy: only 71.9% of trains (7176 out of 9974) contain data.\n"
     ]
    }
   ],
   "source": [
    "fields = [\n",
    "          'XTD10_SA3',             # XGM\n",
    "          *list(pes_map.values()), # PES\n",
    "          # calibrated grating\n",
    "          {'navitar': {'source': 'SA3_XTD10_SPECT/MDL/SPECTROMETER_SCS_NAVITAR:output',\n",
    "                       'key': 'data.intensityDistribution',\n",
    "                       'dim': ['gratingEnergy'],\n",
    "                      },\n",
    "          },\n",
    "          {'energy':\n",
    "          {'source': 'SA3_XTD10_SPECT/MDL/SPECTROMETER_SCS_NAVITAR:output',\n",
    "           'key': 'data.photonEnergy',\n",
    "           'dim': ['gratingEnergy'],\n",
    "          },\n",
    "          }\n",
    "         ]\n",
    "_, data_train = tb.load(proposal, runTrain, fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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": 27,
   "id": "b477bf49-f5ca-4df0-b6ed-a270ee35cd28",
   "metadata": {},
   "outputs": [],
   "source": [
    "channels = tuple(pes_data.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a843e981-e57e-4163-a4e0-310de7181aec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n",
       "<defs>\n",
       "<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n",
       "<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n",
       "<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
       "<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
       "</symbol>\n",
       "<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n",
       "<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n",
       "<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "</symbol>\n",
       "</defs>\n",