Adding jones adapter, docs, vortex retarder, jones-to-mueller conversion

docs/source/how-tos/index.rst

@@ -11,3 +11,4 @@ How-Tos
     Advanced-Interferogram-Processing.ipynb
     GPU and Exascale Computing.ipynb
     Differentiable-Optical-Models.ipynb
+    polarized_propagation.ipynb

docs/source/how-tos/polarized_propagation.ipynb

@@ -0,0 +1,217 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Polarized Propagation\n",
+    "This is a brief how-to on how to use prysm's polarized propagation feature. Users should already be familiar with [Jones Calculus](../tutorials/Jones-Calculus.ipynb), and the [First Diffraction Model](../tutorials/First-Diffraction-Model.ipynb) before going through this how-to. \n",
+    "\n",
+    "When we step outside of the classroom and into the laboratory, we discover that things are not always perfect. In an ideal world, a polarizer is a perfect polarizer, and the degree to which it polarizes doesn't change across the optic. In reality, manufacturing defects can complicate our optical system by introducing unwanted effects. In this how-to, we cover how `prysm` can help you model spatially-varying polarization optics in diffraction problems with polarized field propagation.\n",
+    "\n",
+    "We begin with a simple extension of the Jones Matrix $\\mathbf{J}$ into the spatial domain:\n",
+    "\n",
+    "$$\n",
+    "\\mathbf{J}(x,y) =\n",
+    "\\begin{pmatrix}\n",
+    "J_{xx}(x,y) & J_{xy}(x,y) \\\\\n",
+    "J_{yx}(x,y) & J_{yy}(x,y) \\\\\n",
+    "\\end{pmatrix}\n",
+    "$$\n",
+    "\n",
+    "All this means is that we consider $\\mathbf{J}$ to be a function that varies with position across a given optical element. In `prysm`, polarization-modifying components are simply arrays of dimension `(M,N,2,2)`, which allows their effect to vary spatially, as shown for the case of a linear retarder below."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prysm.x.polarization import linear_retarder\n",
+    "import numpy as np\n",
+    "\n",
+    "retarder = (linear_retarder(retardance=np.pi/2,theta=0,shape=[256,256]))\n",
+    "print(retarder.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Any shape `(M,N,2,2)` complex array can be used as a polarization component.  `x/polarization` contains a function that generates a vector vortex retarder (VVR), a component that is like an azimuthally-varying half wave plate. VVRs allow us to do some really interesting things with light. However for the purposes of this demo, we use it simply to illustrate spatially-varying polarized elements with `prysm`. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Initialize a plotting macro\n",
+    "def plot_jones_matrix(J,title='blank title'):\n",
+    "    k = 1\n",
+    "\n",
+    "    plt.figure()\n",
+    "    plt.suptitle(title)\n",
+    "    for i in range(2):\n",
+    "        for j in range(2):\n",
+    "            plt.subplot(2,2,k)\n",
+    "            plt.imshow(J[...,i,j],vmin=-np.pi,vmax=np.pi)\n",
+    "            plt.colorbar()\n",
+    "            k += 1\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prysm.coordinates import make_xy_grid,cart_to_polar\n",
+    "from prysm.x.polarization import vector_vortex_retarder\n",
+    "\n",
+    "# Generate the VVR, a spatially-varying half-wave plate\n",
+    "x, y = make_xy_grid(256, diameter=1)\n",
+    "r, t = cart_to_polar(x, y)\n",
+    "vvr = vector_vortex_retarder(2,t,retardance=np.pi)\n",
+    "plot_jones_matrix(np.real(vvr),title='Vortex Retarder')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we will put this VVR in front of a perfect lens with a circular aperture to see how this spatially-varying retardance affects image formation. However, to make `prysm.propagation` compatible with polarized fields we need to call `add_jones_propagation` from `x/polarization`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prysm.x.polarization import add_jones_propagation\n",
+    "add_jones_propagation()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The propagation functions that presently support polarized propagation are included in the `supported_propagation_funcs` list in `x/polarization`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prysm.x.polarization import supported_propagation_funcs\n",
+    "print('supported propagation functions = ',supported_propagation_funcs)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "`add_jones_propagation` goes through the supported `prysm.propagation` functions and applies a decorator to them to support propagation of `shape` $\\times$ 2 $\\times$ 2 arrays. We can then go and load in a propagation function to examine the PSF. \n",
+    "\n",
+    "Note that because of the shape required for the matrix multiplication that Jones matrices need, we cannot simply multiply the aperture array `A` by the VVR jones matrix `vvr`. To make this easier, we've added the `apply_polarization_to_field` function that extends the dimensions of scalar field arrays to match the Jones matrix arrays so that they support element-wise multiplication."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prysm.propagation import focus_fixed_sampling\n",
+    "from prysm.geometry import circle\n",
+    "from prysm.x.polarization import apply_polarization_optic\n",
+    "\n",
+    "def propagate(wf):\n",
+    "    wfout = focus_fixed_sampling(wf,\n",
+    "                                 input_dx=5e3/256,\n",
+    "                                 prop_dist=50e3,\n",
+    "                                 wavelength=1,\n",
+    "                                 output_dx=10e-1,\n",
+    "                                 output_samples=256)\n",
+    "    return wfout\n",
+    "\n",
+    "x,y = make_xy_grid(256,diameter=1)\n",
+    "r,t = cart_to_polar(x,y)\n",
+    "\n",
+    "# set up a circular aperture and propagate\n",
+    "A = circle(0.5,r)\n",
+    "a_ref = propagate(A)\n",
+    "\n",
+    "#  multiply A by the polarizing element\n",
+    "A = apply_polarization_optic(A, vvr) \n",
+    "j_out = propagate(A)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To visualize this in irradiance, we need to compute the Mueller matrix from the Jones matrix `j_out`. We can use `jones_to_mueller` to do this rapidly. The [0,0] element of the resultant Mueller matrix represents the response of the optical system to unpolarized light. Below we compare the focal plane irradiances for imaging with a circular aperture (left) and imaging with a circular aperture with a vortex phase (right). The phase of the vortex is such that the on-axis irradiance completely cancels."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prysm.x.polarization import jones_to_mueller\n",
+    "\n",
+    "m_out  = jones_to_mueller(j_out, broadcast=True)\n",
+    "intensity_from_scalar = np.abs(a_ref)**2\n",
+    "intensity_from_mueller = m_out[..., 0, 0]\n",
+    "\n",
+    "plt.figure(figsize=[8,4])\n",
+    "plt.subplot(121)\n",
+    "plt.title('Simple Imaging')\n",
+    "plt.imshow(np.log10(intensity_from_scalar))\n",
+    "plt.subplot(122)\n",
+    "plt.title('Imaging with VVR')\n",
+    "plt.imshow(np.log10(intensity_from_mueller))\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In summary, `x/polarization` enables numerical propagation through optical elements that can be represented as a Jones matrix. These elements are arrays of matrices whose row and column indices are in the last two dimensions of the array. This can be applied to problems that involve optical elements like VVRs, which require polarization for a complete description."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "prysmdev",
+   "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.5"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}