{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Simple planar undulator\n\nBasic example of computation of the field emitted by an undulator of a \nsynchrotron light source.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import modules required for simulations\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\n\nimport pysrw as srw"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We take as example the parameters of an undulator installed in a\nlight source (ALBA-Spain, NCD beamline)\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "energy = 3.0        # GeV\nlambda_u = 0.0216   # m\nN_u = 92\nB_0 = 0.75          # T"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The magnet is created with the default SRW model for planar undulators \n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "und = srw.magnets.UndulatorPlanar(energy=energy, \n                                  undPeriod=lambda_u, numPer=N_u,\n                                  magField=B_0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The deflection parameter and fundamental harmonic of such undulator are \n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "K = und.getK()\nlambda_1 = und.getLambda1()\nprint(f\"Deflection parameter: {K:.2f}\")\nprint(f\"Fundamental harmonic: {lambda_1:.2f} nm\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The magnetic element must be included in a magnetic container.\nIn this simple example, the `dipole` is the only device and the limits for \nthe numerical integration as left as default.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mag_container = srw.magnets.MagnetsContainer([und], lengthPadFraction=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The last input for the simulation is the emitter, instance of \n:py:func:`~pysrw.emitters.ParticleBeam`\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "beam = srw.emitters.ParticleBeam(energy, xPos=0, yPos=0, zPos=-1.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Before computing the wavefront, it is a good practice to obtain and plot \nthe particle trajectory and check that the simulation objects properly \ndefined\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "traj = srw.computeTrajectory(beam, mag_container, sStart=0, sEnd=3)\nsrw.plotTrajectory(traj)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create a 10 mm x 10 mm observation mesh, with a resolution of 100 um, \nplaced 1 m downstream of the source\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "observer = srw.wavefronts.Observer(centerCoord=[0, 0, 10],\n                                   obsXextension=2e-3, \n                                   obsYextension=2e-3,\n                                   obsXres=5e-6,\n                                   obsYres=5e-6)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can finally simulate the wavefront and derive the optical intensity,\nfor example at the fundamental harmonics.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wl = lambda_1\nwfr = srw.computeSrWfrMultiProcess(4, particleBeam=beam, \n                                   magnetsContainer=mag_container,\n                                   observer=observer, wavelength=wl)\nintensity = wfr.getWfrI()\nsrw.plotI(intensity)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The distribution of SR emitted by a particle in a uniform field has an \nanalytical expression. Let's compare it to the simulation result.\nWe extract a vertical slice from the simulation radiation\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "cuts = srw.extractCrossCuts(intensity)\nyax = cuts[\"yax\"]\nint_srw = cuts[\"ydata\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "and compute the theoretical distribution for undulator radiation (tuned at \nfundamental wavelength) as\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "r_p = observer.zPos\ntheta = np.arctan(yax / r_p)\ngamma = beam.nominalParticle.gamma\n\ne_theo = np.sinc((theta**2 / 2 * (N_u * lambda_u) / (wl*1e-9)))\nint_theo = e_theo**2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Simulated and theoretical distributions are finally plotted\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots()\nax.plot(theta*1e3, int_srw / np.max(int_srw))\nax.plot(theta*1e3, int_theo  / np.max(int_theo), label=\"theo\", linestyle=\"--\")\nax.set_xlabel(\"Polar angle [mrad]\")\nax.set_ylabel(\"Intensity [arb. unit]\")\nax.legend()\nplt.show()\nplt.tight_layout()"
      ]
    }
  ],
  "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
}