{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "733ed22f",
   "metadata": {},
   "source": [
    "# Tutorial for trXPD for the HEXTOF instrument at FLASH with background normalization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7eaa93d",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc871acf",
   "metadata": {},
   "source": [
    "### Import necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "368cf206",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "from sed import SedProcessor\n",
    "from sed.dataset import dataset\n",
    "import xarray as xr\n",
    "\n",
    "%matplotlib widget\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage import gaussian_filter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ed214a",
   "metadata": {},
   "source": [
    "### Get data paths\n",
    "\n",
    "If it is your beamtime, you can access both read the raw data and write to processed directory. For the public data, you can not write to processed directory.\n",
    "\n",
    "The paths are such that if you are on Maxwell, it uses those. Otherwise data is downloaded in current directory from Zenodo:\n",
    "https://zenodo.org/records/12609441"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "195737a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "beamtime_dir = \"/asap3/flash/gpfs/pg2/2023/data/11019101\" # on Maxwell\n",
    "if os.path.exists(beamtime_dir) and os.access(beamtime_dir, os.R_OK):\n",
    "    path = beamtime_dir + \"/raw/hdf/offline/fl1user3\"\n",
    "    buffer_path = beamtime_dir + \"/processed/tutorial/\"\n",
    "else:\n",
    "    # data_path can be defined and used to store the data in a specific location\n",
    "    dataset.get(\"W110\") # Put in Path to a storage of at least 10 GByte free space.\n",
    "    path = dataset.dir\n",
    "    buffer_path = path + \"/processed/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30effaac",
   "metadata": {},
   "source": [
    "### Config setup\n",
    "\n",
    "Here we get the path to the config file and setup the relevant directories. This can also be done directly in the config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ca5745a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pick the default configuration file for hextof@FLASH\n",
    "config_file = Path('../sed/config/flash_example_config.yaml')\n",
    "assert config_file.exists()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "347d338c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we setup a dictionary that will be used to override the path configuration\n",
    "config_override = {\n",
    "    \"core\": {\n",
    "        \"beamtime_id\": 11019101,\n",
    "        \"paths\": {\n",
    "            \"data_raw_dir\": path,\n",
    "            \"data_parquet_dir\": buffer_path\n",
    "        },\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee483947",
   "metadata": {},
   "source": [
    "### Prepare Energy Calibration\n",
    "Instead of making completely new energy calibration we can take existing values from the calibration made in the previous tutorial. This allows us to calibrate the conversion between the digital values of the dld and the energy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7417ab20",
   "metadata": {},
   "source": [
    " For this we need to add all those parameters as a dictionary and use them during creation of the processor object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdcbd374",
   "metadata": {},
   "outputs": [],
   "source": [
    "en_cal_config = {\n",
    "    'energy': {\n",
    "        'calibration': {\n",
    "            'E0': -54.971004271795664,\n",
    "            'creation_date': 1718801358.232129,\n",
    "            'd': 0.8096677238144319,\n",
    "            'energy_scale': 'kinetic',\n",
    "            't0': 4.0148196706891397e-07,\n",
    "            'calib_type': 'fit',\n",
    "            'fit_function': '(a0/(x0-a1))**2 + a2',\n",
    "            'coefficients': ([ 8.09667724e-01,  4.01481967e-07, -5.49710043e+01]),\n",
    "            'axis': 0},\n",
    "        'tof': None,\n",
    "        'offsets': {\n",
    "            'constant': -76.5,\n",
    "            'creation_date': 1718801360.817963,\n",
    "            'monochromatorPhotonEnergy': {'preserve_mean': True,'reduction': None,'weight': -1},\n",
    "            'sampleBias': {'preserve_mean': False, 'reduction': None, 'weight': 1},\n",
    "            'tofVoltage': {'preserve_mean': True, 'reduction': None, 'weight': -1}}}}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f070b3bb",
   "metadata": {},
   "source": [
    "## Read data\n",
    "Now we can use those parameters and load our trXPD data using additional config file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39b316d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_number = 44498\n",
    "sp_44498 = SedProcessor(runs=[run_number], folder_config=en_cal_config, config=config_override, system_config=config_file, verbose=True)\n",
    "sp_44498.add_jitter()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2afd119",
   "metadata": {},
   "source": [
    "We can inspect dataframe right after data readout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167aaf25",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_44498.dataframe.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d56f2f83",
   "metadata": {},
   "source": [
    "Now we will do energy calibration, add energy offset, jittering and dld sectors alignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1d12db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_44498.align_dld_sectors()\n",
    "sp_44498.append_energy_axis()\n",
    "sp_44498.add_energy_offset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7addaa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_44498.attributes.metadata['energy_calibration']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c97643f5",
   "metadata": {},
   "source": [
    "We can do the SASE jitter correction, using information from the bam column and do calibration of the pump-probe delay axis, we need to shift the delay stage values to center the pump-probe-time overlap time zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "514e9956",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_44498.add_delay_offset(\n",
    "    constant=-1448, # this is time zero position determined from side band fit\n",
    "    flip_delay_axis=True, # invert the direction of the delay axis\n",
    "    columns=['bam'], # use the bam to offset the values\n",
    "    weights=[-0.001], # bam is in fs, delay in ps\n",
    "    preserve_mean=True # preserve the mean of the delay axis to keep t0 position\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6551542d",
   "metadata": {},
   "source": [
    "### bin in the calibrated energy and corrected delay axis\n",
    "Visualize trXPS data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88055c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = ['energy', 'delayStage']\n",
    "ranges = [[-37.5,-27.5], [-1.5,1.5]]\n",
    "bins = [200,60]\n",
    "res_corr = sp_44498.compute(bins=bins, axes=axes, ranges=ranges, normalize_to_acquisition_time=\"delayStage\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b02865c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,2,figsize=(8,3), layout='constrained')\n",
    "fig.suptitle(f\"Run {run_number}: W 4f, side bands\")\n",
    "res_corr.plot(robust=True, ax=ax[0], cmap='terrain')\n",
    "ax[0].set_title('raw')\n",
    "bg = res_corr.sel(delayStage=slice(-1.3,-1.0)).mean('delayStage')\n",
    "(res_corr-bg).plot(robust=True, ax=ax[1])\n",
    "ax[1].set_title('difference')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43e84eff",
   "metadata": {},
   "source": [
    "## XPD from W4f core level\n",
    "\n",
    "Now we can bin not only in energy but also in both momentum directions to get XPD patterns of different core level line of tungsten."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b768cfc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = ['energy', 'dldPosX', 'dldPosY']\n",
    "ranges = [[-38,-28], [420,900], [420,900]]\n",
    "bins = [100,240,240]\n",
    "res_kx_ky = sp_44498.compute(bins=bins, axes=axes, ranges=ranges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73aed405",
   "metadata": {},
   "outputs": [],
   "source": [
    "## EDC and integration region for XPD\n",
    "plt.figure()\n",
    "res_kx_ky.mean(('dldPosX', 'dldPosY')).plot()\n",
    "plt.vlines([-30.3,-29.9], 0, 2.4, color='r', linestyles='dashed')\n",
    "plt.vlines([-31.4,-31.2], 0, 2.4, color='orange', linestyles='dashed')\n",
    "plt.vlines([-33.6,-33.4], 0, 2.4, color='g', linestyles='dashed')\n",
    "plt.vlines([-37.0,-36.0], 0, 2.4, color='b', linestyles='dashed')\n",
    "plt.title('EDC and integration regions for XPD')\n",
    "plt.show()\n",
    "\n",
    "## XPD plots\n",
    "fig,ax = plt.subplots(2,2,figsize=(9,7), layout='constrained')\n",
    "res_kx_ky.sel(energy=slice(-30.3,-29.9)).mean('energy').plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "ax[0,0].set_title(\"XPD of $1^{st}$ order sidebands\")\n",
    "res_kx_ky.sel(energy=slice(-31.4,-31.2)).mean('energy').plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "ax[0,1].set_title(\"XPD of W4f 7/2 peak\")\n",
    "res_kx_ky.sel(energy=slice(-33.6,-33.4)).mean('energy').plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "ax[1,0].set_title(\"XPD of W4f 5/2 peak\")\n",
    "res_kx_ky.sel(energy=slice(-37.0,-36.0)).mean('energy').plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "ax[1,1].set_title(\"XPD of W5p 3/2 peak\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c30cdbd0",
   "metadata": {},
   "source": [
    "As we can see there is some structure visible, but it looks very similar to each other.\n",
    "We probably have to do some normalization to remove the detector structure/artefacts.\n",
    "The best option is to divide by a flat-field image. The flat-field image can be obtained from a sample that shows no structure under identical measurement conditions.\n",
    "Unfortunately, we don't have such a flat-field image.\n",
    "\n",
    "In this case, we can make a flat-field image from the actual dataset using several different approaches.\n",
    "\n",
    "As a first option, we can integrate in energy over the whole region and use this image as a background.\n",
    "Additionally, we introduce the Gaussian Blur for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e4a7faf",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Background image\n",
    "bgd = res_kx_ky.mean(('energy'))\n",
    "\n",
    "## Apply Gaussian Blur to background image\n",
    "bgd_blur = xr.apply_ufunc(gaussian_filter, bgd, 15)\n",
    "\n",
    "fig,ax = plt.subplots(1,2,figsize=(9,4), layout='constrained')\n",
    "bgd.plot(robust=True, cmap='terrain', ax=ax[0])\n",
    "ax[0].set_title('Background image')\n",
    "bgd_blur.plot(cmap='terrain', ax=ax[1])\n",
    "ax[1].set_title('Gaussian Blur of background image')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa01c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## XPD normalized by background image\n",
    "fig,ax = plt.subplots(2,2,figsize=(9,7), layout='constrained')\n",
    "(res_kx_ky/bgd).sel(energy=slice(-30.3,-29.9)).mean('energy').plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "(res_kx_ky/bgd).sel(energy=slice(-31.4,-31.2)).mean('energy').plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "(res_kx_ky/bgd).sel(energy=slice(-33.6,-33.4)).mean('energy').plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "(res_kx_ky/bgd).sel(energy=slice(-37.0,-36.0)).mean('energy').plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: XPD patterns after background normalization',fontsize='18')\n",
    "\n",
    "## XPD normalized by Gaussian-blurred background image\n",
    "fig,ax = plt.subplots(2,2,figsize=(9,7), layout='constrained')\n",
    "(res_kx_ky/bgd_blur).sel(energy=slice(-30.3,-29.9)).mean('energy').plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "(res_kx_ky/bgd_blur).sel(energy=slice(-31.4,-31.2)).mean('energy').plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "(res_kx_ky/bgd_blur).sel(energy=slice(-33.6,-33.4)).mean('energy').plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "(res_kx_ky/bgd_blur).sel(energy=slice(-37.0,-36.0)).mean('energy').plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: XPD patterns after Gaussian-blurred background normalization',fontsize='18')\n",
    "\n",
    "## XPD normalized by Gaussian-blurred background image and blurred to improve contrast\n",
    "fig,ax = plt.subplots(2,2,figsize=(9,7), layout='constrained')\n",
    "(xr.apply_ufunc(gaussian_filter, res_kx_ky/bgd_blur, 1)).sel(energy=slice(-30.3,-29.9)).mean('energy').plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "(xr.apply_ufunc(gaussian_filter, res_kx_ky/bgd_blur, 1)).sel(energy=slice(-31.4,-31.2)).mean('energy').plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "(xr.apply_ufunc(gaussian_filter, res_kx_ky/bgd_blur, 1)).sel(energy=slice(-33.6,-33.4)).mean('energy').plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "(xr.apply_ufunc(gaussian_filter, res_kx_ky/bgd_blur, 1)).sel(energy=slice(-37.0,-36.0)).mean('energy').plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: resulting Gaussian-blurred XPD patterns',fontsize='18')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f415a21",
   "metadata": {},
   "source": [
    "Sometimes, after this division, you may not be happy with intensity distribution. Thus, other option for background correction is to duplicate the XPD pattern, apply large Gaussian blurring that eliminates the fine structures in the XPD pattern. Then divide the XPD pattern by its blurred version. This process sometimes enhances the visibility of the fine structures a lot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e5607ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "## XPD normalized by Gaussian-blurred background image\n",
    "\n",
    "### Define integration regions for XPD\n",
    "SB = res_kx_ky.sel(energy=slice(-30.3,-29.9)).mean('energy')\n",
    "W_4f_7 = res_kx_ky.sel(energy=slice(-31.4,-31.2)).mean('energy')\n",
    "W_4f_5 = res_kx_ky.sel(energy=slice(-33.6,-33.4)).mean('energy')\n",
    "W_5p = res_kx_ky.sel(energy=slice(-37.0,-36.0)).mean('energy')\n",
    "\n",
    "### Make corresponding Gaussian Blur background\n",
    "SB_blur = xr.apply_ufunc(gaussian_filter, SB, 15)\n",
    "W_4f_7_blur = xr.apply_ufunc(gaussian_filter, W_4f_7, 15)\n",
    "W_4f_5_blur = xr.apply_ufunc(gaussian_filter, W_4f_5, 15)\n",
    "W_5p_blur = xr.apply_ufunc(gaussian_filter, W_5p, 15)\n",
    "\n",
    "### Visualize results\n",
    "fig,ax = plt.subplots(2,2,figsize=(9,7), layout='constrained')\n",
    "(SB/SB_blur).plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "(W_4f_7/W_4f_7_blur).plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "(W_4f_5/W_4f_5_blur).plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "(W_5p/W_5p_blur).plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: XPD patterns after Gaussian Blur normalization',fontsize='18')\n",
    "\n",
    "### Apply Gaussian Blur to resulted images to improve contrast\n",
    "SB_norm = xr.apply_ufunc(gaussian_filter, SB/SB_blur, 1)\n",
    "W_4f_7_norm = xr.apply_ufunc(gaussian_filter, W_4f_7/W_4f_7_blur, 1)\n",
    "W_4f_5_norm = xr.apply_ufunc(gaussian_filter, W_4f_5/W_4f_5_blur, 1)\n",
    "W_5p_norm = xr.apply_ufunc(gaussian_filter, W_5p/W_5p_blur, 1)\n",
    "\n",
    "### Visualize results\n",
    "fig,ax = plt.subplots(2,2,figsize=(9,7), layout='constrained')\n",
    "SB_norm.plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "W_4f_7_norm.plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "W_4f_5_norm.plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "W_5p_norm.plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: XPD patterns after Gauss Blur normalization',fontsize='18')                   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94d41b4f",
   "metadata": {},
   "source": [
    "Third option for background normalization is to use the simultaneously acquired pre-core level region.\n",
    "As an example for W4f 7/2 peak, we define a region on the high energy side of it and integrate in energy to use as a background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c8513b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Define peak and background region on the high energy side of the peak\n",
    "W_4f_7 = res_kx_ky.sel(energy=slice(-31.4,-31.2)).mean('energy')\n",
    "W_4f_7_bgd = res_kx_ky.sel(energy=slice(-32.0,-31.8)).mean('energy')\n",
    "\n",
    "### Make normalization by background, add Gaussian Blur to the resulting image\n",
    "W_4f_7_nrm1 = W_4f_7/(W_4f_7_bgd+W_4f_7_bgd.max()*0.00001)\n",
    "W_4f_7_nrm1_blur = xr.apply_ufunc(gaussian_filter, W_4f_7_nrm1, 1)\n",
    "\n",
    "### Add Gaussian Blur to the background image, normalize by it and add Gaussian Blur to the resulting image\n",
    "W_4f_7_bgd_blur = xr.apply_ufunc(gaussian_filter, W_4f_7_bgd, 15)\n",
    "W_4f_7_nrm2 = W_4f_7/W_4f_7_bgd_blur\n",
    "W_4f_7_nrm2_blur = xr.apply_ufunc(gaussian_filter, W_4f_7_nrm2, 1)\n",
    "\n",
    "### Visualize all steps\n",
    "fig,ax = plt.subplots(4,2,figsize=(9,10), layout='constrained')\n",
    "W_4f_7.plot(robust=True, ax=ax[0,0], cmap='terrain')\n",
    "W_4f_7_bgd.plot(robust=True, ax=ax[0,1], cmap='terrain')\n",
    "W_4f_7_nrm1.plot(robust=True, ax=ax[1,0], cmap='terrain')\n",
    "W_4f_7_nrm1_blur.plot(robust=True, ax=ax[1,1], cmap='terrain')\n",
    "W_4f_7_bgd_blur.plot(robust=True, ax=ax[2,0], cmap='terrain')\n",
    "W_4f_7_nrm2.plot(robust=True, ax=ax[2,1], cmap='terrain')\n",
    "W_4f_7_nrm2_blur.plot(robust=True, ax=ax[3,0], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: XPD patterns of W4f7/2 with pre-core level normalization',fontsize='18')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "807883fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(1,3,figsize=(9,3), layout='constrained')\n",
    "(xr.apply_ufunc(gaussian_filter, res_kx_ky/bgd_blur, 1)).sel(energy=slice(-31.4,-31.2)).mean('energy').plot(robust=True, ax=ax[0], cmap='terrain')\n",
    "W_4f_7_norm.plot(robust=True, ax=ax[1], cmap='terrain')\n",
    "W_4f_7_nrm2_blur.plot(robust=True, ax=ax[2], cmap='terrain')\n",
    "fig.suptitle(f'Run {run_number}: comparison of different normalizations\\nof XPD pattern for W4f 7/2 peak with Gaussian Blur',fontsize='18')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7654516c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}