{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "947e43fe-12b4-48fc-822e-83faaf2b52c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import at\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "if sys.version_info.minor < 9:\n",
    "    from importlib_resources import files, as_file\n",
    "else:\n",
    "    from importlib.resources import files, as_file\n",
    "import matplotlib.pyplot as plt\n",
    "from at.tracking import lattice_pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0a723d43-73e0-440a-be91-26e1038b65f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib osx\n",
    "# Good quality plot in an external window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8c3a5b87-2bad-4695-b036-17bf6bf49105",
   "metadata": {},
   "outputs": [],
   "source": [
    "import at.plot\n",
    "data1 = Path.home() / 'dev/libraries/at/pyat/machine_data'\n",
    "data2 = Path.home() / 'dev/libraries/at/machine_data'\n",
    "data3 = Path.home() / 'ESRFCloud/notebooks/lattice'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "38f7fef8-158f-41d1-bf60-6b4461c9e16b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ring=at.load_lattice(data3 / 'S28F.mat', mat_key='LOW_EMIT_RING_INJ')\n",
    "ring.disable_6d()\n",
    "ringrad=ring.enable_6d(all_pass=None, dipole_pass='auto', copy=True)\n",
    "ringq=ring.enable_6d(all_pass=None, dipole_pass='BndMPoleSymplectic4QuantPass', quadrupole_pass='StrMPoleSymplectic4QuantPass', copy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8de5cdaf-373b-4964-91c2-28fa5d2571de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.00046605467208399255\n"
     ]
    }
   ],
   "source": [
    "vq=np.squeeze(at.lattice_pass(ringq,np.zeros((6,1000)),1,len(ringrad)))\n",
    "meanq = np.mean(vq[4, :])\n",
    "print(meanq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7e227536-1add-4ede-adf3-3e7a1b98acd0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.00042194334321625726\n"
     ]
    }
   ],
   "source": [
    "vr=np.squeeze(at.lattice_pass(ringrad,np.zeros((6,300)),1,len(ringrad)))\n",
    "meanr = np.mean(vr[4, :])\n",
    "print(meanr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "49479699-a62e-49f3-985f-f051c31eb44d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(-0.0014, 400, 'mean(QuantPass): -4.661e-04\\n               RadPass: -4.219e-04')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.hist([vq[4,:], vr[4,:]],bins=30,range=(-1.5e-3, 0), label=['QuantPass', 'RadPass'],rwidth=1.0)     \n",
    "ax.set_xlabel('$\\Delta$p/p for 1 turn')\n",
    "ax.set_ylabel('Counts')\n",
    "ax.set_title('S28F, 1 turn, 1000 samples')\n",
    "ax.legend()\n",
    "ax.text(-0.0014, 400, f\"mean(QuantPass): {meanq:.3e}\\n               RadPass: {meanr:.3e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f249e96a-7071-465e-a146-692f8274de19",
   "metadata": {},
   "outputs": [],
   "source": [
    "rq = np.squeeze(lattice_pass(ringq[:130], np.zeros((6,1000)), refpts=range(130)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "04a74e38-ac97-422a-9a44-a2d93abf2911",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr = np.squeeze(lattice_pass(ringrad, np.zeros((6,5)), refpts=range(130)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "09dd6bda-85bd-4ce9-a23d-02ed2c587c23",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x124d8a130>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(np.mean(rq[4,:,:],axis=0),label='mean(QuantPass)')\n",
    "ax.plot(np.mean(rr[4,:,:],axis=0),label='RadPass')\n",
    "ax.set_xlabel('Element #')\n",
    "ax.set_ylabel('$\\Delta$p/p')\n",
    "ax.set_title('S28F, 130 elements')\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d2ca4c8d-5a32-419f-a801-87dfa40ffdc5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3667"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ringrad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1dfd754-ce75-4f97-9b77-599673d56652",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}