{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Diffraction Profiles\n",
    "\n",
    "This notebook demonstrates how to use **xRHEED** to extract a diffraction profile from a RHEED image.\n",
    "\n",
    "As shown in the **Getting Started** notebook, the first step is to load the `xrheed` library and then import the RHEED image.\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",
    "\n",
    "import xrheed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dir = Path(\"example_data\")\n",
    "image_path = image_dir / \"Si_111_7x7_112_phi_00.raw\"\n",
    "\n",
    "rheed_image = xrheed.load_data(image_path, plugin=\"dsnp_arpes_raw\")\n",
    "print(rheed_image.ri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Data Preparation\n",
    "\n",
    "Before analyzing the RHEED image, it should be properly aligned. This may involve applying a rotation if necessary, and shifting the image horizontally and vertically to position the center accurately.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a copy of the original RHEED image\n",
    "aligned_image = rheed_image.copy()\n",
    "\n",
    "# Apply rotation to correct image alignment\n",
    "aligned_image.ri.rotate(-0.4)\n",
    "\n",
    "# Replace the original image with the rotated version for further analysis\n",
    "rheed_image = aligned_image\n",
    "\n",
    "# Apply screen scaling correction if necessary\n",
    "# Adjust based on exact calibration (here it should be 1.001857)\n",
    "scaling_correction_factor = 1.0\n",
    "rheed_image.ri.screen_scale *= scaling_correction_factor\n",
    "\n",
    "# Set screen ROI\n",
    "rheed_image.ri.screen_roi_width = 60\n",
    "rheed_image.ri.screen_roi_height = 70\n",
    "\n",
    "# Automatically determine and apply the image center after rotation\n",
    "rheed_image.ri.set_center_auto(update_incident_angle=True)\n",
    "\n",
    "# Use automatic levels adjustment\n",
    "rheed_image.ri.plot_image(\n",
    "    auto_levels=1.0, show_center_lines=True, show_specular_spot=True\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Profile Extraction\n",
    "\n",
    "Since the RHEED image is stored as a `DataArray`, a diffraction profile can be easily extracted using the built-in `sel` method, as shown below. \n",
    "\n",
    "However, it is recommended to use the built-in accessor for profile extraction, which will be demonstrated later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_range = (-20, 20)  # in mm\n",
    "y_range = (-10, 0)   # in mm\n",
    "\n",
    "profile = rheed_image.sel(sx=slice(*x_range), sy=slice(*y_range)).mean(\"sy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Profile Extraction Using the `get_profile` Method\n",
    "\n",
    "The `get_profile` method, available through the `ri` accessor, can also be used to extract a diffraction profile by specifying the center point, width, and height of the region.\n",
    "\n",
    "Additionally, this function can plot the RHEED image with the profile region highlighted by setting `plot_origin=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile = rheed_image.ri.get_profile(\n",
    "    center=(0, -5), width=40, height=2, show_origin=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "The `rp` accessor provides basic information about the extracted profile, making it easier to inspect and interpret the data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile.rp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "A RHEED profile retains the attributes of its parent image, ensuring consistency in metadata and coordinate references.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile.ri"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Plotting the Profile\n",
    "\n",
    "The `plot_profile` function, accessible via the `ri` accessor, is used to visualize the extracted profile. It supports additional parameters such as `normalize` and `transform_to_k`.\n",
    "\n",
    "When `transform_to_k=True`, the profile is plotted using a temporary scattering coordinate, $k_y$, which provides a momentum-resolved representation of the diffraction data.\n",
    "\n",
    "> **Note:** In the coordinate system used by the xRHEED project, the screen's horizontal axis (x-coordinate) is parallel to $k_y$ in momentum space. For further details, refer to the **Geometry** section of the documentation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile.rp.plot_profile(\n",
    "    transform_to_k=True, normalize=True, color=\"black\", linewidth=1.0\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Converting the Profile to $k_y$\n",
    "\n",
    "The profile can be permanently converted to momentum space coordinates using the dedicated `ri.convert_to_k` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_k = profile.rp.convert_to_k()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Profile Fitting\n",
    "\n",
    "The diffraction profile can be fitted using the `lmfit` library, as demonstrated in the example below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lmfit.models import LorentzianModel, QuadraticModel\n",
    "\n",
    "# Extract data from profile\n",
    "x = profile_k.coords[\"ky\"].values\n",
    "y = profile_k.values\n",
    "\n",
    "# Preprocess: remove background offset and normalize\n",
    "y = y - np.min(y)\n",
    "y = y / np.max(y)\n",
    "\n",
    "# Define individual models\n",
    "l1 = LorentzianModel(prefix=\"l1_\")\n",
    "l2 = LorentzianModel(prefix=\"l2_\")\n",
    "bkg = QuadraticModel(prefix=\"bkg_\")\n",
    "\n",
    "# Combine into a composite model\n",
    "model = l1 + l2 + bkg\n",
    "\n",
    "# Initialize parameters with reasonable guesses\n",
    "params = model.make_params()\n",
    "params[\"l1_center\"].set(value=-3.0)\n",
    "params[\"l1_amplitude\"].set(value=1.0)\n",
    "params[\"l1_sigma\"].set(value=0.1)\n",
    "\n",
    "params[\"l2_center\"].set(value=3.0)\n",
    "params[\"l2_amplitude\"].set(value=1.0)\n",
    "params[\"l2_sigma\"].set(value=0.1)\n",
    "\n",
    "params[\"bkg_a\"].set(value=-1.0)\n",
    "params[\"bkg_b\"].set(value=0.0)\n",
    "params[\"bkg_c\"].set(value=0.0)\n",
    "\n",
    "# Perform the fit\n",
    "result = model.fit(y, params, x=x)\n",
    "\n",
    "# Plot the fit result\n",
    "result.plot_fit(title=\"RHEED Profile Fit\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate half the distance between the two fitted peaks\n",
    "half_peak_distance = 0.5 * (\n",
    "    result.params[\"l2_center\"].value - result.params[\"l1_center\"].value\n",
    ")\n",
    "\n",
    "print(f\"Half peak distance from the specular reflection: {half_peak_distance:.2f} 1/Å\")\n",
    "\n",
    "# Expected peak separation for Si(111)-(1×1) along the [110] direction\n",
    "expected_distance = 4 * np.pi / 3.84  # 1/A\n",
    "print(f\"Expected peak separation: {expected_distance:.2f} 1/Å\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Fine Adjustment of the Screen Scale\n",
    "\n",
    "The measured distance between the two diffraction peaks can be used to refine the screen scale. If the measured peak separation differs from the expected value, the scale factor may require adjustment to ensure accurate momentum calibration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling_correction = half_peak_distance / expected_distance\n",
    "\n",
    "print(f\"Calculated correction of the screen scale: {scaling_correction:.6f}\")\n",
    "\n",
    "rheed_image.ri.screen_scale *= scaling_correction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## Preparing a Final Plot\n",
    "\n",
    "Generate a final plot with a polished appearance to clearly present the fitted RHEED profile and highlight key features.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate components\n",
    "comps = result.eval_components(x=x)\n",
    "\n",
    "# Create figure\n",
    "fig, ax = plt.subplots(figsize=(5, 4), tight_layout=True)\n",
    "\n",
    "# Plot data and total fit\n",
    "ax.plot(x, y, \"ko\", markersize=2, label=\"Data\")  # black dots for data\n",
    "ax.plot(\n",
    "    x, result.best_fit, color=\"tab:blue\", linestyle=\"-\", linewidth=1, label=\"Total Fit\"\n",
    ")\n",
    "\n",
    "# Plot Lorentzian components as shaded areas\n",
    "for i, color in zip(range(1, 3), [\"tab:blue\", \"tab:green\", \"tab:purple\"]):\n",
    "    comp = comps[f\"l{i}_\"]\n",
    "    ax.fill_between(x, comp, color=color, alpha=0.2, label=f\"Lorentzian {i}\")\n",
    "\n",
    "# Plot background as a dashed line\n",
    "ax.plot(\n",
    "    x, comps[\"bkg_\"], color=\"gray\", linestyle=\"--\", linewidth=1.0, label=\"Background\"\n",
    ")\n",
    "\n",
    "# Labels and styling\n",
    "ax.set_xlabel(\"$k_y$ (1/Å)\", fontsize=10)\n",
    "ax.set_ylabel(\"Normalized Intensity\", fontsize=10)\n",
    "\n",
    "# Legend outside the plot area\n",
    "ax.legend(fontsize=9)\n",
    "ax.set_xlim(-4.5, 4.5)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "## Vertical Profiles\n",
    "\n",
    "By default, profiles are averaged over the `sy` (vertical) direction. This is because, in the defined geometry, `sx` corresponds to the sample's `ky` direction, which is perpendicular to the electron beam.\n",
    "\n",
    "However, profiles can also be computed by averaging or summing over `sx` if the `reduce_over` argument is provided, as shown below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertical_profile = rheed_image.ri.get_profile(\n",
    "    center=(0, -30),\n",
    "    width=5,\n",
    "    height=65,\n",
    "    reduce_over=\"sx\",\n",
    "    method=\"mean\",\n",
    "    show_origin=True,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertical_profile.rp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(5, 4), tight_layout=True)\n",
    "\n",
    "vertical_profile.rp.plot_profile(ax=ax, color=\"k\")\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
}