emittance_from_beamenvelope.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:29.670673Z",
     "start_time": "2021-10-08T17:33:25.662212Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: Cannot change to a different GUI toolkit: notebook. Using qt5 instead.\n"
     ]
    }
   ],
   "source": [
    "%matplotlib qt5\n",
    "%matplotlib notebook\n",
    "\n",
    "import numpy as _np\n",
    "import matplotlib.pyplot as mplt\n",
    "import scipy.linalg as scylin\n",
    "\n",
    "import pyaccel\n",
    "from pymodels import si, bo\n",
    "\n",
    "mplt.rcParams.update({'grid.linestyle': '--', 'grid.alpha': 0.5, 'axes.grid': True, 'font.size': 14})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:30.518147Z",
     "start_time": "2021-10-08T17:33:29.672201Z"
    }
   },
   "outputs": [],
   "source": [
    "acc = si.create_accelerator()\n",
    "famdata = si.get_family_data(acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# acc = bo.create_accelerator()\n",
    "# acc.energy = 3e9\n",
    "# famdata = si.get_family_data(acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:31.105348Z",
     "start_time": "2021-10-08T17:33:30.519429Z"
    }
   },
   "outputs": [],
   "source": [
    "envelope, cumul_mat, bdiff, fixed_point = pyaccel.optics.calc_beamenvelope(acc, full=True, energy_offset=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:31.155446Z",
     "start_time": "2021-10-08T17:33:31.106610Z"
    }
   },
   "outputs": [],
   "source": [
    "m66 = cumul_mat[-1]\n",
    "# # To calculate the emittances along the whole ring uncomment the\n",
    "# # line below:\n",
    "# m66 = _np.linalg.solve(\n",
    "#     self._cumul_mat.transpose(0, 2, 1),\n",
    "#     (self._cumul_mat @ m66).transpose(0, 2, 1)).transpose(0, 2, 1)\n",
    "\n",
    "# Look at section  D.2 of the Ohmi paper to understand this part of the\n",
    "# code on how to get the emmitances:\n",
    "\n",
    "m66 = m66\n",
    "env = envelope[0]\n",
    "bdf = bdiff[-1]\n",
    "\n",
    "# The function numpy.linalg.eig returns the evecs matrix such that:\n",
    "#    evecs^-1 @ m66 @ evecs = np.diag(evals)\n",
    "evals, evecs = _np.linalg.eig(m66)\n",
    "evecsh = evecs.swapaxes(-1, -2).conj()\n",
    "# Notice that the transformation generated by matrix evecs is the\n",
    "# inverse of equation 62 of the Ohmi paper.\n",
    "# So we need to calculate the inverse of evecs:\n",
    "evecsi = _np.linalg.inv(evecs)\n",
    "evecsih = evecsi.swapaxes(-1, -2).conj()\n",
    "\n",
    "# Then, using equation 64, we have:\n",
    "env0r = evecsi @ env @ evecsih\n",
    "emits = _np.diagonal(\n",
    "    env0r, axis1=-1, axis2=-2).real.take([0, 2, 4], axis=-1)\n",
    "\n",
    "# I don't understand why I have to divide the resulting emittances by\n",
    "# the norms of evecsi:\n",
    "norm = _np.linalg.norm(evecsi, axis=-1)\n",
    "emits /= norm.take([0, 2, 4], axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:31.257947Z",
     "start_time": "2021-10-08T17:33:31.157180Z"
    }
   },
   "outputs": [],
   "source": [
    "diff = evecsi @ bdf @ evecsih"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:31.365976Z",
     "start_time": "2021-10-08T17:33:31.263181Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.4893504e-10-1.3233617e-23j, 2.4893504e-10+1.3233685e-23j,\n",
       "       2.0614755e-06-4.2073822e-23j, 2.0614755e-06+4.7148007e-23j,\n",
       "       9.2149696e-38-4.1463335e-54j, 9.2149696e-38+4.1463335e-54j])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_np.set_printoptions(precision=7, suppress=False)\n",
    "-_np.diag(diff)/_np.log(_np.abs(evals))/2 / norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T17:33:31.479707Z",
     "start_time": "2021-10-08T17:33:31.371175Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2.4896065e-10,  2.0617528e-06, -4.4303804e-39])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "emits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T18:25:35.974262Z",
     "start_time": "2021-10-08T18:25:35.967171Z"
    }
   },
   "outputs": [],
   "source": [
    "iqs = famdata['QS']['index'][0][0]\n",
    "acc[iqs].KsL = 0.0046"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T18:25:36.478116Z",
     "start_time": "2021-10-08T18:25:36.137335Z"
    }
   },
   "outputs": [],
   "source": [
    "eqpar = pyaccel.optics.EqParamsFromBeamEnvelope(acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-08T18:44:13.265737Z",
     "start_time": "2021-10-08T18:44:13.096951Z"
    }
   },
   "outputs": [],
   "source": [
    "spos = _np.array(pyaccel.lattice.find_spos(acc, indices='closed'))\n",
    "fig = mplt.figure()\n",
    "gs = mplt.GridSpec(2, 1)\n",
    "ax = fig.add_subplot(gs[0, 0])\n",
    "ay = fig.add_subplot(gs[1, 0], sharex=ax)\n",
    "ax.plot(spos, eqpar.tilt_xyplane/_np.pi * 180, 'o-')\n",
    "ay.plot(spos, 1e12*(eqpar.envelopes[:, 0, 0]- eqpar.envelopes[:, 2, 2]), color='C1')\n",
    "ax.set_xlim([0, 518.4/5])\n",
    "\n",
    "ax.set_ylabel('Angle [°] [m]')\n",
    "ay.set_xlabel('Position [m]')\n",
    "ay.set_ylabel(r'$\\sigma_x^2 - \\sigma_y^2$ [$\\mu m^2$]')\n",
    "mplt.tight_layout()\n",
    "mplt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "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.6.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}