make_dashboard.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "706e8812",
   "metadata": {},
   "source": [
    "# Dashboard creation routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "20ba5219",
   "metadata": {},
   "outputs": [],
   "source": [
    "from impact import Impact\n",
    "from distgen import Generator\n",
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib as mpl\n",
    "#mpl.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('dark_background')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0becfee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageOps, ImageEnhance \n",
    "\n",
    "def fig2data ( fig ):\n",
    "    \"\"\"\n",
    "    @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it\n",
    "    @param fig a matplotlib figure\n",
    "    @return a numpy 3D array of RGBA values\n",
    "    \"\"\"\n",
    "    # draw the renderer\n",
    "    fig.canvas.draw ( )\n",
    "    # Get the RGBA buffer from the figure\n",
    "    w,h = fig.canvas.get_width_height()\n",
    "    buf = np.frombuffer ( fig.canvas.tostring_argb(), dtype=np.uint8 )\n",
    "    buf.shape = ( w, h, 4 )\n",
    "    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode\n",
    "    buf = np.roll ( buf, 3, axis = 2 )\n",
    "    return buf\n",
    "\n",
    "def fig2img ( fig ):\n",
    "    \"\"\"\n",
    "    @brief Convert a Matplotlib figure to a PIL Image in RGBA format and return it\n",
    "    @param fig a matplotlib figure\n",
    "    @return a Python Imaging Library ( PIL ) image\n",
    "    \"\"\"\n",
    "    # put the figure pixmap into a numpy array\n",
    "    buf = fig2data ( fig )\n",
    "    w, h, d = buf.shape\n",
    "    return Image.frombytes( \"RGBA\", ( w ,h ), buf.tobytes( ) )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6b78176f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def iscreen(impact_object, screen='OTR2', k1='x', k2='y', dpi=72, title=None):\n",
    "    fig = impact_object.particles[screen].plot(k1, k2, return_figure=True, figsize=(5,4))\n",
    "    fig.dpi=dpi\n",
    "    \n",
    "    if not title:\n",
    "        title = screen\n",
    "    fig.axes[2].set_title(title)\n",
    "    fig.tight_layout()\n",
    "    return fig2img(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f7f63db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def info_str(impact_object, name=''):\n",
    "    I = impact_object\n",
    "    H = impact_object.header\n",
    "    P = I.particles['initial_particles']\n",
    "    \n",
    "    timestep = I.ele['change_timestep_1']\n",
    "    dt1 = timestep['dt']\n",
    "    s1 = timestep['s']\n",
    "    run_time = I.output['run_info']['run_time'] \n",
    "\n",
    "    \n",
    "    summary=f\"\"\"{name} \n",
    "    \n",
    "LUME-Impact running Impact-T \n",
    "Distgen created particles at the cathode\n",
    "Particles in openPMD-beamphysics format \n",
    "\n",
    "{H['Np']:,} macroparticles\n",
    "{H['Nbunch']} bunch of {I['species']}s\n",
    "total charge: {I['total_charge']*1e12:.1f} pC\n",
    "Processor domain: {H['Nprow']} x {H['Npcol']} = {H['Nprow']*H['Npcol']} CPUs\n",
    "Space charge grid: {H['Nx']} x {H['Ny']} x {H['Nz']}\n",
    "Timestep: {H['Dt']*1e12} ps to {s1} m, \n",
    "          then {dt1*1e12} ps until the end\n",
    "                  \n",
    "Run time: {run_time/60:.1f} min\n",
    "\n",
    "\"\"\"\n",
    "    return summary\n",
    "\n",
    "def itext(impact_object, dpi=72, name=''):\n",
    "    text = info_str(impact_object, name=name)\n",
    "    fig, ax = plt.subplots(figsize=(5,4))\n",
    "    fig.dpi=dpi\n",
    "    fig.tight_layout()\n",
    "    ax.set_axis_off()\n",
    "    ax.text(0.1, 0.5, text, fontsize=13, horizontalalignment='left', verticalalignment='center', transform=ax.transAxes)\n",
    "    return fig2img(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7897c117",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_dashboard(impact_object=None,\n",
    "                   dat=None,\n",
    "                   itime=None,\n",
    "                   outpath='test/',\n",
    "                   screen1='YAG02',\n",
    "                   screen2='YAG03',\n",
    "                   screen3='OTR2',\n",
    "                   ylim=(0,2e-6),\n",
    "                   ylim2=(0,2e-3),\n",
    "                   name='lume-impact-live-demo'\n",
    "                  ):\n",
    "    \"\"\"\n",
    "    Makes a composite dashboard image from data dict\n",
    "    \n",
    "    Returns the path to the figure written\n",
    "    \"\"\"\n",
    "    if impact_object:\n",
    "        I = impact_object   \n",
    "    else:\n",
    "        itime = dat['isotime']\n",
    "        I = Impact.from_archive(dat['outputs']['archive'])\n",
    "        #G = Generator()\n",
    "        #G.load_archive(dat['archive'])\n",
    "    #return I # Debug\n",
    "    \n",
    "    run_time = I.output['run_info']['run_time'] \n",
    "    # Main figure\n",
    "    FIG0 = I.plot(['norm_emit_x','norm_emit_y'], \n",
    "       y2=['sigma_x', 'sigma_y', 'sigma_z'],\n",
    "       \n",
    "       ylim=ylim, ylim2=ylim2,\n",
    "        figsize=(16,8), return_figure=True)\n",
    "    \n",
    "    \n",
    "    n_particle = I.particles['final_particles'].n_particle\n",
    "    \n",
    "    title=f'Acquired settings at {itime}, simulation run time: {run_time/60:5.1f} min'\n",
    "    \n",
    "    FIG0.tight_layout()\n",
    "    FIG0.axes[0].set_title(title)\n",
    "    \n",
    "    DPI = 150 # test\n",
    "    FIG0.dpi=DPI\n",
    "    im0 = fig2img(FIG0)\n",
    "    \n",
    "    \n",
    "    # For short debugging runs\n",
    "    if screen1 not in I.particles:\n",
    "        screen1='initial_particles'\n",
    "        screen2='initial_particles'\n",
    "        screen3='final_particles'\n",
    "\n",
    "        \n",
    "    # info text\n",
    "    #imtext = ImageOps.invert(itext(I, dpi=DPI).convert('RGB'))     \n",
    "    imtext =itext(I, dpi=DPI, name=name)\n",
    "\n",
    "    im1 = iscreen(I, screen=screen1, k1='x', k2='y', dpi=DPI)\n",
    "    im2 = iscreen(I, screen=screen2, k1='x', k2='y', dpi=DPI)\n",
    "    #im3 = iscreen(I, screen=screen3, k1='x', k2='y', dpi=DPI)\n",
    "    im3 = imtext\n",
    "    im4 = iscreen(I, screen=screen3, k1='delta_z', k2='delta_energy', dpi=DPI)\n",
    "    im5 = iscreen(I, screen=screen3, k1='x', k2='y', dpi=DPI)\n",
    "    \n",
    "    im99 = iscreen(I, screen='initial_particles', k1='x', k2='y', dpi=DPI, title='cathode')\n",
    "    \n",
    "    SIZE =  (im0.width + im1.width, im1.height+im2.height+im3.height)\n",
    "    ii = Image.new('RGB', SIZE)\n",
    "    \n",
    "    invim0 = ImageOps.invert(im0.convert('RGB'))\n",
    "    ii.paste(im0, (0, 10))\n",
    "    \n",
    "    ii.paste(im99, (0, im0.height))\n",
    "    ii.paste(im1, (im99.width, im0.height))\n",
    "    ii.paste(im2, (im99.width+im1.width,im0.height))\n",
    "    ii.paste(im3, (im0.width,0))\n",
    "    ii.paste(im4, (im0.width,im3.height))\n",
    "    ii.paste(im5, (im0.width,im4.height+im3.height))\n",
    "    \n",
    "    fname = f'{name}-{itime}-dashboard.png'\n",
    "    fout = os.path.join(outpath, fname)\n",
    "    \n",
    "    # Enhance contrast\n",
    "    #enhancer = ImageEnhance.Brightness(ii) \n",
    "    enhancer = ImageEnhance.Contrast(ii) \n",
    "    iout = enhancer.enhance(1.2)\n",
    "    iout.save(fout)\n",
    "    \n",
    "    return fout\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "063ba356",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#%%capture\n",
    "#I0 = make_dashboard(dat=json.load(open('output/lume-impact-live-demo-2021-04-05T19:13:18-07:00.json')))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}