{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bf2e96ee-063d-408f-815b-59c902e0485a",
   "metadata": {},
   "source": [
    "This tutorial requires that the WavefrontPropaGator (WPG) package is installed. \n",
    "\n",
    "Since the WPG package is a requirement of phenom, we encourage you to install WPG if you wish to use the phenom WPG interface.\n",
    "\n",
    "You can check your installation of WPG, and the subsequent requirement of SRW as :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5c4825e7-c0d2-48bd-be93-a5444cd719d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wpg\n",
    "from wpg import srwlpy as srwl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53a091c2-8378-4e3d-b4af-aa0790fa9407",
   "metadata": {},
   "source": [
    "## Interfacing with WavefrontPropaGator (WPG)\n",
    "The phenom.wpg script provides the functionalities require to convert the source wavefront files to the WavefrontPropaGator (WPG) format."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ec7808c-71b8-45e5-a71b-64bb5a5e1b70",
   "metadata": {},
   "source": [
    "Load the source generated in the any of the previous [tutorials](https://twguest.github.io/phenom/examples.html):\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a2c29d5-1c09-46e5-a371-494a8fd41174",
   "metadata": {},
   "source": [
    "### Converting Source Data to a WPG Wavefront\n",
    "The 'wpg_converter' writes each key of the source to a seperate .h5 file that is WPG readable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fec43de8-d166-4497-b258-ad27449f2d65",
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] Unable to synchronously open file (unable to open file: name = './sase_field.h5', errno = 2, error message = 'No such file or directory', flags = 0, o_flags = 0)",
     "output_type": "error",
     "traceback": [
      "Traceback \u001b[0;36m(most recent call last)\u001b[0m:\n",
      "\u001b[0m  Cell \u001b[1;32mIn[2], line 5\u001b[0m\n    wfr = wpg_converter(save_loc, key = key)\u001b[0m\n",
      "\u001b[0m  File \u001b[1;32m~/phenom/phenom/wpg.py:41\u001b[0m in \u001b[1;35mwpg_converter\u001b[0m\n    with h5.File(save_loc, mode = 'r') as hf:\u001b[0m\n",
      "\u001b[0m  File \u001b[1;32m~/miniconda3/envs/mid/lib/python3.12/site-packages/h5py/_hl/files.py:562\u001b[0m in \u001b[1;35m__init__\u001b[0m\n    fid = make_fid(name, mode, userblock_size, fapl, fcpl, swmr=swmr)\u001b[0m\n",
      "\u001b[0m  File \u001b[1;32m~/miniconda3/envs/mid/lib/python3.12/site-packages/h5py/_hl/files.py:235\u001b[0m in \u001b[1;35mmake_fid\u001b[0m\n    fid = h5f.open(name, flags, fapl=fapl)\u001b[0m\n",
      "\u001b[0m  File \u001b[1;32mh5py/_objects.pyx:54\u001b[0m in \u001b[1;35mh5py._objects.with_phil.wrapper\u001b[0m\n",
      "\u001b[0m  File \u001b[1;32mh5py/_objects.pyx:55\u001b[0m in \u001b[1;35mh5py._objects.with_phil.wrapper\u001b[0m\n",
      "\u001b[0;36m  File \u001b[0;32mh5py/h5f.pyx:102\u001b[0;36m in \u001b[0;35mh5py.h5f.open\u001b[0;36m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m\u001b[0;31m:\u001b[0m [Errno 2] Unable to synchronously open file (unable to open file: name = './sase_field.h5', errno = 2, error message = 'No such file or directory', flags = 0, o_flags = 0)\n"
     ]
    }
   ],
   "source": [
    "from phenom.wpg import wpg_converter\n",
    "save_loc = \"./sase_field.h5\" ### master .h5\n",
    "key =  \"pulse000\" ### pulse_id \n",
    "\n",
    "wfr = wpg_converter(save_loc, key = key)\n",
    "wfr.store_hdf5(\"./wpg_sase.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f211093-b243-49e2-a306-20c8b68c1a7f",
   "metadata": {},
   "source": [
    "With the pulse wavefront now WPG readable, we can apply in to wavefront propagation simulations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec0c8f89-62fc-470c-a646-e2d81beba5a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from wpg.beamline import Beamline\n",
    "from wpg.optical_elements import Drift\n",
    "\n",
    "from wpg.wpg_uti_wf import plot_intensity_map\n",
    "\n",
    "wfr.load_hdf5(\"./wpg_sase.h5\")\n",
    "\n",
    "bl = Beamline()\n",
    "bl.append(Drift(10), propagation_parameters = [0, 0, 1, 0, 0, 2, 1.0, 2, 1.0, 0, 0, 0])\n",
    "bl.propagate(wfr)\n",
    "\n",
    "plot_intensity_map(wfr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dba1b6d-a731-433b-9caa-0de3c66a543f",
   "metadata": {},
   "source": [
    "### Converting Multiple Pulses to WPG Wavefronts\n",
    "We can convert all pulses in a source .h5 to WPG readable wavefronts by iterating of the source .h5:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c63a7f-7558-4b67-9b35-5d2f434763e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py as h5\n",
    "from phenom.wpg import wpg_converter\n",
    "\n",
    "with h5.File(save_loc) as hf:\n",
    "    \n",
    "    for key in list(hf.keys()):\n",
    "        wfr = wpg_converter(save_loc, key = key)\n",
    "        wfr.store_hdf5(\"./wpg_sase_{}.h5\".format(key))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e1081d5-ced7-442c-9501-065e20a83fd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from wpg.wpg_uti_wf import get_intensity_on_axis\n",
    "\n",
    "[E, spectrum] = get_intensity_on_axis(wfr).T\n",
    "\n",
    "fig, ax = plt.subplots(1,1,figsize = (5,5))\n",
    "                       \n",
    "ax.plot(E, spectrum)\n",
    "ax.set_xlabel(\"Energy (keV)\")\n",
    "ax.set_ylabel(\"Intensity (a.u.)\")                       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d47458f-1047-4e37-9e4c-ef0535213464",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (Spyder)",
   "language": "python3",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}