{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8ad4167a-e4e7-498d-909a-c04da9f177ed",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Binning of temperature-dependent ARPES data using time-stamped external temperature data\n",
    "In this example, we pull some temperature-dependent ARPES data from Zenodo, which was recorded as a continuous temperature ramp. We then add the respective temperature information from the respective timestamp/temperature values to the dataframe, and bin the data as function of temperature\n",
    "For performance reasons, best store the data on a locally attached storage (no network drive). This can also be achieved transparently using the included MirrorUtil class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb045e17-fa89-4c11-9d51-7f06e80d96d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import glob\n",
    "\n",
    "import sed\n",
    "from sed.dataset import dataset\n",
    "\n",
    "%matplotlib widget"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "42a6afaa-17dd-4637-ba75-a28c4ead1adf",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f46d54",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.get(\"TaS2\") # Put in Path to a storage of at least 20 GByte free space.\n",
    "data_path = dataset.dir\n",
    "scandir, caldir = dataset.subdirs # scandir contains the data, caldir contains the calibration files\n",
    "\n",
    "# correct timestamps if not correct timezone set\n",
    "tzoffset = os.path.getmtime(scandir + '/Scan0121_1.h5') - 1594998158.0\n",
    "if tzoffset:\n",
    "    for file in glob.glob(scandir +'/*.h5'):\n",
    "        os.utime(file, (os.path.getmtime(file)-tzoffset, os.path.getmtime(file)-tzoffset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f82054",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create sed processor using the config file with time-stamps:\n",
    "sp = sed.SedProcessor(folder=scandir, user_config=\"../sed/config/mpes_example_config.yaml\", time_stamps=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85ac3c83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply jittering to X, Y, t, ADC columns.\n",
    "sp.add_jitter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76bf8aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.bin_and_load_momentum_calibration(df_partitions=10, plane=33, width=3, apply=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f99c6b3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = np.array([[337., 242.], [289., 327.], [187., 344.], [137., 258.], [189., 161.], [289., 158.], [236.0, 250.0]])\n",
    "sp.define_features(features=features, rotation_symmetry=6, include_center=True, apply=True)\n",
    "sp.generate_splinewarp(include_center=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62abfa41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adjust pose alignment, using stored distortion correction\n",
    "sp.pose_adjustment(xtrans=15, ytrans=8, angle=-5, apply=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845f002d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply stored momentum correction\n",
    "sp.apply_momentum_correction()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ae5066",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply stored config momentum calibration\n",
    "sp.apply_momentum_calibration()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb1e2bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply stored config energy correction\n",
    "sp.apply_energy_correction()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de946733",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load energy calibration EDCs\n",
    "scans = np.arange(127,136)\n",
    "voltages = np.arange(22,13,-1)\n",
    "files = [caldir + r'/Scan' + str(num).zfill(4) + '_1.h5' for num in scans]\n",
    "sp.load_bias_series(data_files=files, normalize=True, biases=voltages, ranges=[(64000, 76000)])\n",
    "rg = (65500, 66000)\n",
    "sp.find_bias_peaks(ranges=rg, ref_id=5, infer_others=True, apply=True)\n",
    "sp.calibrate_energy_axis(ref_energy=-0.5, ref_id=4, energy_scale=\"kinetic\", method=\"lmfit\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c470ffd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply stored config energy calibration\n",
    "sp.append_energy_axis()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0943d349",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add time-stamped temperature data\n",
    "# either, directly retrieve data from EPICS archiver instance (within FHI network),\n",
    "#sp.add_time_stamped_data(dest_column=\"T_B\", archiver_channel=\"trARPES:Carving:TEMP-B\")\n",
    "# or use externally provided timestamp/data pairs\n",
    "import h5py\n",
    "with h5py.File(f\"{data_path}/temperature_data.h5\", \"r\") as file:\n",
    "    data = file[\"temperatures\"][()]\n",
    "    time_stamps = file[\"timestamps\"][()]\n",
    "sp.add_time_stamped_data(dest_column=\"sample_temperature\", time_stamps=time_stamps, data=data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c330da64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect calibrated event histogram\n",
    "axes = ['kx', 'ky', 'energy', 'sample_temperature']\n",
    "ranges = [[-3, 3], [-3, 3], [-6, 2], [10, 300]]\n",
    "sp.view_event_histogram(dfpid=80, axes=axes, ranges=ranges)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6902fd56-1456-4da6-83a4-0f3f6b831eb6",
   "metadata": {},
   "source": [
    "## Define the binning ranges and compute calibrated data volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7601cd7-cd51-40a9-8fc7-8b7d32ff15d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = ['kx', 'ky', 'energy', 'sample_temperature']\n",
    "bins = [100, 100, 300, 100]\n",
    "ranges = [[-2, 2], [-2, 2], [-6, 2], [20, 270]]\n",
    "res = sp.compute(bins=bins, axes=axes, ranges=ranges, normalize_to_acquisition_time=\"sample_temperature\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "523794dc",
   "metadata": {},
   "source": [
    "## Some visualization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99d7d136-b677-4c16-bc8f-31ba8216579c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(4, 1, figsize=(4, 12), constrained_layout=True)\n",
    "res.loc[{'energy':slice(-.1, 0)}].sum(axis=(2,3)).T.plot(ax=axs[0])\n",
    "res.loc[{'kx':slice(-.2, .2)}].sum(axis=(0,3)).T.plot(ax=axs[1])\n",
    "res.loc[{'ky':slice(-.2, .2)}].sum(axis=(1,3)).T.plot(ax=axs[2])\n",
    "res.loc[{'kx':slice(-.2, .2), 'ky':slice(-.2, .2), 'energy':slice(-2, 0.2)}].sum(axis=(0,1)).plot(ax=axs[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596a3217",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inspect effect of histogram normalization\n",
    "fig, ax = plt.subplots(1,1)\n",
    "(sp._normalization_histogram/sp._normalization_histogram.sum()).plot(ax=ax)\n",
    "(sp._binned.sum(axis=(0,1,2))/sp._binned.sum(axis=(0,1,2,3))).plot(ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05488944",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remaining fluctuations are an effect of the varying count rate throughout the scan\n",
    "plt.figure()\n",
    "rate, secs = sp.loader.get_count_rate()\n",
    "plt.plot(secs, rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4697a5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalize for intensity around the Gamma point\n",
    "res_norm = res.copy()\n",
    "res_norm = res_norm/res_norm.loc[{'kx':slice(-.3, .3), 'ky':slice(-.3, .3)}].sum(axis=(0,1,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "124b3095",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(4, 1, figsize=(4, 12), constrained_layout=True)\n",
    "res_norm.loc[{'energy':slice(-.1, 0)}].sum(axis=(2,3)).T.plot(ax=axs[0])\n",
    "res_norm.loc[{'kx':slice(-.2, .2)}].sum(axis=(0,3)).T.plot(ax=axs[1])\n",
    "res_norm.loc[{'ky':slice(-.2, .2)}].sum(axis=(1,3)).T.plot(ax=axs[2])\n",
    "res_norm.loc[{'kx':slice(-.2, .2), 'ky':slice(-.2, .2), 'energy':slice(-2, 0.5)}].sum(axis=(0,1)).plot(ax=axs[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a07437",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lower Hubbard band intensity versus temperature\n",
    "plt.figure()\n",
    "res_norm.loc[{'kx':slice(-.2, .2), 'ky':slice(-.2, .2), 'energy':slice(-.6, 0.1)}].sum(axis=(0,1,2)).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15309e42",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}