{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "05402e6d",
   "metadata": {},
   "source": [
    "# Laser from longitudinal and transverse profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a41f95c",
   "metadata": {},
   "source": [
    "In this tutorial, you will learn to create a Laser from an experimental measurement of the laser spectrum and a 2D image of the transverse intensity profile. We'll start by loading all packages required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcf32943-cf31-42c3-995c-2ed75c9761d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from PIL import Image\n",
    "\n",
    "from lasy.backend import to_cpu, xp\n",
    "from lasy.laser import Laser\n",
    "from lasy.profiles.combined_profile import CombinedLongitudinalTransverseProfile\n",
    "from lasy.profiles.longitudinal import LongitudinalProfileFromData\n",
    "from lasy.profiles.transverse import TransverseProfileFromData\n",
    "from lasy.profiles.transverse.hermite_gaussian_profile import (\n",
    "    HermiteGaussianTransverseProfile,\n",
    ")\n",
    "from lasy.utils.mode_decomposition import hermite_gauss_decomposition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e972d0",
   "metadata": {},
   "source": [
    "Next, let's define the physical parameters defining our laser pulse. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a2fd86-a83f-490e-9e10-21d0e3ca7c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "polarization = (1, 0)  # Linearly polarized in the x direction\n",
    "energy_J = 1  # Pulse energy in Joules\n",
    "cal = 0.2e-6  # Camera pixel size in meters. Used for calibration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bda26c20",
   "metadata": {},
   "source": [
    "In what follows, we do a few steps to create a LASY profile. Profile is the object used in LASY to describe the properties of the pulse. This profile is then used to create a Laser object, i.e., a regular mesh with the values of the vector potential of the pulse. This mesh can be 3D for a 3D (xyt) geometry, or can consist of a few 2D grids for cylindrical geometry with mode decomposition."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f88ab1e4-1f5a-4f07-8b7c-4c105b4aa9eb",
   "metadata": {},
   "source": [
    "## Reconstruct longitudinal profile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6321294",
   "metadata": {},
   "source": [
    "We start from an example dataset obtained with a [Frequency-resolved optical grating (FROG)](https://en.wikipedia.org/wiki/Frequency-resolved_optical_gating) measurement. This provides us with the spectrum and spectral phase."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ee0575d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load from online dataset and store in appropriate variables\n",
    "file_longitudinal = (\n",
    "    \"https://github.com/user-attachments/files/17414077/df_intensity_spectral_v3.csv\"\n",
    ")\n",
    "\n",
    "exp_frequency = xp.loadtxt(file_longitudinal, usecols=0, dtype=\"float\")  # Hz\n",
    "exp_spectrum = xp.loadtxt(file_longitudinal, usecols=1, dtype=\"float\")  # Arbitary units\n",
    "exp_phase = xp.loadtxt(file_longitudinal, usecols=2, dtype=\"float\")  # rad"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b2b6ab4",
   "metadata": {},
   "source": [
    "Now, we initialize a LASY LongitudinalProfile from experimentally measured spectrum. The central wavelength is calculated in this step, and user later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e70f703e-15a3-4751-a857-71b17651671a",
   "metadata": {},
   "outputs": [],
   "source": [
    "longitudinal_data = {\n",
    "    \"datatype\": \"spectral\",\n",
    "    \"axis_is_wavelength\": False,\n",
    "    \"axis\": exp_frequency,\n",
    "    \"intensity\": exp_spectrum,\n",
    "    \"phase\": exp_phase,\n",
    "    \"dt\": 1e-15,\n",
    "}\n",
    "\n",
    "# Create the longitudinal profile. The temporal range is from -200 to +200 femtoseconds\n",
    "longitudinal_profile = LongitudinalProfileFromData(\n",
    "    longitudinal_data, lo=-200e-15, hi=200e-15\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb7c9541",
   "metadata": {},
   "source": [
    "Plot the longitudinal profile data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0085e86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot both the temporal and spectral data\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 4), tight_layout=True)\n",
    "\n",
    "# Spectral data\n",
    "exp_spectrum /= xp.max(exp_spectrum)  # Normalize the spectrum\n",
    "color = \"tab:red\"\n",
    "ax[0].set_xlabel(\"Frequency (Hz)\", fontsize=12)\n",
    "ax[0].set_ylabel(\"Spectral intensity (AU)\", color=color, fontsize=12)\n",
    "ax[0].plot(to_cpu(exp_frequency), to_cpu(exp_spectrum), color=color)\n",
    "ax[0].tick_params(axis=\"y\", labelcolor=color)\n",
    "ax[0].set_title(\"Measured data (Spectral space)\", fontsize=15)\n",
    "ax0 = ax[0].twinx()\n",
    "color = \"tab:blue\"\n",
    "ax0.set_ylabel(\"Spectral phase (radian)\", color=color, fontsize=12)\n",
    "ax0.plot(to_cpu(exp_frequency), to_cpu(exp_phase), color=color)\n",
    "ax0.tick_params(axis=\"y\", labelcolor=color)\n",
    "\n",
    "\n",
    "# Temporal data\n",
    "color = \"tab:red\"\n",
    "ax[1].set_xlabel(\"Time (s)\", fontsize=12)\n",
    "ax[1].set_ylabel(\"Amplitude (AU)\", color=color, fontsize=12)\n",
    "ax[1].plot(\n",
    "    to_cpu(longitudinal_profile.time),\n",
    "    to_cpu(xp.sqrt(longitudinal_profile.temporal_intensity)),\n",
    "    color=color,\n",
    ")\n",
    "ax[1].tick_params(axis=\"y\", labelcolor=color)\n",
    "ax[1].set_title(\"Reconstructed data (Temporal space)\", fontsize=15)\n",
    "ax1 = ax[1].twinx()\n",
    "color = \"tab:blue\"\n",
    "ax1.set_ylabel(\"Temporal phase (radian)\", color=color, fontsize=12)\n",
    "ax1.plot(\n",
    "    to_cpu(longitudinal_profile.time),\n",
    "    to_cpu(longitudinal_profile.temporal_phase),\n",
    "    color=color,\n",
    ")\n",
    "ax1.tick_params(axis=\"y\", labelcolor=color)\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ba9ced4-cc19-48ba-9fec-ddcb6e77128b",
   "metadata": {},
   "source": [
    "## Reconstruct transverse profile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d133d0",
   "metadata": {},
   "source": [
    "For the following reconstruction, the data is provided in the .png image format, which can be read using pillow package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcfff530-c04b-4d28-977c-8143ffe87c62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the transverse profile of the laser pulse, and perform minor cleaning\n",
    "!curl https://user-images.githubusercontent.com/27694869/228038930-d6ab03b1-a726-4b41-a378-5f4a83dc3778.png -o transverse_profile.png\n",
    "file_transverse = \"transverse_profile.png\"\n",
    "img = Image.open(file_transverse)\n",
    "intensity_data = xp.array(img)\n",
    "intensity_scale = xp.max(intensity_data)  # Maximum value of the intensity\n",
    "\n",
    "intensity_data[intensity_data < intensity_scale / 100] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a563b9c0",
   "metadata": {},
   "source": [
    "Create a LASY TransverseProfile from our experimental data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac17f984",
   "metadata": {},
   "outputs": [],
   "source": [
    "nx, ny = intensity_data.shape\n",
    "lo = (0, 0)  # Lower bounds in x and y\n",
    "hi = (ny * cal, nx * cal)  # Upper bounds in x and y\n",
    "\n",
    "# Create the transverse profile. This also centers the data by default\n",
    "transverse_profile = TransverseProfileFromData(\n",
    "    intensity_data, [lo[0], lo[1]], [hi[0], hi[1]]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe9c0b5f",
   "metadata": {},
   "source": [
    "Plot the original transverse profile from the file. The laser occupies a small fraction of the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b0d5eb-ad0a-4777-ae82-d612eba69fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "cax = ax.imshow(\n",
    "    to_cpu(intensity_data),\n",
    "    aspect=\"auto\",\n",
    "    extent=to_cpu(xp.array([lo[0], hi[0], lo[1], hi[1]]) * 1e6),\n",
    ")\n",
    "color_bar = fig.colorbar(cax)\n",
    "color_bar.set_label(r\"Fluence  (J/cm$^2$)\")\n",
    "ax.set_xlabel(\"x ($ \\\\mu m $)\")\n",
    "ax.set_ylabel(\"y ($ \\\\mu m $)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd87054b-ab2c-4591-be03-c6bfeae859a1",
   "metadata": {},
   "source": [
    "## Combine longitudinal and transverse profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab6ecbd2-0b4e-4251-8709-ce741dc4174e",
   "metadata": {},
   "outputs": [],
   "source": [
    "laser_profile_raw = CombinedLongitudinalTransverseProfile(\n",
    "    wavelength=longitudinal_profile.lambda0,\n",
    "    pol=polarization,\n",
    "    laser_energy=energy_J,\n",
    "    long_profile=longitudinal_profile,\n",
    "    trans_profile=transverse_profile,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d55f0ea-a2ea-4157-895f-5d5681403851",
   "metadata": {},
   "source": [
    "## Clean/denoise the profile\n",
    "LASY functions can be used for denoising/cleaning. Here, the measured profile is decomposed into Hermite-Gauss modes. The cleaning is done by keeping only the first few modes. \n",
    "Take a look at the following [example](https://github.com/LASY-org/lasy/blob/13f0e4515493deca36c1375be1d9e83c7e379d42/examples/example_modal_decomposition_data.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82a8a391-c4f1-473f-928e-cc18524e1cec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Maximum Hermite-Gauss mode index in x and y\n",
    "m_max = 10\n",
    "n_max = 10\n",
    "\n",
    "# Calculate the decomposition and waist of the laser pulse\n",
    "modeCoeffs, waistX, waistY = hermite_gauss_decomposition(\n",
    "    transverse_profile,\n",
    "    longitudinal_profile.lambda0,\n",
    "    m_max=m_max,\n",
    "    n_max=n_max,\n",
    "    res=cal,\n",
    "    lo=[-2e-4, -2e-4],\n",
    "    hi=[2e-4, 2e-4],\n",
    ")\n",
    "\n",
    "# Create a num profile, summing over the first few modes\n",
    "energy_frac = 0\n",
    "for i, mode_key in enumerate(list(modeCoeffs)):\n",
    "    tmp_transverse_profile = HermiteGaussianTransverseProfile(\n",
    "        waistX, waistY, mode_key[0], mode_key[1], longitudinal_profile.lambda0\n",
    "    )\n",
    "    energy_frac += modeCoeffs[mode_key] ** 2  # Energy fraction of the mode\n",
    "    if i == 0:  # First mode (0,0)\n",
    "        laser_profile_cleaned = modeCoeffs[\n",
    "            mode_key\n",
    "        ] * CombinedLongitudinalTransverseProfile(\n",
    "            longitudinal_profile.lambda0,\n",
    "            polarization,\n",
    "            longitudinal_profile,\n",
    "            tmp_transverse_profile,\n",
    "            laser_energy=energy_J,\n",
    "        )\n",
    "    else:  # All other modes\n",
    "        laser_profile_cleaned += modeCoeffs[\n",
    "            mode_key\n",
    "        ] * CombinedLongitudinalTransverseProfile(\n",
    "            longitudinal_profile.lambda0,\n",
    "            polarization,\n",
    "            longitudinal_profile,\n",
    "            tmp_transverse_profile,\n",
    "            laser_energy=energy_J,\n",
    "        )\n",
    "\n",
    "# Energy loss due to decomposition\n",
    "energy_loss = 1 - energy_frac\n",
    "print(f\"Energy loss: {energy_loss * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb44ccf-ad2b-4a9b-bc19-9f615d057335",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the original and denoised profiles\n",
    "# Create a grid for plotting\n",
    "x = xp.linspace(-5 * waistX, 5 * waistX, 500)\n",
    "X, Y = xp.meshgrid(x, x)\n",
    "\n",
    "# Determine the figure parameters\n",
    "fig, ax = plt.subplots(1, 3, figsize=(12, 4), tight_layout=True)\n",
    "fig.suptitle(\n",
    "    \"Hermite-Gauss Reconstruction using m_max = %i, n_max = %i\" % (m_max, n_max)\n",
    ")\n",
    "pltextent = xp.array([xp.min(x), xp.max(x), xp.min(x), xp.max(x)]) * 1e6  # in microns\n",
    "cbar_labels = [\"Intensity (norm.)\", \"Intensity (norm.)\", \"|Intensity Error| (%)\"]\n",
    "title = [\"Original Profile\", \"Reconstructed Profile\", \"Error\"]\n",
    "scale = [\"linear\", \"log\"]\n",
    "lw = [1.0, 2.5]\n",
    "\n",
    "\n",
    "# Original profile\n",
    "prof1 = xp.abs(laser_profile_raw.evaluate(X, Y, 0)) ** 2\n",
    "maxInten = xp.max(prof1)\n",
    "prof1 /= maxInten\n",
    "\n",
    "# Reconstructed profile\n",
    "prof2 = xp.abs(laser_profile_cleaned.evaluate(X, Y, 0)) ** 2\n",
    "prof2 /= maxInten\n",
    "\n",
    "# Normalized error\n",
    "prof3 = (prof1 - prof2) / xp.max(prof1)\n",
    "prof = [prof1, prof2, prof3]\n",
    "\n",
    "fig2, ax2 = plt.subplots(2, 1, figsize=(8, 6), tight_layout=True)\n",
    "fig2.suptitle(\n",
    "    \"Hermite-Gauss Reconstruction using m_max = %i, n_max = %i\" % (m_max, n_max)\n",
    ")\n",
    "\n",
    "for i in range(3):\n",
    "    divider = make_axes_locatable(ax[i])\n",
    "    ax_cb = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "    pl = ax[i].imshow(\n",
    "        to_cpu(100 * xp.abs(prof[i])), cmap=\"magma\", extent=to_cpu(pltextent)\n",
    "    )\n",
    "    cbar = fig.colorbar(pl, cax=ax_cb)\n",
    "    cbar.set_label(cbar_labels[i])\n",
    "    ax[i].set_xlim(pltextent[0], pltextent[1])\n",
    "    ax[i].set_ylim(pltextent[2], pltextent[3])\n",
    "    ax[i].set_xlabel(\"x ($ \\\\mu m $)\")\n",
    "    ax[i].set_ylabel(\"y ($ \\\\mu m $)\")\n",
    "    ax[i].set_title(title[i])\n",
    "    if i < 2:\n",
    "        ax2[i].plot(\n",
    "            to_cpu(x * 1e6),\n",
    "            to_cpu(prof2[int(len(x) / 2), :]),\n",
    "            label=title[1],\n",
    "            color=(1, 0.5, 0.5),\n",
    "            lw=2.5,\n",
    "        )\n",
    "        ax2[i].plot(\n",
    "            to_cpu(x * 1e6),\n",
    "            to_cpu(prof1[int(len(x) / 2), :]),\n",
    "            label=title[0],\n",
    "            color=(0.3, 0.3, 0.3),\n",
    "            lw=1.0,\n",
    "        )\n",
    "        ax2[i].legend()\n",
    "        ax2[1].set_yscale(scale[i])\n",
    "        ax2[i].set_xlim(pltextent[0], pltextent[1])\n",
    "        ax2[i].set_xlabel(\"x ($ \\\\mu m $)\")\n",
    "        ax2[i].set_ylabel(\"Intensity (norm.)\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "103c9501-3e91-4f9b-a0bb-d3a5e2320d0c",
   "metadata": {},
   "source": [
    "## Create a laser"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db52cd0e",
   "metadata": {},
   "source": [
    "Now the hard part is done! From the cleaned profile, we create a LASY Laser object (3D cartesian or cylindrical geometry) and write to a file compliant with the openPMD standard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b0f1ffd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, 3D geometry\n",
    "dimensions = \"xyt\"  # Use 3D geometry\n",
    "lo = (-40e-6, -20e-6, -150e-15)  # Lower bounds of the simulation box\n",
    "hi = (40e-6, 20e-6, 150e-15)  # Upper bounds of the simulation box\n",
    "num_points = (\n",
    "    50,\n",
    "    50,\n",
    "    20,\n",
    ")  # Number of points in each dimension. Use (300, 300, 200) for production.\n",
    "\n",
    "# Constructing the object using 3D geometry might take a while to run depending on the hardware used.\n",
    "laser_xyt = Laser(dimensions, lo, hi, num_points, laser_profile_cleaned)  # Laser\n",
    "laser_xyt.normalize(energy_J * energy_frac)  # Normalize the laser energy\n",
    "laser_xyt.show()\n",
    "\n",
    "# Save the laser object to a file\n",
    "laser_xyt.write_to_file(\"Laser_xyt_denoised\", \"h5\", save_as_vector_potential=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07be1792",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Then, cylindrical geometry. Here, we also propagate the pulse backwards by 0.5 mm.\n",
    "dimensions = \"rt\"  # Use cylindrical geometry\n",
    "lo = (0, -150e-15)\n",
    "hi = (40e-6, 150e-15)\n",
    "num_points = (500, 400)\n",
    "\n",
    "laser = Laser(dimensions, lo, hi, num_points, laser_profile_cleaned)\n",
    "laser.normalize(energy_J * energy_frac)\n",
    "laser.propagate(-0.5e-3)  # Propagate the laser pulse backwards\n",
    "laser.show()\n",
    "\n",
    "# Save the laser object to a file\n",
    "laser.write_to_file(\"Laser_rt_propagated\", \"h5\", save_as_vector_potential=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}