day4_genesis_input_output.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![alt text](images/uspas.png)\n",
    "# VUV and X-ray Free Electron Lasers\n",
    "# Running Genesis with lume-genesis\n",
    "#### In this session, we will use lume-genesis to do the following:\n",
    "- load input files\n",
    "- change input files\n",
    "- run a genesis simulation\n",
    "- load output results\n",
    "- plot output results\n",
    "\n",
    "Some comments and parameter descriptions taken from the Genesis manual: http://genesis.web.psi.ch/download.html   \n",
    "Parts of this notebook taken from lume-genesis examples here: https://github.com/slaclab/lume-genesis/tree/master/examples  \n",
    "\n",
    "##### Instructors: D. Nguyen, P. Anisimov, N. Neveu\n",
    "##### Teaching Assistant: Y.S. Li\n",
    "----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing python packages we will use\n",
    "import os, h5py\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# from lume-genesis\n",
    "from genesis import Genesis\n",
    "from genesis import parsers, lattice\n",
    "from genesis.parsers import parse_beam_file\n",
    "from genesis.writers import write_beam_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Make a run directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.getcwd() # your simulation will run here now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lume-genesis will create and use a temporary directory unless a working directory is specified.\n",
    "# I prefer to make a new run directory in the current working directory\n",
    "rundir = os.getcwd()+'/test_run' \n",
    "if not os.path.exists(rundir):\n",
    "    os.makedirs(rundir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Initialize lume-genesis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the run, make genesis object with some template input file\n",
    "gen = Genesis('/home/vagrant/jupyter/FEL2021/examples/lcls_sase.in', verbose=True, workdir=rundir, use_tempdir=False)\n",
    "# If use_tempdir=True, your files will run in a ~/tmp folder that is not in or visible from the working directory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Genesis() call above loaded a premade genesis input and lattice file.   \n",
    "After running this command, we can now look at what is in the input file and make changes in Python.   \n",
    "All values in 'gen.input' are parameters in the input file. \n",
    "Take a look at the information from the input file in the following cells:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Look at loaded input with Python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# All inputs in Python are fields from the genesis input file\n",
    "# if no value in input file template, default values are filled in\n",
    "gen.input.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.beam # empty because no beam is loaded \n",
    "# you can supply genesis an initial beam distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.param.keys()\n",
    "\n",
    "# Pick a few parameters each, and find them in the manual.\n",
    "# Summarize what they are used for. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, look at the lattice info.   \n",
    "\n",
    "There are two keys in the lattice information: \n",
    "- eles  = elements\n",
    "- param = parameters "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list of elements and parameters\n",
    "# s = end of elements, final positions\n",
    "# lume-genesis fills in empty spaces for overlap\n",
    "gen.input['lattice'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.lattice['param'] # lattice information"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genesis lattice files consist of components in the beamline (from manual):\n",
    "\n",
    "- AW - Main magnetic field (undulator)\n",
    "- AD - Drift section\n",
    "- QF - Quadrupole strength\n",
    "- QX - Quadrupole offset in x\n",
    "- QY - Quadrupole offset in y\n",
    "- SL - Solenoid strength\n",
    "- CX - Corrector strength in x\n",
    "- CY - Corrector strength in y\n",
    "\n",
    "\n",
    "Information included in the element descriptions: \n",
    "- strength = magnet strength\n",
    "- L = length of the structure in measure of the unit length\n",
    "- d = distance to previous element\n",
    "- s = distance to the previous element of the same kind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.lattice['eles']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Output settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genesis can output many types of data. There are several flags for data output\n",
    "\n",
    "- itdp    - time dependant simulation \n",
    "- idump   - wavefront output \n",
    "- idmpfld - turn on field output\n",
    "- idmppar - turn on particle distribution otuput\n",
    "- ippart  - \"Write the particle distribution to file at each IPPARTth integration step\"\n",
    "- ipradi  - \"Write the field distribution to file at each IPRADIth integration step\"\n",
    "\n",
    "**Note** for list above, parameter = 1 is on, and = 0 is off, unless otherwise noted in manual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dump wavefront output\n",
    "# set to 0 for no wavefront output\n",
    "gen.param['idump'] = 1\n",
    "\n",
    "# Set to 0, for time independant simulation\n",
    "# This is a good starting point to check if things look reasonable\n",
    "gen.param['itdp']  = 0\n",
    "\n",
    "# # Turn on history\n",
    "gen.param['ippart'] = 10\n",
    "gen.param['ipradi'] = 10\n",
    "\n",
    "# Change number of slices \n",
    "# gen['nslice'] = 1\n",
    "\n",
    "# # Turn on/off field output\n",
    "gen['idmpfld'] = 1\n",
    "\n",
    "# # Turn on particle output\n",
    "# gen['idmppar'] = 1\n",
    "\n",
    "# Change number of particles\n",
    "# gen['npart'] = 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure the runs with changes made above\n",
    "gen.configure()\n",
    "# gen.write_input()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Run lcls genesis file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.output['run_info'] # gives info about run time and location"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The run above is by default single core.   \n",
    "You can run also run using mpi on up to 10 cores: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!which genesis2-mpi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # This command will run the genesis simulation \n",
    "# gen.binary_prefixes = ['mpirun', '-n', '2']\n",
    "# gen.configure()\n",
    "# gen.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gen.output['run_info']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Archiving the data\n",
    "\n",
    "Saving the data to an h5 file allows you to reload with lume-genesis, and save most data in one file.   \n",
    "You can reload your data after closing jupyterlab, and don't need to leave the browser window open.\n",
    "\n",
    "Some notes on output files:  \n",
    "fld - field history file (can get large)  \n",
    "par - can get very large if many slices  \n",
    "dlf (wavefront), dpa (phase space) are final field and particle files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Archive data to h5 so that you can reload it later\n",
    "gen.archive('test_run/lcls_sase.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This will only work if 'idumpfld' option is non-zero\n",
    "gen.write_wavefront('test_run/lcls_sase_wavefront.h5') # save wavefront data to seperate h5 file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7. Loading data from archive file & looking at output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.load_archive('test_run/lcls_sase.h5')\n",
    "# Output data is now saved in gen object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.output.keys() # three categories of saved data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.output['run_info'] #information about simulation run time and location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen.output['param'].keys() # input parameters to simulation\n",
    "# should be similar or same as input params\n",
    "# some numbers filled in during run with defaults"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are the available data types\n",
    "gen.output['data'].keys() # all the output data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 8. Plotting results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get z values in sim\n",
    "zlist = gen.output['data']['z'] # 1D array\n",
    "zlist.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get power. This is a 2d array of: slice, z\n",
    "power = gen.output['data']['power']\n",
    "power.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(zlist, power[0]/1e9);\n",
    "plt.xlabel('Z (m)'); plt.ylabel('Power (GW)');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(zlist,np.log(power[0]))\n",
    "plt.xlabel('z along undulator')\n",
    "plt.ylabel('Power (W)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the wavefront file from above\n",
    "wf  = h5py.File('test_run/lcls_sase_wavefront.h5', 'r') \n",
    "dfl = wf['data/000000/meshes/electricField/x']\n",
    "wf['data/000000/meshes/electricField/x']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfl.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get parameters from .out file\n",
    "params = gen.output['param']\n",
    "my_ncar = params['ncar']\n",
    "my_dgrid = params['dgrid']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_ncar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_dgrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Field phase at end, slice 0\n",
    "def plot_field(dat, dgrid):\n",
    "    ndat = np.angle(dat)\n",
    "    plt.imshow(ndat, extent = [1000*dgrid*i for i in [-1,1,-1,1]])\n",
    "    plt.xlabel('x (mm)')\n",
    "    plt.ylabel('y (mm)')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_field(dfl[:, :, 0], my_dgrid )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you run w/ time dependance, you can plot field at multiple points\n",
    "# plot_field(dfl[:, :, 100], my_dgrid )\n",
    "# plot_field(dfl[:, :, 200], my_dgrid )\n",
    "# plot_field(dfl[:, :, 300], my_dgrid )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "py3"
  },
  "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}