{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "### Iteration across coordinates\n",
    "\n",
    "`iter_coords` allows you to iterate across the coordinates of a `DataArray`, without also iterating across the data values. This can be useful if you would like to transform the coordinates before selection, or would only like access to the coordinates. \n",
    "\n",
    "Here is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import arpes\n",
    "\n",
    "# Set the random seed so that you get the same numbers\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "\n",
    "import xarray as xr\n",
    "import arpes.xarray_extensions  # so .G is in scope\n",
    "\n",
    "test_data = xr.DataArray(\n",
    "    np.random.random((3, 3)),\n",
    "    coords={\"X\": [0, 1, 2], \"Y\": [-5, -4, -3]},\n",
    "    dims=[\"X\", \"Y\"],\n",
    ")\n",
    "test_data.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for coordinate in test_data.G.iter_coords():\n",
    "    print(coordinate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "You can also iterate simultaneously over the coordinates and their indices in the data with `enumerate_iter_coords`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for index, coordinate in test_data.G.enumerate_iter_coords():\n",
    "    print(index, coordinate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Raveling/Flattening\n",
    "\n",
    "It is sometimes necessary to have access to the data in a flat format where each of the coordinates has the full size of the data. The most common usecase is in preparing an isosurface plot, but this functionality is also used internally in the coordinate conversion code.\n",
    "\n",
    "The return value is a dictionary, with keys equal to all the dimension names, plus a special key \"data\" for the values of the array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from arpes.io import example_data\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data = example_data.temperature_dependence.spectrum.sel(\n",
    "    eV=slice(-0.08, 0.05), phi=slice(-0.22, None)\n",
    ").sum(\"eV\")\n",
    "\n",
    "raveled = data.G.ravel()\n",
    "fig = plt.figure(figsize=(7, 7))\n",
    "ax = fig.add_subplot(projection=\"3d\")\n",
    "\n",
    "ax.plot_trisurf(\n",
    "    raveled[\"temperature\"],  # use temperature as the X coordinates\n",
    "    raveled[\"phi\"],  # use phi as the Y coordinates\n",
    "    data.values.T.ravel(),  # use the intensity as the Z coordinate\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Functional Programming Primitives: `filter` and `map`\n",
    "\n",
    "You can `filter` or conditionally remove some of a datasets contents. To do this over coordinates on a dataset according to a function/sieve which accepts the coordinate and data value, you can use `filter_coord`. The sieving function should accept two arguments, the coordinate and the cut at that coordinate respectively. You can specify which coordinate or coordinates are iterated across when filtering using the `coordinate_name` parameter.\n",
    "\n",
    "As a simple, example, we can remove all the odd valued coordinates along Y:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data.G.filter_coord(\"Y\", lambda y, _: y % 2 == 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Functional programming can also be used to modify data. With `map` we can apply a function onto a `DataArray`'s values. You can use this to add one to all of the elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data.G.map(lambda v: v + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "Additionally, we can simultaneously iterate and apply a function onto a specified dimension of the data with `map_axes`. Here we can use this to ensure that the rows along `Y` have unit norm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data.G.map_axes(\"Y\", lambda v, c: v / np.linalg.norm(v))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Shifting\n",
    "\n",
    "Suppose you have a bundle of spaghetti in your hand with varying lengths. You might want to align them so that they all meet in a flat plane at the tops of the strands. In general, you will have to shift each a different amount depending on the length of each strand, and its initial position in your hand.\n",
    "\n",
    "A similar problem presents itself in multidimensional data. You might want to shift 1D or 2D \"strands\" of data by differing amounts along an axis. One practical use case in ARPES is to [align the chemical potential](/fermi-edge-correction) to take into account the spectrometer calibration and shape of the spectrometer entrance slit. Using the curve fitting data we explored in the previous section, we can align the data as a function of the temperature so that all the Fermi momenta are at the same index:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# first, get the same cut and fermi angles/momenta from the previous page\n",
    "# this is reproduced for clarity and so you can run the whole notebook\n",
    "# please feel free to skip...\n",
    "\n",
    "from lmfit.models import LinearModel, LorentzianModel\n",
    "\n",
    "temp_dep = example_data.temperature_dependence\n",
    "near_ef = temp_dep.sel(eV=slice(-0.05, 0.05), phi=slice(-0.2, None)).sum(\"eV\").spectrum\n",
    "model = LinearModel(prefix=\"a_\") + LorentzianModel(prefix=\"b_\")\n",
    "lorents_params = LorentzianModel(prefix=\"b_\").guess(\n",
    "    near_ef.sel(temperature=20, method=\"nearest\").values, near_ef.coords[\"phi\"].values\n",
    ")\n",
    "\n",
    "phis = phis = near_ef.S.modelfit(\"phi\", model, params=lorents_params).modelfit_results.F.p(\n",
    "    \"b_center\"\n",
    ")\n",
    "\n",
    "\n",
    "# ...to here\n",
    "fig, ax = plt.subplots(1, 3, figsize=(13, 4))\n",
    "near_ef.S.plot(ax=ax[0])\n",
    "near_ef.G.shift_by(phis - phis.mean(), shift_axis=\"phi\").S.plot(ax=ax[1])\n",
    "near_ef.G.shift_by(phis - phis.mean(), shift_axis=\"phi\", extend_coords=True).S.plot(ax=ax[2])\n",
    "\n",
    "ax[0].set_title(\"Original data\")\n",
    "ax[1].set_title(\"Shifted to align Fermi angle\")\n",
    "ax[2].set_title(\"Shifted to align Fermi angle (with coord extension\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}