{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Custom magnet\n\nExample to show the simulation of a arbitrary magnetic field. A longitudinal \ngradient dipole from ESRF-EBS is implemented using the \n:py:class:`~pysrw.magnets.CustomMagnet`.\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": [
        "A longitudinal gradient dipole is a device where the magnetic field increases\nvaries the propagation. The ideal device should feature a smooth variation of\nthe magnetic field (continuous gradient dipole). This option is often not \npractical and, in real scenarios, it is more convenient to create a series \nof discrete magnetic steps.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "energy = 6      # GeV\nL_eff = 1.784   # total device length\nB_start = 0.65  # the magnetic field at the magnet entrance\nB_end = 0.17    # the magnetic field at the magnet exit\nsegments = 5    # number of steps of the magnetic field"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A 3D map of the field components is required to define the device field.\nSince this particular case involve only a longitudinal change of the field,\nwe can create field arrays with only one bin along the transverse direction.\nThe longitudinal profile is a sequence of decreasing steps. \nWe create a profile with 3 points per segment. This is not the \nfinal resolution of the field as SRW will interpolate the provided field map\non the discretization step of the trajectory computation. Notice also that\nthe actual extension of the magnetic field region is not the final one, \nas the field map will be adjusted to match the physical extension specified \nin the magnet definition\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nx = 1\nny = 1\nnz = segments * 3\n\nB_profile = np.repeat(np.linspace(B_start, B_end, segments), nz // segments)\n\nBz = np.zeros((nx, ny, nz))\nBx = np.zeros((nx, ny, nz))\nBy = np.reshape(B_profile, (nx, ny, nz))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We create the magnetic object. The field map defined above is passed as a \nlist with the three components. The argument `ranges` define the actual \nextension of the magnetic region. In this case, we will get a 1 m x 1m \nuniform field at each transverse plane and the field maps are adjusted to \nfit in a region of length `L_eff`. We also specify the strategy that is later \nused to interpolate the coarse mesh of the field map on the much finer \ndiscretization of the trajectory.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dipole = srw.magnets.CustomMagnet(magField=[Bx, By, Bz], \n                                  ranges=[1,1,L_eff],\n                                  interp=\"bilinear\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The device is wrapped as usual in a container and the beam is created with\nall zero initial conditions at the dipole entrance\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mag_container = srw.magnets.MagnetsContainer([dipole])\nbeam = srw.emitters.ParticleBeam(energy, xPos=0, yPos=0, zPos=-1.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We finally compute the trajectory and plot a detailed view of the resulting\ncustom field profile and trajectory. Note in particular how the user-defined\nfield points are interpolated to create a continous field profile.\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)\n\nfig, ax_mag = plt.subplots()\nax_mag.plot(np.linspace(-.5, .5, nz) * L_eff, B_profile, \n            linewidth=0, marker=\"x\", color=\"red\", label=\"Input field\")\nax_mag.plot(traj[\"z\"], traj[\"By\"], color=\"red\", label=\"Interp. field\")\nax_trj = ax_mag.twinx()\nax_trj.plot(traj[\"z\"], traj[\"x\"] * 1e3, color=\"black\")\nax_mag.set_xlabel(\"z position [m]\")\nax_mag.set_ylabel(\"By field [T]\", color=\"red\")\nax_trj.set_ylabel(\"x position [mm]\", color=\"black\")\nax_mag.legend()\nplt.tight_layout()\nplt.show()"
      ]
    }
  ],
  "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
}