{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e55d344d-fbf8-495b-a88d-db0c89d41d0e",
   "metadata": {},
   "source": [
    "# How to create a poincare plot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47fb8e79-cc9b-4f7c-af0e-55a31f112c6c",
   "metadata": {},
   "source": [
    "## Import everything we might need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f34dd740-f678-4cdf-bbb8-5c370c340e23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import xemc3\n",
    "import matplotlib.pyplot as plt\n",
    "import tqdm\n",
    "\n",
    "# Matplotlib setup\n",
    "import setup_plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a7842c7-ea97-41d1-a370-3d4ff60889c4",
   "metadata": {},
   "source": [
    "## Get some data\n",
    "\n",
    "Feel free to modify this, to "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66af6ca3-ff34-42b6-8e35-e21a272f1a69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use local helper function to get some data\n",
    "from get_data import load_example_data\n",
    "\n",
    "ds = load_example_data()\n",
    "# If you want to use your own data use something like\n",
    "# ds = xemc3.load.all(\"path/to/mydata/\")\n",
    "# or if you have converted it already to a netcdf file\n",
    "# ds = xr.open_dataset(\"path/to/mydata.nc\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f2ae519-aee2-4dcb-8619-95bdfe7f58a5",
   "metadata": {},
   "source": [
    "## Do the tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c08dad6-b7af-4e7c-947a-4bcd6118a130",
   "metadata": {},
   "outputs": [],
   "source": [
    "pnc0 = ds.emc3.tracer.poincare_phi0(\n",
    "    [(x, 0.0) for x in np.linspace(5.556, 5.9191, 20)], 1000\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38179be1-5d90-470c-9d11-0aa37c6be65b",
   "metadata": {},
   "source": [
    "## Plot it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e746b75-d77a-4ab5-929d-9ee049ad6473",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(pnc0[..., 0], pnc0[..., 1], s=0.1, color=\"black\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63d368ea-e8db-47d1-ba97-17eb59080095",
   "metadata": {},
   "source": [
    "## Trace to other cross sections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af1df02e-559c-4dc7-a130-62a40f85b95f",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi_indices = np.arange(5) * 9\n",
    "pncs = ds.emc3.tracer.trace_to_phi_index(pnc0.reshape(-1, 2), phi_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d733214-d334-448e-973f-10df4215555b",
   "metadata": {},
   "source": [
    "## Plot them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac658571-aa96-45e9-b17c-f7aab77213f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "phis = ds.emc3[\"phi_corners\"][phi_indices]\n",
    "for phi, pnc in zip(phis, pncs):\n",
    "    plt.figure()\n",
    "    pnc = np.array(pnc).T\n",
    "    plt.scatter(*pnc, s=0.1, color=\"black\")\n",
    "    plt.title(f\"phi = {phi.values}\")\n",
    "    plt.gca().set_aspect(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abb1860f-ffae-45bf-a1c5-52c32d50a91f",
   "metadata": {},
   "source": [
    "## Find fix points (e.g. x-points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42f00edb-111b-492f-908f-0d21b6b5edba",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.emc3.tracer.fix_point([5.5, 0.8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28d3dfb-ad04-43b5-bc20-17c183c7e8fc",
   "metadata": {},
   "source": [
    "## Scan for fix points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e9421bd-ad5a-447a-9f37-32b8cd4146f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xemc3.core.flt_emc_like import FixPointConvergenceError, OutOfDomainError\n",
    "\n",
    "found = []\n",
    "\n",
    "for x in tqdm.auto.tqdm(np.linspace(5.5, 6.3, 20)):\n",
    "    for y in np.linspace(0, 1.2, 40):\n",
    "        try:\n",
    "            new = ds.emc3.tracer.fix_point([x, y])\n",
    "            for old in found:\n",
    "                if np.sum((old - new) ** 2) < 1e-5:\n",
    "                    break\n",
    "            else:\n",
    "                print(f\"Found: {new}\")\n",
    "                found.append(new)\n",
    "        except (FixPointConvergenceError, OutOfDomainError):\n",
    "            # Not close to a fix point\n",
    "            pass\n",
    "        except IndexError:\n",
    "            # Not in domain\n",
    "            pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb3239b-d4da-478d-9c7d-f496a39d9a17",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(pnc0[..., 0], pnc0[..., 1], s=0.1, color=\"black\")\n",
    "plt.scatter(*np.array(found).T, s=1, color=\"red\")"
   ]
  }
 ],
 "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.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}