{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8ad4167a-e4e7-498d-909a-c04da9f177ed",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Distortion correction with orthorhombic symmetry\n",
    "This example showcases how to use the distortion correction workflow with landmarks that are not at symmetry-equivalent positions, such as for orthorhombic systems with different in-plane axis parameters."
   ]
  },
  {
   "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",
    "\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\n",
    "For this example, we use the example data from WSe2. Even though the system is hexagonal, we will use it for demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f46d54",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.get(\"WSe2\") # Put in Path to a storage of at least 20 GByte free space.\n",
    "data_path = dataset.dir # This is the path to the data\n",
    "scandir, _ = dataset.subdirs # scandir contains the data, _ contains the calibration files"
   ]
  },
  {
   "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, verbose=True)\n",
    "sp.add_jitter()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f10ec6c",
   "metadata": {},
   "source": [
    "Get slice for momentum calibration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ff23621",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.bin_and_load_momentum_calibration(df_partitions=100, plane=203, width=10, apply=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fee3ca76",
   "metadata": {},
   "source": [
    "## Feature definition:\n",
    "We will describe the symmetry of the system with a 4-fold symmetry, and select two K points and two M points as symmetry points (as well as the Gamma point)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd9666c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = np.array([[252., 355.], [361., 251.], [250., 144.], [156., 247.], [254., 247.]])\n",
    "sp.define_features(features=features, rotation_symmetry=4, include_center=True, apply=True)\n",
    "# Manual selection: Use a GUI tool to select peaks:\n",
    "# sp.define_features(rotation_symmetry=4, include_center=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f7519ff8",
   "metadata": {},
   "source": [
    "## Spline-warp generation: \n",
    "For the spline-warp generation, we need to tell the algorithm the difference in length of Gamma-K and Gamma-M. This we can do using the ascale parameter, which can either be a single number (the ratio), or a list of length ``rotation_symmetry`` defining the relative length of the respective vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d27cd7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma_m = np.pi/3.28\n",
    "gamma_k = 2/np.sqrt(3)*np.pi/3.28\n",
    "# Option 1: Ratio of the two distances:\n",
    "#sp.generate_splinewarp(include_center=True, ascale=gamma_k/gamma_m)\n",
    "# Option 2: List of distances:\n",
    "sp.generate_splinewarp(include_center=True, ascale=[gamma_m, gamma_k, gamma_m, gamma_k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62abfa41",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.pose_adjustment(xtrans=4, ytrans=7, angle=1, apply=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845f002d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.apply_momentum_correction()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d9810488",
   "metadata": {},
   "source": [
    "## Momentum calibration with orthorhombic axes\n",
    "For the momentum calibration using symmetry points with non-equal distances, the option ``equiscale`` can be used:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a358f07d",
   "metadata": {},
   "outputs": [],
   "source": [
    "point_a = [256, 155]\n",
    "point_b = [370, 256]\n",
    "sp.calibrate_momentum_axes(point_a=point_a, point_b=point_b, k_coord_a=[0, gamma_m], k_coord_b=[gamma_k, 0], equiscale=False, apply=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ae5066",
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.apply_momentum_calibration()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6902fd56-1456-4da6-83a4-0f3f6b831eb6",
   "metadata": {},
   "source": [
    "## Bin the top of the valence band"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7601cd7-cd51-40a9-8fc7-8b7d32ff15d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = ['kx', 'ky']\n",
    "bins = [100, 100]\n",
    "ranges = [[-2, 2], [-2, 2]]\n",
    "res = sp.compute(bins=bins, axes=axes, ranges=ranges, filter=[{\"col\":\"t\", \"lower_bound\": 66100, \"upper_bound\": 66300}])\n",
    "plt.figure()\n",
    "res.T.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f73f099",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "728003ee06929e5fa5ff815d1b96bf487266025e4b7440930c6bf4536d02d243"
  },
  "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
}