{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# LHC source\n\nEvolution of the wavefront emitted by the LHC source during the energy ramp.\nThis examples shows how to setup a simulation with protons and to use one of \nthe predefined magnetic provided by the PySRW module.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We import the usual packages, along with the `animation` submodule of \nmatploitlib needed to create the sequence of the wavefront evolution.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nimport matplotlib.animation as animation\n\nimport pysrw as srw"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The PySRW library is initialized with 'proton'.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "srw.CONFIG.initParticleType(\"proton\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The magnetic container can be directly created calling the function \n:py:func:`~pysrw.magnets.predefinedLHCsource`, passing the beam energy. \nIn order to include all devices of the LHC source, we leave the `sourceFlag`\nparameter to its default value 'D3undD3'. Note that, by convention, the \norigin of the longitudinal coordinate (z=0) is set at the undulator centre.\nWe initialize the beam at z=-3, leaving both transverse positions and angles\nto zero.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "energy = 450    # GeV\n\nmag_container = srw.magnets.predefinedLHCsource(energy)\nbeam = srw.emitters.ParticleBeam(energy, xPos=0, yPos=0, zPos=-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The trajectory is simulated to check the good instantiaiton of the objects.\nWe track the particle between s=-130 and s=30, to include both D3 dipoles.\nAs expected from the choice of the beam initial conditions, the beam travels\nalong the longitudinal axis in the straight section between the dipoles.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "traj = srw.computeTrajectory(beam, mag_container, sStart=-130, sEnd=30)\nsrw.plotTrajectory(traj)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We simulate the distribution of the emitted wavefront. Since we will not \npropagate it through an optical system, we can define a rather coarse \nobservation mesh to make the simulation faster.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wl = 560        # nm\nobserver = srw.wavefronts.Observer(centerCoord=[0, 0, 27.24],\n                                   obsXextension=50e-3, \n                                   obsYextension=50e-3,\n                                   obsXres=200e-6,\n                                   obsYres=200e-6)\n\n\nwfr = srw.computeSrWfrMultiProcess(numCores=4, \n                        particleBeam=beam, magnetsContainer=mag_container,\n                        observer=observer, wavelength=wl)\nintensity = wfr.getWfrI()\nsrw.plotI(intensity)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we can iterate the above caluclation changing the beam energy \nand following the evolution of the spatial intensity as a function of \nthe beam energy (and at the fixed wavelength of 560 nm).\nThis simulation clearly shows the undulator radiation that dominates the \nemission at injection. As the energy increases, the visible undulator \nradiation drifts to larger polar angles outside of our mesh. Above 1 TeV, \ndipole radiation appears first in form of edge radiation and later as core\nradiation. Circular fringes due to the interference between the two D3s are\nalso clearly visible.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "energies = [450, 550, 800, 1600, 3300, 6800]\n\nframes = []\nfig, ax = plt.subplots()\nax.set_xlabel(\"x position [mm]\")\nax.set_xlabel(\"y position [mm]\")\nextent = np.array([intensity[\"xax\"][0], intensity[\"xax\"][-1], \n                   intensity[\"yax\"][0], intensity[\"xax\"][-1]]) * 1e3\n\nfor energy in energies:\n\n    mag_container = srw.magnets.predefinedLHCsource(energy)\n    beam = srw.emitters.ParticleBeam(energy, xPos=0, yPos=0, zPos=-3)\n\n    wfr = srw.computeSrWfrMultiProcess(numCores=4, \n                        particleBeam=beam, magnetsContainer=mag_container,\n                        observer=observer, wavelength=wl)\n    intensity = wfr.getWfrI()\n\n    frame = ax.imshow(intensity[\"data\"], cmap=\"jet\", extent=extent)\n    annotation = ax.annotate(f\"Energy: {energy:.0f} GeV\",\n                             (.05,.95), xycoords=\"axes fraction\", \n                             color=\"white\", fontsize=12)\n    frames.append([frame, annotation])\n \nanim = animation.ArtistAnimation(fig, frames, \n                                 interval=500, blit=True, repeat_delay=1000)"
      ]
    }
  ],
  "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.9.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}