{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# RHEED Image Stacks\n",
    "\n",
    "This notebook demonstrates how to load and process a sequence of RHEED images acquired during molecular beam epitaxy (MBE)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "import xrheed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Loading a Series of RHEED Images\n",
    "In this example, we work with a sequence of BMP images located in the `deposition` directory. To load the data, we first need to prepare the following:\n",
    "\n",
    "- A list of file paths (preferably as `pathlib.Path` objects)\n",
    "- A name for the stack dimension\n",
    "- A list or vector of precise coverage values corresponding to each image in the sequence\n",
    "\n",
    "> **Note:** The allowed stack dimension options are defined in `constants.CANONICAL_STACK_DIMS`. \n",
    "\n",
    "> Valid values include: `alpha`, `beta`, `coverage`, `time`, `temperature`, `current`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dir = Path(\"example_data/deposition\")\n",
    "\n",
    "# Get all BMP files and sort them numerically by the number in the filename\n",
    "image_paths = sorted(\n",
    "    image_dir.glob(\"*.BMP\"),\n",
    "    key=lambda p: int(p.stem.split(\"_\")[-1])\n",
    ")\n",
    "\n",
    "# Typically the list of paths could be generated in a simpler way:\n",
    "# image_paths = list(image_dir.glob(\"*.BMP\"))\n",
    "\n",
    "n_stack = len(image_paths)\n",
    "\n",
    "stack_dim = 'coverage'\n",
    "deposition_rate = 2.54 # Hz per image\n",
    "\n",
    "stack_coordinate = np.linspace(0.0, n_stack * deposition_rate, n_stack)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "The data can be loaded using the same `xrheed.load_data()` function by providing a list of paths instead of a single path. In addition, the `stack_dim` and `stack_coords` parameters must be specified.\n",
    "\n",
    "To correctly interpret the data format, the `plugin` argument must also be provided. This should point to a specific plugin capable of handling the structure and metadata of the dataset being loaded.\n",
    "\n",
    "As clearly visible those images show significant shift in y direction, that is why we first apply the manual shift and later use automatic adjustment.\n",
    "\n",
    "Next, we manually set the center of each image using the `set_center_manual()` function. If numeric values for `center_x` and `center_y` are provided, the same shift is applied uniformly to all images in the stack.  \n",
    "Optionally, if a list or vector of `center_x` values is provided (with length equal to the number of images), each image will be shifted according to its corresponding value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "rheed_data = xrheed.load_data(\n",
    "    image_paths,\n",
    "    plugin=\"dsnp_arpes_bmp\",\n",
    "    stack_dim=stack_dim,\n",
    "    stack_coords=stack_coordinate,\n",
    ")\n",
    "\n",
    "rheed_data.ri.screen_roi_width = 40\n",
    "rheed_data.ri.screen_roi_height = 50\n",
    "\n",
    "# Plot the image to inspect the alignment\n",
    "rheed_data.ri.plot_image(auto_levels=1.0)\n",
    "\n",
    "center_x = -1.0\n",
    "center_y = -8.0\n",
    "plt.plot(center_x, center_y, '+')\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# Set approximate center of the images\n",
    "rheed_data.ri.set_center_manual(center_x, center_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "rheed_data.ri.set_center_auto()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Plot the Images from the Stack\n",
    "The `plot_image()` function displays the first image from the stack by default. Alternatively, a specific image can be shown by providing the `stack_index` keyword argument.\n",
    "\n",
    "If `stack_index` is set, the corresponding image at that index in the stack will be plotted instead of the default first image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "rheed_data.ri.plot_image(\n",
    "    show_center_lines=True,\n",
    "    auto_levels=0.2,\n",
    "    stack_index=0,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "The `plot_images()` function from the `plotting` module can be used to display all images in the stack.\n",
    "\n",
    "To visualize a subset of the images, standard xarray slicing can be applied. For example, to show every third image in the stack, use slicing as demonstrated below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xrheed.plotting import plot_images\n",
    "\n",
    "plot_images(rheed_data.isel(coverage=slice(None, None, 3)), ncols=2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## RHEED Stack Profiles\n",
    "\n",
    "The `get_profile` method, available through the `ri` accessor, returns a 2D dataset extracted from a 3D RHEED stack. By default, it integrates the data along the screen's Y direction, similar to how profiles are obtained from a single RHEED image.\n",
    "\n",
    "> **Note:** If `show_origin` is set to `True`, it is also possible to pass the `stack_index` argument to visualize the origin area. In this case, the origin will be marked by a rectangle on the selected image from the stack, rather than on the default first image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_2d = rheed_data.ri.get_profile(\n",
    "    center=(0, -7), \n",
    "    width=50, \n",
    "    height=15, \n",
    "    show_origin=True, \n",
    "    stack_index=25\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Plotting a 2D Profile\n",
    "\n",
    "The returned profile is an `xarray.DataArray` and can be visualized using its built-in `.plot()` method, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_2d.plot(vmin=10, vmax=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Extracting a Single Profile\n",
    "\n",
    "Optionally, it is possible to extract a single profile from a 2D stack using standard slicing or selection operations. This allows for flexible access to specific regions or lines within the dataset without invoking the full `get_profile` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_2d.isel(coverage=1).plot()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## 1D Profile from a 3D Stack\n",
    "\n",
    "The default behavior of the `get_profile` method can be modified by providing the `reduce_over` keyword argument. Allowed options are: `\"sy\"`, `\"sx\"`, and `\"both\"`.\n",
    "\n",
    "If the `\"both\"` option is selected, the profile is integrated over the entire selection area. This is particularly useful for spot analysis during MBE growth, for example, to record *RHEED oscillations*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_1d = rheed_data.ri.get_profile(\n",
    "    center=(-8.5, -9),\n",
    "    width=5,\n",
    "    height=5,\n",
    "    reduce_over=\"both\",\n",
    "    show_origin=True,\n",
    "    stack_index=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "Finally, the profile, returned as an `xarray.DataArray`, can be plotted using its internal `.plot()` method.\n",
    "Alternatively, it can be visualized using the `plot_profile()` method available via the `rp` accessor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_1d.rp.plot_profile()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}