{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Performing Fermi Edge Corrections\n",
    "\n",
    "The most straightforward way to correct the Fermi edge, either due to monochromator miscalibration in a photon energy scan or due to using a straight slit on a hemispherical analyzer is just to broadcast an edge.\n",
    "\n",
    "In the case of correcting for the slit shape, it may be helpful to further fit a model for the edge shape, like a quadratic, so that a smooth correction is applied across the detector.\n",
    "\n",
    "## Correcting monochromator miscalibration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lmfit.models import ConstantModel\n",
    "import arpes\n",
    "from arpes.fits import AffineBroadenedFD\n",
    "from arpes.io import example_data\n",
    "\n",
    "photon_energy = example_data.photon_energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_data = photon_energy.sel(phi=slice(-0.28, -0.15), eV=slice(-0.1, 0.1)).sum(\"phi\").spectrum\n",
    "edge_data.T.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "model = AffineBroadenedFD() + ConstantModel()\n",
    "# fermi_param = AffineBroadenedFD().guess(data=edge_data.sel(hv=70, method=\"nearest\").values, x=edge_data.coords[\"eV\"].values)\n",
    "\n",
    "param = AffineBroadenedFD().make_params(\n",
    "    center=0.0, width=0.02, sigma=0.02, lin_slope=0, const_bkg=15000\n",
    ")\n",
    "fit_results = edge_data.S.modelfit(\"eV\", model, params=param)\n",
    "# fit_results = broadcast_model(AffineBroadenedFD, edge_data, \"hv\")\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "edge_data.T.plot(ax=ax)\n",
    "# ax.scatter(*fit_results.results.F.p(\"center\").G.to_arrays(), color=\"red\")\n",
    "ax.scatter(\n",
    "    fit_results.modelfit_results.F.p(\"center\").coords[\"hv\"],\n",
    "    fit_results.modelfit_results.F.p(\"center\").values,\n",
    "    color=\"red\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Now, we can perform the shift to correct the edge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "corrected_photon_energy = photon_energy.spectrum.G.shift_by(\n",
    "    fit_results.modelfit_results.F.p(\"center\"), shift_axis=\"eV\", shift_coords=True\n",
    ")  # Note that G.shift_by is applied to \"spectrum\" (i.e. xr.DataArray)\n",
    "\n",
    "corrected_edge_data = corrected_photon_energy.sel(phi=slice(-0.28, -0.15),\n",
    "                                                  eV=slice(-0.1, 0.1)).sum(\n",
    "    \"phi\"\n",
    ")\n",
    "\n",
    "model = AffineBroadenedFD() + ConstantModel()\n",
    "param = AffineBroadenedFD().make_params(\n",
    "    center=-0.025, width=0.02, sigma=0.02, lin_slope=0, const_bkg=15000,\n",
    ")\n",
    "results_check = corrected_edge_data.S.modelfit(\"eV\", model, params=param)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "corrected_edge_data.T.plot(ax=ax)\n",
    "# ax.scatter(*results_check.results.F.p(\"center\").G.to_arrays(), color=\"red\")\n",
    "ax.scatter(\n",
    "    results_check.modelfit_results.F.p(\"center\").coords[\"hv\"],\n",
    "    results_check.modelfit_results.F.p(\"center\").values,\n",
    "    color=\"red\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "As we can see, the edge is now uniform across photon energy and correctly zero referenced.\n",
    "\n",
    "A caveat to be aware of when shifting data is whether to make data adjustments only or to use coordinate adjustments as well. Coordinate adjustments (above: `shift_coords=True`) are useful when the shift is very large. If the coords are not allowed to compensate for some of the shift in that context, large portions of data will be shifted out of the array extent and be replaced by `np.nan` or `0`.\n",
    "\n",
    "However, if coordinates no longer agree between two pieces of data, we will not be able to perform array operations involving both of them, because of their incompatible coordinates. \n",
    "\n",
    "The correct behavior is context dependent and requires you to consider what analysis you are trying to do."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Correcting the curved Fermi edges\n",
    "\n",
    "Let's now turn to an example addressing the uneven energy calibration arising from the use of a straight slit in ARPES data. \n",
    "\n",
    "We will shortly turn to the question of momentum conversion, but we will want to have this issue corrected before converting. The correction we need to apply is a function of the detector angle `phi`, so it will not be a constant function of any momentum coordinate, in general.\n",
    "\n",
    "First, let's assess the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from arpes.io import example_data\n",
    "\n",
    "cut = example_data.map.sum(\"theta\").spectrum\n",
    "cut = cut.sel(eV=slice(-0.2, 0.1), phi=slice(-0.25, 0.3))\n",
    "cut.S.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Now, let's fit edges to this data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from lmfit.models import QuadraticModel, ConstantModel\n",
    "from arpes.fits.fit_models import AffineBroadenedFD\n",
    "\n",
    "model = AffineBroadenedFD() + ConstantModel()\n",
    "params = AffineBroadenedFD().make_params(\n",
    "    center=0,\n",
    "    width=0.005,\n",
    "    sigma=0.02,\n",
    "    const_bkg=200000,\n",
    "    lin_slope=0,\n",
    ")\n",
    "\n",
    "\n",
    "fit_results = cut.S.modelfit(\"eV\", model, params=params)\n",
    "\n",
    "fit_results.modelfit_results.F.plot_param(\"center\")\n",
    "plt.gca().set_ylim([-0.05, 0.05])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "We could either refine these fits a little by setting some constraints, or we can make a smooth correction by fitting a quadratic to these edge locations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "quad_mod = (\n",
    "    fit_results.modelfit_results.F.p(\"center\")\n",
    "    .S.modelfit(\"phi\", QuadraticModel())\n",
    "    .modelfit_results.item()\n",
    ")\n",
    "\n",
    "# quad_mod = QuadraticModel().guess_fit(fit_results.results.F.p(\"center\"))\n",
    "quad_mod_plot = quad_mod.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "Now, we can shift the data exactly like we did before. In order to get the shift amount to apply at each `phi`, we just evaluate our quadratic at these `phi`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "fmap = example_data.map.spectrum\n",
    "edge = quad_mod.eval(x=fmap.phi)\n",
    "\n",
    "corrected_map = fmap.G.shift_by(edge, shift_axis=\"eV\", by_axis=\"phi\")\n",
    "extend_corrected_map = fmap.G.shift_by(edge,\n",
    "                                       shift_axis=\"eV\",\n",
    "                                       by_axis=\"phi\",\n",
    "                                       extend_coords=True)\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(10, 5))\n",
    "corrected_map.isel(theta=10).S.plot(ax=axes[0])\n",
    "fmap.isel(theta=10).S.plot(ax=axes[1])\n",
    "extend_corrected_map.isel(theta=10).S.plot(ax=axes[2])\n",
    "axes[0].set_title(\"Corrected data\")\n",
    "axes[1].set_title(\"Original data\")\n",
    "axes[2].set_title(\"Corrected data, coordinate extended\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "Much better. With this in order, we can now consider momentum conversion."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "1. Correct the Fermi edge on the Bi2Se3 cut data `example_data.cut`. How does the presence of the dispersing surface state affect the range you should use?\n",
    "2. Perform a random shift of a Fermi edge using `.G.shift_by`. Then, try to correct it. How close is your recovered correction to the original shift?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}