{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Multi-Image Transformation\n",
    "\n",
    "This example notebook demonstrates how to load and transform a series of RHEED images taken during sequential azimuthal rotations. \n",
    "\n",
    "The rotation starts at $\\alpha = 0^\\circ$, where the electron beam is aligned along the $[11\\bar{2}]$ direction, and ends at $\\alpha = 30^\\circ$, corresponding to alignment along the $[1\\bar{1}0]$ direction.\n"
   ]
  },
  {
   "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",
    "import xarray as xr\n",
    "\n",
    "import xrheed\n",
    "from xrheed.plotting.overview import plot_images\n",
    "from xrheed.preparation.filters import high_pass_filter\n",
    "from xrheed.preparation import (\n",
    "    find_horizontal_center,\n",
    "    find_incident_angle,\n",
    "    find_vertical_center,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Image Preprocessing and Correction\n",
    "\n",
    "To begin, we load a series of images and apply the necessary corrections.  \n",
    "Due to minor manipulator inaccuracies, each image may exhibit slight variations in positional shifts along the $S_x$ and $S_y$ axes.\n",
    "\n",
    "Typically, these variations change linearly during the rotation.  \n",
    "\n",
    "For instance, the shift along $S_y$ could be prepared as a linear space between the values recorded for the first and last images.  \n",
    "\n",
    "The same approach applies to the azimuthal angle, α.\n",
    "\n",
    "The example shown below is using a semiautomatic alignment where we use some of the images to get the center position and incident angle and later use those values to calculate linear components applied for all images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dir = Path(\"example_data\")\n",
    "image_paths = list(sorted(image_dir.glob(\"Si_111_r3Ag_thA_phi*.raw\")))\n",
    "\n",
    "n_images = len(image_paths)\n",
    "\n",
    "# Load data into a stack\n",
    "arheed = xrheed.load_data(\n",
    "    image_paths, \n",
    "    plugin=\"dsnp_arpes_raw\", \n",
    "    stack_dim=\"frame\", \n",
    "    stack_coords=np.arange(n_images)\n",
    ")\n",
    "\n",
    "# Set screen ROI\n",
    "arheed.ri.screen_roi_width = 30\n",
    "arheed.ri.screen_roi_height = 30\n",
    "\n",
    "# Select some of images where transmission spots is visible and use them for guided alignment\n",
    "selected_indices = [0, 1, 2, 3]\n",
    "all_indices = np.arange(n_images)\n",
    "\n",
    "center_x_samples = []\n",
    "center_y_samples = []\n",
    "incident_angle_samples = []\n",
    "\n",
    "for idx in selected_indices:\n",
    "    # Work on copy of an image - do not shift the original stack\n",
    "    rheed = arheed[idx].copy()\n",
    "\n",
    "    # Get ROI image\n",
    "    roi_image = rheed.ri.get_roi_image()\n",
    "\n",
    "    # Find center coordinates\n",
    "    center_x = find_horizontal_center(roi_image)\n",
    "    center_y = find_vertical_center(roi_image, center_x=center_x)\n",
    "\n",
    "    # Update ROI center\n",
    "    rheed.ri.set_center_manual(center_x, center_y)\n",
    "\n",
    "    # Recompute ROI and incident angle\n",
    "    roi_image = rheed.ri.get_roi_image()\n",
    "    incident_angle = find_incident_angle(roi_image)\n",
    "    rheed.ri.incident_angle = incident_angle\n",
    "\n",
    "    # Store results\n",
    "    center_x_samples.append(center_x)\n",
    "    center_y_samples.append(center_y)\n",
    "    incident_angle_samples.append(incident_angle)\n",
    "\n",
    "    # Plot diagnostics\n",
    "    rheed.ri.plot_image(\n",
    "        show_center_lines=True,\n",
    "        show_specular_spot=True,\n",
    "        auto_levels=0.5\n",
    "    )\n",
    "    plt.title(\"\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Linear fits for interpolation across all images\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "slope, intercept = np.polyfit(selected_indices, center_x_samples, 1)\n",
    "center_x_fit = all_indices * slope + intercept\n",
    "\n",
    "slope, intercept = np.polyfit(selected_indices, center_y_samples, 1)\n",
    "center_y_fit = all_indices * slope + intercept\n",
    "\n",
    "slope, intercept = np.polyfit(selected_indices, incident_angle_samples, 1)\n",
    "incident_angle_fit = all_indices * slope + intercept\n",
    "\n",
    "arheed.ri.set_center_manual(center_x=center_x_fit, center_y=center_y_fit)\n",
    "\n",
    "alpha_coords = np.linspace(0.0, 29.6, n_images)\n",
    "beta_coords = incident_angle_fit\n",
    "\n",
    "arheed = arheed.assign_coords(\n",
    "    alpha=(\"frame\", alpha_coords),\n",
    "    beta=(\"frame\", beta_coords),\n",
    ")\n",
    "\n",
    "# Set screen ROI\n",
    "arheed.ri.screen_roi_width = 30\n",
    "arheed.ri.screen_roi_height = 40\n",
    "\n",
    "# Optionally use high pass filter (could be applied to the whole stack)\n",
    "arheed = high_pass_filter(arheed, sigma=2.0, threshold=0.7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Prepare a Lattice Object for the Expected Reconstruction\n",
    "\n",
    "To verify the alignment, we create a lattice object representing the expected (√3 × √3) R30° reconstruction.  \n",
    "\n",
    "This lattice will later be used by the `Ewald` object for spot calculation and overlay.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xrheed.kinematics import Lattice\n",
    "\n",
    "# Si(111)-(1x1)\n",
    "si_111_1x1 = Lattice.from_surface_hex(a=3.84, label=\"(1x1)\")\n",
    "\n",
    "# Si(111)-(r3xr3)R30\n",
    "si_111_r3 = Lattice.from_surface_hex(a=3.84 * np.sqrt(3), label=\"(√3 x √3)\")\n",
    "si_111_r3.rotate(30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Generate Ewald Object and Overlay Expected Spot Positions\n",
    "\n",
    "We generate an `Ewald` objects and use it to overlay the expected spot positions on each image.\n",
    "\n",
    "This allows for visual verification of the alignment and consistency across the dataset.\n",
    "\n",
    "> **Note**: The `Ewald` object could take a stack of images. In such case a `stack_index` property could be used to select a particular image from a stack.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xrheed.kinematics import Ewald\n",
    "\n",
    "arheed.ri.screen_roi_width = 60\n",
    "arheed.ri.screen_roi_height = 60\n",
    "\n",
    "ew = Ewald(si_111_r3, arheed)\n",
    "\n",
    "fig, axs = plt.subplots(n_images, 1, sharex=True, figsize=(4.0, 14))\n",
    "\n",
    "for idx in range(n_images):\n",
    "    ax = axs[idx]\n",
    "    # Select the image from the stack\n",
    "    ew.stack_index = idx\n",
    "\n",
    "    ew.plot(\n",
    "        ax=ax, show_image=True,\n",
    "        vmin=3, vmax=35,\n",
    "        marker=\"o\", facecolor=\"none\", edgecolor=\"c\", s=50,\n",
    "    )\n",
    "    ax.set_title(\"\")\n",
    "    if idx < n_images - 1:\n",
    "        ax.set_xlabel(\"\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Transform Image Stack to kx–ky Space\n",
    "\n",
    "If the alignment appears satisfactory, we proceed by transforming the whole image stack into the kx–ky space.  \n",
    "\n",
    "This step uses the arguments `rotate=True` and `point_symmetry=True`.  \n",
    "\n",
    "- The `rotate=True` option applies a rotation around the image center by the azimuthal angle α.  \n",
    "- The `point_symmetry=True` option adds a mirrored image rotated by 180°, which is valid for most crystallographic structures.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xrheed.conversion import transform_stack_to_kxky\n",
    "\n",
    "kxky_stack = transform_stack_to_kxky(arheed, point_symmetry=True, rotate=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Compute the Mean Along the Azimuthal Dimension\n",
    "\n",
    "Calculate the mean of the dataset for all frames.  \n",
    "This operation reduces the data to a single representative image, averaging over all angular orientations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "kxky_stack_mean = kxky_stack.mean(dim=\"frame\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Prepare the Final Plot\n",
    "\n",
    "Finally, we generate an image that displays the mean intensity across all frames, along with overlays of the (1×1) and (√3 × √3) R30° reciprocal lattices for comparison.\n",
    "\n",
    "As observed, the diffraction spots are elongated along several directions.  \n",
    "This elongation is a characteristic feature of RHEED, resulting from the finite domain size, which affects the shape of the diffraction spots.\n",
    "\n",
    "Nonetheless, in this representation, it is straightforward to identify specific surface reconstructions, their symmetries, and mutual orientations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(5, 5))\n",
    "\n",
    "kxky_stack_mean.plot(ax=ax, cmap=\"gray\", add_colorbar=False, vmin=2, vmax=20)\n",
    "\n",
    "si_111_r3.plot_reciprocal(ax=ax, facecolor=\"none\", edgecolor=\"c\", s=150)\n",
    "si_111_1x1.plot_reciprocal(ax=ax, facecolor=\"none\", edgecolor=\"m\", s=50)\n",
    "\n",
    "ax.set_aspect(1)\n",
    "ax.set_facecolor(\"black\")\n",
    "ax.set_xlim(-2.5, 2.5)\n",
    "ax.set_ylim(-2.5, 2.5)\n",
    "ax.legend(loc=\"upper right\")\n",
    "ax.set_title(\"\")\n",
    "\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
}