{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9296ccf6-4e83-467f-afe5-a74e6820e110",
   "metadata": {},
   "source": [
    "# Collins propagator and focal scan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac770df8-1c85-43dc-87f8-f3cab5745ced",
   "metadata": {},
   "source": [
    "In this tutorial, we simulate the propagation of a focusing continuous-wave (CW) beam in free-space. This gives an overview of the functionality of the `CollinsSFFTPropagator` class which defines a diffraction integral expressed in terms of an optical ray matrix. The optical ray matrix is defined in the `ABCD` class. \n",
    "\n",
    "The `CollinsSFFTPropagator` is based on a single fast-Fourier transform (SFFT) and is used to go between the near-field (NF) and the far-field (FF) in imaging systems.\n",
    "\n",
    "This tutorial gives an overview of how to use `ABCD` to define the optical ray matrices for basic components such as thin-lenses and drift sections in vacuum. A Gaussian continous wave (CW) beam is propagated to the first plane of the focal scan using the `CollinsSFFTPropagator`, using a lens and free-space drift defined with `ABCD` matrices. This is repeated for subsequent planes in the focal scan.\n",
    "\n",
    "First we import a few relevant modules: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90cb0c07-1378-41c3-bb98-d2dde6a0538c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from scipy.constants import c, epsilon_0\n",
    "\n",
    "from lasy.backend import to_cpu, xp\n",
    "from lasy.laser import Laser\n",
    "from lasy.profiles import CombinedLongitudinalTransverseProfile\n",
    "from lasy.profiles.longitudinal import ContinuousWaveProfile\n",
    "from lasy.profiles.transverse import GaussianTransverseProfile\n",
    "from lasy.propagators import ABCD, CollinsSFFTPropagator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c64938e-7312-4b4c-87df-fc06c363351f",
   "metadata": {},
   "source": [
    "Next, we define the beam to be CW with a tophat profile in the transverse plane:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e8bb7fe-a3e0-4409-9574-1a3ec09406aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "peak_fluence = 1.0e4  # J/m^2\n",
    "spot_size = 10e-3\n",
    "wavelength = 800e-9\n",
    "omega0 = 2 * xp.pi * c / wavelength\n",
    "pol = (1, 0)\n",
    "\n",
    "long_prof = ContinuousWaveProfile(wavelength)\n",
    "tran_prof = GaussianTransverseProfile(spot_size)\n",
    "laser_profile = CombinedLongitudinalTransverseProfile(\n",
    "    wavelength, pol, long_prof, tran_prof, peak_fluence=peak_fluence\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c88a7ef-5a9c-4556-a3f7-0ed9cbddb4d3",
   "metadata": {},
   "source": [
    "We now define the grids at the input plane and the full laser object in terms of the grid and the laser:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc731f4b-cd52-4437-84d6-f115c9b0e1ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "dimensions = \"xyt\"  # Use cylindrical geometry\n",
    "lo = (-15.0 * spot_size, -15.0 * spot_size, None)  # Lower bounds of the simulation box\n",
    "hi = (15.0 * spot_size, 15.0 * spot_size, None)  # Upper bounds of the simulation box\n",
    "num_points = (256, 256, 1)  # Number of points in each dimension\n",
    "\n",
    "laser = Laser(dimensions, lo, hi, num_points, laser_profile)\n",
    "laser.show(envelope_type=\"intensity\")  # In this case this is the fluence (CW beam)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "097cee34-cb52-4f73-a98d-a70801e344c1",
   "metadata": {},
   "source": [
    "We will now set up the focal scan by estimating the Rayleigh length and using this to define the axial grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9608070-7cf7-4272-b5bf-bc64c0bf9472",
   "metadata": {},
   "outputs": [],
   "source": [
    "focal_length = 300e-3\n",
    "z_R = focal_length**2 / (\n",
    "    xp.pi * spot_size**2 / wavelength\n",
    ")  # The estimated Rayleigh length after the lens\n",
    "spot_size_focus = (\n",
    "    laser.profile.lambda0 * focal_length / (xp.pi * spot_size)\n",
    ")  # Estimated focal spot-size after the lens\n",
    "limits = [-10.0 * z_R, 10.0 * z_R]  # Set the depth limits for the focal scan\n",
    "N_points = 200\n",
    "\n",
    "# Estimate of fluence at input plane and focus (assuming a Gaussian pulse)\n",
    "E_0 = peak_fluence * xp.pi * spot_size**2 / 2\n",
    "print(E_0)\n",
    "F_in = 2.0 * E_0 / (xp.pi * (spot_size) ** 2) / 1e4\n",
    "F_out = 2.0 * E_0 / (xp.pi * (spot_size_focus) ** 2) / 1e4\n",
    "print(\n",
    "    \"Fluence at input plane [J cm-2]: %0.1e\" % (F_in),\n",
    "    \"\\nFluence at focus [J cm-2]: %0.1e\" % (F_out),\n",
    ")\n",
    "\n",
    "z_grid = (\n",
    "    xp.linspace(limits[0], limits[1], N_points) + focal_length\n",
    ")  # Absolute axial distances from the input plane of the lens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd257fa-9f45-4ab1-b338-cf84ccc52119",
   "metadata": {},
   "source": [
    "We will now make a copy of the laser and define both the SFFT propagator that we will use in our focal scan. We then initialise the ABCD matrix and add a thin-lens of the correct focal length:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f04bee0-c1ce-482c-b1a9-22899e83ab8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "prop = CollinsSFFTPropagator(dimensions, omega0)\n",
    "abcd = ABCD()\n",
    "abcd.add_lens(focal_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27485f78-fabb-414d-b372-b1c4fb87d5bb",
   "metadata": {},
   "source": [
    "Now we will perform the focal scan (see comments within the loop for an explanation of intermediate steps):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "434e31be-90a8-4cd9-a55c-cf96c15f79a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "focal_scan = []  # Create empty list to store the focal scan data\n",
    "\n",
    "for i, z in enumerate(z_grid):\n",
    "    laser_input = copy.deepcopy(laser)\n",
    "    if i == 0:\n",
    "        abcd.add_vacuum(z)\n",
    "        prop.propagate(laser_input.grid, abcd)\n",
    "        grid_out = laser_input.grid\n",
    "    else:\n",
    "        abcd.add_vacuum(z - z_grid[i - 1])\n",
    "        prop.propagate(laser_input.grid, abcd, grid_out=grid_out)\n",
    "\n",
    "    focal_scan.append(laser_input)  # Storing the full focal scan data\n",
    "\n",
    "    field = laser_input.grid.get_temporal_field()[\n",
    "        num_points[0] // 2, :, num_points[2] // 2\n",
    "    ]\n",
    "    lineout = epsilon_0 * c / 2 * xp.abs(field) ** 2 / 1e4\n",
    "\n",
    "    # Stack these lineouts to plot a cross-section of the focal scan\n",
    "    if i == 0:\n",
    "        scanImg = lineout\n",
    "    else:\n",
    "        scanImg = xp.vstack((scanImg, lineout))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a464fdbc-a8bd-46c7-ad37-12fbbe33395a",
   "metadata": {},
   "source": [
    "We will now get the output grids from the simulation and plot a cross-section of the focal scan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8988b1e-b07c-4db9-955e-209c82375733",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the grids from the simulation\n",
    "x_grid = focal_scan[0].grid.axes[\n",
    "    0\n",
    "]  # Output axes as automatically calculated by the CollinsSFFTPropagator\n",
    "y_grid = focal_scan[0].grid.axes[1]\n",
    "extent = [z_grid[0] * 1e3, z_grid[-1] * 1e3, y_grid[0] * 1e6, y_grid[-1] * 1e6]\n",
    "\n",
    "# Make plot of focal scan cross-section\n",
    "fig, ax = plt.subplots(1, 1, figsize=(7, 5), tight_layout=True)\n",
    "ax_divider = make_axes_locatable(ax)\n",
    "im = ax.imshow(\n",
    "    to_cpu(scanImg.T), aspect=\"auto\", interpolation=\"none\", cmap=\"Reds\", extent=extent\n",
    ")\n",
    "cax = ax_divider.append_axes(\"right\", size=\"3%\", pad=\"2%\")\n",
    "cb = fig.colorbar(im, cax=cax)\n",
    "cb.set_label(r\"$|E|$ (V/m)\")\n",
    "ax.set_xlabel(\"z (mm)\")\n",
    "ax.set_ylabel(r\"r ($\\mu m$)\")\n",
    "ax.set_ylim(-30, 30)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "044aab5c-0971-4ce7-85e1-96962cacfe91",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}